Overhauled Code Interpreter tool by wrapping it in a new class to make it robust to agent output

mycoder/src/inner/config/agents.yaml

@@ -6,6 +6,10 @@ coder:
     First you plan how the code will work, then you write the code, then you run it and check the output.
     You should check schema of uploaded files to understand their structure and proceed with analyses more
     efficiently.
+    When you use the Code Interpreter tool, assign the final markdown report to a variable named `result`.
+    Do not rely on print statements alone. If you need dataframe info text, capture `df.info()` with `io.StringIO()`.
+    If you compute correlations, use numeric-only columns.
+    Once code execution succeeds and you have the output you need, stop rerunning the tool and write the final answer.
     In some cases, no code generation is required, and you can simply provide the necessary results in markdown format.
   backstory: >
     You're a seasoned python developer and bioinformatician with a knack for writing clean, efficient code.

mycoder/src/inner/config/tasks.yaml

@@ -4,6 +4,8 @@ coding_task:
 
     Make sure to save any plots generated by code as 'output/plot.png'.
     Do not use seaborn for any plots/visualizations.
+    If you use Code Interpreter, assign the final markdown report to a variable named `result`.
+    Do not rely on print statements alone.
   expected_output: >
     A markdown file that includes the code itself, any comments/explanation, 
     and the output of the code. Alternatively, if no code generation is required,

mycoder/src/inner/crew.py

@@ -16,9 +16,12 @@ class Coder():
 
     @agent
     def coder(self) -> Agent:
-        from src.inner.tools.custom_tool import ExtractSchema
-        from crewai_tools import CodeInterpreterTool
-        code_interpreter = CodeInterpreterTool(unsafe_mode=True)
+        from src.inner.tools.custom_tool import (
+            ExtractSchema,
+            ResultAwareCodeInterpreterTool,
+        )
+
+        code_interpreter = ResultAwareCodeInterpreterTool(unsafe_mode=True)
         return Agent(
             config=self.agents_config['coder'],
             verbose=True,

mycoder/src/inner/tools/custom_tool.py

@@ -1,8 +1,12 @@
-import pandas as pd
-import numpy as np
+import io
+from contextlib import redirect_stderr, redirect_stdout
+from typing import Any, Dict, Type
 
+import numpy as np
+import pandas as pd
 from crewai.tools import BaseTool
-from typing import Any, Dict, Type
+from crewai_tools import CodeInterpreterTool
+from crewai_tools.tools.code_interpreter_tool.code_interpreter_tool import SandboxPython
 from pydantic import BaseModel, Field
 
 
@@ -32,3 +36,69 @@ class ExtractSchema(BaseTool):
             "numeric_cols": df.select_dtypes(include=[np.number]).columns.tolist(),
             "categorical_cols": df.select_dtypes(include=['object']).columns.tolist()
         }
+
+
+class ResultAwareCodeInterpreterTool(CodeInterpreterTool):
+    """Returns `result` when present and falls back to captured stdout/stderr."""
+
+    @staticmethod
+    def _resolve_execution_output(
+        exec_locals: Dict[str, Any],
+        stdout_buffer: io.StringIO,
+        stderr_buffer: io.StringIO,
+    ) -> str:
+        result = exec_locals.get("result")
+        if result is not None:
+            return str(result)
+
+        stdout_text = stdout_buffer.getvalue().strip()
+        stderr_text = stderr_buffer.getvalue().strip()
+
+        if stdout_text and stderr_text:
+            return f"{stdout_text}\n\n[stderr]\n{stderr_text}"
+        if stdout_text:
+            return stdout_text
+        if stderr_text:
+            return stderr_text
+        return "Code executed successfully, but it did not set `result` or print output."
+
+    @staticmethod
+    def run_code_in_restricted_sandbox(code: str) -> str:
+        stdout_buffer = io.StringIO()
+        stderr_buffer = io.StringIO()
+        exec_locals: Dict[str, Any] = {}
+
+        try:
+            with redirect_stdout(stdout_buffer), redirect_stderr(stderr_buffer):
+                SandboxPython.exec(code=code, locals_=exec_locals)
+        except Exception as e:
+            return f"An error occurred: {e!s}"
+
+        return ResultAwareCodeInterpreterTool._resolve_execution_output(
+            exec_locals,
+            stdout_buffer,
+            stderr_buffer,
+        )
+
+    @staticmethod
+    def run_code_unsafe(code: str, libraries_used: list[str]) -> str:
+        stdout_buffer = io.StringIO()
+        stderr_buffer = io.StringIO()
+
+        for library in libraries_used:
+            import os
+
+            os.system(f"pip install {library}")  # noqa: S605
+
+        try:
+            exec_locals: Dict[str, Any] = {}
+            with redirect_stdout(stdout_buffer), redirect_stderr(stderr_buffer):
+                exec(code, {}, exec_locals)  # noqa: S102
+        except Exception as e:
+            return f"An error occurred: {e!s}"
+
+        return ResultAwareCodeInterpreterTool._resolve_execution_output(
+            exec_locals,
+            stdout_buffer,
+            stderr_buffer,
+        )