fix(core): 修复 Tavily 其他5种search方法 API Key 和轮询失败未自动切换的问题#8896
fix(core): 修复 Tavily 其他5种search方法 API Key 和轮询失败未自动切换的问题#8896NayukiChiba wants to merge 3 commits into
Conversation
- 新增可重试状态码集合 _RETRYABLE_HTTP_STATUSES,支持 401/403/429/432 - 重构 _tavily_search 等搜索函数,遍历所有 Key 并在可重试错误时自动切换 - 更新 _KeyRotator 文档注释,明确多协程安全的 Round-Robin 实现 - 补充 Tavily 多 Key 失败切换的单元测试 Closes AstrBotDevs#8886
There was a problem hiding this comment.
Hey - I've found 2 issues, and left some high level feedback:
- The failover loop (checking keys, handling
_RETRYABLE_HTTP_STATUSES, trackinglast_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
_resetKeyRotatorsfixture is defined withautouse=Falsebut none of the new tests use it (they manually resetindex); either wire the fixture into the relevant tests or remove it to avoid dead code. - You are reusing the same
_RETRYABLE_HTTP_STATUSESset 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>Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
| @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 |
There was a problem hiding this comment.
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 = 0To fully implement the review suggestion and avoid future state leakage:
- In all async Tavily tests (and any tests that currently manually reset
tools._TAVILY_KEY_ROTATOR.indexor other*_KEY_ROTATOR.indexvalues), remove those manual resets. The fixture is nowautouse=True, so every test in this module will start and end with all rotator indices reset. - 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.
| @pytest.mark.asyncio | ||
| async def test_tavily_search_key_failover_on_quota_exceeded_432( |
There was a problem hiding this comment.
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:
- The Tavily API key configuration is exposed as
tools.websearch_tavily_key(a list of keys), and_tavily_searchreads from it. If the actual configuration lives elsewhere (e.g.tools.config.websearch_tavily_keyor environment variables), update themonkeypatch.setattr(...)target accordingly to ensure_tavily_searchsees an empty/missing key list. _tavily_searchhas a signature compatible withawait 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.
There was a problem hiding this comment.
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.
| async def get(self, provider_settings: dict) -> str: | ||
| """获取当前轮到的 Key 并推进索引(取模循环)""" | ||
| keys = provider_settings.get(self.setting_name, []) | ||
| if not keys: | ||
| raise ValueError( |
There was a problem hiding this comment.
潜在的 IndexError 越界风险
当用户在 AstrBot 运行期间动态修改并减少 API Key 列表时,self.index 可能会大于当前新 keys 列表的长度,从而在执行 key = keys[self.index] 时抛出 IndexError: list index out of range 异常。
建议在获取 Key 之前,先对 self.index 进行安全取模或重置,以确保其始终在有效范围内。
| 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 |
- 修补 _KeyRotator.get_key() 运行时热更新减少 key 列表后可能的索引越界 - 将 _resetKeyRotators fixture 设为 autouse,确保每个测试自动重置轮询索引 - 新增未配置 Tavily API Key 时抛出 ValueError 的测试用例 - 移除多个测试函数中冗余的 tools._TAVILY_KEY_ROTATOR.index = 0 手动重置
- 在 _KeyRotator.get_key() 中增加 index 取模校验,防止越界 - 避免运行时通过配置热更新减少 API key 数量导致的崩溃 - 提升搜索工具的稳定性和容错能力
|
我感觉可以把这几个search抽象一下,后续如果增添更多search方法的话再抽象一下吧 |
Modifications / 改动点
Closes #8886
Screenshots or Test Results / 运行截图或测试结果
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.txtandpyproject.toml./ 我确保没有引入新依赖库,或者引入了新依赖库的同时将其添加到
requirements.txt和pyproject.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:
Enhancements:
Tests: