From 8a31cf01f23a0c189457e2185790b71409db1b31 Mon Sep 17 00:00:00 2001 From: NayukiChiba Date: Fri, 3 Jul 2026 22:53:22 +0800 Subject: [PATCH] fix(core): improve web search API key failover Merge PR #8896 --- astrbot/core/tools/web_search_tools.py | 486 ++++++++++++++++++------- tests/unit/test_web_search_tools.py | 227 ++++++++++++ 2 files changed, 572 insertions(+), 141 deletions(-) diff --git a/astrbot/core/tools/web_search_tools.py b/astrbot/core/tools/web_search_tools.py index 191f42c9e..8b6adbb77 100644 --- a/astrbot/core/tools/web_search_tools.py +++ b/astrbot/core/tools/web_search_tools.py @@ -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( diff --git a/tests/unit/test_web_search_tools.py b/tests/unit/test_web_search_tools.py index 3eb96f930..b3f51acbc 100644 --- a/tests/unit/test_web_search_tools.py +++ b/tests/unit/test_web_search_tools.py @@ -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(