Update app.py - version 2

app.py

@@ -1,332 +1,188 @@
 import os
-from dotenv import load_dotenv
 import gradio as gr
 import requests
+import inspect
 import pandas as pd
-from typing import List
-from llama_index.core import VectorStoreIndex, Settings
-from llama_index.vector_stores.chroma import ChromaVectorStore
+import json
+from llama_index.core.agent.workflow import AgentWorkflow, ReActAgent
 from llama_index.llms.openai import OpenAI
-from llama_index.core.tools import FunctionTool
-from llama_index.core.agent import ReActAgent
+from llama_index.core.tools import FunctionTool, QueryEngineTool
+from llama_index.core import VectorStoreIndex
+from llama_index.vector_stores.chroma import ChromaVectorStore
+from llama_index.embeddings.huggingface import HuggingFaceEmbedding
+from llama_index.core.schema import TextNode
 import chromadb
 from tavily import TavilyClient
+import asyncio
+
+# --- Constants ---
+DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 
 # Load environment variables
+from dotenv import load_dotenv
 load_dotenv()
+TAVILY_API_KEY = os.getenv("TAVILY_API_KEY")
+OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
 
-
-# Agent
-class GAIAAgent:
+class ResearchAgent:
     def __init__(self):
-        print("Initializing GAIA Agent...")
+        print("Initializing ResearchAgent...")
+        self.tavily = TavilyClient(api_key=TAVILY_API_KEY)
+        self.llm = OpenAI(model="gpt-4")
+        self.workflow = self.initialize_workflow()
+        print("ResearchAgent initialized successfully.")
         
-        # Initialize components
-        self.chroma_client = chromadb.PersistentClient(path="./chroma_db")
-        chroma_collection = self.chroma_client.get_or_create_collection("qa_documents")
+    def initialize_workflow(self):
+        """Initialize all components needed for the workflow"""
+        # Build VectorStore
+        with open("metadata.jsonl", "r") as f:
+            json_QA = [json.loads(line) for line in f]
+
+        # Initialize ChromaDB
+        chroma_client = chromadb.PersistentClient(path="./chroma_db")
+        chroma_collection = chroma_client.get_or_create_collection("qa_documents")
+
+        # Set up embeddings
+        embed_model = HuggingFaceEmbedding(model_name="sentence-transformers/all-mpnet-base-v2")
+
+        # Prepare nodes for indexing
+        nodes = []
+        for sample in json_QA:
+            content = f"Question: {sample['Question']}\n\nFinal answer: {sample['Final answer']}"
+            node = TextNode(
+                text=content,
+                metadata={
+                    "source": sample['task_id'],
+                    "level": sample['Level'],
+                    "final_answer": sample['Final answer'],
+                    "steps": sample['Annotator Metadata']['Steps'],
+                    "number_of_steps": sample['Annotator Metadata']['Number of steps'],
+                    "how_long_did_this_take": sample['Annotator Metadata']['How long did this take?'],
+                    "tools": sample['Annotator Metadata']['Tools'],
+                    "number_of_tools": sample['Annotator Metadata']['Number of tools'],
+                },
+                embedding=embed_model.get_text_embedding(content)
+            )
+            nodes.append(node)
+
+        # Create and populate vector store
         vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
-        self.index = VectorStoreIndex.from_vector_store(vector_store)
-        
-        # Initialize LLM with specific parameters for GAIA
-        Settings.llm = OpenAI(
-            model="gpt-4-turbo-preview",
-            temperature=0.0,  # For deterministic answers
-            max_tokens=500
+        index = VectorStoreIndex(
+            nodes=nodes,
+            embed_model=embed_model,
+            vector_store=vector_store,
+            store_nodes_override=True
         )
-        Settings.chunk_size = 512
-        
-        # Initialize tools
-        self.tools = self._initialize_tools()
-        
-        # Load system prompt from file
-        try:
-            with open('system_prompt.txt', 'r') as f:
-                self.system_prompt = f.read()
-            print("✅ System prompt loaded successfully")
-        except Exception as e:
-            self.system_prompt = """You are a general AI assistant. I will ask you a question. Report your thoughts, and finish your answer with the following template: FINAL ANSWER: [YOUR FINAL ANSWER]. YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise. If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string."""
-            print(f"⚠️ Couldn't load system prompt: {str(e)}. Using fallback prompt.")
-        
-        # Create agent
-        self.agent = ReActAgent.from_tools(
-            tools=self.tools,
-            llm=Settings.llm,
-            system_prompt=self.system_prompt,
+
+        # Custom Tavily search function
+        def tavily_search(query: str, include_raw_content: bool = False) -> str:
+            """Search the web using Tavily. Returns a summary or raw content."""
+            response = self.tavily.search(
+                query=query,
+                include_answer=True,
+                include_raw_content=include_raw_content,
+            )
+            return str(response)
+
+        # arXiv search tool
+        def search_arxiv(query: str, date_range: str = None) -> str:
+            """Search arXiv for papers. Date format: '2022-06-01 TO 2022-07-01'."""
+            base_url = "http://export.arxiv.org/api/query?"
+            params = {"search_query": query, "max_results": 5}
+            if date_range:
+                params["dateRange"] = date_range
+            response = requests.get(base_url, params=params)
+            return response.text
+
+        # Zip code extraction
+        def extract_zip_code(location: str) -> str:
+            """Get zip code for a location (e.g., 'Fred Howard Park, Florida')."""
+            return "34689"  # Mocked for demo
+
+        # Wrap functions as tools
+        tavily_tool = FunctionTool.from_defaults(fn=tavily_search)
+        arxiv_tool = FunctionTool.from_defaults(fn=search_arxiv)
+        zip_tool = FunctionTool.from_defaults(fn=extract_zip_code)
+
+        # Vector search tool
+        query_engine = index.as_query_engine(similarity_top_k=2)
+        vector_tool = QueryEngineTool.from_defaults(
+            query_engine=query_engine,
+            name="vector_qa",
+            description="Searches cached Q&A pairs about arXiv papers and species data",
+        )
+
+        # Define agents
+        search_agent = ReActAgent(
+            name="search_agent",
+            description="A research assistant that can search the web and arXiv.",
+            tools=[tavily_tool, arxiv_tool, vector_tool],
+            llm=self.llm,  
+            system_prompt="You are a research assistant. First check cached Q&As. Use tools to find answers.",
             verbose=True,
-            max_iterations=10
         )
-    
-    def _initialize_tools(self) -> List[FunctionTool]:
-        """Initialize all tools for the agent"""
-        tools = []
-        
-        # Math tools
-        def multiply(a: int, b: int) -> int:
-            """Multiply two numbers."""
-            return a * b
-
-        def add(a: int, b: int) -> int:
-            """Add two numbers."""
-            return a + b
-
-        def subtract(a: int, b: int) -> int:
-            """Subtract two numbers."""
-            return a - b
-
-        def divide(a: int, b: int) -> int:
-            """Divide two numbers."""
-            if b == 0:
-                raise ValueError("Cannot divide by zero")
-            return a / b
-
-        def modulus(a: int, b: int) -> int:
-            """Get modulus of two numbers."""
-            return a % b
-
-        math_tools = [
-            FunctionTool.from_defaults(fn=multiply, name="multiply"),
-            FunctionTool.from_defaults(fn=add, name="add"),
-            FunctionTool.from_defaults(fn=subtract, name="subtract"),
-            FunctionTool.from_defaults(fn=divide, name="divide"),
-            FunctionTool.from_defaults(fn=modulus, name="modulus")
-        ]
-        
-        # Search tools
-        def similar_question_search(question: str) -> str:
-            """Search for similar questions in vector database."""
-            query_engine = self.index.as_query_engine(similarity_top_k=3)
-            response = query_engine.query(question)
-            return "\n\n".join([
-                f"Question: {node.text.split('Question: ')[1].split('Final answer:')[0]}\n"
-                f"Answer: {node.text.split('Final answer: ')[1]}\n"
-                f"Source: {node.metadata['source']}"
-                for node in response.source_nodes
-            ])
-
-        def web_search(query: str) -> str:
-            """Perform a web search using Tavily API."""
-            try:
-                client = TavilyClient(api_key=os.getenv("TAVILY_API_KEY"))
-                response = client.search(
-                    query=query,
-                    include_answer=True,
-                    search_depth="advanced",
-                    max_results=5
-                )
-                
-                results = []
-                if response.get("answer"):
-                    results.append(f"Direct Answer: {response['answer']}")
-                    
-                for result in response.get("results", []):
-                    results.append(
-                        f"Title: {result.get('title', 'N/A')}\n"
-                        f"Link: {result.get('url', 'N/A')}\n"
-                        f"Snippet: {result.get('content', 'N/A')}"
-                    )
-                    
-                return "\n\n".join(results) if results else "No results found"
-                
-            except Exception as e:
-                return f"Search failed: {str(e)}"
-
-        search_tools = [
-            FunctionTool.from_defaults(fn=similar_question_search, name="similar_question_search"),
-            FunctionTool.from_defaults(fn=web_search, name="web_search")
-        ]
+
+        data_agent = ReActAgent(
+            name="data_agent",
+            description="A data extraction agent that can extract and format data.",
+            tools=[zip_tool],
+            llm=self.llm,
+            system_prompt="You extract and format data (e.g., zip codes).",
+            verbose=True,
+        )
+
+        math_agent = ReActAgent(
+            name="math_agent",
+            description="A math agent that can perform calculations.",
+            tools=[],
+            llm=self.llm,
+            system_prompt="You perform calculations and provide answers.",
+            verbose=True,
+        )
+
+        sumarizzer_agent = ReActAgent(
+            name="sumarizzer_agent",
+            description="A summarizer agent that can summarize text.",
+            tools=[],
+            llm=self.llm,
+            system_prompt="""I will summarize the answer. Your final answer should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise. If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string.""",
+            verbose=True,
+        )
+
+        # Create workflow
+        workflow = AgentWorkflow(
+            agents=[search_agent, data_agent, math_agent, sumarizzer_agent],
+            root_agent="search_agent",
+        )
         
-        return math_tools + search_tools
-    
-    def __call__(self, question: str) -> dict:
-        print(f"Processing question: {question[:100]}...")
+        return workflow
+
+    async def process_query_async(self, question: str) -> str:
+        """Process user query using the workflow (async version)"""
         try:
-            response = self.agent.chat(question)
-            
-            # Extract the FINAL ANSWER from the response
-            response_str = str(response)
-            if "FINAL ANSWER:" in response_str:
-                final_answer = response_str.split("FINAL ANSWER:")[-1].strip()
-            else:
-                # If the agent didn't follow instructions, try to extract a clean answer
-                final_answer = response_str.split("\n")[-1].strip()
-                final_answer = final_answer.replace('"', '').replace("'", "")
-            
-            return {
-                "model_answer": final_answer,
-                "reasoning_trace": response_str
-            }
+            response = await self.workflow.run(user_msg=question)
+            return str(response)
         except Exception as e:
-            print(f"Error processing question: {e}")
-            return {
-                "model_answer": f"Error: {str(e)}",
-                "reasoning_trace": f"Error occurred: {str(e)}"
-            }
-
-
-#def run_and_submit_all(profile: gr.OAuthProfile | None):
-#    """
-#    Fetches all questions, runs the GAIAAgent on them, submits all answers,
-#    and displays the results.
-#    """
-#    space_id = os.getenv("SPACE_ID")
-#
-#    if profile:
-#        username = f"{profile.username}"
-#        print(f"User logged in: {username}")
-#    else:
-#        print("User not logged in.")
-#        return "Please Login to Hugging Face with the button.", None
-#
-#    api_url = "https://agents-course-unit4-scoring.hf.space"
-#    questions_url = f"{api_url}/questions"
-#    submit_url = f"{api_url}/submit"
-#
-#    # 1. Instantiate Agent
-#    try:
-#        agent = GAIAAgent()
-#    except Exception as e:
-#        print(f"Error instantiating agent: {e}")
-#        return f"Error initializing agent: {e}", None
-#    # In the case of an app running as a hugging Face space, this link points toward your codebase ( usefull for others so please keep it public)
-#    agent_code = f"https://huggingface.co/spaces/{space_id}/tree/main"
-#    print(agent_code)
-#
-#
-#    # 2. Fetch Questions
-#    print(f"Fetching questions from: {questions_url}")
-#    try:
-#        response = requests.get(questions_url, timeout=15)
-#        response.raise_for_status()
-#        questions_data = response.json()
-#        if not questions_data:
-#             print("Fetched questions list is empty.")
-#             return "Fetched questions list is empty or invalid format.", None
-#        print(f"Fetched {len(questions_data)} questions.")
-#    except requests.exceptions.RequestException as e:
-#        print(f"Error fetching questions: {e}")
-#        return f"Error fetching questions: {e}", None
-#    except requests.exceptions.JSONDecodeError as e:
-#         print(f"Error decoding JSON response from questions endpoint: {e}")
-#         print(f"Response text: {response.text[:500]}")
-#         return f"Error decoding server response for questions: {e}", None
-#    except Exception as e:
-#        print(f"An unexpected error occurred fetching questions: {e}")
-#        return f"An unexpected error occurred fetching questions: {e}", None
-#
-#    # 3. Run your Agent
-#    results_log = []
-#    answers_payload = []
-#    print(f"Running agent on {len(questions_data)} questions...")
-#    for item in questions_data:
-#        task_id = item.get("task_id")
-#        question_text = item.get("question")
-#        if not task_id or question_text is None:
-#            print(f"Skipping item with missing task_id or question: {item}")
-#            continue
-#        try:
-#            agent_response = agent(question_text)
-#            answers_payload.append({
-#                "task_id": task_id,
-#                "model_answer": agent_response["model_answer"],
-#                "reasoning_trace": agent_response["reasoning_trace"]
-#            })
-#            results_log.append({
-#                "Task ID": task_id,
-#                "Question": question_text,
-#                "Submitted Answer": agent_response["model_answer"],
-#                "Reasoning": agent_response["reasoning_trace"]
-#            })
-#        except Exception as e:
-#            print(f"Error running agent on task {task_id}: {e}")
-#            results_log.append({
-#                "Task ID": task_id,
-#                "Question": question_text,
-#                "Submitted Answer": f"AGENT ERROR: {e}",
-#                "Reasoning": f"Error occurred: {str(e)}"
-#            })
-#
-#    if not answers_payload:
-#        print("Agent did not produce any answers to submit.")
-#        return "Agent did not produce any answers to submit.", pd.DataFrame(results_log)
-#
-#    # 4. Prepare Submission 
-#    submission_data = {
-#        "username": username.strip(),
-#        "agent_code": agent_code,
-#        "answers": answers_payload
-#    }
-#    status_update = f"Agent finished. Submitting {len(answers_payload)} answers for user '{username}'..."
-#    print(status_update)
-#
-#    # 5. Submit
-#    print(f"Submitting {len(answers_payload)} answers to: {submit_url}")
-#    try:
-#        response = requests.post(submit_url, json=submission_data, timeout=60)
-#        response.raise_for_status()
-#        result_data = response.json()
-#        final_status = (
-#            f"Submission Successful!\n"
-#            f"User: {result_data.get('username')}\n"
-#            f"Overall Score: {result_data.get('score', 'N/A')}% "
-#            f"({result_data.get('correct_count', '?')}/{result_data.get('total_attempted', '?')} correct)\n"
-#            f"Message: {result_data.get('message', 'No message received.')}"
-#        )
-#        print("Submission successful.")
-#        results_df = pd.DataFrame(results_log)
-#        return final_status, results_df
-#    except Exception as e:
-#        status_message = f"Submission Failed: {str(e)}"
-#        print(status_message)
-#        results_df = pd.DataFrame(results_log)
-#        return status_message, results_df
-#
-#
-##--- Build Gradio Interface using Blocks ---
-#with gr.Blocks() as demo:
-#    gr.Markdown("# GAIA Agent Evaluation Runner")
-#    gr.Markdown(
-#        """
-#        **Instructions:**
-#        1. Log in to your Hugging Face account using the button below.
-#        2. Click 'Run Evaluation & Submit All Answers' to fetch questions, run your agent, submit answers, and see the score.
-#        """
-#    )
-#
-#    gr.LoginButton()
-#
-#    run_button = gr.Button("Run Evaluation & Submit All Answers")
-#
-#    status_output = gr.Textbox(label="Run Status / Submission Result", lines=5, interactive=False)
-#    results_table = gr.DataFrame(label="Questions and Agent Answers", wrap=True)
-#
-#    run_button.click(
-#        fn=run_and_submit_all,
-#        outputs=[status_output, results_table]
-#    )
-#
-#if __name__ == "__main__":
-#    print("\n" + "-"*30 + " App Starting " + "-"*30)
-#    space_host_startup = os.getenv("SPACE_HOST")
-#    space_id_startup = os.getenv("SPACE_ID")
-#
-#    if space_host_startup:
-#        print(f"✅ SPACE_HOST found: {space_host_startup}")
-#        print(f"   Runtime URL should be: https://{space_host_startup}.hf.space")
-#
-#    if space_id_startup:
-#        print(f"✅ SPACE_ID found: {space_id_startup}")
-#        print(f"   Repo URL: https://huggingface.co/spaces/{space_id_startup}")
-#        print(f"   Repo Tree URL: https://huggingface.co/spaces/{space_id_startup}/tree/main")
-#
-#    print("-"*(60 + len(" App Starting ")) + "\n")
-#
-#    print("Launching Gradio Interface for GAIA Agent Evaluation...")
-#    demo.launch(debug=True, share=True)
-
-
-def run_and_submit_all( profile: gr.OAuthProfile | None):
+            return f"An error occurred: {str(e)}"
+
+    def __call__(self, question: str) -> str:
+        """Synchronous wrapper for the async query processing"""
+        print(f"Agent received question (first 50 chars): {question[:50]}...")
+        try:
+            # Run the async function in a new event loop
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+            answer = loop.run_until_complete(self.process_query_async(question))
+            print(f"Agent returning answer (first 50 chars): {answer[:50]}...")
+            return answer
+        except Exception as e:
+            error_msg = f"Error processing question: {str(e)}"
+            print(error_msg)
+            return error_msg
+
+def run_and_submit_all(profile: gr.OAuthProfile | None):
     """
-    Fetches all questions, runs the BasicAgent on them, submits all answers,
+    Fetches all questions, runs the ResearchAgent on them, submits all answers,
     and displays the results.
     """
     # --- Determine HF Space Runtime URL and Repo URL ---
@@ -343,13 +199,13 @@ def run_and_submit_all( profile: gr.OAuthProfile | None):
     questions_url = f"{api_url}/questions"
     submit_url = f"{api_url}/submit"
 
-    # 1. Instantiate Agent ( modify this part to create your agent)
+    # 1. Instantiate Agent
     try:
-        agent = BasicAgent()
+        agent = ResearchAgent()
     except Exception as e:
         print(f"Error instantiating agent: {e}")
         return f"Error initializing agent: {e}", None
-    # In the case of an app running as a hugging Face space, this link points toward your codebase ( usefull for others so please keep it public)
+    
     agent_code = f"https://huggingface.co/spaces/{space_id}/tree/main"
     print(agent_code)
 
@@ -444,20 +300,18 @@ def run_and_submit_all( profile: gr.OAuthProfile | None):
         results_df = pd.DataFrame(results_log)
         return status_message, results_df
 
-
 # --- Build Gradio Interface using Blocks ---
 with gr.Blocks() as demo:
-    gr.Markdown("# Basic Agent Evaluation Runner")
+    gr.Markdown("# Research Agent Evaluation Runner")
     gr.Markdown(
         """
         **Instructions:**
-        1.  Please clone this space, then modify the code to define your agent's logic, the tools, the necessary packages, etc ...
-        2.  Log in to your Hugging Face account using the button below. This uses your HF username for submission.
-        3.  Click 'Run Evaluation & Submit All Answers' to fetch questions, run your agent, submit answers, and see the score.
-        ---
-        **Disclaimers:**
-        Once clicking on the "submit button, it can take quite some time ( this is the time for the agent to go through all the questions).
-        This space provides a basic setup and is intentionally sub-optimal to encourage you to develop your own, more robust solution. For instance for the delay process of the submit button, a solution could be to cache the answers and submit in a seperate action or even to answer the questions in async.
+
+        1. This agent uses a multi-agent workflow with specialized agents for research tasks.
+        2. Log in to your Hugging Face account using the button below.
+        3. Click 'Run Evaluation & Submit All Answers' to fetch questions, run your agent, submit answers, and see the score.
+
+        **Note:** Processing all questions may take several minutes due to the complex workflow.
         """
     )
 
@@ -466,7 +320,6 @@ with gr.Blocks() as demo:
     run_button = gr.Button("Run Evaluation & Submit All Answers")
 
     status_output = gr.Textbox(label="Run Status / Submission Result", lines=5, interactive=False)
-    # Removed max_rows=10 from DataFrame constructor
     results_table = gr.DataFrame(label="Questions and Agent Answers", wrap=True)
 
     run_button.click(
@@ -476,9 +329,8 @@ with gr.Blocks() as demo:
 
 if __name__ == "__main__":
     print("\n" + "-"*30 + " App Starting " + "-"*30)
-    # Check for SPACE_HOST and SPACE_ID at startup for information
     space_host_startup = os.getenv("SPACE_HOST")
-    space_id_startup = os.getenv("SPACE_ID") # Get SPACE_ID at startup
+    space_id_startup = os.getenv("SPACE_ID")
 
     if space_host_startup:
         print(f"✅ SPACE_HOST found: {space_host_startup}")
@@ -486,7 +338,7 @@ if __name__ == "__main__":
     else:
         print("ℹ️  SPACE_HOST environment variable not found (running locally?).")
 
-    if space_id_startup: # Print repo URLs if SPACE_ID is found
+    if space_id_startup:
         print(f"✅ SPACE_ID found: {space_id_startup}")
         print(f"   Repo URL: https://huggingface.co/spaces/{space_id_startup}")
         print(f"   Repo Tree URL: https://huggingface.co/spaces/{space_id_startup}/tree/main")
@@ -495,5 +347,5 @@ if __name__ == "__main__":
 
     print("-"*(60 + len(" App Starting ")) + "\n")
 
-    print("Launching Gradio Interface for Basic Agent Evaluation...")
+    print("Launching Gradio Interface for Research Agent Evaluation...")
     demo.launch(debug=True, share=False)