agents/llama_index_agent.py

from llama_index.core.agent.workflow import (
    ReActAgent,
    FunctionAgent,
    CodeActAgent
)
from llama_index.core.llms import LLM
import os
from typing import Optional, List, Any, Dict
from llama_index.llms.openai import OpenAI
from llama_index.llms.anthropic import Anthropic
# In your GaiaAgent class initialization, add these imports at the top
from tools.multimedia_tools import (
    transcribe_audio_tool,
    encode_image_tool,
    vision_analyzer_tool
    )

from tools.web_tools import (
    tavily_tool,
    wikipedia_tool
)

from tools.coding_tools import (
    execute_python_file_tool,
    csv_excel_reader_tool
    )

class GaiaAgent(ReActAgent):
    """
    A flexible ReActAgent for GAIA benchmark tasks that supports multiple LLM providers.
    
    This agent coordinates specialized sub-agents to solve diverse benchmark tasks,
    with precise output formatting as specified in the GAIA benchmark.
    """
    
    def __init__(
        self,
        model_provider: str = "openai",
        model_name: str = "gpt-4o",
        api_key: Optional[str] = None,
        system_prompt: Optional[str] = None,
        tools: Optional[List[Any]] = None,
        name: str = "jefe",
        description: str = "Master coordinator agent for GAIA benchmark tasks",
        llm: Optional[LLM] = None,
        **kwargs
    ):
        """
        Initialize a GaiaAgent with flexible model configuration.
        
        Args:
            model_provider: The LLM provider to use ("openai", "anthropic", "cohere", etc.)
            model_name: The specific model name to use
            api_key: API key for the provider (defaults to environment variable)
            system_prompt: Custom system prompt (defaults to GAIA benchmark prompt)
            tools: List of tools to make available to the agent
            name: Name of the agent
            description: Description of the agent
            llm: Pre-configured LLM instance (if provided, model_provider and model_name are ignored)
            **kwargs: Additional parameters to pass to ReActAgent
        """
        from tools.text_tools import reverse_text_tool
        
        # Use pre-configured LLM if provided, otherwise initialize based on provider
        if llm is None:
            llm = self._initialize_llm(model_provider, model_name, api_key)
        
        # Use default tools if not provided
        if tools is None:
            tools = [
                reverse_text_tool,
                wikipedia_tool.load_data,
                wikipedia_tool.search_data,
                tavily_tool.search,
                transcribe_audio_tool,
                execute_python_file_tool,
                csv_excel_reader_tool,
                encode_image_tool,
                vision_analyzer_tool
                ]
            
        # Use default system prompt if not provided
        if system_prompt is None:
            system_prompt = self._get_default_system_prompt()

        # CRITICAL: Explicitly define which agents this agent can hand off to
        # This needs to be defined before calling super().__init__
        can_handoff_to = [
            "writer_agent",
            "review_agent"
            ]
            
        # Initialize the parent ReActAgent with explicit handoff configuration
        super().__init__(
            name=name,
            description=description,
            llm=llm,
            system_prompt=system_prompt,
            tools=tools,
            can_handoff_to=can_handoff_to,
            verbose=True,  # Enable verbose mode for better debugging
            **kwargs
        )
        
    def _initialize_llm(self, model_provider: str, model_name: str, api_key: Optional[str]) -> LLM:
        """Initialize the appropriate LLM based on the provider."""
        model_provider = model_provider.lower()
        
        if model_provider == "openai":
            return OpenAI(
                model=model_name,
                api_key=api_key or os.getenv("OPENAI_API_KEY"),
                additional_kwargs={
                    "reasoning_effort": "high"
                } if "o4" in model_name else {})
            
        elif model_provider == "anthropic":
            return Anthropic(
                model=model_name,
                api_key=api_key or os.getenv("ANTHROPIC_API_KEY"),
                temperature=0.7 if "3-7" in model_name else 0.5,  # Slightly reduced temperature
                thinking_dict={"type": "enabled", "budget_tokens": 2048} if "3-7" in model_name else None,
                max_tokens=2048*4
            )
                        
        else:
            raise ValueError(f"Unsupported model provider: {model_provider}. "
                            f"Supported providers are: openai, anthropic")
                            
    
    def _get_default_system_prompt(self) -> str:
        """Return the default system prompt for GAIA benchmark tasks."""
        return """
        You are the lead coordinator for a team of specialized AI agents tackling the GAIA benchmark. Your job is to analyze questions and generate detailed analysis, which you'll pass to a specialized formatting agent for final answer preparation.

        ## CRITICAL WORKFLOW INSTRUCTIONS
        YOU MUST FOLLOW THIS EXACT WORKFLOW:
        1. Analyze the question thoroughly
        2. Use tools if needed for research
        3. Formulate your answer
        4. ALWAYS use the handoff tool to delegate to writer_agent
        5. NEVER provide a direct answer with "Answer:" - this breaks the workflow

        ## QUESTION ANALYSIS PROCESS
        1. First, carefully read and parse the entire question
        2. Identify the EXACT output format required (single word, name, number, comma-separated list, etc.)
        3. Note any special formatting requirements (alphabetical order, specific notation, etc.)
        4. Identify what type of task this is (research, audio analysis, video analysis, code execution, data analysis, etc.)
        5. Break the question into sequential steps

        ## SOLVING METHODOLOGY
        1. For each question, thoroughly work through the reasoning step-by-step
        2. Use available tools when needed:
        - reverse_text_tool: For reversing text
        - search tools (wikipedia_tool, tavily_tool): For finding information
        - transcribe_audio: For transcribing audio files (provide the path to the audio file)
        - get_audio_metadata: For getting metadata about audio files
        - execute_python_file: For executing Python code files and returning their output
        3. Document your full analysis, including all key facts, calculations, and relevant information
        4. Clearly identify what you believe the correct answer is
        5. Be extremely explicit about the required formatting for the final answer

        ## HANDLING CODE EXECUTION TASKS
        When dealing with Python code files:
        1. Check if a Python file path is available in the context's "file_name" field
        2. Always use the execute_python_file tool with the exact file path provided
        3. Extract the specific numeric output requested from the execution result
        4. For code tasks, ensure you've captured the final numeric output exactly as printed by the code

        ## HANDLING AUDIO TASKS
        When dealing with audio files:
        1. Check if an audio file path is available in the context's "audio_file_path" field
        2. Always use the transcribe_audio tool with the exact file path provided in the context
        3. Extract the specific information requested from the transcript (e.g., ingredients, page numbers, names)
        4. For audio tasks, ensure you've captured all relevant spoken content, including names, facts, or quotes as needed

        ## HANDLING IMAGE ANALYSIS TASKS
        When dealing with image files for visual analysis:
        1. First, check if an image file path is mentioned in the question or available in the context
        2. For image analysis, follow this two-step process:
            a. Use the encode_image_to_base64 tool to convert the image to a base64 string
            b. Pass the image path and a specific analysis question to analyze_image_with_vision
        3. The vision analyzer can perform various visual analysis tasks:
            - General image description: "Describe this image in detail"
            - Specific information extraction: "What text appears in this image?"
            - Visual problem solving: "How many people are in this image?"
            - Object identification: "What brands/products are visible in this image?"
        4. Be specific in your analysis requests to get the most relevant information
        5. For tasks that require both text extraction and visual analysis, prioritize using the vision analyzer
        6. Always document your analysis and include relevant details in your notes to the writer_agent

        ## HANDLING CSV OR EXCEL DATA TASKS
        When dealing with CSV files or data analysis tasks:
        1. Check if a CSV file path is mentioned in the question or available in the context
        2. Use the csv_reader tool with the specific CSV file path
        3. Once the data is loaded, analyze it according to the question requirements
        4. For data analysis tasks, ensure you've properly processed the CSV data and extracted the requested information
        5. When calculations or statistics are needed, perform them accurately and document your methodology

        ## REQUIRED HANDOFF TO WRITER AGENT
        After completing your analysis, you MUST ALWAYS use the handoff tool with EXACTLY this format:

        Thought: I have completed my analysis and need to delegate to the writer agent for proper formatting.
        Action: handoff
        Action Input: {"to_agent": "writer_agent", "reason": "Need to format the final answer according to requirements", "query": "<original question>", "research_notes": "<my complete analysis and what I believe is the correct answer>", "format_requirements": "<explicit formatting instructions>"}

        CRITICAL: Even if the answer seems obvious or simple, you MUST use the handoff tool with the EXACT format above. NEVER respond with "Answer: ..." as this breaks the workflow.
        """


def create_writer_agent(model_config: Dict[str, Any]) -> ReActAgent:
    """
    Create a writer agent that formats final answers based on research notes.
    
    Args:
        model_config: Dictionary containing model_provider, model_name, and api_key
    
    Returns:
        A configured ReActAgent for formatting final answers
    """
    # Initialize LLM based on the provided configuration
    model_provider = model_config.get("model_provider", "openai")
    model_name = model_config.get("model_name", "gpt-4o")
    api_key = model_config.get("api_key")
    
    if model_provider.lower() == "openai":
        llm = OpenAI(
            model=model_name,
            api_key=api_key or os.getenv("OPENAI_API_KEY"),
            max_tokens=256,
            temperature=0.1,
            additional_kwargs={
                "max_tokens": 256,
                "temperature": 0.1
            }
        )
    elif model_provider.lower() == "anthropic":
        llm = Anthropic(
            model=model_name,
            api_key=api_key or os.getenv("ANTHROPIC_API_KEY"),
            temperature=0.1,
            thinking_dict={"type": "enabled", "budget_tokens": 1024} if "3-7" in model_name else None,
            max_tokens=1024
        )
    else:
        raise ValueError(f"Unsupported model provider for writer agent: {model_provider}")
    
    # Create and return the writer agent with updated system prompt
    return ReActAgent(
        name="writer_agent",
        description="Formats the final answer based on research notes for GAIA benchmark questions",
        system_prompt="""
        You are a specialized formatting agent for the GAIA benchmark. Your job is to take the research from the main agent and format the answer according to the benchmark requirements.

        ## CRITICAL WORKFLOW INSTRUCTIONS
        YOU MUST FOLLOW THIS EXACT WORKFLOW:
        1. Review the main agent's research and format requirements
        2. Format the answer according to specifications
        3. ALWAYS use the handoff tool to delegate to review_agent
        4. NEVER provide a direct answer with "Answer:" - this breaks the workflow

        ## YOUR ROLE
        You will receive:
        - query: The original question
        - research_notes: The main agent's complete analysis and reasoning
        - format_requirements: How the answer should be formatted

        ## FORMATTING RULES
        1. Format the answer according to the instructions in the `query` and `format_requirements`
        2. Your answers will be always as minimal as necessary to answer the question
        3. Remove unnecessary characters, spaces, or wording
        4. If asked for a name, provide **ONLY** the name
        5. If asked for a number, provide the **ONLY** number
        6. If asked for a list, format it exactly as specified

        ## REQUIRED HANDOFF TO REVIEW AGENT
        After formatting your answer, you MUST ALWAYS use the handoff tool with EXACTLY this format:

        Thought: I have formatted the answer and need to delegate to the review agent for final verification.
        Action: handoff
        Action Input: {"to_agent": "review_agent", "reason": "Need final verification and cleanup", "query": "<original question>", "formatted_answer": "<my formatted answer>", "format_requirements": "<explicit formatting requirements>"}

        CRITICAL: Even if the answer seems perfectly formatted already, you MUST use the handoff tool with the EXACT format above. NEVER respond with "Answer: ..." as this breaks the workflow.
        Action: handoff
        Action Input: {"to_agent": "review_agent", "reason": "Final answer to be sent", "query": "<original question>", "formatted_answer": "<my formatted answer>", "format_requirements": "<explicit formatting requirements>"}
        """,
        llm=llm,
        can_handoff_to=["review_agent"],
        verbose=True  # Enable verbose mode for better debugging
    )


def create_review_agent(model_config: Dict[str, Any]) -> ReActAgent:
    """
    Create a review agent that ensures the final answer follows exact formatting requirements.
    
    Args:
        model_config: Dictionary containing model_provider, model_name, and api_key
    
    Returns:
        A configured ReActAgent for final answer review and formatting
    """
    # Initialize LLM based on the provided configuration
    model_provider = model_config.get("model_provider", "openai")
    model_name = model_config.get("model_name", "gpt-4o-mini")
    api_key = model_config.get("api_key")
    
    if model_provider.lower() == "openai":
        llm = OpenAI(
            model=model_name,
            api_key=api_key or os.getenv("OPENAI_API_KEY"),
            max_tokens=128,
            temperature=0.0,  # Use 0 temperature for deterministic output
            additional_kwargs={
                "max_tokens": 128,
                "temperature": 0.0
            }
        )
    elif model_provider.lower() == "anthropic":
        llm = Anthropic(
            model=model_name,
            api_key=api_key or os.getenv("ANTHROPIC_API_KEY"),
            temperature=0.0,  # Use 0 temperature for deterministic output
            thinking_dict={"type": "enabled", "budget_tokens": 1024} if "3-7" in model_name else None,
            max_tokens=128  # Keep token limit low for final answers
        )
    else:
        raise ValueError(f"Unsupported model provider for review agent: {model_provider}")
    
    # Create and return the review agent
    return ReActAgent(
        name="review_agent",
        description="Ensures the final answer is formatted exactly as required, removing any unnecessary information",
        system_prompt="""
        You are the final review agent for the GAIA benchmark. Your ONLY job is to ensure the answer is in the EXACT format required. This is EXTREMELY important for benchmark scoring.

        ## YOUR ROLE
        You will receive:
        - query: The original question
        - formatted_answer: The answer formatted by the writer agent
        - format_requirements: How the answer should be formatted

        ## CRITICAL RULES
        1. Your ENTIRE response must be ONLY the final answer - NOTHING ELSE
        2. Remove ALL of the following:
           - Explanations like "The answer is..." or "I found that..."
           - Quotation marks (unless explicitly required)
           - Punctuation at the end (unless explicitly required)
           - Unnecessary whitespace
        3. If no specific format is mentioned, make the answer as minimal as possible:
           - For names/words: just the name/word (e.g., "Paris")
           - For numbers: just the number (e.g., "42")
           - For lists: comma-separated values (e.g., "apple, banana, cherry")
        4. NEVER add ANY commentary, explanation, or additional information
        5. Double-check for exact formatting requirements like:
           - Numerical format (e.g., "42" vs "forty-two")
           - Case sensitivity (e.g., "PARIS" vs "Paris")
           - List formatting (e.g., comma-separated vs numbered)

        ## OUTPUT FORMAT
        Your entire response should be ONLY the final, formatted answer. Do not include any additional text, explanation, or reasoning.

        ## OUTPUT EXAMPLES
        - Input: "The answer is Alexander."
          Output: Alexander
        - Input: "The result is 42 because..."
          Output: 42
        - Input: "The capital of France is Paris."
          Output: Paris
        - Input: "I found that it's eleven."
          Output: eleven
        - Input: "These actors starred in the film: Tom Hanks, Meg Ryan, and Bill Pullman."
          Output: Tom Hanks, Meg Ryan, Bill Pullman
        - Input: "She published studio albums "Album 1", "Album 2", "Album 3", so in total 3."
          Output: 3
        - Input: Misa Criolla (2000), Acústico (2002), Corazón Libre (2005), Cantora 1 (2009), and Cantora 2 (2009)
          Output: 5

        REMEMBER: Your ENTIRE response should be just the bare answer with NOTHING else.

        Hand the answer off to `main_agent` generate the final answer
        Action: handoff
        Action Input: {"to_agent": "jefe", "reason": "Final answer to be sent", "query": "<original question>", "formatted_answer": "<my formatted answer>", "format_requirements": "<explicit formatting requirements>"}
        """,
        llm=llm,
        can_handoff_to=["jefe"],
    )