pydantic-rpc 0.6.1__py3-none-any.whl → 0.8.0__py3-none-any.whl

pydantic_rpc/mcp/__init__.py

@@ -0,0 +1,5 @@
+"""MCP (Model Context Protocol) support for pydantic-rpc."""
+
+from .exporter import MCPExporter
+
+__all__ = ["MCPExporter"]

pydantic_rpc/mcp/converter.py

@@ -0,0 +1,115 @@
+"""Type conversion utilities for MCP integration."""
+
+import datetime
+import enum
+import inspect
+from collections.abc import AsyncIterator
+from typing import Any, Callable, Union, get_args, get_origin
+
+from pydantic import BaseModel
+
+_SIMPLE_TYPE_MAP = {
+    int: {"type": "integer"},
+    float: {"type": "number"},
+    str: {"type": "string"},
+    bool: {"type": "boolean"},
+    bytes: {"type": "string", "format": "byte"},
+    datetime.datetime: {"type": "string", "format": "date-time"},
+    datetime.timedelta: {"type": "string", "format": "duration"},
+}
+
+
+def is_streaming_return(return_type: Any) -> bool:
+    """Check if the return type is a streaming response (AsyncIterator)."""
+    origin = get_origin(return_type)
+    return origin is AsyncIterator
+
+
+def python_type_to_json_type(python_type: Any) -> dict[str, Any]:
+    """Convert Python type to JSON Schema type."""
+    if python_type in _SIMPLE_TYPE_MAP:
+        return _SIMPLE_TYPE_MAP[python_type].copy()
+
+    match python_type:
+        case t if get_origin(t) is list:
+            item_type = get_args(t)[0]
+            return {"type": "array", "items": python_type_to_json_type(item_type)}
+        case t if get_origin(t) is dict:
+            _, value_type = get_args(t)
+            return {
+                "type": "object",
+                "additionalProperties": python_type_to_json_type(value_type),
+            }
+        case t if get_origin(t) is Union:
+            union_args = get_args(t)
+            non_none_types = [arg for arg in union_args if arg is not type(None)]
+            if len(non_none_types) == 1:
+                schema = python_type_to_json_type(non_none_types[0])
+                schema["nullable"] = True
+                return schema
+            else:
+                return {
+                    "oneOf": [python_type_to_json_type(arg) for arg in non_none_types]
+                }
+        case t if inspect.isclass(t) and issubclass(t, BaseModel):
+            return t.model_json_schema()
+        case t if inspect.isclass(t) and issubclass(t, enum.Enum):
+            return {"type": "string", "enum": [e.value for e in t]}
+        case _:
+            # Default to object type for unknown types
+            return {"type": "object"}
+
+
+def extract_method_info(method: Callable[..., Any]) -> dict[str, Any]:
+    """Extract method information for MCP tool definition."""
+    sig = inspect.signature(method)
+    doc = inspect.getdoc(method) or ""
+
+    # Get parameter types (skip 'self' for instance methods)
+    params = list(sig.parameters.values())
+    if params and params[0].name in ("self", "cls"):
+        params = params[1:]
+
+    # Extract input type
+    input_type = None
+    if params:
+        input_type = params[0].annotation
+
+    # Extract return type
+    return_type = sig.return_annotation
+
+    # Build parameter schema
+    parameters_schema = {}
+    if input_type and input_type != inspect._empty:
+        if inspect.isclass(input_type) and issubclass(input_type, BaseModel):
+            # Use Pydantic's built-in schema generation
+            parameters_schema = input_type.model_json_schema()
+        else:
+            parameters_schema = python_type_to_json_type(input_type)
+
+    # Build response schema
+    response_schema = {}
+    if return_type and return_type != inspect._empty:
+        if is_streaming_return(return_type):
+            # For streaming responses, extract the inner type
+            inner_type = get_args(return_type)[0]
+            response_schema = {
+                "type": "object",
+                "properties": {
+                    "stream": {
+                        "type": "array",
+                        "items": python_type_to_json_type(inner_type),
+                    }
+                },
+            }
+        elif inspect.isclass(return_type) and issubclass(return_type, BaseModel):
+            response_schema = return_type.model_json_schema()
+        else:
+            response_schema = python_type_to_json_type(return_type)
+
+    return {
+        "description": doc,
+        "parameters": parameters_schema,
+        "response": response_schema,
+        "is_streaming": is_streaming_return(return_type),
+    }

pydantic_rpc/mcp/exporter.py

@@ -0,0 +1,283 @@
+"""MCP exporter for pydantic-rpc services using the official MCP SDK."""
+
+import asyncio
+import inspect
+from collections.abc import Callable
+from typing import Any
+
+from pydantic import BaseModel
+
+try:
+    from mcp.server import InitializationOptions, Server
+    from mcp.server.sse import SseServerTransport
+    from mcp.server.stdio import stdio_server
+    from mcp.types import (
+        Content,
+        ServerCapabilities,
+        TextContent,
+        Tool,
+    )
+except ImportError:
+    raise ImportError("mcp is required for MCP support. Install with: pip install mcp")
+
+from .converter import extract_method_info
+
+
+class MCPExporter:
+    """Export pydantic-rpc services as MCP tools using the official MCP SDK."""
+
+    def __init__(
+        self,
+        service_obj: object,
+        name: str | None = None,
+        description: str | None = None,
+    ):
+        """Initialize MCPExporter with a service object.
+
+        Args:
+            service_obj: The service object containing RPC methods to export as MCP tools.
+            name: Name for the MCP server (defaults to service class name)
+            description: Description for the MCP server
+        """
+        self.service: object = service_obj
+        self.name: str = name or service_obj.__class__.__name__
+        self.description: str = (
+            description or f"MCP tools from {service_obj.__class__.__name__}"
+        )
+
+        # Create MCP Server instance
+        self.server: Server[Any] = Server(
+            self.name, version="1.0.0", instructions=self.description
+        )
+
+        # Store tools for later reference
+        self.tools: dict[str, tuple[Tool, Any]] = {}
+
+        # SSE transport instance (created lazily)
+        self._sse_transport: SseServerTransport | None = None
+
+        # Register handlers
+        self._register_handlers()
+
+        # Extract and store tools
+        self._extract_tools()
+
+    def _register_handlers(self):
+        """Register MCP protocol handlers."""
+
+        @self.server.list_tools()
+        async def handle_list_tools() -> list[Tool]:  # pyright:ignore[reportUnusedFunction]
+            """List all available tools."""
+            return [tool for tool, _ in self.tools.values()]
+
+        @self.server.call_tool()
+        async def handle_call_tool(  # pyright:ignore[reportUnusedFunction]
+            name: str, arguments: dict[str, Any] | None
+        ) -> list[Content]:
+            """Execute a tool."""
+            if name not in self.tools:
+                raise ValueError(f"Unknown tool: {name}")
+
+            _, method = self.tools[name]
+
+            # Extract arguments from the request
+            args = arguments or {}
+
+            # Call the method
+            if inspect.iscoroutinefunction(method):
+                result = await method(**args)
+            else:
+                result = method(**args)
+
+            # Convert result to TextContent
+            if isinstance(result, BaseModel):
+                content = result.model_dump_json(indent=2)
+            else:
+                content = str(result)
+
+            return [TextContent(type="text", text=content)]
+
+    def _extract_tools(self):
+        """Extract tools from the service object."""
+        for method_name, method in inspect.getmembers(self.service, inspect.ismethod):
+            # Skip private methods
+            if method_name.startswith("_"):
+                continue
+
+            # Exclude methods from external modules (like pytest fixtures)
+            method_module = inspect.getmodule(method)
+            if method_module and not method_module.__name__.startswith(
+                self.service.__class__.__module__
+            ):
+                continue
+
+            tool_name = method_name.lower()
+
+            # Get method info
+            try:
+                method_info = extract_method_info(method)
+            except Exception:
+                # Skip methods that can't be processed
+                continue
+
+            # Create Tool definition
+            tool = Tool(
+                name=tool_name,
+                description=method_info["description"] or f"Execute {method_name}",
+                inputSchema=method_info["parameters"],
+            )
+
+            # Create a wrapper that handles Pydantic model parameters
+            sig = inspect.signature(method)
+            params = list(sig.parameters.values())
+            if params and params[0].name in ("self", "cls"):
+                params = params[1:]
+
+            if params and params[0].annotation != inspect._empty:
+                param_type = params[0].annotation
+                if inspect.isclass(param_type) and issubclass(param_type, BaseModel):
+                    # For Pydantic models, create a wrapper that constructs the model
+                    if inspect.iscoroutinefunction(method):
+
+                        def make_async_wrapper(
+                            m: Callable[..., Any], pt: type[BaseModel]
+                        ) -> Callable[..., Any]:
+                            async def wrapped_method(**kwargs: Any) -> Any:
+                                request = pt(**kwargs)
+                                return await m(request)
+
+                            return wrapped_method
+
+                        self.tools[tool_name] = (
+                            tool,
+                            make_async_wrapper(method, param_type),
+                        )
+                    else:
+
+                        def make_sync_wrapper(
+                            m: Callable[..., Any], pt: type[BaseModel]
+                        ) -> Callable[..., Any]:
+                            def wrapped_method(**kwargs: Any) -> Any:
+                                request = pt(**kwargs)
+                                return m(request)
+
+                            return wrapped_method
+
+                        self.tools[tool_name] = (
+                            tool,
+                            make_sync_wrapper(method, param_type),
+                        )
+                else:
+                    # For non-Pydantic types, use the method directly
+                    self.tools[tool_name] = (tool, method)
+            else:
+                # No parameters
+                self.tools[tool_name] = (tool, method)
+
+    def run_stdio(self):
+        """Run the MCP server in stdio mode."""
+        asyncio.run(self._run_stdio())
+
+    async def _run_stdio(self):
+        """Async implementation of stdio server."""
+        async with stdio_server() as (read_stream, write_stream):
+            init_options = InitializationOptions(
+                server_name=self.name,
+                server_version="1.0.0",
+                capabilities=ServerCapabilities(),
+            )
+            await self.server.run(read_stream, write_stream, init_options)
+
+    def get_asgi_app(self, path: str = "/mcp"):
+        """Get the ASGI app for HTTP/SSE transport.
+
+        Args:
+            path: The base path for MCP endpoints (default: "/mcp")
+
+        Returns:
+            An ASGI application that can be mounted or run directly.
+        """
+        _ = path
+        try:
+            from starlette.applications import Starlette
+            from starlette.routing import Mount, Route
+        except ImportError:
+            raise ImportError(
+                "starlette is required for HTTP/SSE transport. "
+                "Install with: pip install starlette"
+            )
+
+        # Create SSE transport if not already created
+        if self._sse_transport is None:
+            self._sse_transport = SseServerTransport("/messages/")
+
+        # Get transport (guaranteed to be non-None after above check)
+        sse_transport = self._sse_transport
+
+        # Create SSE endpoint handler
+        async def handle_sse(request: Any) -> None:
+            # Use ASGI interface directly
+            scope = request.scope
+            receive = request.receive
+            send = request._send
+
+            async with sse_transport.connect_sse(scope, receive, send) as (
+                read_stream,
+                write_stream,
+            ):
+                init_options = InitializationOptions(
+                    server_name=self.name,
+                    server_version="1.0.0",
+                    capabilities=ServerCapabilities(),
+                )
+                await self.server.run(read_stream, write_stream, init_options)
+
+        # Create Starlette app with routes
+        app = Starlette(
+            routes=[
+                Route("/sse", endpoint=handle_sse, methods=["GET"]),
+                Mount("/messages/", app=sse_transport.handle_post_message),
+            ]
+        )
+
+        return app
+
+    def mount_to_asgi(self, asgi_app: Any, path: str = "/mcp"):
+        """Mount MCP endpoints to an existing ASGI application.
+
+        Works with Starlette/FastAPI applications and pydantic-rpc ASGI apps.
+
+        Args:
+            asgi_app: The ASGI application to mount to
+            path: The path prefix for MCP endpoints (default: "/mcp")
+        """
+        mcp_asgi = self.get_asgi_app(path)
+
+        # Check if the app has a mount method (Starlette/FastAPI)
+        if hasattr(asgi_app, "mount"):
+            asgi_app.mount(path, mcp_asgi)
+        # Check if it's a pydantic-rpc ASGI app (ASGIApp/ConnecpyASGIApp)
+        elif hasattr(asgi_app, "_app"):
+            original_app = asgi_app._app
+
+            async def wrapped_app(
+                scope: dict[str, Any],
+                receive: Callable[[], Any],
+                send: Callable[[Any], Any],
+            ) -> None:
+                if scope["type"] == "http" and scope["path"].startswith(path):
+                    # Create a new scope with adjusted path
+                    scope = dict(scope)
+                    scope["path"] = scope["path"][len(path) :]
+                    if not scope["path"]:
+                        scope["path"] = "/"
+                    await mcp_asgi(scope, receive, send)
+                else:
+                    await original_app(scope, receive, send)
+
+            asgi_app._app = wrapped_app
+        else:
+            raise ValueError(
+                "Unable to mount MCP to the provided ASGI app. "
+                "The app must have either a 'mount' method or '_app' attribute."
+            )