Fix MCP prefix tool build

Author: wu-clan

chat_runtime.py

@@ -4,7 +4,7 @@
 from ag_ui.core import RunAgentInput, RunErrorEvent
 from pydantic_ai import Agent, AgentRunResult, BinaryImage, ModelRequest, ModelResponse, TextPart, UserPromptPart
 from pydantic_ai.builtin_tools import AbstractBuiltinTool, CodeExecutionTool, ImageGenerationTool
-from pydantic_ai.capabilities import AbstractCapability, BuiltinTool, Thinking, Toolset
+from pydantic_ai.capabilities import AbstractCapability, BuiltinTool, Thinking
 from pydantic_core import to_jsonable_python
 from sqlalchemy.ext.asyncio import AsyncSession
 from starlette.responses import StreamingResponse
@@ -94,9 +94,7 @@ async def build_agent_tools_capabilities(
     if 'thinking' in forwarded_props.model_fields_set and forwarded_props.thinking is not None:
         capabilities.append(Thinking(forwarded_props.thinking))
     if forwarded_props.mcp_ids:
-        capabilities.extend(
-            Toolset(toolset) for toolset in await mcp_service.get_toolsets(db=db, mcp_ids=forwarded_props.mcp_ids)
-        )
+        capabilities.extend(await mcp_service.get_capabilities(db=db, mcp_ids=forwarded_props.mcp_ids))
 
     tools, search_capabilities = build_chat_search_tools(
         web_search=forwarded_props.web_search,

plugin.toml

@@ -1,6 +1,6 @@
 [plugin]
 summary = 'AI'
-version = '0.1.0'
+version = '0.1.1'
 description = '为系统提供 AI 赋能'
 author = 'wu-clan'
 tags = ['ai', 'mcp', 'agent']

service/mcp_service.py

@@ -3,19 +3,18 @@
 from collections.abc import Sequence
 from typing import Any
 
+from pydantic_ai.capabilities import AbstractCapability, PrefixTools, Toolset
 from pydantic_ai.mcp import MCPServerSSE, MCPServerStdio, MCPServerStreamableHTTP
-from pydantic_ai.toolsets import AbstractToolset, PrefixedToolset
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from backend.common.exception import errors
 from backend.common.pagination import paging_data
 from backend.plugin.ai.crud.crud_mcp import mcp_dao
+from backend.plugin.ai.dataclasses import ChatAgentDeps
 from backend.plugin.ai.enums import McpType
 from backend.plugin.ai.model import Mcp
 from backend.plugin.ai.schema.mcp import CreateMcpParam, UpdateMcpParam
 
-McpToolset = AbstractToolset[Any]
-
 
 class McpService:
     """MCP 服务类"""
@@ -45,16 +44,16 @@ async def get_all(*, db: AsyncSession) -> Sequence[Mcp]:
         return await mcp_dao.get_all(db)
 
     @staticmethod
-    async def get_toolsets(*, db: AsyncSession, mcp_ids: list[int]) -> list[McpToolset]:
+    async def get_capabilities(*, db: AsyncSession, mcp_ids: list[int]) -> list[AbstractCapability[ChatAgentDeps]]:
         """
-        获取 MCP 工具集
+        获取 MCP 能力
 
         :param db: 数据库会话
         :param mcp_ids: MCP ID 列表
         :return:
         """
         mcps = await mcp_dao.get_by_ids(db, mcp_ids)
-        toolsets: list[McpToolset] = []
+        capabilities: list[AbstractCapability[ChatAgentDeps]] = []
         for mcp in mcps:
             headers = json.loads(mcp.headers) if isinstance(mcp.headers, str) else (mcp.headers or {})
             if not isinstance(headers, dict):
@@ -91,9 +90,8 @@ async def get_toolsets(*, db: AsyncSession, mcp_ids: list[int]) -> list[McpTools
                     timeout=mcp.timeout,
                     read_timeout=mcp.read_timeout,
                 )
-            # 此举是为了为避免 MCP 工具名称冲突
-            toolsets.append(PrefixedToolset(toolset, prefix=f'mcp_{mcp.id}'))
-        return toolsets
+            capabilities.append(PrefixTools(Toolset(toolset), prefix=f'mcp_{mcp.id}'))
+        return capabilities
 
     @staticmethod
     async def get_list(*, db: AsyncSession, name: str | None, type: int | None) -> dict[str, Any]: