Update agent.py

agent.py

@@ -1,28 +1,28 @@
 import os
 import asyncio
 import re
-from openai import OpenAI
+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
 
-# --- Custom DuckDuckGo Search Tool ---
+# --- Custom DuckDuckGo Tool ---
 class DuckDuckGoSearchTool:
     def __init__(self):
         self.metadata = {
             "name": "duckduckgo_search",
-            "description": "Search web via DuckDuckGo and return brief summaries."
+            "description": "Search the web via DuckDuckGo."
         }
     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'))
+                return "\n".join(r.get('body','') for r in results if r.get('body'))
         except Exception as e:
             return f"ERROR: {e}"
 
-# --- Other Tools ---
-
+# --- Tools ---
 def eval_python_code(code: str) -> str:
     try:
         return str(eval(code, {"__builtins__": {}}))
@@ -32,36 +32,32 @@ def eval_python_code(code: str) -> str:
 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| 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)
+    ans = re.sub(r'(?i)final answer:?\s*', '', answer).strip()
+    ans = re.sub(r'(?i)i(\'?m| cannot| can\'t| unable| apologize| not available).*', '', ans).strip()
+    if ans.startswith('"') and ans.endswith('"'):
+        ans = ans[1:-1]
+    if ans.startswith('[') and ans.endswith(']'):
+        ans = ans[1:-1]
+    if not re.match(r'^[A-Za-z]+\.$', ans):
+        ans = re.sub(r'\.$', '', ans)
 
     if question:
-        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 re.search(r'how many|number of|at bats|total sales|albums|output.*python|highest number', question, re.I):
+            m = re.search(r'(\$?\d[\d,\.]*)', ans)
+            if m: return m.group(1).replace(',', '')
         if 'first name' in question:
-            return answer.split()[0]
+            return ans.split()[0]
         if 'surname' in question:
-            return answer.split()[-1]
+            return ans.split()[-1]
         if 'city' in question:
-            return answer.split()[0]
+            return ans.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 list_q:
-            items = [x.strip('",.').lower() for x in re.split(r'[,\n]', answer) if x.strip()]
+            c = re.search(r'[A-Z0-9]{3,}', ans)
+            if c: return c.group(0)
+        if re.search(r'list|comma.*separated|page numbers', question, re.I):
+            items = [x.strip('",.').lower() for x in re.split(r'[,\n]', ans) if x.strip()]
             if 'page numbers' in question:
-                nums = sorted(int(x) for x in re.findall(r'\d+', answer))
+                nums = sorted(int(x) for x in re.findall(r'\d+', ans))
                 return ', '.join(str(n) for n in nums)
             if 'ingredient' in question or 'vegetable' in question or 'grocery' in question:
                 merged, skip = [], False
@@ -76,28 +72,27 @@ def format_gaia_answer(answer: str, question: str = "") -> str:
                         merged.append(x)
                 return ', '.join(sorted(set(merged)))
             return ', '.join(items)
+    return ans.strip().rstrip('.')
 
-    return answer.strip().rstrip('.')
-
-# --- LLM & Tools Setup ---
+# --- LLM setup ---
 llm = OpenAI(model="gpt-4o", api_key=os.getenv("OPENAI_API_KEY"))
 
 tools = [
-    FunctionTool.from_defaults(DuckDuckGoSearchTool(), name="duckduckgo_search", description="Searches the web via DuckDuckGo"),
+    FunctionTool.from_defaults(DuckDuckGoSearchTool(), 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="Strict GAIA output formatting")
+    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. Use tools and always output only the final answer in strict format—no explanation or apology.",
+    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:
-    result = await agent.achat(question)
-    return result.response
+    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))