163 lines
4.6 KiB
Python
163 lines
4.6 KiB
Python
import argparse
|
|
import yaml
|
|
from dotenv import load_dotenv
|
|
from tqdm import tqdm
|
|
from test_data import RESOURCE_TYPES, DB_SCHEMA
|
|
|
|
load_dotenv()
|
|
|
|
import openai
|
|
|
|
import re
|
|
|
|
|
|
from typing import TypedDict, Tuple
|
|
|
|
|
|
class Literal(str):
|
|
pass
|
|
|
|
|
|
def literal_presenter(dumper, data):
|
|
return dumper.represent_scalar("tag:yaml.org,2002:str", data, style="|")
|
|
|
|
|
|
yaml.add_representer(Literal, literal_presenter)
|
|
|
|
|
|
class Prompt(TypedDict):
|
|
prompt: str
|
|
|
|
|
|
class PromptsConfig(TypedDict):
|
|
prompts: dict[str, Prompt]
|
|
system: str
|
|
|
|
|
|
class Query(TypedDict):
|
|
description: str
|
|
type: str
|
|
lang: str
|
|
code: str
|
|
error: str
|
|
|
|
|
|
def get_prompts(
|
|
prompts_path: str,
|
|
) -> Tuple[PromptsConfig, PromptsConfig, PromptsConfig]:
|
|
GEN_CONFIG = None
|
|
EDIT_CONFIG = None
|
|
FIX_CONFIG = None
|
|
with open(prompts_path + "/gen.yaml") as f:
|
|
GEN_CONFIG = yaml.safe_load(f)
|
|
with open(prompts_path + "/edit.yaml") as f:
|
|
EDIT_CONFIG = yaml.safe_load(f)
|
|
with open(prompts_path + "/fix.yaml") as f:
|
|
FIX_CONFIG = yaml.safe_load(f)
|
|
return GEN_CONFIG, EDIT_CONFIG, FIX_CONFIG
|
|
|
|
|
|
def get_queries(queries_path: str) -> list[Query]:
|
|
with open(queries_path) as f:
|
|
return yaml.safe_load(f)
|
|
|
|
|
|
def prepare_prompt(
|
|
query: Query,
|
|
GEN_CONFIG: PromptsConfig,
|
|
EDIT_CONFIG: PromptsConfig,
|
|
FIX_CONFIG: PromptsConfig,
|
|
):
|
|
system = ""
|
|
prompt = ""
|
|
template_prompt = ""
|
|
if query["type"] == "gen":
|
|
system = GEN_CONFIG["system"]
|
|
template_prompt = GEN_CONFIG["prompts"][query["lang"]]["prompt"]
|
|
prompt = template_prompt.replace("{description}", query["description"])
|
|
elif query["type"] == "edit":
|
|
system = EDIT_CONFIG["system"]
|
|
template_prompt = EDIT_CONFIG["prompts"][query["lang"]]["prompt"]
|
|
prompt = template_prompt.replace("{description}", query["description"]).replace(
|
|
"{code}", query["code"]
|
|
)
|
|
elif query["type"] == "fix":
|
|
system = FIX_CONFIG["system"]
|
|
template_prompt = FIX_CONFIG["prompts"][query["lang"]]["prompt"]
|
|
prompt = template_prompt.replace("{error}", query["error"]).replace(
|
|
"{code}", query["code"]
|
|
)
|
|
|
|
if query["lang"] in ["deno", "bun", "nativets"]:
|
|
prompt = prompt.replace("{resourceTypes}", RESOURCE_TYPES["typescript"])
|
|
elif query["lang"] == "python3":
|
|
prompt = prompt.replace("{resourceTypes}", RESOURCE_TYPES["python"])
|
|
elif query["lang"] == "php":
|
|
prompt = prompt.replace("{resourceTypes}", RESOURCE_TYPES["php"])
|
|
|
|
if query["lang"] in ["postgresql"]:
|
|
prompt = (
|
|
prompt
|
|
+ "\nHere's the database schema, each column is in the format [name, type, required, default?]: <dbschema>\n"
|
|
+ DB_SCHEMA
|
|
+ "\n</dbschema>"
|
|
)
|
|
|
|
return system, prompt, template_prompt
|
|
|
|
|
|
def format_literal(answer: str):
|
|
return re.sub("[^\\S\n]+\n", "\n", answer).replace("\t", " ")
|
|
|
|
|
|
def gen_samples(queries_path: str, answers_path: str, prompts_path: str):
|
|
GEN_CONFIG, EDIT_CONFIG, FIX_CONFIG = get_prompts(prompts_path)
|
|
|
|
queries = get_queries(queries_path)
|
|
|
|
answers = []
|
|
|
|
for query in tqdm(queries):
|
|
(system, prompt, template_prompt) = prepare_prompt(
|
|
query, GEN_CONFIG, EDIT_CONFIG, FIX_CONFIG
|
|
)
|
|
client = openai.OpenAI()
|
|
chat_completion = client.chat.completions.create(
|
|
model="gpt-4o-2024-05-13",
|
|
messages=[
|
|
{"role": "system", "content": system},
|
|
{"role": "user", "content": prompt},
|
|
],
|
|
temperature=0,
|
|
max_tokens=2048,
|
|
seed=42,
|
|
)
|
|
|
|
answer = {
|
|
**query,
|
|
"answer": Literal(format_literal(chat_completion.choices[0].message.content)), # type: ignore
|
|
"template_system": Literal(format_literal(system)),
|
|
"template_prompt": Literal(format_literal(template_prompt)),
|
|
}
|
|
|
|
if "code" in query:
|
|
answer["code"] = Literal(format_literal(query["code"]))
|
|
|
|
answers.append(answer)
|
|
|
|
with open(answers_path, "w") as f:
|
|
yaml.dump(answers, f)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(description="Process some integers.")
|
|
parser.add_argument("--queries_path", type=str, default="./sample_queries.yaml")
|
|
parser.add_argument("--answers_path", type=str, default="./sample_answers.yaml")
|
|
parser.add_argument(
|
|
"--prompts_path",
|
|
type=str,
|
|
default="../frontend/src/lib/components/copilot/prompts",
|
|
)
|
|
args = parser.parse_args()
|
|
gen_samples(args.queries_path, args.answers_path, args.prompts_path)
|