code_interpreter.py

"""
Code Interpreter Tool
Allows the agent to execute Python code for calculations and data processing.
"""

import io
import sys
from smolagents import Tool
from contextlib import redirect_stdout, redirect_stderr

class CodeInterpreterTool(Tool):
    name = "code_interpreter"
    description = "Executes Python code and returns the output. Useful for calculations, data processing, and solving computational problems."
    inputs = {
        "code": {
            "type": "string",
            "description": "The Python code to execute. Should be a complete, runnable code snippet."
        }
    }
    output_type = "string"

    def forward(self, code: str) -> str:
        """
        Execute Python code safely and return the output.
        
        Args:
            code: Python code string to execute
            
        Returns:
            Output from code execution or error message
        """
        # Security: Block dangerous operations
        dangerous_keywords = [
            "import os", "import sys", "import subprocess",
            "__import__", "eval", "exec", "open(",
            "file(", "input(", "raw_input(",
        ]
        
        code_lower = code.lower()
        for keyword in dangerous_keywords:
            if keyword.lower() in code_lower:
                return f"Error: Potentially dangerous operation '{keyword}' is not allowed."
        
        try:
            # Capture stdout and stderr
            stdout_capture = io.StringIO()
            stderr_capture = io.StringIO()
            
            # Redirect output
            with redirect_stdout(stdout_capture), redirect_stderr(stderr_capture):
                # Execute the code
                exec(code, {"__builtins__": __builtins__})
            
            stdout_output = stdout_capture.getvalue()
            stderr_output = stderr_capture.getvalue()
            
            if stderr_output:
                return f"Error: {stderr_output}"
            
            if stdout_output:
                return stdout_output.strip()
            else:
                return "Code executed successfully (no output)."
                
        except Exception as e:
            return f"Error executing code: {str(e)}"