fz-mirrors/AstrBot
1
0
Fork 0
mirror of https://github.com/AstrBotDevs/AstrBot synced 2026-07-15 17:30:13 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(core): improve web search API key failover

Browse Source 
Merge PR #8896
This commit is contained in:
NayukiChiba
2026-07-03 22:53:22 +08:00
committed by GitHub
parent 6d798908ad
commit 8a31cf01f2
2 changed files with 572 additions and 141 deletions
Download Patch File Download Diff File Expand all files Collapse all files

486
astrbot/core/tools/web_search_tools.py
View File

@@ -60,12 +60,29 @@ class SearchResult:
@std_dataclass
class _KeyRotator:
"""Concurrency-safe round-robin API key rotator.
Each call returns the current key and advances the index. Search functions
combine this with failover loops to retry with the next configured key.
"""
setting_name: str
provider_name: str
index: int = 0
lock: asyncio.Lock = field(default_factory=asyncio.Lock)
async def get(self, provider_settings: dict) -> str:
"""Return the current key and advance the round-robin index.
Args:
provider_settings: Provider settings containing API key lists.
Returns:
The API key selected for this call.
Raises:
ValueError: If the configured key list is empty or missing.
"""
keys = provider_settings.get(self.setting_name, [])
if not keys:
raise ValueError(
@@ -73,13 +90,21 @@ class _KeyRotator:
)
async with self.lock:
if self.index >= len(keys):
self.index = 0
# Keep the index valid if runtime config reloads shrink the key list.
self.index = self.index % len(keys)
key = keys[self.index]
self.index = (self.index + 1) % len(keys)
return key
# Retry with the next API key when these HTTP statuses indicate key-specific
# auth, quota, or rate-limit failures.
# 401 - Unauthorized, usually invalid or expired key.
# 403 - Forbidden, usually disabled key.
# 429 - Rate limited.
# 432 - Tavily quota exceeded.
_RETRYABLE_HTTP_STATUSES: frozenset[int] = frozenset({401, 403, 429, 432})
_TAVILY_KEY_ROTATOR = _KeyRotator("websearch_tavily_key", "Tavily")
_BOCHA_KEY_ROTATOR = _KeyRotator("websearch_bocha_key", "BoCha")
_BRAVE_KEY_ROTATOR = _KeyRotator("websearch_brave_key", "Brave")
@@ -153,193 +178,372 @@ async def _tavily_search(
provider_settings: dict,
payload: dict,
) -> list[SearchResult]:
tavily_key = await _TAVILY_KEY_ROTATOR.get(provider_settings)
header = {
"Authorization": f"Bearer {tavily_key}",
"Content-Type": "application/json",
}
async with aiohttp.ClientSession(trust_env=True) as session:
async with session.post(
"https://api.tavily.com/search",
json=payload,
headers=header,
) as response:
if response.status != 200:
"""Call the Tavily Search API with API key failover.
Args:
provider_settings: Provider settings containing Tavily API keys.
payload: Request payload for the Tavily search endpoint.
Returns:
Normalized search results.
Raises:
ValueError: If Tavily API keys are not configured.
Exception: If the request fails after all retryable keys are exhausted,
or if a non-retryable HTTP error is returned.
"""
keys = provider_settings.get("websearch_tavily_key", [])
if not keys:
raise ValueError("Error: Tavily API key is not configured in AstrBot.")
# Retry key-specific failures with the next key, but fail fast for
# non-retryable errors such as server-side 5xx responses.
last_error = None
for _ in range(len(keys)):
tavily_key = await _TAVILY_KEY_ROTATOR.get(provider_settings)
header = {
"Authorization": f"Bearer {tavily_key}",
"Content-Type": "application/json",
}
async with aiohttp.ClientSession(trust_env=True) as session:
async with session.post(
"https://api.tavily.com/search",
json=payload,
headers=header,
) as response:
if response.status == 200:
data = await response.json()
return [
SearchResult(
title=item.get("title"),
url=item.get("url"),
snippet=item.get("content"),
favicon=item.get("favicon"),
)
for item in data.get("results", [])
]
reason = await response.text()
# Retryable errors are saved so the final failure is meaningful.
if response.status in _RETRYABLE_HTTP_STATUSES:
last_error = Exception(
f"Tavily web search failed: {reason}, status: {response.status}",
)
continue
raise Exception(
f"Tavily web search failed: {reason}, status: {response.status}",
)
data = await response.json()
return [
SearchResult(
title=item.get("title"),
url=item.get("url"),
snippet=item.get("content"),
favicon=item.get("favicon"),
)
for item in data.get("results", [])
]
if last_error is not None:
raise last_error
raise Exception("Tavily web search failed with all configured keys.")
async def _tavily_extract(provider_settings: dict, payload: dict) -> list[dict]:
tavily_key = await _TAVILY_KEY_ROTATOR.get(provider_settings)
header = {
"Authorization": f"Bearer {tavily_key}",
"Content-Type": "application/json",
}
async with aiohttp.ClientSession(trust_env=True) as session:
async with session.post(
"https://api.tavily.com/extract",
json=payload,
headers=header,
) as response:
if response.status != 200:
"""Call the Tavily Extract API with API key failover.
Args:
provider_settings: Provider settings containing Tavily API keys.
payload: Request payload for the Tavily extract endpoint.
Returns:
Raw Tavily extraction results.
Raises:
ValueError: If Tavily API keys are not configured or no results are
returned.
Exception: If the request fails after all retryable keys are exhausted,
or if a non-retryable HTTP error is returned.
"""
keys = provider_settings.get("websearch_tavily_key", [])
if not keys:
raise ValueError("Error: Tavily API key is not configured in AstrBot.")
last_error = None
for _ in range(len(keys)):
tavily_key = await _TAVILY_KEY_ROTATOR.get(provider_settings)
header = {
"Authorization": f"Bearer {tavily_key}",
"Content-Type": "application/json",
}
async with aiohttp.ClientSession(trust_env=True) as session:
async with session.post(
"https://api.tavily.com/extract",
json=payload,
headers=header,
) as response:
if response.status == 200:
data = await response.json()
results: list[dict] = data.get("results", [])
if not results:
raise ValueError(
"Error: Tavily web searcher does not return any results."
)
return results
reason = await response.text()
if response.status in _RETRYABLE_HTTP_STATUSES:
last_error = Exception(
f"Tavily web search failed: {reason}, status: {response.status}",
)
continue
raise Exception(
f"Tavily web search failed: {reason}, status: {response.status}",
)
data = await response.json()
results: list[dict] = data.get("results", [])
if not results:
raise ValueError(
"Error: Tavily web searcher does not return any results."
)
return results
if last_error is not None:
raise last_error
raise Exception("Tavily web extract failed with all configured keys.")
async def _bocha_search(
provider_settings: dict,
payload: dict,
) -> list[SearchResult]:
bocha_key = await _BOCHA_KEY_ROTATOR.get(provider_settings)
header = {
"Authorization": f"Bearer {bocha_key}",
"Content-Type": "application/json",
# Explicitly disable brotli encoding to avoid aiohttp >= 3.13.3 brotli
# decompression incompatibility (TypeError: process() takes exactly 1 argument).
# See: https://github.com/aio-libs/aiohttp/issues/11898
"Accept-Encoding": "gzip, deflate",
}
async with aiohttp.ClientSession(trust_env=True) as session:
async with session.post(
"https://api.bochaai.com/v1/web-search",
json=payload,
headers=header,
) as response:
if response.status != 200:
"""Call the BoCha Search API with API key failover.
Args:
provider_settings: Provider settings containing BoCha API keys.
payload: Request payload for the BoCha search endpoint.
Returns:
Normalized search results.
Raises:
ValueError: If BoCha API keys are not configured.
Exception: If the request fails after all retryable keys are exhausted,
or if a non-retryable HTTP error is returned.
"""
keys = provider_settings.get("websearch_bocha_key", [])
if not keys:
raise ValueError("Error: BoCha API key is not configured in AstrBot.")
last_error = None
for _ in range(len(keys)):
bocha_key = await _BOCHA_KEY_ROTATOR.get(provider_settings)
header = {
"Authorization": f"Bearer {bocha_key}",
"Content-Type": "application/json",
# Explicitly disable brotli encoding to avoid aiohttp >= 3.13.3
# decompression incompatibility.
# See: https://github.com/aio-libs/aiohttp/issues/11898
"Accept-Encoding": "gzip, deflate",
}
async with aiohttp.ClientSession(trust_env=True) as session:
async with session.post(
"https://api.bochaai.com/v1/web-search",
json=payload,
headers=header,
) as response:
if response.status == 200:
data = await response.json()
rows = data["data"]["webPages"]["value"]
return [
SearchResult(
title=item.get("name"),
url=item.get("url"),
snippet=item.get("snippet"),
favicon=item.get("siteIcon"),
)
for item in rows
]
reason = await response.text()
if response.status in _RETRYABLE_HTTP_STATUSES:
last_error = Exception(
f"BoCha web search failed: {reason}, status: {response.status}",
)
continue
raise Exception(
f"BoCha web search failed: {reason}, status: {response.status}",
)
data = await response.json()
rows = data["data"]["webPages"]["value"]
return [
SearchResult(
title=item.get("name"),
url=item.get("url"),
snippet=item.get("snippet"),
favicon=item.get("siteIcon"),
)
for item in rows
]
if last_error is not None:
raise last_error
raise Exception("BoCha web search failed with all configured keys.")
async def _brave_search(
provider_settings: dict,
payload: dict,
) -> list[SearchResult]:
brave_key = await _BRAVE_KEY_ROTATOR.get(provider_settings)
header = {
"Accept": "application/json",
"X-Subscription-Token": brave_key,
}
async with aiohttp.ClientSession(trust_env=True) as session:
async with session.get(
"https://api.search.brave.com/res/v1/web/search",
params=payload,
headers=header,
) as response:
if response.status != 200:
"""Call the Brave Search API with API key failover.
Args:
provider_settings: Provider settings containing Brave API keys.
payload: Request payload for the Brave search endpoint.
Returns:
Normalized search results.
Raises:
ValueError: If Brave API keys are not configured.
Exception: If the request fails after all retryable keys are exhausted,
or if a non-retryable HTTP error is returned.
"""
keys = provider_settings.get("websearch_brave_key", [])
if not keys:
raise ValueError("Error: Brave API key is not configured in AstrBot.")
last_error = None
for _ in range(len(keys)):
brave_key = await _BRAVE_KEY_ROTATOR.get(provider_settings)
header = {
"Accept": "application/json",
"X-Subscription-Token": brave_key,
}
async with aiohttp.ClientSession(trust_env=True) as session:
async with session.get(
"https://api.search.brave.com/res/v1/web/search",
params=payload,
headers=header,
) as response:
if response.status == 200:
data = await response.json()
rows = data.get("web", {}).get("results", [])
return [
SearchResult(
title=item.get("title", ""),
url=item.get("url", ""),
snippet=item.get("description", ""),
)
for item in rows
]
reason = await response.text()
if response.status in _RETRYABLE_HTTP_STATUSES:
last_error = Exception(
f"Brave web search failed: {reason}, status: {response.status}",
)
continue
raise Exception(
f"Brave web search failed: {reason}, status: {response.status}",
)
data = await response.json()
rows = data.get("web", {}).get("results", [])
return [
SearchResult(
title=item.get("title", ""),
url=item.get("url", ""),
snippet=item.get("description", ""),
)
for item in rows
]
if last_error is not None:
raise last_error
raise Exception("Brave web search failed with all configured keys.")
async def _firecrawl_search(
provider_settings: dict,
payload: dict,
) -> list[SearchResult]:
firecrawl_key = await _FIRECRAWL_KEY_ROTATOR.get(provider_settings)
header = {
"Authorization": f"Bearer {firecrawl_key}",
"Content-Type": "application/json",
}
async with aiohttp.ClientSession(trust_env=True) as session:
async with session.post(
"https://api.firecrawl.dev/v2/search",
json=payload,
headers=header,
) as response:
if response.status != 200:
"""Call the Firecrawl Search API with API key failover.
Args:
provider_settings: Provider settings containing Firecrawl API keys.
payload: Request payload for the Firecrawl search endpoint.
Returns:
Normalized search results.
Raises:
ValueError: If Firecrawl API keys are not configured.
Exception: If the request fails after all retryable keys are exhausted,
or if a non-retryable HTTP error is returned.
"""
keys = provider_settings.get("websearch_firecrawl_key", [])
if not keys:
raise ValueError("Error: Firecrawl API key is not configured in AstrBot.")
last_error = None
for _ in range(len(keys)):
firecrawl_key = await _FIRECRAWL_KEY_ROTATOR.get(provider_settings)
header = {
"Authorization": f"Bearer {firecrawl_key}",
"Content-Type": "application/json",
}
async with aiohttp.ClientSession(trust_env=True) as session:
async with session.post(
"https://api.firecrawl.dev/v2/search",
json=payload,
headers=header,
) as response:
if response.status == 200:
data = await response.json()
rows = data.get("data", [])
if isinstance(rows, dict):
rows = rows.get("web", [])
return [
SearchResult(
title=item.get("title", ""),
url=item.get("url", ""),
snippet=(
item.get("description")
or item.get("snippet")
or item.get("markdown")
or ""
),
)
for item in rows
if item.get("url")
]
reason = await response.text()
if response.status in _RETRYABLE_HTTP_STATUSES:
last_error = Exception(
f"Firecrawl web search failed: {reason}, status: {response.status}",
)
continue
raise Exception(
f"Firecrawl web search failed: {reason}, status: {response.status}",
)
data = await response.json()
rows = data.get("data", [])
if isinstance(rows, dict):
rows = rows.get("web", [])
return [
SearchResult(
title=item.get("title", ""),
url=item.get("url", ""),
snippet=(
item.get("description")
or item.get("snippet")
or item.get("markdown")
or ""
),
)
for item in rows
if item.get("url")
]
if last_error is not None:
raise last_error
raise Exception("Firecrawl web search failed with all configured keys.")
async def _firecrawl_scrape(provider_settings: dict, payload: dict) -> dict:
firecrawl_key = await _FIRECRAWL_KEY_ROTATOR.get(provider_settings)
header = {
"Authorization": f"Bearer {firecrawl_key}",
"Content-Type": "application/json",
}
async with aiohttp.ClientSession(trust_env=True) as session:
async with session.post(
"https://api.firecrawl.dev/v2/scrape",
json=payload,
headers=header,
) as response:
if response.status != 200:
"""Call the Firecrawl Scrape API with API key failover.
Args:
provider_settings: Provider settings containing Firecrawl API keys.
payload: Request payload for the Firecrawl scrape endpoint.
Returns:
Raw Firecrawl scrape result data.
Raises:
ValueError: If Firecrawl API keys are not configured or no result data
is returned.
Exception: If the request fails after all retryable keys are exhausted,
or if a non-retryable HTTP error is returned.
"""
keys = provider_settings.get("websearch_firecrawl_key", [])
if not keys:
raise ValueError("Error: Firecrawl API key is not configured in AstrBot.")
last_error = None
for _ in range(len(keys)):
firecrawl_key = await _FIRECRAWL_KEY_ROTATOR.get(provider_settings)
header = {
"Authorization": f"Bearer {firecrawl_key}",
"Content-Type": "application/json",
}
async with aiohttp.ClientSession(trust_env=True) as session:
async with session.post(
"https://api.firecrawl.dev/v2/scrape",
json=payload,
headers=header,
) as response:
if response.status == 200:
data = await response.json()
result = data.get("data", {})
if not result:
raise ValueError(
"Error: Firecrawl web scraper does not return any results."
)
return result
reason = await response.text()
if response.status in _RETRYABLE_HTTP_STATUSES:
last_error = Exception(
f"Firecrawl web scraper failed: {reason}, status: {response.status}",
)
continue
raise Exception(
f"Firecrawl web scraper failed: {reason}, status: {response.status}",
)
data = await response.json()
result = data.get("data", {})
if not result:
raise ValueError(
"Error: Firecrawl web scraper does not return any results."
)
return result
if last_error is not None:
raise last_error
raise Exception("Firecrawl web scraper failed with all configured keys.")
async def _baidu_search(

227
tests/unit/test_web_search_tools.py
View File

@@ -371,6 +371,233 @@ class _FakeFirecrawlSession:
return self.response
class _CycleSession:
"""Return the next response for each post() call in key rotation tests."""
def __init__(self, responses: list):
self.responses = responses
self.cursor = 0
self.trust_env = None
self.entered = False
self.exited = False
self.calls: list[dict] = []
async def __aenter__(self):
self.entered = True
return self
async def __aexit__(self, exc_type, exc, tb):
self.exited = True
return None
def post(self, url, json, headers):
resp = self.responses[self.cursor]
self.cursor = (self.cursor + 1) % len(self.responses)
self.calls.append({"url": url, "json": json, "headers": headers})
return resp
class _TavilyResponse:
"""Fake HTTP response for Tavily API tests."""
def __init__(self, status=200, jsonData=None, textData=""):
self.status = status
self.jsonData = jsonData or {}
self.textData = textData
async def __aenter__(self):
return self
async def __aexit__(self, exc_type, exc, tb):
return None
async def json(self):
return self.jsonData
async def text(self):
return self.textData
@pytest.fixture(autouse=True)
def _resetKeyRotators():
"""Reset KeyRotator indexes to avoid state leakage between tests."""
tools._TAVILY_KEY_ROTATOR.index = 0
tools._BOCHA_KEY_ROTATOR.index = 0
tools._BRAVE_KEY_ROTATOR.index = 0
tools._FIRECRAWL_KEY_ROTATOR.index = 0
yield
tools._TAVILY_KEY_ROTATOR.index = 0
tools._BOCHA_KEY_ROTATOR.index = 0
tools._BRAVE_KEY_ROTATOR.index = 0
tools._FIRECRAWL_KEY_ROTATOR.index = 0
# ---------------------------------------------------------------------------
# Issue #8886: Tavily key rotation did not fail over to the next key.
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_tavily_search_raises_value_error_when_no_key_configured():
"""Raise ValueError when no Tavily API key is configured."""
with pytest.raises(
ValueError,
match="Error: Tavily API key is not configured in AstrBot.",
):
await tools._tavily_search({}, {"query": "test"})
@pytest.mark.asyncio
async def test_tavily_search_key_failover_on_quota_exceeded_432(
monkeypatch,
):
"""Fail over to the second key when the first key returns 432."""
session = _CycleSession(
[
_TavilyResponse(
status=432,
textData='{"detail":{"error":"quota exceeded"}}',
),
_TavilyResponse(
status=200,
jsonData={
"results": [
{"title": "AstrBot", "url": "https://example.com", "content": "OK"}
]
},
),
]
)
def fakeClientSession(*, trust_env):
session.trust_env = trust_env
return session
monkeypatch.setattr(tools.aiohttp, "ClientSession", fakeClientSession)
providerSettings = {"websearch_tavily_key": ["bad-key", "good-key"]}
results = await tools._tavily_search(providerSettings, {"query": "test"})
assert len(results) == 1
assert results[0].title == "AstrBot"
assert results[0].url == "https://example.com"
assert len(session.calls) == 2 # Both keys were attempted.
@pytest.mark.asyncio
async def test_tavily_search_key_failover_on_rate_limited_429(
monkeypatch,
):
"""Fail over to the second key when the first key returns 429."""
session = _CycleSession(
[
_TavilyResponse(
status=429,
textData='{"detail":{"error":"rate limited"}}',
),
_TavilyResponse(
status=200,
jsonData={
"results": [
{"title": "RateLimitOK", "url": "https://example2.com", "content": "OK"}
]
},
),
]
)
def fakeClientSession(*, trust_env):
session.trust_env = trust_env
return session
monkeypatch.setattr(tools.aiohttp, "ClientSession", fakeClientSession)
providerSettings = {"websearch_tavily_key": ["rate-limited-key", "good-key"]}
results = await tools._tavily_search(providerSettings, {"query": "test"})
assert len(results) == 1
assert results[0].title == "RateLimitOK"
assert len(session.calls) == 2 # Both keys were attempted.
@pytest.mark.asyncio
async def test_tavily_search_fails_when_all_keys_exhausted_8886(
monkeypatch,
):
"""Raise the last error when all keys are exhausted."""
# Both responses are retryable failures.
session = _CycleSession(
[
_TavilyResponse(
status=432,
textData='{"detail":{"error":"quota exceeded"}}',
),
_TavilyResponse(
status=429,
textData='{"detail":{"error":"rate limited"}}',
),
]
)
def fakeClientSession(*, trust_env):
session.trust_env = trust_env
return session
monkeypatch.setattr(tools.aiohttp, "ClientSession", fakeClientSession)
providerSettings = {"websearch_tavily_key": ["bad-key-1", "bad-key-2"]}
with pytest.raises(
Exception,
match="Tavily web search failed",
):
await tools._tavily_search(providerSettings, {"query": "test"})
assert len(session.calls) == 2 # Both keys were attempted.
@pytest.mark.asyncio
async def test_tavily_search_does_not_failover_on_server_error_500(
monkeypatch,
):
"""Raise immediately for non-key-related errors such as 500 responses."""
session = _CycleSession(
[
_TavilyResponse(
status=500,
textData='{"error":"internal server error"}',
),
_TavilyResponse(
status=200,
jsonData={
"results": [
{"title": "OK", "url": "https://example.com", "content": "OK"}
]
},
),
]
)
def fakeClientSession(*, trust_env):
session.trust_env = trust_env
return session
monkeypatch.setattr(tools.aiohttp, "ClientSession", fakeClientSession)
providerSettings = {"websearch_tavily_key": ["key-1", "key-2"]}
with pytest.raises(
Exception,
match="Tavily web search failed.*status: 500",
):
await tools._tavily_search(providerSettings, {"query": "test"})
# Only one key is attempted because 500 is not retryable.
assert len(session.calls) == 1
def _context_with_provider_settings(provider_settings):
config = {"provider_settings": provider_settings}
agent_context = SimpleNamespace(