smartify-ai 0.1.0__py3-none-any.whl

smartify/guardrails/breakers.py

@@ -0,0 +1,409 @@
+"""Breaker system for Smartify grid execution.
+
+Breakers are safety limits that can trip during execution:
+- tokens: Total token usage limits
+- time: Execution time limits  
+- cost: Cost limits in USD
+- requests: Rate limiting (requests per minute, concurrent agents)
+
+When a breaker trips, a configured action is taken:
+- notify: Log warning, continue execution
+- pause: Pause grid, await resume
+- cooldown: Wait before continuing
+- require_approval: Request human approval
+- stop: Stop grid execution
+- downgrade: Switch to cheaper model
+- block: Block and fail immediately
+"""
+
+import asyncio
+import logging
+import time
+from dataclasses import dataclass, field
+from datetime import datetime
+from enum import Enum
+from typing import Any, Callable, Dict, List, Optional, Awaitable
+
+from smartify.models.grid import (
+    BreakerSpec,
+    BreakerActions,
+    TripAction,
+    BreakerStatus,
+    TokenLimits,
+    TimeLimits,
+    CostLimits,
+    RequestLimits,
+)
+
+
+logger = logging.getLogger(__name__)
+
+
+class BreakerType(str, Enum):
+    """Types of breakers."""
+    TOKENS = "tokens"
+    TIME = "time"
+    COST = "cost"
+    REQUESTS = "requests"
+
+
+@dataclass
+class BreakerTrip:
+    """Record of a breaker trip event."""
+    breaker_type: BreakerType
+    action: TripAction
+    reason: str
+    current_value: float
+    limit_value: float
+    timestamp: datetime = field(default_factory=datetime.now)
+    resolved: bool = False
+
+
+@dataclass
+class RateLimitState:
+    """Tracks request rate limiting state."""
+    window_start: float = field(default_factory=time.time)
+    request_count: int = 0
+    concurrent_count: int = 0
+    
+    def reset_window(self) -> None:
+        """Reset the rate limit window."""
+        self.window_start = time.time()
+        self.request_count = 0
+
+
+@dataclass  
+class BreakerState:
+    """Current state of all breakers."""
+    # Cumulative tracking
+    total_tokens: int = 0
+    total_cost: float = 0.0
+    start_time: Optional[float] = None
+    
+    # Rate limiting
+    rate_limit: RateLimitState = field(default_factory=RateLimitState)
+    
+    # Trip history
+    trips: List[BreakerTrip] = field(default_factory=list)
+    
+    # Current status
+    status: BreakerStatus = BreakerStatus.OK
+    paused: bool = False
+    cooldown_until: Optional[float] = None
+    awaiting_approval: bool = False
+    downgraded: bool = False
+    
+    @property
+    def elapsed_seconds(self) -> float:
+        """Get elapsed execution time in seconds."""
+        if self.start_time is None:
+            return 0.0
+        return time.time() - self.start_time
+
+
+class BreakerManager:
+    """Manages breaker checking and trip actions during grid execution.
+    
+    Usage:
+        manager = BreakerManager(spec, actions)
+        manager.start()  # Start timing
+        
+        # Before each node execution:
+        await manager.check_and_enforce()
+        
+        # After LLM call:
+        manager.record_usage(tokens=1500, cost=0.02)
+        
+        # After node completion:
+        manager.record_request_complete()
+    """
+    
+    def __init__(
+        self,
+        spec: BreakerSpec,
+        actions: BreakerActions,
+        on_pause: Optional[Callable[[], Awaitable[None]]] = None,
+        on_approval_required: Optional[Callable[[str], Awaitable[bool]]] = None,
+        on_downgrade: Optional[Callable[[], Awaitable[str]]] = None,
+    ):
+        """Initialize the breaker manager.
+        
+        Args:
+            spec: Breaker configuration (limits)
+            actions: Trip actions configuration
+            on_pause: Async callback when pause action triggered
+            on_approval_required: Async callback for approval (returns True if approved)
+            on_downgrade: Async callback for downgrade (returns new model name)
+        """
+        self.spec = spec
+        self.actions = actions
+        self.state = BreakerState()
+        
+        # Callbacks
+        self._on_pause = on_pause
+        self._on_approval_required = on_approval_required
+        self._on_downgrade = on_downgrade
+    
+    def start(self) -> None:
+        """Start the breaker timers."""
+        self.state.start_time = time.time()
+        self.state.rate_limit = RateLimitState()
+        logger.debug("Breaker manager started")
+    
+    def record_usage(self, tokens: int = 0, cost: float = 0.0) -> None:
+        """Record token and cost usage."""
+        self.state.total_tokens += tokens
+        self.state.total_cost += cost
+        logger.debug(f"Usage recorded: +{tokens} tokens, +${cost:.4f}")
+    
+    def record_request_start(self) -> None:
+        """Record start of a request (for concurrent tracking)."""
+        self.state.rate_limit.concurrent_count += 1
+    
+    def record_request_complete(self) -> None:
+        """Record completion of a request."""
+        self.state.rate_limit.request_count += 1
+        self.state.rate_limit.concurrent_count = max(0, self.state.rate_limit.concurrent_count - 1)
+    
+    async def check_and_enforce(self) -> Optional[BreakerTrip]:
+        """Check all breakers and enforce trip actions.
+        
+        Returns:
+            BreakerTrip if a breaker tripped, None otherwise.
+            
+        Raises:
+            BreakerError: If block action triggered
+        """
+        # Check cooldown first
+        if self.state.cooldown_until:
+            remaining = self.state.cooldown_until - time.time()
+            if remaining > 0:
+                logger.info(f"In cooldown, waiting {remaining:.1f}s")
+                await asyncio.sleep(remaining)
+            self.state.cooldown_until = None
+        
+        # Check if paused
+        if self.state.paused:
+            logger.info("Execution paused, awaiting resume")
+            while self.state.paused:
+                await asyncio.sleep(0.5)
+        
+        # Check each breaker type
+        trip = None
+        
+        trip = trip or await self._check_tokens()
+        trip = trip or await self._check_time()
+        trip = trip or await self._check_cost()
+        trip = trip or await self._check_requests()
+        
+        return trip
+    
+    async def _check_tokens(self) -> Optional[BreakerTrip]:
+        """Check token limits."""
+        if not self.spec.tokens:
+            return None
+        
+        limits = self.spec.tokens
+        limit = limits.maxTotalTokensPerRun
+        
+        # Check warning threshold (80% of limit)
+        warn_at = int(limit * 0.8)
+        if self.state.total_tokens >= warn_at:
+            if self.state.status == BreakerStatus.OK:
+                self.state.status = BreakerStatus.WARNING
+                logger.warning(f"Token warning: {self.state.total_tokens}/{limit}")
+        
+        # Check hard limit
+        if self.state.total_tokens >= limit:
+            return await self._trip(
+                BreakerType.TOKENS,
+                self.actions.onTokensLimit,
+                f"Token limit exceeded: {self.state.total_tokens}/{limit}",
+                self.state.total_tokens,
+                limit,
+            )
+        
+        return None
+    
+    async def _check_time(self) -> Optional[BreakerTrip]:
+        """Check time limits."""
+        if not self.spec.time:
+            return None
+        
+        limits = self.spec.time
+        elapsed = self.state.elapsed_seconds
+        max_seconds = limits.maxRuntimeSeconds
+        
+        # Check warning threshold (80% of limit)
+        warn_at = max_seconds * 0.8
+        if elapsed >= warn_at:
+            if self.state.status == BreakerStatus.OK:
+                self.state.status = BreakerStatus.WARNING
+                logger.warning(f"Time warning: {elapsed:.1f}s/{max_seconds}s")
+        
+        # Check hard limit
+        if elapsed >= max_seconds:
+            return await self._trip(
+                BreakerType.TIME,
+                self.actions.onTimeLimit,
+                f"Time limit exceeded: {elapsed:.1f}s/{max_seconds}s",
+                elapsed,
+                max_seconds,
+            )
+        
+        return None
+    
+    async def _check_cost(self) -> Optional[BreakerTrip]:
+        """Check cost limits."""
+        if not self.spec.cost:
+            return None
+        
+        limits = self.spec.cost
+        limit = limits.maxCostPerRun
+        
+        # Check warning threshold (80% of limit)
+        warn_at = limit * 0.8
+        if self.state.total_cost >= warn_at:
+            if self.state.status == BreakerStatus.OK:
+                self.state.status = BreakerStatus.WARNING
+                logger.warning(f"Cost warning: ${self.state.total_cost:.4f}/${limit}")
+        
+        # Check hard limit
+        if self.state.total_cost >= limit:
+            return await self._trip(
+                BreakerType.COST,
+                self.actions.onCostLimit,
+                f"Cost limit exceeded: ${self.state.total_cost:.4f}/${limit}",
+                self.state.total_cost,
+                limit,
+            )
+        
+        return None
+    
+    async def _check_requests(self) -> Optional[BreakerTrip]:
+        """Check request rate limits."""
+        if not self.spec.requests:
+            return None
+        
+        limits = self.spec.requests
+        rl = self.state.rate_limit
+        
+        # Reset window if needed (1 minute windows)
+        if time.time() - rl.window_start >= 60:
+            rl.reset_window()
+        
+        # Check requests per minute
+        if limits.maxRequestsPerMinute and rl.request_count >= limits.maxRequestsPerMinute:
+            return await self._trip(
+                BreakerType.REQUESTS,
+                self.actions.onRequestsLimit,
+                f"Rate limit exceeded: {rl.request_count}/{limits.maxRequestsPerMinute} req/min",
+                rl.request_count,
+                limits.maxRequestsPerMinute,
+            )
+        
+        # Check concurrent agents
+        if limits.maxConcurrentAgents and rl.concurrent_count >= limits.maxConcurrentAgents:
+            return await self._trip(
+                BreakerType.REQUESTS,
+                self.actions.onRequestsLimit,
+                f"Concurrent limit exceeded: {rl.concurrent_count}/{limits.maxConcurrentAgents}",
+                rl.concurrent_count,
+                limits.maxConcurrentAgents,
+            )
+        
+        return None
+    
+    async def _trip(
+        self,
+        breaker_type: BreakerType,
+        action: TripAction,
+        reason: str,
+        current: float,
+        limit: float,
+    ) -> BreakerTrip:
+        """Handle a breaker trip."""
+        trip = BreakerTrip(
+            breaker_type=breaker_type,
+            action=action,
+            reason=reason,
+            current_value=current,
+            limit_value=limit,
+        )
+        self.state.trips.append(trip)
+        self.state.status = BreakerStatus.TRIPPED
+        
+        logger.warning(f"Breaker tripped: {reason} → action: {action.value}")
+        
+        # Execute the trip action
+        await self._execute_action(action, reason)
+        
+        return trip
+    
+    async def _execute_action(self, action: TripAction, reason: str) -> None:
+        """Execute a trip action."""
+        if action == TripAction.NOTIFY:
+            # Just log, already done
+            pass
+        
+        elif action == TripAction.PAUSE:
+            self.state.paused = True
+            if self._on_pause:
+                await self._on_pause()
+        
+        elif action == TripAction.COOLDOWN:
+            cooldown_secs = self.actions.cooldownSeconds
+            self.state.cooldown_until = time.time() + cooldown_secs
+            logger.info(f"Entering cooldown for {cooldown_secs}s")
+            await asyncio.sleep(cooldown_secs)
+            self.state.cooldown_until = None
+        
+        elif action == TripAction.REQUIRE_APPROVAL:
+            self.state.awaiting_approval = True
+            if self._on_approval_required:
+                approved = await self._on_approval_required(reason)
+                if not approved:
+                    raise BreakerError(f"Approval denied: {reason}")
+            self.state.awaiting_approval = False
+        
+        elif action == TripAction.STOP:
+            raise BreakerError(f"Breaker stop: {reason}")
+        
+        elif action == TripAction.DOWNGRADE:
+            self.state.downgraded = True
+            if self._on_downgrade:
+                new_model = await self._on_downgrade()
+                logger.info(f"Downgraded to model: {new_model}")
+        
+        elif action == TripAction.BLOCK:
+            raise BreakerError(f"Breaker blocked: {reason}")
+    
+    def resume(self) -> None:
+        """Resume paused execution."""
+        self.state.paused = False
+        logger.info("Execution resumed")
+    
+    def get_summary(self) -> Dict[str, Any]:
+        """Get a summary of breaker state."""
+        return {
+            "status": self.state.status.value,
+            "total_tokens": self.state.total_tokens,
+            "total_cost": self.state.total_cost,
+            "elapsed_seconds": self.state.elapsed_seconds,
+            "trips": [
+                {
+                    "type": t.breaker_type.value,
+                    "action": t.action.value,
+                    "reason": t.reason,
+                    "timestamp": t.timestamp.isoformat(),
+                }
+                for t in self.state.trips
+            ],
+            "paused": self.state.paused,
+            "downgraded": self.state.downgraded,
+        }
+
+
+class BreakerError(Exception):
+    """Raised when a breaker trips with stop/block action."""
+    pass

smartify/models/__init__.py

@@ -0,0 +1,61 @@
+"""Pydantic models for Smartify Grid specification."""
+
+from smartify.models.grid import (
+    # Enums
+    NodeKind,
+    TriggerType,
+    ExecutionMode,
+    AutonomyMode,
+    TripAction,
+    BreakerStatus,
+    GridState,
+    WorkspaceType,
+    ArtifactType,
+    AggregateStrategy,
+    # Specs
+    GridSpec,
+    MetadataSpec,
+    TopologySpec,
+    NodeSpec,
+    EdgeSpec,
+    TriggerSpec,
+    InputSpec,
+    EnvironmentSpec,
+    OutputsSpec,
+    AgentSpec,
+    ToolsSpec,
+    GuardrailsSpec,
+    WorkspaceSpec,
+    BreakerSpec,
+    BreakerActions,
+)
+
+__all__ = [
+    # Enums
+    "NodeKind",
+    "TriggerType",
+    "ExecutionMode",
+    "AutonomyMode",
+    "TripAction",
+    "BreakerStatus",
+    "GridState",
+    "WorkspaceType",
+    "ArtifactType",
+    "AggregateStrategy",
+    # Specs
+    "GridSpec",
+    "MetadataSpec",
+    "TopologySpec",
+    "NodeSpec",
+    "EdgeSpec",
+    "TriggerSpec",
+    "InputSpec",
+    "EnvironmentSpec",
+    "OutputsSpec",
+    "AgentSpec",
+    "ToolsSpec",
+    "GuardrailsSpec",
+    "WorkspaceSpec",
+    "BreakerSpec",
+    "BreakerActions",
+]