File size: 2,584 Bytes
c8a4550 |
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 |
"""Tool system for Data Analyzer Agent.
Provides execute_code function for running Python code in E2B sandbox,
along with tool schema and routing infrastructure.
"""
import json
from typing import Callable
from e2b_code_interpreter import Sandbox
def execute_code(code: str, sbx: Sandbox):
"""Execute Python code in E2B sandbox and return results with metadata.
Args:
code: Python code string to execute
sbx: E2B Sandbox instance
Returns:
Tuple of (execution.to_json(), metadata dict)
- execution.to_json() contains results and errors
- metadata contains images list with PNG data extracted from results
"""
execution = sbx.run_code(code)
metadata = {}
# Extract PNG data from execution results into metadata
if hasattr(execution, "results") and execution.results:
for result in execution.results:
if getattr(result, "png", None):
metadata.setdefault("images", []).append(result.png)
result.png = None # Clear PNG from result to avoid duplication
if execution.error:
metadata["error"] = str(execution.error)
return execution.to_json(), metadata
execute_code_schema = {
"type": "function",
"name": "execute_code",
"description": "Execute Python code and return result",
"parameters": {
"type": "object",
"properties": {
"code": {"type": "string", "description": "Python code"}
},
"required": ["code"],
"additionalProperties": False
}
}
tools = {"execute_code": execute_code}
def execute_tool(name: str, args: str, tools: dict, **kwargs):
"""Route tool calls from LLM to implementation functions.
Args:
name: Tool name to execute
args: JSON string with tool arguments
tools: Dict mapping tool names to functions
**kwargs: Additional parameters passed to tool function (e.g., sbx)
Returns:
Tuple of (result dict, metadata dict)
- result contains execution output or error
- metadata contains additional data like images
"""
try:
args_dict = json.loads(args)
if name not in tools:
return {"error": f"Tool {name} does not exist."}, {}
return tools[name](**args_dict, **kwargs)
except json.JSONDecodeError as e:
return {"error": f"Failed to parse JSON arguments: {str(e)}"}, {}
except KeyError as e:
return {"error": f"Missing key in arguments: {str(e)}"}, {}
except Exception as e:
return {"error": str(e)}, {}
|