Add types

This commit is contained in:
Anton Osika
2023-06-25 11:05:43 +02:00
parent 2a39cc4cd7
commit 22051774d5

View File

@@ -12,16 +12,13 @@ from gpt_engineer.chat_to_files import to_files
from gpt_engineer.db import DBs from gpt_engineer.db import DBs
def setup_sys_prompt(dbs): def setup_sys_prompt(dbs: DBs) -> str:
return ( return (
dbs.preprompts["generate"] + "\nUseful to know:\n" + dbs.preprompts["philosophy"] dbs.preprompts["generate"] + "\nUseful to know:\n" + dbs.preprompts["philosophy"]
) )
Step = TypeVar("Step", bound=Callable[[AI, DBs], List[dict]]) def get_prompt(dbs: DBs) -> str:
def get_prompt(dbs):
"""While we migrate we have this fallback getter""" """While we migrate we have this fallback getter"""
assert ( assert (
"prompt" in dbs.input or "main_prompt" in dbs.input "prompt" in dbs.input or "main_prompt" in dbs.input
@@ -37,14 +34,18 @@ def get_prompt(dbs):
return dbs.input["prompt"] return dbs.input["prompt"]
def simple_gen(ai: AI, dbs: DBs): # All steps below have this signature
Step = TypeVar("Step", bound=Callable[[AI, DBs], List[dict]])
def simple_gen(ai: AI, dbs: DBs) -> List[dict]:
"""Run the AI on the main prompt and save the results""" """Run the AI on the main prompt and save the results"""
messages = ai.start(setup_sys_prompt(dbs), get_prompt(dbs)) messages = ai.start(setup_sys_prompt(dbs), get_prompt(dbs))
to_files(messages[-1]["content"], dbs.workspace) to_files(messages[-1]["content"], dbs.workspace)
return messages return messages
def clarify(ai: AI, dbs: DBs): def clarify(ai: AI, dbs: DBs) -> List[dict]:
""" """
Ask the user if they want to clarify anything and save the results to the workspace Ask the user if they want to clarify anything and save the results to the workspace
""" """
@@ -83,7 +84,7 @@ def clarify(ai: AI, dbs: DBs):
return messages return messages
def gen_spec(ai: AI, dbs: DBs): def gen_spec(ai: AI, dbs: DBs) -> List[dict]:
""" """
Generate a spec from the main prompt + clarifications and save the results to Generate a spec from the main prompt + clarifications and save the results to
the workspace the workspace
@@ -100,7 +101,7 @@ def gen_spec(ai: AI, dbs: DBs):
return messages return messages
def respec(ai: AI, dbs: DBs): def respec(ai: AI, dbs: DBs) -> List[dict]:
messages = json.loads(dbs.logs[gen_spec.__name__]) messages = json.loads(dbs.logs[gen_spec.__name__])
messages += [ai.fsystem(dbs.preprompts["respec"])] messages += [ai.fsystem(dbs.preprompts["respec"])]
@@ -121,7 +122,7 @@ def respec(ai: AI, dbs: DBs):
return messages return messages
def gen_unit_tests(ai: AI, dbs: DBs): def gen_unit_tests(ai: AI, dbs: DBs) -> List[dict]:
""" """
Generate unit tests based on the specification, that should work. Generate unit tests based on the specification, that should work.
""" """
@@ -139,8 +140,8 @@ def gen_unit_tests(ai: AI, dbs: DBs):
return messages return messages
def gen_clarified_code(ai: AI, dbs: DBs): def gen_clarified_code(ai: AI, dbs: DBs) -> List[dict]:
# get the messages from previous step """Takes clarification and generates code"""
messages = json.loads(dbs.logs[clarify.__name__]) messages = json.loads(dbs.logs[clarify.__name__])
@@ -153,7 +154,7 @@ def gen_clarified_code(ai: AI, dbs: DBs):
return messages return messages
def gen_code(ai: AI, dbs: DBs): def gen_code(ai: AI, dbs: DBs) -> List[dict]:
# get the messages from previous step # get the messages from previous step
messages = [ messages = [
@@ -167,7 +168,7 @@ def gen_code(ai: AI, dbs: DBs):
return messages return messages
def execute_entrypoint(ai, dbs): def execute_entrypoint(ai: AI, dbs: DBs) -> List[dict]:
command = dbs.workspace["run.sh"] command = dbs.workspace["run.sh"]
print("Do you want to execute this code?") print("Do you want to execute this code?")