kailash 0.6.0__py3-none-any.whl → 0.6.1__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/middleware/auth/access_control.py

@@ -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/event_store.py

@@ -13,7 +13,7 @@ import logging
 import time
 import uuid
 from dataclasses import dataclass, field
-from datetime import datetime
+from datetime import UTC, datetime
 from enum import Enum
 from typing import Any, AsyncIterator, Callable, Dict, List, Optional
 
@@ -55,7 +55,7 @@ class RequestEvent:
     event_id: str = field(default_factory=lambda: f"evt_{uuid.uuid4().hex[:12]}")
     event_type: EventType = EventType.REQUEST_CREATED
     request_id: str = ""
-    timestamp: datetime = field(default_factory=datetime.utcnow)
+    timestamp: datetime = field(default_factory=lambda: datetime.now(UTC))
     sequence_number: int = 0
     data: Dict[str, Any] = field(default_factory=dict)
     metadata: Dict[str, Any] = field(default_factory=dict)
@@ -103,6 +103,7 @@ class EventStore:
         # In-memory buffer
         self._buffer: List[RequestEvent] = []
         self._buffer_lock = asyncio.Lock()
+        self._flush_in_progress = False
 
         # Event stream
         self._event_stream: List[RequestEvent] = []
@@ -120,7 +121,16 @@ class EventStore:
         self.flush_count = 0
 
         # Start flush task
-        self._flush_task = asyncio.create_task(self._flush_loop())
+        try:
+            self._flush_task = asyncio.create_task(self._flush_loop())
+        except RuntimeError:
+            # If no event loop is running, defer task creation
+            self._flush_task = None
+
+    async def _ensure_flush_task(self):
+        """Ensure the flush task is running."""
+        if self._flush_task is None:
+            self._flush_task = asyncio.create_task(self._flush_loop())
 
     async def append(
         self,
@@ -130,6 +140,9 @@ class EventStore:
         metadata: Optional[Dict[str, Any]] = None,
     ) -> RequestEvent:
         """Append an event to the store."""
+        # Ensure flush task is running
+        await self._ensure_flush_task()
+
         async with self._buffer_lock:
             # Get next sequence number
             sequence = self._sequences.get(request_id, 0)
@@ -148,19 +161,27 @@ class EventStore:
             self._buffer.append(event)
             self.event_count += 1
 
-            # Flush if buffer is full
-            if len(self._buffer) >= self.batch_size:
-                await self._flush_buffer()
+            # Check if we need to flush (but don't flush inside the lock)
+            needs_flush = len(self._buffer) >= self.batch_size
 
-            # Apply projections
-            await self._apply_projections(event)
+        # Apply projections outside the lock
+        await self._apply_projections(event)
 
-            logger.debug(
-                f"Appended event {event.event_type.value} for request {request_id} "
-                f"(seq: {sequence})"
-            )
+        # Flush if needed (outside the lock to avoid deadlock)
+        if needs_flush and not self._flush_in_progress:
+            # Set flag to prevent concurrent flushes
+            self._flush_in_progress = True
+            try:
+                await self._flush_buffer()
+            finally:
+                self._flush_in_progress = False
 
-            return event
+        logger.debug(
+            f"Appended event {event.event_type.value} for request {request_id} "
+            f"(seq: {sequence})"
+        )
+
+        return event
 
     async def get_events(
         self,
@@ -233,6 +254,9 @@ class EventStore:
         follow: bool = False,
     ) -> AsyncIterator[RequestEvent]:
         """Stream events as they occur."""
+        # Ensure buffer is flushed before streaming
+        await self._flush_buffer()
+
         last_index = 0
 
         while True:
@@ -294,12 +318,19 @@ class EventStore:
 
     async def _flush_buffer(self) -> None:
         """Flush event buffer to storage."""
-        async with self._buffer_lock:
-            if not self._buffer:
-                return
-
-            events_to_flush = self._buffer.copy()
-            self._buffer.clear()
+        # Acquire lock with timeout to prevent deadlock
+        try:
+            # Use wait_for to add timeout on lock acquisition
+            async with asyncio.timeout(1.0):  # 1 second timeout
+                async with self._buffer_lock:
+                    if not self._buffer:
+                        return
+
+                    events_to_flush = self._buffer.copy()
+                    self._buffer.clear()
+        except asyncio.TimeoutError:
+            logger.warning("Timeout acquiring buffer lock during flush")
+            return
 
         # Add to in-memory stream
         async with self._stream_lock:
@@ -317,10 +348,16 @@ class EventStore:
         while True:
             try:
                 await asyncio.sleep(self.flush_interval)
-                await self._flush_buffer()
+                if not self._flush_in_progress:
+                    self._flush_in_progress = True
+                    try:
+                        await self._flush_buffer()
+                    finally:
+                        self._flush_in_progress = False
             except asyncio.CancelledError:
                 # Final flush before shutdown
-                await self._flush_buffer()
+                if not self._flush_in_progress:
+                    await self._flush_buffer()
                 break
             except Exception as e:
                 logger.error(f"Flush error: {e}")
@@ -388,11 +425,14 @@ class EventStore:
 
     async def close(self) -> None:
         """Close event store and flush remaining events."""
-        self._flush_task.cancel()
-        try:
-            await self._flush_task
-        except asyncio.CancelledError:
-            pass
+        if self._flush_task is not None:
+            self._flush_task.cancel()
+            try:
+                await self._flush_task
+            except asyncio.CancelledError:
+                pass
+        # Final flush
+        await self._flush_buffer()
 
 
 # Example projection handlers