refactor: enhance web search functionality with engine fallback and retry mechanism

This commit is contained in:
ca-ke 2025-03-13 14:17:57 -03:00
parent c4d628cc4e
commit cba275d405

View File

@ -4,14 +4,13 @@ from typing import List
from app.tool.base import BaseTool from app.tool.base import BaseTool
from app.config import config from app.config import config
from app.tool.search import WebSearchEngine, BaiduSearchEngine, GoogleSearchEngine, DuckDuckGoSearchEngine from app.tool.search import WebSearchEngine, BaiduSearchEngine, GoogleSearchEngine, DuckDuckGoSearchEngine
from tenacity import retry, stop_after_attempt, wait_exponential
class WebSearch(BaseTool): class WebSearch(BaseTool):
name: str = "web_search" name: str = "web_search"
description: str = """Perform a web search and return a list of relevant links. description: str = """Perform a web search and return a list of relevant links.
Use this tool when you need to find information on the web, get up-to-date data, or research specific topics. This function attempts to use the primary search engine API to get up-to-date results.
The tool returns a list of URLs that match the search query. If an error occurs, it falls back to an alternative search engine."""
"""
parameters: dict = { parameters: dict = {
"type": "object", "type": "object",
"properties": { "properties": {
@ -44,20 +43,48 @@ The tool returns a list of URLs that match the search query.
Returns: Returns:
List[str]: A list of URLs matching the search query. List[str]: A list of URLs matching the search query.
""" """
# Run the search in a thread pool to prevent blocking engine_order = self._get_engine_order()
for engine_name in engine_order:
engine = self._search_engine[engine_name]
try:
links = await self._perform_search_with_engine(engine, query, num_results)
if links:
return links
except Exception as e:
print(f"Search engine '{engine_name}' failed with error: {e}")
return []
def _get_engine_order(self) -> List[str]:
"""
Determines the order in which to try search engines.
Preferred engine is first (based on configuration), followed by the remaining engines.
Returns:
List[str]: Ordered list of search engine names.
"""
preferred = "google"
if config.search_config and config.search_config.engine:
preferred = config.search_config.engine.lower()
engine_order = []
if preferred in self._search_engine:
engine_order.append(preferred)
for key in self._search_engine:
if key not in engine_order:
engine_order.append(key)
return engine_order
@retry(
stop=stop_after_attempt(3),
wait=wait_exponential(multiplier=1, min=1, max=10),
)
async def _perform_search_with_engine(
self,
engine: WebSearchEngine,
query: str,
num_results: int,
) -> List[str]:
loop = asyncio.get_event_loop() loop = asyncio.get_event_loop()
search_engine = self.get_search_engine() return await loop.run_in_executor(
links = await loop.run_in_executor( None, lambda: list(engine.perform_search(query, num_results=num_results))
None, lambda: list(search_engine.perform_search(query, num_results=num_results))
) )
return links
def get_search_engine(self) -> WebSearchEngine:
"""Determines the search engine to use based on the configuration."""
default_engine = self._search_engine.get("google")
if config.search_config is None:
return default_engine
else:
engine = config.search_config.engine.lower()
return self._search_engine.get(engine, default_engine)