Improvements

.gitignore

@@ -1,3 +1,6 @@
+# project specific
+example/*
+
 # See https://help.github.com/ignore-files/ for more about ignoring files.
 
 # Byte-compiled / optimized / DLL files

README.md

@@ -36,7 +36,7 @@ With an api key that has GPT4 access run:
 
 
 **Run**:
-- Create an empty folder. If inside the repo, you can run:
+- Create a new empty folder with a `main_prompt` file (or copy the example folder `cp -r example/ my-new-project`)
   - `cp -r projects/example/ projects/my-new-project`
 - Fill in the `main_prompt` file in your new folder
 - Run: `gpt-engineer projects/my-new-project`

gpt_engineer/ai/ai.py

@@ -0,0 +1,59 @@
+from abc import ABC, abstractmethod
+from typing import Any, Dict, List, Optional, Tuple
+
+from gpt_engineer.models import Message, Role
+
+
+class AI(ABC):
+    """Abstract class for AI models. Any LLM model for use in gpt-engineer must satisfy
+    the following interface.
+    """
+
+    @abstractmethod
+    def __init__(self, **kwargs: Dict[str, 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: List[Tuple[Role, Message]]
+    ) -> List[Tuple[Role, Message]]:
+        """
+        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: List[Tuple[Role, Message]], user_prompt: Optional[Message] = None
+    ) -> List[Tuple[Role, Message]]:
+        """
+        Asks the AI model to respond to the user prompt after ingesting some initial
+            messages.
+
+        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.
+
+        Returns:
+            List[Dict[str, str]]: Returns the chat completion response along with
+              previous messages.
+        """
+        pass

gpt_engineer/ai/gpt.py

@@ -0,0 +1,71 @@
+import logging
+
+from typing import Any, Dict, List, Optional, Tuple
+
+import openai
+
+from gpt_engineer.ai.ai import AI
+from gpt_engineer.models import Message, Role
+
+logging.basicConfig(level=logging.INFO)
+
+
+class GPT(AI):
+    def __init__(self, **kwargs: Dict[str, Any]) -> None:
+        self.kwargs = kwargs
+        self._model_check_and_fallback()
+
+    def start(
+        self, initial_conversation: List[Tuple[Role, Message]]
+    ) -> List[Tuple[Role, Message]]:
+        return self.next(initial_conversation)
+
+    def next(
+        self, messages: List[Tuple[Role, Message]], user_prompt: Optional[Message] = None
+    ) -> List[Tuple[Role, Message]]:
+        if user_prompt:
+            messages.append((Role.USER, user_prompt))
+
+        response = openai.ChatCompletion.create(
+            messages=[self._format_message(r, m) for r, m in messages],
+            stream=True,
+            **self.kwargs,
+        )
+
+        chat = []
+        for chunk in response:
+            delta = chunk["choices"][0]["delta"]
+            msg = delta.get("content", "")
+            logging.info(msg)
+            chat.append(msg)
+
+        messages.append((Role.ASSISTANT, Message(content="".join(chat))))
+        return messages
+
+    def _format_message(self, role: Role, 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": role.value, "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"

gpt_engineer/ai/models.py

@@ -0,0 +1,24 @@
+from enum import Enum
+from typing import Any, Callable, Dict
+
+from .ai import AI
+from .gpt import GPT
+from .test import TestAI
+
+
+class ModelName(str, Enum):
+    GPT4 = "gpt-4"
+    GPT35_TURBO = "gpt-3.5-turbo"
+    TEST = "test"
+
+
+models: Dict[
+    ModelName,
+    Callable[
+        [
+            Any,
+        ],
+        AI,
+    ],
+] = {ModelName.GPT4: GPT, ModelName.GPT35_TURBO: GPT, ModelName.TEST: TestAI}
+default_model_name: ModelName = ModelName.GPT4

gpt_engineer/ai/test.py

@@ -0,0 +1,21 @@
+from typing import Any, Dict, List, Optional, Tuple
+
+from gpt_engineer.ai.ai import AI
+from gpt_engineer.models import Message, Role
+
+
+class TestAI(AI):
+    """A simple AI that tests the code's functionality."""
+
+    def __init__(self, **kwargs: Dict[str, Any]) -> None:
+        pass
+
+    def start(
+        self, initial_conversation: List[Tuple[Role, Message]]
+    ) -> List[Tuple[Role, Message]]:
+        return [(Role.ASSISTANT, Message("hello world"))]
+
+    def next(
+        self, messages: List[Tuple[Role, Message]], user_prompt: Optional[Message] = None
+    ) -> List[Tuple[Role, Message]]:
+        return [(Role.ASSISTANT, Message("Unto the next world"))]

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
@@ -25,7 +27,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)

gpt_engineer/main.py

@@ -7,7 +7,7 @@
 
 import typer
 
-from gpt_engineer.ai import AI
+from gpt_engineer.ai.models import models
 from gpt_engineer.db import DB, DBs
 from gpt_engineer.steps import STEPS
 
@@ -41,10 +41,8 @@ def main(
         shutil.rmtree(memory_path, ignore_errors=True)
         shutil.rmtree(workspace_path, ignore_errors=True)
 
-    ai = AI(
-        model=model,
-        temperature=temperature,
-    )
+    kwargs = {"model": model, "temperature": temperature}
+    ai = models[model](**kwargs)
 
     dbs = DBs(
         memory=DB(memory_path),

gpt_engineer/models.py

@@ -0,0 +1,19 @@
+from dataclasses import dataclass
+from enum import Enum
+
+
+class Step(Enum):
+    CLARIFY = "clarify"
+    RUN_CLARIFIED = "run_clarified"
+
+
+@dataclass
+class Message:
+    content: str
+
+
+class Role(str, Enum):
+    SYSTEM = "system"
+    USER = "user"
+    ASSISTANT = "assistant"
+    FUNCTION = "function"