agents/code_agent.py

import os
import logging
from dotenv import load_dotenv

from llama_index.core.agent.workflow import CodeActAgent, ReActAgent
from llama_index.core.tools import FunctionTool
from llama_index.llms.google_genai import GoogleGenAI
from llama_index.llms.openai import OpenAI
from llama_index.tools.code_interpreter import CodeInterpreterToolSpec

# Load environment variables
load_dotenv()

# Setup logging
logger = logging.getLogger(__name__)

# Helper function to load prompt from file
def load_prompt_from_file(filename: str, default_prompt: str) -> str:
    """Loads a prompt from a text file."""
    try:
        # Assuming the prompt file is in the same directory as the agent script
        script_dir = os.path.dirname(__file__)
        prompt_path = os.path.join(script_dir, filename)
        with open(prompt_path, "r") as f:
            prompt = f.read()
            logger.info(f"Successfully loaded prompt from {prompt_path}")
            return prompt
    except FileNotFoundError:
        logger.warning(f"Prompt file {filename} not found at {prompt_path}. Using default.")
        return default_prompt
    except Exception as e:
        logger.error(f"Error loading prompt file {filename}: {e}", exc_info=True)
        return default_prompt

def generate_python_code(prompt: str) -> str:
    """
    Generate valid Python code from a natural language description using a configured LLM.
    Args:
        prompt (str): A clear description of the desired Python code functionality.
    Returns:
        str: A string containing the generated Python code.
    Raises:
        ValueError: If required API key is not set.
        Exception: If the LLM call fails.
    """
    logger.info(f"Generating Python code for prompt: {prompt[:100]}...")
    
    # Configuration for code generation LLM
    gen_llm_model = os.getenv("CODE_GEN_LLM_MODEL", "o4-mini")
    gen_api_key_env = os.getenv("CODE_GEN_API_KEY_ENV", "ALPAFLOW_OPENAI_API_KEY")
    gen_api_key = os.getenv(gen_api_key_env)

    if not gen_api_key:
        logger.error(f"{gen_api_key_env} not found in environment variables for code generation LLM.")
        raise ValueError(f"{gen_api_key_env} must be set for code generation")

    # Load the prompt template
    default_gen_prompt_template = ("You are a helpful assistant that writes Python code. "
                                   "You will be given a prompt and you must generate Python code based on that prompt. "
                                   "You must only generate Python code and nothing else. "
                                   "Do not include any explanations or any other text. "
                                   "Do not use any markdown. \n"
                                   "Prompt: {prompt} \n"
                                   "Code:\n")
    gen_prompt_template = load_prompt_from_file("../prompts/code_gen_prompt.txt", default_gen_prompt_template)
    input_prompt = gen_prompt_template.format(prompt=prompt)

    try:
        llm = OpenAI(
            model=gen_llm_model,
            api_key=gen_api_key
        )
        logger.info(f"Using code generation LLM: {gen_llm_model}")
        generated_code = llm.complete(input_prompt)
        logger.info("Code generation successful.")
        return generated_code.text
    except Exception as e:
        logger.error(f"LLM call failed during code generation: {e}", exc_info=True)
        raise # Re-raise the exception to be handled by the agent/workflow

# --- Tool Definitions ---

python_code_generator_tool = FunctionTool.from_defaults(
    fn=generate_python_code,
    name="python_code_generator",
    description=(
        "Generates executable Python code based on a natural language prompt. "
        "Input: prompt string. Output: Python code string."
    ),
)

# Use LlamaIndex's built-in Code Interpreter Tool Spec for safe execution
# This assumes the necessary environment (e.g., docker) for the spec is available
try:
    code_interpreter_spec = CodeInterpreterToolSpec()
    # Get the tool(s) from the spec. It might return multiple tools.
    code_interpreter_tools = code_interpreter_spec.to_tool_list()
    if not code_interpreter_tools:
        raise RuntimeError("CodeInterpreterToolSpec did not return any tools.")
    # Assuming the primary tool is the first one, or find by name if necessary
    code_interpreter_tool = next((t for t in code_interpreter_tools if t.metadata.name == "code_interpreter"), None)
    if code_interpreter_tool is None:
         raise RuntimeError("Could not find 'code_interpreter' tool in CodeInterpreterToolSpec results.")
    logger.info("CodeInterpreterToolSpec initialized successfully.")
except Exception as e:
    logger.error(f"Failed to initialize CodeInterpreterToolSpec: {e}", exc_info=True)
    # Fallback: Define a dummy tool or raise error to prevent agent start?
    # For now, let initialization fail if the safe interpreter isn't available.
    raise RuntimeError("CodeInterpreterToolSpec failed to initialize. Cannot create code_agent.") from e

# --- REMOVED SimpleCodeExecutor --- 
# The SimpleCodeExecutor class that used subprocess has been entirely removed 
# due to severe security risks. Execution MUST go through the CodeInterpreterToolSpec.

# --- Agent Initialization ---

def initialize_code_agent() -> ReActAgent:
    """Initializes the CodeActAgent, configured for safe code execution."""
    logger.info("Initializing CodeAgent...")

    # Configuration for the agent's main LLM
    agent_llm_model = os.getenv("CODE_AGENT_LLM_MODEL", "models/gemini-1.5-pro")
    gemini_api_key = os.getenv("GEMINI_API_KEY")

    if not gemini_api_key:
        logger.error("GEMINI_API_KEY not found in environment variables for CodeAgent.")
        raise ValueError("GEMINI_API_KEY must be set for CodeAgent")

    try:
        llm = GoogleGenAI(
            api_key=gemini_api_key,
            model=agent_llm_model,
        )
        logger.info(f"Using agent LLM: {agent_llm_model}")

        # Load system prompt (consider loading from file)
        default_system_prompt = """\
        You are CodeAgent, a specialist in generating and executing Python code. Your mission:
        
        1. **Thought**: Think step-by-step before acting and state your reasoning.
        2. **Code Generation**: To produce code, call `python_code_generator` with a concise, unambiguous prompt. Review the generated code for correctness and safety.
        3. **Execution & Testing**: To execute or test code, call `code_interpreter`. Provide the complete code snippet. Analyze its output (stdout, stderr, result) to verify functionality and debug errors.
        4. **Iteration**: If execution fails or the result is incorrect, analyze the error, think about the fix, generate corrected code using `python_code_generator`, and execute again using `code_interpreter`.
        5. **Tool Use**: Always adhere strictly to each tool’s input/output format.
        6. **Final Output**: Once the code works correctly and achieves the goal, output *only* the final functional code or the final execution result, as appropriate for the task.
        7. **Hand-Off**: If further logical reasoning or verification is needed, delegate to **reasoning_agent**. Otherwise, pass your final output to **planner_agent** for synthesis.
        """
        # system_prompt = load_prompt_from_file("code_agent_system_prompt.txt", default_system_prompt)
        system_prompt = default_system_prompt # Using inline for now

        agent = ReActAgent(
            name="code_agent",
            description=(
                "Generates Python code using `python_code_generator` and executes it safely using `code_interpreter`. "
                "Iteratively debugs and refines code based on execution results."
            ),
            # REMOVED: code_execute_fn - Execution is handled by the code_interpreter tool via the agent loop.
            tools=[
                python_code_generator_tool,
                code_interpreter_tool, # Use the safe tool from the spec
            ],
            llm=llm,
            system_prompt=system_prompt,
            can_handoff_to=["planner_agent", "reasoning_agent"],
        )
        logger.info("CodeAgent initialized successfully.")

        return agent

    except Exception as e:
        logger.error(f"Error during CodeAgent initialization: {e}", exc_info=True)
        raise

# Example usage (for testing if run directly)
if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    logger.info("Running code_agent.py directly for testing...")

    # Ensure API keys are set for testing
    required_keys = ["GEMINI_API_KEY", os.getenv("CODE_GEN_API_KEY_ENV", "ALPAFLOW_OPENAI_API_KEY")]
    missing_keys = [key for key in required_keys if not os.getenv(key)]
    if missing_keys:
        print(f"Error: Required environment variable(s) not set: {', '.join(missing_keys)}. Cannot run test.")
    else:
        try:
            test_agent = initialize_code_agent()
            print("Code Agent initialized successfully for testing.")
            # Example test (requires user interaction or pre-defined task)
            # result = test_agent.chat("Write and execute python code to print 'hello world'")
            # print(f"Test query result: {result}")
        except Exception as e:
            print(f"Error during testing: {e}")