Update agent.py

agent.py

@@ -1,96 +1,107 @@
 import os
 import asyncio
 import re
-from typing import Any
-
-from llama_index.llms.openai import OpenAI
+from openai import OpenAI
 from llama_index.core.agent.react import ReActAgent
 from llama_index.core.tools import FunctionTool
+from duckduckgo_search import DDGS
+
+# --- Custom DuckDuckGo Search Tool ---
+class DuckDuckGoSearchTool:
+    def __init__(self):
+        self.metadata = {
+            "name": "duckduckgo_search",
+            "description": "Search web via DuckDuckGo and return brief summaries."
+        }
+    def __call__(self, query: str) -> str:
+        try:
+            with DDGS() as ddg:
+                results = ddg.text(query=query, region="wt-wt", max_results=3)
+                return "\n".join(r.get('body', '') for r in results if r.get('body'))
+        except Exception as e:
+            return f"ERROR: {e}"
 
-# Correct import for LlamaIndex >= 0.10
-from llama_index.tools.duckduckgo_search import DuckDuckGoSearchTool
+# --- Other Tools ---
 
-# Simple tool: Evaluate Python code for math/code questions
 def eval_python_code(code: str) -> str:
     try:
         return str(eval(code, {"__builtins__": {}}))
     except Exception as e:
         return f"ERROR: {e}"
 
-# Strict output formatting for GAIA
 def format_gaia_answer(answer: str, question: str = "") -> str:
     if not answer:
         return ""
     answer = re.sub(r'(?i)final answer:?\s*', '', answer).strip()
-    answer = re.sub(r'(?i)i(\'?m| cannot| can\'t| unable to| apologize| not available|process the file).*', '', answer).strip()
-    if answer.startswith('"') and answer.endswith('"'): answer = answer[1:-1]
-    if answer.startswith('[') and answer.endswith(']'): answer = answer[1:-1]
-    if not re.match(r'^[A-Za-z]+\.$', answer): answer = re.sub(r'\.$', '', answer)
+    answer = re.sub(r'(?i)i(\'?m| cannot| can\'t| unable| apologize| not available|process the file).*', '', answer).strip()
+    if answer.startswith('"') and answer.endswith('"'):
+        answer = answer[1:-1]
+    if answer.startswith('[') and answer.endswith(']'):
+        answer = answer[1:-1]
+    if not re.match(r'^[A-Za-z]+\.$', answer):
+        answer = re.sub(r'\.$', '', answer)
+
     if question:
-        if re.search(r'how many|number of|at bats|total sales|albums|output.*python|highest number', question, re.I):
+        num_q = re.search(r'how many|number of|at bats|total sales|albums|output.*python|highest number', question, re.I)
+        list_q = re.search(r'list|comma.*separated|page numbers', question, re.I)
+        if num_q:
             num = re.search(r'(\$?\d[\d,\.]*)', answer)
-            if num: return num.group(1).replace(',', '')
-        if 'first name' in question: return answer.split()[0]
-        if 'surname' in question: return answer.split()[-1]
-        if 'city' in question: return answer.split()[0]
+            if num:
+                return num.group(1).replace(',', '')
+        if 'first name' in question:
+            return answer.split()[0]
+        if 'surname' in question:
+            return answer.split()[-1]
+        if 'city' in question:
+            return answer.split()[0]
         if re.search(r'IOC country code|award number|NASA', question, re.I):
             code = re.search(r'[A-Z0-9]{3,}', answer)
-            if code: return code.group(0)
-        if re.search(r'list|comma.*separated|page numbers', question, re.I):
+            if code:
+                return code.group(0)
+        if list_q:
             items = [x.strip('",.').lower() for x in re.split(r'[,\n]', answer) if x.strip()]
             if 'page numbers' in question:
-                nums = [int(x) for x in re.findall(r'\d+', answer)]
-                return ', '.join(str(n) for n in sorted(nums))
-            if 'ingredient' in question or 'vegetable' in question:
-                merged = []
-                skip = False
-                for i, item in enumerate(items):
-                    if skip: skip = False; continue
-                    if i+1 < len(items) and item in ['sweet', 'green', 'lemon', 'ripe', 'whole', 'fresh']:
-                        merged.append(f"{item} {items[i+1]}")
+                nums = sorted(int(x) for x in re.findall(r'\d+', answer))
+                return ', '.join(str(n) for n in nums)
+            if 'ingredient' in question or 'vegetable' in question or 'grocery' in question:
+                merged, skip = [], False
+                for i, x in enumerate(items):
+                    if skip:
+                        skip = False
+                        continue
+                    if i+1 < len(items) and x in ['sweet','green','lemon','ripe','whole','fresh']:
+                        merged.append(f"{x} {items[i+1]}")
                         skip = True
-                    else: merged.append(item)
-                merged = sorted(set(merged))
-                return ', '.join(merged)
+                    else:
+                        merged.append(x)
+                return ', '.join(sorted(set(merged)))
             return ', '.join(items)
-    return answer.strip().rstrip('.').strip()
 
-# LLM setup
-llm = OpenAI(model="gpt-4o", api_key=os.environ.get("OPENAI_API_KEY"))
+    return answer.strip().rstrip('.')
+
+# --- LLM & Tools Setup ---
+llm = OpenAI(model="gpt-4o", api_key=os.getenv("OPENAI_API_KEY"))
 
-# Tool registry
 tools = [
-    DuckDuckGoSearchTool(),
-    FunctionTool.from_defaults(
-        eval_python_code,
-        name="python_eval",
-        description="Evaluate simple Python code and return result as string. Use for math or code output."
-    ),
-    FunctionTool.from_defaults(
-        format_gaia_answer,
-        name="format_gaia_answer",
-        description="Postprocess and enforce strict GAIA format on answers given a question."
-    ),
+    FunctionTool.from_defaults(DuckDuckGoSearchTool(), name="duckduckgo_search", description="Searches the web via DuckDuckGo"),
+    FunctionTool.from_defaults(eval_python_code, name="python_eval", description="Evaluate Python code"),
+    FunctionTool.from_defaults(format_gaia_answer, name="format_gaia_answer", description="Strict GAIA output formatting")
 ]
 
-# Main agent
 agent = ReActAgent.from_tools(
     tools=tools,
     llm=llm,
-    system_prompt="You are a helpful GAIA benchmark agent. For every question, use the best tools available and always return only the final answer in the strict GAIA-required format—never explain, never apologize.",
+    system_prompt="You're a GAIA benchmark agent. Use tools and always output only the final answer in strict format—no explanation or apology.",
     verbose=False
 )
 
-# Async entrypoint
 async def answer_question(question: str, task_id: str = None, file_path: str = None) -> str:
     result = await agent.achat(question)
     return result.response
 
-# Synchronous wrapper
 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))
 
-# For compatibility with app.py (GAIAAgent class)
 class GaiaAgent:
-    def __call__(self, question: str, task_id: str = None, file_path: str = None) -> str:
-        return answer_question_sync(question, task_id, file_path)
+    def __call__(self, question: str, task_id: str = None) -> str:
+        return answer_question_sync(question, task_id)