chibi-bot 1.6.0b0__py3-none-any.whl

chibi/services/providers/tools/mcp_management.py

@@ -0,0 +1,274 @@
+import os
+from typing import Any, Callable, Coroutine, cast
+
+from dotenv import dotenv_values
+from loguru import logger
+from mcp import ClientSession
+from mcp.types import CallToolResult
+from openai.types.chat import ChatCompletionToolParam
+from openai.types.shared_params import FunctionDefinition
+
+from chibi.config import application_settings
+from chibi.services.mcp.manager import MCPManager
+from chibi.services.providers.tools.tool import ChibiTool, RegisteredChibiTools
+from chibi.services.providers.tools.utils import AdditionalOptions
+from chibi.services.providers.utils import escape_and_truncate
+
+JsonNode = dict[str, Any] | list[Any] | str | int | float | bool | None
+
+
+def _clean_schema(schema: JsonNode) -> JsonNode:
+    """Recursively cleans the schema for LLM provider compatibility.
+
+    Removes 'additionalProperties', '$schema', and other meta-fields that cause validation errors.
+    """
+    if isinstance(schema, dict):
+        # Fields to exclude from the schema
+        exclude = {"additionalProperties", "$schema"}
+        # Also clean nested objects in 'properties' and 'items'
+        cleaned = {k: _clean_schema(v) for k, v in schema.items() if k not in exclude}
+        return cleaned
+    if isinstance(schema, list):
+        return [_clean_schema(i) for i in schema]
+    return schema
+
+
+def create_wrapper(
+    server_name: str, original_tool_name: str, chibi_tool_name: str
+) -> Callable[..., Coroutine[Any, Any, dict[str, Any]]]:
+    async def wrapper(**kwargs_inner: Any) -> dict[str, Any]:
+        tool_args = {k: v for k, v in kwargs_inner.items() if k not in list(AdditionalOptions.__annotations__.keys())}
+
+        logger.log(
+            "TOOL",
+            (
+                f"[{kwargs_inner.get('model', 'Unknown model')}] "
+                f"Calling MCP tool '{chibi_tool_name}' with args: {escape_and_truncate(tool_args)}"
+            ),
+        )
+        res: CallToolResult = await MCPManager.call_tool(
+            server_name=server_name, tool_name=original_tool_name, arguments=tool_args
+        )
+
+        if len(res.content) == 1:
+            output = res.content[0].model_dump()
+        else:
+            output = {"result": [content.model_dump() for content in res.content]}
+
+        logger.log(
+            "TOOL",
+            (
+                f"[{kwargs_inner.get('model', 'Unknown model')}] "
+                f"MCP tool '{chibi_tool_name}' returned: {escape_and_truncate(output)}",
+            ),
+        )
+        return output
+
+    return wrapper
+
+
+async def register_tools_from_mcp_session(
+    mcp_session: ClientSession, server_name: str, transport: str
+) -> dict[str, str]:
+    try:
+        tools_result = await mcp_session.list_tools()
+
+        mcp_tools = []
+        for tool in tools_result.tools:
+            chibi_tool_name = f"mcp_{server_name}_{transport}_{tool.name}"
+            class_attributes = {
+                "register": True,
+                "definition": ChatCompletionToolParam(
+                    type="function",
+                    function=FunctionDefinition(
+                        name=chibi_tool_name,
+                        description=f"MCP Server: {server_name}. {tool.description or ''}",
+                        parameters=cast(dict[str, object], _clean_schema(tool.inputSchema)),
+                    ),
+                ),
+                "name": chibi_tool_name,
+                "function": staticmethod(create_wrapper(server_name, tool.name, chibi_tool_name)),
+            }
+            type(f"MCP{server_name.capitalize()}{transport.capitalize()}Tool", (ChibiTool,), class_attributes)
+            mcp_tools.append(chibi_tool_name)
+    except Exception as e:
+        logger.error(f"Failed to register tools for MCP server '{server_name}': {e}")
+        await MCPManager.disconnect(server_name)
+        return {"status": "error", "message": f"Failed to register tools: {e}"}
+    MCPManager.associate_tools_with_server(server_name=server_name, tool_names=mcp_tools)
+
+    return {
+        "status": "success",
+        "message": (
+            f"Successfully connected to MCP server '{server_name}'. "
+            f"Registered {len(mcp_tools)} tools: {', '.join(mcp_tools)}"
+        ),
+    }
+
+
+class InitializeStdioMCPServer(ChibiTool):
+    register = application_settings.enable_mcp_stdio
+    name = "initialize_stdio_mcp_server"
+    definition = ChatCompletionToolParam(
+        type="function",
+        function=FunctionDefinition(
+            name="initialize_stdio_mcp_server",
+            description="Connect to an MCP server via stdio and register its tools dynamically.",
+            parameters={
+                "type": "object",
+                "properties": {
+                    "server_name": {
+                        "type": "string",
+                        "description": "Unique name for this server session (e.g., 'sqlite', 'github').",
+                    },
+                    "command": {
+                        "type": "string",
+                        "description": "Command to start the MCP server",
+                    },
+                    "args": {
+                        "type": "array",
+                        "items": {"type": "string"},
+                        "description": "Arguments for the command",
+                    },
+                    "env": {
+                        "type": "object",
+                        "description": "Optional environment variables",
+                    },
+                    "secret_envs": {
+                        "type": "array",
+                        "items": {"type": "string"},
+                        "description": (
+                            "A list of environment variables that need to be passed to the mcp server "
+                            "without direct access to their contents."
+                        ),
+                    },
+                    "env_file_path": {"type": "string", "description": "Absolute path to the env-file."},
+                },
+                "required": ["server_name", "command", "args"],
+            },
+        ),
+    )
+
+    @classmethod
+    async def function(
+        cls,
+        server_name: str,
+        command: str,
+        args: list[str],
+        env: dict[str, str] | None = None,
+        secret_envs: list[str] | None = None,
+        env_file_path: str | None = None,
+        **kwargs: Any,
+    ) -> dict[str, Any]:
+        """Connect to an MCP server via stdio and register its tools dynamically.
+
+        Args:
+            server_name: Unique name for the session.
+            command: Command to start the server.
+            args: Command arguments.
+            env: Optional environment variables.
+            secret_envs: Optional list of environment variable names.
+            env_file_path: Optional path to the env file.
+
+        Returns:
+            dict[str, Any]: Success or error message.
+        """
+        if secret_envs:
+            env = env or {}
+            for secret_env in secret_envs:
+                env[secret_env] = os.environ[secret_env]
+
+        if env_file_path:
+            if not os.path.exists(env_file_path):
+                raise ValueError(f"The file {env_file_path} does not exist.")
+            env = env or {}
+            envs_from_file = dotenv_values(env_file_path)
+            envs_with_values_from_file = {k: v for k, v in envs_from_file.items() if v}
+            if envs_with_values_from_file:
+                env.update(envs_with_values_from_file)
+
+        session = await MCPManager.connect_stdio(name=server_name, command=command, args=args, env=env)
+        return await register_tools_from_mcp_session(mcp_session=session, server_name=server_name, transport="stdio")
+
+
+class InitializeSseMCPServer(ChibiTool):
+    register = application_settings.enable_mcp_sse
+    name = "initialize_sse_mcp_server"
+    definition = ChatCompletionToolParam(
+        type="function",
+        function=FunctionDefinition(
+            name="initialize_sse_mcp_server",
+            description="Connect to an MCP server via SSE and register its tools dynamically.",
+            parameters={
+                "type": "object",
+                "properties": {
+                    "server_name": {
+                        "type": "string",
+                        "description": "Unique name for this server session (e.g., 'sqlite', 'github').",
+                    },
+                    "url": {"type": "string", "description": "SSE endpoint URL (required for sse)."},
+                },
+                "required": ["server_name", "url"],
+            },
+        ),
+    )
+
+    @classmethod
+    async def function(
+        cls,
+        server_name: str,
+        url: str,
+        **kwargs: Any,
+    ) -> dict[str, Any]:
+        """Executes the tool.
+
+        Args:
+            server_name: Unique name for the session.
+            url: SSE endpoint URL (sse).
+
+        Returns:
+            dict[str, Any]: Success or error message.
+        """
+        try:
+            session = await MCPManager.connect_sse(server_name, url)
+        except Exception:
+            raise
+        return await register_tools_from_mcp_session(mcp_session=session, server_name=server_name, transport="sse")
+
+
+class DeinitializeMCPServer(ChibiTool):
+    """Disconnects from an MCP server and removes its tools."""
+
+    register = application_settings.enable_mcp_stdio or application_settings.enable_mcp_sse
+    name = "deinitialize_mcp_server"
+    definition = ChatCompletionToolParam(
+        type="function",
+        function=FunctionDefinition(
+            name="deinitialize_mcp_server",
+            description="Disconnect from an MCP server and unregister its tools.",
+            parameters={
+                "type": "object",
+                "properties": {"server_name": {"type": "string", "description": "Name of the server to disconnect."}},
+                "required": ["server_name"],
+            },
+        ),
+    )
+
+    @classmethod
+    async def function(cls, server_name: str, **kwargs: Any) -> dict[str, Any]:
+        """Executes the tool.
+
+        Args:
+            server_name: Name of the server to disconnect.
+
+        Returns:
+            dict[str, Any]: Success or error message.
+        """
+        tools_to_deregister = await MCPManager.disconnect(server_name)
+        if tools_to_deregister:
+            RegisteredChibiTools.deregister_tools(tools_to_deregister)
+
+        return {
+            "status": "success",
+            "message": f"Disconnected from MCP server '{server_name}' and removed its tools.",
+        }

chibi/services/providers/tools/mcp_simple.py

@@ -0,0 +1,72 @@
+from typing import Any, Unpack
+
+from mcp import ClientSession, StdioServerParameters
+from mcp.client.stdio import stdio_client
+
+from chibi.services.providers.tools.tool import ChibiTool
+from chibi.services.providers.tools.utils import AdditionalOptions
+
+
+class McpEchoTool(ChibiTool):
+    """Test/demo only"""
+
+    register = False
+    name = "mcp_echo"
+    definition = {
+        "type": "function",
+        "function": {
+            "name": "mcp_echo",
+            "description": "Send a message to a local MCP echo server and get the response.",
+            "parameters": {
+                "type": "object",
+                "properties": {"message": {"type": "string", "description": "The message to echo back"}},
+                "required": ["message"],
+            },
+        },
+    }
+
+    @classmethod
+    async def function(cls, message: str, **kwargs: Unpack[AdditionalOptions]) -> dict[str, Any]:
+        server_script = """
+import asyncio
+import sys
+from mcp.server.fastmcp import FastMCP
+
+mcp = FastMCP("Echo")
+
+@mcp.tool()
+def echo(message: str) -> str:
+    return f"Echo from MCP: {message}"
+
+if __name__ == "__main__":
+    mcp.run(transport='stdio')
+"""
+        import os
+        import tempfile
+
+        with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
+            f.write(server_script)
+            server_path = f.name
+
+        try:
+            server_params = StdioServerParameters(
+                command="python",
+                args=[server_path],
+            )
+
+            async with stdio_client(server_params) as (read, write):
+                async with ClientSession(read, write) as session:
+                    await session.initialize()
+                    result = await session.call_tool("echo", arguments={"message": message})
+
+                    if result.content and len(result.content) > 0:
+                        content_item = result.content[0]
+                        if content_item.type == "text":
+                            return {"response": content_item.text}
+                    return {"response": "No content returned"}
+
+        except Exception as e:
+            return {"error": f"MCP Error: {str(e)}"}
+        finally:
+            if os.path.exists(server_path):
+                os.unlink(server_path)