Update agent.py

agent.py

@@ -1,15 +1,14 @@
-# agent.py — final version without circular imports
+# agent.py — GAIA-ready with FunctionCallingAgent and improved tools
 
 import os
 import asyncio
 from llama_index.llms.openai import OpenAI
-from llama_index.core.agent.react.base import ReActAgent
+from llama_index.core.agent import FunctionCallingAgent
 from llama_index.core.tools import FunctionTool
 
-from langchain_community.tools.wikipedia.tool import WikipediaQueryRun
 from langchain_community.utilities.wikipedia import WikipediaAPIWrapper
-from langchain_experimental.tools.python.tool import PythonREPLTool
 from langchain_community.document_loaders import YoutubeLoader
+from langchain_experimental.tools.python.tool import PythonREPLTool
 
 import whisper
 import openpyxl
@@ -20,34 +19,51 @@ if os.getenv("OPENAI_API_KEY"):
 else:
     print("⚠️ Missing OPENAI_API_KEY – LLM may fail")
 
-# Tools definitions
-api_wrapper = WikipediaAPIWrapper(top_k_results=1, doc_content_chars_max=1000)
+# Tool 1 — Wikipedia
+wiki_api = WikipediaAPIWrapper(top_k_results=1, doc_content_chars_max=1000)
+
+def search_wikipedia(query: str) -> str:
+    """Search Wikipedia for a given query and return relevant summary."""
+    return wiki_api.run(query)
 
-def wikipedia_search(query: str) -> str:
-    return WikipediaQueryRun(api_wrapper=api_wrapper).run({"query": query})
+# Tool 2 — Python with output
+python_tool = PythonREPLTool()
+
+def run_python_code(code: str) -> str:
+    """Run Python code and return printed result."""
+    try:
+        if "print(" not in code:
+            code = f"print({code})"
+        return python_tool.run(code)
+    except Exception as e:
+        return f"[PYTHON ERROR] {e}"
 
-def run_python_with_output(code: str) -> str:
-    if "print(" not in code:
-        code = f"print({code})"
-    return PythonREPLTool().run(code)
+# Tool 3 — YouTube transcript
 
 def get_youtube_transcript(url: str) -> str:
+    """Get transcript from YouTube video."""
     try:
         loader = YoutubeLoader.from_youtube_url(url, add_video_info=False)
         docs = loader.load()
         return " ".join(d.page_content for d in docs)
     except Exception as e:
-        return "[YOUTUBE ERROR] " + str(e)
+        return f"[YOUTUBE ERROR] {e}"
+
+# Tool 4 — Whisper transcription
 
 def transcribe_audio(file_path: str) -> str:
+    """Transcribe an MP3 file to text using Whisper."""
     try:
         model = whisper.load_model("base")
         res = model.transcribe(file_path)
         return res["text"]
     except Exception as e:
-        return "[AUDIO ERROR] " + str(e)
+        return f"[AUDIO ERROR] {e}"
+
+# Tool 5 — Excel parser
 
 def extract_excel_total_food_sales(file_path: str) -> str:
+    """Sum sales from Excel where category is 'food'."""
     try:
         wb = openpyxl.load_workbook(file_path)
         sheet = wb.active
@@ -57,31 +73,41 @@ def extract_excel_total_food_sales(file_path: str) -> str:
                 total += float(amount or 0)
         return f"${total:.2f}"
     except Exception as e:
-        return "[EXCEL ERROR] " + str(e)
+        return f"[EXCEL ERROR] {e}"
 
-# Assemble tools
+# Assemble tools with proper descriptions
 TOOLS = [
-    FunctionTool.from_defaults(wikipedia_search),
-    FunctionTool.from_defaults(run_python_with_output),
-    FunctionTool.from_defaults(get_youtube_transcript),
-    FunctionTool.from_defaults(transcribe_audio),
-    FunctionTool.from_defaults(extract_excel_total_food_sales),
+    FunctionTool.from_defaults(search_wikipedia, name="search_wikipedia", description="Search Wikipedia for facts and lists."),
+    FunctionTool.from_defaults(run_python_code, name="run_python", description="Run Python code for logic, math, or set processing."),
+    FunctionTool.from_defaults(get_youtube_transcript, name="get_youtube_transcript", description="Fetch transcript from YouTube video by URL."),
+    FunctionTool.from_defaults(transcribe_audio, name="transcribe_audio", description="Transcribe MP3 audio file using Whisper."),
+    FunctionTool.from_defaults(extract_excel_total_food_sales, name="extract_excel_total_food_sales", description="Sum total sales from Excel where category is 'food'.")
 ]
 
-# LLM and Agent
+# Create agent with improved system prompt
 llm = OpenAI(model="gpt-4")
-agent = ReActAgent.from_tools(
+
+agent = FunctionCallingAgent.from_tools(
     tools=TOOLS,
     llm=llm,
     verbose=True,
     system_prompt="""
-You are an expert AI assistant on the GAIA benchmark.
+You are a highly capable AI agent taking the GAIA benchmark test.
+
+You have access to the following tools:
+- Wikipedia search for factual lookups
+- Python runner for math, logic, or text analysis
+- YouTube transcript fetcher (via URL)
+- Audio transcriber (Whisper, MP3)
+- Excel food sales analyzer
 
-Use available tools (Wikipedia, Python, YouTube transcript, audio, Excel).
-Output ONLY the final answer. No reasoning or commentary.
-Format exactly as requested (list, number, name, chess move, currency).
-If tool fails, output "Tool not available".
-""",
+Rules:
+1. Always try to use a tool if relevant.
+2. Return ONLY the final answer in the requested format.
+3. Do not guess. If a tool fails, say "Tool not available".
+4. Follow formats strictly: comma-separated lists, numeric values, chess notation, names only, etc.
+5. Avoid all explanation unless requested.
+"""
 )
 
 def answer_question_sync(question: str) -> str: