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": []
} |