chuk-tool-processor 0.1.2__py3-none-any.whl → 0.1.4__py3-none-any.whl

chuk_tool_processor/__init__.py

@@ -1 +0,0 @@
-# chuk_tool_processor/__init__.py

chuk_tool_processor/mcp/__init__.py

@@ -0,0 +1,21 @@
+# chuk_tool_processor/mcp/__init__.py
+"""
+MCP integration for CHUK Tool Processor.
+"""
+from chuk_tool_processor.mcp.transport import MCPBaseTransport, StdioTransport, SSETransport
+from chuk_tool_processor.mcp.stream_manager import StreamManager
+from chuk_tool_processor.mcp.mcp_tool import MCPTool
+from chuk_tool_processor.mcp.register_mcp_tools import register_mcp_tools
+from chuk_tool_processor.mcp.setup_mcp_stdio import setup_mcp_stdio
+from chuk_tool_processor.mcp.setup_mcp_sse import setup_mcp_sse
+
+__all__ = [
+    "MCPBaseTransport",
+    "StdioTransport",
+    "SSETransport",
+    "StreamManager",
+    "MCPTool",
+    "register_mcp_tools",
+    "setup_mcp_stdio",
+    "setup_mcp_sse"
+]

chuk_tool_processor/mcp/mcp_tool.py

@@ -0,0 +1,53 @@
+# chuk_tool_processor/mcp/mcp_tool.py
+"""
+MCP tool that uses StreamManager for execution.
+"""
+
+from typing import Any
+
+from chuk_tool_processor.mcp.stream_manager import StreamManager
+from chuk_tool_processor.logging import get_logger
+
+logger = get_logger("chuk_tool_processor.mcp.mcp_tool")
+
+class MCPTool:
+    """
+    MCP tool that uses StreamManager for execution.
+    
+    This tool handles both namespaced and non-namespaced execution.
+    """
+    
+    def __init__(self, tool_name: str, stream_manager: StreamManager):
+        """
+        Initialize the MCP tool.
+        
+        Args:
+            tool_name: Name of the MCP tool
+            stream_manager: StreamManager instance
+        """
+        self.tool_name = tool_name
+        self.stream_manager = stream_manager
+        
+    async def execute(self, **kwargs: Any) -> Any:
+        """
+        Execute the tool using StreamManager.
+        
+        Args:
+            **kwargs: Tool arguments
+            
+        Returns:
+            Tool result
+        """
+        logger.debug(f"Executing MCP tool {self.tool_name}")
+        
+        result = await self.stream_manager.call_tool(
+            tool_name=self.tool_name,
+            arguments=kwargs
+        )
+        
+        if result.get("isError"):
+            error_msg = result.get("error", "Unknown error")
+            logger.error(f"Error executing MCP tool {self.tool_name}: {error_msg}")
+            raise RuntimeError(error_msg)
+            
+        return result.get("content")

chuk_tool_processor/mcp/register_mcp_tools.py

@@ -0,0 +1,82 @@
+# chuk_tool_processor/mcp/register_mcp_tools.py
+"""
+Registration functions for MCP tools.
+"""
+
+from typing import List, Dict, Any
+
+from chuk_tool_processor.mcp.mcp_tool import MCPTool
+from chuk_tool_processor.mcp.stream_manager import StreamManager
+from chuk_tool_processor.registry.provider import ToolRegistryProvider
+from chuk_tool_processor.logging import get_logger
+
+logger = get_logger("chuk_tool_processor.mcp.register")
+
+
+def register_mcp_tools(
+    stream_manager: StreamManager,
+    namespace: str = "mcp"
+) -> List[str]:
+    """
+    Register MCP tools with the CHUK registry.
+    
+    Args:
+        stream_manager: StreamManager instance
+        namespace: Namespace for the tools
+        
+    Returns:
+        List of registered tool names
+    """
+    registry = ToolRegistryProvider.get_registry()
+    registered_tools = []
+    
+    # Get all tools from StreamManager
+    mcp_tools = stream_manager.get_all_tools()
+    
+    for tool_def in mcp_tools:
+        tool_name = tool_def.get("name")
+        if not tool_name:
+            logger.warning("Tool definition missing name")
+            continue
+            
+        description = tool_def.get("description", f"MCP tool: {tool_name}")
+        
+        try:
+            # Create tool
+            tool = MCPTool(tool_name, stream_manager)
+            
+            # Register with registry under the original name in the given namespace
+            registry.register_tool(
+                tool,
+                name=tool_name,
+                namespace=namespace,
+                metadata={
+                    "description": description,
+                    "is_async": True,
+                    "tags": {"mcp", "remote"},
+                    "argument_schema": tool_def.get("inputSchema", {})
+                }
+            )
+            
+            # Also register the tool in the default namespace with the namespaced name
+            # This allows calling the tool as either "echo" or "stdio.echo" from parsers
+            namespaced_tool_name = f"{namespace}.{tool_name}"
+            registry.register_tool(
+                tool,
+                name=namespaced_tool_name,
+                namespace="default",
+                metadata={
+                    "description": description,
+                    "is_async": True,
+                    "tags": {"mcp", "remote", "namespaced"},
+                    "argument_schema": tool_def.get("inputSchema", {})
+                }
+            )
+            
+            registered_tools.append(tool_name)
+            logger.info(f"Registered MCP tool '{tool_name}' in namespace '{namespace}' (also as '{namespaced_tool_name}' in default)")
+            
+        except Exception as e:
+            logger.error(f"Error registering MCP tool '{tool_name}': {e}")
+    
+    return registered_tools

chuk_tool_processor/mcp/setup_mcp_sse.py

@@ -0,0 +1,74 @@
+# chuk_tool_processor/mcp/setup_mcp_sse.py
+"""
+Setup function for SSE transport MCP integration.
+"""
+from __future__ import annotations
+
+from typing import Dict, List, Optional
+
+from chuk_tool_processor.core.processor import ToolProcessor
+from chuk_tool_processor.mcp.stream_manager import StreamManager
+from chuk_tool_processor.mcp.register_mcp_tools import register_mcp_tools
+from chuk_tool_processor.logging import get_logger
+
+logger = get_logger("chuk_tool_processor.mcp.setup_sse")
+
+
+async def setup_mcp_sse(
+    servers: List[Dict[str, str]],
+    server_names: Optional[Dict[int, str]] = None,
+    default_timeout: float = 10.0,
+    max_concurrency: Optional[int] = None,
+    enable_caching: bool = True,
+    cache_ttl: int = 300,
+    enable_rate_limiting: bool = False,
+    global_rate_limit: Optional[int] = None,
+    tool_rate_limits: Optional[Dict[str, tuple]] = None,
+    enable_retries: bool = True,
+    max_retries: int = 3,
+    namespace: str = "mcp"
+) -> tuple[ToolProcessor, StreamManager]:
+    """
+    Set up MCP with SSE transport and CHUK Tool Processor.
+    
+    Args:
+        servers: List of server configurations with "name" and "url" keys
+        server_names: Optional mapping of server indices to names
+        default_timeout: Default timeout for tool execution
+        max_concurrency: Maximum concurrent executions
+        enable_caching: Whether to enable caching
+        cache_ttl: Cache TTL in seconds
+        enable_rate_limiting: Whether to enable rate limiting
+        global_rate_limit: Global rate limit (requests per minute)
+        tool_rate_limits: Per-tool rate limits
+        enable_retries: Whether to enable retries
+        max_retries: Maximum retry attempts
+        namespace: Namespace for MCP tools
+        
+    Returns:
+        Tuple of (processor, stream_manager)
+    """
+    # Create and initialize StreamManager with SSE transport
+    stream_manager = await StreamManager.create_with_sse(
+        servers=servers,
+        server_names=server_names
+    )
+    
+    # Register MCP tools
+    registered_tools = register_mcp_tools(stream_manager, namespace)
+    
+    # Create processor
+    processor = ToolProcessor(
+        default_timeout=default_timeout,
+        max_concurrency=max_concurrency,
+        enable_caching=enable_caching,
+        cache_ttl=cache_ttl,
+        enable_rate_limiting=enable_rate_limiting,
+        global_rate_limit=global_rate_limit,
+        tool_rate_limits=tool_rate_limits,
+        enable_retries=enable_retries,
+        max_retries=max_retries
+    )
+    
+    logger.info(f"Set up MCP (SSE) with {len(registered_tools)} tools")
+    return processor, stream_manager

chuk_tool_processor/mcp/setup_mcp_stdio.py

@@ -0,0 +1,78 @@
+# chuk_tool_processor/mcp/setup_mcp_stdio.py
+"""
+Setup function for stdio transport MCP integration.
+"""
+from __future__ import annotations
+
+from typing import Dict, List, Optional
+
+from chuk_tool_processor.core.processor import ToolProcessor
+from chuk_tool_processor.mcp.stream_manager import StreamManager
+from chuk_tool_processor.mcp.register_mcp_tools import register_mcp_tools
+from chuk_tool_processor.logging import get_logger
+
+logger = get_logger("chuk_tool_processor.mcp.setup_stdio")
+
+
+async def setup_mcp_stdio(
+    config_file: str,
+    servers: List[str],
+    server_names: Optional[Dict[int, str]] = None,
+    default_timeout: float = 10.0,
+    max_concurrency: Optional[int] = None,
+    enable_caching: bool = True,
+    cache_ttl: int = 300,
+    enable_rate_limiting: bool = False,
+    global_rate_limit: Optional[int] = None,
+    tool_rate_limits: Optional[Dict[str, tuple]] = None,
+    enable_retries: bool = True,
+    max_retries: int = 3,
+    namespace: str = "mcp"
+) -> tuple[ToolProcessor, StreamManager]:
+    """
+    Set up MCP with stdio transport and CHUK Tool Processor.
+    
+    Args:
+        config_file: Path to the config file
+        servers: List of server names to connect to
+        server_names: Optional mapping of server indices to names
+        default_timeout: Default timeout for tool execution
+        max_concurrency: Maximum concurrent executions
+        enable_caching: Whether to enable caching
+        cache_ttl: Cache TTL in seconds
+        enable_rate_limiting: Whether to enable rate limiting
+        global_rate_limit: Global rate limit (requests per minute)
+        tool_rate_limits: Per-tool rate limits
+        enable_retries: Whether to enable retries
+        max_retries: Maximum retry attempts
+        namespace: Namespace for MCP tools
+        
+    Returns:
+        Tuple of (processor, stream_manager)
+    """
+    # Create and initialize StreamManager with stdio transport
+    stream_manager = await StreamManager.create(
+        config_file=config_file,
+        servers=servers,
+        server_names=server_names,
+        transport_type="stdio"
+    )
+    
+    # Register MCP tools
+    registered_tools = register_mcp_tools(stream_manager, namespace)
+    
+    # Create processor
+    processor = ToolProcessor(
+        default_timeout=default_timeout,
+        max_concurrency=max_concurrency,
+        enable_caching=enable_caching,
+        cache_ttl=cache_ttl,
+        enable_rate_limiting=enable_rate_limiting,
+        global_rate_limit=global_rate_limit,
+        tool_rate_limits=tool_rate_limits,
+        enable_retries=enable_retries,
+        max_retries=max_retries
+    )
+    
+    logger.info(f"Set up MCP (stdio) with {len(registered_tools)} tools")
+    return processor, stream_manager

chuk_tool_processor/mcp/stream_manager.py

@@ -0,0 +1,293 @@
+# chuk_tool_processor/mcp/stream_manager.py
+"""
+StreamManager for CHUK Tool Processor.
+"""
+from __future__ import annotations
+
+import asyncio
+from typing import Any, Dict, List, Optional, Tuple
+
+# --------------------------------------------------------------------------- #
+#  CHUK imports                                                               #
+# --------------------------------------------------------------------------- #
+from chuk_mcp.config import load_config
+from chuk_tool_processor.mcp.transport import (
+    MCPBaseTransport,
+    StdioTransport,
+    SSETransport,
+)
+from chuk_tool_processor.logging import get_logger
+
+logger = get_logger("chuk_tool_processor.mcp.stream_manager")
+
+
+class StreamManager:
+    """
+    Manager for MCP server streams with support for multiple transport types.
+    """
+
+    # ------------------------------------------------------------------ #
+    #  construction                                                      #
+    # ------------------------------------------------------------------ #
+    def __init__(self) -> None:
+        self.transports: Dict[str, MCPBaseTransport] = {}
+        self.server_info: List[Dict[str, Any]] = []
+        self.tool_to_server_map: Dict[str, str] = {}
+        self.server_names: Dict[int, str] = {}
+        self.all_tools: List[Dict[str, Any]] = []
+        self._lock = asyncio.Lock()
+
+    # ------------------------------------------------------------------ #
+    #  factory helpers                                                   #
+    # ------------------------------------------------------------------ #
+    @classmethod
+    async def create(
+        cls,
+        config_file: str,
+        servers: List[str],
+        server_names: Optional[Dict[int, str]] = None,
+        transport_type: str = "stdio",
+    ) -> "StreamManager":
+        inst = cls()
+        await inst.initialize(config_file, servers, server_names, transport_type)
+        return inst
+
+    @classmethod
+    async def create_with_sse(
+        cls,
+        servers: List[Dict[str, str]],
+        server_names: Optional[Dict[int, str]] = None,
+    ) -> "StreamManager":
+        inst = cls()
+        await inst.initialize_with_sse(servers, server_names)
+        return inst
+
+    # ------------------------------------------------------------------ #
+    #  initialisation – stdio / sse                                      #
+    # ------------------------------------------------------------------ #
+    async def initialize(
+        self,
+        config_file: str,
+        servers: List[str],
+        server_names: Optional[Dict[int, str]] = None,
+        transport_type: str = "stdio",
+    ) -> None:
+        async with self._lock:
+            self.server_names = server_names or {}
+
+            for idx, server_name in enumerate(servers):
+                try:
+                    if transport_type == "stdio":
+                        params = await load_config(config_file, server_name)
+                        transport: MCPBaseTransport = StdioTransport(params)
+                    elif transport_type == "sse":
+                        transport = SSETransport("http://localhost:8000")
+                    else:
+                        logger.error("Unsupported transport type: %s", transport_type)
+                        continue
+
+                    if not await transport.initialize():
+                        logger.error("Failed to init %s", server_name)
+                        continue
+
+                    #  store transport
+                    self.transports[server_name] = transport
+
+                    #  ping + gather tools
+                    status = "Up" if await transport.send_ping() else "Down"
+                    tools = await transport.get_tools()
+
+                    for t in tools:
+                        name = t.get("name")
+                        if name:
+                            self.tool_to_server_map[name] = server_name
+                    self.all_tools.extend(tools)
+
+                    self.server_info.append(
+                        {
+                            "id": idx,
+                            "name": server_name,
+                            "tools": len(tools),
+                            "status": status,
+                        }
+                    )
+                    logger.info("Initialised %s – %d tool(s)", server_name, len(tools))
+                except Exception as exc:  # noqa: BLE001
+                    logger.error("Error initialising %s: %s", server_name, exc)
+
+            logger.info(
+                "StreamManager ready – %d server(s), %d tool(s)",
+                len(self.transports),
+                len(self.all_tools),
+            )
+
+    async def initialize_with_sse(
+        self,
+        servers: List[Dict[str, str]],
+        server_names: Optional[Dict[int, str]] = None,
+    ) -> None:
+        async with self._lock:
+            self.server_names = server_names or {}
+
+            for idx, cfg in enumerate(servers):
+                name, url = cfg.get("name"), cfg.get("url")
+                if not (name and url):
+                    logger.error("Bad server config: %s", cfg)
+                    continue
+                try:
+                    transport = SSETransport(url, cfg.get("api_key"))
+                    if not await transport.initialize():
+                        logger.error("Failed to init SSE %s", name)
+                        continue
+
+                    self.transports[name] = transport
+                    status = "Up" if await transport.send_ping() else "Down"
+                    tools = await transport.get_tools()
+
+                    for t in tools:
+                        tname = t.get("name")
+                        if tname:
+                            self.tool_to_server_map[tname] = name
+                    self.all_tools.extend(tools)
+
+                    self.server_info.append(
+                        {"id": idx, "name": name, "tools": len(tools), "status": status}
+                    )
+                    logger.info("Initialised SSE %s – %d tool(s)", name, len(tools))
+                except Exception as exc:  # noqa: BLE001
+                    logger.error("Error initialising SSE %s: %s", name, exc)
+
+            logger.info(
+                "StreamManager ready – %d SSE server(s), %d tool(s)",
+                len(self.transports),
+                len(self.all_tools),
+            )
+
+    # ------------------------------------------------------------------ #
+    #  queries                                                           #
+    # ------------------------------------------------------------------ #
+    def get_all_tools(self) -> List[Dict[str, Any]]:
+        return self.all_tools
+
+    def get_server_for_tool(self, tool_name: str) -> Optional[str]:
+        return self.tool_to_server_map.get(tool_name)
+
+    def get_server_info(self) -> List[Dict[str, Any]]:
+        return self.server_info
+
+    # ------------------------------------------------------------------ #
+    #  EXTRA HELPERS – ping / resources / prompts                        #
+    # ------------------------------------------------------------------ #
+    async def ping_servers(self) -> List[Dict[str, Any]]:
+        async def _ping_one(name: str, tr: MCPBaseTransport):
+            try:
+                ok = await tr.send_ping()
+            except Exception:  # pragma: no cover
+                ok = False
+            return {"server": name, "ok": ok}
+
+        return await asyncio.gather(*(_ping_one(n, t) for n, t in self.transports.items()))
+
+    async def list_resources(self) -> List[Dict[str, Any]]:
+        out: List[Dict[str, Any]] = []
+
+        async def _one(name: str, tr: MCPBaseTransport):
+            if not hasattr(tr, "list_resources"):
+                return
+            try:
+                res = await tr.list_resources()  # type: ignore[attr-defined]
+                # accept either {"resources": [...]} **or** a plain list
+                resources = (
+                    res.get("resources", []) if isinstance(res, dict) else res
+                )
+                for item in resources:
+                    item = dict(item)
+                    item["server"] = name
+                    out.append(item)
+            except Exception as exc:
+                logger.debug("resources/list failed for %s: %s", name, exc)
+
+        await asyncio.gather(*(_one(n, t) for n, t in self.transports.items()))
+        return out
+
+    async def list_prompts(self) -> List[Dict[str, Any]]:
+        out: List[Dict[str, Any]] = []
+
+        async def _one(name: str, tr: MCPBaseTransport):
+            if not hasattr(tr, "list_prompts"):
+                return
+            try:
+                res = await tr.list_prompts()  # type: ignore[attr-defined]
+                prompts = res.get("prompts", []) if isinstance(res, dict) else res
+                for item in prompts:
+                    item = dict(item)
+                    item["server"] = name
+                    out.append(item)
+            except Exception as exc:
+                logger.debug("prompts/list failed for %s: %s", name, exc)
+
+        await asyncio.gather(*(_one(n, t) for n, t in self.transports.items()))
+        return out
+
+    # ------------------------------------------------------------------ #
+    #  tool execution                                                    #
+    # ------------------------------------------------------------------ #
+    async def call_tool(
+        self,
+        tool_name: str,
+        arguments: Dict[str, Any],
+        server_name: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        server_name = server_name or self.get_server_for_tool(tool_name)
+        if not server_name or server_name not in self.transports:
+            # wording kept exactly for unit-test expectation
+            return {
+                "isError": True,
+                "error": f"No server found for tool: {tool_name}",
+            }
+        return await self.transports[server_name].call_tool(tool_name, arguments)
+
+    # ------------------------------------------------------------------ #
+    #  shutdown                                                          #
+    # ------------------------------------------------------------------ #
+    async def close(self) -> None:
+        tasks = [tr.close() for tr in self.transports.values()]
+        if tasks:
+            try:
+                await asyncio.gather(*tasks)
+            except asyncio.CancelledError:  # pragma: no cover
+                pass
+            except Exception as exc:  # noqa: BLE001
+                logger.error("Error during close: %s", exc)
+
+        self.transports.clear()
+        self.server_info.clear()
+        self.tool_to_server_map.clear()
+        self.all_tools.clear()
+
+    # ------------------------------------------------------------------ #
+    #  backwards-compat: streams helper                                  #
+    # ------------------------------------------------------------------ #
+    def get_streams(self) -> List[Tuple[Any, Any]]:
+        """
+        Return a list of ``(read_stream, write_stream)`` tuples for **all**
+        transports.  Older CLI commands rely on this helper.
+        """
+        pairs: List[Tuple[Any, Any]] = []
+
+        for tr in self.transports.values():
+            if hasattr(tr, "get_streams") and callable(tr.get_streams):
+                pairs.extend(tr.get_streams())  # type: ignore[arg-type]
+                continue
+
+            rd = getattr(tr, "read_stream", None)
+            wr = getattr(tr, "write_stream", None)
+            if rd and wr:
+                pairs.append((rd, wr))
+
+        return pairs
+
+    # convenience alias
+    @property
+    def streams(self) -> List[Tuple[Any, Any]]:  # pragma: no cover
+        return self.get_streams()

chuk_tool_processor/mcp/transport/__init__.py

@@ -0,0 +1,14 @@
+# chuk_tool_processor/mcp/transport/__init__.py
+"""
+MCP transport implementations.
+"""
+
+from .base_transport import MCPBaseTransport
+from .stdio_transport import StdioTransport
+from .sse_transport import SSETransport
+
+__all__ = [
+    "MCPBaseTransport",
+    "StdioTransport",
+    "SSETransport"
+]