kailash 0.6.1__py3-none-any.whl → 0.6.3__py3-none-any.whl

kailash/__init__.py

@@ -33,7 +33,7 @@ except ImportError:
 # For backward compatibility
 WorkflowGraph = Workflow
 
-__version__ = "0.6.1"
+__version__ = "0.6.3"
 
 __all__ = [
     # Core workflow components

kailash/core/actors/connection_actor.py

@@ -47,7 +47,7 @@ class Message:
     type: MessageType = MessageType.QUERY
     payload: Any = None
     reply_to: Optional[asyncio.Queue] = None
-    timestamp: datetime = field(default_factory=datetime.utcnow)
+    timestamp: datetime = field(default_factory=lambda: datetime.now(UTC))
 
 
 @dataclass
@@ -70,8 +70,8 @@ class ConnectionStats:
     total_execution_time: float = 0.0
     health_checks_passed: int = 0
     health_checks_failed: int = 0
-    created_at: datetime = field(default_factory=datetime.utcnow)
-    last_used_at: datetime = field(default_factory=datetime.utcnow)
+    created_at: datetime = field(default_factory=lambda: datetime.now(UTC))
+    last_used_at: datetime = field(default_factory=lambda: datetime.now(UTC))
     health_score: float = 100.0
 
 

kailash/gateway/api.py

@@ -11,7 +11,7 @@ from typing import Any, Dict, List, Optional, Union
 
 from fastapi import APIRouter, BackgroundTasks, Depends, FastAPI, HTTPException
 from fastapi.responses import JSONResponse
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, ConfigDict, Field
 
 from ..resources.registry import ResourceRegistry
 from .enhanced_gateway import (
@@ -37,14 +37,15 @@ class ResourceReferenceModel(BaseModel):
         None, description="Reference to credentials secret"
     )
 
-    class Config:
-        schema_extra = {
+    model_config = ConfigDict(
+        json_schema_extra={
             "example": {
                 "type": "database",
                 "config": {"host": "localhost", "port": 5432, "database": "myapp"},
                 "credentials_ref": "db_credentials",
             }
         }
+    )
 
 
 class WorkflowRequestModel(BaseModel):
@@ -59,8 +60,8 @@ class WorkflowRequestModel(BaseModel):
         None, description="Additional context variables"
     )
 
-    class Config:
-        schema_extra = {
+    model_config = ConfigDict(
+        json_schema_extra={
             "example": {
                 "inputs": {"user_id": 123, "action": "process"},
                 "resources": {
@@ -74,6 +75,7 @@ class WorkflowRequestModel(BaseModel):
                 "context": {"environment": "production", "trace_id": "abc123"},
             }
         }
+    )
 
 
 class WorkflowResponseModel(BaseModel):

kailash/gateway/enhanced_gateway.py

@@ -40,7 +40,7 @@ class WorkflowRequest:
     inputs: Dict[str, Any] = field(default_factory=dict)
     resources: Dict[str, Union[str, ResourceReference]] = field(default_factory=dict)
     context: Dict[str, Any] = field(default_factory=dict)
-    timestamp: datetime = field(default_factory=datetime.utcnow)
+    timestamp: datetime = field(default_factory=lambda: datetime.now(UTC))
 
     def to_dict(self) -> Dict[str, Any]:
         """Convert to JSON-serializable dict."""

kailash/{mcp → mcp_server}/__init__.py

@@ -12,8 +12,8 @@ Design Philosophy:
 
 Key Components:
     - MCPClient: Connects to MCP servers for tool and resource access
-    - MCPServer: Base framework for creating custom MCP servers
-    - SimpleMCPServer: Quick server creation for basic use cases
+    - MCPServer: Main production-ready server with all features
+    - MCPServerBase: Abstract base class for custom server implementations
 
 Upstream Dependencies:
     - Official Anthropic MCP Python SDK for protocol implementation
@@ -28,26 +28,31 @@ Downstream Consumers:
 Examples:
     Basic MCP client usage:
 
-    >>> from kailash.mcp import MCPClient
+    >>> from kailash.mcp_server import MCPClient
     >>> client = MCPClient()
     >>> tools = await client.discover_tools(server_config)
     >>> result = await client.call_tool(server_config, "search", {"query": "AI"})
 
     Simple MCP server creation:
 
-    >>> from kailash.mcp import SimpleMCPServer
-    >>> server = SimpleMCPServer("my-tools")
+    >>> from kailash.mcp_server import MCPServer
+    >>> server = MCPServer("my-tools")
     >>> @server.tool()
     ... def calculate(a: int, b: int) -> int:
     ...     return a + b
-    >>> server.start()
+    >>> server.run()
 """
 
 from .client import MCPClient
-from .server_enhanced import MCPServer, SimpleMCPServer
+
+# For backward compatibility
+from .server import EnhancedMCPServer, MCPServer, MCPServerBase, SimpleMCPServer
 
 __all__ = [
     "MCPClient",
     "MCPServer",
+    "MCPServerBase",
+    # Backward compatibility
     "SimpleMCPServer",
+    "EnhancedMCPServer",
 ]

kailash/{mcp → mcp_server}/ai_registry_server.py

@@ -5,7 +5,7 @@ AI Registry MCP Server using Anthropic's Official MCP Python SDK.
 This creates a real MCP server that exposes AI Registry tools following
 the actual Model Context Protocol specification.
 
-Run as: python -m kailash.mcp.ai_registry_server
+Run as: python -m kailash.mcp_server.ai_registry_server
 """
 
 import asyncio
@@ -708,5 +708,5 @@ if __name__ == "__main__":
 
 # For module execution
 def run_server():
-    """Entry point for python -m kailash.mcp.ai_registry_server"""
+    """Entry point for python -m kailash.mcp_server.ai_registry_server"""
     asyncio.run(main())

kailash/{mcp/server_enhanced.py → mcp_server/server.py}

@@ -1,16 +1,39 @@
 """
-Enhanced MCP Server with production-ready capabilities.
-
-This module provides an enhanced MCP server that includes caching, configuration,
-metrics, and other production features by default, while maintaining compatibility
-with the official Anthropic FastMCP framework.
+MCP Server Framework with production-ready capabilities.
+
+This module provides both basic and enhanced MCP server implementations using
+the official FastMCP framework from Anthropic. Servers run as long-lived
+services that expose tools, resources, and prompts to MCP clients.
+
+Basic Usage:
+    Abstract base class for custom servers:
+
+    >>> class MyServer(MCPServerBase):
+    ...     def setup(self):
+    ...         @self.add_tool()
+    ...         def calculate(a: int, b: int) -> int:
+    ...             return a + b
+    >>> server = MyServer("calculator")
+    >>> server.start()
+
+Production Usage:
+    Main server with all production features:
+
+    >>> from kailash.mcp_server import MCPServer
+    >>> server = MCPServer("my-server", enable_cache=True)
+    >>> @server.tool(cache_key="search", cache_ttl=600)
+    ... def search(query: str) -> dict:
+    ...     return {"results": f"Found data for {query}"}
+    >>> server.run()
 """
 
 import asyncio
 import functools
 import logging
+from abc import ABC, abstractmethod
+from collections.abc import Callable
 from pathlib import Path
-from typing import Any, Callable, Dict, Optional, TypeVar, Union
+from typing import Any, Dict, Optional, TypeVar, Union
 
 from .utils import CacheManager, ConfigManager, MetricsCollector, format_response
 
@@ -19,18 +42,199 @@ logger = logging.getLogger(__name__)
 F = TypeVar("F", bound=Callable[..., Any])
 
 
+class MCPServerBase(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(MCPServerBase):
+        ...     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.
+        """
+
+    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 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 EnhancedMCPServer:
     """
-    Production-ready MCP server with enhanced capabilities.
+    Production-ready MCP server (available as SimpleMCPServer).
 
-    Features included by default:
-    - Caching with TTL support
+    This is the main concrete MCP server implementation with all production
+    features available. Features can be enabled/disabled as needed.
+
+    Features available:
+    - Caching with TTL support (enable_cache=True)
+    - Metrics collection and monitoring (enable_metrics=True)
+    - Response formatting utilities (enable_formatting=True)
     - Hierarchical configuration management
-    - Metrics collection and monitoring
-    - Response formatting utilities
     - Error handling and logging
 
-    All features can be disabled if not needed.
+    Examples:
+        Basic usage (recommended):
+        >>> from kailash.mcp_server import MCPServer
+        >>> server = MCPServer("my-server")
+        >>> @server.tool()
+        ... def search(query: str) -> dict:
+        ...     return {"results": f"Found: {query}"}
+        >>> server.run()
+
+        With production features enabled:
+        >>> server = MCPServer("my-server", enable_cache=True, enable_metrics=True)
+        >>> @server.tool(cache_key="search", cache_ttl=600)
+        ... def search(query: str) -> dict:
+        ...     return {"results": f"Found: {query}"}
+        >>> server.run()
     """
 
     def __init__(
@@ -104,15 +308,21 @@ class EnhancedMCPServer:
             return
 
         try:
-            from mcp.server.fastmcp import FastMCP
+            # Now we can safely import from external mcp.server (no namespace collision)
+            from mcp.server import FastMCP
 
             self._mcp = FastMCP(self.name)
             logger.info(f"Initialized FastMCP server: {self.name}")
-        except ImportError:
+        except ImportError as e:
+            logger.error(
+                f"FastMCP import failed with: {e}. Details: {type(e).__name__}"
+            )
             logger.error(
                 "FastMCP not available. Install with: pip install 'mcp[server]'"
             )
-            raise
+            raise ImportError(
+                "FastMCP not available. Install with: pip install 'mcp[server]'"
+            ) from e
 
     def tool(
         self,
@@ -413,37 +623,10 @@ class EnhancedMCPServer:
             self._running = False
 
 
-# For backward compatibility, make EnhancedMCPServer the default MCPServer
+# Clean public API design:
+# - MCPServerBase: Abstract base for custom implementations (e.g., AIRegistryServer)
+# - MCPServer: Main concrete server with all production features
+# - SimpleMCPServer: Alias for backward compatibility
+# - EnhancedMCPServer: Alias for backward compatibility
 MCPServer = EnhancedMCPServer
-
-
-class SimpleMCPServer(EnhancedMCPServer):
-    """
-    Simplified MCP server with minimal configuration.
-
-    This inherits all enhanced capabilities but disables some features
-    by default for simpler use cases.
-    """
-
-    def __init__(self, name: str, description: str = ""):
-        """
-        Initialize simple MCP server.
-
-        Args:
-            name: Server name
-            description: Server description
-        """
-        # Initialize with some features disabled for simplicity
-        super().__init__(
-            name=name,
-            enable_cache=False,  # Disable cache by default
-            enable_metrics=False,  # Disable metrics by default
-            enable_formatting=True,  # Keep formatting for better output
-        )
-
-        self.description = description
-
-        # Update config for simple use
-        self.config.update(
-            {"server": {"name": name, "description": description, "version": "1.0.0"}}
-        )
+SimpleMCPServer = EnhancedMCPServer

kailash/{mcp → mcp_server}/servers/ai_registry.py

@@ -9,10 +9,10 @@ import os
 from pathlib import Path
 from typing import Any
 
-from kailash.mcp.server import MCPServer
+from kailash.mcp_server.server import MCPServerBase
 
 
-class AIRegistryServer(MCPServer):
+class AIRegistryServer(MCPServerBase):
     """MCP server for AI use case registry.
 
     Provides tools and resources for exploring AI use cases from

kailash/{mcp → mcp_server}/utils/__init__.py

@@ -10,12 +10,7 @@ This module provides production-ready utilities for MCP servers including:
 
 from .cache import CacheManager, LRUCache, cached_query
 from .config import ConfigManager
-from .formatters import (
-    format_response,
-    json_formatter,
-    markdown_formatter,
-    search_formatter,
-)
+from .formatters import format_response, json_formatter, markdown_formatter, search_formatter
 from .metrics import MetricsCollector
 
 __all__ = [

kailash/middleware/auth/access_control.py

@@ -60,11 +60,11 @@ class MiddlewareAccessControlManager:
         self.enable_audit = enable_audit
 
         # Kailash nodes for operations
-        self.user_mgmt_node = UserManagementNode("middleware_user_mgmt")
-        self.role_mgmt_node = RoleManagementNode("middleware_role_mgmt")
-        self.permission_check_node = PermissionCheckNode("middleware_perm_check")
-        self.audit_node = AuditLogNode("middleware_audit") if enable_audit else None
-        self.security_event_node = SecurityEventNode("middleware_security")
+        self.user_mgmt_node = UserManagementNode()
+        self.role_mgmt_node = RoleManagementNode()
+        self.permission_check_node = PermissionCheckNode()
+        self.audit_node = AuditLogNode() if enable_audit else None
+        self.security_event_node = SecurityEventNode()
 
     async def check_session_access(
         self, user_context: UserContext, session_id: str, action: str = "access"

kailash/middleware/gateway/checkpoint_manager.py

@@ -175,8 +175,32 @@ class CheckpointManager:
         compression_enabled: bool = True,
         compression_threshold_bytes: int = 1024,  # 1KB
         retention_hours: int = 24,
+        # Backward compatibility parameter
+        storage: Optional[DiskStorage] = None,
     ):
-        """Initialize checkpoint manager."""
+        """Initialize checkpoint manager.
+
+        Args:
+            memory_storage: Memory storage backend (optional)
+            disk_storage: Disk storage backend (optional)
+            cloud_storage: Cloud storage backend (optional)
+            compression_enabled: Enable compression for large checkpoints
+            compression_threshold_bytes: Minimum size for compression
+            retention_hours: Hours to retain checkpoints
+            storage: DEPRECATED - Use disk_storage instead
+        """
+        # Handle backward compatibility
+        if storage is not None:
+            import warnings
+
+            warnings.warn(
+                "The 'storage' parameter is deprecated. Use 'disk_storage' instead.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+            if disk_storage is None:
+                disk_storage = storage
+
         self.memory_storage = memory_storage or MemoryStorage()
         self.disk_storage = disk_storage or DiskStorage()
         self.cloud_storage = cloud_storage  # Optional cloud backend
@@ -189,11 +213,23 @@ class CheckpointManager:
         self.load_count = 0
         self.compression_ratio_sum = 0.0
 
-        # Start garbage collection task
-        self._gc_task = asyncio.create_task(self._garbage_collection_loop())
+        # Initialize garbage collection task (will be started when first used)
+        self._gc_task = None
+        self._gc_started = False
+
+    def _ensure_gc_started(self):
+        """Ensure garbage collection task is started (lazy initialization)."""
+        if not self._gc_started:
+            try:
+                self._gc_task = asyncio.create_task(self._garbage_collection_loop())
+                self._gc_started = True
+            except RuntimeError:
+                # No event loop running, GC will be started later
+                pass
 
     async def save_checkpoint(self, checkpoint: Checkpoint) -> None:
         """Save checkpoint to storage."""
+        self._ensure_gc_started()
         start_time = time.time()
 
         # Serialize checkpoint
@@ -391,8 +427,9 @@ class CheckpointManager:
 
     async def close(self) -> None:
         """Close checkpoint manager and cleanup."""
-        self._gc_task.cancel()
-        try:
-            await self._gc_task
-        except asyncio.CancelledError:
-            pass
+        if self._gc_task is not None:
+            self._gc_task.cancel()
+            try:
+                await self._gc_task
+            except asyncio.CancelledError:
+                pass

kailash/middleware/mcp/client_integration.py

@@ -20,7 +20,7 @@ from kailash.workflow.builder import WorkflowBuilder
 
 # Import existing Kailash MCP components
 try:
-    from kailash.mcp import MCPClient
+    from kailash.mcp_server import MCPClient
 
     _KAILASH_MCP_AVAILABLE = True
 except ImportError:

kailash/middleware/mcp/enhanced_server.py

@@ -24,8 +24,8 @@ from kailash.workflow.builder import WorkflowBuilder
 
 # Import existing Kailash MCP components
 try:
-    from kailash.mcp import MCPServer, SimpleMCPServer
-    from kailash.mcp.utils import CacheManager, ConfigManager, MetricsCollector
+    from kailash.mcp_server import MCPServer, SimpleMCPServer
+    from kailash.mcp_server.utils import CacheManager, ConfigManager, MetricsCollector
 
     _KAILASH_MCP_AVAILABLE = True
 except ImportError: