agents/executor_agent.py

# 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.")