-
-
Notifications
You must be signed in to change notification settings - Fork 1.8k
Added New web search tool: Brave Search #6847
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: master
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -24,6 +24,7 @@ class Main(star.Star): | |
| "web_search_tavily", | ||
| "tavily_extract_web_page", | ||
| "web_search_bocha", | ||
| "web_search_brave", | ||
| ] | ||
|
|
||
| def __init__(self, context: star.Context) -> None: | ||
|
|
@@ -33,6 +34,8 @@ def __init__(self, context: star.Context) -> None: | |
|
|
||
| self.bocha_key_index = 0 | ||
| self.bocha_key_lock = asyncio.Lock() | ||
| self.brave_key_index = 0 | ||
| self.brave_key_lock = asyncio.Lock() | ||
|
|
||
| # 将 str 类型的 key 迁移至 list[str],并保存 | ||
| cfg = self.context.get_config() | ||
|
|
@@ -57,6 +60,14 @@ def __init__(self, context: star.Context) -> None: | |
| provider_settings["websearch_bocha_key"] = [] | ||
| cfg.save_config() | ||
|
|
||
| brave_key = provider_settings.get("websearch_brave_key") | ||
| if isinstance(brave_key, str): | ||
| if brave_key: | ||
| provider_settings["websearch_brave_key"] = [brave_key] | ||
| else: | ||
| provider_settings["websearch_brave_key"] = [] | ||
| cfg.save_config() | ||
|
|
||
| self.bing_search = Bing() | ||
| self.sogo_search = Sogo() | ||
| self.baidu_initialized = False | ||
|
|
@@ -430,6 +441,50 @@ async def _web_search_bocha( | |
| results.append(result) | ||
| return results | ||
|
|
||
| async def _get_brave_key(self, cfg: AstrBotConfig) -> str: | ||
| """并发安全的从列表中获取并轮换 Brave API 密钥。""" | ||
| brave_keys = cfg.get("provider_settings", {}).get("websearch_brave_key", []) | ||
|
|
||
| async with self.brave_key_lock: | ||
| key = brave_keys[self.brave_key_index] | ||
| self.brave_key_index = (self.brave_key_index + 1) % len(brave_keys) | ||
| return key | ||
|
|
||
| async def _web_search_brave( | ||
| self, | ||
| cfg: AstrBotConfig, | ||
| payload: dict, | ||
| ) -> list[SearchResult]: | ||
| """使用 Brave 搜索引擎进行搜索""" | ||
| brave_key = await self._get_brave_key(cfg) | ||
| url = "https://api.search.brave.com/res/v1/web/search" | ||
| header = { | ||
| "Accept": "application/json", | ||
| "X-Subscription-Token": brave_key, | ||
| } | ||
| async with aiohttp.ClientSession(trust_env=True) as session: | ||
| async with session.get( | ||
| url, | ||
| params=payload, | ||
| headers=header, | ||
| ) as response: | ||
| if response.status != 200: | ||
| reason = await response.text() | ||
| raise Exception( | ||
| f"Brave web search failed: {reason}, status: {response.status}", | ||
| ) | ||
| data = await response.json() | ||
| rows = data.get("web", {}).get("results", []) | ||
| results = [] | ||
| for item in rows: | ||
| result = SearchResult( | ||
| title=item.get("title", ""), | ||
| url=item.get("url", ""), | ||
| snippet=item.get("description", ""), | ||
| ) | ||
| results.append(result) | ||
| return results | ||
|
|
||
| @llm_tool("web_search_bocha") | ||
| async def search_from_bocha( | ||
| self, | ||
|
|
@@ -537,6 +592,64 @@ async def search_from_bocha( | |
| ret = json.dumps({"results": ret_ls}, ensure_ascii=False) | ||
| return ret | ||
|
|
||
| @llm_tool("web_search_brave") | ||
| async def search_from_brave( | ||
| self, | ||
| event: AstrMessageEvent, | ||
| query: str, | ||
| count: int = 10, | ||
| country: str = "US", | ||
| search_lang: str = "zh-hans", | ||
| freshness: str = "", | ||
| ) -> str: | ||
| """ | ||
| A web search tool based on Brave Search API. | ||
|
|
||
| Args: | ||
| query(string): Required. Search query. | ||
| count(number): Optional. Number of results to return. Range: 1–20. Default is 10. | ||
| country(string): Optional. Country code for region-specific results (e.g., "US", "CN"). | ||
| search_lang(string): Optional. Brave language code (e.g., "zh-hans", "en", "en-gb"). | ||
| freshness(string): Optional. "day", "week", "month", "year". | ||
| """ | ||
| logger.info(f"web_searcher - search_from_brave: {query}") | ||
| cfg = self.context.get_config(umo=event.unified_msg_origin) | ||
| if not cfg.get("provider_settings", {}).get("websearch_brave_key", []): | ||
| raise ValueError("Error: Brave API key is not configured in AstrBot.") | ||
|
Comment on lines
+617
to
+618
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This check for the Brave API key is redundant. The Additionally, the error message here is in English, while the one in |
||
|
|
||
| if count < 1: | ||
| count = 1 | ||
| if count > 20: | ||
| count = 20 | ||
|
|
||
| payload = { | ||
| "q": query, | ||
| "count": count, | ||
| "country": country, | ||
| "search_lang": search_lang, | ||
| } | ||
| if freshness in ["day", "week", "month", "year"]: | ||
| payload["freshness"] = freshness | ||
|
|
||
| results = await self._web_search_brave(cfg, payload) | ||
| if not results: | ||
| return "Error: Brave web searcher does not return any results." | ||
|
|
||
| ret_ls = [] | ||
| ref_uuid = str(uuid.uuid4())[:4] | ||
| for idx, result in enumerate(results, 1): | ||
| index = f"{ref_uuid}.{idx}" | ||
| ret_ls.append( | ||
| { | ||
| "title": f"{result.title}", | ||
| "url": f"{result.url}", | ||
| "snippet": f"{result.snippet}", | ||
| "index": index, | ||
| } | ||
| ) | ||
| ret = json.dumps({"results": ret_ls}, ensure_ascii=False) | ||
| return ret | ||
|
|
||
| @filter.on_llm_request(priority=-10000) | ||
| async def edit_web_search_tools( | ||
| self, | ||
|
|
@@ -575,6 +688,7 @@ async def edit_web_search_tools( | |
| tool_set.remove_tool("tavily_extract_web_page") | ||
| tool_set.remove_tool("AIsearch") | ||
| tool_set.remove_tool("web_search_bocha") | ||
| tool_set.remove_tool("web_search_brave") | ||
| elif provider == "tavily": | ||
| web_search_tavily = func_tool_mgr.get_func("web_search_tavily") | ||
| tavily_extract_web_page = func_tool_mgr.get_func("tavily_extract_web_page") | ||
|
|
@@ -586,6 +700,7 @@ async def edit_web_search_tools( | |
| tool_set.remove_tool("fetch_url") | ||
| tool_set.remove_tool("AIsearch") | ||
| tool_set.remove_tool("web_search_bocha") | ||
| tool_set.remove_tool("web_search_brave") | ||
| elif provider == "baidu_ai_search": | ||
| try: | ||
| await self.ensure_baidu_ai_search_mcp(event.unified_msg_origin) | ||
|
|
@@ -597,6 +712,7 @@ async def edit_web_search_tools( | |
| tool_set.remove_tool("web_search_tavily") | ||
| tool_set.remove_tool("tavily_extract_web_page") | ||
| tool_set.remove_tool("web_search_bocha") | ||
| tool_set.remove_tool("web_search_brave") | ||
| except Exception as e: | ||
| logger.error(f"Cannot Initialize Baidu AI Search MCP Server: {e}") | ||
| elif provider == "bocha": | ||
|
|
@@ -608,3 +724,14 @@ async def edit_web_search_tools( | |
| tool_set.remove_tool("AIsearch") | ||
| tool_set.remove_tool("web_search_tavily") | ||
| tool_set.remove_tool("tavily_extract_web_page") | ||
| tool_set.remove_tool("web_search_brave") | ||
| elif provider == "brave": | ||
| web_search_brave = func_tool_mgr.get_func("web_search_brave") | ||
| if web_search_brave and web_search_brave.active: | ||
| tool_set.add_tool(web_search_brave) | ||
| tool_set.remove_tool("web_search") | ||
| tool_set.remove_tool("fetch_url") | ||
| tool_set.remove_tool("AIsearch") | ||
| tool_set.remove_tool("web_search_tavily") | ||
| tool_set.remove_tool("tavily_extract_web_page") | ||
| tool_set.remove_tool("web_search_bocha") | ||
| Original file line number | Diff line number | Diff line change | ||||
|---|---|---|---|---|---|---|
|
|
@@ -60,6 +60,8 @@ AstrBot 默认配置如下: | |||||
| "web_search": False, | ||||||
| "websearch_provider": "default", | ||||||
| "websearch_tavily_key": [], | ||||||
| "websearch_bocha_key": [], | ||||||
| "websearch_brave_key": [], | ||||||
| "web_search_link": False, | ||||||
| "display_reasoning_text": False, | ||||||
| "identifier": False, | ||||||
|
|
@@ -286,16 +288,27 @@ ID 白名单。填写后,将只处理所填写的 ID 发来的消息事件。 | |||||
|
|
||||||
| #### `provider_settings.websearch_provider` | ||||||
|
|
||||||
| 网页搜索提供商类型。默认为 `default`。目前支持 `default` 和 `tavily`。 | ||||||
| 网页搜索提供商类型。默认为 `default`。目前支持 `default`、`tavily`、`bocha`、`baidu_ai_search`、`brave`。 | ||||||
|
|
||||||
| - `default`:能访问 Google 时效果最佳。如果 Google 访问失败,程序会依次访问 Bing, Sogo 搜索引擎。 | ||||||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. issue (typo): Possible typo: "Sogo" should likely be "Sogou" (or its Chinese name). This likely refers to the search engine “Sogou” (搜狗). Please change “Sogo” to “Sogou” or its Chinese brand name for correctness.
Suggested change
|
||||||
|
|
||||||
| - `tavily`:使用 Tavily 搜索引擎。 | ||||||
| - `bocha`:使用 BoCha 搜索引擎。 | ||||||
| - `baidu_ai_search`:使用百度 AI Search(MCP)。 | ||||||
| - `brave`:使用 Brave Search API。 | ||||||
|
|
||||||
| #### `provider_settings.websearch_tavily_key` | ||||||
|
|
||||||
| Tavily 搜索引擎的 API Key 列表。使用 `tavily` 作为网页搜索提供商时需要填写。 | ||||||
|
|
||||||
| #### `provider_settings.websearch_bocha_key` | ||||||
|
|
||||||
| BoCha 搜索引擎的 API Key 列表。使用 `bocha` 作为网页搜索提供商时需要填写。 | ||||||
|
|
||||||
| #### `provider_settings.websearch_brave_key` | ||||||
|
|
||||||
| Brave 搜索引擎的 API Key 列表。使用 `brave` 作为网页搜索提供商时需要填写。 | ||||||
|
|
||||||
| #### `provider_settings.web_search_link` | ||||||
|
|
||||||
| 是否在回复中提示模型附上搜索结果的链接。默认为 `false`。 | ||||||
|
|
||||||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
suggestion: Avoid duplicated configuration checks and inconsistent error messages for missing Brave key.
This case is handled both here and in
_get_brave_key, with different error messages. Please choose a single place to perform the validation and define the message—either centralize it in_get_brave_keyand letsearch_from_braverely on that, or have_get_brave_keyassume a non-empty list and keep the check only here. A single validation path will keep behavior consistent and easier to maintain.Suggested implementation:
These changes assume that
_get_brave_key:If
_get_brave_keydoes not yet validate the key or has an undesired error message, you should:_get_brave_keyto whatever you want to expose to callers likesearch_from_brave.