refactor: use DuckDuckGo and Brave as tools instead of retrievers

This commit is contained in:
Siddhant Rai
2025-06-23 09:22:17 +05:30
parent f0908af3c0
commit a9e4583695
8 changed files with 161 additions and 333 deletions

View File

@@ -25,27 +25,35 @@ class BraveSearchTool(Tool):
else:
raise ValueError(f"Unknown action: {action_name}")
def _web_search(self, query, country="ALL", search_lang="en", count=10,
offset=0, safesearch="off", freshness=None,
result_filter=None, extra_snippets=False, summary=False):
def _web_search(
self,
query,
country="ALL",
search_lang="en",
count=10,
offset=0,
safesearch="off",
freshness=None,
result_filter=None,
extra_snippets=False,
summary=False,
):
"""
Performs a web search using the Brave Search API.
"""
print(f"Performing Brave web search for: {query}")
url = f"{self.base_url}/web/search"
# Build query parameters
params = {
"q": query,
"country": country,
"search_lang": search_lang,
"count": min(count, 20),
"offset": min(offset, 9),
"safesearch": safesearch
"safesearch": safesearch,
}
# Add optional parameters only if they have values
if freshness:
params["freshness"] = freshness
if result_filter:
@@ -54,68 +62,69 @@ class BraveSearchTool(Tool):
params["extra_snippets"] = 1
if summary:
params["summary"] = 1
# Set up headers
headers = {
"Accept": "application/json",
"Accept-Encoding": "gzip",
"X-Subscription-Token": self.token
"X-Subscription-Token": self.token,
}
# Make the request
response = requests.get(url, params=params, headers=headers)
if response.status_code == 200:
return {
"status_code": response.status_code,
"results": response.json(),
"message": "Search completed successfully."
"message": "Search completed successfully.",
}
else:
return {
"status_code": response.status_code,
"message": f"Search failed with status code: {response.status_code}."
"message": f"Search failed with status code: {response.status_code}.",
}
def _image_search(self, query, country="ALL", search_lang="en", count=5,
safesearch="off", spellcheck=False):
def _image_search(
self,
query,
country="ALL",
search_lang="en",
count=5,
safesearch="off",
spellcheck=False,
):
"""
Performs an image search using the Brave Search API.
"""
print(f"Performing Brave image search for: {query}")
url = f"{self.base_url}/images/search"
# Build query parameters
params = {
"q": query,
"country": country,
"search_lang": search_lang,
"count": min(count, 100), # API max is 100
"safesearch": safesearch,
"spellcheck": 1 if spellcheck else 0
"spellcheck": 1 if spellcheck else 0,
}
# Set up headers
headers = {
"Accept": "application/json",
"Accept-Encoding": "gzip",
"X-Subscription-Token": self.token
"X-Subscription-Token": self.token,
}
# Make the request
response = requests.get(url, params=params, headers=headers)
if response.status_code == 200:
return {
"status_code": response.status_code,
"results": response.json(),
"message": "Image search completed successfully."
"message": "Image search completed successfully.",
}
else:
return {
"status_code": response.status_code,
"message": f"Image search failed with status code: {response.status_code}."
"message": f"Image search failed with status code: {response.status_code}.",
}
def get_actions_metadata(self):
@@ -130,42 +139,14 @@ class BraveSearchTool(Tool):
"type": "string",
"description": "The search query (max 400 characters, 50 words)",
},
# "country": {
# "type": "string",
# "description": "The 2-character country code (default: US)",
# },
"search_lang": {
"type": "string",
"description": "The search language preference (default: en)",
},
# "count": {
# "type": "integer",
# "description": "Number of results to return (max 20, default: 10)",
# },
# "offset": {
# "type": "integer",
# "description": "Pagination offset (max 9, default: 0)",
# },
# "safesearch": {
# "type": "string",
# "description": "Filter level for adult content (off, moderate, strict)",
# },
"freshness": {
"type": "string",
"description": "Time filter for results (pd: last 24h, pw: last week, pm: last month, py: last year)",
},
# "result_filter": {
# "type": "string",
# "description": "Comma-delimited list of result types to include",
# },
# "extra_snippets": {
# "type": "boolean",
# "description": "Get additional excerpts from result pages",
# },
# "summary": {
# "type": "boolean",
# "description": "Enable summary generation in search results",
# }
},
"required": ["query"],
"additionalProperties": False,
@@ -181,37 +162,21 @@ class BraveSearchTool(Tool):
"type": "string",
"description": "The search query (max 400 characters, 50 words)",
},
# "country": {
# "type": "string",
# "description": "The 2-character country code (default: US)",
# },
# "search_lang": {
# "type": "string",
# "description": "The search language preference (default: en)",
# },
"count": {
"type": "integer",
"description": "Number of results to return (max 100, default: 5)",
},
# "safesearch": {
# "type": "string",
# "description": "Filter level for adult content (off, strict). Default: strict",
# },
# "spellcheck": {
# "type": "boolean",
# "description": "Whether to spellcheck provided query (default: true)",
# }
},
"required": ["query"],
"additionalProperties": False,
},
}
},
]
def get_config_requirements(self):
return {
"token": {
"type": "string",
"description": "Brave Search API key for authentication"
"type": "string",
"description": "Brave Search API key for authentication",
},
}
}