feat: add MCP support (#20716)

Co-authored-by: QuantumGhost <obelisk.reg+git@gmail.com>

api/core/tools/mcp_tool/provider.py

@@ -0,0 +1,130 @@
+import json
+from typing import Any
+
+from core.mcp.types import Tool as RemoteMCPTool
+from core.tools.__base.tool_provider import ToolProviderController
+from core.tools.__base.tool_runtime import ToolRuntime
+from core.tools.entities.common_entities import I18nObject
+from core.tools.entities.tool_entities import (
+    ToolDescription,
+    ToolEntity,
+    ToolIdentity,
+    ToolProviderEntityWithPlugin,
+    ToolProviderIdentity,
+    ToolProviderType,
+)
+from core.tools.mcp_tool.tool import MCPTool
+from models.tools import MCPToolProvider
+from services.tools.tools_transform_service import ToolTransformService
+
+
+class MCPToolProviderController(ToolProviderController):
+    provider_id: str
+    entity: ToolProviderEntityWithPlugin
+
+    def __init__(self, entity: ToolProviderEntityWithPlugin, provider_id: str, tenant_id: str, server_url: str) -> None:
+        super().__init__(entity)
+        self.entity = entity
+        self.tenant_id = tenant_id
+        self.provider_id = provider_id
+        self.server_url = server_url
+
+    @property
+    def provider_type(self) -> ToolProviderType:
+        """
+        returns the type of the provider
+
+        :return: type of the provider
+        """
+        return ToolProviderType.MCP
+
+    @classmethod
+    def _from_db(cls, db_provider: MCPToolProvider) -> "MCPToolProviderController":
+        """
+        from db provider
+        """
+        tools = []
+        tools_data = json.loads(db_provider.tools)
+        remote_mcp_tools = [RemoteMCPTool(**tool) for tool in tools_data]
+        user = db_provider.load_user()
+        tools = [
+            ToolEntity(
+                identity=ToolIdentity(
+                    author=user.name if user else "Anonymous",
+                    name=remote_mcp_tool.name,
+                    label=I18nObject(en_US=remote_mcp_tool.name, zh_Hans=remote_mcp_tool.name),
+                    provider=db_provider.server_identifier,
+                    icon=db_provider.icon,
+                ),
+                parameters=ToolTransformService.convert_mcp_schema_to_parameter(remote_mcp_tool.inputSchema),
+                description=ToolDescription(
+                    human=I18nObject(
+                        en_US=remote_mcp_tool.description or "", zh_Hans=remote_mcp_tool.description or ""
+                    ),
+                    llm=remote_mcp_tool.description or "",
+                ),
+                output_schema=None,
+                has_runtime_parameters=len(remote_mcp_tool.inputSchema) > 0,
+            )
+            for remote_mcp_tool in remote_mcp_tools
+        ]
+
+        return cls(
+            entity=ToolProviderEntityWithPlugin(
+                identity=ToolProviderIdentity(
+                    author=user.name if user else "Anonymous",
+                    name=db_provider.name,
+                    label=I18nObject(en_US=db_provider.name, zh_Hans=db_provider.name),
+                    description=I18nObject(en_US="", zh_Hans=""),
+                    icon=db_provider.icon,
+                ),
+                plugin_id=None,
+                credentials_schema=[],
+                tools=tools,
+            ),
+            provider_id=db_provider.server_identifier or "",
+            tenant_id=db_provider.tenant_id or "",
+            server_url=db_provider.decrypted_server_url,
+        )
+
+    def _validate_credentials(self, user_id: str, credentials: dict[str, Any]) -> None:
+        """
+        validate the credentials of the provider
+        """
+        pass
+
+    def get_tool(self, tool_name: str) -> MCPTool:  # type: ignore
+        """
+        return tool with given name
+        """
+        tool_entity = next(
+            (tool_entity for tool_entity in self.entity.tools if tool_entity.identity.name == tool_name), None
+        )
+
+        if not tool_entity:
+            raise ValueError(f"Tool with name {tool_name} not found")
+
+        return MCPTool(
+            entity=tool_entity,
+            runtime=ToolRuntime(tenant_id=self.tenant_id),
+            tenant_id=self.tenant_id,
+            icon=self.entity.identity.icon,
+            server_url=self.server_url,
+            provider_id=self.provider_id,
+        )
+
+    def get_tools(self) -> list[MCPTool]:  # type: ignore
+        """
+        get all tools
+        """
+        return [
+            MCPTool(
+                entity=tool_entity,
+                runtime=ToolRuntime(tenant_id=self.tenant_id),
+                tenant_id=self.tenant_id,
+                icon=self.entity.identity.icon,
+                server_url=self.server_url,
+                provider_id=self.provider_id,
+            )
+            for tool_entity in self.entity.tools
+        ]

api/core/tools/mcp_tool/tool.py

@@ -0,0 +1,92 @@
+import base64
+import json
+from collections.abc import Generator
+from typing import Any, Optional
+
+from core.mcp.error import MCPAuthError, MCPConnectionError
+from core.mcp.mcp_client import MCPClient
+from core.mcp.types import ImageContent, TextContent
+from core.tools.__base.tool import Tool
+from core.tools.__base.tool_runtime import ToolRuntime
+from core.tools.entities.tool_entities import ToolEntity, ToolInvokeMessage, ToolParameter, ToolProviderType
+
+
+class MCPTool(Tool):
+    tenant_id: str
+    icon: str
+    runtime_parameters: Optional[list[ToolParameter]]
+    server_url: str
+    provider_id: str
+
+    def __init__(
+        self, entity: ToolEntity, runtime: ToolRuntime, tenant_id: str, icon: str, server_url: str, provider_id: str
+    ) -> None:
+        super().__init__(entity, runtime)
+        self.tenant_id = tenant_id
+        self.icon = icon
+        self.runtime_parameters = None
+        self.server_url = server_url
+        self.provider_id = provider_id
+
+    def tool_provider_type(self) -> ToolProviderType:
+        return ToolProviderType.MCP
+
+    def _invoke(
+        self,
+        user_id: str,
+        tool_parameters: dict[str, Any],
+        conversation_id: Optional[str] = None,
+        app_id: Optional[str] = None,
+        message_id: Optional[str] = None,
+    ) -> Generator[ToolInvokeMessage, None, None]:
+        from core.tools.errors import ToolInvokeError
+
+        try:
+            with MCPClient(self.server_url, self.provider_id, self.tenant_id, authed=True) as mcp_client:
+                tool_parameters = self._handle_none_parameter(tool_parameters)
+                result = mcp_client.invoke_tool(tool_name=self.entity.identity.name, tool_args=tool_parameters)
+        except MCPAuthError as e:
+            raise ToolInvokeError("Please auth the tool first") from e
+        except MCPConnectionError as e:
+            raise ToolInvokeError(f"Failed to connect to MCP server: {e}") from e
+        except Exception as e:
+            raise ToolInvokeError(f"Failed to invoke tool: {e}") from e
+
+        for content in result.content:
+            if isinstance(content, TextContent):
+                try:
+                    content_json = json.loads(content.text)
+                    if isinstance(content_json, dict):
+                        yield self.create_json_message(content_json)
+                    elif isinstance(content_json, list):
+                        for item in content_json:
+                            yield self.create_json_message(item)
+                    else:
+                        yield self.create_text_message(content.text)
+                except json.JSONDecodeError:
+                    yield self.create_text_message(content.text)
+
+            elif isinstance(content, ImageContent):
+                yield self.create_blob_message(
+                    blob=base64.b64decode(content.data), meta={"mime_type": content.mimeType}
+                )
+
+    def fork_tool_runtime(self, runtime: ToolRuntime) -> "MCPTool":
+        return MCPTool(
+            entity=self.entity,
+            runtime=runtime,
+            tenant_id=self.tenant_id,
+            icon=self.icon,
+            server_url=self.server_url,
+            provider_id=self.provider_id,
+        )
+
+    def _handle_none_parameter(self, parameter: dict[str, Any]) -> dict[str, Any]:
+        """
+        in mcp tool invoke, if the parameter is empty, it will be set to None
+        """
+        return {
+            key: value
+            for key, value in parameter.items()
+            if value is not None and not (isinstance(value, str) and value.strip() == "")
+        }