Update tools/WebSearchTool.py

tools/WebSearchTool.py

@@ -1,66 +1,42 @@
 import os
 from smolagents import Tool
-from googleapiclient.discovery import build
-class GoogleSearchTool(Tool):
+from tavily import TavilyClient
+class TavilySearchTool(Tool):
     """
-    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.
+    General-purpose Tavily web search.
+    Use this FIRST before Wikipedia or other tools.
+    Returns summarized results with sources.
     """
-    # 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
+    
+    name = "tavily_search"
+    description = (
+        "Use Tavily to search the web for up-to-date information. "
+        "Returns summarized content and source links."
+    )
+
     inputs = {
-        "query": {"type": "string", "description": "The search term to look up."}
+        "query": {"type": "string", "description": "Search query."}
     }
+
     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")
-        # 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
-        )
+        api_key = os.getenv("TAVILY_API_KEY")
+        if not api_key:
+            raise ValueError("Missing TAVILY_API_KEY in environment variables.")
+        self.client = TavilyClient(api_key=api_key)
+
     def forward(self, query: str) -> str:
-        """
-        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 of search results, or an error message.
-        """
-        print(f"Executing Google search for: '{query}'")
         try:
-            # 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
+            response = self.client.search(query=query, max_results=5)
+            out = []
+            for r in response.get("results", []):
+                out.append(
+                    f"TITLE: {r.get('title')}\n"
+                    f"CONTENT: {r.get('content')}\n"
+                    f"SOURCE: {r.get('url')}"
                 )
-            # Join the results with a clear separator
-            return "\n\n---SEPARATOR---\n\n".join(search_results)
+            return "\n\n---SEPARATOR---\n\n".join(out)
         except Exception as e:
-            # Provide an informative error message upon API failure
-            return f"Error during Google Search API call: {e}"
+            return f"Tavily search error: {e}"