kubiya-control-plane-api 0.9.15__py3-none-any.whl

control_plane_api/worker/metrics.py

@@ -0,0 +1,237 @@
+"""
+Prometheus Metrics Collection for Agent Runtime
+
+Collects and exposes metrics for monitoring agent-runtime executions.
+"""
+
+import time
+import structlog
+from typing import Optional
+
+logger = structlog.get_logger(__name__)
+
+# Try to import prometheus_client, but make it optional
+try:
+    from prometheus_client import Counter, Histogram, Gauge, Info
+    PROMETHEUS_AVAILABLE = True
+except ImportError:
+    logger.warning("prometheus_client_not_installed", message="Metrics collection disabled. Install with: pip install prometheus-client")
+    PROMETHEUS_AVAILABLE = False
+
+
+# Metrics definitions (only if prometheus available)
+if PROMETHEUS_AVAILABLE:
+    # Execution counter
+    execution_count = Counter(
+        "agent_runtime_executions_total",
+        "Total number of agent executions",
+        ["status", "runtime_type"],
+    )
+
+    # Execution duration histogram
+    execution_duration = Histogram(
+        "agent_runtime_execution_duration_seconds",
+        "Execution duration in seconds",
+        buckets=[0.1, 0.5, 1.0, 2.0, 5.0, 10.0, 30.0, 60.0],
+    )
+
+    # Active executions gauge
+    active_executions = Gauge(
+        "agent_runtime_active_executions",
+        "Number of currently active executions",
+    )
+
+    # Token usage counter
+    token_usage = Counter(
+        "agent_runtime_tokens_total",
+        "Total number of tokens used",
+        ["type"],  # input, output
+    )
+
+    # Tool execution counter
+    tool_executions = Counter(
+        "agent_runtime_tool_executions_total",
+        "Total number of tool executions",
+        ["tool_name", "status"],
+    )
+
+    # Agent runtime info
+    runtime_info = Info(
+        "agent_runtime",
+        "Agent runtime information",
+    )
+
+    # Health check counter
+    health_checks = Counter(
+        "agent_runtime_health_checks_total",
+        "Total number of health checks",
+        ["status"],  # healthy, unhealthy
+    )
+
+    # Restart counter
+    restarts = Counter(
+        "agent_runtime_restarts_total",
+        "Total number of agent-runtime server restarts",
+    )
+
+
+class MetricsCollector:
+    """Collects metrics for agent-runtime executions."""
+
+    def __init__(self, enabled: bool = True):
+        """
+        Initialize metrics collector.
+
+        Args:
+            enabled: Enable metrics collection (default: True)
+        """
+        self.enabled = enabled and PROMETHEUS_AVAILABLE
+
+        if self.enabled:
+            logger.info("metrics_collector_enabled")
+            # Set runtime info
+            runtime_info.info({
+                "version": "0.1.0",  # TODO: Get from version module
+                "runtime_type": "agent_runtime",
+            })
+        else:
+            if not PROMETHEUS_AVAILABLE:
+                logger.info("metrics_collector_disabled_prometheus_not_available")
+            else:
+                logger.info("metrics_collector_disabled")
+
+    def record_execution_start(self):
+        """Record the start of an execution."""
+        if not self.enabled:
+            return
+
+        active_executions.inc()
+
+    def record_execution_end(
+        self,
+        duration: float,
+        status: str,
+        runtime_type: str = "agent_runtime",
+    ):
+        """
+        Record the end of an execution.
+
+        Args:
+            duration: Execution duration in seconds
+            status: Execution status (success, error)
+            runtime_type: Runtime type (default: agent_runtime)
+        """
+        if not self.enabled:
+            return
+
+        active_executions.dec()
+        execution_count.labels(status=status, runtime_type=runtime_type).inc()
+        execution_duration.observe(duration)
+
+    def record_token_usage(self, input_tokens: int, output_tokens: int):
+        """
+        Record token usage.
+
+        Args:
+            input_tokens: Number of input tokens
+            output_tokens: Number of output tokens
+        """
+        if not self.enabled:
+            return
+
+        token_usage.labels(type="input").inc(input_tokens)
+        token_usage.labels(type="output").inc(output_tokens)
+
+    def record_tool_execution(self, tool_name: str, status: str):
+        """
+        Record a tool execution.
+
+        Args:
+            tool_name: Name of the tool
+            status: Execution status (success, error)
+        """
+        if not self.enabled:
+            return
+
+        tool_executions.labels(tool_name=tool_name, status=status).inc()
+
+    def record_health_check(self, is_healthy: bool):
+        """
+        Record a health check.
+
+        Args:
+            is_healthy: Whether the health check passed
+        """
+        if not self.enabled:
+            return
+
+        status = "healthy" if is_healthy else "unhealthy"
+        health_checks.labels(status=status).inc()
+
+    def record_restart(self):
+        """Record an agent-runtime server restart."""
+        if not self.enabled:
+            return
+
+        restarts.inc()
+
+
+class ExecutionMetricsContext:
+    """Context manager for collecting execution metrics."""
+
+    def __init__(
+        self,
+        collector: MetricsCollector,
+        runtime_type: str = "agent_runtime",
+    ):
+        """
+        Initialize context manager.
+
+        Args:
+            collector: MetricsCollector instance
+            runtime_type: Runtime type (default: agent_runtime)
+        """
+        self.collector = collector
+        self.runtime_type = runtime_type
+        self.start_time: Optional[float] = None
+        self.status: str = "error"  # Default to error, set to success explicitly
+
+    def __enter__(self):
+        """Enter context - start metrics collection."""
+        self.start_time = time.time()
+        self.collector.record_execution_start()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Exit context - record metrics."""
+        if self.start_time is not None:
+            duration = time.time() - self.start_time
+            self.collector.record_execution_end(
+                duration=duration,
+                status=self.status,
+                runtime_type=self.runtime_type,
+            )
+
+    def set_success(self):
+        """Mark execution as successful."""
+        self.status = "success"
+
+    def set_error(self):
+        """Mark execution as error."""
+        self.status = "error"
+
+
+# Global metrics collector instance
+_global_collector: Optional[MetricsCollector] = None
+
+
+def get_metrics_collector() -> MetricsCollector:
+    """Get or create global metrics collector."""
+    global _global_collector
+
+    if _global_collector is None:
+        import os
+        enabled = os.environ.get("AGENT_RUNTIME_METRICS_ENABLED", "true").lower() in ("true", "1", "yes")
+        _global_collector = MetricsCollector(enabled=enabled)
+
+    return _global_collector

control_plane_api/worker/models/__init__.py

@@ -0,0 +1 @@
+"""Data models for worker activities"""

control_plane_api/worker/models/error_events.py

@@ -0,0 +1,105 @@
+"""
+Error event models for standardized error handling and streaming.
+
+This module provides consistent error event structures for publishing
+errors from worker components to the Control Plane and UI.
+"""
+
+from enum import Enum
+from dataclasses import dataclass
+from typing import Optional, Dict, Any, List
+from datetime import datetime, timezone
+
+
+class ErrorSeverity(Enum):
+    """Error severity levels for categorization and UI display."""
+    CRITICAL = "critical"  # Execution completely failed
+    ERROR = "error"        # Operation failed but execution may continue
+    WARNING = "warning"    # Issue detected but non-fatal
+    INFO = "info"          # Informational, not an error
+
+
+class ErrorCategory(Enum):
+    """Error categories for filtering and analytics."""
+    RUNTIME_INIT = "runtime_initialization"
+    SKILL_LOADING = "skill_loading"
+    MCP_CONNECTION = "mcp_connection"
+    TOOL_EXECUTION = "tool_execution"
+    API_COMMUNICATION = "api_communication"
+    TIMEOUT = "timeout"
+    VALIDATION = "validation"
+    MODEL_ERROR = "llm_model_error"
+    AUTHENTICATION = "authentication"
+    NETWORK = "network"
+    CONFIGURATION = "configuration"
+    UNKNOWN = "unknown"
+
+
+@dataclass
+class ErrorContext:
+    """Contextual information about where/when error occurred."""
+    execution_id: str
+    timestamp: str
+    stage: str  # "initialization", "skill_loading", "execution", "cleanup"
+    component: str  # "runtime", "skill_factory", "mcp_server", "workflow"
+    operation: Optional[str] = None  # Specific operation that failed
+    retry_count: Optional[int] = None
+
+
+@dataclass
+class ErrorDetails:
+    """Detailed error information."""
+    error_type: str  # Exception class name
+    error_message: str  # User-friendly message
+    technical_message: str  # Full exception message
+    stack_trace: Optional[str] = None  # Stack trace (truncated)
+    code_location: Optional[str] = None  # File:line where error occurred
+
+
+@dataclass
+class ErrorRecovery:
+    """Recovery and remediation information."""
+    is_retryable: bool
+    retry_suggested: bool
+    recovery_actions: List[str]  # User-actionable steps
+    documentation_url: Optional[str] = None
+
+
+@dataclass
+class ErrorEvent:
+    """Complete error event structure."""
+    severity: ErrorSeverity
+    category: ErrorCategory
+    context: ErrorContext
+    details: ErrorDetails
+    recovery: ErrorRecovery
+    metadata: Optional[Dict[str, Any]] = None
+
+    def to_event_data(self) -> Dict[str, Any]:
+        """Convert to event data dict for publishing."""
+        return {
+            "severity": self.severity.value,
+            "category": self.category.value,
+            "context": {
+                "execution_id": self.context.execution_id,
+                "timestamp": self.context.timestamp,
+                "stage": self.context.stage,
+                "component": self.context.component,
+                "operation": self.context.operation,
+                "retry_count": self.context.retry_count,
+            },
+            "details": {
+                "error_type": self.details.error_type,
+                "error_message": self.details.error_message,
+                "technical_message": self.details.technical_message,
+                "stack_trace": self.details.stack_trace,
+                "code_location": self.details.code_location,
+            },
+            "recovery": {
+                "is_retryable": self.recovery.is_retryable,
+                "retry_suggested": self.recovery.retry_suggested,
+                "recovery_actions": self.recovery.recovery_actions,
+                "documentation_url": self.recovery.documentation_url,
+            },
+            "metadata": self.metadata or {},
+        }

control_plane_api/worker/models/inputs.py

@@ -0,0 +1,89 @@
+"""Input dataclasses for all Temporal activities"""
+
+from dataclasses import dataclass
+from typing import Optional, List, Dict, Any
+
+
+@dataclass
+class AgentExecutionInput:
+    """Input for agent execution activity"""
+    execution_id: str
+    agent_id: str
+    organization_id: str
+    prompt: str
+    system_prompt: Optional[str] = None
+    model_id: Optional[str] = None
+    model_config: Optional[Dict[str, Any]] = None
+    mcp_servers: Optional[Dict[str, Any]] = None
+    session_id: Optional[str] = None
+    user_id: Optional[str] = None
+
+    def __post_init__(self):
+        if self.model_config is None:
+            self.model_config = {}
+        if self.mcp_servers is None:
+            self.mcp_servers = {}
+
+
+@dataclass
+class TeamExecutionInput:
+    """Input for team execution activity"""
+    execution_id: str
+    team_id: str
+    organization_id: str
+    prompt: str
+    system_prompt: Optional[str] = None
+    model_id: Optional[str] = None
+    model_config: Optional[Dict[str, Any]] = None
+    mcp_servers: Optional[Dict[str, Any]] = None
+    session_id: Optional[str] = None
+    user_id: Optional[str] = None
+    team_config: Optional[Dict[str, Any]] = None
+
+    def __post_init__(self):
+        if self.model_config is None:
+            self.model_config = {}
+        if self.mcp_servers is None:
+            self.mcp_servers = {}
+        if self.team_config is None:
+            self.team_config = {}
+
+
+@dataclass
+class CancelExecutionInput:
+    """Input for cancellation activity"""
+    execution_id: str
+
+
+@dataclass
+class UpdateExecutionStatusInput:
+    """Input for status update activity"""
+    execution_id: str
+    status: str
+    started_at: Optional[str] = None
+    completed_at: Optional[str] = None
+    response: Optional[str] = None
+    error_message: Optional[str] = None
+    usage: Optional[Dict[str, Any]] = None
+    execution_metadata: Optional[Dict[str, Any]] = None
+
+    def __post_init__(self):
+        if self.usage is None:
+            self.usage = {}
+        if self.execution_metadata is None:
+            self.execution_metadata = {}
+
+
+@dataclass
+class UpdateAgentStatusInput:
+    """Input for agent status update"""
+    agent_id: str
+    organization_id: str
+    status: str
+    last_active_at: str
+    error_message: Optional[str] = None
+    state: Optional[Dict[str, Any]] = None
+
+    def __post_init__(self):
+        if self.state is None:
+            self.state = {}

control_plane_api/worker/runtimes/__init__.py

@@ -0,0 +1,35 @@
+"""
+Runtime abstraction layer for agent execution.
+
+This package provides a pluggable runtime system that allows agents to be
+powered by different frameworks (Agno, Claude Code SDK, etc.) without changing
+the core workflow and activity logic.
+"""
+
+from .base import (
+    RuntimeType,
+    RuntimeExecutionResult,
+    RuntimeExecutionContext,
+    RuntimeCapabilities,
+    BaseRuntime,
+    RuntimeRegistry,
+)
+from .factory import RuntimeFactory
+from .agno import AgnoRuntime
+from .claude_code import ClaudeCodeRuntime
+
+# Backward compatibility: DefaultRuntime is now AgnoRuntime
+DefaultRuntime = AgnoRuntime
+
+__all__ = [
+    "RuntimeType",
+    "RuntimeExecutionResult",
+    "RuntimeExecutionContext",
+    "RuntimeCapabilities",
+    "BaseRuntime",
+    "RuntimeRegistry",
+    "RuntimeFactory",
+    "AgnoRuntime",
+    "DefaultRuntime",  # Backward compatibility alias
+    "ClaudeCodeRuntime",
+]