Update tools/WebSearchTool.py

tools/WebSearchTool.py

@@ -1,80 +1,66 @@
-import requests
+import os
 from smolagents import Tool
-import json
-
-class DuckDuckGoSearchTool(Tool):
+from googleapiclient.discovery import build
+class GoogleSearchTool(Tool):
     """
-    A keyless web search tool using the DuckDuckGo Instant Answer API. 
-    This is a final fallback option when other robust search APIs are unavailable.
-    It primarily returns a single, instant answer or the top result snippet.
+    A tool to perform web search Google searches using the Custom Search Engine (CSE) API.
+    Use this tool first to get the necessary information for calculation or further tools.
     """
-    name = "duckduckgo_search"
-    description = "Use the DuckDuckGo API to fetch instant answers or the top search snippet for general knowledge questions. This tool is quick but may not provide multiple detailed results."
-    
+    # 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."
+    # Define the required input structure for the agent framework
     inputs = {
         "query": {"type": "string", "description": "The search term to look up."}
     }
     output_type = "string"
-
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
-        # DuckDuckGo API endpoint for instant answers (format=json, no_html, skip_disambig)
-        self.endpoint = "https://api.duckduckgo.com/"
-
+        # Retrieve credentials from environment variables
+        self.api_key = os.getenv("GOOGLE_API_KEY")
+        self.cse_id = os.getenv("GOOGLE_CSE_ID")
+        # Check for mandatory credentials
+        if not self.api_key or not self.cse_id:
+            raise ValueError("GOOGLE_API_KEY or GOOGLE_CSE_ID secret not found. Check environment variables.")
+        # 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 DuckDuckGo Instant Answer query.
-        
+        Executes a Google search query and formats the top results (up to 3).
+        The output is formatted to prioritize the Title and Snippet (informational content)
+        over the URL to address the confusion that it was only returning links.
         Args:
             query: The search term provided by the agent.
-            
         Returns:
-            A formatted string containing the instant answer or top result.
+            A formatted string of search results, or an error message.
         """
-        print(f"Executing DuckDuckGo Search for: '{query}'")
-        
-        params = {
-            "q": query,
-            "format": "json",
-            "no_html": "1",          # Ensure content is plain text
-            "skip_disambig": "1"     # Skip pages that offer multiple results
-        }
-
+        print(f"Executing Google search for: '{query}'")
         try:
-            # Make the API request
-            response = requests.get(self.endpoint, params=params, timeout=10)
-            response.raise_for_status() 
-            data = response.json()
-
-            # 1. Prioritize the Instant Answer (a direct, concise answer)
-            instant_answer = data.get('AbstractText')
-            instant_source = data.get('AbstractURL')
-            
-            if instant_answer and instant_source:
-                return (
-                    f"RESULT 1 (Instant Answer): '{query}'\n"
-                    f"CONTENT: {instant_answer}\n"
-                    f"SOURCE: {instant_source}"
+            # Execute the search request for up to 3 results
+            res = self.service.cse().list(
+                q=query,
+                cx=self.cse_id,
+                num=3
+            ).execute()
+            items = res.get('items', [])
+            if not items:
+                # Return the specific failure message expected by the agent
+                return "XX record info: No results found."
+            search_results = []
+            for i, item in enumerate(items):
+                # We've adjusted the format here to put the snippet first,
+                # which is usually the most important information for the agent.
+                search_results.append(
+                    f"RESULT {i+1}: '{item.get('title')}'\n"
+                    f"CONTENT: {item.get('snippet')}\n"
+                    f"SOURCE: {item.get('link')}" # The URL is now clearly labeled as the SOURCE
                 )
-
-            # 2. Fallback to the top result snippet if no instant answer is found
-            related_topics = data.get('RelatedTopics', [])
-            if related_topics and isinstance(related_topics[0], dict):
-                # We expect the first related topic to be the main search result snippet
-                top_topic = related_topics[0]
-                text = top_topic.get('Text', 'No snippet available.')
-                url = top_topic.get('FirstURL', 'N/A')
-                
-                return (
-                    f"RESULT 1 (Top Snippet): '{query}'\n"
-                    f"CONTENT: {text}\n"
-                    f"SOURCE: {url}"
-                )
-
-            # If nothing useful is found:
-            return "XX record info: No results found."
-
-        except requests.exceptions.RequestException as e:
-            return f"Error during DuckDuckGo Search API Request: {e}"
+            # Join the results with a clear separator
+            return "\n\n---SEPARATOR---\n\n".join(search_results)
         except Exception as e:
-            return f"Error processing DuckDuckGo results: {e}"
+            # Provide an informative error message upon API failure
+            return f"Error during Google Search API call: {e}"