extending WebSearchTool

tools.py

@@ -1,10 +1,10 @@
-from smolagents import Tool
+from smolagents import WebSearchTool
 import requests
 import time
 from typing import List, Dict, Any, Optional
 
 
-class MySearchTool(Tool):
+class MySearchTool(WebSearchTool):
     name = "Search tool to avoid rate limit"
     description = "It receives a question and returns a list of search results from various sources. It uses DuckDuckGo as the primary source"
     inputs = {
@@ -15,11 +15,11 @@ class MySearchTool(Tool):
     }
     output_type = "string"
 
-    def __init__(self, query: str):
-        self.query = query
-        self.is_initialized = True
+    def __init__(self, max_results=5):
+        super().__init__()
+        self.max_results = max_results
 
-    def forward(self, max_results: int = 5) -> List[Dict[str, Any]]:
+    def search_duckduckgo(self, query: str) -> list:
         """
         Free search using DuckDuckGo Instant Answer API
 
@@ -34,7 +34,7 @@ class MySearchTool(Tool):
             # DuckDuckGo Instant Answer API (free)
             url = "https://api.duckduckgo.com/"
             params = {
-                "q": self.query,
+                "q": query,
                 "format": "json",
                 "no_html": "1",
                 "skip_disambig": "1",
@@ -42,8 +42,9 @@ class MySearchTool(Tool):
 
             response = requests.get(url, params=params, timeout=10)
             response.raise_for_status()
-            time.sleep(2)  # 2-second delay
             data = response.json()
+            time.sleep(2)  # 2-second delay
+
             results = []
 
             # Process abstract
@@ -58,7 +59,9 @@ class MySearchTool(Tool):
                 )
 
             # Process related topics
-            for topic in data.get("RelatedTopics", [])[: max_results - len(results)]:
+            for topic in data.get("RelatedTopics", [])[
+                : self.max_results - len(results)
+            ]:
                 if isinstance(topic, dict) and "Text" in topic:
                     results.append(
                         {
@@ -69,7 +72,7 @@ class MySearchTool(Tool):
                         }
                     )
 
-            return results[:max_results]
+            return results[: self.max_results]
 
         except Exception as e:
             print(f"DuckDuckGo search failed: {str(e)}")