AlfredAgent / tools /web_search.py
Etqad's picture
Upload agent
051eb6b verified
raw
history blame
2.48 kB
from typing import Any, Optional
from smolagents.tools import Tool
import duckduckgo_search
import time
class DuckDuckGoSearchTool(Tool):
"""Web search tool that performs searches using the DuckDuckGo search engine.
Args:
max_results (`int`, default `10`): Maximum number of search results to return.
rate_limit (`float`, default `1.0`): Maximum queries per second. Set to `None` to disable rate limiting.
**kwargs: Additional keyword arguments for the `DDGS` client.
Examples:
```python
>>> from smolagents import DuckDuckGoSearchTool
>>> web_search_tool = DuckDuckGoSearchTool(max_results=5, rate_limit=2.0)
>>> results = web_search_tool("Hugging Face")
>>> print(results)
```
"""
name = "web_search"
description = "Performs a duckduckgo web search based on your query (think a Google search) then returns the top search results."
inputs = {'query': {'type': 'string', 'description': 'The search query to perform.'}}
output_type = "string"
def __init__(self, max_results: int = 10, rate_limit: float | None = 1.0, **kwargs):
super().__init__()
self.max_results = max_results
self.rate_limit = rate_limit
self._min_interval = 1.0 / rate_limit if rate_limit else 0.0
self._last_request_time = 0.0
try:
from duckduckgo_search import DDGS
except ImportError as e:
raise ImportError(
"You must install package `duckduckgo_search` to run this tool: for instance run `pip install duckduckgo-search`."
) from e
self.ddgs = DDGS(**kwargs)
def forward(self, query: str) -> str:
self._enforce_rate_limit()
results = self.ddgs.text(query, max_results=self.max_results)
if len(results) == 0:
raise Exception("No results found! Try a less restrictive/shorter query.")
postprocessed_results = [f"[{result['title']}]({result['href']})\n{result['body']}" for result in results]
return "## Search Results\n\n" + "\n\n".join(postprocessed_results)
def _enforce_rate_limit(self) -> None:
import time
# No rate limit enforced
if not self.rate_limit:
return
now = time.time()
elapsed = now - self._last_request_time
if elapsed < self._min_interval:
time.sleep(self._min_interval - elapsed)
self._last_request_time = time.time()