fix: fall back to next Tavily key when web search fails#9072
Conversation
When multiple Tavily keys are configured, a failed key (invalid / out of quota / rate limited) made the search error out instead of trying the next key. Add `_KeyRotator.ordered_keys()` which returns all keys from the current round-robin position, and have the Tavily search loop over them, moving to the next key on failure and only raising once every key fails.
There was a problem hiding this comment.
Hey - I've left some high level feedback:
- In
_tavily_search, you create a newaiohttp.ClientSessionfor each key; consider reusing a single session across all key attempts in the loop to avoid unnecessary connection setup overhead. last_errorin_tavily_searchis initialized to a generic exception but is always overwritten or bypassed byordered_keysraising; you can simplify by initializing it toNoneand only raising after the loop if an actual error was captured.
Prompt for AI Agents
Please address the comments from this code review:
## Overall Comments
- In `_tavily_search`, you create a new `aiohttp.ClientSession` for each key; consider reusing a single session across all key attempts in the loop to avoid unnecessary connection setup overhead.
- `last_error` in `_tavily_search` is initialized to a generic exception but is always overwritten or bypassed by `ordered_keys` raising; you can simplify by initializing it to `None` and only raising after the loop if an actual error was captured.Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
There was a problem hiding this comment.
Code Review
This pull request introduces a key rotation fallback mechanism for Tavily web searches, allowing the system to try alternative API keys sequentially if one fails. It adds an ordered_keys method to the _KeyRotator class and updates _tavily_search to loop through these keys, along with a corresponding unit test. The feedback suggests refactoring this fallback logic into a reusable helper method (execute_with_fallback) within _KeyRotator to avoid code duplication and simplify the provider-specific search functions.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| async def ordered_keys(self, provider_settings: dict) -> list[str]: | ||
| """All configured keys, ordered from the current rotation position. | ||
|
|
||
| Lets a caller fall through to the next key when one fails (invalid, | ||
| out of quota, rate limited) instead of giving up on the first error, | ||
| while keeping the round-robin starting point consistent across calls. | ||
| """ | ||
| 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 % len(keys) | ||
| self.index = (self.index + 1) % len(keys) | ||
| return keys[start:] + keys[:start] |
There was a problem hiding this comment.
To avoid code duplication and make the fallback logic reusable across all search providers (such as BoCha, Brave, Firecrawl, Exa, etc.), we can encapsulate the fallback execution logic directly inside _KeyRotator as a helper method. This keeps the provider-specific search functions clean and focused only on their request/response handling.
async def ordered_keys(self, provider_settings: dict) -> list[str]:
"""All configured keys, ordered from the current rotation position.
Lets a caller fall through to the next key when one fails (invalid,
out of quota, rate limited) instead of giving up on the first error,
while keeping the round-robin starting point consistent across calls.
"""
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 % len(keys)
self.index = (self.index + 1) % len(keys)
return keys[start:] + keys[:start]
async def execute_with_fallback(self, provider_settings: dict, func):
"""Execute a function with fallback to the next keys if it fails."""
last_error = Exception(
f"Error: {self.provider_name} API key is not configured in AstrBot."
)
for key in await self.ordered_keys(provider_settings):
try:
return await func(key)
except Exception as e:
last_error = e
logger.warning(f"{self.provider_name} key failed, trying the next one: {e}")
raise last_errorReferences
- When implementing similar functionality for different cases, refactor the logic into a shared helper function to avoid code duplication.
| last_error: Exception = Exception( | ||
| "Error: Tavily API key is not configured in AstrBot." | ||
| ) | ||
| for tavily_key in await _TAVILY_KEY_ROTATOR.ordered_keys(provider_settings): | ||
| header = { | ||
| "Authorization": f"Bearer {tavily_key}", | ||
| "Content-Type": "application/json", | ||
| } | ||
| try: | ||
| 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: | ||
| reason = await response.text() | ||
| 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", []) | ||
| ] | ||
| except Exception as e: | ||
| last_error = e | ||
| logger.warning(f"Tavily key failed, trying the next one: {e}") | ||
| raise last_error |
There was a problem hiding this comment.
Using the new execute_with_fallback helper method on _KeyRotator, we can simplify _tavily_search significantly. This pattern can also be easily applied to other search providers and _tavily_extract to resolve the same fallback issue globally.
async def _search(tavily_key: str) -> list[SearchResult]:
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:
reason = await response.text()
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", [])
]
return await _TAVILY_KEY_ROTATOR.execute_with_fallback(provider_settings, _search)References
- When implementing similar functionality for different cases, refactor the logic into a shared helper function to avoid code duplication.
Address review feedback: move the per-key failover loop into a reusable _KeyRotator.execute_with_fallback so other web-search providers can adopt the same behaviour, and simplify the error handling.
|
感谢 review!已按建议调整:
本 PR 仍先聚焦 issue 里报的 Tavily 搜索路径。如果希望在本 PR 内就把其余 provider 一并接入 |
问题
#8886:配置多个 Tavily Key 后,如果本次轮到的 Key 已失效 / 额度耗尽 / 被限流,搜索会直接报错,而不会继续尝试下一个 Key。
原因
_tavily_search只调用一次_TAVILY_KEY_ROTATOR.get()取一个 Key,响应非 200 就直接raise,没有失败转移。改动
_KeyRotator增加ordered_keys():从当前轮询位置返回全部 Key(保留原有 round-robin 起点),供调用方在失败时依次尝试。_tavily_search改为遍历这些 Key:成功即返回;某个 Key 失败则记录日志并尝试下一个;全部失败才抛出最后一个错误。tests/unit/test_web_search_tools.py)。Closes #8886
Summary by Sourcery
Add key rotation fallback for Tavily web search to continue trying configured keys when one fails.
Bug Fixes:
Enhancements:
Tests: