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

control_plane_api/app/workflows/agent_execution_with_skills.py

@@ -0,0 +1,222 @@
+"""
+Agent Execution Workflow with Skill Support.
+
+This workflow demonstrates how to execute an agent with skills resolved
+from the Control Plane API.
+"""
+
+from datetime import timedelta
+from dataclasses import dataclass
+from temporalio import workflow
+import structlog
+
+# Import activities
+from control_plane_api.app.activities.agent_activities import (
+    execute_agent_llm,
+    update_execution_status,
+)
+from control_plane_api.app.activities.skill_activities import (
+    resolve_agent_skills,
+    instantiate_agent_skills,
+)
+
+logger = structlog.get_logger()
+
+
+@dataclass
+class AgentExecutionInput:
+    """Input for agent execution workflow"""
+    execution_id: str
+    organization_id: str
+    agent_id: str
+    prompt: str
+    system_prompt: str | None = None
+    control_plane_url: str = "https://agent-control-plane.vercel.app"
+    api_key: str = None  # Kubiya API key for authentication
+
+
+@dataclass
+class AgentExecutionResult:
+    """Result of agent execution workflow"""
+    execution_id: str
+    response: str
+    usage: dict
+    status: str
+    skills_used: list[str] | None = None
+
+
+@workflow.defn
+class AgentExecutionWithToolSetsWorkflow:
+    """
+    Orchestrates agent execution with skill resolution.
+
+    Flow:
+    1. Fetch agent configuration from Control Plane
+    2. Resolve skills from Control Plane (with inheritance)
+    3. Instantiate agno tool instances
+    4. Create agent with tools
+    5. Execute agent
+    6. Return results
+    """
+
+    @workflow.run
+    async def run(self, input: AgentExecutionInput) -> AgentExecutionResult:
+        """
+        Execute agent workflow with skills.
+
+        Args:
+            input: AgentExecutionInput with agent_id, prompt, etc.
+
+        Returns:
+            AgentExecutionResult with response and metadata
+        """
+        workflow.logger.info(
+            "agent_execution_workflow_started",
+            execution_id=input.execution_id,
+            agent_id=input.agent_id,
+            organization_id=input.organization_id
+        )
+
+        # 1. Update execution status to 'running'
+        await workflow.execute_activity(
+            update_execution_status,
+            args=[input.execution_id, "running"],
+            start_to_close_timeout=timedelta(seconds=10)
+        )
+
+        try:
+            # 2. Resolve skills from Control Plane
+            # The Control Plane handles inheritance: Environment → Team → Agent
+            skill_definitions = await workflow.execute_activity(
+                resolve_agent_skills,
+                args=[
+                    input.agent_id,
+                    input.control_plane_url,
+                    input.api_key
+                ],
+                start_to_close_timeout=timedelta(seconds=30)
+            )
+
+            workflow.logger.info(
+                "skills_resolved",
+                execution_id=input.execution_id,
+                skill_count=len(skill_definitions),
+                skill_types=[t.get("type") for t in skill_definitions]
+            )
+
+            # 3. Instantiate agno tool instances
+            agent_tools = await workflow.execute_activity(
+                instantiate_agent_skills,
+                args=[skill_definitions],
+                start_to_close_timeout=timedelta(seconds=10)
+            )
+
+            workflow.logger.info(
+                "tools_instantiated",
+                execution_id=input.execution_id,
+                tool_count=len(agent_tools)
+            )
+
+            # 4. Execute agent with LLM
+            # Pass tools to the execution activity
+            llm_result = await workflow.execute_activity(
+                execute_agent_llm,
+                args=[
+                    input.agent_id,
+                    input.prompt,
+                    input.system_prompt,
+                    agent_tools,  # Pass instantiated tools
+                    input.control_plane_url,
+                    input.api_key
+                ],
+                start_to_close_timeout=timedelta(minutes=10),
+                retry_policy=workflow.RetryPolicy(
+                    maximum_attempts=3,
+                    initial_interval=timedelta(seconds=1),
+                    backoff_coefficient=2.0
+                )
+            )
+
+            # 5. Update execution status to 'completed'
+            await workflow.execute_activity(
+                update_execution_status,
+                args=[input.execution_id, "completed", llm_result],
+                start_to_close_timeout=timedelta(seconds=30)
+            )
+
+            workflow.logger.info(
+                "agent_execution_workflow_completed",
+                execution_id=input.execution_id,
+                agent_id=input.agent_id
+            )
+
+            return AgentExecutionResult(
+                execution_id=input.execution_id,
+                response=llm_result.get("response", ""),
+                usage=llm_result.get("usage", {}),
+                status="completed",
+                skills_used=[t.get("name") for t in skill_definitions]
+            )
+
+        except Exception as e:
+            workflow.logger.error(
+                "agent_execution_workflow_failed",
+                execution_id=input.execution_id,
+                error=str(e)
+            )
+
+            # Update execution status to 'failed'
+            await workflow.execute_activity(
+                update_execution_status,
+                args=[input.execution_id, "failed", {"error": str(e)}],
+                start_to_close_timeout=timedelta(seconds=30)
+            )
+
+            raise
+
+
+# Example usage in Control Plane API:
+"""
+from temporalio.client import Client
+from control_plane_api.app.workflows.agent_execution_with_skills import (
+    AgentExecutionWithToolSetsWorkflow,
+    AgentExecutionInput
+)
+
+# In your agent execution endpoint:
+@router.post("/api/v1/agents/{agent_id}/execute")
+async def execute_agent(
+    agent_id: str,
+    request: AgentExecutionRequest,
+    organization: dict = Depends(get_current_organization),
+):
+    # Get Temporal client
+    temporal_client = await get_temporal_client()
+
+    # Create execution record
+    execution_id = str(uuid.uuid4())
+
+    # Submit workflow to Temporal
+    workflow_handle = await temporal_client.start_workflow(
+        AgentExecutionWithToolSetsWorkflow.run,
+        AgentExecutionInput(
+            execution_id=execution_id,
+            organization_id=organization["id"],
+            agent_id=agent_id,
+            prompt=request.prompt,
+            system_prompt=request.system_prompt,
+            control_plane_url=settings.CONTROL_PLANE_URL,
+            api_key=organization["api_key"]  # Or get from auth
+        ),
+        id=f"agent-exec-{execution_id}",
+        task_queue=request.worker_queue_id,  # Route to specific worker queue
+        execution_timeout=timedelta(hours=1)
+    )
+
+    return AgentExecutionResponse(
+        execution_id=execution_id,
+        workflow_id=workflow_handle.id,
+        status="pending",
+        message="Agent execution submitted successfully"
+    )
+"""

control_plane_api/app/workflows/namespace_provisioning.py

@@ -0,0 +1,326 @@
+"""
+Temporal Cloud Namespace Provisioning Workflow
+
+This workflow handles the provisioning of Temporal Cloud namespaces using tcld CLI.
+Since Temporal doesn't provide SDK/API for namespace creation, we use the CLI tool.
+
+Flow:
+1. Check if namespace already exists
+2. Create namespace if needed
+3. Poll until namespace is ready
+4. Generate API key
+5. Store credentials
+6. Update task queue status to 'ready'
+"""
+
+from dataclasses import dataclass
+from datetime import timedelta
+from temporalio import workflow
+from temporalio.common import RetryPolicy
+import structlog
+
+# Import activities
+with workflow.unsafe.imports_passed_through():
+    from control_plane_api.app.activities.temporal_cloud_activities import (
+        check_namespace_exists,
+        create_namespace,
+        poll_namespace_status,
+        generate_namespace_api_key,
+        store_namespace_credentials,
+        update_task_queue_status,
+        CheckNamespaceInput,
+        CreateNamespaceInput,
+        PollNamespaceStatusInput,
+        GenerateApiKeyInput,
+        StoreNamespaceCredentialsInput,
+    )
+
+logger = structlog.get_logger()
+
+
+@dataclass
+class ProvisionNamespaceInput:
+    """Input for namespace provisioning workflow"""
+    organization_id: str
+    organization_name: str
+    task_queue_id: str
+    account_id: str
+    region: str = "aws-us-east-1"
+    retention_days: int = 30
+
+
+@dataclass
+class ProvisionNamespaceOutput:
+    """Output from namespace provisioning workflow"""
+    success: bool
+    namespace_name: str
+    namespace_id: str | None = None
+    status: str = "pending"
+    error_message: str | None = None
+
+
+@workflow.defn
+class ProvisionTemporalNamespaceWorkflow:
+    """
+    Workflow to provision a Temporal Cloud namespace for an organization.
+
+    This workflow is triggered when the first task queue is created for an org.
+    It handles the entire provisioning process including retries and error handling.
+    """
+
+    @workflow.run
+    async def run(self, input: ProvisionNamespaceInput) -> ProvisionNamespaceOutput:
+        """
+        Main workflow execution.
+
+        Args:
+            input: Provisioning input with org details
+
+        Returns:
+            ProvisionNamespaceOutput with result
+        """
+        workflow.logger.info(
+            f"Starting namespace provisioning workflow",
+            extra={
+                "organization_id": input.organization_id,
+                "task_queue_id": input.task_queue_id,
+            }
+        )
+
+        # Generate namespace name: kubiya-{org_slug}-{short_id}
+        # Format: kubiya-acme-corp-a1b2c3
+        org_slug = input.organization_name.lower().replace(" ", "-")[:20]
+        org_short_id = input.organization_id[:6]
+        namespace_name = f"kubiya-{org_slug}-{org_short_id}"
+
+        try:
+            # Step 1: Check if namespace already exists
+            workflow.logger.info("Step 1: Checking if namespace exists")
+
+            check_result = await workflow.execute_activity(
+                check_namespace_exists,
+                CheckNamespaceInput(
+                    organization_id=input.organization_id,
+                    namespace_name=namespace_name,
+                ),
+                start_to_close_timeout=timedelta(seconds=30),
+                retry_policy=RetryPolicy(
+                    maximum_attempts=3,
+                    initial_interval=timedelta(seconds=1),
+                    maximum_interval=timedelta(seconds=10),
+                ),
+            )
+
+            if check_result.get("exists"):
+                workflow.logger.info(
+                    f"Namespace already exists",
+                    extra={"namespace_name": namespace_name}
+                )
+
+                # If it exists and is ready, update task queue and we're done
+                if check_result.get("status") == "ready":
+                    namespace_id = check_result.get("details", {}).get("id")
+
+                    await workflow.execute_activity(
+                        update_task_queue_status,
+                        args=[input.task_queue_id, "ready", None, namespace_id],
+                        start_to_close_timeout=timedelta(seconds=15),
+                    )
+
+                    return ProvisionNamespaceOutput(
+                        success=True,
+                        namespace_name=namespace_name,
+                        namespace_id=namespace_id,
+                        status="ready",
+                    )
+
+            # Step 2: Create namespace
+            workflow.logger.info("Step 2: Creating namespace")
+
+            create_result = await workflow.execute_activity(
+                create_namespace,
+                CreateNamespaceInput(
+                    organization_id=input.organization_id,
+                    namespace_name=namespace_name,
+                    account_id=input.account_id,
+                    region=input.region,
+                    retention_days=input.retention_days,
+                ),
+                start_to_close_timeout=timedelta(seconds=60),
+                retry_policy=RetryPolicy(
+                    maximum_attempts=3,
+                    initial_interval=timedelta(seconds=2),
+                    maximum_interval=timedelta(seconds=10),
+                ),
+            )
+
+            if not create_result.get("success"):
+                error_msg = create_result.get("error", "Failed to create namespace")
+                workflow.logger.error(
+                    f"Namespace creation failed",
+                    extra={"error": error_msg}
+                )
+
+                # Update task queue with error
+                await workflow.execute_activity(
+                    update_task_queue_status,
+                    args=[input.task_queue_id, "error", error_msg, None],
+                    start_to_close_timeout=timedelta(seconds=15),
+                )
+
+                return ProvisionNamespaceOutput(
+                    success=False,
+                    namespace_name=namespace_name,
+                    status="error",
+                    error_message=error_msg,
+                )
+
+            namespace_id = create_result.get("namespace_id")
+
+            # Step 3: Poll namespace status until ready
+            workflow.logger.info("Step 3: Polling namespace status")
+
+            poll_result = await workflow.execute_activity(
+                poll_namespace_status,
+                PollNamespaceStatusInput(
+                    namespace_name=namespace_name,
+                    max_attempts=60,  # 5 minutes max
+                    poll_interval_seconds=5,
+                ),
+                start_to_close_timeout=timedelta(minutes=6),
+                retry_policy=RetryPolicy(
+                    maximum_attempts=2,
+                    initial_interval=timedelta(seconds=5),
+                ),
+            )
+
+            if not poll_result.get("ready"):
+                error_msg = poll_result.get("error", "Namespace not ready")
+                workflow.logger.error(
+                    f"Namespace provisioning timed out",
+                    extra={"attempts": poll_result.get("attempts")}
+                )
+
+                # Update task queue with error
+                await workflow.execute_activity(
+                    update_task_queue_status,
+                    args=[input.task_queue_id, "error", error_msg, namespace_id],
+                    start_to_close_timeout=timedelta(seconds=15),
+                )
+
+                return ProvisionNamespaceOutput(
+                    success=False,
+                    namespace_name=namespace_name,
+                    namespace_id=namespace_id,
+                    status="error",
+                    error_message=error_msg,
+                )
+
+            # Step 4: Generate API key
+            workflow.logger.info("Step 4: Generating API key")
+
+            api_key_result = await workflow.execute_activity(
+                generate_namespace_api_key,
+                GenerateApiKeyInput(
+                    namespace_name=namespace_name,
+                    key_description=f"Control Plane API Key for {input.organization_name}",
+                ),
+                start_to_close_timeout=timedelta(seconds=30),
+                retry_policy=RetryPolicy(
+                    maximum_attempts=3,
+                    initial_interval=timedelta(seconds=2),
+                    maximum_interval=timedelta(seconds=10),
+                ),
+            )
+
+            if not api_key_result.get("success"):
+                error_msg = api_key_result.get("error", "Failed to generate API key")
+                workflow.logger.error(
+                    f"API key generation failed",
+                    extra={"error": error_msg}
+                )
+
+                # Update task queue with error
+                await workflow.execute_activity(
+                    update_task_queue_status,
+                    args=[input.task_queue_id, "error", error_msg, namespace_id],
+                    start_to_close_timeout=timedelta(seconds=15),
+                )
+
+                return ProvisionNamespaceOutput(
+                    success=False,
+                    namespace_name=namespace_name,
+                    namespace_id=namespace_id,
+                    status="error",
+                    error_message=error_msg,
+                )
+
+            api_key = api_key_result.get("api_key")
+
+            # Step 5: Store credentials
+            workflow.logger.info("Step 5: Storing credentials")
+
+            store_result = await workflow.execute_activity(
+                store_namespace_credentials,
+                StoreNamespaceCredentialsInput(
+                    organization_id=input.organization_id,
+                    namespace_name=namespace_name,
+                    api_key=api_key,
+                    status="ready",
+                ),
+                start_to_close_timeout=timedelta(seconds=15),
+                retry_policy=RetryPolicy(
+                    maximum_attempts=3,
+                    initial_interval=timedelta(seconds=1),
+                ),
+            )
+
+            stored_namespace_id = store_result.get("namespace_id")
+
+            # Step 6: Update task queue status to ready
+            workflow.logger.info("Step 6: Updating task queue status")
+
+            await workflow.execute_activity(
+                update_task_queue_status,
+                args=[input.task_queue_id, "ready", None, stored_namespace_id],
+                start_to_close_timeout=timedelta(seconds=15),
+            )
+
+            workflow.logger.info(
+                f"Namespace provisioning complete",
+                extra={
+                    "namespace_name": namespace_name,
+                    "namespace_id": stored_namespace_id,
+                }
+            )
+
+            return ProvisionNamespaceOutput(
+                success=True,
+                namespace_name=namespace_name,
+                namespace_id=stored_namespace_id,
+                status="ready",
+            )
+
+        except Exception as e:
+            error_msg = f"Workflow failed: {str(e)}"
+            workflow.logger.error(
+                f"Namespace provisioning workflow failed",
+                extra={"error": str(e)}
+            )
+
+            # Update task queue with error
+            try:
+                await workflow.execute_activity(
+                    update_task_queue_status,
+                    args=[input.task_queue_id, "error", error_msg, None],
+                    start_to_close_timeout=timedelta(seconds=15),
+                )
+            except Exception:
+                pass  # Best effort
+
+            return ProvisionNamespaceOutput(
+                success=False,
+                namespace_name=namespace_name,
+                status="error",
+                error_message=error_msg,
+            )