Skip to content

fix: fall back to next Tavily key when web search fails#9072

Open
kldsjfas wants to merge 2 commits into
AstrBotDevs:masterfrom
kldsjfas:fix/8886-tavily-key-rotation
Open

fix: fall back to next Tavily key when web search fails#9072
kldsjfas wants to merge 2 commits into
AstrBotDevs:masterfrom
kldsjfas:fix/8886-tavily-key-rotation

Conversation

@kldsjfas

@kldsjfas kldsjfas commented Jun 29, 2026

Copy link
Copy Markdown

问题

#8886:配置多个 Tavily Key 后,如果本次轮到的 Key 已失效 / 额度耗尽 / 被限流,搜索会直接报错,而不会继续尝试下一个 Key。

原因

_tavily_search 只调用一次 _TAVILY_KEY_ROTATOR.get() 取一个 Key,响应非 200 就直接 raise,没有失败转移。

改动

  • _KeyRotator 增加 ordered_keys():从当前轮询位置返回全部 Key(保留原有 round-robin 起点),供调用方在失败时依次尝试。
  • _tavily_search 改为遍历这些 Key:成功即返回;某个 Key 失败则记录日志并尝试下一个;全部失败才抛出最后一个错误。
  • 补充单元测试:第一个 Key 返回 401、第二个可用时,搜索能自动改用第二个 Key 成功(tests/unit/test_web_search_tools.py)。

说明:_tavily_extract 以及 BoCha / Brave / Firecrawl / Exa 的搜索也共用 _KeyRotator、存在同样的「失败不切换」问题。本 PR 先聚焦 issue 报的 Tavily 搜索路径;ordered_keys() 是通用的,如果希望在本 PR 里一并覆盖其余 provider,我可以补上。

Closes #8886

Summary by Sourcery

Add key rotation fallback for Tavily web search to continue trying configured keys when one fails.

Bug Fixes:

  • Ensure Tavily web search falls back to subsequent API keys when the current key is invalid or fails instead of failing immediately.

Enhancements:

  • Expose ordered key rotation from the current position via _KeyRotator.ordered_keys to support failure fallback while preserving round-robin behavior.

Tests:

  • Add unit test verifying Tavily search retries with the next key when the first key returns an error and succeeds with a subsequent key.

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.
@dosubot dosubot Bot added the size:M This PR changes 30-99 lines, ignoring generated files. label Jun 29, 2026

@sourcery-ai sourcery-ai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hey - I've left some high level feedback:

  • 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.
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.

Sourcery is free for open source - if you like our reviews please consider sharing them ✨
Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment on lines +82 to +97
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]

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

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_error
References
  1. When implementing similar functionality for different cases, refactor the logic into a shared helper function to avoid code duplication.

Comment thread astrbot/core/tools/web_search_tools.py Outdated
Comment on lines +173 to +206
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

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

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
  1. 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.
@kldsjfas

Copy link
Copy Markdown
Author

感谢 review!已按建议调整:

  • 把「逐个 key 失败转移」的逻辑抽成了 _KeyRotator.execute_with_fallback(provider_settings, fn)_tavily_search 改为调用它。这样 BoCha / Brave / Firecrawl / Exa 以及 _tavily_extract 之后都能直接复用同一套失败转移逻辑,不再重复。
  • 简化了 last_error 的处理(先置 None,循环结束后再抛)。

本 PR 仍先聚焦 issue 里报的 Tavily 搜索路径。如果希望在本 PR 内就把其余 provider 一并接入 execute_with_fallback,我可以补上,麻烦告知~

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

size:M This PR changes 30-99 lines, ignoring generated files.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Bug]Tavily 多 Key 轮询不会在失败时切换下一个 Key

1 participant