Refactor

gpt_engineer/applications/cli/cli_agent.py

@@ -1,12 +1,17 @@
 from gpt_engineer.core.code import Code
 from gpt_engineer.core.base_version_manager import BaseVersionManager
 from gpt_engineer.core.ai import AI
-from gpt_engineer.core.default.steps import gen_code, gen_entrypoint, execute_entrypoint
+from gpt_engineer.core.default.steps import (
+    gen_code,
+    gen_entrypoint,
+    execute_entrypoint,
+    improve,
+)
 from gpt_engineer.core.base_repository import BaseRepository
 from gpt_engineer.core.default.on_disk_repository import OnDiskRepository
 from gpt_engineer.core.base_execution_env import BaseExecutionEnv
 from gpt_engineer.core.default.on_disk_execution_env import OnDiskExecutionEnv
-from gpt_engineer.core.default.paths import memory_path
+from gpt_engineer.core.default.paths import memory_path, ENTRYPOINT_FILE
 from gpt_engineer.core.base_agent import BaseAgent
 from gpt_engineer.applications.cli.learning import human_review
 
@@ -78,5 +83,11 @@ def init(self, prompt: str) -> Code:
         human_review(self.memory)
         return code
 
-    def improve(self, prompt: str, code) -> Code:
-        pass
+    def improve(self, prompt: str, code: Code) -> Code:
+        code = improve(self.ai, prompt, code)
+        if not ENTRYPOINT_FILE in code:
+            entrypoint = gen_entrypoint(self.ai, code, self.memory)
+            code = Code(code | entrypoint)
+        execute_entrypoint(self.execution_env, code)
+        human_review(self.memory)
+        return code

gpt_engineer/applications/cli/learning.py

@@ -56,6 +56,7 @@
     OnDiskRepository,
     FileRepositories,
 )
+from gpt_engineer.core.base_repository import BaseRepository
 from gpt_engineer.core.domain import Step
 
 
@@ -282,7 +283,7 @@ def get_session() -> str:
         return "ephemeral_" + str(random.randint(0, 2**32))
 
 
-def human_review(memory: OnDiskRepository):
+def human_review(memory: BaseRepository):
     """
     Collects human feedback on the code and stores it in memory.
 

gpt_engineer/core/ai.py

@@ -89,7 +89,13 @@ class AI:
 
     """
 
-    def __init__(self, model_name="gpt-4", temperature=0.1, azure_endpoint=""):
+    def __init__(
+        self,
+        model_name="gpt-4-1106-preview",
+        temperature=0.1,
+        azure_endpoint="",
+        streaming=True,
+    ):
         """
         Initialize the AI class.
 
@@ -103,7 +109,7 @@ def __init__(self, model_name="gpt-4", temperature=0.1, azure_endpoint=""):
         self.temperature = temperature
         self.azure_endpoint = azure_endpoint
         self.model_name = self._check_model_access_and_fallback(model_name)
-
+        self.streaming = streaming
         self.llm = self._create_chat_model()
         self.token_usage_log = TokenUsageLog(model_name)
 
@@ -306,13 +312,13 @@ def _create_chat_model(self) -> BaseChatModel:
                 openai_api_version=os.getenv("OPENAI_API_VERSION", "2023-05-15"),
                 deployment_name=self.model_name,
                 openai_api_type="azure",
-                streaming=True,
+                streaming=self.streaming,
             )
 
         return ChatOpenAI(
             model=self.model_name,
             temperature=self.temperature,
-            streaming=True,
+            streaming=self.streaming,
             client=openai.ChatCompletion,
         )
 

gpt_engineer/core/base_agent.py

@@ -5,46 +5,6 @@
 
 
 class BaseAgent(ABC):
-    """
-    The `Agent` class is responsible for managing the lifecycle of code generation and improvement.
-
-    Attributes:
-        path (str): The file path where the `Agent` will operate, used for version management and
-                    file operations.
-        version_manager (BaseVersionManager): An object that adheres to the VersionManagerInterface,
-                        responsible for version control of the generated code. Defaults to `VersionManager`
-                        if not provided. PROBABLY GIT SHOULD BE USED IN THE DEFAULT
-        step_bundle (StepBundleInterface): Workflows of code generation steps that define the behavior of gen_code and
-        improve.
-        ai (AI): Manages calls to the LLM.
-
-    Methods:
-        __init__(self, path: str, version_manager: VersionManagerInterface = None,
-                 step_bundle: StepBundleInterface = None, ai: AI = None):
-            Initializes a new instance of the Agent class with the provided path, version manager,
-            step bundle, and AI. It falls back to default instances if specific components are not provided.
-
-        init(self, prompt: str) -> Code:
-            Generates a new piece of code using the AI and step bundle based on the provided prompt.
-            It also snapshots the generated code using the version manager.
-
-            Parameters:
-                prompt (str): A string prompt that guides the code generation process.
-
-            Returns:
-                Code: An instance of the `Code` class containing the generated code.
-
-        improve(self, prompt: str) -> Code:
-            Improves an existing piece of code using the AI and step bundle based on the provided prompt.
-            It also snapshots the improved code using the version manager.
-
-            Parameters:
-                prompt (str): A string prompt that guides the code improvement process.
-
-            Returns:
-                Code: An instance of the `Code` class containing the improved code.
-    """
-
     @abstractmethod
     def init(self, prompt: str) -> Code:
         pass

gpt_engineer/core/chat_to_files.py

@@ -26,7 +26,6 @@
 - apply_edits: Applies file edits to a workspace.
 """
 
-import os
 import re
 import logging
 
@@ -89,10 +88,7 @@ def parse_chat(chat) -> List[Tuple[str, str]]:
     return files
 
 
-
-
-
-def overwrite_files_with_edits(chat: str, code: Code):
+def overwrite_code_with_edits(chat: str, code: Code):
     edits = parse_edits(chat)
     apply_edits(edits, code)
 

gpt_engineer/core/code.py

@@ -7,10 +7,10 @@
 class Code(dict):
     def __setitem__(self, key, value):
         if not isinstance(key, str | Path):
-            raise TypeError("Keys must be strings")
+            raise TypeError("Keys must be strings or Path's")
         if not isinstance(value, str):
             raise TypeError("Values must be strings")
-        super()[key] = value
+        super().__setitem__(key, value)
 
     def to_chat(self):
         def format_file_to_input(file_name: str, file_content: str) -> str:
@@ -30,10 +30,16 @@ def format_file_to_input(file_name: str, file_content: str) -> str:
                 The formatted file string.
             """
             file_str = f"""
-            {file_name}
-            ```
-            {file_content}
+{file_name}
+```
+{file_content}
             ```
             """
             return file_str
-        return "\n".join([format_file_to_input(file_name, file_content) + "\n" for file_name, file_content in self.items()])
+
+        return "\n".join(
+            [
+                format_file_to_input(file_name, file_content) + "\n"
+                for file_name, file_content in self.items()
+            ]
+        )

gpt_engineer/core/default/lean_agent.py

@@ -1,56 +1,19 @@
 from gpt_engineer.core.code import Code
-from gpt_engineer.core.base_version_manager import BaseVersionManager
 from gpt_engineer.core.ai import AI
-from gpt_engineer.core.default.steps import gen_code, gen_entrypoint, execute_entrypoint
+from gpt_engineer.core.default.steps import (
+    gen_code,
+    gen_entrypoint,
+    improve,
+)
 from gpt_engineer.core.base_repository import BaseRepository
 from gpt_engineer.core.default.on_disk_repository import OnDiskRepository
 from gpt_engineer.core.base_execution_env import BaseExecutionEnv
 from gpt_engineer.core.default.on_disk_execution_env import OnDiskExecutionEnv
-from gpt_engineer.core.default.paths import memory_path
+from gpt_engineer.core.default.paths import memory_path, ENTRYPOINT_FILE
 from gpt_engineer.core.base_agent import BaseAgent
 
 
-class Agent(BaseAgent):
-    """
-    The `Agent` class is responsible for managing the lifecycle of code generation and improvement.
-
-    Attributes:
-        path (str): The file path where the `Agent` will operate, used for version management and
-                    file operations.
-        version_manager (BaseVersionManager): An object that adheres to the VersionManagerInterface,
-                        responsible for version control of the generated code. Defaults to `VersionManager`
-                        if not provided. PROBABLY GIT SHOULD BE USED IN THE DEFAULT
-        step_bundle (StepBundleInterface): Workflows of code generation steps that define the behavior of gen_code and
-        improve.
-        ai (AI): Manages calls to the LLM.
-
-    Methods:
-        __init__(self, path: str, version_manager: VersionManagerInterface = None,
-                 step_bundle: StepBundleInterface = None, ai: AI = None):
-            Initializes a new instance of the Agent class with the provided path, version manager,
-            step bundle, and AI. It falls back to default instances if specific components are not provided.
-
-        init(self, prompt: str) -> Code:
-            Generates a new piece of code using the AI and step bundle based on the provided prompt.
-            It also snapshots the generated code using the version manager.
-
-            Parameters:
-                prompt (str): A string prompt that guides the code generation process.
-
-            Returns:
-                Code: An instance of the `Code` class containing the generated code.
-
-        improve(self, prompt: str) -> Code:
-            Improves an existing piece of code using the AI and step bundle based on the provided prompt.
-            It also snapshots the improved code using the version manager.
-
-            Parameters:
-                prompt (str): A string prompt that guides the code improvement process.
-
-            Returns:
-                Code: An instance of the `Code` class containing the improved code.
-    """
-
+class LeanAgent(BaseAgent):
     def __init__(
         self,
         memory: BaseRepository,
@@ -73,8 +36,13 @@ def init(self, prompt: str) -> Code:
         code = gen_code(self.ai, prompt, self.memory)
         entrypoint = gen_entrypoint(self.ai, code, self.memory)
         code = Code(code | entrypoint)
-        execute_entrypoint(self.execution_env, code)
+        self.execution_env.execute_program(code)
         return code
 
     def improve(self, prompt: str, code: Code) -> Code:
-        pass
+        code = improve(self.ai, prompt, code)
+        if not ENTRYPOINT_FILE in code:
+            entrypoint = gen_entrypoint(self.ai, code, self.memory)
+            code = Code(code | entrypoint)
+        self.execution_env.execute_program(code)
+        return code

gpt_engineer/core/default/steps.py

@@ -1,6 +1,6 @@
 from gpt_engineer.core.code import Code
 from gpt_engineer.core.ai import AI
-from gpt_engineer.core.chat_to_files import parse_chat, overwrite_files_with_edits#, format_file_to_input
+from gpt_engineer.core.chat_to_files import parse_chat, overwrite_code_with_edits
 from gpt_engineer.core.default.paths import (
     ENTRYPOINT_FILE,
     CODE_GEN_LOG_FILE,
@@ -35,7 +35,7 @@ def curr_fn() -> str:
     return inspect.stack()[1].function
 
 
-def setup_sys_prompt(db: OnDiskRepository) -> str:
+def setup_sys_prompt(preprompts: OnDiskRepository) -> str:
     """
     Constructs a system prompt for the AI based on predefined instructions and philosophies.
 
@@ -45,16 +45,16 @@ def setup_sys_prompt(db: OnDiskRepository) -> str:
     "philosophy" taken from the given DBs object.
 
     Parameters:
-    - dbs (DBs): The database object containing pre-defined prompts and instructions.
+    - preprompts (DBs): The database object containing pre-defined prompts and instructions.
 
     Returns:
     - str: The constructed system prompt for the AI.
     """
     return (
-        db["roadmap"]
-        + db["generate"].replace("FILE_FORMAT", db["file_format"])
+        preprompts["roadmap"]
+        + preprompts["generate"].replace("FILE_FORMAT", preprompts["file_format"])
         + "\nUseful to know:\n"
-        + db["philosophy"]
+        + preprompts["philosophy"]
     )
 
 
@@ -69,7 +69,7 @@ def gen_code(ai: AI, prompt: str, memory: BaseRepository) -> Code:
 
     Parameters:
     - ai (AI): An instance of the AI model.
-    - dbs (DBs): An instance containing the database configurations, including system and
+    - preprompts (DBs): An instance containing the database configurations, including system and
       input prompts, and file formatting preferences.
 
     Returns:
@@ -79,8 +79,8 @@ def gen_code(ai: AI, prompt: str, memory: BaseRepository) -> Code:
     The function assumes the `ai.start` method and the `to_files` utility are correctly
     set up and functional. Ensure these prerequisites are in place before invoking `simple_gen`.
     """
-    db = OnDiskRepository(PREPROMPTS_PATH)
-    messages = ai.start(setup_sys_prompt(db), prompt, step_name=curr_fn())
+    preprompts = OnDiskRepository(PREPROMPTS_PATH)
+    messages = ai.start(setup_sys_prompt(preprompts), prompt, step_name=curr_fn())
     chat = messages[-1].content.strip()
     memory[CODE_GEN_LOG_FILE] = chat
     files = parse_chat(chat)
@@ -99,7 +99,7 @@ def gen_entrypoint(ai: AI, code: Code, memory: BaseRepository) -> Code:
 
     Parameters:
     - ai (AI): An instance of the AI model.
-    - dbs (DBs): An instance containing the database configurations and workspace
+    - prepromptss (DBs): An instance containing the database configurations and workspace
       information, particularly the 'all_output.txt' which contains details about the
       codebase on disk.
 
@@ -155,8 +155,6 @@ def execute_entrypoint(execution_env: BaseExecutionEnv, code: Code) -> None:
     Parameters:
     - ai (AI): An instance of the AI model, not directly used in this function but
       included for consistency with other functions.
-    - dbs (DBs): An instance containing the database configurations and workspace
-      information.
 
     Returns:
     - List[dict]: An empty list. This function does not produce a list of messages
@@ -205,7 +203,7 @@ def execute_entrypoint(execution_env: BaseExecutionEnv, code: Code) -> None:
     execution_env.execute_program(code)
 
 
-def setup_sys_prompt_existing_code(db: OnDiskRepository) -> str:
+def setup_sys_prompt_existing_code(preprompts: OnDiskRepository) -> str:
     """
     Constructs a system prompt for the AI focused on improving an existing codebase.
 
@@ -215,17 +213,18 @@ def setup_sys_prompt_existing_code(db: OnDiskRepository) -> str:
     "philosophy" taken from the given DBs object.
 
     Parameters:
-    - dbs (DBs): The database object containing pre-defined prompts and instructions.
+    - preprompts (DBs): The database object containing pre-defined prompts and instructions.
 
     Returns:
     - str: The constructed system prompt focused on existing code improvement for the AI.
     """
     return (
-        db.preprompts["improve"].replace("FILE_FORMAT", db.preprompts["file_format"])
+        preprompts["improve"].replace("FILE_FORMAT", preprompts["file_format"])
         + "\nUseful to know:\n"
-        + db.preprompts["philosophy"]
+        + preprompts["philosophy"]
     )
 
+
 def improve(ai: AI, prompt: str, code: Code) -> Code:
     """
     Process and improve the code from a specified set of existing files based on a user prompt.
@@ -271,5 +270,5 @@ def improve(ai: AI, prompt: str, code: Code) -> Code:
 
     messages = ai.next(messages, step_name=curr_fn())
 
-    overwrite_files_with_edits(messages[-1].content.strip(), code)
-    return messages
+    overwrite_code_with_edits(messages[-1].content.strip(), code)
+    return code