Skip to content

fix(core): 修复 Tavily 其他5种search方法 API Key 和轮询失败未自动切换的问题#8896

Open
NayukiChiba wants to merge 3 commits into
AstrBotDevs:masterfrom
NayukiChiba:fix/8886-tavily-polling
Open

fix(core): 修复 Tavily 其他5种search方法 API Key 和轮询失败未自动切换的问题#8896
NayukiChiba wants to merge 3 commits into
AstrBotDevs:masterfrom
NayukiChiba:fix/8886-tavily-polling

Conversation

@NayukiChiba

@NayukiChiba NayukiChiba commented Jun 19, 2026

Copy link
Copy Markdown
Contributor

Modifications / 改动点

  • 新增可重试状态码集合 _RETRYABLE_HTTP_STATUSES,支持 401/403/429/432
  • 重构 _tavily_search 等搜索函数,遍历所有 Key 并在可重试错误时自动切换
  • 更新 _KeyRotator 文档注释,明确多协程安全的 Round-Robin 实现
  • 补充 Tavily 多 Key 失败切换的单元测试

Closes #8886

  • This is NOT a breaking change. / 这不是一个破坏性变更。

Screenshots or Test Results / 运行截图或测试结果

image

Checklist / 检查清单

  • 😊 If there are new features added in the PR, I have discussed it with the authors through issues/emails, etc.
    / 如果 PR 中有新加入的功能,已经通过 Issue / 邮件等方式和作者讨论过。

  • 👀 My changes have been well-tested, and "Verification Steps" and "Screenshots" have been provided above.
    / 我的更改经过了良好的测试,并已在上方提供了“验证步骤”和“运行截图”

  • 🤓 I have ensured that no new dependencies are introduced, OR if new dependencies are introduced, they have been added to the appropriate locations in requirements.txt and pyproject.toml.
    / 我确保没有引入新依赖库,或者引入了新依赖库的同时将其添加到 requirements.txtpyproject.toml 文件相应位置。

  • 😮 My changes do not introduce malicious code.
    / 我的更改没有引入恶意代码。

Summary by Sourcery

Improve web search providers to support multi-key failover with clearer key rotation behavior and retryable error handling.

Bug Fixes:

  • Fix Tavily search so that it correctly rotates to the next API key on specific failure statuses instead of getting stuck on a bad key.

Enhancements:

  • Introduce a shared set of retryable HTTP status codes and apply it across Tavily, BoCha, Brave, and Firecrawl search/scrape helpers to drive key failover.
  • Document the KeyRotator as a coroutine-safe round-robin API key rotator and clarify its intended usage in failover loops.

Tests:

  • Add unit tests and test helpers to validate Tavily multi-key failover behavior, including retryable and non-retryable error scenarios.

- 新增可重试状态码集合 _RETRYABLE_HTTP_STATUSES,支持 401/403/429/432
- 重构 _tavily_search 等搜索函数,遍历所有 Key 并在可重试错误时自动切换
- 更新 _KeyRotator 文档注释,明确多协程安全的 Round-Robin 实现
- 补充 Tavily 多 Key 失败切换的单元测试

Closes AstrBotDevs#8886
@dosubot dosubot Bot added size:L This PR changes 100-499 lines, ignoring generated files. area:core The bug / feature is about astrbot's core, backend labels Jun 19, 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 found 2 issues, and left some high level feedback:

  • The failover loop (checking keys, handling _RETRYABLE_HTTP_STATUSES, tracking last_error) is duplicated across all provider functions; consider extracting a shared helper (e.g., _request_with_key_failover(...)) to centralize this logic and reduce future maintenance risk.
  • The _resetKeyRotators fixture is defined with autouse=False but none of the new tests use it (they manually reset index); either wire the fixture into the relevant tests or remove it to avoid dead code.
  • You are reusing the same _RETRYABLE_HTTP_STATUSES set for Tavily, BoCha, Brave, and Firecrawl, including Tavily-specific code 432; if 432 is not meaningful for the other providers, consider per-provider retryable status sets to make behavior clearer and easier to adjust.
Prompt for AI Agents
Please address the comments from this code review:

## Overall Comments
- The failover loop (checking keys, handling `_RETRYABLE_HTTP_STATUSES`, tracking `last_error`) is duplicated across all provider functions; consider extracting a shared helper (e.g., `_request_with_key_failover(...)`) to centralize this logic and reduce future maintenance risk.
- The `_resetKeyRotators` fixture is defined with `autouse=False` but none of the new tests use it (they manually reset `index`); either wire the fixture into the relevant tests or remove it to avoid dead code.
- You are reusing the same `_RETRYABLE_HTTP_STATUSES` set for Tavily, BoCha, Brave, and Firecrawl, including Tavily-specific code 432; if 432 is not meaningful for the other providers, consider per-provider retryable status sets to make behavior clearer and easier to adjust.

## Individual Comments

### Comment 1
<location path="tests/unit/test_web_search_tools.py" line_range="421-427" />
<code_context>
+        return self.textData
+
+
+@pytest.fixture(autouse=False)
+def _resetKeyRotators():
+    """重置所有 KeyRotator 的索引,避免测试间状态干扰"""
+    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
</code_context>
<issue_to_address>
**suggestion (testing):** Consider actually using `_resetKeyRotators` in the new async tests to avoid future state leakage between tests.

Currently each Tavily test manually resets only `_TAVILY_KEY_ROTATOR.index`, while the `_resetKeyRotators` fixture is unused and would also cover BoCha/Brave/Firecrawl and future state on `_KeyRotator`. Please either wire the fixture into these async tests (e.g., add `_resetKeyRotators` as a parameter and drop the manual index resets) or remove the unused fixture if you prefer explicit per-test resets.

Suggested implementation:

```python
@pytest.fixture(autouse=True)
def _resetKeyRotators():
    """重置所有 KeyRotator 的索引,避免测试间状态干扰"""
    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

```

To fully implement the review suggestion and avoid future state leakage:

1. In all async Tavily tests (and any tests that currently manually reset `tools._TAVILY_KEY_ROTATOR.index` or other `*_KEY_ROTATOR.index` values), remove those manual resets. The fixture is now `autouse=True`, so every test in this module will start and end with all rotator indices reset.
2. If there are any tests that rely on non-zero index values being preserved across tests, they will now be isolated; if that behavior is desired, those tests should explicitly set the index within the test body instead of relying on cross-test state.
</issue_to_address>

### Comment 2
<location path="tests/unit/test_web_search_tools.py" line_range="440-441" />
<code_context>
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_tavily_search_key_failover_on_quota_exceeded_432(
+    monkeypatch,
+):
</code_context>
<issue_to_address>
**suggestion (testing):** Consider adding a test for the "no Tavily key configured" ValueError branch.

Since `_tavily_search` now raises `ValueError("Error: Tavily API key is not configured in AstrBot.")` when `websearch_tavily_key` is missing or empty, please add an async test that calls it with an empty/missing key list and asserts that this `ValueError` (and message) is raised, to cover that configuration error path.

Suggested implementation:

```python
# ---------------------------------------------------------------------------
# Issue #8886: Tavily 多 Key 轮询不会在失败时切换到下一个 Key
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_tavily_search_raises_when_no_key_configured(monkeypatch):
    """当未配置 Tavily API Key 时,应抛出明确的配置错误"""
    tools._TAVILY_KEY_ROTATOR.index = 0

    # 模拟未配置 / 空的 key 列表
    # 注意:如果实际配置路径不同,请在此处调整 monkeypatch 的目标
    monkeypatch.setattr(tools, "websearch_tavily_key", [], raising=False)

    with pytest.raises(
        ValueError,
        match="Error: Tavily API key is not configured in AstrBot.",
    ):
        await tools._tavily_search(
            query="test",
            session=_CycleSession([]),
        )


@pytest.mark.asyncio
async def test_tavily_search_key_failover_on_quota_exceeded_432(
    monkeypatch,
):

```

The new test assumes:
1. The Tavily API key configuration is exposed as `tools.websearch_tavily_key` (a list of keys), and `_tavily_search` reads from it. If the actual configuration lives elsewhere (e.g. `tools.config.websearch_tavily_key` or environment variables), update the `monkeypatch.setattr(...)` target accordingly to ensure `_tavily_search` sees an empty/missing key list.
2. `_tavily_search` has a signature compatible with `await tools._tavily_search(query="test", session=_CycleSession([]))`. If the real signature differs (e.g. extra required parameters or different names), adjust the call to match the actual function definition while still passing an empty-key configuration via monkeypatch.
</issue_to_address>

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.

Comment thread tests/unit/test_web_search_tools.py Outdated
Comment on lines +421 to +427
@pytest.fixture(autouse=False)
def _resetKeyRotators():
"""重置所有 KeyRotator 的索引,避免测试间状态干扰"""
tools._TAVILY_KEY_ROTATOR.index = 0
tools._BOCHA_KEY_ROTATOR.index = 0
tools._BRAVE_KEY_ROTATOR.index = 0
tools._FIRECRAWL_KEY_ROTATOR.index = 0

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.

suggestion (testing): Consider actually using _resetKeyRotators in the new async tests to avoid future state leakage between tests.

Currently each Tavily test manually resets only _TAVILY_KEY_ROTATOR.index, while the _resetKeyRotators fixture is unused and would also cover BoCha/Brave/Firecrawl and future state on _KeyRotator. Please either wire the fixture into these async tests (e.g., add _resetKeyRotators as a parameter and drop the manual index resets) or remove the unused fixture if you prefer explicit per-test resets.

Suggested implementation:

@pytest.fixture(autouse=True)
def _resetKeyRotators():
    """重置所有 KeyRotator 的索引,避免测试间状态干扰"""
    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

To fully implement the review suggestion and avoid future state leakage:

  1. In all async Tavily tests (and any tests that currently manually reset tools._TAVILY_KEY_ROTATOR.index or other *_KEY_ROTATOR.index values), remove those manual resets. The fixture is now autouse=True, so every test in this module will start and end with all rotator indices reset.
  2. If there are any tests that rely on non-zero index values being preserved across tests, they will now be isolated; if that behavior is desired, those tests should explicitly set the index within the test body instead of relying on cross-test state.

Comment on lines +440 to +441
@pytest.mark.asyncio
async def test_tavily_search_key_failover_on_quota_exceeded_432(

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.

suggestion (testing): Consider adding a test for the "no Tavily key configured" ValueError branch.

Since _tavily_search now raises ValueError("Error: Tavily API key is not configured in AstrBot.") when websearch_tavily_key is missing or empty, please add an async test that calls it with an empty/missing key list and asserts that this ValueError (and message) is raised, to cover that configuration error path.

Suggested implementation:

# ---------------------------------------------------------------------------
# Issue #8886: Tavily 多 Key 轮询不会在失败时切换到下一个 Key
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_tavily_search_raises_when_no_key_configured(monkeypatch):
    """当未配置 Tavily API Key 时,应抛出明确的配置错误"""
    tools._TAVILY_KEY_ROTATOR.index = 0

    # 模拟未配置 / 空的 key 列表
    # 注意:如果实际配置路径不同,请在此处调整 monkeypatch 的目标
    monkeypatch.setattr(tools, "websearch_tavily_key", [], raising=False)

    with pytest.raises(
        ValueError,
        match="Error: Tavily API key is not configured in AstrBot.",
    ):
        await tools._tavily_search(
            query="test",
            session=_CycleSession([]),
        )


@pytest.mark.asyncio
async def test_tavily_search_key_failover_on_quota_exceeded_432(
    monkeypatch,
):

The new test assumes:

  1. The Tavily API key configuration is exposed as tools.websearch_tavily_key (a list of keys), and _tavily_search reads from it. If the actual configuration lives elsewhere (e.g. tools.config.websearch_tavily_key or environment variables), update the monkeypatch.setattr(...) target accordingly to ensure _tavily_search sees an empty/missing key list.
  2. _tavily_search has a signature compatible with await tools._tavily_search(query="test", session=_CycleSession([])). If the real signature differs (e.g. extra required parameters or different names), adjust the call to match the actual function definition while still passing an empty-key configuration via monkeypatch.

@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 implements automatic API key rotation and failover mechanisms for web search providers (Tavily, BoCha, Brave, Firecrawl) when encountering retryable HTTP status codes (401, 403, 429, 432), accompanied by unit tests to verify this behavior. The review feedback points out a potential IndexError in _KeyRotator.get if the list of keys is dynamically reduced, and provides a code suggestion to safely handle index wrapping.

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 68 to 72
async def get(self, provider_settings: dict) -> str:
"""获取当前轮到的 Key 并推进索引(取模循环)"""
keys = provider_settings.get(self.setting_name, [])
if not keys:
raise ValueError(

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.

high

潜在的 IndexError 越界风险

当用户在 AstrBot 运行期间动态修改并减少 API Key 列表时,self.index 可能会大于当前新 keys 列表的长度,从而在执行 key = keys[self.index] 时抛出 IndexError: list index out of range 异常。

建议在获取 Key 之前,先对 self.index 进行安全取模或重置,以确保其始终在有效范围内。

Suggested change
async def get(self, provider_settings: dict) -> str:
"""获取当前轮到的 Key 并推进索引(取模循环)"""
keys = provider_settings.get(self.setting_name, [])
if not keys:
raise ValueError(
async def get(self, provider_settings: dict) -> str:
"""获取当前轮到的 Key 并推进索引(取模循环)"""
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:
# 确保 index 始终在当前 keys 长度范围内,防止动态修改配置导致越界
self.index = self.index % len(keys)
key = keys[self.index]
self.index = (self.index + 1) % len(keys)
return key

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

?好变态的要求,加一下吧

- 修补 _KeyRotator.get_key() 运行时热更新减少 key 列表后可能的索引越界
- 将 _resetKeyRotators fixture 设为 autouse,确保每个测试自动重置轮询索引
- 新增未配置 Tavily API Key 时抛出 ValueError 的测试用例
- 移除多个测试函数中冗余的 tools._TAVILY_KEY_ROTATOR.index = 0 手动重置
- 在 _KeyRotator.get_key() 中增加 index 取模校验,防止越界
- 避免运行时通过配置热更新减少 API key 数量导致的崩溃
- 提升搜索工具的稳定性和容错能力
@NayukiChiba

Copy link
Copy Markdown
Contributor Author

我感觉可以把这几个search抽象一下,后续如果增添更多search方法的话再抽象一下吧

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

Labels

area:core The bug / feature is about astrbot's core, backend size:L This PR changes 100-499 lines, ignoring generated files.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

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

1 participant