diff --git a/.gitignore b/.gitignore index 3aedcd0701..db7be7dfa2 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,4 @@ + # See https://help.github.com/ignore-files/ for more about ignoring files. # Byte-compiled / optimized / DLL files diff --git a/gpt_engineer/ai.py b/gpt_engineer/ai.py index d837a07d3f..ea4c768868 100644 --- a/gpt_engineer/ai.py +++ b/gpt_engineer/ai.py @@ -2,53 +2,81 @@ import logging +from abc import ABC, abstractmethod +from typing import Any, Dict, Optional + import openai -logger = logging.getLogger(__name__) +from gpt_engineer.models import Message, Messages, Role -class AI: - def __init__(self, model="gpt-4", temperature=0.1): - self.temperature = temperature +class AI(ABC): + """Abstract class for AI models. Any LLM model for use in gpt-engineer must satisfy + the following interface. + """ - try: - openai.Model.retrieve(model) - self.model = model - except openai.InvalidRequestError: - print( - f"Model {model} not available for provided API key. Reverting " - "to gpt-3.5-turbo. Sign up for the GPT-4 wait list here: " - "https://openai.com/waitlist/gpt-4-api" - ) - self.model = "gpt-3.5-turbo" + @abstractmethod + def __init__(self, **kwargs: Any) -> None: + """ + Initialization for the AI model. + + Args: + **kwargs: Variable length argument list for model configuration. **kwargs + can/will depend on AI subclass. + """ + pass + + @abstractmethod + def start(self, initial_conversation: Messages) -> Messages: + """ + Initializes the conversation with specific system and user messages. + + Args: + initial_conversation (List[Tuple[Role, Message]]): The initial messages + given to the + AI. Generally these are prompt to the system about their task, qa and/or + philosophy. + + Returns: + List[Dict[str, str]]: Returns the next set of conversation. + """ + pass + + @abstractmethod + def next(self, messages: Messages, user_prompt: Optional[str] = None) -> Messages: + """ + Asks the AI model to respond to the user prompt after ingesting some initial + messages. - def start(self, system, user): - messages = [ - {"role": "system", "content": system}, - {"role": "user", "content": user}, - ] + Args: + messages (List[Dict[str, str]]): The list of messages to be used for + chat completion. + user_prompt (str, optional): Additional prompt to be added to messages. + Defaults to None. - return self.next(messages) + Returns: + List[Dict[str, str]]: Returns the chat completion response along with + previous messages. + """ + pass - def fsystem(self, msg): - return {"role": "system", "content": msg} - def fuser(self, msg): - return {"role": "user", "content": msg} +class GPT(AI): + def __init__(self, **kwargs: Any) -> None: + self.kwargs = kwargs + self._model_check_and_fallback() - def fassistant(self, msg): - return {"role": "assistant", "content": msg} + def start(self, initial_conversation: Messages) -> Messages: + return self.next(initial_conversation) - def next(self, messages: list[dict[str, str]], prompt=None): - if prompt: - messages += [{"role": "user", "content": prompt}] + def next(self, messages: Messages, user_prompt: Optional[str] = None) -> Messages: + if user_prompt: + messages.messages.append(Message(user_prompt, Role.USER)) - logger.debug(f"Creating a new chat completion: {messages}") response = openai.ChatCompletion.create( - messages=messages, + messages=[self._format_message(m) for m in messages.messages], stream=True, - model=self.model, - temperature=self.temperature, + **self.kwargs, ) chat = [] @@ -57,7 +85,34 @@ def next(self, messages: list[dict[str, str]], prompt=None): msg = delta.get("content", "") print(msg, end="") chat.append(msg) - print() - messages += [{"role": "assistant", "content": "".join(chat)}] - logger.debug(f"Chat completion finished: {messages}") + + messages.messages.append(Message("".join(chat), Role.ASSISTANT)) return messages + + def _format_message(self, msg: Message) -> Dict[str, str]: + """ + Formats the message as per role. + + Args: + role (str): The role to be used for the message. + msg (str): The message content. + + Returns: + Dict[str, str]: A dictionary containing the role and content. + """ + return {"role": msg.role, "content": msg.content} + + def _model_check_and_fallback(self) -> None: + """ + Checks if the desired model is available; if not, it falls back to a default + model. + """ + try: + openai.Model.retrieve(self.kwargs.get("model", "gpt-4")) + except openai.error.InvalidRequestError: + logging.info( + "Model gpt-4 not available for provided api key reverting " + "to gpt-3.5.turbo. Sign up for the gpt-4 wait list here: " + "https://openai.com/waitlist/gpt-4-api" + ) + self.kwargs["model"] = "gpt-3.5-turbo" diff --git a/gpt_engineer/chat_to_files.py b/gpt_engineer/chat_to_files.py index 8e251cbdfd..1d56e1c21b 100644 --- a/gpt_engineer/chat_to_files.py +++ b/gpt_engineer/chat_to_files.py @@ -1,5 +1,7 @@ import re +from gpt_engineer.db import DB + def parse_chat(chat): # -> List[Tuple[str, str]]: # Get all ``` blocks and preceding filenames @@ -34,7 +36,7 @@ def parse_chat(chat): # -> List[Tuple[str, str]]: return files -def to_files(chat, workspace): +def to_files(chat, workspace: DB): workspace["all_output.txt"] = chat files = parse_chat(chat) diff --git a/gpt_engineer/main.py b/gpt_engineer/main.py index cface530fc..75095214ef 100644 --- a/gpt_engineer/main.py +++ b/gpt_engineer/main.py @@ -1,4 +1,3 @@ -import json import logging import os import shutil @@ -7,10 +6,9 @@ import typer -from gpt_engineer import steps -from gpt_engineer.ai import AI +from gpt_engineer.ai import GPT from gpt_engineer.db import DB, DBs -from gpt_engineer.steps import STEPS +from gpt_engineer.steps import STEPS, Config app = typer.Typer() @@ -21,8 +19,8 @@ def main( delete_existing: bool = typer.Argument(False, help="delete existing files"), model: str = "gpt-4", temperature: float = 0.1, - steps_config: steps.Config = typer.Option( - steps.Config.DEFAULT, "--steps", "-s", help="decide which steps to run" + steps_config: Config = typer.Option( + Config.DEFAULT, "--steps", "-s", help="decide which steps to run" ), verbose: bool = typer.Option(False, "--verbose", "-v"), run_prefix: str = typer.Option( @@ -44,10 +42,7 @@ def main( shutil.rmtree(memory_path, ignore_errors=True) shutil.rmtree(workspace_path, ignore_errors=True) - ai = AI( - model=model, - temperature=temperature, - ) + ai = GPT(model=model, temperature=temperature) dbs = DBs( memory=DB(memory_path), @@ -59,7 +54,7 @@ def main( for step in STEPS[steps_config]: messages = step(ai, dbs) - dbs.logs[step.__name__] = json.dumps(messages) + dbs.logs[step.__name__] = messages.to_json() if __name__ == "__main__": diff --git a/gpt_engineer/models.py b/gpt_engineer/models.py new file mode 100644 index 0000000000..5f14087b04 --- /dev/null +++ b/gpt_engineer/models.py @@ -0,0 +1,28 @@ +from dataclasses import dataclass +from enum import Enum +from typing import List + +from dataclasses_json import DataClassJsonMixin, dataclass_json + + +class Role(str, Enum): + SYSTEM = "system" + USER = "user" + ASSISTANT = "assistant" + FUNCTION = "function" + + +@dataclass_json +@dataclass +class Message(DataClassJsonMixin): + content: str + role: Role + + +@dataclass_json +@dataclass +class Messages(DataClassJsonMixin): + messages: List[Message] + + def last_message_content(self) -> str: + return self.messages[-1].content diff --git a/gpt_engineer/steps.py b/gpt_engineer/steps.py index 3370c9490c..ce88a27eca 100644 --- a/gpt_engineer/steps.py +++ b/gpt_engineer/steps.py @@ -1,144 +1,188 @@ -import json import re import subprocess from enum import Enum +from typing import Callable, Dict, List from gpt_engineer.ai import AI from gpt_engineer.chat_to_files import to_files from gpt_engineer.db import DBs +from gpt_engineer.models import Message, Messages, Role -def setup_sys_prompt(dbs): +def setup_sys_prompt(dbs: DBs) -> str: + """Constructs the system setup prompt.""" return dbs.identity["generate"] + "\nUseful to know:\n" + dbs.identity["philosophy"] -def simple_gen(ai: AI, dbs: DBs): - """Run the AI on the main prompt and save the results""" - messages = ai.start( - setup_sys_prompt(dbs), - dbs.input["main_prompt"], - ) - to_files(messages[-1]["content"], dbs.workspace) - return messages +### Steps -def clarify(ai: AI, dbs: DBs): +def clarify(ai: AI, dbs: DBs) -> Messages: """ - 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. """ - messages = [ai.fsystem(dbs.identity["qa"])] - user = dbs.input["main_prompt"] + user_msg = dbs.input["main_prompt"] + messages = Messages([Message(role=Role.SYSTEM, content=dbs.identity["qa"])]) while True: - messages = ai.next(messages, user) + messages = ai.next(messages=messages, user_prompt=user_msg) - if messages[-1]["content"].strip().lower().startswith("no"): + response_content: str = messages.last_message_content() + if response_content.strip().lower().startswith("no"): break print() - user = input('(answer in text, or "c" to move on)\n') + user_msg = input('(answer in text, or "c" to move on)\n') print() - if not user or user == "c": + if not user_msg or user_msg == "c": break - user += ( + user_msg += ( "\n\n" "Is anything else unclear? If yes, only answer in the form:\n" "{remaining unclear areas} remaining questions.\n" "{Next question}\n" 'If everything is sufficiently clear, only answer "no".' ) - print() return messages -def gen_spec(ai: AI, dbs: DBs): +def simple_gen(ai: AI, dbs: DBs) -> Messages: + """Run the AI on the main prompt and save the results""" + messages = ai.start( + Messages( + [ + Message(role=Role.SYSTEM, content=setup_sys_prompt(dbs)), + Message(role=Role.USER, content=dbs.input["main_prompt"]), + ] + ) + ) + to_files(messages.last_message_content(), dbs.workspace) + return messages + + +def run_clarified(ai: AI, dbs: DBs) -> Messages: + """ + Run the AI using the messages clarified in the previous step. + """ + # get the messages from previous step + messages = ai.next( + Messages( + [Message(role=Role.SYSTEM, content=setup_sys_prompt(dbs))] + + Messages.from_json(dbs.logs[clarify.__name__]).messages[1:] + ), + user_prompt=dbs.identity["use_qa"], + ) + to_files(messages.last_message_content(), dbs.workspace) + return messages + + +def gen_spec(ai: AI, dbs: DBs) -> Messages: """ Generate a spec from the main prompt + clarifications and save the results to the workspace """ - messages = [ - ai.fsystem(setup_sys_prompt(dbs)), - ai.fsystem(f"Instructions: {dbs.input['main_prompt']}"), - ] - - messages = ai.next(messages, dbs.identity["spec"]) + messages = Messages( + [ + Message(role=Role.SYSTEM, content=setup_sys_prompt(dbs)), + Message( + role=Role.SYSTEM, content=f"Instructions: {dbs.input['main_prompt']}" + ), + ] + ) + messages = ai.next(messages, user_prompt=dbs.identity["spec"]) - dbs.memory["specification"] = messages[-1]["content"] + dbs.memory["specification"] = messages.last_message_content() return messages -def respec(ai: AI, dbs: DBs): - messages = json.loads(dbs.logs[gen_spec.__name__]) - messages += [ai.fsystem(dbs.identity["respec"])] +def respec(ai: AI, dbs: DBs) -> Messages: + messages = Messages( + Messages.from_json(dbs.logs[gen_spec.__name__]).messages + + [Message(role=Role.SYSTEM, content=dbs.identity["respec"])] + ) messages = ai.next(messages) messages = ai.next( messages, - ( + user_prompt=( "Based on the conversation so far, please reiterate the specification for " - "the program. " - "If there are things that can be improved, please incorporate the " - "improvements. " - "If you are satisfied with the specification, just write out the " - "specification word by word again." + "the program. If there are things that can be improved, please incorporate " + "the improvements. If you are satisfied with the specification, just write " + "out the specification word by word again." ), ) - dbs.memory["specification"] = messages[-1]["content"] + dbs.memory["specification"] = messages.last_message_content() return messages -def gen_unit_tests(ai: AI, dbs: DBs): +def gen_unit_tests(ai: AI, dbs: DBs) -> Messages: """ Generate unit tests based on the specification, that should work. """ - messages = [ - ai.fsystem(setup_sys_prompt(dbs)), - ai.fuser(f"Instructions: {dbs.input['main_prompt']}"), - ai.fuser(f"Specification:\n\n{dbs.memory['specification']}"), - ] + messages = Messages( + [ + Message(role=Role.SYSTEM, content=setup_sys_prompt(dbs)), + Message(role=Role.USER, content=f"Instructions: {dbs.input['main_prompt']}"), + Message( + role=Role.USER, content=f"Specification:\n\n{dbs.memory['specification']}" + ), + ] + ) - messages = ai.next(messages, dbs.identity["unit_tests"]) + messages = ai.next(messages, user_prompt=dbs.identity["unit_tests"]) - dbs.memory["unit_tests"] = messages[-1]["content"] + dbs.memory["unit_tests"] = messages.last_message_content() to_files(dbs.memory["unit_tests"], dbs.workspace) return messages -def gen_clarified_code(ai: AI, dbs: DBs): +def gen_clarified_code(ai: AI, dbs: DBs) -> Messages: # get the messages from previous step - messages = json.loads(dbs.logs[clarify.__name__]) - - messages = [ - ai.fsystem(setup_sys_prompt(dbs)), - ] + messages[1:] - messages = ai.next(messages, dbs.identity["use_qa"]) + messages = ai.next( + Messages( + [Message(role=Role.SYSTEM, content=setup_sys_prompt(dbs))] + + Messages.from_json(dbs.logs[clarify.__name__]).messages[1:] + ), + user_prompt=dbs.identity["use_qa"], + ) - to_files(messages[-1]["content"], dbs.workspace) + to_files(messages.last_message_content(), dbs.workspace) return messages -def gen_code(ai: AI, dbs: DBs): +def gen_code(ai: AI, dbs: DBs) -> Messages: # get the messages from previous step - messages = [ - ai.fsystem(setup_sys_prompt(dbs)), - ai.fuser(f"Instructions: {dbs.input['main_prompt']}"), - ai.fuser(f"Specification:\n\n{dbs.memory['specification']}"), - ai.fuser(f"Unit tests:\n\n{dbs.memory['unit_tests']}"), - ] - messages = ai.next(messages, dbs.identity["use_qa"]) - to_files(messages[-1]["content"], dbs.workspace) + messages = ai.next( + Messages( + [ + Message(role=Role.SYSTEM, content=setup_sys_prompt(dbs)), + Message( + role=Role.USER, content=f"Instructions: {dbs.input['main_prompt']}" + ), + Message( + role=Role.USER, + content=f"Specification:\n\n{dbs.memory['specification']}", + ), + Message( + role=Role.USER, content=f"Unit tests:\n\n{dbs.memory['unit_tests']}" + ), + ] + ), + dbs.identity["use_qa"], + ) + to_files(messages.last_message_content(), dbs.workspace) return messages -def execute_entrypoint(ai, dbs): +def execute_entrypoint(ai, dbs) -> Messages: command = dbs.workspace["run.sh"] print("Do you want to execute this code?") @@ -149,7 +193,7 @@ def execute_entrypoint(ai, dbs): print() if input() not in ["", "y", "yes"]: print("Ok, not executing the code.") - return [] + return Messages(messages=[]) print("Executing the code...") print( "\033[92m" # green color @@ -159,55 +203,69 @@ def execute_entrypoint(ai, dbs): ) print() subprocess.run("bash run.sh", shell=True, cwd=dbs.workspace.path) - return [] + return Messages(messages=[]) -def gen_entrypoint(ai, dbs): +def gen_entrypoint(ai, dbs) -> Messages: messages = ai.start( - system=( - "You will get information about a codebase that is currently on disk in " - "the current folder.\n" - "From this you will answer with code blocks that includes all the necessary " - "unix terminal commands to " - "a) install dependencies " - "b) run all necessary parts of the codebase (in parallell if necessary).\n" - "Do not install globally. Do not use sudo.\n" - "Do not explain the code, just give the commands.\n" - "Do not use placeholders, use example values (like . for a folder argument) " - "if necessary.\n" - ), - user="Information about the codebase:\n\n" + dbs.workspace["all_output.txt"], + Messages( + [ + Message( + role=Role.SYSTEM, + content=( + "You will get information about a codebase that is currently on " + "disk in the current folder.\n From this you will answer with " + "code blocks that includes all the necessary unix terminal " + "commands to a) install dependencies b) run all necessary parts " + "of the codebase (in parallell if necessary).\n Do not install " + "globally. Do not use sudo.\n Do not explain the code, just give" + " the commands.\n Do not use placeholders, use example values " + "(like . for a folder argument) if necessary.\n" + ), + ), + Message( + role=Role.USER, + content=( + "Information about the codebase:\n\n" + + dbs.workspace["all_output.txt"] + ), + ), + ] + ) ) print() - regex = r"```\S*\n(.+?)```" - matches = re.finditer(regex, messages[-1]["content"], re.DOTALL) + matches = re.finditer(regex, messages.last_message_content(), re.DOTALL) dbs.workspace["run.sh"] = "\n".join(match.group(1) for match in matches) return messages -def use_feedback(ai: AI, dbs: DBs): - messages = [ - ai.fsystem(setup_sys_prompt(dbs)), - ai.fuser(f"Instructions: {dbs.input['main_prompt']}"), - ai.fassistant(dbs.workspace["all_output.txt"]), - ai.fsystem(dbs.identity["use_feedback"]), - ] - messages = ai.next(messages, dbs.input["feedback"]) - to_files(messages[-1]["content"], dbs.workspace) +def use_feedback(ai: AI, dbs: DBs) -> Messages: + messages = Messages( + [ + Message(role=Role.SYSTEM, content=setup_sys_prompt(dbs)), + Message(role=Role.USER, content=f"Instructions: {dbs.input['main_prompt']}"), + Message(role=Role.ASSISTANT, content=dbs.workspace["all_output.txt"]), + Message(role=Role.SYSTEM, content=dbs.identity["use_feedback"]), + ] + ) + messages = ai.next(messages, user_prompt=dbs.memory["feedback"]) + to_files(messages.last_message_content(), dbs.workspace) return messages -def fix_code(ai: AI, dbs: DBs): - code_ouput = json.loads(dbs.logs[gen_code.__name__])[-1]["content"] - messages = [ - ai.fsystem(setup_sys_prompt(dbs)), - ai.fuser(f"Instructions: {dbs.input['main_prompt']}"), - ai.fuser(code_ouput), - ai.fsystem(dbs.identity["fix_code"]), - ] - messages = ai.next(messages, "Please fix any errors in the code above.") - to_files(messages[-1]["content"], dbs.workspace) +def fix_code(ai: AI, dbs: DBs) -> Messages: + code_output = Messages.from_json(dbs.logs[gen_code.__name__]).last_message_content() + messages = Messages( + [ + Message(role=Role.SYSTEM, content=setup_sys_prompt(dbs)), + Message(role=Role.USER, content=f"Instructions: {dbs.input['main_prompt']}"), + Message(role=Role.USER, content=code_output), + Message(role=Role.SYSTEM, content=dbs.identity["fix_code"]), + ] + ) + messages = ai.next(messages, user_prompt="Please fix any errors in the code above.") + to_files(messages.last_message_content(), dbs.workspace) return messages @@ -224,7 +282,9 @@ class Config(str, Enum): # Different configs of what steps to run -STEPS = { +StepFunction = Callable[[AI, DBs], Messages] + +STEPS: Dict[Config, List[StepFunction]] = { Config.DEFAULT: [ clarify, gen_clarified_code, diff --git a/pyproject.toml b/pyproject.toml index c1ca348b10..cbfd94f12f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -9,6 +9,7 @@ readme = "README.md" requires-python = ">=3" dependencies = [ 'black == 23.3.0', + 'dataclasses-json == 0.5.8', 'mypy == 1.3.0', 'openai == 0.27.8', 'pre-commit == 3.3.3', @@ -24,6 +25,7 @@ gpt-engineer = 'gpt_engineer.main:app' [tool.setuptools] packages = ["gpt_engineer"] + # https://beta.ruff.rs/docs/configuration/#using-rufftoml [tool.ruff] select = ["F", "E", "W", "I001"] diff --git a/requirements.txt b/requirements.txt index ffc5c5338e..a9ab7ec4e1 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,9 +1,9 @@ black==23.3.0 +dataclasses-json==0.5.8 mypy==1.3.0 openai==0.27.8 pre-commit==3.3.3 pytest==7.3.1 ruff==0.0.272 termcolor==2.3.0 -typer==0.9.0 - +typer==0.9.0 \ No newline at end of file diff --git a/scripts/rerun_edited_message_logs.py b/scripts/rerun_edited_message_logs.py index 2bf4409dc1..a485f4dd8e 100644 --- a/scripts/rerun_edited_message_logs.py +++ b/scripts/rerun_edited_message_logs.py @@ -1,10 +1,11 @@ import json import pathlib +from typing import Optional + import typer -from gpt_engineer.ai import AI -from gpt_engineer.chat_to_files import to_files +from gpt_engineer.ai import GPT app = typer.Typer() @@ -12,11 +13,11 @@ @app.command() def main( messages_path: str, - out_path: str | None = None, + out_path: Optional[str] = None, model: str = "gpt-4", temperature: float = 0.1, ): - ai = AI( + ai = GPT( model=model, temperature=temperature, ) @@ -27,9 +28,8 @@ def main( messages = ai.next(messages) if out_path: - to_files(messages[-1]["content"], out_path) with open(pathlib.Path(out_path) / "all_output.txt", "w") as f: - json.dump(messages[-1]["content"], f) + json.dump(messages.last_message_content(), f) if __name__ == "__main__": diff --git a/tests/test_ai.py b/tests/test_ai.py index acc017e13b..5c30d096e3 100644 --- a/tests/test_ai.py +++ b/tests/test_ai.py @@ -1,9 +1,9 @@ import pytest -from gpt_engineer.ai import AI +from gpt_engineer.ai import GPT @pytest.mark.xfail(reason="Constructor assumes API access") def test_ai(): - AI() + GPT() # TODO Assert that methods behave and not only constructor.