agents/code_agent.py

import os
import logging

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

# 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 = os.getenv("OPENAI_API_KEY")

    if not gen_api_key:
        raise ValueError("OPENAI_API_KEY environment variable is not set.")

    # 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,
            reasoning_effort="high",
            temperature=0.1,
            max_tokens=16384
        )
        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", "gemini-2.5-pro-preview-03-25")
    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,
            temperature=0.10
        )
        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.
            
            **Special Instructions for Chess-Related Tasks**:
            - Prioritize using the Stockfish engine to solve chess problems. Ubuntu installation: `sudo apt-get install stockfish` so path is `/usr/games/stockfish`
            - Use `python-chess` to represent boards, generate and validate moves, and parse PGN/FEN.
            
            **Available Python Packages**:
            
            - beautifulsoup4: HTML/XML parsing and lightweight web scraping  
            - certifi: Mozilla CA bundle for secure TLS/SSL requests  
            - datasets: Hugging Face dataset loading and streaming  
            - duckdb: In‑process OLAP SQL engine (analytics, Parquet, Arrow)  
            - ffmpeg-python: Wrapper around FFmpeg for audio/video operations  
            - gradio[oauth]: Rapid web‑UI prototyping with optional OAuth  
            - helium: High‑level Selenium / browser automation toolkit  
            - huggingface: Interact with Hugging Face Hub models, datasets, spaces  
            - imageio: Read and write images, GIFs, MP4s, volumes, etc.  
            - matplotlib: 2‑D plotting (figures, axes, annotations)  
            - numpy: N‑dimensional arrays and vectorized math  
            - openai-whisper: Speech‑to‑text transcription  
            - opencv-python: Computer vision, image/video processing  
            - openpyxl: Excel .xlsx read/write, styles, formulas  
            - pandas: DataFrames, time series, CSV/Parquet I/O  
            - pyarrow: Apache Arrow tables, Parquet, Flight RPC  
            - pygame: Simple 2‑D game/graphics engine (SDL based)  
            - python-chess: Chess move generation, PGN/FEN handling, engine UCI integration  
            - requests: HTTP/HTTPS client with sessions and retries  
            - scikit-learn: Machine‑learning algorithms, preprocessing, pipelines  
            - scipy: Scientific computing, optimization, signal processing  
            - seaborn: Statistical visualization on top of matplotlib  
            - sqlalchemy: SQL ORM and core engine for many databases  
            - statsmodels: Econometrics and statistical modeling (GLM, ARIMA)  
            - stockfish: UCI interface to Stockfish chess engine  
            - sympy: Symbolic math, algebra, calculus CAS  
            - youtube-transcript-api: Fetch YouTube video transcripts via API  
            - yt-dlp: Download videos/playlists from YouTube and other sites  
            """

        system_prompt = load_prompt_from_file("code_agent_system_prompt.txt", default_system_prompt)

        agent = ReActAgent(
            name="code_agent",
            description=(
                "Generates Python code using `python_code_generator` and executes it safely with "
                "`code_interpreter`, then iteratively debugs and refines the code from run-time feedback.\n\n"
                "The agent can leverage the following pre-installed packages:\n"
                "- beautifulsoup4>=4.13.4 : HTML/XML parsing and lightweight web scraping\n"
                "- certifi>=2025.4.26     : Mozilla CA bundle for secure TLS/SSL requests\n"
                "- datasets>=3.5.1        : Hugging Face dataset loading and streaming\n"
                "- duckdb>=1.2.2          : In‑process OLAP SQL engine (analytics, Parquet, Arrow)\n"
                "- ffmpeg-python>=0.2.0   : Wrapper around FFmpeg for audio/video operations\n"
                "- gradio[oauth]>=5.28.0  : Rapid web‑UI prototyping with optional OAuth\n"
                "- helium>=5.1.1          : High‑level Selenium / browser automation toolkit\n"
                "- huggingface>=0.0.1     : Interact with Hugging Face Hub models, datasets, spaces\n"
                "- imageio>=2.37.0        : Read and write images, GIFs, MP4s, volumes, etc.\n"
                "- matplotlib>=3.10.1     : 2‑D plotting (figures, axes, annotations)\n"
                "- numpy>=2.2.5           : N‑dimensional arrays and vectorized math\n"
                "- openai-whisper>=20240930 : Speech‑to‑text transcription\n"
                "- opencv-python>=4.11.0.86 : Computer vision, image/video processing\n"
                "- openpyxl>=3.1.5        : Excel .xlsx read/write, styles, formulas\n"
                "- pandas>=2.2.3          : DataFrames, time series, CSV/Parquet I/O\n"
                "- pyarrow>=20.0.0        : Apache Arrow tables, Parquet, Flight RPC\n"
                "- pygame>=2.6.1          : Simple 2‑D game/graphics engine (SDL based)\n"
                "- python-chess>=1.999    : Chess move generation, PGN/FEN handling, engines\n"
                "- requests>=2.32.3       : HTTP/HTTPS client with sessions and retries\n"
                "- scikit-learn>=1.6.1    : Machine‑learning algorithms, preprocessing, pipelines\n"
                "- scipy>=1.15.2          : Scientific computing, optimization, signal processing\n"
                "- seaborn>=0.13.2        : Statistical visualization on top of matplotlib\n"
                "- sqlalchemy>=2.0.40     : SQL ORM and core engine for many databases\n"
                "- statsmodels>=0.14.4    : Econometrics and statistical modeling (GLM, ARIMA)\n"
                "- stockfish==3.28.0      : UCI interface to Stockfish chess engine\n"
                "- sympy>=1.14.0          : Symbolic math, algebra, calculus CAS\n"
                "- youtube-transcript-api>=1.0.3 : Fetch YouTube video transcripts via API\n"
                "- yt-dlp>=2025.3.31      : Download videos/playlists from YouTube and other sites\n\n"
                "Additionally, the `stockfish` package enables the agent to solve chess problems by analyzing positions, "
                "identifying tactical motifs, and calculating optimal move sequences, making it a valuable tool for chess training and analysis."
            ),
            # 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}")