code_agent.py

import os
import tempfile
import subprocess
import logging
from typing import Optional, Dict, Any

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class CodeInterpreter:
    """Simple code interpreter for executing Python code safely."""
    
    def __init__(self, working_dir: Optional[str] = None):
        """Initialize the code interpreter with a working directory."""
        if working_dir is None:
            self.working_dir = tempfile.mkdtemp()
        else:
            self.working_dir = working_dir
            os.makedirs(working_dir, exist_ok=True)
        
        logger.info(f"Initialized CodeInterpreter with working directory: {self.working_dir}")
    
    def execute(self, code: str, language: str = "python") -> Dict[str, Any]:
        """Execute code and return results."""
        try:
            if language.lower() != "python":
                return {
                    "status": "error",
                    "stdout": "",
                    "stderr": f"Language '{language}' not supported. Only Python is supported.",
                    "plots": [],
                    "dataframes": []
                }
            
            # Create a temporary file for the code
            with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False, dir=self.working_dir) as f:
                f.write(code)
                temp_file = f.name
            
            try:
                # Execute the code
                result = subprocess.run(
                    ["python", temp_file],
                    capture_output=True,
                    text=True,
                    timeout=30,  # 30 second timeout
                    cwd=self.working_dir
                )
                
                status = "success" if result.returncode == 0 else "error"
                
                return {
                    "status": status,
                    "stdout": result.stdout,
                    "stderr": result.stderr,
                    "plots": [],  # Could be extended to detect plot files
                    "dataframes": []  # Could be extended to detect CSV outputs
                }
                
            finally:
                # Clean up the temporary file
                try:
                    os.unlink(temp_file)
                except OSError:
                    pass
                    
        except subprocess.TimeoutExpired:
            return {
                "status": "error",
                "stdout": "",
                "stderr": "Code execution timed out (30 seconds)",
                "plots": [],
                "dataframes": []
            }
        except Exception as e:
            logger.error(f"Error executing code: {str(e)}")
            return {
                "status": "error",
                "stdout": "",
                "stderr": str(e),
                "plots": [],
                "dataframes": []
            }

class CodeInterpreterTool:
    """Tool wrapper for the code interpreter."""
    
    def __init__(self, working_directory: Optional[str] = None):
        """Initialize the code interpreter tool."""
        # Use absolute path without special characters
        default_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "code_outputs"))
        self.interpreter = CodeInterpreter(
            working_dir=working_directory or default_dir
        )

    def execute(self, code: str, language: str = "python") -> Dict[str, Any]:
        """Execute code and return results."""
        try:
            logger.info(f"Executing {language} code")
            result = self.interpreter.execute(code, language)
            
            # Format the response
            response = {
                "status": result["status"],
                "output": result["stdout"],
                "error": result["stderr"] if result["status"] == "error" else None,
                "plots": result.get("plots", []),
                "dataframes": result.get("dataframes", [])
            }
            
            return response
        except Exception as e:
            logger.error(f"Error executing code: {str(e)}")
            return {
                "status": "error",
                "error": str(e),
                "output": None,
                "plots": [],
                "dataframes": []
            }