Update agent.py

agent.py

@@ -1,16 +1,12 @@
 import os
-import asyncio
 import re
 from openai import OpenAI as OpenAIClient
-from llama_index.llms.openai import OpenAI
-from llama_index.core.agent.react import ReActAgent
-from llama_index.core.tools import FunctionTool
 from duckduckgo_search import DDGS
 
 def duckduckgo_search(query: str) -> str:
     try:
         with DDGS() as ddg:
-            results = ddg.text(query=query, region="wt-wt", max_results=3)
+            results = ddg.text(query=query, region="wt-wt", max_results=5)
             return "\n".join(r.get('body', '') for r in results if r.get('body'))
     except Exception as e:
         return f"ERROR: {e}"
@@ -65,28 +61,51 @@ def format_gaia_answer(answer: str, question: str = "") -> str:
             return ', '.join(items)
     return ans.strip().rstrip('.')
 
-llm = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))  # no model= param
-
-tools = [
-    FunctionTool.from_defaults(duckduckgo_search, name="duckduckgo_search", description="Web search tool"),
-    FunctionTool.from_defaults(eval_python_code, name="python_eval", description="Evaluate Python code"),
-    FunctionTool.from_defaults(format_gaia_answer, name="format_gaia_answer", description="GAIA output formatting")
-]
-
-agent = ReActAgent.from_tools(
-    tools=tools,
-    llm=llm,
-    system_prompt="You're a GAIA benchmark agent. Always use tools and strictly output only the final answer in GAIA format—no apologies or explanations.",
-    verbose=False
-)
-
-async def answer_question(question: str, task_id: str = None, file_path: str = None) -> str:
-    resp = await agent.achat(question)
-    return resp.response
-
-def answer_question_sync(question: str, task_id: str = None, file_path: str = None) -> str:
-    return asyncio.run(answer_question(question, task_id, file_path))
-
 class GaiaAgent:
+    def __init__(self):
+        self.llm = OpenAIClient(api_key=os.getenv("OPENAI_API_KEY"))
+
     def __call__(self, question: str, task_id: str = None) -> str:
-        return answer_question_sync(question, task_id)
+        # Route to tools by keyword
+        if any(kw in question.lower() for kw in ["who", "when", "what", "which", "how many", "number", "name", "albums", "surname", "at bats", "nasa", "city", "winner", "code"]):
+            web_result = duckduckgo_search(question)
+            llm_answer = self.llm.chat.completions.create(
+                model="gpt-4o",
+                messages=[
+                    {"role": "system", "content": "You are a research assistant. Based on the following web search results and question, answer strictly and concisely for the GAIA benchmark. Only the answer, no explanations."},
+                    {"role": "user", "content": f"Web search results:\n{web_result}\n\nQuestion: {question}"}
+                ],
+                temperature=0.0,
+                max_tokens=256,
+            ).choices[0].message.content.strip()
+            return format_gaia_answer(llm_answer, question)
+        # Code/math
+        if "output" in question.lower() and "python" in question.lower():
+            code_match = re.search(r'```python(.*?)```', question, re.DOTALL)
+            code = code_match.group(1) if code_match else ""
+            result = eval_python_code(code)
+            return format_gaia_answer(result, question)
+        # List/ingredients/vegetables
+        if "list" in question.lower() or "ingredient" in question.lower() or "vegetable" in question.lower():
+            web_result = duckduckgo_search(question)
+            llm_answer = self.llm.chat.completions.create(
+                model="gpt-4o",
+                messages=[
+                    {"role": "system", "content": "You are a research assistant. Based on the following web search results and question, answer strictly and concisely for the GAIA benchmark. Only the answer, no explanations."},
+                    {"role": "user", "content": f"Web search results:\n{web_result}\n\nQuestion: {question}"}
+                ],
+                temperature=0.0,
+                max_tokens=256,
+            ).choices[0].message.content.strip()
+            return format_gaia_answer(llm_answer, question)
+        # Fallback
+        llm_answer = self.llm.chat.completions.create(
+            model="gpt-4o",
+            messages=[
+                {"role": "system", "content": "You are a research assistant. Answer strictly and concisely for the GAIA benchmark. Only the answer, no explanations."},
+                {"role": "user", "content": question}
+            ],
+            temperature=0.0,
+            max_tokens=256,
+        ).choices[0].message.content.strip()
+        return format_gaia_answer(llm_answer, question)