kubiya-control-plane-api 0.1.0__py3-none-any.whl → 0.3.4__py3-none-any.whl

control_plane_api/worker/runtimes/factory.py

@@ -0,0 +1,173 @@
+"""
+Runtime factory using the registry system.
+
+This module provides a simplified factory that delegates to RuntimeRegistry.
+"""
+
+from typing import TYPE_CHECKING, List
+import structlog
+
+from .base import RuntimeType, BaseRuntime, RuntimeRegistry
+
+if TYPE_CHECKING:
+    from control_plane_client import ControlPlaneClient
+    from services.cancellation_manager import CancellationManager
+
+logger = structlog.get_logger(__name__)
+
+
+class RuntimeFactory:
+    """
+    Factory for creating runtime instances.
+
+    This is a thin wrapper around RuntimeRegistry that provides
+    backward compatibility and convenience methods.
+    """
+
+    @staticmethod
+    def create_runtime(
+        runtime_type: RuntimeType,
+        control_plane_client: "ControlPlaneClient",
+        cancellation_manager: "CancellationManager",
+        **kwargs,
+    ) -> BaseRuntime:
+        """
+        Create a runtime instance using the registry.
+
+        Args:
+            runtime_type: Type of runtime to create
+            control_plane_client: Client for Control Plane API
+            cancellation_manager: Manager for execution cancellation
+            **kwargs: Additional runtime-specific configuration
+
+        Returns:
+            BaseRuntime instance
+
+        Raises:
+            ValueError: If runtime_type is not supported
+
+        Example:
+            >>> factory = RuntimeFactory()
+            >>> runtime = factory.create_runtime(
+            ...     RuntimeType.CLAUDE_CODE,
+            ...     control_plane_client,
+            ...     cancellation_manager
+            ... )
+        """
+        logger.info(
+            "creating_runtime",
+            runtime_type=runtime_type.value,
+            has_kwargs=bool(kwargs),
+        )
+
+        return RuntimeRegistry.create(
+            runtime_type=runtime_type,
+            control_plane_client=control_plane_client,
+            cancellation_manager=cancellation_manager,
+            **kwargs,
+        )
+
+    @staticmethod
+    def get_default_runtime_type() -> RuntimeType:
+        """
+        Get the default runtime type.
+
+        This is used when no runtime is explicitly specified in agent config.
+
+        Returns:
+            Default RuntimeType (RuntimeType.DEFAULT)
+        """
+        return RuntimeType.DEFAULT
+
+    @staticmethod
+    def get_supported_runtimes() -> List[RuntimeType]:
+        """
+        Get list of supported runtimes from registry.
+
+        Returns:
+            List of RuntimeType enum values
+        """
+        return RuntimeRegistry.list_available()
+
+    @staticmethod
+    def parse_runtime_type(runtime_str: str) -> RuntimeType:
+        """
+        Parse runtime type from string with fallback to default.
+
+        Args:
+            runtime_str: Runtime type as string
+
+        Returns:
+            RuntimeType enum value, defaults to DEFAULT if invalid
+
+        Example:
+            >>> RuntimeFactory.parse_runtime_type("claude_code")
+            RuntimeType.CLAUDE_CODE
+            >>> RuntimeFactory.parse_runtime_type("invalid")
+            RuntimeType.DEFAULT
+        """
+        try:
+            return RuntimeType(runtime_str)
+        except ValueError:
+            logger.warning(
+                "invalid_runtime_type_fallback",
+                runtime_str=runtime_str,
+                default=RuntimeType.DEFAULT.value,
+            )
+            return RuntimeType.DEFAULT
+
+    @staticmethod
+    def get_runtime_info_all() -> dict:
+        """
+        Get information about all available runtimes.
+
+        Returns:
+            Dict mapping runtime type to info dict
+        """
+        return RuntimeRegistry.get_runtime_info_all()
+
+    @staticmethod
+    def validate_runtime_config(
+        runtime_type: RuntimeType, config: dict
+    ) -> tuple[bool, str]:
+        """
+        Validate runtime-specific configuration.
+
+        Args:
+            runtime_type: Type of runtime
+            config: Configuration dict to validate
+
+        Returns:
+            Tuple of (is_valid, error_message)
+
+        Example:
+            >>> is_valid, error = RuntimeFactory.validate_runtime_config(
+            ...     RuntimeType.CLAUDE_CODE,
+            ...     {"allowed_tools": ["Bash"]}
+            ... )
+        """
+        # Get runtime class from registry
+        try:
+            runtime_class = RuntimeRegistry.get(runtime_type)
+
+            # Create temporary instance to validate (with mocks)
+            from unittest.mock import MagicMock
+            temp_runtime = runtime_class(
+                control_plane_client=MagicMock(),
+                cancellation_manager=MagicMock(),
+            )
+
+            # Use runtime's validation if available
+            if hasattr(temp_runtime, 'validate_config'):
+                try:
+                    temp_runtime.validate_config(config)
+                    return True, ""
+                except ValueError as e:
+                    return False, str(e)
+
+            return True, ""
+
+        except ValueError as e:
+            return False, f"Unknown runtime type: {str(e)}"
+        except Exception as e:
+            return False, f"Validation error: {str(e)}"

control_plane_api/worker/runtimes/validation.py

@@ -0,0 +1,93 @@
+"""
+Runtime validation system with model compatibility and requirements checking.
+
+This module re-exports validation functions from the shared validation module
+and provides worker-specific validation helpers for RuntimeExecutionContext.
+"""
+
+from typing import List, Optional, Dict, Any
+from control_plane_api.worker.runtimes.base import RuntimeType, RuntimeExecutionContext
+
+# Import shared validation logic that both API and worker can use
+from control_plane_api.app.lib.validation import (
+    validate_agent_for_runtime as _validate_agent_for_runtime,
+    get_runtime_requirements_info as _get_runtime_requirements_info,
+    list_all_runtime_requirements as _list_all_runtime_requirements,
+    RUNTIME_REQUIREMENTS as _RUNTIME_REQUIREMENTS,
+)
+
+
+# Re-export for backward compatibility
+validate_agent_for_runtime = _validate_agent_for_runtime
+get_runtime_requirements_info = _get_runtime_requirements_info
+list_all_runtime_requirements = _list_all_runtime_requirements
+
+
+class ValidationError(Exception):
+    """Raised when validation fails."""
+
+    def __init__(self, message: str, field: Optional[str] = None, details: Optional[Dict[str, Any]] = None):
+        self.message = message
+        self.field = field
+        self.details = details or {}
+        super().__init__(message)
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dict for API responses."""
+        result = {"error": self.message}
+        if self.field:
+            result["field"] = self.field
+        if self.details:
+            result["details"] = self.details
+        return result
+
+
+def validate_execution_context(context: RuntimeExecutionContext) -> tuple[bool, List[str]]:
+    """
+    Validate RuntimeExecutionContext for worker execution.
+
+    This is a worker-specific helper that validates the full execution context
+    including conversation history, skills, and runtime-specific requirements.
+
+    Args:
+        context: Runtime execution context
+
+    Returns:
+        Tuple of (is_valid, error_messages)
+    """
+    errors = []
+
+    # Get runtime requirements from shared module
+    runtime_type_str = context.runtime_type.value if isinstance(context.runtime_type, RuntimeType) else str(context.runtime_type)
+    requirements = _RUNTIME_REQUIREMENTS.get(runtime_type_str)
+
+    if not requirements:
+        # No requirements registered - allow by default
+        return True, []
+
+    # Validate model
+    is_valid, error = requirements.validate_model(context.model_id)
+    if not is_valid:
+        errors.append(error)
+
+    # Validate config
+    config_errors = requirements.validate_config(context.agent_config)
+    errors.extend(config_errors)
+
+    # Validate system prompt if required
+    if requirements.requires_system_prompt and not context.system_prompt:
+        errors.append("System prompt is required for this runtime")
+
+    # Validate tools if required
+    if requirements.requires_tools and not context.skills:
+        errors.append("At least one skill is required for this runtime")
+
+    # Validate history length
+    if requirements.max_history_length and context.conversation_history:
+        if len(context.conversation_history) > requirements.max_history_length:
+            errors.append(
+                f"Conversation history too long ({len(context.conversation_history)} messages). "
+                f"Maximum: {requirements.max_history_length}"
+            )
+
+    return len(errors) == 0, errors

control_plane_api/worker/services/__init__.py

@@ -0,0 +1 @@
+"""Business logic services for worker operations"""

control_plane_api/worker/services/agent_executor.py

@@ -0,0 +1,422 @@
+"""Agent executor service - handles agent execution business logic"""
+
+from typing import Dict, Any, Optional, List
+from datetime import datetime, timezone
+import structlog
+import asyncio
+import os
+
+from agno.agent import Agent
+from agno.models.litellm import LiteLLM
+
+from control_plane_api.worker.control_plane_client import ControlPlaneClient
+from control_plane_api.worker.services.session_service import SessionService
+from control_plane_api.worker.services.cancellation_manager import CancellationManager
+from control_plane_api.worker.services.skill_factory import SkillFactory
+from control_plane_api.worker.utils.streaming_utils import StreamingHelper
+
+logger = structlog.get_logger()
+
+
+class AgentExecutorService:
+    """
+    Service for executing agents with full session management and cancellation support.
+
+    This service orchestrates:
+    - Session loading and restoration
+    - Agent creation with LiteLLM configuration
+    - Skill instantiation
+    - Streaming execution with real-time updates
+    - Session persistence
+    - Cancellation support via CancellationManager
+    """
+
+    def __init__(
+        self,
+        control_plane: ControlPlaneClient,
+        session_service: SessionService,
+        cancellation_manager: CancellationManager
+    ):
+        self.control_plane = control_plane
+        self.session_service = session_service
+        self.cancellation_manager = cancellation_manager
+
+    async def execute(self, input: Any) -> Dict[str, Any]:
+        """
+        Execute an agent with full session management and streaming.
+
+        Args:
+            input: AgentExecutionInput with execution details
+
+        Returns:
+            Dict with response, usage, success flag, etc.
+        """
+        execution_id = input.execution_id
+
+        print("\n" + "="*80)
+        print("🤖 AGENT EXECUTION START")
+        print("="*80)
+        print(f"Execution ID: {execution_id}")
+        print(f"Agent ID: {input.agent_id}")
+        print(f"Organization: {input.organization_id}")
+        print(f"Model: {input.model_id or 'default'}")
+        print(f"Session ID: {input.session_id}")
+        print(f"Prompt: {input.prompt[:100]}..." if len(input.prompt) > 100 else f"Prompt: {input.prompt}")
+        print("="*80 + "\n")
+
+        logger.info(
+            "agent_execution_start",
+            execution_id=execution_id[:8],
+            agent_id=input.agent_id,
+            session_id=input.session_id
+        )
+
+        try:
+            # STEP 1: Load session history
+            session_history = self.session_service.load_session(
+                execution_id=execution_id,
+                session_id=input.session_id
+            )
+
+            if session_history:
+                print(f"✅ Loaded {len(session_history)} messages from previous session\n")
+            else:
+                print("ℹ️  Starting new conversation\n")
+
+            # STEP 2: Build conversation context for Agno
+            conversation_context = self.session_service.build_conversation_context(session_history)
+
+            # STEP 3: Get LiteLLM configuration
+            litellm_api_base = os.getenv("LITELLM_API_BASE", "https://llm-proxy.kubiya.ai")
+            litellm_api_key = os.getenv("LITELLM_API_KEY")
+
+            if not litellm_api_key:
+                raise ValueError("LITELLM_API_KEY environment variable not set")
+
+            model = input.model_id or os.environ.get("LITELLM_DEFAULT_MODEL", "kubiya/claude-sonnet-4")
+
+            # STEP 4: Fetch and instantiate skills
+            skills = []
+            if input.agent_id:
+                print(f"🔧 Fetching skills from Control Plane...")
+                try:
+                    skill_configs = self.control_plane.get_skills(input.agent_id)
+                    if skill_configs:
+                        print(f"✅ Resolved {len(skill_configs)} skills")
+                        print(f"   Types: {[t.get('type') for t in skill_configs]}")
+                        print(f"   Names: {[t.get('name') for t in skill_configs]}\n")
+
+                        skills = SkillFactory.create_skills_from_list(skill_configs)
+
+                        if skills:
+                            print(f"✅ Instantiated {len(skills)} skill(s)\n")
+                    else:
+                        print(f"⚠️  No skills found\n")
+                except Exception as e:
+                    print(f"❌ Error fetching skills: {str(e)}\n")
+                    logger.error("skill_fetch_error", error=str(e))
+
+            # STEP 5: Create agent with streaming helper
+            print(f"\n🤖 Creating Agno Agent:")
+            print(f"   Model: {model}")
+            print(f"   Skills: {len(skills)}")
+
+            # Create streaming helper for this execution
+            streaming_helper = StreamingHelper(
+                control_plane_client=self.control_plane,
+                execution_id=execution_id
+            )
+
+            # Create tool hook for real-time updates
+            def tool_hook(name: str = None, function_name: str = None, function=None, arguments: dict = None, **kwargs):
+                """Hook to capture tool execution for real-time streaming"""
+                tool_name = name or function_name or "unknown"
+                tool_args = arguments or {}
+
+                # Generate unique tool execution ID
+                import time
+                tool_execution_id = f"{tool_name}_{int(time.time() * 1000000)}"
+
+                print(f"   🔧 Tool Starting: {tool_name} (ID: {tool_execution_id})")
+
+                # Publish tool start event
+                streaming_helper.publish_tool_start(
+                    tool_name=tool_name,
+                    tool_execution_id=tool_execution_id,
+                    tool_args=tool_args,
+                    source="agent"
+                )
+
+                # Execute the tool
+                result = None
+                error = None
+                try:
+                    if function and callable(function):
+                        result = function(**tool_args) if tool_args else function()
+                    else:
+                        raise ValueError(f"Function not callable: {function}")
+
+                    status = "success"
+                    print(f"   ✅ Tool Success: {tool_name}")
+
+                except Exception as e:
+                    error = e
+                    status = "failed"
+                    print(f"   ❌ Tool Failed: {tool_name} - {str(e)}")
+
+                # Publish tool completion event
+                streaming_helper.publish_tool_complete(
+                    tool_name=tool_name,
+                    tool_execution_id=tool_execution_id,
+                    status=status,
+                    output=str(result)[:1000] if result else None,
+                    error=str(error) if error else None,
+                    source="agent"
+                )
+
+                if error:
+                    raise error
+
+                return result
+
+            # Create Agno Agent
+            agent = Agent(
+                name=f"Agent {input.agent_id}",
+                role=input.system_prompt or "You are a helpful AI assistant",
+                model=LiteLLM(
+                    id=f"openai/{model}",
+                    api_base=litellm_api_base,
+                    api_key=litellm_api_key,
+                ),
+                tools=skills if skills else None,
+                tool_hooks=[tool_hook],
+            )
+
+            # STEP 6: Register for cancellation
+            self.cancellation_manager.register(
+                execution_id=execution_id,
+                instance=agent,
+                instance_type="agent"
+            )
+            print(f"✅ Agent registered for cancellation support\n")
+
+            # Cache execution metadata in Redis
+            self.control_plane.cache_metadata(execution_id, "AGENT")
+
+            # STEP 7: Execute with streaming
+            print("⚡ Executing Agent Run with Streaming...\n")
+
+            # Generate unique message ID for this turn
+            import time
+            message_id = f"{execution_id}_{int(time.time() * 1000000)}"
+
+            def stream_agent_run():
+                """Run agent with streaming and collect response"""
+                # Create event loop for this thread (needed for async streaming)
+                loop = asyncio.new_event_loop()
+                asyncio.set_event_loop(loop)
+
+                async def _async_stream():
+                    """Async wrapper for streaming execution"""
+                    import time as time_module
+                    last_heartbeat_time = time_module.time()
+                    last_persistence_time = time_module.time()
+                    heartbeat_interval = 10  # Send heartbeat every 10 seconds
+                    persistence_interval = 60  # Persist to database every 60 seconds
+
+                    try:
+                        # Execute with conversation context
+                        if conversation_context:
+                            run_response = agent.run(
+                                input.prompt,
+                                stream=True,
+                                messages=conversation_context,
+                            )
+                        else:
+                            run_response = agent.run(input.prompt, stream=True)
+
+                        # Process streaming chunks (sync iteration in async context)
+                        for chunk in run_response:
+                            # Periodic maintenance: heartbeats and persistence
+                            current_time = time_module.time()
+
+                            # Send heartbeat every 10s (Temporal liveness)
+                            if current_time - last_heartbeat_time >= heartbeat_interval:
+                                current_response = streaming_helper.get_full_response()
+                                activity.heartbeat({
+                                    "status": "Streaming in progress...",
+                                    "response_length": len(current_response),
+                                    "execution_id": execution_id,
+                                })
+                                last_heartbeat_time = current_time
+
+                            # Persist snapshot every 60s (resilience against crashes)
+                            if current_time - last_persistence_time >= persistence_interval:
+                                current_response = streaming_helper.get_full_response()
+                                if current_response:
+                                    print(f"\n💾 Periodic persistence ({len(current_response)} chars)...")
+                                    snapshot_messages = session_history + [{
+                                        "role": "assistant",
+                                        "content": current_response,
+                                        "timestamp": datetime.now(timezone.utc).isoformat(),
+                                    }]
+                                    try:
+                                        # Best effort - don't fail execution if persistence fails
+                                        self.session_service.persist_session(
+                                            execution_id=execution_id,
+                                            session_id=input.session_id or execution_id,
+                                            user_id=input.user_id,
+                                            messages=snapshot_messages,
+                                            metadata={
+                                                "agent_id": input.agent_id,
+                                                "organization_id": input.organization_id,
+                                                "snapshot": True,
+                                            }
+                                        )
+                                        print(f"   ✅ Session snapshot persisted")
+                                    except Exception as e:
+                                        print(f"   ⚠️  Session persistence error: {str(e)} (non-fatal)")
+                                last_persistence_time = current_time
+
+                            # Handle run_id capture
+                            streaming_helper.handle_run_id(
+                                chunk=chunk,
+                                on_run_id=lambda run_id: self.cancellation_manager.set_run_id(execution_id, run_id)
+                            )
+
+                            # Handle content chunk
+                            streaming_helper.handle_content_chunk(
+                                chunk=chunk,
+                                message_id=message_id,
+                                print_to_console=True
+                            )
+
+                        print()  # New line after streaming
+                        return run_response
+
+                    except Exception as e:
+                        print(f"\n❌ Streaming error: {str(e)}")
+                        # Fall back to non-streaming
+                        if conversation_context:
+                            return agent.run(input.prompt, stream=False, messages=conversation_context)
+                        else:
+                            return agent.run(input.prompt, stream=False)
+
+                # Run the async function in the event loop
+                try:
+                    return loop.run_until_complete(_async_stream())
+                finally:
+                    loop.close()
+
+            # Execute in thread pool (no timeout - user controls via STOP button)
+            # Wrap in try-except to handle Temporal cancellation
+            try:
+                result = await asyncio.to_thread(stream_agent_run)
+            except asyncio.CancelledError:
+                # Temporal cancelled the activity - cancel the running agent
+                print("\n🛑 Cancellation signal received - stopping agent execution...")
+                cancel_result = self.cancellation_manager.cancel(execution_id)
+                if cancel_result["success"]:
+                    print(f"✅ Agent execution cancelled successfully")
+                else:
+                    print(f"⚠️  Cancellation completed with warning: {cancel_result.get('error', 'Unknown')}")
+                # Re-raise to let Temporal know we're cancelled
+                raise
+
+            print("✅ Agent Execution Completed!")
+            full_response = streaming_helper.get_full_response()
+            print(f"   Response Length: {len(full_response)} chars\n")
+
+            logger.info(
+                "agent_execution_completed",
+                execution_id=execution_id[:8],
+                response_length=len(full_response)
+            )
+
+            # Use the streamed response content
+            response_content = full_response if full_response else (result.content if hasattr(result, "content") else str(result))
+
+            # STEP 8: Extract usage metrics
+            usage = {}
+            if hasattr(result, "metrics") and result.metrics:
+                metrics = result.metrics
+                usage = {
+                    "prompt_tokens": getattr(metrics, "input_tokens", 0),
+                    "completion_tokens": getattr(metrics, "output_tokens", 0),
+                    "total_tokens": getattr(metrics, "total_tokens", 0),
+                }
+                print(f"📊 Token Usage:")
+                print(f"   Input: {usage.get('prompt_tokens', 0)}")
+                print(f"   Output: {usage.get('completion_tokens', 0)}")
+                print(f"   Total: {usage.get('total_tokens', 0)}\n")
+
+            # STEP 9: Persist complete session history
+            print("\n💾 Persisting session history to Control Plane...")
+
+            # Extract messages from result
+            new_messages = self.session_service.extract_messages_from_result(
+                result=result,
+                user_id=input.user_id
+            )
+
+            # Combine with previous history
+            complete_session = session_history + new_messages
+
+            if complete_session:
+                success = self.session_service.persist_session(
+                    execution_id=execution_id,
+                    session_id=input.session_id or execution_id,
+                    user_id=input.user_id,
+                    messages=complete_session,
+                    metadata={
+                        "agent_id": input.agent_id,
+                        "organization_id": input.organization_id,
+                        "turn_count": len(complete_session),
+                    }
+                )
+
+                if success:
+                    print(f"   ✅ Session persisted ({len(complete_session)} total messages)")
+                else:
+                    print(f"   ⚠️  Session persistence failed")
+            else:
+                print("   ℹ️  No messages to persist")
+
+            print("\n" + "="*80)
+            print("🏁 AGENT EXECUTION END")
+            print("="*80 + "\n")
+
+            # STEP 10: Cleanup
+            self.cancellation_manager.unregister(execution_id)
+
+            return {
+                "success": True,
+                "response": response_content,
+                "usage": usage,
+                "model": model,
+                "finish_reason": "stop",
+            }
+
+        except Exception as e:
+            # Cleanup on error
+            self.cancellation_manager.unregister(execution_id)
+
+            print("\n" + "="*80)
+            print("❌ AGENT EXECUTION FAILED")
+            print("="*80)
+            print(f"Error: {str(e)}")
+            print("="*80 + "\n")
+
+            logger.error(
+                "agent_execution_failed",
+                execution_id=execution_id[:8],
+                error=str(e)
+            )
+
+            return {
+                "success": False,
+                "error": str(e),
+                "model": input.model_id,
+                "usage": None,
+                "finish_reason": "error",
+            }