kailash 0.6.0__py3-none-any.whl → 0.6.2__py3-none-any.whl

kailash/core/resilience/circuit_breaker.py

@@ -0,0 +1,382 @@
+"""Circuit Breaker pattern implementation for connection management.
+
+This module implements the Circuit Breaker pattern to prevent cascading failures
+in connection pools and database operations. It provides automatic failure detection,
+recovery testing, and graceful degradation.
+
+The circuit breaker has three states:
+- CLOSED: Normal operation, requests pass through
+- OPEN: Failures detected, requests fail fast
+- HALF_OPEN: Testing recovery, limited requests allowed
+
+Example:
+    >>> breaker = ConnectionCircuitBreaker(
+    ...     failure_threshold=5,
+    ...     recovery_timeout=60,
+    ...     half_open_requests=3
+    ... )
+    >>>
+    >>> # Wrap connection operations
+    >>> async with breaker.call() as protected:
+    ...     result = await connection.execute(query)
+"""
+
+import asyncio
+import logging
+import time
+from collections import deque
+from dataclasses import dataclass, field
+from datetime import datetime, timedelta
+from enum import Enum
+from typing import Any, Callable, Dict, Generic, List, Optional, TypeVar
+
+logger = logging.getLogger(__name__)
+
+T = TypeVar("T")
+
+
+class CircuitState(Enum):
+    """Circuit breaker states."""
+
+    CLOSED = "closed"  # Normal operation
+    OPEN = "open"  # Failing fast
+    HALF_OPEN = "half_open"  # Testing recovery
+
+
+class CircuitBreakerError(Exception):
+    """Raised when circuit breaker is open."""
+
+    pass
+
+
+@dataclass
+class CircuitBreakerConfig:
+    """Configuration for circuit breaker behavior."""
+
+    failure_threshold: int = 5  # Failures before opening
+    success_threshold: int = 3  # Successes to close from half-open
+    recovery_timeout: int = 60  # Seconds before trying half-open
+    half_open_requests: int = 3  # Requests allowed in half-open
+    error_rate_threshold: float = 0.5  # Error rate to trigger open
+    window_size: int = 100  # Rolling window for error rate
+    excluded_exceptions: List[type] = field(default_factory=list)  # Don't count these
+
+
+@dataclass
+class CircuitBreakerMetrics:
+    """Metrics tracking for circuit breaker."""
+
+    total_calls: int = 0
+    successful_calls: int = 0
+    failed_calls: int = 0
+    rejected_calls: int = 0
+    state_transitions: List[Dict[str, Any]] = field(default_factory=list)
+    last_failure_time: Optional[float] = None
+    consecutive_failures: int = 0
+    consecutive_successes: int = 0
+
+    def record_success(self):
+        """Record successful call."""
+        self.total_calls += 1
+        self.successful_calls += 1
+        self.consecutive_successes += 1
+        self.consecutive_failures = 0
+
+    def record_failure(self):
+        """Record failed call."""
+        self.total_calls += 1
+        self.failed_calls += 1
+        self.consecutive_failures += 1
+        self.consecutive_successes = 0
+        self.last_failure_time = time.time()
+
+    def record_rejection(self):
+        """Record rejected call (circuit open)."""
+        self.rejected_calls += 1
+
+    def get_error_rate(self) -> float:
+        """Calculate current error rate."""
+        if self.total_calls == 0:
+            return 0.0
+        return self.failed_calls / self.total_calls
+
+
+class ConnectionCircuitBreaker(Generic[T]):
+    """Circuit breaker for database connections and operations.
+
+    Monitors failures and prevents cascading failures by failing fast
+    when error threshold is reached. Automatically tests recovery
+    after timeout period.
+    """
+
+    def __init__(self, config: Optional[CircuitBreakerConfig] = None):
+        """Initialize circuit breaker with configuration."""
+        self.config = config or CircuitBreakerConfig()
+        self.state = CircuitState.CLOSED
+        self.metrics = CircuitBreakerMetrics()
+        self._lock = asyncio.Lock()
+        self._half_open_requests = 0
+        self._last_state_change = time.time()
+        self._rolling_window = deque(maxlen=self.config.window_size)
+        self._listeners: List[Callable] = []
+
+    async def call(self, func: Callable[..., T], *args, **kwargs) -> T:
+        """Execute function with circuit breaker protection.
+
+        Args:
+            func: Async function to protect
+            *args: Function arguments
+            **kwargs: Function keyword arguments
+
+        Returns:
+            Function result
+
+        Raises:
+            CircuitBreakerError: If circuit is open
+            Exception: If function fails
+        """
+        async with self._lock:
+            # Check if we should transition states
+            await self._check_state_transition()
+
+            if self.state == CircuitState.OPEN:
+                self.metrics.record_rejection()
+                raise CircuitBreakerError(
+                    f"Circuit breaker is OPEN. "
+                    f"Rejected after {self.metrics.consecutive_failures} failures. "
+                    f"Will retry in {self._time_until_recovery():.1f}s"
+                )
+
+            if self.state == CircuitState.HALF_OPEN:
+                if self._half_open_requests >= self.config.half_open_requests:
+                    self.metrics.record_rejection()
+                    raise CircuitBreakerError(
+                        "Circuit breaker is HALF_OPEN but request limit reached"
+                    )
+                self._half_open_requests += 1
+
+        # Execute the function
+        start_time = time.time()
+        try:
+            result = await func(*args, **kwargs)
+            await self._record_success()
+            return result
+        except Exception as e:
+            # Check if this exception should be counted
+            if not any(
+                isinstance(e, exc_type) for exc_type in self.config.excluded_exceptions
+            ):
+                await self._record_failure(e)
+            raise
+
+    async def _check_state_transition(self):
+        """Check if state should transition based on metrics."""
+        current_time = time.time()
+
+        if self.state == CircuitState.CLOSED:
+            # Check if we should open
+            if self._should_open():
+                await self._transition_to(CircuitState.OPEN)
+
+        elif self.state == CircuitState.OPEN:
+            # Check if we should try recovery
+            time_since_open = current_time - self._last_state_change
+            if time_since_open >= self.config.recovery_timeout:
+                await self._transition_to(CircuitState.HALF_OPEN)
+                self._half_open_requests = 0
+
+        elif self.state == CircuitState.HALF_OPEN:
+            # This is handled after request execution
+            pass
+
+    def _should_open(self) -> bool:
+        """Determine if circuit should open based on failures."""
+        # Check consecutive failures
+        if self.metrics.consecutive_failures >= self.config.failure_threshold:
+            return True
+
+        # Check error rate in rolling window
+        if len(self._rolling_window) >= self.config.window_size / 2:
+            error_count = sum(1 for success in self._rolling_window if not success)
+            error_rate = error_count / len(self._rolling_window)
+            if error_rate >= self.config.error_rate_threshold:
+                return True
+
+        return False
+
+    async def _record_success(self):
+        """Record successful execution."""
+        async with self._lock:
+            self.metrics.record_success()
+            self._rolling_window.append(True)
+
+            if self.state == CircuitState.HALF_OPEN:
+                if self.metrics.consecutive_successes >= self.config.success_threshold:
+                    await self._transition_to(CircuitState.CLOSED)
+
+    async def _record_failure(self, error: Exception):
+        """Record failed execution."""
+        async with self._lock:
+            self.metrics.record_failure()
+            self._rolling_window.append(False)
+
+            if self.state == CircuitState.HALF_OPEN:
+                # Single failure in half-open goes back to open
+                await self._transition_to(CircuitState.OPEN)
+            elif self.state == CircuitState.CLOSED:
+                # Check if we should open the circuit
+                if self._should_open():
+                    await self._transition_to(CircuitState.OPEN)
+
+            logger.warning(
+                f"Circuit breaker recorded failure: {type(error).__name__}: {error}"
+            )
+
+    async def _transition_to(self, new_state: CircuitState):
+        """Transition to new state and notify listeners."""
+        old_state = self.state
+        self.state = new_state
+        self._last_state_change = time.time()
+
+        # Reset counters on state change
+        if new_state == CircuitState.CLOSED:
+            self.metrics.consecutive_failures = 0
+        elif new_state == CircuitState.OPEN:
+            self.metrics.consecutive_successes = 0
+
+        # Record transition
+        self.metrics.state_transitions.append(
+            {
+                "from": old_state.value,
+                "to": new_state.value,
+                "timestamp": datetime.now().isoformat(),
+                "reason": self._get_transition_reason(old_state, new_state),
+            }
+        )
+
+        logger.info(
+            f"Circuit breaker transitioned from {old_state.value} to {new_state.value}"
+        )
+
+        # Notify listeners
+        for listener in self._listeners:
+            try:
+                await listener(old_state, new_state, self.metrics)
+            except Exception as e:
+                logger.error(f"Error notifying circuit breaker listener: {e}")
+
+    def _get_transition_reason(
+        self, old_state: CircuitState, new_state: CircuitState
+    ) -> str:
+        """Get human-readable reason for state transition."""
+        if old_state == CircuitState.CLOSED and new_state == CircuitState.OPEN:
+            return f"Failure threshold reached ({self.metrics.consecutive_failures} failures)"
+        elif old_state == CircuitState.OPEN and new_state == CircuitState.HALF_OPEN:
+            return f"Recovery timeout elapsed ({self.config.recovery_timeout}s)"
+        elif old_state == CircuitState.HALF_OPEN and new_state == CircuitState.CLOSED:
+            return f"Success threshold reached ({self.metrics.consecutive_successes} successes)"
+        elif old_state == CircuitState.HALF_OPEN and new_state == CircuitState.OPEN:
+            return "Failure during recovery test"
+        return "Unknown reason"
+
+    def _time_until_recovery(self) -> float:
+        """Calculate seconds until recovery attempt."""
+        if self.state != CircuitState.OPEN:
+            return 0.0
+        elapsed = time.time() - self._last_state_change
+        remaining = self.config.recovery_timeout - elapsed
+        return max(0.0, remaining)
+
+    async def force_open(self, reason: str = "Manual override"):
+        """Manually open the circuit breaker."""
+        async with self._lock:
+            if self.state != CircuitState.OPEN:
+                logger.warning(f"Manually opening circuit breaker: {reason}")
+                await self._transition_to(CircuitState.OPEN)
+
+    async def force_close(self, reason: str = "Manual override"):
+        """Manually close the circuit breaker."""
+        async with self._lock:
+            if self.state != CircuitState.CLOSED:
+                logger.warning(f"Manually closing circuit breaker: {reason}")
+                self.metrics.consecutive_failures = 0
+                self.metrics.consecutive_successes = 0
+                await self._transition_to(CircuitState.CLOSED)
+
+    async def reset(self):
+        """Reset circuit breaker to initial state."""
+        async with self._lock:
+            self.state = CircuitState.CLOSED
+            self.metrics = CircuitBreakerMetrics()
+            self._rolling_window.clear()
+            self._half_open_requests = 0
+            self._last_state_change = time.time()
+            logger.info("Circuit breaker reset to initial state")
+
+    def add_listener(self, listener: Callable):
+        """Add state change listener."""
+        self._listeners.append(listener)
+
+    def remove_listener(self, listener: Callable):
+        """Remove state change listener."""
+        if listener in self._listeners:
+            self._listeners.remove(listener)
+
+    def get_status(self) -> Dict[str, Any]:
+        """Get current circuit breaker status."""
+        return {
+            "state": self.state.value,
+            "metrics": {
+                "total_calls": self.metrics.total_calls,
+                "successful_calls": self.metrics.successful_calls,
+                "failed_calls": self.metrics.failed_calls,
+                "rejected_calls": self.metrics.rejected_calls,
+                "error_rate": self.metrics.get_error_rate(),
+                "consecutive_failures": self.metrics.consecutive_failures,
+                "consecutive_successes": self.metrics.consecutive_successes,
+            },
+            "config": {
+                "failure_threshold": self.config.failure_threshold,
+                "success_threshold": self.config.success_threshold,
+                "recovery_timeout": self.config.recovery_timeout,
+                "error_rate_threshold": self.config.error_rate_threshold,
+            },
+            "time_until_recovery": (
+                self._time_until_recovery() if self.state == CircuitState.OPEN else None
+            ),
+            "state_transitions": self.metrics.state_transitions[
+                -5:
+            ],  # Last 5 transitions
+        }
+
+
+class CircuitBreakerManager:
+    """Manages multiple circuit breakers for different resources."""
+
+    def __init__(self):
+        """Initialize circuit breaker manager."""
+        self._breakers: Dict[str, ConnectionCircuitBreaker] = {}
+        self._default_config = CircuitBreakerConfig()
+
+    def get_or_create(
+        self, name: str, config: Optional[CircuitBreakerConfig] = None
+    ) -> ConnectionCircuitBreaker:
+        """Get existing or create new circuit breaker."""
+        if name not in self._breakers:
+            self._breakers[name] = ConnectionCircuitBreaker(
+                config or self._default_config
+            )
+        return self._breakers[name]
+
+    def get_all_status(self) -> Dict[str, Dict[str, Any]]:
+        """Get status of all circuit breakers."""
+        return {name: breaker.get_status() for name, breaker in self._breakers.items()}
+
+    async def reset_all(self):
+        """Reset all circuit breakers."""
+        for breaker in self._breakers.values():
+            await breaker.reset()
+
+    def set_default_config(self, config: CircuitBreakerConfig):
+        """Set default configuration for new breakers."""
+        self._default_config = config

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/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"
@@ -72,7 +72,7 @@ class MiddlewareAccessControlManager:
         """Check if user can access a specific session."""
 
         # Use Kailash permission check node
-        result = self.permission_check_node.process(
+        result = self.permission_check_node.execute(
             {
                 "user_context": user_context,
                 "resource_type": "session",
@@ -114,7 +114,7 @@ class MiddlewareAccessControlManager:
 
         # Audit logging using Kailash audit node
         if self.enable_audit and self.audit_node:
-            self.audit_node.process(
+            self.audit_node.execute(
                 {
                     "event_type": "workflow_access_check",
                     "user_id": user_context.user_id,
@@ -192,7 +192,7 @@ class MiddlewareAccessControlManager:
     ) -> Dict[str, Any]:
         """Assign role to user using Kailash role management node."""
 
-        result = self.role_mgmt_node.process(
+        result = self.role_mgmt_node.execute(
             {
                 "action": "assign_role",
                 "user_id": user_id,
@@ -239,7 +239,7 @@ class MiddlewareAccessControlManager:
 
         # Audit the rule creation
         if self.enable_audit and self.audit_node:
-            self.audit_node.process(
+            self.audit_node.execute(
                 {
                     "event_type": "permission_rule_created",
                     "rule_data": rule_data,
@@ -366,7 +366,7 @@ class MiddlewareAuthenticationMiddleware:
         try:
             # This would typically validate JWT token
             # For now, simulating with credential manager
-            cred_result = self.credential_manager.process(
+            cred_result = self.credential_manager.execute(
                 {"action": "validate_token", "token": token}
             )
 
@@ -388,7 +388,7 @@ class MiddlewareAuthenticationMiddleware:
 
         except Exception as e:
             # Log security event using Kailash security event node
-            self.access_manager.security_event_node.process(
+            self.access_manager.security_event_node.execute(
                 {
                     "event_type": "authentication_failure",
                     "error": str(e),

kailash/middleware/communication/ai_chat.py

@@ -371,7 +371,7 @@ EXPLANATION:
 
         try:
             result = await asyncio.to_thread(
-                self.llm_node.process, messages=[{"role": "user", "content": prompt}]
+                self.llm_node.execute, messages=[{"role": "user", "content": prompt}]
             )
 
             # Extract content from response
@@ -847,10 +847,10 @@ What would you like to work on? Just describe what you want to accomplish and I'
         """Store chat message with embedding in vector database."""
         try:
             # Generate embedding
-            embedding_result = await self.embedding_node.process({"text": content})
+            embedding_result = self.embedding_node.execute(text=content)
 
             # Store in database (simplified for now)
-            await self.vector_db.process(
+            self.vector_db.execute(
                 {
                     "query": "INSERT INTO chat_messages (id, session_id, user_id, content, role, timestamp) VALUES (?, ?, ?, ?, ?, ?)",
                     "parameters": [
@@ -875,10 +875,10 @@ What would you like to work on? Just describe what you want to accomplish and I'
         """Find similar past conversations using vector search."""
         try:
             # Generate query embedding
-            query_embedding = await self.embedding_node.process({"text": query})
+            query_embedding = self.embedding_node.execute(text=query)
 
             # Search for similar messages (simplified for now)
-            search_result = await self.vector_db.process(
+            search_result = self.vector_db.execute(
                 {
                     "query": "SELECT * FROM chat_messages WHERE role = 'user' ORDER BY timestamp DESC LIMIT ?",
                     "parameters": [limit * 2],
@@ -930,7 +930,7 @@ What would you like to work on? Just describe what you want to accomplish and I'
 
         try:
             # Generate query embedding
-            query_embedding = await self.embedding_node.process({"text": query})
+            query_embedding = self.embedding_node.execute(text=query)
 
             # Prepare filters
             filters = {}
@@ -948,7 +948,7 @@ What would you like to work on? Just describe what you want to accomplish and I'
             query_parts.append("ORDER BY timestamp DESC LIMIT ?")
             params.append(limit)
 
-            search_result = await self.vector_db.process(
+            search_result = self.vector_db.execute(
                 {"query": " ".join(query_parts), "parameters": params}
             )
 

kailash/middleware/communication/api_gateway.py

@@ -225,10 +225,8 @@ class APIGateway:
         # Data transformer for request/response formatting
         self.data_transformer = DataTransformer(
             name="gateway_transformer",
-            transformations=[
-                {"type": "validate", "schema": "api_response"},
-                {"type": "add_field", "field": "timestamp", "value": "now()"},
-            ],
+            # Transformations will be provided at runtime
+            transformations=[],
         )
 
         # Credential manager for gateway security
@@ -362,17 +360,9 @@ class APIGateway:
                     "active": session.active,
                 }
 
-                transformed = await self.data_transformer.process(
-                    {
-                        "data": response_data,
-                        "transformations": [
-                            {
-                                "type": "add_field",
-                                "field": "api_version",
-                                "value": self.version,
-                            }
-                        ],
-                    }
+                transformed = self.data_transformer.execute(
+                    data=response_data,
+                    transformations=[f"{{**data, 'api_version': '{self.version}'}}"],
                 )
 
                 return SessionResponse(**transformed["result"])

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