File size: 4,935 Bytes
557ed35 8d66edb 557ed35 8d66edb 557ed35 8d66edb 557ed35 8d66edb 557ed35 8d66edb 557ed35 8d66edb 557ed35 a320529 |
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 |
# agents/executor_agent.py
import io
import contextlib
import json
# IMPORTANT: We will dynamically import the tools module when needed for execution
# to ensure it uses the version from the tools/ directory.
# import tools.agent_tools as agent_tools # Avoid top-level import for now
class ExecutorAgent:
def __init__(self, openai_api_key: str = None):
print("ExecutorAgent initialized.")
self.openai_api_key = openai_api_key
def execute_code(self, python_code: str) -> dict:
print(f"ExecutorAgent received code for execution:\n{python_code}")
# Dynamically import agent_tools from the tools directory
# This assumes main.py is run from the project root.
import sys
import os
# Add project root to sys.path to allow `import tools.agent_tools`
# This might be needed if executor_agent.py itself is run from a different context later,
# but for now, assuming standard Python module resolution from root where main.py is.
# script_dir = os.path.dirname(os.path.abspath(__file__))
# project_root = os.path.abspath(os.path.join(script_dir, ".."))
# if project_root not in sys.path:
# sys.path.insert(0, project_root)
try:
# Ensure tools.agent_tools can be imported relative to project root
import tools.agent_tools as agent_tools_module
except ImportError as e:
return {
"execution_output": f"ExecutorAgent Error: Could not import agent_tools module. Ensure it's in tools/ and __init__.py might be needed in tools/. Error: {e}",
"execution_status": "ERROR: ImportFailure"
}
except Exception as e:
return {
"execution_output": f"ExecutorAgent Error: Unexpected error during tools import. Error: {e}",
"execution_status": "ERROR: ImportFailure"
}
# Create a restricted global scope for exec()
# Only allow access to the agent_tools module (aliased as 'tools') and builtins
restricted_globals = {
"__builtins__": __builtins__, # Standard builtins (print, len, etc.)
"tools": agent_tools_module,
"json": json,
"api_key": self.openai_api_key
}
# No separate locals, exec will use restricted_globals as locals too
# Create usage collector for tracking OpenAI API calls in agent_tools
import builtins
builtins.__agent_usage_collector__ = []
captured_output = io.StringIO()
try:
with contextlib.redirect_stdout(captured_output):
exec(python_code, restricted_globals)
output_str = captured_output.getvalue()
# Extract collected usage info
usage_list = builtins.__agent_usage_collector__
aggregated_usage = {
'prompt_tokens': sum(u.get('prompt_tokens', 0) for u in usage_list),
'completion_tokens': sum(u.get('completion_tokens', 0) for u in usage_list),
'total_tokens': sum(u.get('total_tokens', 0) for u in usage_list)
}
return {
"execution_output": output_str.strip() if output_str else "(No output printed by code)",
"execution_status": "SUCCESS",
"usage": aggregated_usage
}
except Exception as e:
error_details = f"{type(e).__name__}: {str(e)}"
# Try to get traceback if possible, though might be complex to format cleanly here
# Extract usage even on error (API calls may have occurred before failure)
usage_list = builtins.__agent_usage_collector__
aggregated_usage = {
'prompt_tokens': sum(u.get('prompt_tokens', 0) for u in usage_list),
'completion_tokens': sum(u.get('completion_tokens', 0) for u in usage_list),
'total_tokens': sum(u.get('total_tokens', 0) for u in usage_list)
}
return {
"execution_output": f"Execution Error!\n{error_details}",
"execution_status": f"ERROR: {type(e).__name__}",
"usage": aggregated_usage
}
if __name__ == '__main__':
# For testing individual agent if needed
# executor = ExecutorAgent()
# result = executor.execute_code("print(tools.get_biorxiv_paper_url())")
# print(result)
# result_error = executor.execute_code("print(tools.non_existent_tool())")
# print(result_error)
# result_unsafe = executor.execute_code("import os\nprint('dangerous')")
# print(result_unsafe) # Should fail at exec if globals are well-restricted from direct os import
print("ExecutorAgent should be orchestrated by the ManagerAgent.") |