From 72aea2d3f33d4d26cb48370ad00482c595b664d7 Mon Sep 17 00:00:00 2001 From: RC-CHN <67079377+RC-CHN@users.noreply.github.com> Date: Fri, 12 Sep 2025 15:03:47 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E5=85=81=E8=AE=B8=E6=B7=BB=E5=8A=A0?= =?UTF-8?q?=E5=A4=9A=E4=B8=AA=20tavily=20API=20Key=20=E8=BF=9B=E8=A1=8C?= =?UTF-8?q?=E8=BD=AE=E8=AF=A2=20(#2725)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat: 允许添加多个tavily API Key进行轮询 * perf: 并发安全的从列表中获取并轮换Tavily API密钥 * fix: 自动迁移旧版 websearch_tavily_key 为列表格式并保存 --- astrbot/core/config/default.py | 6 +++-- packages/web_searcher/main.py | 40 +++++++++++++++++++++++++--------- 2 files changed, 34 insertions(+), 12 deletions(-) diff --git a/astrbot/core/config/default.py b/astrbot/core/config/default.py index 60a2dae78..12b06bdcc 100644 --- a/astrbot/core/config/default.py +++ b/astrbot/core/config/default.py @@ -56,7 +56,7 @@ DEFAULT_CONFIG = { "wake_prefix": "", "web_search": False, "websearch_provider": "default", - "websearch_tavily_key": "", + "websearch_tavily_key": [], "web_search_link": False, "display_reasoning_text": False, "identifier": False, @@ -1938,7 +1938,9 @@ CONFIG_METADATA_3 = { }, "provider_settings.websearch_tavily_key": { "description": "Tavily API Key", - "type": "string", + "type": "list", + "items": {"type": "string"}, + "hint": "可添加多个 Key 进行轮询。", "condition": { "provider_settings.websearch_provider": "tavily", }, diff --git a/packages/web_searcher/main.py b/packages/web_searcher/main.py index 526461c53..015d8e403 100644 --- a/packages/web_searcher/main.py +++ b/packages/web_searcher/main.py @@ -26,6 +26,21 @@ class Main(star.Star): def __init__(self, context: star.Context) -> None: self.context = context + self.tavily_key_index = 0 + self.tavily_key_lock = asyncio.Lock() + + # 将 str 类型的 key 迁移至 list[str]，并保存 + cfg = self.context.get_config() + provider_settings = cfg.get("provider_settings") + if provider_settings: + tavily_key = provider_settings.get("websearch_tavily_key") + if isinstance(tavily_key, str): + logger.info("检测到旧版 websearch_tavily_key (字符串格式)，自动迁移为列表格式并保存。") + if tavily_key: + provider_settings["websearch_tavily_key"] = [tavily_key] + else: + provider_settings["websearch_tavily_key"] = [] + cfg.save_config() self.bing_search = Bing() self.sogo_search = Sogo() @@ -94,13 +109,22 @@ class Main(star.Star): return results + async def _get_tavily_key(self, cfg: AstrBotConfig) -> str: + """并发安全的从列表中获取并轮换Tavily API密钥。""" + tavily_keys = cfg.get("provider_settings", {}).get("websearch_tavily_key", []) + if not tavily_keys: + raise ValueError("错误：Tavily API密钥未在AstrBot中配置。") + + async with self.tavily_key_lock: + key = tavily_keys[self.tavily_key_index] + self.tavily_key_index = (self.tavily_key_index + 1) % len(tavily_keys) + return key + async def _web_search_tavily( self, cfg: AstrBotConfig, payload: dict ) -> list[SearchResult]: """使用 Tavily 搜索引擎进行搜索""" - tavily_key = cfg.get("provider_settings", {}).get("websearch_tavily_key", None) - if not tavily_key: - raise ValueError("Error: Tavily API key is not configured in AstrBot.") + tavily_key = await self._get_tavily_key(cfg) url = "https://api.tavily.com/search" header = { "Authorization": f"Bearer {tavily_key}", @@ -128,9 +152,7 @@ class Main(star.Star): async def _extract_tavily(self, cfg: AstrBotConfig, payload: dict) -> list[dict]: """使用 Tavily 提取网页内容""" - tavily_key = cfg.get("provider_settings", {}).get("websearch_tavily_key", None) - if not tavily_key: - raise ValueError("Error: Tavily API key is not configured in AstrBot.") + tavily_key = await self._get_tavily_key(cfg) url = "https://api.tavily.com/extract" header = { "Authorization": f"Bearer {tavily_key}", @@ -235,8 +257,7 @@ class Main(star.Star): logger.info(f"web_searcher - search_from_tavily: {query}") cfg = self.context.get_config(umo=event.unified_msg_origin) websearch_link = cfg["provider_settings"].get("web_search_link", False) - tavily_key = cfg.get("provider_settings", {}).get("websearch_tavily_key", None) - if not tavily_key: + if not cfg.get("provider_settings", {}).get("websearch_tavily_key", []): raise ValueError("Error: Tavily API key is not configured in AstrBot.") # build payload @@ -288,8 +309,7 @@ class Main(star.Star): extract_depth(string): Optional. The depth of the extraction, must be one of 'basic', 'advanced'. Default is "basic". """ cfg = self.context.get_config(umo=event.unified_msg_origin) - tavily_key = cfg.get("provider_settings", {}).get("websearch_tavily_key", None) - if not tavily_key: + if not cfg.get("provider_settings", {}).get("websearch_tavily_key", []): raise ValueError("Error: Tavily API key is not configured in AstrBot.") if not url: