Update agent.py

agent.py

@@ -1,43 +1,79 @@
-# agent.py — stable ReActAgent version with API key check and error logging
+# agent.py — full GAIA-ready agent with tools for web, audio, Excel, Python
 
-import asyncio
 import os
+import asyncio
 from llama_index.llms.openai import OpenAI
-from llama_index.core.tools import FunctionTool
 from llama_index.core.agent.react.base import ReActAgent
+from llama_index.core.tools import FunctionTool
 
-from langchain_community.tools import DuckDuckGoSearchRun, WikipediaQueryRun
+from langchain_community.tools.wikipedia.tool import WikipediaTool
 from langchain_experimental.tools.python.tool import PythonREPLTool
 from langchain_community.document_loaders import YoutubeLoader
 
-# Optional log check for Hugging Face Secrets
+import openai_whisper as whisper
+import openpyxl
+
+# Confirm OpenAI API key
 if os.getenv("OPENAI_API_KEY"):
     print("✅ Detected OPENAI_API_KEY in environment")
 else:
-    print("⚠️  Missing OPENAI_API_KEY — LLM may fail")
-
-# Tool wrappers
-def search_duckduckgo(query: str) -> str:
-    return DuckDuckGoSearchRun().run(query)
+    print("⚠️  Missing OPENAI_API_KEY — LLM may not work")
 
-def search_wikipedia(query: str) -> str:
-    return WikipediaQueryRun(api_wrapper=None).run(query)
+# --- Web tools ---
+def wikipedia_search(query: str) -> str:
+    return WikipediaTool().run(query)
 
-def run_python(code: str) -> str:
-    return PythonREPLTool().run(code)
+# --- 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}"
 
+# --- YouTube (fallback placeholder) ---
 def get_youtube_transcript(url: str) -> str:
-    loader = YoutubeLoader.from_youtube_url(url, add_video_info=False)
-    docs = loader.load()
-    return " ".join(doc.page_content for doc in docs)
+    try:
+        loader = YoutubeLoader.from_youtube_url(url, add_video_info=False)
+        docs = loader.load()
+        return " ".join(doc.page_content for doc in docs)
+    except Exception as e:
+        return f"[YOUTUBE ERROR] {e}"
+
+# --- Whisper transcription ---
+def transcribe_audio(file_path: str) -> str:
+    try:
+        model = whisper.load_model("base")
+        result = model.transcribe(file_path)
+        return result['text']
+    except Exception as e:
+        return f"[AUDIO ERROR] {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)
+        return f"${total:.2f}"
+    except Exception as e:
+        return f"[EXCEL ERROR] {e}"
 
+# --- Tool list ---
 TOOLS = [
-    FunctionTool.from_defaults(search_duckduckgo),
-    FunctionTool.from_defaults(search_wikipedia),
-    FunctionTool.from_defaults(run_python),
+    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),
 ]
 
+# --- LLM and Agent ---
 llm = OpenAI(model="gpt-4")
 
 agent = ReActAgent.from_tools(
@@ -47,16 +83,22 @@ agent = ReActAgent.from_tools(
     system_prompt="""
 You are an expert AI assistant participating in the GAIA benchmark.
 
-Your task is to answer questions as precisely and exactly as possible. Each answer is evaluated automatically—strict formatting matters.
+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 explanation or commentary.
-2. Format exactly as requested (lists, names, numbers, chess moves, currency).
-3. Use tools as needed — no guessing.
+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".
 """
 )
 
-# Safe sync wrapper
+# --- Run function ---
 def answer_question_sync(question: str) -> str:
     try:
         response = agent.chat(question)
@@ -70,6 +112,5 @@ def answer_question_sync(question: str) -> str:
         print(f"❌ Exception while answering: {e}")
         return f"[ERROR] {e}"
 
-# Async wrapper for FastAPI/Gradio
 async def answer_question(question: str) -> str:
     return answer_question_sync(question)