Update tools/WebSearchTool.py

tools/WebSearchTool.py

@@ -1,24 +1,63 @@
+import os
 from smolagents import Tool
-from duckduckgo_search import DDGS
-class DuckDuckGoSearchTool(Tool):
-    name = "duckduckgo_search"
-    description = "Fetch a short answer from Wikipedia; fallback to DuckDuckGo if needed."
+from googleapiclient.discovery import build
+
+class GoogleSearchTool(Tool):
+    # Use a descriptive name to guide the agent
+    name = "google_search"
+    # Update the description to reflect the new functionality
+    description = "Use Google to find current information and general knowledge. Returns a snippet and URL."
+    
     inputs = {
-        "query": {"type": "string", "description": "Search query"}
+        "query": {"type": "string", "description": "The search term to look up."}
     }
     output_type = "string"
 
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        # Retrieve credentials from environment variables
+        self.api_key = os.getenv("GOOGLE_API_KEY")
+        self.cse_id = os.getenv("GOOGLE_CSE_ID")
+        
+        if not self.api_key or not self.cse_id:
+            # Raise an error if credentials are missing
+            raise ValueError("GOOGLE_API_KEY or GOOGLE_CSE_ID secret not found.")
+
+        # Initialize the Google Custom Search service
+        # 'customsearch' is the API name, 'v1' is the version
+        self.service = build(
+            "customsearch", "v1", developerKey=self.api_key
+        )
+
     def forward(self, query: str) -> str:
+        """
+        Executes a Google search query and formats the top results.
+        """
         try:
-            with DDGS() as ddgs:
-                results = ddgs.text(query, max_results=3)  # returns a generator
-                results = list(results)
+            # Execute the search request for up to 3 results
+            res = self.service.cse().list(
+                q=query,
+                cx=self.cse_id,
+                num=3
+            ).execute()
             
-            if not results:
-                return "No results found."
+            items = res.get('items', [])
             
-            # Return the first result title as concise answer
-            return results[0].get("title", "No results found.")
-        
+            if not items:
+                # Return the specific failure message expected by the agent
+                return "XX record info: No results found."
+
+            search_results = []
+            for item in items:
+                # Format the output clearly for the LLM
+                search_results.append(
+                    f"Title: {item.get('title')}\n"
+                    f"Snippet: {item.get('snippet')}\n"
+                    f"URL: {item.get('link')}"
+                )
+            
+            return "\n---\n".join(search_results)
+
         except Exception as e:
-            return f"Error during search: {e}"
+            # Provide an informative error message
+            return f"Error during Google Search API call: {e}"