Integrate Morph API for code editing

.env.template

@@ -3,3 +3,4 @@
 # OPENAI_API_KEY=Your personal OpenAI API key from https://platform.openai.com/account/api-keys
 OPENAI_API_KEY=...
 ANTHROPIC_API_KEY=...
+MORPH_API_KEY=...

gpt_engineer/core/default/steps.py

@@ -25,16 +25,22 @@
 setup_sys_prompt_existing_code : function
     Sets up the system prompt for improving existing code.
 
-
 improve : function
     Improves the code based on user input and returns the updated files.
+
+_apply_morph_edit : function
+    Send code editing instructions to Morph API and return the merged code.
 """
 
 import inspect
 import io
 import re
 import sys
 import traceback
+import json
+import os
+from openai import OpenAI
+
 
 from pathlib import Path
 from typing import List, MutableMapping, Union
@@ -95,7 +101,7 @@ def setup_sys_prompt(preprompts: MutableMapping[Union[str, Path], str]) -> str:
 
 
 def setup_sys_prompt_existing_code(
-    preprompts: MutableMapping[Union[str, Path], str]
+    preprompts: MutableMapping[Union[str, Path], str],
 ) -> str:
     """
     Sets up the system prompt for improving existing code.
@@ -268,6 +274,36 @@ def execute_entrypoint(
     return files_dict
 
 
+def apply_morph_edit(instructions: str, initial_code: str, code_edit: str) -> str:
+    """
+    Send code editing instructions to Morph API and return the merged code.
+
+    Args:
+        instructions (str): The instructions for the edit.
+        initial_code (str): The original code to modify.
+        code_edit (str): The proposed code changes.
+
+    Returns:
+        str: The merged/edited code returned by Morph.
+    """
+    client = OpenAI(
+        api_key=os.getenv("MORPH_API_KEY"), base_url="https://api.morphllm.com/v1"
+    )
+
+    response = client.chat.completions.create(
+        model="morph-v3-large",
+        messages=[
+            {
+                "role": "user",
+                "content": f"<instruction>{instructions}</instruction>\n<code>{initial_code}</code>\n<update>{code_edit}</update>",
+            }
+        ],
+    )
+
+    merged_code = response.choices[0].message.content
+    return merged_code
+
+
 def improve_fn(
     ai: AI,
     prompt: Prompt,
@@ -316,50 +352,25 @@ def _improve_loop(
     ai: AI, files_dict: FilesDict, memory: BaseMemory, messages: List, diff_timeout=3
 ) -> FilesDict:
     messages = ai.next(messages, step_name=curr_fn())
-    files_dict, errors = salvage_correct_hunks(
-        messages, files_dict, memory, diff_timeout=diff_timeout
-    )
-
-    retries = 0
-    while errors and retries < MAX_EDIT_REFINEMENT_STEPS:
-        messages.append(
-            HumanMessage(
-                content="Some previously produced diffs were not on the requested format, or the code part was not found in the code. Details:\n"
-                + "\n".join(errors)
-                + "\n Only rewrite the problematic diffs, making sure that the failing ones are now on the correct format and can be found in the code. Make sure to not repeat past mistakes. \n"
+    # Log AI raw response
+    ai_response = messages[-1].content.strip()
+    data = json.loads(ai_response)
+
+    for file in data:
+        if file["target_file"] not in files_dict:
+            files_dict[file["target_file"]] = file["code_edit"]
+        else:
+            # Use Morph to apply the edit
+            merged_code = apply_morph_edit(
+                instructions=file["instructions"],
+                initial_code=files_dict[file["target_file"]],
+                code_edit=file["code_edit"],
             )
-        )
-        messages = ai.next(messages, step_name=curr_fn())
-        files_dict, errors = salvage_correct_hunks(
-            messages, files_dict, memory, diff_timeout
-        )
-        retries += 1
+            files_dict[file["target_file"]] = merged_code
 
     return files_dict
 
 
-def salvage_correct_hunks(
-    messages: List, files_dict: FilesDict, memory: BaseMemory, diff_timeout=3
-) -> tuple[FilesDict, List[str]]:
-    error_messages = []
-    ai_response = messages[-1].content.strip()
-
-    diffs = parse_diffs(ai_response, diff_timeout=diff_timeout)
-    # validate and correct diffs
-
-    for _, diff in diffs.items():
-        # if diff is a new file, validation and correction is unnecessary
-        if not diff.is_new_file():
-            problems = diff.validate_and_correct(
-                file_to_lines_dict(files_dict[diff.filename_pre])
-            )
-            error_messages.extend(problems)
-    files_dict = apply_diffs(diffs, files_dict)
-    memory.log(IMPROVE_LOG_FILE, "\n\n".join(x.pretty_repr() for x in messages))
-    memory.log(DIFF_LOG_FILE, "\n\n".join(error_messages))
-    return files_dict, error_messages
-
-
 class Tee(object):
     def __init__(self, *files):
         self.files = files

gpt_engineer/preprompts/file_format_diff

@@ -1,41 +1,39 @@
-You will output the content of each file necessary to achieve the goal, including ALL code.
-Output requested code changes and new code in the unified "git diff" syntax. Example:
-
-```diff
---- example.txt
-+++ example.txt
-@@ -6,3 +6,4 @@
-     line content A
-     line content B
-+    new line added
--    original line X
-+    modified line X with changes
-@@ -26,4 +27,5 @@
-         condition check:
--            action for condition A
-+            if certain condition is met:
-+                alternative action for condition A
-         another condition check:
--            action for condition B
-+            modified action for condition B
-```
-
-Example of a git diff creating a new file:
-
-```diff
---- /dev/null
-+++ new_file.txt
-@@ -0,0 +1,3 @@
-+First example line
-+
-+Last example line
-```
-
-RULES:
--A program will apply the diffs you generate exactly to the code, so diffs must be precise and unambiguous!
--Every diff must be fenced with triple backtick ```.
--The file names at the beginning of a diff, (lines starting with --- and +++) is the relative path to the file before and after the diff.
--LINES TO BE REMOVED (starting with single -) AND LINES TO BE RETAIN (no starting symbol) HAVE TO REPLICATE THE DIFFED HUNK OF THE CODE EXACTLY LINE BY LINE. KEEP THE NUMBER OF RETAIN LINES SMALL IF POSSIBLE.
--EACH LINE IN THE SOURCE FILES STARTS WITH A LINE NUMBER, WHICH IS NOT PART OF THE SOURCE CODE. NEVER TRANSFER THESE LINE NUMBERS TO THE DIFF HUNKS.
--AVOID STARTING A HUNK WITH AN EMPTY LINE.
--ENSURE ALL CHANGES ARE PROVIDED IN A SINGLE DIFF CHUNK PER FILE TO PREVENT MULTIPLE DIFFS ON THE SAME FILE.
+You must produce your output as a JSON array of objects. Each object must have exactly three fields:
+
+1. **target_file**: the relative path of the file to edit.
+2. **instructions**: a brief description of the changes made to this file.
+3. **code_edit**: the edited content of the file, showing all changes in context. Use the special comment `// ... existing code ...` to represent unchanged spans of code.
+
+Example JSON output: 
+[
+  {
+    "target_file": "path/to/file.ext",
+    "instructions": "Brief summary of changes",
+    "code_edit": "// ... existing code ...\nFIRST_EDIT\n// ... existing code ...\nSECOND_EDIT\n// ... existing code ..."
+  }
+]
+
+Use this tool to make an edit to an existing file.
+
+This will be read by a less intelligent model, which will quickly apply the edit. You should make it clear what the edit is, while also minimizing the unchanged code you write.
+When writing the edit, you should specify each edit in sequence, with the special comment // ... existing code ... to represent unchanged code in between edited lines.
+
+For example:
+
+// ... existing code ...
+FIRST_EDIT
+// ... existing code ...
+SECOND_EDIT
+// ... existing code ...
+THIRD_EDIT
+// ... existing code ...
+
+You should still bias towards repeating as few lines of the original file as possible to convey the change.
+But, each edit should contain minimally sufficient context of unchanged lines around the code you're editing to resolve ambiguity.
+DO NOT omit spans of pre-existing code (or comments) without using the // ... existing code ... comment to indicate its absence. If you omit the existing code comment, the model may inadvertently delete these lines.
+If you plan on deleting a section, you must provide context before and after to delete it. If the initial code is ```code \n Block 1 \n Block 2 \n Block 3 \n code```, and you want to remove Block 2, you would output ```// ... existing code ... \n Block 1 \n  Block 3 \n // ... existing code ...```.
+Make sure it is clear what the edit should be, and where it should be applied.
+Make edits to a file in a single edit_file call instead of multiple edit_file calls to the same file. The apply model can handle many distinct edits at once.
+
+
+RULES: Only return a JSON object nothing else dont even use markdown syntax and strictly no other text or instructions or formatting

gpt_engineer/preprompts/improve

@@ -1,5 +1,5 @@
 Think step by step and reason yourself to the correct decisions to make sure we get it right.
-Make changes to existing code and implement new code in the unified git diff syntax. When implementing new code, First lay out the names of the core classes, functions, methods that will be necessary, As well as a quick comment on their purpose.
+Make changes to existing code. When implementing new code, First lay out the names of the core classes, functions, methods that will be necessary, As well as a quick comment on their purpose.
 
 FILE_FORMAT
 

tests/applications/cli/test_cli_agent.py

@@ -1,7 +1,9 @@
 import os
 import tempfile
+import json
 
 import pytest
+from unittest.mock import patch
 
 from langchain.schema import AIMessage
 
@@ -112,43 +114,45 @@ def test_init_clarified_gen_config(monkeypatch):
     assert code[outfile].strip() == "Hello World!"
 
 
-def test_improve_standard_config(monkeypatch):
+@patch("gpt_engineer.core.default.steps.apply_morph_edit")
+def test_improve_standard_config(mock_apply_morph, monkeypatch): # 2. Pytest gives you the mock as an argument
+    # 3. Tell your fake function what to return when it's called
+    mock_apply_morph.return_value = "!dlroW olleH"
+
+    # The rest of the test setup
     monkeypatch.setattr("builtins.input", lambda _: "y")
     temp_dir = tempfile.mkdtemp()
     code = FilesDict(
         {
-            "main.py": "def write_hello_world_to_file(filename):\n    \"\"\"\n    Writes 'Hello World!' to the specified file.\n    \n    :param filename: The name of the file to write to.\n    \"\"\"\n    with open(filename, 'w') as file:\n        file.write('Hello World!')\n\nif __name__ == \"__main__\":\n    output_filename = 'output.txt'\n    write_hello_world_to_file(output_filename)",
+            "main.py": "def write_hello_world_to_file(filename):\n    file.write('Hello World!')",
             "requirements.txt": "# No dependencies required",
             "run.sh": "python3 main.py\n",
         }
     )
     memory = DiskMemory(memory_path(temp_dir))
-    # version_manager = GitVersionManager(temp_dir)
     execution_env = DiskExecutionEnv()
-    mock_ai = MockAI(
-        [
-            AIMessage(
-                "```diff\n--- main.py\n+++ main.py\n@@ -7,3 +7,3 @@\n     with open(filename, 'w') as file:\n-        file.write('Hello World!')\n+        file.write('!dlroW olleH')\n```"
-            )
-        ]
-    )
+
+    # The mock for the *first* AI call (this part you had correct)
+    mock_response_data = [{
+        "target_file": "main.py",
+        "instructions": "Reverse the string written to the file",
+        "code_edit": "...", # The content here doesn't matter, as the call is mocked
+    }]
+    mock_ai = MockAI([AIMessage(content=json.dumps(mock_response_data))])
+
     cli_agent = CliAgent.with_default_config(memory, execution_env, ai=mock_ai)
 
-    code = cli_agent.improve(
+    # When this improve function is called, it will use your MOCKED apply_morph_edit
+    improved_code = cli_agent.improve(
         code,
-        Prompt(
-            "Change the program so that it prints '!dlroW olleH' instead of 'Hello World!'"
-        ),
+        Prompt("Change the program so that it prints '!dlroW olleH' instead of 'Hello World!'"),
     )
 
-    env = DiskExecutionEnv()
-    env.upload(code).run(f"bash {ENTRYPOINT_FILE}")
-    code = env.download()
-
-    outfile = "output.txt"
-    assert outfile in code
-    assert code[outfile] == "!dlroW olleH"
-
+    # 4. The result of the `improve` call is now the direct output of your mock
+    assert improved_code["main.py"] == "!dlroW olleH"
+
+    # 5. (Best Practice) We can also check that the fake API call was made
+    mock_apply_morph.assert_called_once()
 
 if __name__ == "__main__":
     pytest.main()

tests/core/default/test_simple_agent.py

@@ -1,6 +1,7 @@
 import tempfile
 
 import pytest
+import json
 
 from langchain.schema import AIMessage
 
@@ -47,11 +48,16 @@ def test_improve():
             "run.sh": "python3 main.py\n",
         }
     )
+    mock_response_data = [
+        {
+            "target_file": "main.py",
+            "instructions": "Reverse the string written to the file",
+            "code_edit": "with open(filename, 'w') as file:\n    file.write('!dlroW olleH')",
+        }
+    ]
     mock_ai = MockAI(
         [
-            AIMessage(
-                "```diff\n--- main.py\n+++ main.py\n@@ -7,3 +7,3 @@\n     with open(filename, 'w') as file:\n-        file.write('Hello World!')\n+        file.write('!dlroW olleH')\n```"
-            )
+            AIMessage(content=json.dumps(mock_response_data))
         ]
     )
     lean_agent = SimpleAgent.with_default_config(temp_dir, mock_ai)