From 788ceb97212111db29ea263c0f87891b66bc663f Mon Sep 17 00:00:00 2001 From: piexian Date: Thu, 20 Nov 2025 08:05:42 +0800 Subject: [PATCH 1/4] =?UTF-8?q?=E6=B7=BB=E5=8A=A0=E9=98=BF=E9=87=8C?= =?UTF-8?q?=E7=99=BE=E7=82=BC=E9=87=8D=E6=8E=92=E5=BA=8F=E6=A8=A1=E5=9E=8B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- astrbot/core/config/default.py | 29 +++ astrbot/core/provider/manager.py | 4 + .../provider/sources/bailian_rerank_source.py | 171 ++++++++++++++++++ 3 files changed, 204 insertions(+) create mode 100644 astrbot/core/provider/sources/bailian_rerank_source.py diff --git a/astrbot/core/config/default.py b/astrbot/core/config/default.py index 528611d2d..da0fefbf2 100644 --- a/astrbot/core/config/default.py +++ b/astrbot/core/config/default.py @@ -1308,6 +1308,20 @@ CONFIG_METADATA_2 = { "timeout": 20, "launch_model_if_not_running": False, }, + "阿里云百炼重排序": { + "id": "bailian_rerank", + "type": "bailian_rerank", + "provider": "bailian", + "provider_type": "rerank", + "enable": True, + "rerank_api_key": "", + "rerank_api_base": "https://dashscope.aliyuncs.com/api/v1/services/rerank/text-rerank/text-rerank", + "rerank_model": "qwen3-rerank", + "timeout": 30, + "top_n": 3, + "return_documents": False, + "instruct": "", + }, "Xinference STT": { "id": "xinference_stt", "type": "xinference_stt", @@ -1342,6 +1356,21 @@ CONFIG_METADATA_2 = { "description": "重排序模型名称", "type": "string", }, + "top_n": { + "description": "返回排序后的top_n个文档", + "type": "int", + "hint": "默认返回全部文档。如果指定的值大于文档总数，将返回全部文档。", + }, + "return_documents": { + "description": "是否在排序结果中返回文档原文", + "type": "bool", + "hint": "默认值false，以减少网络传输开销。", + }, + "instruct": { + "description": "自定义排序任务类型说明", + "type": "string", + "hint": "仅在使用 qwen3-rerank 模型时生效。建议使用英文撰写。", + }, "launch_model_if_not_running": { "description": "模型未运行时自动启动", "type": "bool", diff --git a/astrbot/core/provider/manager.py b/astrbot/core/provider/manager.py index 320c98d4e..ec2550415 100644 --- a/astrbot/core/provider/manager.py +++ b/astrbot/core/provider/manager.py @@ -331,6 +331,10 @@ class ProviderManager: from .sources.xinference_rerank_source import ( XinferenceRerankProvider as XinferenceRerankProvider, ) + case "bailian_rerank": + from .sources.bailian_rerank_source import ( + BailianRerankProvider as BailianRerankProvider, + ) except (ImportError, ModuleNotFoundError) as e: logger.critical( f"加载 {provider_config['type']}({provider_config['id']}) 提供商适配器失败：{e}。可能是因为有未安装的依赖。", diff --git a/astrbot/core/provider/sources/bailian_rerank_source.py b/astrbot/core/provider/sources/bailian_rerank_source.py new file mode 100644 index 000000000..e32b4ea0d --- /dev/null +++ b/astrbot/core/provider/sources/bailian_rerank_source.py @@ -0,0 +1,171 @@ +import os + +import aiohttp + +from astrbot import logger + +from ..entities import ProviderType, RerankResult +from ..provider import RerankProvider +from ..register import register_provider_adapter + + +@register_provider_adapter( + "bailian_rerank", "阿里云百炼文本排序适配器", provider_type=ProviderType.RERANK +) +class BailianRerankProvider(RerankProvider): + """阿里云百炼文本重排序适配器""" + + def __init__(self, provider_config: dict, provider_settings: dict) -> None: + super().__init__(provider_config, provider_settings) + self.provider_config = provider_config + self.provider_settings = provider_settings + + # API配置 + self.api_key = provider_config.get("rerank_api_key", "") + if not self.api_key: + self.api_key = os.getenv("DASHSCOPE_API_KEY", "") + + if not self.api_key: + raise ValueError( + "阿里云百炼 API Key 不能为空，请在配置中设置 rerank_api_key 或设置环境变量 DASHSCOPE_API_KEY" + ) + + self.model = provider_config.get("rerank_model", "qwen3-rerank") + self.timeout = provider_config.get("timeout", 30) + self.default_top_n = provider_config.get("top_n", 5) + self.return_documents = provider_config.get("return_documents", False) + self.instruct = provider_config.get("instruct", "") + + self.base_url = provider_config.get( + "rerank_api_base", + "https://dashscope.aliyuncs.com/api/v1/services/rerank/text-rerank/text-rerank", + ) + + # 设置HTTP客户端 + headers = { + "Authorization": f"Bearer {self.api_key}", + "Content-Type": "application/json", + } + + self.client = aiohttp.ClientSession( + headers=headers, timeout=aiohttp.ClientTimeout(total=self.timeout) + ) + + # 设置模型名称 + self.set_model(self.model) + + logger.info(f"AstrBot 百炼 Rerank 初始化完成。模型: {self.model}") + + async def rerank( + self, + query: str, + documents: list[str], + top_n: int | None = None, + ) -> list[RerankResult]: + """ + 对文档进行重排序 + + Args: + query: 查询文本 + documents: 待排序的文档列表 + top_n: 返回前N个结果，如果为None则使用配置中的默认值 + + Returns: + 重排序结果列表 + """ + if not documents: + logger.warning("文档列表为空，返回空结果") + return [] + + if not query.strip(): + logger.warning("查询文本为空，返回空结果") + return [] + + # 检查限制 + if len(documents) > 500: + logger.warning( + f"文档数量({len(documents)})超过限制(500)，将截断前500个文档" + ) + documents = documents[:500] + + # 使用传入的top_n或默认配置 + final_top_n = top_n if top_n is not None else self.default_top_n + + try: + # 构建请求载荷 + payload = { + "model": self.model, + "input": {"query": query, "documents": documents}, + } + + # 添加可选参数 + parameters = {} + if final_top_n is not None: + parameters["top_n"] = final_top_n + if self.return_documents: + parameters["return_documents"] = True + if self.instruct and self.model == "qwen3-rerank": + parameters["instruct"] = self.instruct + + if parameters: + payload["parameters"] = parameters + + logger.debug( + f"百炼 Rerank 请求: query='{query[:50]}...', 文档数量={len(documents)}" + ) + + # 发送请求 + async with self.client.post(self.base_url, json=payload) as response: + response.raise_for_status() + response_data = await response.json() + + # 检查响应状态 + if "code" in response_data and response_data["code"] != "200": + error_msg = response_data.get("message", "未知错误") + raise Exception( + f"百炼 API 返回错误: {response_data['code']} - {error_msg}" + ) + + # 解析结果 + output = response_data.get("output", {}) + results = output.get("results", []) + + if not results: + logger.warning(f"百炼 Rerank 返回空结果: {response_data}") + return [] + + # 转换为RerankResult对象 + rerank_results = [] + for result in results: + rerank_result = RerankResult( + index=result["index"], relevance_score=result["relevance_score"] + ) + rerank_results.append(rerank_result) + + logger.debug(f"百炼 Rerank 成功返回 {len(rerank_results)} 个结果") + + # 记录使用量信息 + usage = response_data.get("usage", {}) + total_tokens = usage.get("total_tokens", 0) + if total_tokens > 0: + logger.debug(f"百炼 Rerank 消耗 Token 数量: {total_tokens}") + + return rerank_results + + except aiohttp.ClientError as e: + logger.error(f"百炼 Rerank 网络请求失败: {e}") + raise Exception(f"网络请求失败: {e}") + except Exception as e: + logger.error(f"百炼 Rerank 处理失败: {e}") + raise Exception(f"重排序失败: {e}") + + async def terminate(self) -> None: + """关闭HTTP客户端会话""" + if self.client: + logger.info("关闭 百炼 Rerank 客户端会话") + try: + await self.client.close() + except Exception as e: + logger.error(f"关闭 百炼 Rerank 客户端时出错: {e}") + finally: + self.client = None From 2ada1deb9a8f74aeb6a224d67baea649abceaa93 Mon Sep 17 00:00:00 2001 From: piexian Date: Thu, 20 Nov 2025 08:31:50 +0800 Subject: [PATCH 2/4] =?UTF-8?q?=E4=BF=AE=E5=A4=8D=E6=96=87=E6=A1=A3?= =?UTF-8?q?=E8=BF=94=E5=9B=9E=E8=AF=BB=E5=8F=96=E9=97=AE=E9=A2=98?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- astrbot/core/config/default.py | 6 ------ astrbot/core/provider/sources/bailian_rerank_source.py | 7 +++++-- 2 files changed, 5 insertions(+), 8 deletions(-) diff --git a/astrbot/core/config/default.py b/astrbot/core/config/default.py index da0fefbf2..e37646a04 100644 --- a/astrbot/core/config/default.py +++ b/astrbot/core/config/default.py @@ -1318,7 +1318,6 @@ CONFIG_METADATA_2 = { "rerank_api_base": "https://dashscope.aliyuncs.com/api/v1/services/rerank/text-rerank/text-rerank", "rerank_model": "qwen3-rerank", "timeout": 30, - "top_n": 3, "return_documents": False, "instruct": "", }, @@ -1356,11 +1355,6 @@ CONFIG_METADATA_2 = { "description": "重排序模型名称", "type": "string", }, - "top_n": { - "description": "返回排序后的top_n个文档", - "type": "int", - "hint": "默认返回全部文档。如果指定的值大于文档总数，将返回全部文档。", - }, "return_documents": { "description": "是否在排序结果中返回文档原文", "type": "bool", diff --git a/astrbot/core/provider/sources/bailian_rerank_source.py b/astrbot/core/provider/sources/bailian_rerank_source.py index e32b4ea0d..9953699d4 100644 --- a/astrbot/core/provider/sources/bailian_rerank_source.py +++ b/astrbot/core/provider/sources/bailian_rerank_source.py @@ -32,7 +32,10 @@ class BailianRerankProvider(RerankProvider): self.model = provider_config.get("rerank_model", "qwen3-rerank") self.timeout = provider_config.get("timeout", 30) - self.default_top_n = provider_config.get("top_n", 5) + # 自动读取知识库配置的 kb_final_top_k，如果没有则使用配置中的 top_n + self.default_top_n = provider_settings.get( + "kb_final_top_k" + ) or provider_config.get("top_n", 5) self.return_documents = provider_config.get("return_documents", False) self.instruct = provider_config.get("instruct", "") @@ -88,7 +91,7 @@ class BailianRerankProvider(RerankProvider): ) documents = documents[:500] - # 使用传入的top_n或默认配置 + # 优先使用传入的top_n参数（来自知识库配置），如果没有才使用默认配置 final_top_n = top_n if top_n is not None else self.default_top_n try: From 234ce93dc1761574c72358f48b067da0230c57f5 Mon Sep 17 00:00:00 2001 From: piexian Date: Fri, 21 Nov 2025 04:07:45 +0800 Subject: [PATCH 3/4] =?UTF-8?q?refactor(bailian=5Frerank):=20=E4=BC=98?= =?UTF-8?q?=E5=8C=96=E4=BB=A3=E7=A0=81=E8=B4=A8=E9=87=8F=E5=92=8C=E9=94=99?= =?UTF-8?q?=E8=AF=AF=E5=A4=84=E7=90=86?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 移除未使用的 os 导入 - 简化 API Key 验证逻辑 - 优化 top_n 参数处理，优先使用传入值 - 改进错误处理，使用 RuntimeError 替代通用 Exception - 添加异常链保持原始错误上下文 --- .../provider/sources/bailian_rerank_source.py | 35 +++++++------------ 1 file changed, 12 insertions(+), 23 deletions(-) diff --git a/astrbot/core/provider/sources/bailian_rerank_source.py b/astrbot/core/provider/sources/bailian_rerank_source.py index 9953699d4..bbe6ab82b 100644 --- a/astrbot/core/provider/sources/bailian_rerank_source.py +++ b/astrbot/core/provider/sources/bailian_rerank_source.py @@ -1,5 +1,3 @@ -import os - import aiohttp from astrbot import logger @@ -13,7 +11,7 @@ from ..register import register_provider_adapter "bailian_rerank", "阿里云百炼文本排序适配器", provider_type=ProviderType.RERANK ) class BailianRerankProvider(RerankProvider): - """阿里云百炼文本重排序适配器""" + """阿里云百炼文本重排序适配器.""" def __init__(self, provider_config: dict, provider_settings: dict) -> None: super().__init__(provider_config, provider_settings) @@ -23,19 +21,10 @@ class BailianRerankProvider(RerankProvider): # API配置 self.api_key = provider_config.get("rerank_api_key", "") if not self.api_key: - self.api_key = os.getenv("DASHSCOPE_API_KEY", "") - - if not self.api_key: - raise ValueError( - "阿里云百炼 API Key 不能为空，请在配置中设置 rerank_api_key 或设置环境变量 DASHSCOPE_API_KEY" - ) + raise ValueError("阿里云百炼 API Key 不能为空。") self.model = provider_config.get("rerank_model", "qwen3-rerank") self.timeout = provider_config.get("timeout", 30) - # 自动读取知识库配置的 kb_final_top_k，如果没有则使用配置中的 top_n - self.default_top_n = provider_settings.get( - "kb_final_top_k" - ) or provider_config.get("top_n", 5) self.return_documents = provider_config.get("return_documents", False) self.instruct = provider_config.get("instruct", "") @@ -71,7 +60,7 @@ class BailianRerankProvider(RerankProvider): Args: query: 查询文本 documents: 待排序的文档列表 - top_n: 返回前N个结果，如果为None则使用配置中的默认值 + top_n: 返回前N个结果，如果为None则返回所有重排序结果 Returns: 重排序结果列表 @@ -91,9 +80,6 @@ class BailianRerankProvider(RerankProvider): ) documents = documents[:500] - # 优先使用传入的top_n参数（来自知识库配置），如果没有才使用默认配置 - final_top_n = top_n if top_n is not None else self.default_top_n - try: # 构建请求载荷 payload = { @@ -103,8 +89,8 @@ class BailianRerankProvider(RerankProvider): # 添加可选参数 parameters = {} - if final_top_n is not None: - parameters["top_n"] = final_top_n + if top_n is not None and top_n > 0: + parameters["top_n"] = top_n if self.return_documents: parameters["return_documents"] = True if self.instruct and self.model == "qwen3-rerank": @@ -125,9 +111,10 @@ class BailianRerankProvider(RerankProvider): # 检查响应状态 if "code" in response_data and response_data["code"] != "200": error_msg = response_data.get("message", "未知错误") - raise Exception( + api_error_msg = ( f"百炼 API 返回错误: {response_data['code']} - {error_msg}" ) + raise RuntimeError(api_error_msg) # 解析结果 output = response_data.get("output", {}) @@ -156,14 +143,16 @@ class BailianRerankProvider(RerankProvider): return rerank_results except aiohttp.ClientError as e: + error_msg = f"网络请求失败: {e}" logger.error(f"百炼 Rerank 网络请求失败: {e}") - raise Exception(f"网络请求失败: {e}") + raise RuntimeError(error_msg) from e except Exception as e: + error_msg = f"重排序失败: {e}" logger.error(f"百炼 Rerank 处理失败: {e}") - raise Exception(f"重排序失败: {e}") + raise RuntimeError(error_msg) from e async def terminate(self) -> None: - """关闭HTTP客户端会话""" + """关闭HTTP客户端会话.""" if self.client: logger.info("关闭 百炼 Rerank 客户端会话") try: From 2e608cdc098f6e7fd7b51225e909c3b83bc9098c Mon Sep 17 00:00:00 2001 From: piexian Date: Fri, 21 Nov 2025 05:34:18 +0800 Subject: [PATCH 4/4] =?UTF-8?q?refactor(bailian=5Frerank):=20=E4=BF=AE?= =?UTF-8?q?=E5=A4=8D=E8=AF=AF=E5=88=A0=E9=99=A4=E5=B9=B6=E4=BC=98=E5=8C=96?= =?UTF-8?q?top=5Fn=E5=8F=82=E6=95=B0=E5=A4=84=E7=90=86?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 移除不合理的知识库配置读取逻辑 - 添加os模块导入（用于读取环境变量） - 抽取辅助函数：_build_payload()、_parse_results()、_log_usage() - 添加自定义异常类：BailianRerankError、BailianAPIError、BailianNetworkError - 使用.get()安全访问API响应字段，避免KeyError - 使用raise ... from e保持异常链 --- .../provider/sources/bailian_rerank_source.py | 177 +++++++++++++----- 1 file changed, 125 insertions(+), 52 deletions(-) diff --git a/astrbot/core/provider/sources/bailian_rerank_source.py b/astrbot/core/provider/sources/bailian_rerank_source.py index bbe6ab82b..e6f6f1a4d 100644 --- a/astrbot/core/provider/sources/bailian_rerank_source.py +++ b/astrbot/core/provider/sources/bailian_rerank_source.py @@ -1,3 +1,5 @@ +import os + import aiohttp from astrbot import logger @@ -7,6 +9,24 @@ from ..provider import RerankProvider from ..register import register_provider_adapter +class BailianRerankError(Exception): + """百炼重排序服务异常基类""" + + pass + + +class BailianAPIError(BailianRerankError): + """百炼API返回错误""" + + pass + + +class BailianNetworkError(BailianRerankError): + """百炼网络请求错误""" + + pass + + @register_provider_adapter( "bailian_rerank", "阿里云百炼文本排序适配器", provider_type=ProviderType.RERANK ) @@ -19,7 +39,9 @@ class BailianRerankProvider(RerankProvider): self.provider_settings = provider_settings # API配置 - self.api_key = provider_config.get("rerank_api_key", "") + self.api_key = provider_config.get("rerank_api_key") or os.getenv( + "DASHSCOPE_API_KEY", "" + ) if not self.api_key: raise ValueError("阿里云百炼 API Key 不能为空。") @@ -48,6 +70,96 @@ class BailianRerankProvider(RerankProvider): logger.info(f"AstrBot 百炼 Rerank 初始化完成。模型: {self.model}") + def _build_payload( + self, query: str, documents: list[str], top_n: int | None + ) -> dict: + """构建请求载荷 + + Args: + query: 查询文本 + documents: 文档列表 + top_n: 返回前N个结果，如果为None则返回所有结果 + + Returns: + 请求载荷字典 + """ + base = {"model": self.model, "input": {"query": query, "documents": documents}} + + params = { + k: v + for k, v in [ + ("top_n", top_n if top_n is not None and top_n > 0 else None), + ("return_documents", True if self.return_documents else None), + ( + "instruct", + self.instruct + if self.instruct and self.model == "qwen3-rerank" + else None, + ), + ] + if v is not None + } + + if params: + base["parameters"] = params + + return base + + def _parse_results(self, data: dict) -> list[RerankResult]: + """解析API响应结果 + + Args: + data: API响应数据 + + Returns: + 重排序结果列表 + + Raises: + BailianAPIError: API返回错误 + KeyError: 结果缺少必要字段 + """ + # 检查响应状态 + if data.get("code", "200") != "200": + raise BailianAPIError( + f"百炼 API 错误: {data.get('code')} – {data.get('message', '')}" + ) + + results = data.get("output", {}).get("results", []) + if not results: + logger.warning(f"百炼 Rerank 返回空结果: {data}") + return [] + + # 转换为RerankResult对象，使用.get()避免KeyError + rerank_results = [] + for idx, result in enumerate(results): + try: + index = result.get("index", idx) + relevance_score = result.get("relevance_score", 0.0) + + if relevance_score is None: + logger.warning(f"结果 {idx} 缺少 relevance_score，使用默认值 0.0") + relevance_score = 0.0 + + rerank_result = RerankResult( + index=index, relevance_score=relevance_score + ) + rerank_results.append(rerank_result) + except Exception as e: + logger.warning(f"解析结果 {idx} 时出错: {e}, result={result}") + continue + + return rerank_results + + def _log_usage(self, data: dict) -> None: + """记录使用量信息 + + Args: + data: API响应数据 + """ + tokens = data.get("usage", {}).get("total_tokens", 0) + if tokens > 0: + logger.debug(f"百炼 Rerank 消耗 Token: {tokens}") + async def rerank( self, query: str, @@ -60,7 +172,7 @@ class BailianRerankProvider(RerankProvider): Args: query: 查询文本 documents: 待排序的文档列表 - top_n: 返回前N个结果，如果为None则返回所有重排序结果 + top_n: 返回前N个结果，如果为None则使用配置中的默认值 Returns: 重排序结果列表 @@ -81,23 +193,8 @@ class BailianRerankProvider(RerankProvider): documents = documents[:500] try: - # 构建请求载荷 - payload = { - "model": self.model, - "input": {"query": query, "documents": documents}, - } - - # 添加可选参数 - parameters = {} - if top_n is not None and top_n > 0: - parameters["top_n"] = top_n - if self.return_documents: - parameters["return_documents"] = True - if self.instruct and self.model == "qwen3-rerank": - parameters["instruct"] = self.instruct - - if parameters: - payload["parameters"] = parameters + # 构建请求载荷，如果top_n为None则返回所有重排序结果 + payload = self._build_payload(query, documents, top_n) logger.debug( f"百炼 Rerank 请求: query='{query[:50]}...', 文档数量={len(documents)}" @@ -108,48 +205,24 @@ class BailianRerankProvider(RerankProvider): response.raise_for_status() response_data = await response.json() - # 检查响应状态 - if "code" in response_data and response_data["code"] != "200": - error_msg = response_data.get("message", "未知错误") - api_error_msg = ( - f"百炼 API 返回错误: {response_data['code']} - {error_msg}" - ) - raise RuntimeError(api_error_msg) + # 解析结果并记录使用量 + results = self._parse_results(response_data) + self._log_usage(response_data) - # 解析结果 - output = response_data.get("output", {}) - results = output.get("results", []) + logger.debug(f"百炼 Rerank 成功返回 {len(results)} 个结果") - if not results: - logger.warning(f"百炼 Rerank 返回空结果: {response_data}") - return [] - - # 转换为RerankResult对象 - rerank_results = [] - for result in results: - rerank_result = RerankResult( - index=result["index"], relevance_score=result["relevance_score"] - ) - rerank_results.append(rerank_result) - - logger.debug(f"百炼 Rerank 成功返回 {len(rerank_results)} 个结果") - - # 记录使用量信息 - usage = response_data.get("usage", {}) - total_tokens = usage.get("total_tokens", 0) - if total_tokens > 0: - logger.debug(f"百炼 Rerank 消耗 Token 数量: {total_tokens}") - - return rerank_results + return results except aiohttp.ClientError as e: error_msg = f"网络请求失败: {e}" logger.error(f"百炼 Rerank 网络请求失败: {e}") - raise RuntimeError(error_msg) from e + raise BailianNetworkError(error_msg) from e + except BailianRerankError: + raise except Exception as e: error_msg = f"重排序失败: {e}" logger.error(f"百炼 Rerank 处理失败: {e}") - raise RuntimeError(error_msg) from e + raise BailianRerankError(error_msg) from e async def terminate(self) -> None: """关闭HTTP客户端会话."""