File size: 4,390 Bytes
04ffb15
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
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": []
            }