kailash 0.1.5__py3-none-any.whl → 0.2.1__py3-none-any.whl

kailash/mcp/client_new.py

@@ -0,0 +1,334 @@
+"""
+MCP Client - Clean implementation using official Anthropic MCP SDK.
+
+This is NOT a node - it's a utility class used by LLM agents to interact with MCP servers.
+"""
+
+import logging
+import os
+from typing import Any, Dict, List, Optional, Union
+
+# Will use official MCP SDK when available
+try:
+    from mcp import ClientSession, StdioServerParameters
+    from mcp.client.stdio import stdio_client
+
+    MCP_AVAILABLE = True
+except ImportError:
+    MCP_AVAILABLE = False
+
+
+class MCPClient:
+    """
+    Clean MCP client for connecting to Model Context Protocol servers.
+
+    This client is used internally by LLMAgentNode and other components
+    that need MCP capabilities. It provides a simple interface over the
+    official MCP SDK.
+
+    Examples:
+        >>> client = MCPClient()
+        >>>
+        >>> # Connect to a server
+        >>> async with client.connect_stdio("python", ["-m", "my_mcp_server"]) as session:
+        ...     # Discover available tools
+        ...     tools = await client.discover_tools(session)
+        ...
+        ...     # Call a tool
+        ...     result = await client.call_tool(
+        ...         session,
+        ...         "search",
+        ...         {"query": "AI in healthcare"}
+        ...     )
+    """
+
+    def __init__(self):
+        """Initialize MCP client."""
+        self.logger = logging.getLogger(__name__)
+
+    async def connect_stdio(
+        self, command: str, args: List[str], env: Optional[Dict[str, str]] = None
+    ):
+        """
+        Connect to an MCP server via stdio transport.
+
+        Args:
+            command: Command to run the server
+            args: Arguments for the command
+            env: Environment variables
+
+        Returns:
+            Context manager yielding ClientSession
+
+        Raises:
+            ImportError: If MCP SDK is not available
+            RuntimeError: If connection fails
+        """
+        if not MCP_AVAILABLE:
+            raise ImportError("MCP SDK not available. Install with: pip install mcp")
+
+        # Merge environment
+        server_env = os.environ.copy()
+        if env:
+            server_env.update(env)
+
+        # Create server parameters
+        server_params = StdioServerParameters(
+            command=command, args=args, env=server_env
+        )
+
+        try:
+            # Connect to server
+            async with stdio_client(server_params) as (read_stream, write_stream):
+                async with ClientSession(read_stream, write_stream) as session:
+                    # Initialize session
+                    await session.initialize()
+                    yield session
+        except Exception as e:
+            self.logger.error(f"Failed to connect to MCP server: {e}")
+            raise RuntimeError(f"MCP connection failed: {e}")
+
+    async def discover_tools(self, session: "ClientSession") -> List[Dict[str, Any]]:
+        """
+        Discover available tools from an MCP server.
+
+        Args:
+            session: Active MCP client session
+
+        Returns:
+            List of tool definitions with name, description, and schema
+        """
+        try:
+            result = await session.list_tools()
+            tools = []
+
+            for tool in result.tools:
+                tools.append(
+                    {
+                        "name": tool.name,
+                        "description": tool.description,
+                        "inputSchema": tool.inputSchema,
+                    }
+                )
+
+            return tools
+
+        except Exception as e:
+            self.logger.error(f"Failed to discover tools: {e}")
+            return []
+
+    async def call_tool(
+        self, session: "ClientSession", name: str, arguments: Dict[str, Any]
+    ) -> Any:
+        """
+        Call a tool on the MCP server.
+
+        Args:
+            session: Active MCP client session
+            name: Tool name
+            arguments: Tool arguments
+
+        Returns:
+            Tool execution result
+        """
+        try:
+            result = await session.call_tool(name=name, arguments=arguments)
+
+            # Extract content from result
+            if hasattr(result, "content"):
+                content = []
+                for item in result.content:
+                    if hasattr(item, "text"):
+                        content.append({"type": "text", "text": item.text})
+                    else:
+                        content.append(str(item))
+                return content
+            else:
+                return str(result)
+
+        except Exception as e:
+            self.logger.error(f"Failed to call tool '{name}': {e}")
+            raise
+
+    async def list_resources(self, session: "ClientSession") -> List[Dict[str, Any]]:
+        """
+        List available resources from an MCP server.
+
+        Args:
+            session: Active MCP client session
+
+        Returns:
+            List of resource definitions
+        """
+        try:
+            result = await session.list_resources()
+            resources = []
+
+            for resource in result.resources:
+                resources.append(
+                    {
+                        "uri": resource.uri,
+                        "name": resource.name,
+                        "description": resource.description,
+                        "mimeType": resource.mimeType,
+                    }
+                )
+
+            return resources
+
+        except Exception as e:
+            self.logger.error(f"Failed to list resources: {e}")
+            return []
+
+    async def read_resource(self, session: "ClientSession", uri: str) -> Any:
+        """
+        Read a specific resource from an MCP server.
+
+        Args:
+            session: Active MCP client session
+            uri: Resource URI
+
+        Returns:
+            Resource content
+        """
+        try:
+            result = await session.read_resource(uri=uri)
+
+            # Extract content
+            if hasattr(result, "contents"):
+                content = []
+                for item in result.contents:
+                    if hasattr(item, "text"):
+                        content.append({"type": "text", "text": item.text})
+                    elif hasattr(item, "blob"):
+                        content.append({"type": "blob", "data": item.blob})
+                    else:
+                        content.append(str(item))
+                return content
+            else:
+                return str(result)
+
+        except Exception as e:
+            self.logger.error(f"Failed to read resource '{uri}': {e}")
+            raise
+
+    async def list_prompts(self, session: "ClientSession") -> List[Dict[str, Any]]:
+        """
+        List available prompts from an MCP server.
+
+        Args:
+            session: Active MCP client session
+
+        Returns:
+            List of prompt definitions
+        """
+        try:
+            result = await session.list_prompts()
+            prompts = []
+
+            for prompt in result.prompts:
+                prompt_dict = {
+                    "name": prompt.name,
+                    "description": prompt.description,
+                    "arguments": [],
+                }
+
+                if hasattr(prompt, "arguments"):
+                    for arg in prompt.arguments:
+                        prompt_dict["arguments"].append(
+                            {
+                                "name": arg.name,
+                                "description": arg.description,
+                                "required": arg.required,
+                            }
+                        )
+
+                prompts.append(prompt_dict)
+
+            return prompts
+
+        except Exception as e:
+            self.logger.error(f"Failed to list prompts: {e}")
+            return []
+
+    async def get_prompt(
+        self, session: "ClientSession", name: str, arguments: Dict[str, Any]
+    ) -> Dict[str, Any]:
+        """
+        Get a prompt from an MCP server.
+
+        Args:
+            session: Active MCP client session
+            name: Prompt name
+            arguments: Prompt arguments
+
+        Returns:
+            Prompt with messages
+        """
+        try:
+            result = await session.get_prompt(name=name, arguments=arguments)
+
+            # Extract messages
+            messages = []
+            if hasattr(result, "messages"):
+                for msg in result.messages:
+                    messages.append(
+                        {
+                            "role": msg.role,
+                            "content": (
+                                msg.content.text
+                                if hasattr(msg.content, "text")
+                                else str(msg.content)
+                            ),
+                        }
+                    )
+
+            return {"name": name, "messages": messages, "arguments": arguments}
+
+        except Exception as e:
+            self.logger.error(f"Failed to get prompt '{name}': {e}")
+            raise
+
+
+# Convenience functions for LLM agents
+async def discover_and_prepare_tools(
+    mcp_servers: List[Union[str, Dict[str, Any]]]
+) -> List[Dict[str, Any]]:
+    """
+    Discover tools from multiple MCP servers and prepare them for LLM use.
+
+    Args:
+        mcp_servers: List of server URLs or configurations
+
+    Returns:
+        List of tool definitions ready for LLM function calling
+    """
+    client = MCPClient()
+    all_tools = []
+
+    for server in mcp_servers:
+        try:
+            # Parse server configuration
+            if isinstance(server, str):
+                # Simple URL format - not supported in stdio
+                continue
+
+            if server.get("transport") == "stdio":
+                command = server.get("command", "python")
+                args = server.get("args", [])
+                env = server.get("env", {})
+
+                async with client.connect_stdio(command, args, env) as session:
+                    tools = await client.discover_tools(session)
+
+                    # Tag tools with server info
+                    for tool in tools:
+                        tool["server"] = server.get("name", "mcp_server")
+                        tool["server_config"] = server
+
+                    all_tools.extend(tools)
+
+        except Exception as e:
+            logging.warning(f"Failed to discover tools from {server}: {e}")
+
+    return all_tools

kailash/mcp/server.py

@@ -0,0 +1,293 @@
+"""MCP Server Framework using official Anthropic SDK.
+
+This module provides a comprehensive framework for creating MCP servers using
+the official FastMCP framework from Anthropic. Servers run as long-lived
+services that expose tools, resources, and prompts to MCP clients, enabling
+dynamic capability extension for AI workflows.
+
+Note:
+    This module requires the FastMCP framework to be installed.
+    Install with: pip install 'mcp[server]'
+
+Examples:
+    Basic server with tools:
+
+    >>> from kailash.mcp.server import MCPServer
+    >>> class MyServer(MCPServer):
+    ...     def setup(self):
+    ...         @self.add_tool()
+    ...         def calculate(a: int, b: int) -> int:
+    ...             return a + b
+    >>> server = MyServer("calculator", port=8080)
+    >>> server.start()
+
+    Quick server creation:
+
+    >>> from kailash.mcp.server import SimpleMCPServer
+    >>> server = SimpleMCPServer("my-tools")
+    >>> @server.tool()
+    ... def search(query: str) -> list:
+    ...     return [f"Result for {query}"]
+    >>> server.start()
+"""
+
+import logging
+from abc import ABC, abstractmethod
+from typing import Callable
+
+logger = logging.getLogger(__name__)
+
+
+class MCPServer(ABC):
+    """Base class for MCP servers using FastMCP.
+
+    This provides a framework for creating MCP servers that expose
+    tools, resources, and prompts via the Model Context Protocol.
+
+    Examples:
+        Creating a custom server:
+
+        >>> class MyServer(MCPServer):
+        ...     def setup(self):
+        ...         @self.add_tool()
+        ...         def search(query: str) -> str:
+        ...             return f"Results for: {query}"
+        ...         @self.add_resource("data://example")
+        ...         def get_example():
+        ...             return "Example data"
+        >>> server = MyServer("my-server", port=8080)
+        >>> server.start()  # Runs until stopped
+    """
+
+    def __init__(self, name: str, port: int = 8080, host: str = "localhost"):
+        """Initialize the MCP server.
+
+        Args:
+            name: Name of the server.
+            port: Port to listen on (default: 8080).
+            host: Host to bind to (default: "localhost").
+        """
+        self.name = name
+        self.port = port
+        self.host = host
+        self._mcp = None
+        self._running = False
+
+    @abstractmethod
+    def setup(self):
+        """Setup server tools, resources, and prompts.
+
+        This method should be implemented by subclasses to define
+        the server's capabilities using decorators.
+
+        Note:
+            Use @self.add_tool(), @self.add_resource(uri), and
+            @self.add_prompt(name) decorators to register capabilities.
+        """
+        pass
+
+    def add_tool(self):
+        """Decorator to add a tool to the server.
+
+        Returns:
+            Function decorator for registering tools.
+
+        Examples:
+            >>> @server.add_tool()
+            ... def calculate(a: int, b: int) -> int:
+            ...     '''Add two numbers'''
+            ...     return a + b
+        """
+
+        def decorator(func: Callable):
+            if self._mcp is None:
+                self._init_mcp()
+
+            # Use FastMCP's tool decorator
+            return self._mcp.tool()(func)
+
+        return decorator
+
+    def add_resource(self, uri: str):
+        """Decorator to add a resource to the server.
+
+        Args:
+            uri: URI pattern for the resource (supports wildcards).
+
+        Returns:
+            Function decorator for registering resources.
+
+        Examples:
+            >>> @server.add_resource("file:///data/*")
+            ... def get_file(path: str) -> str:
+            ...     return f"Content of {path}"
+        """
+
+        def decorator(func: Callable):
+            if self._mcp is None:
+                self._init_mcp()
+
+            # Use FastMCP's resource decorator
+            return self._mcp.resource(uri)(func)
+
+        return decorator
+
+    def add_prompt(self, name: str):
+        """Decorator to add a prompt template to the server.
+
+        Args:
+            name: Name of the prompt.
+
+        Returns:
+            Function decorator for registering prompts.
+
+        Examples:
+            >>> @server.add_prompt("analyze")
+            ... def analyze_prompt(data: str) -> str:
+            ...     return f"Please analyze the following data: {data}"
+        """
+
+        def decorator(func: Callable):
+            if self._mcp is None:
+                self._init_mcp()
+
+            # Use FastMCP's prompt decorator
+            return self._mcp.prompt(name)(func)
+
+        return decorator
+
+    def _init_mcp(self):
+        """Initialize the FastMCP instance."""
+        try:
+            from mcp.server.fastmcp import FastMCP
+
+            self._mcp = FastMCP(self.name)
+        except ImportError:
+            logger.error(
+                "FastMCP not available. Install with: pip install 'mcp[server]'"
+            )
+            raise
+
+    def start(self):
+        """Start the MCP server.
+
+        This runs the server as a long-lived process until stopped.
+
+        Raises:
+            ImportError: If FastMCP is not available.
+            Exception: If server fails to start.
+        """
+        if self._mcp is None:
+            self._init_mcp()
+
+        # Run setup to register tools/resources
+        self.setup()
+
+        logger.info(f"Starting MCP server '{self.name}' on {self.host}:{self.port}")
+        self._running = True
+
+        try:
+            # Run the FastMCP server
+            logger.info("Running FastMCP server in stdio mode")
+            self._mcp.run()
+        except Exception as e:
+            logger.error(f"Failed to start server: {e}")
+            raise
+        finally:
+            self._running = False
+
+    def stop(self):
+        """Stop the MCP server."""
+        logger.info(f"Stopping MCP server '{self.name}'")
+        self._running = False
+        # In a real implementation, we'd need to handle graceful shutdown
+
+
+class SimpleMCPServer(MCPServer):
+    """Simple MCP server for basic use cases.
+
+    This provides an easy way to create MCP servers without subclassing.
+
+    Examples:
+        >>> server = SimpleMCPServer("my-server")
+        >>> @server.tool()
+        ... def add(a: int, b: int) -> int:
+        ...     return a + b
+        >>> server.start()
+    """
+
+    def __init__(self, name: str, port: int = 8080, host: str = "localhost"):
+        """Initialize the simple MCP server.
+
+        Args:
+            name: Name of the server.
+            port: Port to listen on (default: 8080).
+            host: Host to bind to (default: "localhost").
+        """
+        super().__init__(name, port, host)
+        self._tools = []
+        self._resources = []
+        self._prompts = []
+
+    def tool(self):
+        """Decorator to add a tool.
+
+        Returns:
+            Function decorator for registering tools.
+        """
+
+        def decorator(func):
+            self._tools.append(func)
+            return func
+
+        return decorator
+
+    def resource(self, uri: str):
+        """Decorator to add a resource.
+
+        Args:
+            uri: URI pattern for the resource.
+
+        Returns:
+            Function decorator for registering resources.
+        """
+
+        def decorator(func):
+            self._resources.append((uri, func))
+            return func
+
+        return decorator
+
+    def prompt(self, name: str):
+        """Decorator to add a prompt.
+
+        Args:
+            name: Name of the prompt.
+
+        Returns:
+            Function decorator for registering prompts.
+        """
+
+        def decorator(func):
+            self._prompts.append((name, func))
+            return func
+
+        return decorator
+
+    def setup(self):
+        """Setup the server with registered components.
+
+        Registers all tools, resources, and prompts that were decorated
+        before calling start().
+        """
+        # Register all tools
+        for tool_func in self._tools:
+            self.add_tool()(tool_func)
+
+        # Register all resources
+        for uri, resource_func in self._resources:
+            self.add_resource(uri)(resource_func)
+
+        # Register all prompts
+        for name, prompt_func in self._prompts:
+            self.add_prompt(name)(prompt_func)