AstrBot 4.5.0__py3-none-any.whl → 4.5.2__py3-none-any.whl

astrbot/core/provider/func_tool_manager.py

@@ -1,17 +1,18 @@
 from __future__ import annotations
+
+import asyncio
 import json
 import os
-import asyncio
+from collections.abc import Awaitable, Callable
+from typing import Any
+
 import aiohttp
 
-from typing import Dict, List, Awaitable, Callable, Any
 from astrbot import logger
 from astrbot.core import sp
-
+from astrbot.core.agent.mcp_client import MCPClient, MCPTool
+from astrbot.core.agent.tool import FunctionTool, ToolSet
 from astrbot.core.utils.astrbot_path import get_astrbot_data_path
-from astrbot.core.agent.mcp_client import MCPClient
-from astrbot.core.agent.tool import ToolSet, FunctionTool
-
 
 DEFAULT_MCP_CONFIG = {"mcpServers": {}}
 
@@ -30,7 +31,7 @@ FuncTool = FunctionTool
 
 def _prepare_config(config: dict) -> dict:
     """准备配置，处理嵌套格式"""
-    if "mcpServers" in config and config["mcpServers"]:
+    if config.get("mcpServers"):
         first_key = next(iter(config["mcpServers"]))
         config = config["mcpServers"][first_key]
     config.pop("active", None)
@@ -72,8 +73,7 @@ async def _quick_test_mcp_connection(config: dict) -> tuple[bool, str]:
                 ) as response:
                     if response.status == 200:
                         return True, ""
-                    else:
-                        return False, f"HTTP {response.status}: {response.reason}"
+                    return False, f"HTTP {response.status}: {response.reason}"
             else:
                 async with session.get(
                     url,
@@ -85,8 +85,7 @@ async def _quick_test_mcp_connection(config: dict) -> tuple[bool, str]:
                 ) as response:
                     if response.status == 200:
                         return True, ""
-                    else:
-                        return False, f"HTTP {response.status}: {response.reason}"
+                    return False, f"HTTP {response.status}: {response.reason}"
 
     except asyncio.TimeoutError:
         return False, f"连接超时: {timeout}秒"
@@ -96,10 +95,10 @@ async def _quick_test_mcp_connection(config: dict) -> tuple[bool, str]:
 
 class FunctionToolManager:
     def __init__(self) -> None:
-        self.func_list: List[FuncTool] = []
-        self.mcp_client_dict: Dict[str, MCPClient] = {}
+        self.func_list: list[FuncTool] = []
+        self.mcp_client_dict: dict[str, MCPClient] = {}
         """MCP 服务列表"""
-        self.mcp_client_event: Dict[str, asyncio.Event] = {}
+        self.mcp_client_event: dict[str, asyncio.Event] = {}
 
     def empty(self) -> bool:
         return len(self.func_list) == 0
@@ -150,14 +149,12 @@ class FunctionToolManager:
                 func_args=func_args,
                 desc=desc,
                 handler=handler,
-            )
+            ),
         )
         logger.info(f"添加函数调用工具: {name}")
 
     def remove_func(self, name: str) -> None:
-        """
-        删除一个函数调用工具。
-        """
+        """删除一个函数调用工具。"""
         for i, f in enumerate(self.func_list):
             if f.name == name:
                 self.func_list.pop(i)
@@ -202,16 +199,16 @@ class FunctionToolManager:
             logger.info(f"未找到 MCP 服务配置文件，已创建默认配置文件 {mcp_json_file}")
             return
 
-        mcp_server_json_obj: Dict[str, Dict] = json.load(
-            open(mcp_json_file, "r", encoding="utf-8")
+        mcp_server_json_obj: dict[str, dict] = json.load(
+            open(mcp_json_file, encoding="utf-8"),
         )["mcpServers"]
 
-        for name in mcp_server_json_obj.keys():
+        for name in mcp_server_json_obj:
             cfg = mcp_server_json_obj[name]
             if cfg.get("active", True):
                 event = asyncio.Event()
                 asyncio.create_task(
-                    self._init_mcp_client_task_wrapper(name, cfg, event)
+                    self._init_mcp_client_task_wrapper(name, cfg, event),
                 )
                 self.mcp_client_event[name] = event
 
@@ -257,18 +254,15 @@ class FunctionToolManager:
         self.func_list = [
             f
             for f in self.func_list
-            if not (f.origin == "mcp" and f.mcp_server_name == name)
+            if not (isinstance(f, MCPTool) and f.mcp_server_name == name)
         ]
 
         # 将 MCP 工具转换为 FuncTool 并添加到 func_list
         for tool in mcp_client.tools:
-            func_tool = FuncTool(
-                name=tool.name,
-                parameters=tool.inputSchema,
-                description=tool.description,
-                origin="mcp",
-                mcp_server_name=name,
+            func_tool = MCPTool(
+                mcp_tool=tool,
                 mcp_client=mcp_client,
+                mcp_server_name=name,
             )
             self.func_list.append(func_tool)
 
@@ -287,7 +281,7 @@ class FunctionToolManager:
             self.func_list = [
                 f
                 for f in self.func_list
-                if not (f.origin == "mcp" and f.mcp_server_name == name)
+                if not (isinstance(f, MCPTool) and f.mcp_server_name == name)
             ]
             logger.info(f"已关闭 MCP 服务 {name}")
 
@@ -325,9 +319,11 @@ class FunctionToolManager:
             event (asyncio.Event): Event to signal when the MCP client is ready.
             ready_future (asyncio.Future): Future to signal when the MCP client is ready.
             timeout (int): Timeout for the initialization.
+
         Raises:
             TimeoutError: If the initialization does not complete within the specified timeout.
             Exception: If there is an error during initialization.
+
         """
         if not event:
             event = asyncio.Event()
@@ -336,7 +332,7 @@ class FunctionToolManager:
         if name in self.mcp_client_dict:
             return
         asyncio.create_task(
-            self._init_mcp_client_task_wrapper(name, config, event, ready_future)
+            self._init_mcp_client_task_wrapper(name, config, event, ready_future),
         )
         try:
             await asyncio.wait_for(ready_future, timeout=timeout)
@@ -349,13 +345,16 @@ class FunctionToolManager:
                 raise exc
 
     async def disable_mcp_server(
-        self, name: str | None = None, timeout: float = 10
+        self,
+        name: str | None = None,
+        timeout: float = 10,
     ) -> None:
         """Disable an MCP server by its name.
 
         Args:
             name (str): The name of the MCP server to disable. If None, ALL MCP servers will be disabled.
             timeout (int): Timeout.
+
         """
         if name:
             if name not in self.mcp_client_event:
@@ -372,7 +371,7 @@ class FunctionToolManager:
                 self.func_list = [
                     f
                     for f in self.func_list
-                    if f.origin != "mcp" or f.mcp_server_name != name
+                    if not (isinstance(f, MCPTool) and f.mcp_server_name == name)
                 ]
         else:
             running_events = [
@@ -386,30 +385,26 @@ class FunctionToolManager:
             finally:
                 self.mcp_client_event.clear()
                 self.mcp_client_dict.clear()
-                self.func_list = [f for f in self.func_list if f.origin != "mcp"]
+                self.func_list = [
+                    f for f in self.func_list if not isinstance(f, MCPTool)
+                ]
 
     def get_func_desc_openai_style(self, omit_empty_parameter_field=False) -> list:
-        """
-        获得 OpenAI API 风格的**已经激活**的工具描述
-        """
+        """获得 OpenAI API 风格的**已经激活**的工具描述"""
         tools = [f for f in self.func_list if f.active]
         toolset = ToolSet(tools)
         return toolset.openai_schema(
-            omit_empty_parameter_field=omit_empty_parameter_field
+            omit_empty_parameter_field=omit_empty_parameter_field,
         )
 
     def get_func_desc_anthropic_style(self) -> list:
-        """
-        获得 Anthropic API 风格的**已经激活**的工具描述
-        """
+        """获得 Anthropic API 风格的**已经激活**的工具描述"""
         tools = [f for f in self.func_list if f.active]
         toolset = ToolSet(tools)
         return toolset.anthropic_schema()
 
     def get_func_desc_google_genai_style(self) -> dict:
-        """
-        获得 Google GenAI API 风格的**已经激活**的工具描述
-        """
+        """获得 Google GenAI API 风格的**已经激活**的工具描述"""
         tools = [f for f in self.func_list if f.active]
         toolset = ToolSet(tools)
         return toolset.google_schema()
@@ -418,13 +413,18 @@ class FunctionToolManager:
         """停用一个已经注册的函数调用工具。
 
         Returns:
-            如果没找到，会返回 False"""
+            如果没找到，会返回 False
+
+        """
         func_tool = self.get_func(name)
         if func_tool is not None:
             func_tool.active = False
 
             inactivated_llm_tools: list = sp.get(
-                "inactivated_llm_tools", [], scope="global", scope_id="global"
+                "inactivated_llm_tools",
+                [],
+                scope="global",
+                scope_id="global",
             )
             if name not in inactivated_llm_tools:
                 inactivated_llm_tools.append(name)
@@ -445,13 +445,16 @@ class FunctionToolManager:
             if func_tool.handler_module_path in star_map:
                 if not star_map[func_tool.handler_module_path].activated:
                     raise ValueError(
-                        f"此函数调用工具所属的插件 {star_map[func_tool.handler_module_path].name} 已被禁用，请先在管理面板启用再激活此工具。"
+                        f"此函数调用工具所属的插件 {star_map[func_tool.handler_module_path].name} 已被禁用，请先在管理面板启用再激活此工具。",
                     )
 
             func_tool.active = True
 
             inactivated_llm_tools: list = sp.get(
-                "inactivated_llm_tools", [], scope="global", scope_id="global"
+                "inactivated_llm_tools",
+                [],
+                scope="global",
+                scope_id="global",
             )
             if name in inactivated_llm_tools:
                 inactivated_llm_tools.remove(name)
@@ -479,7 +482,7 @@ class FunctionToolManager:
             return DEFAULT_MCP_CONFIG
 
         try:
-            with open(self.mcp_config_path, "r", encoding="utf-8") as f:
+            with open(self.mcp_config_path, encoding="utf-8") as f:
                 return json.load(f)
         except Exception as e:
             logger.error(f"加载 MCP 配置失败: {e}")
@@ -509,7 +512,8 @@ class FunctionToolManager:
                     if response.status == 200:
                         data = await response.json()
                         mcp_server_list = data.get("data", {}).get(
-                            "mcp_server_list", []
+                            "mcp_server_list",
+                            [],
                         )
                         local_mcp_config = self.load_mcp_config()
 
@@ -541,23 +545,23 @@ class FunctionToolManager:
                                     self.enable_mcp_server(
                                         name=name,
                                         config=local_mcp_config["mcpServers"][name],
-                                    )
+                                    ),
                                 )
                             await asyncio.gather(*tasks)
                             logger.info(
-                                f"从 ModelScope 同步了 {synced_count} 个 MCP 服务器"
+                                f"从 ModelScope 同步了 {synced_count} 个 MCP 服务器",
                             )
                         else:
                             logger.warning("没有找到可用的 ModelScope MCP 服务器")
                     else:
                         raise Exception(
-                            f"ModelScope API 请求失败: HTTP {response.status}"
+                            f"ModelScope API 请求失败: HTTP {response.status}",
                         )
 
         except aiohttp.ClientError as e:
-            raise Exception(f"网络连接错误: {str(e)}")
+            raise Exception(f"网络连接错误: {e!s}")
         except Exception as e:
-            raise Exception(f"同步 ModelScope MCP 服务器时发生错误: {str(e)}")
+            raise Exception(f"同步 ModelScope MCP 服务器时发生错误: {e!s}")
 
     def __str__(self):
         return str(self.func_list)

astrbot/core/provider/manager.py

@@ -5,16 +5,16 @@ from astrbot.core import logger, sp
 from astrbot.core.astrbot_config_mgr import AstrBotConfigManager
 from astrbot.core.db import BaseDatabase
 
+from ..persona_mgr import PersonaManager
 from .entities import ProviderType
 from .provider import (
+    EmbeddingProvider,
     Provider,
+    RerankProvider,
     STTProvider,
     TTSProvider,
-    EmbeddingProvider,
-    RerankProvider,
 )
 from .register import llm_tools, provider_cls_map
-from ..persona_mgr import PersonaManager
 
 
 class ProviderManager:
@@ -76,7 +76,10 @@ class ProviderManager:
         return self.persona_mgr.selected_default_persona_v3
 
     async def set_provider(
-        self, provider_id: str, provider_type: ProviderType, umo: str | None = None
+        self,
+        provider_id: str,
+        provider_type: ProviderType,
+        umo: str | None = None,
     ):
         """设置提供商。
 
@@ -86,6 +89,7 @@ class ProviderManager:
             umo (str, optional): 用户会话 ID，用于提供商会话隔离。
 
         Version 4.0.0: 这个版本下已经默认隔离提供商
+
         """
         if provider_id not in self.inst_map:
             raise ValueError(f"提供商 {provider_id} 不存在，无法设置。")
@@ -100,17 +104,20 @@ class ProviderManager:
 
         prov = self.inst_map[provider_id]
         if provider_type == ProviderType.TEXT_TO_SPEECH and isinstance(
-            prov, TTSProvider
+            prov,
+            TTSProvider,
         ):
             self.curr_tts_provider_inst = prov
             sp.put("curr_provider_tts", provider_id, scope="global", scope_id="global")
         elif provider_type == ProviderType.SPEECH_TO_TEXT and isinstance(
-            prov, STTProvider
+            prov,
+            STTProvider,
         ):
             self.curr_stt_provider_inst = prov
             sp.put("curr_provider_stt", provider_id, scope="global", scope_id="global")
         elif provider_type == ProviderType.CHAT_COMPLETION and isinstance(
-            prov, Provider
+            prov,
+            Provider,
         ):
             self.curr_provider_inst = prov
             sp.put("curr_provider", provider_id, scope="global", scope_id="global")
@@ -120,7 +127,9 @@ class ProviderManager:
         return self.inst_map.get(provider_id)
 
     def get_using_provider(
-        self, provider_type: ProviderType, umo=None
+        self,
+        provider_type: ProviderType,
+        umo=None,
     ) -> Provider | STTProvider | TTSProvider | None:
         """获取正在使用的提供商实例。
 
@@ -130,6 +139,7 @@ class ProviderManager:
 
         Returns:
             Provider: 正在使用的提供商实例。
+
         """
         provider = None
         if umo:
@@ -219,7 +229,7 @@ class ProviderManager:
             return
 
         logger.info(
-            f"载入 {provider_config['type']}({provider_config['id']}) 服务提供商 ..."
+            f"载入 {provider_config['type']}({provider_config['id']}) 服务提供商 ...",
         )
 
         # 动态导入
@@ -259,6 +269,10 @@ class ProviderManager:
                     from .sources.whisper_selfhosted_source import (
                         ProviderOpenAIWhisperSelfHost as ProviderOpenAIWhisperSelfHost,
                     )
+                case "xinference_stt":
+                    from .sources.xinference_stt_provider import (
+                        ProviderXinferenceSTT as ProviderXinferenceSTT,
+                    )
                 case "openai_tts_api":
                     from .sources.openai_tts_api_source import (
                         ProviderOpenAITTSAPI as ProviderOpenAITTSAPI,
@@ -311,20 +325,24 @@ class ProviderManager:
                     from .sources.vllm_rerank_source import (
                         VLLMRerankProvider as VLLMRerankProvider,
                     )
+                case "xinference_rerank":
+                    from .sources.xinference_rerank_source import (
+                        XinferenceRerankProvider as XinferenceRerankProvider,
+                    )
         except (ImportError, ModuleNotFoundError) as e:
             logger.critical(
-                f"加载 {provider_config['type']}({provider_config['id']}) 提供商适配器失败：{e}。可能是因为有未安装的依赖。"
+                f"加载 {provider_config['type']}({provider_config['id']}) 提供商适配器失败：{e}。可能是因为有未安装的依赖。",
             )
             return
         except Exception as e:
             logger.critical(
-                f"加载 {provider_config['type']}({provider_config['id']}) 提供商适配器失败：{e}。未知原因"
+                f"加载 {provider_config['type']}({provider_config['id']}) 提供商适配器失败：{e}。未知原因",
             )
             return
 
         if provider_config["type"] not in provider_cls_map:
             logger.error(
-                f"未找到适用于 {provider_config['type']}({provider_config['id']}) 的提供商适配器，请检查是否已经安装或者名称填写错误。已跳过。"
+                f"未找到适用于 {provider_config['type']}({provider_config['id']}) 的提供商适配器，请检查是否已经安装或者名称填写错误。已跳过。",
             )
             return
 
@@ -350,7 +368,7 @@ class ProviderManager:
                 ):
                     self.curr_stt_provider_inst = inst
                     logger.info(
-                        f"已选择 {provider_config['type']}({provider_config['id']}) 作为当前语音转文本提供商适配器。"
+                        f"已选择 {provider_config['type']}({provider_config['id']}) 作为当前语音转文本提供商适配器。",
                     )
                 if not self.curr_stt_provider_inst:
                     self.curr_stt_provider_inst = inst
@@ -366,7 +384,7 @@ class ProviderManager:
                 if self.provider_settings.get("provider_id") == provider_config["id"]:
                     self.curr_tts_provider_inst = inst
                     logger.info(
-                        f"已选择 {provider_config['type']}({provider_config['id']}) 作为当前文本转语音提供商适配器。"
+                        f"已选择 {provider_config['type']}({provider_config['id']}) 作为当前文本转语音提供商适配器。",
                     )
                 if not self.curr_tts_provider_inst:
                     self.curr_tts_provider_inst = inst
@@ -389,7 +407,7 @@ class ProviderManager:
                 ):
                     self.curr_provider_inst = inst
                     logger.info(
-                        f"已选择 {provider_config['type']}({provider_config['id']}) 作为当前提供商适配器。"
+                        f"已选择 {provider_config['type']}({provider_config['id']}) 作为当前提供商适配器。",
                     )
                 if not self.curr_provider_inst:
                     self.curr_provider_inst = inst
@@ -408,10 +426,10 @@ class ProviderManager:
             self.inst_map[provider_config["id"]] = inst
         except Exception as e:
             logger.error(
-                f"实例化 {provider_config['type']}({provider_config['id']}) 提供商适配器失败：{e}"
+                f"实例化 {provider_config['type']}({provider_config['id']}) 提供商适配器失败：{e}",
             )
             raise Exception(
-                f"实例化 {provider_config['type']}({provider_config['id']}) 提供商适配器失败：{e}"
+                f"实例化 {provider_config['type']}({provider_config['id']}) 提供商适配器失败：{e}",
             )
 
     async def reload(self, provider_config: dict):
@@ -431,7 +449,7 @@ class ProviderManager:
         elif self.curr_provider_inst is None and len(self.provider_insts) > 0:
             self.curr_provider_inst = self.provider_insts[0]
             logger.info(
-                f"自动选择 {self.curr_provider_inst.meta().id} 作为当前提供商适配器。"
+                f"自动选择 {self.curr_provider_inst.meta().id} 作为当前提供商适配器。",
             )
 
         if len(self.stt_provider_insts) == 0:
@@ -439,7 +457,7 @@ class ProviderManager:
         elif self.curr_stt_provider_inst is None and len(self.stt_provider_insts) > 0:
             self.curr_stt_provider_inst = self.stt_provider_insts[0]
             logger.info(
-                f"自动选择 {self.curr_stt_provider_inst.meta().id} 作为当前语音转文本提供商适配器。"
+                f"自动选择 {self.curr_stt_provider_inst.meta().id} 作为当前语音转文本提供商适配器。",
             )
 
         if len(self.tts_provider_insts) == 0:
@@ -447,7 +465,7 @@ class ProviderManager:
         elif self.curr_tts_provider_inst is None and len(self.tts_provider_insts) > 0:
             self.curr_tts_provider_inst = self.tts_provider_insts[0]
             logger.info(
-                f"自动选择 {self.curr_tts_provider_inst.meta().id} 作为当前文本转语音提供商适配器。"
+                f"自动选择 {self.curr_tts_provider_inst.meta().id} 作为当前文本转语音提供商适配器。",
             )
 
     def get_insts(self):
@@ -456,7 +474,7 @@ class ProviderManager:
     async def terminate_provider(self, provider_id: str):
         if provider_id in self.inst_map:
             logger.info(
-                f"终止 {provider_id} 提供商适配器({len(self.provider_insts)}, {len(self.stt_provider_insts)}, {len(self.tts_provider_insts)}) ..."
+                f"终止 {provider_id} 提供商适配器({len(self.provider_insts)}, {len(self.stt_provider_insts)}, {len(self.tts_provider_insts)}) ...",
             )
 
             if self.inst_map[provider_id] in self.provider_insts:
@@ -483,7 +501,7 @@ class ProviderManager:
                 await self.inst_map[provider_id].terminate()  # type: ignore
 
             logger.info(
-                f"{provider_id} 提供商适配器已终止({len(self.provider_insts)}, {len(self.stt_provider_insts)}, {len(self.tts_provider_insts)})"
+                f"{provider_id} 提供商适配器已终止({len(self.provider_insts)}, {len(self.stt_provider_insts)}, {len(self.tts_provider_insts)})",
             )
             del self.inst_map[provider_id]