Update agent.py

agent.py

@@ -12,26 +12,31 @@ from docx import Document
 import fitz  # PyMuPDF
 import requests
 
-class BraveSearchTool:
-    def __init__(self, api_key: str):
+class GoogleCustomSearchTool:
+    def __init__(self, api_key: str, cse_id: str):
         self.api_key = api_key
-        self.base_url = "https://api.search.brave.com/res/v1/web/search"
+        self.cse_id = cse_id
+        self.base_url = "https://www.googleapis.com/customsearch/v1"
 
     def run(self, query: str) -> str:
         try:
             response = requests.get(
                 self.base_url,
-                headers={"Accept": "application/json", "X-Subscription-Token": self.api_key},
-                params={"q": query}
+                params={
+                    "key": self.api_key,
+                    "cx": self.cse_id,
+                    "q": query,
+                },
+                timeout=10,
             )
             response.raise_for_status()
-            results = response.json().get("web", {}).get("results", [])
+            results = response.json().get("items", [])
             if results:
-                return results[0].get("title", "") + ": " + results[0].get("url", "")
+                return results[0].get("title", "") + ": " + results[0].get("link", "")
             else:
                 return "No results found."
         except Exception as e:
-            return f"BraveSearchTool ERROR: {str(e)}"
+            return f"GoogleCustomSearchTool ERROR: {str(e)}"
 
 def classify_question_type(question: str) -> str:
     q = question.lower()
@@ -47,11 +52,13 @@ def classify_question_type(question: str) -> str:
 class Agent:
     def __init__(self):
         gemini_key = os.getenv("GEMINI_API_KEY")
-        brave_key = os.getenv("BRAVE_SEARCH_API_KEY")
+        gcs_key = os.getenv("GOOGLE_API_KEY")
+        gcs_cx = os.getenv("GOOGLE_CSE_ID")
+
         if not gemini_key:
             raise ValueError("GEMINI_API_KEY not found in environment variables.")
-        if not brave_key:
-            raise ValueError("BRAVE_SEARCH_API_KEY not found in environment variables.")
+        if not gcs_key or not gcs_cx:
+            raise ValueError("GOOGLE_API_KEY or GOOGLE_CSE_ID not found in environment variables.")
 
         llm = ChatGoogleGenerativeAI(
             model="gemini-1.5-pro",
@@ -71,9 +78,9 @@ class Agent:
                 description="Useful for solving math and logical problems through Python."
             ),
             Tool(
-                name="Brave Search",
-                func=BraveSearchTool(api_key=brave_key).run,
-                description="Useful for factual and current event queries using Brave search engine."
+                name="Google Custom Search",
+                func=GoogleCustomSearchTool(api_key=gcs_key, cse_id=gcs_cx).run,
+                description="Useful for factual queries using Google Custom Search."
             )
         ]
 
@@ -94,7 +101,7 @@ class Agent:
 
         type_prefix = f"[Task Type: {task_type.upper()}]\n\n"
         system_prompt = (
-            "You are a member of a multidisciplinary research institute... [truncated for brevity, kept same]"
+            "You are a member of a multidisciplinary research institute. If a file cannot be loaded, do not abandon the task. Use your best judgment based on the task and file name. The file may be unavailable by design; this is part of the test. Always attempt to reason based on partial or inferred data."
         )
 
         file_summary = ""
@@ -125,10 +132,11 @@ class Agent:
                     else:
                         summaries.append(f"{fname}: Unsupported file type {ext}")
                 except Exception as fe:
-                    summaries.append(f"{fname}: ERROR reading file ({fe})")
+                    guessed_type = "spreadsheet" if ext in [".csv", ".tsv", ".xlsx"] else "document" if ext in [".pdf", ".docx"] else "data file"
+                    summaries.append(f"{fname}: Could not load, but based on the file extension, we assume it is a {guessed_type}. Please attempt to reason based on the task.")
 
             file_summary = "\n\n".join(summaries)
-            full_prompt = type_prefix + system_prompt + file_summary + f"\n\nTASK:\n{question}"
+            full_prompt = type_prefix + system_prompt + "\n\n" + file_summary + f"\n\nTASK:\n{question}"
             result = self.agent.run(full_prompt)
             return result.strip()
         except Exception as e: