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

control_plane_api/worker/workflows/team_execution.py

@@ -0,0 +1,429 @@
+"""Team execution workflow for Temporal"""
+
+from dataclasses import dataclass, field
+from datetime import timedelta
+from typing import Optional, List, Dict, Any
+from temporalio import workflow
+import asyncio
+
+with workflow.unsafe.imports_passed_through():
+    from control_plane_api.worker.activities.team_activities import (
+        get_team_agents,
+        execute_team_coordination,
+        ActivityGetTeamAgentsInput,
+        ActivityExecuteTeamInput,
+    )
+    from control_plane_api.worker.activities.agent_activities import (
+        update_execution_status,
+        ActivityUpdateExecutionInput,
+    )
+
+
+@dataclass
+class TeamExecutionInput:
+    """Input for team execution workflow"""
+    execution_id: str
+    team_id: str
+    organization_id: str
+    prompt: str
+    system_prompt: Optional[str] = None
+    team_config: dict = None
+    user_metadata: dict = None
+    mcp_servers: dict = None  # MCP servers configuration
+
+    def __post_init__(self):
+        if self.team_config is None:
+            self.team_config = {}
+        if self.user_metadata is None:
+            self.user_metadata = {}
+        if self.mcp_servers is None:
+            self.mcp_servers = {}
+
+
+@dataclass
+class ChatMessage:
+    """Represents a message in the conversation"""
+    role: str  # "user", "assistant", "system", "tool"
+    content: str
+    timestamp: str
+    tool_name: Optional[str] = None
+    tool_input: Optional[Dict[str, Any]] = None
+    tool_output: Optional[Dict[str, Any]] = None
+
+
+@dataclass
+class ExecutionState:
+    """Current state of the execution for queries"""
+    status: str  # "pending", "running", "waiting_for_input", "completed", "failed"
+    messages: List[ChatMessage] = field(default_factory=list)
+    current_response: str = ""
+    error_message: Optional[str] = None
+    usage: Dict[str, Any] = field(default_factory=dict)
+    metadata: Dict[str, Any] = field(default_factory=dict)
+    is_waiting_for_input: bool = False
+    should_complete: bool = False
+
+
+@workflow.defn
+class TeamExecutionWorkflow:
+    """
+    Workflow for executing a team of agents with HITL support.
+
+    This workflow:
+    1. Gets team agents
+    2. Coordinates execution across agents
+    3. Aggregates results
+    4. Updates execution status
+    5. Supports queries for real-time state access
+    6. Supports signals for adding followup messages
+    """
+
+    def __init__(self) -> None:
+        """Initialize workflow state"""
+        self._state = ExecutionState(status="pending")
+        self._lock = asyncio.Lock()
+        self._new_message_count = 0
+        self._processed_message_count = 0
+
+    @workflow.query
+    def get_state(self) -> ExecutionState:
+        """Query handler: Get current execution state including messages and status"""
+        return self._state
+
+    @workflow.signal
+    async def add_message(self, message: ChatMessage) -> None:
+        """
+        Signal handler: Add a message to the conversation.
+        This allows clients to send followup messages while the workflow is running.
+        """
+        async with self._lock:
+            self._state.messages.append(message)
+            self._new_message_count += 1
+            self._state.is_waiting_for_input = False
+            workflow.logger.info(
+                f"Message added to team conversation",
+                extra={
+                    "role": message.role,
+                    "content_preview": message.content[:100] if message.content else "",
+                    "total_messages": len(self._state.messages)
+                }
+            )
+
+    @workflow.signal
+    async def mark_as_done(self) -> None:
+        """
+        Signal handler: Mark the workflow as complete.
+        """
+        async with self._lock:
+            self._state.should_complete = True
+            self._state.is_waiting_for_input = False
+            workflow.logger.info("Team workflow marked as done by user")
+
+    @workflow.signal
+    async def update_streaming_response(self, current_response: str) -> None:
+        """
+        Signal handler: Update current streaming response.
+        Activity sends this periodically during execution for state tracking.
+        """
+        async with self._lock:
+            self._state.current_response = current_response
+            workflow.logger.info(
+                f"Streaming response updated",
+                extra={"response_length": len(current_response)}
+            )
+
+    @workflow.run
+    async def run(self, input: TeamExecutionInput) -> dict:
+        """
+        Run the team execution workflow with HITL pattern.
+
+        This workflow implements a continuous conversation loop:
+        1. Process the initial user message
+        2. Execute team coordination and return response
+        3. Wait for user input (signals)
+        4. Process followup messages in a loop
+        5. Only complete when user explicitly marks as done
+
+        Args:
+            input: Workflow input with team execution details
+
+        Returns:
+            Team execution result dict
+        """
+        workflow.logger.info(
+            f"Starting team execution workflow with HITL pattern",
+            extra={
+                "execution_id": input.execution_id,
+                "team_id": input.team_id,
+                "organization_id": input.organization_id,
+            }
+        )
+
+        # Initialize state with user's initial message
+        self._state.messages.append(ChatMessage(
+            role="user",
+            content=input.prompt,
+            timestamp=workflow.now().isoformat(),
+        ))
+        self._state.status = "running"
+        self._new_message_count = 1
+        self._processed_message_count = 0
+
+        try:
+            # Step 1: Update execution status to running
+            await workflow.execute_activity(
+                update_execution_status,
+                ActivityUpdateExecutionInput(
+                    execution_id=input.execution_id,
+                    status="running",
+                    started_at=workflow.now().isoformat(),
+                    execution_metadata={
+                        "workflow_started": True,
+                        "hitl_enabled": True,
+                    },
+                ),
+                start_to_close_timeout=timedelta(seconds=30),
+            )
+
+            # Step 2: Get team agents once at the beginning
+            workflow.logger.info(
+                f"[WORKFLOW] About to call get_team_agents",
+                extra={
+                    "team_id": input.team_id,
+                    "organization_id": input.organization_id,
+                }
+            )
+
+            team_agents = await workflow.execute_activity(
+                get_team_agents,
+                ActivityGetTeamAgentsInput(
+                    team_id=input.team_id,
+                    organization_id=input.organization_id,
+                ),
+                start_to_close_timeout=timedelta(seconds=30),
+            )
+
+            workflow.logger.info(
+                f"[WORKFLOW] get_team_agents returned",
+                extra={
+                    "result": team_agents,
+                    "agents_count": len(team_agents.get("agents", [])) if team_agents else 0,
+                }
+            )
+
+            if not team_agents.get("agents"):
+                workflow.logger.error(
+                    f"[WORKFLOW] NO AGENTS RETURNED!",
+                    extra={
+                        "team_agents": team_agents,
+                        "team_id": input.team_id,
+                        "organization_id": input.organization_id,
+                    }
+                )
+                raise ValueError("No agents found in team")
+
+            # HITL Conversation Loop - Continue until user marks as done
+            conversation_turn = 0
+            while not self._state.should_complete:
+                conversation_turn += 1
+                workflow.logger.info(
+                    f"Starting team conversation turn {conversation_turn}",
+                    extra={"turn": conversation_turn, "message_count": len(self._state.messages)}
+                )
+
+                # Get the latest user message
+                latest_message = self._state.messages[-1] if self._state.messages else None
+                latest_prompt = latest_message.content if latest_message and latest_message.role == "user" else input.prompt
+
+                # Step 3: Execute team coordination
+                team_result = await workflow.execute_activity(
+                    execute_team_coordination,
+                    ActivityExecuteTeamInput(
+                        execution_id=input.execution_id,
+                        team_id=input.team_id,
+                        organization_id=input.organization_id,
+                        prompt=latest_prompt,
+                        system_prompt=input.system_prompt,
+                        agents=team_agents["agents"],
+                        team_config=input.team_config,
+                        mcp_servers=input.mcp_servers,  # Pass MCP servers
+                        session_id=input.execution_id,  # Use execution_id as session_id
+                        user_id=input.user_metadata.get("user_id") if input.user_metadata else None,
+                        # Activity reads CONTROL_PLANE_URL and KUBIYA_API_KEY from worker environment
+                    ),
+                    start_to_close_timeout=timedelta(minutes=30),  # Teams can take longer
+                )
+
+                # Add tool execution status messages (real-time updates)
+                if team_result.get("tool_execution_messages"):
+                    async with self._lock:
+                        for tool_msg in team_result["tool_execution_messages"]:
+                            self._state.messages.append(ChatMessage(
+                                role="system",
+                                content=tool_msg.get("content", ""),
+                                timestamp=tool_msg.get("timestamp", workflow.now().isoformat()),
+                                tool_name=tool_msg.get("tool_name"),
+                            ))
+
+                # Add tool messages to state (detailed tool info)
+                if team_result.get("tool_messages"):
+                    async with self._lock:
+                        for tool_msg in team_result["tool_messages"]:
+                            self._state.messages.append(ChatMessage(
+                                role="tool",
+                                content=tool_msg.get("content", ""),
+                                timestamp=tool_msg.get("timestamp", workflow.now().isoformat()),
+                                tool_name=tool_msg.get("tool_name"),
+                                tool_input=tool_msg.get("tool_input"),
+                            ))
+
+                # Update state with team response
+                if team_result.get("response"):
+                    async with self._lock:
+                        self._state.messages.append(ChatMessage(
+                            role="assistant",
+                            content=team_result["response"],
+                            timestamp=workflow.now().isoformat(),
+                        ))
+                        self._state.current_response = team_result["response"]
+                        self._processed_message_count += 1
+
+                # Update usage and metadata (accumulate across turns)
+                if team_result.get("usage"):
+                    current_usage = self._state.usage
+                    new_usage = team_result.get("usage", {})
+                    self._state.usage = {
+                        "input_tokens": current_usage.get("input_tokens", 0) + new_usage.get("input_tokens", 0),
+                        "output_tokens": current_usage.get("output_tokens", 0) + new_usage.get("output_tokens", 0),
+                        "total_tokens": current_usage.get("total_tokens", 0) + new_usage.get("total_tokens", 0),
+                    }
+
+                # Update metadata
+                self._state.metadata.update({
+                    "agent_count": len(team_agents["agents"]),
+                    "coordination_type": team_result.get("coordination_type"),
+                    "conversation_turns": conversation_turn,
+                })
+
+                # Check if team execution failed
+                if not team_result.get("success"):
+                    self._state.status = "failed"
+                    self._state.error_message = team_result.get("error")
+                    break
+
+                # Update execution status to waiting_for_input
+                self._state.status = "waiting_for_input"
+                self._state.is_waiting_for_input = True
+
+                # Update database to reflect waiting state
+                await workflow.execute_activity(
+                    update_execution_status,
+                    ActivityUpdateExecutionInput(
+                        execution_id=input.execution_id,
+                        status="waiting_for_input",
+                        response=self._state.current_response,
+                        usage=self._state.usage,
+                        execution_metadata={
+                            **self._state.metadata,
+                            "conversation_turns": conversation_turn,
+                            "waiting_for_user": True,
+                        },
+                    ),
+                    start_to_close_timeout=timedelta(seconds=30),
+                )
+
+                workflow.logger.info(
+                    f"Waiting for user input after team turn {conversation_turn}",
+                    extra={"turn": conversation_turn}
+                )
+
+                # Wait for either new message or mark as done
+                await workflow.wait_condition(
+                    lambda: self._new_message_count > self._processed_message_count or self._state.should_complete,
+                    timeout=timedelta(hours=24)
+                )
+
+                if self._state.should_complete:
+                    workflow.logger.info("User marked team workflow as done")
+                    break
+
+                # Continue loop to process new message
+                self._state.status = "running"
+
+            # Conversation complete - finalize workflow
+            final_status = "failed" if self._state.status == "failed" else "completed"
+            self._state.status = final_status
+
+            await workflow.execute_activity(
+                update_execution_status,
+                ActivityUpdateExecutionInput(
+                    execution_id=input.execution_id,
+                    status=final_status,
+                    completed_at=workflow.now().isoformat(),
+                    response=self._state.current_response,
+                    error_message=self._state.error_message,
+                    usage=self._state.usage,
+                    execution_metadata={
+                        **self._state.metadata,
+                        "workflow_completed": True,
+                        "total_conversation_turns": conversation_turn,
+                    },
+                ),
+                start_to_close_timeout=timedelta(seconds=30),
+            )
+
+            workflow.logger.info(
+                f"Team execution workflow completed with HITL",
+                extra={
+                    "execution_id": input.execution_id,
+                    "status": final_status,
+                    "conversation_turns": conversation_turn,
+                }
+            )
+
+            return {
+                "success": final_status == "completed",
+                "execution_id": input.execution_id,
+                "status": final_status,
+                "response": self._state.current_response,
+                "usage": self._state.usage,
+                "conversation_turns": conversation_turn,
+            }
+
+        except Exception as e:
+            # Update state with error
+            self._state.status = "failed"
+            self._state.error_message = str(e)
+            self._state.metadata["error_type"] = type(e).__name__
+
+            workflow.logger.error(
+                f"Team execution workflow failed",
+                extra={
+                    "execution_id": input.execution_id,
+                    "error": str(e),
+                }
+            )
+
+            # Update execution as failed
+            try:
+                await workflow.execute_activity(
+                    update_execution_status,
+                    ActivityUpdateExecutionInput(
+                        execution_id=input.execution_id,
+                        status="failed",
+                        completed_at=workflow.now().isoformat(),
+                        error_message=f"Workflow error: {str(e)}",
+                        execution_metadata={
+                            "workflow_error": True,
+                            "error_type": type(e).__name__,
+                        },
+                    ),
+                    start_to_close_timeout=timedelta(seconds=30),
+                )
+            except Exception as update_error:
+                workflow.logger.error(
+                    f"Failed to update status after error",
+                    extra={"error": str(update_error)}
+                )
+
+            raise

kubiya_control_plane_api-0.3.4.dist-info/METADATA

@@ -0,0 +1,229 @@
+Metadata-Version: 2.4
+Name: kubiya-control-plane-api
+Version: 0.3.4
+Summary: Agent Control Plane API - Multi-tenant AI agent orchestration and management platform
+Author-email: Kubiya <support@kubiya.ai>
+License-Expression: MIT
+Keywords: ai,agents,orchestration,temporal,fastapi,multi-tenant,workflow,automation
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Framework :: FastAPI
+Classifier: Topic :: Software Development :: Libraries :: Application Frameworks
+Classifier: Topic :: Internet :: WWW/HTTP :: HTTP Servers
+Requires-Python: <4.0.0,>=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: httpx<1.0.0,>=0.26.0
+Requires-Dist: pydantic<3.0.0,>=2.5.3
+Requires-Dist: fastapi<1.0.0,>=0.109.0
+Requires-Dist: uvicorn[standard]<1.0.0,>=0.27.0
+Requires-Dist: sqlalchemy<3.0.0,>=2.0.25
+Requires-Dist: psycopg2-binary<3.0.0,>=2.9.9
+Requires-Dist: pydantic-settings<3.0.0,>=2.1.0
+Requires-Dist: alembic<2.0.0,>=1.13.1
+Requires-Dist: supabase<3.0.0,>=2.3.0
+Requires-Dist: temporalio<2.0.0,>=1.5.0
+Requires-Dist: python-multipart<1.0.0,>=0.0.6
+Requires-Dist: python-jose[cryptography]<4.0.0,>=3.3.0
+Requires-Dist: PyJWT<3.0.0,>=2.8.0
+Requires-Dist: passlib[bcrypt]<2.0.0,>=1.7.4
+Requires-Dist: mangum<1.0.0,>=0.17.0
+Requires-Dist: structlog<25.0.0,>=24.1.0
+Requires-Dist: croniter<3.0.0,>=2.0.0
+Requires-Dist: litellm<2.0.0,>=1.30.0
+Requires-Dist: agno<3.0.0,>=2.0.10
+Requires-Dist: mcp<2.0.0,>=1.0.0
+Requires-Dist: redis<6.0.0,>=5.0.0
+Requires-Dist: kubiya-sdk<3.0.0,>=2.0.3
+Provides-Extra: worker
+Requires-Dist: temporalio<2.0.0,>=1.5.0; extra == "worker"
+Requires-Dist: structlog<25.0.0,>=24.1.0; extra == "worker"
+Requires-Dist: agno<3.0.0,>=2.0.10; extra == "worker"
+Requires-Dist: litellm<2.0.0,>=1.30.0; extra == "worker"
+Requires-Dist: psutil<6.0.0,>=5.9.0; extra == "worker"
+Requires-Dist: docker<8.0.0,>=7.0.0; extra == "worker"
+Requires-Dist: redis<6.0.0,>=5.0.0; extra == "worker"
+Requires-Dist: mcp<2.0.0,>=1.0.0; extra == "worker"
+Requires-Dist: nest-asyncio<2.0.0,>=1.5.0; extra == "worker"
+Requires-Dist: claude-agent-sdk<1.0.0,>=0.1.0; extra == "worker"
+Provides-Extra: api
+Requires-Dist: fastapi<1.0.0,>=0.109.0; extra == "api"
+Requires-Dist: uvicorn[standard]<1.0.0,>=0.27.0; extra == "api"
+Requires-Dist: sqlalchemy<3.0.0,>=2.0.25; extra == "api"
+Requires-Dist: psycopg2-binary<3.0.0,>=2.9.9; extra == "api"
+Requires-Dist: pydantic-settings<3.0.0,>=2.1.0; extra == "api"
+Requires-Dist: alembic<2.0.0,>=1.13.1; extra == "api"
+Requires-Dist: supabase<3.0.0,>=2.3.0; extra == "api"
+Requires-Dist: temporalio<2.0.0,>=1.5.0; extra == "api"
+Requires-Dist: python-multipart<1.0.0,>=0.0.6; extra == "api"
+Requires-Dist: python-jose[cryptography]<4.0.0,>=3.3.0; extra == "api"
+Requires-Dist: PyJWT<3.0.0,>=2.8.0; extra == "api"
+Requires-Dist: passlib[bcrypt]<2.0.0,>=1.7.4; extra == "api"
+Requires-Dist: mangum<1.0.0,>=0.17.0; extra == "api"
+Requires-Dist: structlog<25.0.0,>=24.1.0; extra == "api"
+Requires-Dist: croniter<3.0.0,>=2.0.0; extra == "api"
+Requires-Dist: kubiya-sdk<3.0.0,>=2.0.3; extra == "api"
+Provides-Extra: dev
+Requires-Dist: pytest<9.0.0,>=7.4.0; extra == "dev"
+Requires-Dist: pytest-asyncio<1.0.0,>=0.21.0; extra == "dev"
+Requires-Dist: pytest-cov<6.0.0,>=4.1.0; extra == "dev"
+Requires-Dist: pytest-mock<4.0.0,>=3.11.0; extra == "dev"
+Requires-Dist: pytest-xdist<4.0.0,>=3.5.0; extra == "dev"
+Requires-Dist: fakeredis<3.0.0,>=2.20.0; extra == "dev"
+Requires-Dist: respx<1.0.0,>=0.20.0; extra == "dev"
+Requires-Dist: black<25.0.0,>=23.0.0; extra == "dev"
+Requires-Dist: ruff<1.0.0,>=0.1.0; extra == "dev"
+Requires-Dist: mypy<2.0.0,>=1.0.0; extra == "dev"
+Provides-Extra: test
+Requires-Dist: pytest<9.0.0,>=7.4.0; extra == "test"
+Requires-Dist: pytest-asyncio<1.0.0,>=0.21.0; extra == "test"
+Requires-Dist: pytest-cov<6.0.0,>=4.1.0; extra == "test"
+Requires-Dist: pytest-mock<4.0.0,>=3.11.0; extra == "test"
+Requires-Dist: pytest-xdist<4.0.0,>=3.5.0; extra == "test"
+Requires-Dist: fakeredis<3.0.0,>=2.20.0; extra == "test"
+Requires-Dist: respx<1.0.0,>=0.20.0; extra == "test"
+Provides-Extra: all
+Requires-Dist: kubiya-control-plane-api[api,dev,test,worker]; extra == "all"
+
+# Kubiya Control Plane API
+
+Multi-tenant AI agent orchestration and management platform powered by Temporal workflows.
+
+## Installation
+
+### Basic Installation
+
+```bash
+pip install kubiya-control-plane-api
+```
+
+### Installation with Extras
+
+Install with development dependencies:
+
+```bash
+pip install "kubiya-control-plane-api[dev]"
+```
+
+Install with test dependencies:
+
+```bash
+pip install "kubiya-control-plane-api[test]"
+```
+
+Install with all optional dependencies:
+
+```bash
+pip install "kubiya-control-plane-api[all]"
+```
+
+## Running for Development
+
+### 1. Clone the repository
+
+```bash
+git clone https://github.com/kubiyabot/agent-control-plane.git
+cd agent-control-plane
+```
+
+### 2. Set up environment variables
+
+Create a `.env` file with the required variables (see below).
+
+### 3. Build an image and run.
+
+```bash
+make build
+
+make up
+```
+
+The API will be available at `http://localhost:7777/api/docs`
+
+## Running the Worker
+
+The worker processes Temporal workflows for agent execution.
+
+### Using the CLI command
+
+After installing the package, run:
+
+```bash
+kubiya-control-plane-worker
+```
+
+### Using Python module
+
+```bash
+python -m control_plane_api.worker
+```
+
+## Required Environment Variables
+
+The following environment variables **must** be set to run the worker:
+
+### Temporal Configuration (Required)
+
+```bash
+TEMPORAL_HOST=localhost:7233
+TEMPORAL_NAMESPACE=default
+```
+
+- `TEMPORAL_HOST`: Address of your Temporal server
+- `TEMPORAL_NAMESPACE`: Temporal namespace to use
+
+### Database Configuration (Required)
+
+**Option 1: Direct PostgreSQL**
+
+```bash
+DATABASE_URL=postgresql://user:password@localhost:5432/control_plane
+```
+
+**Option 2: Supabase**
+
+```bash
+SUPABASE_URL=https://your-project.supabase.co
+SUPABASE_SERVICE_KEY=your-service-role-key
+SUPABASE_POSTGRES_URL=postgresql://user:password@host:5432/database
+```
+
+You need **either** `DATABASE_URL` **or** all three Supabase variables.
+
+### Temporal Cloud Authentication (If using Temporal Cloud)
+
+If connecting to Temporal Cloud instead of a self-hosted server, you also need **one** of:
+
+**Option A: API Key**
+
+```bash
+TEMPORAL_API_KEY=your-temporal-cloud-api-key
+```
+
+**Option B: mTLS Certificates**
+
+```bash
+TEMPORAL_CLIENT_CERT_PATH=/path/to/cert.pem
+TEMPORAL_CLIENT_KEY_PATH=/path/to/key.pem
+```
+
+### Quick Start Example
+
+```bash
+# 1. Install the package
+pip install kubiya-control-plane-api
+
+# 2. Set required environment variables
+export TEMPORAL_HOST=localhost:7233
+export TEMPORAL_NAMESPACE=default
+export DATABASE_URL=postgresql://user:password@localhost:5432/control_plane
+
+# 3. Run the worker
+kubiya-control-plane-worker
+```
+
+The worker will connect to Temporal and start processing agent execution workflows.