added model selection to main agent class

agents/llama_index_agent.py

@@ -1,65 +1,141 @@
-from llama_index.core.agent.workflow import (
-    AgentWorkflow,
-    ReActAgent,
-    FunctionAgent
-)
-from tools.text_tools import reverse_text_tool
+from llama_index.core.agent.workflow import ReActAgent
+from llama_index.core.llms import LLM
+import os
+from typing import Optional, List, Any
 from llama_index.llms.openai import OpenAI
-import os 
+from llama_index.llms.anthropic import Anthropic
 
-openai = OpenAI(model="gpt-4o", api_key=os.getenv("OPENAI_API_KEY"))
-
-main_agent = ReActAgent(
-    name="jefe",
-    description="Agent that will receive the queries, understand them, and send them to the correct agents to do the job",
-    llm=openai,
-    system_prompt="""
-        You are a ReActAgent that has a team of AI agents available to solve
-        questions and challenges from the GAIA Benchmark. 
+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]
+            
+        # Use default system prompt if not provided
+        if system_prompt is None:
+            system_prompt = self._get_default_system_prompt()
+            
+        # Initialize the parent ReActAgent
+        super().__init__(
+            name=name,
+            description=description,
+            llm=llm,
+            system_prompt=system_prompt,
+            tools=tools,
+            **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"))
+            
+        elif model_provider == "anthropic":
+            return Anthropic(model=model_name, api_key=api_key or os.getenv("ANTHROPIC_API_KEY"))
+            
+        elif model_provider == "cohere":
+            from llama_index.llms.cohere import Cohere
+            return Cohere(model=model_name, api_key=api_key or os.getenv("COHERE_API_KEY"))
+            
+        elif model_provider == "huggingface":
+            from llama_index.llms.huggingface import HuggingFaceLLM
+            return HuggingFaceLLM(model_name=model_name, tokenizer_name=model_name)
+            
+        elif model_provider == "llama":
+            from llama_index.llms.llama_cpp import LlamaCPP
+            return LlamaCPP(model_path=model_name)
+            
+        else:
+            raise ValueError(f"Unsupported model provider: {model_provider}. "
+                            f"Supported providers are: openai, anthropic, cohere, huggingface, llama")
+                            
+    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 each question with extreme precision, determine the exact format required for the answer, break the task into logical steps, and either solve it yourself or delegate to the appropriate specialized agents.
 
-        You must very carefully read the questions, understand them, and divide them into steps.
-        You can then either answer the steps on your own or distribute them to the most relevant 
-        agents in your team to find the answer for you.
+        ## 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
 
-        At the end, once you gather 
-        
-        The questions will be given to you following the format:
-        ```
-            {
-                'task_id': '5a0c1adf-205e-4841-a666-7c3ef95def9d',
-                'question': 'What is the first name of the only Malko Competition recipient from the 20th Century (after 1977) whose nationality on record is a country that no longer exists?',
-                'Level': '1',
-                'file_name': ''
-            }
-        ```
+        ## DELEGATION GUIDELINES
+        - video_analyst: Use for all YouTube video analysis, visual content identification, or scene description
+        - audio_analyst: Use for transcribing audio files, identifying speakers, or extracting information from recordings
+        - researcher: Use for factual queries, literature searches, finding specific information in papers or websites
+        - code_analyst: Use for executing, debugging or analyzing code snippets
+        - excel_analyst: Use for analyzing spreadsheets, calculating values, or extracting data from Excel files
 
-        If the question has a file attached, the other agents in your team will have the tools to open and 
-        analyze them.
+        ## CRITICAL RESPONSE RULES
+        - NEVER include explanations in your final answer
+        - NEVER include phrases like "the answer is" or "the result is"
+        - Return EXACTLY what was asked for - no more, no less
+        - If asked for a name, return ONLY the name
+        - If asked for a number, return ONLY the number
+        - If asked for a list, format it EXACTLY as specified (comma-separated, alphabetical, etc.)
+        - Double-check your answer against the exact output requirements before submitting
 
-        Once you have all the intermediate steps and you can provide the final answer, make sure that
-        you are doing so EXACTLY as the answer format is defined in the query.
+        ## EXAMPLES OF PROPER RESPONSES:
+        Question: "What is the first name of the scientist who discovered penicillin?"
+        Correct answer: Alexander
 
-        You also have access to your own tools: 
-        * `reverse_text_tool` --> Reverses the input text
+        Question: "List the prime numbers between 10 and 20 in ascending order."
+        Correct answer: 11, 13, 17, 19
 
-        Send as final answer your last answer formated as expected in the instructions of the question.
+        Question: "If you understand this sentence, write the opposite of the word 'right' as the answer."
+        Correct answer: left
 
-        IMPORTANT:
-        Answer ONLY the minimal text that is asked in the query. Not a single character more.
+        Question: "How many at bats did the Yankee with the most walks in the 1977 regular season have that same season?"
+        Correct answer: 572
 
-        For example:
+        For questions with reverse text:
+        1. Use your reverse_text_tool to process the text
+        2. Understand the instruction in the reversed text
+        3. Follow the instruction exactly
 
-        If you are asked for the first name of the discoverer of gravity, your answer should be: `Isaac`
+        After you have the final answer, verify one last time that it meets ALL formatting requirements from the question before submitting.
 
-    """,
-    can_handoff_to=[
-        "video_analyst",
-        "audio_analyst",
-        "researcher",
-        "code_analyst",
-        "excel_analyst"
-        ],
-    tools=[
-        reverse_text_tool
-    ]
-)
+        IMPORTANT: Your value is in providing PRECISELY what was asked for - not in showing your work or explaining how you got there.
+        """

app.py

@@ -3,7 +3,7 @@ import gradio as gr
 import requests
 import inspect
 import pandas as pd
-from agents.llama_index_agent import main_agent
+from agents.llama_index_agent import GaiaAgent
 import asyncio
 # (Keep Constants as is)
 # --- Constants ---
@@ -11,17 +11,45 @@ DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 
 # --- Basic Agent Definition ---
 # ----- THIS IS WERE YOU CAN BUILD WHAT YOU WANT ------
+
+CLAUDE = {
+    "model_provider": "anthropic",
+    "model_name": "claude-3-7-sonnet-20250219"
+}
+OPENAI = {
+    "model_provider": "openai",
+    "model_name": "gpt-4o"
+}
 class BasicAgent:
-    def __init__(self):
-        print("BasicAgent initialized.")
+    def __init__(
+            self,
+            model_provider="openai",
+            model_name="gpt-4o",
+            api_key=None
+            ):
+        """
+        Initialize the BasicAgent with configurable model settings.
+        
+        Args:
+            model_provider: LLM provider to use (openai, anthropic, etc.)
+            model_name: Specific model to use
+            api_key: Optional API key (defaults to environment variable)
+        """
+        self.agent = GaiaAgent(**CLAUDE)
+        print(f"BasicAgent initialized with {model_provider} {model_name}.")
+        
     def __call__(self, question: str) -> str:
+        """Process a GAIA benchmark question and return the formatted answer."""
         print(f"Agent received question (first 50 chars): {question[:50]}...")
+        
         async def agentic_main():
-            response = await main_agent.run(question)
+            response = await self.agent.run(question)
             return response
+            
         response = asyncio.run(agentic_main())
-        print(f"Agent returning answer: {response.response.blocks[-1].text}")
-        return response.response.blocks[-1].text
+        final_answer = response.response.blocks[-1].text
+        print(f"Agent returning answer: {final_answer}")
+        return final_answer
 
 def run_and_submit_all( profile: gr.OAuthProfile | None):
     """