Update agent.py

agent.py

@@ -1,52 +1,52 @@
+import os
+import re
+import requests
+from openai import OpenAI
+from duckduckgo_search import DDGS
+
+PROMPT = (
+    "You are a general AI assistant. I will ask you a question. "
+    "Report your thoughts, and finish your answer with the following template: "
+    "FINAL ANSWER: [YOUR FINAL ANSWER]. YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. "
+    "If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise. "
+    "If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. "
+    "If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string."
+)
+
+class BasicAgent:
+    def __init__(self):
+        self.llm = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
+        print("BasicAgent initialized.")
+
+    def web_search(self, query: str, max_results: int = 5) -> str:
+        try:
+            with DDGS() as ddgs:
+                results = list(ddgs.text(query, max_results=max_results))
+            if not results:
+                return ""
+            formatted_results = ""
+            for i, result in enumerate(results, 1):
+                title = result.get('title', '')
+                body = result.get('body', '')
+                href = result.get('href', '')
+                formatted_results += f"{i}. {title}\n   URL: {href}\n   Description: {body}\n\n"
+            return formatted_results
+        except Exception as e:
+            return ""
+
+    def __call__(self, question: str, task_id: str = None) -> str:
+        search_snippet = self.web_search(question)
+        prompt = PROMPT + f"\n\nWeb search results:\n{search_snippet}\n\nQuestion: {question}"
+        response = self.llm.chat.completions.create(
+            model="gpt-4o",
+            messages=[{"role": "system", "content": prompt}],
+            temperature=0.0,
+            max_tokens=512,
+        )
         answer = response.choices[0].message.content.strip()
         final_line = ""
         for line in answer.splitlines():
             if line.strip().lower().startswith("final answer:"):
                 final_line = line.split(":", 1)[-1].strip(" .\"'")
                 break
-
-        bads = [
-            "", "unknown", "unable to determine", "unable to provide page numbers",
-            "unable to access video content directly", "unable to analyze video content",
-            "unable to determine without code", "unable to determine without file",
-            "follow the steps to locate the paper and find the nasa award number in the acknowledgment section",
-            "i am unable to view images or access external content directly", "unable to determine without access to the file",
-            "no results found", "n/a", "[your final answer]", "i'm sorry", "i apologize"
-        ]
-        if final_line.lower() in bads or final_line.lower().startswith("unable") or final_line.lower().startswith("i'm sorry") or final_line.lower().startswith("i apologize"):
-            # --- Try to extract a plausible answer from web or file ---
-            # Example: For numbers
-            numbers = re.findall(r'\b\d{2,}\b', search_snippet)
-            if numbers:
-                return numbers[0]
-            # Example: For possible names (capitalize words)
-            words = re.findall(r'\b[A-Z][a-z]{2,}\b', search_snippet)
-            if words:
-                return words[0]
-            # Example: For Excel, code, or file extraction, return the result
-            if file_result:
-                file_numbers = re.findall(r'\b\d{2,}\b', file_result)
-                if file_numbers:
-                    return file_numbers[0]
-                file_words = re.findall(r'\b[A-Z][a-z]{2,}\b', file_result)
-                if file_words:
-                    return file_words[0]
-            # --- Try to re-ask the LLM to answer "without apologies" ---
-            retry_prompt = (
-                "Based ONLY on the search results and/or file content above, return a direct answer to the question. "
-                "If you do not know, make your best plausible guess. Do NOT apologize or say you cannot assist. "
-                f"File: {file_result}\n\nWeb: {search_snippet}\n\nQuestion: {question}\nFINAL ANSWER:"
-            )
-            response2 = self.llm.chat.completions.create(
-                model="gpt-4o",
-                messages=[{"role": "system", "content": retry_prompt}],
-                temperature=0.1,
-                max_tokens=128,
-            )
-            retry_answer = response2.choices[0].message.content.strip()
-            for line in retry_answer.splitlines():
-                if line.strip().lower().startswith("final answer:"):
-                    return line.split(":", 1)[-1].strip(" .\"'")
-            if retry_answer:
-                return retry_answer.strip(" .\"'")
-        return final_line
+        return final_line or answer