From 88b35298e2c0244d6d7ff13e1dd6ec01f205c2c9 Mon Sep 17 00:00:00 2001
From: Dillon Laird <dillonalaird@gmail.com>
Date: Mon, 9 Sep 2024 09:18:40 -0700
Subject: [PATCH] added function to strip results from remote code execution
 call

---
 vision_agent/agent/vision_agent_coder.py | 21 +++++++++++++++------
 vision_agent/utils/execute.py            | 22 ++++++++++++----------
 2 files changed, 27 insertions(+), 16 deletions(-)

diff --git a/vision_agent/agent/vision_agent_coder.py b/vision_agent/agent/vision_agent_coder.py
index 3c11c730..195806cc 100644
--- a/vision_agent/agent/vision_agent_coder.py
+++ b/vision_agent/agent/vision_agent_coder.py
@@ -174,9 +174,10 @@ def pick_plan(
         }
     )
     tool_output = code_interpreter.exec_isolation(DefaultImports.prepend_imports(code))
-    tool_output_str = ""
-    if len(tool_output.text().strip()) > 0:
-        tool_output_str = tool_output.text().strip()
+    # Because of the way we trace function calls the trace information ends up in the
+    # results. We don't want to show this info to the LLM so we don't include it in the
+    # tool_output_str.
+    tool_output_str = tool_output.text(include_results=False).strip()
 
     if verbosity == 2:
         _print_code("Initial code and tests:", code)
@@ -203,7 +204,7 @@ def pick_plan(
             docstring=tool_info,
             plans=plan_str,
             previous_attempts=PREVIOUS_FAILED.format(
-                code=code, error=tool_output.text()
+                code=code, error="\n".join(tool_output_str.splitlines()[-50:])
             ),
             media=media,
         )
@@ -232,7 +233,7 @@ def pick_plan(
                 "status": "completed" if tool_output.success else "failed",
             }
         )
-        tool_output_str = tool_output.text().strip()
+        tool_output_str = tool_output.text(include_results=False).strip()
 
         if verbosity == 2:
             _print_code("Code and test after attempted fix:", code)
@@ -412,6 +413,7 @@ def write_and_test_code(
             working_memory,
             debugger,
             code_interpreter,
+            tool_info,
             code,
             test,
             result,
@@ -437,6 +439,7 @@ def debug_code(
     working_memory: List[Dict[str, str]],
     debugger: LMM,
     code_interpreter: CodeInterpreter,
+    tool_info: str,
     code: str,
     test: str,
     result: Execution,
@@ -461,9 +464,15 @@ def debug_code(
             # followed by code each wrapped in markdown blocks.
             fixed_code_and_test_str = debugger(
                 FIX_BUG.format(
+                    docstring=tool_info,
                     code=code,
                     tests=test,
-                    result="\n".join(result.text().splitlines()[-50:]),
+                    # Because of the way we trace function calls the trace information
+                    # ends up in the results. We don't want to show this info to the
+                    # LLM so we don't include it in the tool_output_str.
+                    result="\n".join(
+                        result.text(include_results=False).splitlines()[-50:]
+                    ),
                     feedback=format_memory(working_memory + new_working_memory),
                 ),
                 stream=False,
diff --git a/vision_agent/utils/execute.py b/vision_agent/utils/execute.py
index 7967367e..c2e0e652 100644
--- a/vision_agent/utils/execute.py
+++ b/vision_agent/utils/execute.py
@@ -292,7 +292,7 @@ class Config:
     error: Optional[Error] = None
     "Error object if an error occurred, None otherwise."
 
-    def text(self, include_logs: bool = True) -> str:
+    def text(self, include_logs: bool = True, include_results: bool = True) -> str:
         """Returns the text representation of this object, i.e. including the main
         result or the error traceback, optionally along with the logs (stdout, stderr).
         """
@@ -300,15 +300,17 @@ def text(self, include_logs: bool = True) -> str:
         if self.error:
             return prefix + "\n----- Error -----\n" + self.error.traceback
 
-        result_str = [
-            (
-                f"----- Final output -----\n{res.text}"
-                if res.is_main_result
-                else f"----- Intermediate output-----\n{res.text}"
-            )
-            for res in self.results
-        ]
-        return prefix + "\n" + "\n".join(result_str)
+        if include_results:
+            result_str = [
+                (
+                    f"----- Final output -----\n{res.text}"
+                    if res.is_main_result
+                    else f"----- Intermediate output-----\n{res.text}"
+                )
+                for res in self.results
+            ]
+            return prefix + "\n" + "\n".join(result_str)
+        return prefix
 
     @property
     def success(self) -> bool: