Update agent.py

agent.py

@@ -1,70 +1,66 @@
-# agent.py — full GAIA-ready agent with tools for web, audio, Excel, Python
+# agent.py — full GAIA-ready agent with working WikipediaQueryRun + 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.tools import FunctionTool
 
-from langchain_community.tools.wikipedia.tool import WikipediaTool
+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
 
-import openai_whisper as whisper
+import whisper
 import openpyxl
 
-# Confirm OpenAI API key
+# Check OpenAI key
 if os.getenv("OPENAI_API_KEY"):
-    print("✅ Detected OPENAI_API_KEY in environment")
+    print("✅ Detected OPENAI_API_KEY")
 else:
-    print("⚠️  Missing OPENAI_API_KEY — LLM may not work")
+    print("⚠️ Missing OPENAI_API_KEY – LLM may fail")
+
+# Tools definitions
 
-# --- Web tools ---
+api_wrapper = WikipediaAPIWrapper(top_k_results=1, doc_content_chars_max=1000)
 def wikipedia_search(query: str) -> str:
-    return WikipediaTool().run(query)
+    return WikipediaQueryRun(api_wrapper=api_wrapper).run({"query": query})
 
-# --- Python with output ---
 def run_python_with_output(code: str) -> str:
-    try:
-        if "print(" not in code:
-            code = f"print({code})" if not code.strip().endswith("\n") else f"print({code.strip()})"
-        return PythonREPLTool().run(code)
-    except Exception as e:
-        return f"[PYTHON ERROR] {e}"
+    if "print(" not in code:
+        code = f"print({code})"
+    return PythonREPLTool().run(code)
 
-# --- YouTube (fallback placeholder) ---
 def get_youtube_transcript(url: str) -> str:
     try:
         loader = YoutubeLoader.from_youtube_url(url, add_video_info=False)
         docs = loader.load()
-        return " ".join(doc.page_content for doc in docs)
+        return " ".join(d.page_content for d in docs)
     except Exception as e:
-        return f"[YOUTUBE ERROR] {e}"
+        return "[YOUTUBE ERROR] " + str(e)
 
-# --- Whisper transcription ---
 def transcribe_audio(file_path: str) -> str:
     try:
         model = whisper.load_model("base")
-        result = model.transcribe(file_path)
-        return result['text']
+        res = model.transcribe(file_path)
+        return res["text"]
     except Exception as e:
-        return f"[AUDIO ERROR] {e}"
+        return "[AUDIO ERROR] " + str(e)
 
-# --- Excel sales extraction ---
 def extract_excel_total_food_sales(file_path: str) -> str:
     try:
         wb = openpyxl.load_workbook(file_path)
         sheet = wb.active
-        total = 0
-        for row in sheet.iter_rows(min_row=2, values_only=True):
-            category, amount = row[1], row[2]
-            if isinstance(category, str) and 'food' in category.lower():
-                total += float(amount)
+        total = 0.0
+        for _, category, amount in sheet.iter_rows(min_row=2, values_only=True):
+            if isinstance(category, str) and "food" in category.lower():
+                total += float(amount or 0)
         return f"${total:.2f}"
     except Exception as e:
-        return f"[EXCEL ERROR] {e}"
+        return "[EXCEL ERROR] " + str(e)
 
-# --- Tool list ---
+# Assemble tools
 TOOLS = [
     FunctionTool.from_defaults(wikipedia_search),
     FunctionTool.from_defaults(run_python_with_output),
@@ -73,44 +69,31 @@ TOOLS = [
     FunctionTool.from_defaults(extract_excel_total_food_sales),
 ]
 
-# --- LLM and Agent ---
+# LLM and Agent
 llm = OpenAI(model="gpt-4")
-
 agent = ReActAgent.from_tools(
     tools=TOOLS,
     llm=llm,
     verbose=True,
     system_prompt="""
-You are an expert AI assistant participating in the GAIA benchmark.
-
-Your goal is to answer 20 diverse questions using available tools:
-- Wikipedia search
-- Python code runner
-- YouTube transcript
-- MP3 transcription (Whisper)
-- Excel analysis
-
-Rules:
-1. Output only the FINAL answer. No explanations.
-2. Format must match expected output exactly: comma-separated lists, plain names, numeric values, algebraic notation.
-3. Use tools smartly. Don't guess when tools can help.
-4. If tools fail (e.g., YouTube blocked), say clearly: "Tool not available".
-"""
+You are an expert AI assistant on the GAIA benchmark.
+
+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".
+""",
 )
 
-# --- Run function ---
 def answer_question_sync(question: str) -> str:
     try:
-        response = agent.chat(question)
-        if hasattr(response, "response") and hasattr(response.response, "content"):
-            return response.response.content.strip()
-        elif isinstance(response, str):
-            return response.strip()
-        else:
-            return str(response)
+        resp = agent.chat(question)
+        if hasattr(resp, "response") and hasattr(resp.response, "content"):
+            return resp.response.content.strip()
+        return str(resp).strip()
     except Exception as e:
-        print(f"❌ Exception while answering: {e}")
-        return f"[ERROR] {e}"
+        print("❌ Agent exception:", e)
+        return "[ERROR] " + str(e)
 
 async def answer_question(question: str) -> str:
     return answer_question_sync(question)