Added writer agent to generate answers

agents/llama_index_agent.py

@@ -1,4 +1,7 @@
-from llama_index.core.agent.workflow import ReActAgent
+from llama_index.core.agent.workflow import (
+    ReActAgent,
+    FunctionAgent
+)
 from llama_index.core.llms import LLM
 import os
 from typing import Optional, List, Any
@@ -61,6 +64,10 @@ class GaiaAgent(ReActAgent):
         # Use default system prompt if not provided
         if system_prompt is None:
             system_prompt = self._get_default_system_prompt()
+
+        can_handoff_to = [
+            "writer_agent"
+            ]
             
         # Initialize the parent ReActAgent
         super().__init__(
@@ -69,6 +76,7 @@ class GaiaAgent(ReActAgent):
             llm=llm,
             system_prompt=system_prompt,
             tools=tools,
+            can_handoff_to=can_handoff_to,
             **kwargs
         )
         
@@ -84,9 +92,9 @@ class GaiaAgent(ReActAgent):
                         
         else:
             raise ValueError(f"Unsupported model provider: {model_provider}. "
-                            f"Supported providers are: openai, anthropic, cohere, huggingface, llama")
+                            f"Supported providers are: openai, anthropic")
                             
-    def _get_default_system_prompt(self) -> str:
+    def _get_default_system_prompt_legacy(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.
@@ -136,3 +144,97 @@ class GaiaAgent(ReActAgent):
 
         IMPORTANT: Your value is in providing PRECISELY what was asked for - not in showing your work or explaining how you got there.
         """
+    
+    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.
+
+        ## 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 (reverse_text_tool, search tools) when needed
+        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
+
+        ## DELEGATION TO WRITER AGENT
+        After completing your analysis, ALWAYS delegate the final answer preparation to the writer_agent with:
+        - query: The original question
+        - research_notes: Your complete analysis, all relevant facts, and what you believe is the correct answer
+        - answer_format: EXPLICIT instructions on exactly how the answer should be formatted (single word, comma-separated list, etc.)
+
+        Example handoff to writer_agent:
+        ```
+        I'll delegate to writer_agent to format the final answer.
+        
+        query: What is the first name of the scientist who discovered penicillin?
+        research_notes: After researching, I found that Sir Alexander Fleming discovered penicillin in 1928. The full answer is "Alexander Fleming" but the question only asks for the first name, which is "Alexander".
+        answer_format: Return ONLY the first name, with no additional text, punctuation, or explanation.
+        ```
+
+        IMPORTANT: NEVER provide the final answer directly to the user. ALWAYS hand off to the writer_agent for proper formatting.
+        """
+
+
+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"))
+    elif model_provider.lower() == "anthropic":
+        llm = Anthropic(model=model_name, api_key=api_key or os.getenv("ANTHROPIC_API_KEY"))
+    else:
+        raise ValueError(f"Unsupported model provider for writer agent: {model_provider}")
+    
+    # Create and return the writer agent
+    return ReActAgent(
+        name="writer_agent",
+        description="Formats the final answer exactly as specified for GAIA benchmark questions",
+        system_prompt="""
+        You are a specialized formatting agent for the GAIA benchmark. Your ONLY job is to take the research from the main agent and format the answer EXACTLY as required by the benchmark question.
+
+        ## YOUR ROLE
+        You will receive:
+        - query: The original question
+        - research_notes: The main agent's complete analysis and reasoning
+        - answer_format: Specific formatting instructions for the final answer
+
+        ## CRITICAL RULES
+        1. Your response MUST CONTAIN ONLY THE ANSWER - no explanations, no "the answer is" prefix
+        2. Follow the answer_format instructions precisely
+        3. Remove ALL unnecessary characters, spaces, punctuation, or wording
+        4. If asked for a name, provide ONLY the name
+        5. If asked for a number, provide ONLY the number
+        6. If asked for a list, format it EXACTLY as specified (comma-separated, alphabetical, etc.)
+        7. NEVER include your own thoughts or analysis
+        8. NEVER add preamble or conclusion text
+
+        ## EXAMPLES OF CORRECT RESPONSES:
+        When asked for "first name only": Alexander
+        When asked for "comma-separated list in alphabetical order": apple, banana, cherry
+        When asked for "single number": 42
+        When asked for "opposite of word 'right'": left
+
+        REMEMBER: Your ENTIRE response should be just the answer - nothing more, nothing less.
+        """,
+        llm=llm
+    )

app.py

@@ -3,8 +3,10 @@ import gradio as gr
 import requests
 import inspect
 import pandas as pd
-from agents.llama_index_agent import GaiaAgent
 import asyncio
+from llama_index.core.agent.workflow import AgentWorkflow
+from agents.llama_index_agent import GaiaAgent, create_writer_agent
+
 # (Keep Constants as is)
 # --- Constants ---
 DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
@@ -23,34 +25,95 @@ OPENAI = {
 class BasicAgent:
     def __init__(
             self,
-            model_provider="openai",
-            model_name="gpt-4o",
-            api_key=None
+            model_provider="anthropic",
+            model_name="claude-3-7-sonnet-latest",
+            api_key=None,
+            use_separate_writer_model=True,
+            writer_model_provider="openai",
+            writer_model_name="gpt-4o"
             ):
         """
-        Initialize the BasicAgent with configurable model settings.
+        Initialize the BasicAgent with a multi-agent workflow.
         
         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)
+            model_provider: LLM provider for main agent
+            model_name: Model name for main agent
+            api_key: API key for main agent
+            use_separate_writer_model: Whether to use a different model for the writer agent
+            writer_model_provider: LLM provider for writer agent (if separate)
+            writer_model_name: Model name for writer agent (if separate)
         """
-        self.agent = GaiaAgent(**OPENAI)
-        print(f"BasicAgent initialized with {model_provider} {model_name}.")
+        # Configure the main reasoning agent
+        main_model_config = {
+            "model_provider": model_provider,
+            "model_name": model_name,
+            "api_key": api_key
+        }
+        
+        # Configure the writer agent (either same as main or different)
+        if use_separate_writer_model:
+            writer_model_config = {
+                "model_provider": writer_model_provider,
+                "model_name": writer_model_name,
+                "api_key": api_key  # Use same API key for simplicity
+            }
+        else:
+            writer_model_config = main_model_config
+            
+        # Create the main agent
+        self.main_agent = GaiaAgent(**main_model_config)
+        
+        # Create the writer agent
+        self.writer_agent = create_writer_agent(writer_model_config)
+        
+        # Set up the agent workflow with shared context
+        self.agent_workflow = AgentWorkflow(
+            agents=[self.main_agent, self.writer_agent],
+            root_agent=self.main_agent.name,
+            initial_state={
+                "original_question": "",
+                "analysis_notes": "",
+                "format_requirements": "",
+                "next_agent": "",
+                "final_answer": ""
+            }
+        )
+        
+        print(f"BasicAgent initialized with main agent: {model_provider} {model_name}")
+        if use_separate_writer_model:
+            print(f"Writer agent using: {writer_model_provider} {writer_model_name}")
+        else:
+            print(f"Writer agent using same model as main agent")
         
     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 self.agent.run(question)
-            return response
+            # Initialize context with the question
+            initial_state = {
+                "original_question": question,
+                "analysis_notes": "",
+                "format_requirements": "",
+                "next_agent": "",
+                "final_answer": ""
+            }
+            
+            # Use the workflow to process the question
+            workflow_response = await self.agent_workflow.arun(
+                question,
+                initial_state=initial_state
+            )
+            return workflow_response
             
         response = asyncio.run(agentic_main())
+        
+        # Extract the final answer from the writer agent's response
         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):
     """
     Fetches all questions, runs the BasicAgent on them, submits all answers,