Update tools/google_search_tool.py

tools/google_search_tool.py

@@ -1,36 +1,37 @@
-from smolagents import Tool
+from tools.base_tool import BaseTool
 from googleapiclient.discovery import build
 import os
 
-
-class GoogleSearchTool(Tool):
+class GoogleSearchTool(BaseTool):
     name = "web_search"
-    description = "Performs a Google search and returns top results in markdown format."
+    description = "Performs a Google web search for a query and returns the top results in markdown format."
 
     inputs = {
         "query": {
             "type": "string",
-            "description": "Search string to query the web.",
+            "description": "Search query for Google."
         }
     }
 
     output_type = "string"
     skip_forward_signature_validation = True
 
-    def __init__(self, api_key: str | None = None, search_engine_id: str | None = None, num_results: int = 10, **kwargs):
-        from dotenv import load_dotenv
-        load_dotenv()
-
-        self.api_key = api_key or os.getenv("GOOGLE_SEARCH_API_KEY")
-        self.search_engine_id = search_engine_id or os.getenv("GOOGLE_SEARCH_ENGINE_ID")
+    def __init__(
+        self,
+        api_key: str | None = None,
+        search_engine_id: str | None = None,
+        num_results: int = 10,
+        **kwargs,
+    ):
+        api_key = api_key or os.getenv("GOOGLE_SEARCH_API_KEY")
+        search_engine_id = search_engine_id or os.getenv("GOOGLE_SEARCH_ENGINE_ID")
 
-        if not self.api_key:
-            raise EnvironmentError("GOOGLE_SEARCH_API_KEY is not configured.")
-        if not self.search_engine_id:
-            raise EnvironmentError("GOOGLE_SEARCH_ENGINE_ID is not configured.")
+        if not api_key or not search_engine_id:
+            raise EnvironmentError("Missing Google Search API credentials.")
 
-        self.search = build("customsearch", "v1", developerKey=self.api_key).cse()
-        self.max_results = num_results
+        self.search_service = build("customsearch", "v1", developerKey=api_key).cse()
+        self.engine_id = search_engine_id
+        self.num_results = num_results
 
         super().__init__(**kwargs)
 
@@ -38,41 +39,40 @@ class GoogleSearchTool(Tool):
         return {}
 
     def forward(self, query: str, *args, **kwargs) -> str:
-        params = {
+        search_args = {
             "q": query,
-            "cx": self.search_engine_id,
+            "cx": self.engine_id,
             "fields": "items(link,title,snippet)",
-            "num": self.max_results,
+            "num": self.num_results,
         }
+        search_args.update(self._collect_params(*args, **kwargs))
 
-        params.update(self._collect_params(*args, **kwargs))
-        results = self.search.list(**params).execute()
-
-        if "items" not in results:
-            return "No search results found."
+        response = self.search_service.list(**search_args).execute()
+        items = response.get("items", [])
 
-        return "\n\n".join(
-            f"{item['title']}\n{item['link']}\n{item['snippet']}" for item in results["items"]
-        )
+        if not items:
+            return "No results found."
 
+        results = [f"**{item['title']}**\n{item['link']}\n{item['snippet']}" for item in items]
+        return "\n\n".join(results)
 
 class GoogleSiteSearchTool(GoogleSearchTool):
     name = "site_search"
-    description = "Performs a Google search scoped to a specific domain, such as wikipedia.org or arxiv.org."
+    description = "Searches a specific domain for a query using Google and returns results."
 
     inputs = {
         "query": {
             "type": "string",
-            "description": "Search string to query.",
+            "description": "Search query."
         },
         "site": {
             "type": "string",
-            "description": "Domain to restrict search to (e.g., reddit.com, wikipedia.org).",
+            "description": "Domain to restrict search to."
         },
     }
 
     def _collect_params(self, site: str) -> dict:
         return {
             "siteSearch": site,
-            "siteSearchFilter": "i",
+            "siteSearchFilter": "i"
         }