diff --git a/astrbot/core/tools/web_search_tools.py b/astrbot/core/tools/web_search_tools.py index 7af425045c..0966dabcbe 100644 --- a/astrbot/core/tools/web_search_tools.py +++ b/astrbot/core/tools/web_search_tools.py @@ -77,6 +77,23 @@ async def get(self, provider_settings: dict) -> str: self.index = (self.index + 1) % len(keys) return key + async def iter_keys(self, provider_settings: dict) -> list[str]: + """Return every configured key in rotation order (current index first). + + Used for failover: callers try the keys in turn and move on to the next + one when a key is invalid, out of quota, or rate-limited. + """ + keys = provider_settings.get(self.setting_name, []) + if not keys: + raise ValueError( + f"Error: {self.provider_name} API key is not configured in AstrBot." + ) + + async with self.lock: + start = self.index + self.index = (self.index + 1) % len(keys) + return [keys[(start + i) % len(keys)] for i in range(len(keys))] + _TAVILY_KEY_ROTATOR = _KeyRotator("websearch_tavily_key", "Tavily") _BOCHA_KEY_ROTATOR = _KeyRotator("websearch_bocha_key", "BoCha") @@ -147,11 +164,10 @@ def _search_result_payload(results: list[SearchResult]) -> str: return json.dumps({"results": ret_ls}, ensure_ascii=False) -async def _tavily_search( - provider_settings: dict, +async def _tavily_request_once( + tavily_key: str, payload: dict, ) -> list[SearchResult]: - tavily_key = await _TAVILY_KEY_ROTATOR.get(provider_settings) header = { "Authorization": f"Bearer {tavily_key}", "Content-Type": "application/json", @@ -179,6 +195,22 @@ async def _tavily_search( ] +async def _tavily_search( + provider_settings: dict, + payload: dict, +) -> list[SearchResult]: + keys = await _TAVILY_KEY_ROTATOR.iter_keys(provider_settings) + last_exc: Exception | None = None + for tavily_key in keys: + try: + return await _tavily_request_once(tavily_key, payload) + except Exception as e: + # Key invalid / out of quota / rate-limited: fall through to the next key. + last_exc = e + assert last_exc is not None # iter_keys raises when no keys are configured + raise last_exc + + async def _tavily_extract(provider_settings: dict, payload: dict) -> list[dict]: tavily_key = await _TAVILY_KEY_ROTATOR.get(provider_settings) header = { diff --git a/tests/unit/test_web_search_tools.py b/tests/unit/test_web_search_tools.py index 3eb96f9302..325085cf59 100644 --- a/tests/unit/test_web_search_tools.py +++ b/tests/unit/test_web_search_tools.py @@ -513,3 +513,54 @@ def fake_client_session(*, trust_env): {"websearch_exa_key": ["exa-key"]}, {"ids": ["https://example.com"]}, ) + + +@pytest.mark.asyncio +async def test_iter_keys_returns_rotation_order(): + rotator = tools._KeyRotator("websearch_tavily_key", "Tavily") + settings = {"websearch_tavily_key": ["k1", "k2", "k3"]} + + assert await rotator.iter_keys(settings) == ["k1", "k2", "k3"] + # the starting point advances so a different key leads the next call + assert await rotator.iter_keys(settings) == ["k2", "k3", "k1"] + + +@pytest.mark.asyncio +async def test_iter_keys_raises_when_unconfigured(): + rotator = tools._KeyRotator("websearch_tavily_key", "Tavily") + with pytest.raises(ValueError): + await rotator.iter_keys({"websearch_tavily_key": []}) + + +@pytest.mark.asyncio +async def test_tavily_search_falls_over_to_next_key(monkeypatch): + tried = [] + + async def fake_request_once(key, payload): + tried.append(key) + if key == "bad-key": + raise Exception("Tavily web search failed: quota exceeded, status: 429") + return [tools.SearchResult(title="t", url="u", snippet="s")] + + monkeypatch.setattr(tools, "_tavily_request_once", fake_request_once) + monkeypatch.setattr(tools._TAVILY_KEY_ROTATOR, "index", 0) + settings = {"websearch_tavily_key": ["bad-key", "good-key"]} + + results = await tools._tavily_search(settings, {"query": "x"}) + + # the first key failed, so the search must have moved on to the second key + assert tried == ["bad-key", "good-key"] + assert results[0].title == "t" + + +@pytest.mark.asyncio +async def test_tavily_search_raises_last_error_when_all_keys_fail(monkeypatch): + async def fake_request_once(key, payload): + raise Exception(f"Tavily web search failed for {key}") + + monkeypatch.setattr(tools, "_tavily_request_once", fake_request_once) + monkeypatch.setattr(tools._TAVILY_KEY_ROTATOR, "index", 0) + settings = {"websearch_tavily_key": ["k1", "k2"]} + + with pytest.raises(Exception, match="Tavily web search failed for k2"): + await tools._tavily_search(settings, {"query": "x"})