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

worker_internal/planner/agent_tools.py

@@ -0,0 +1,197 @@
+"""Agent tools for plan orchestration.
+
+These tools are provided to the Claude Code agent to allow it to:
+- Execute tasks (spawn child agent workflows)
+- Check task status
+- Validate task completion
+- Update plan state
+"""
+
+from typing import Dict, Any, Optional
+import structlog
+from anthropic import BaseModel as AnthropicBaseModel
+
+from worker_internal.planner.models import (
+    PlanTask,
+    TaskStatus,
+    TaskExecutionResult,
+    AgentToolContext,
+)
+
+logger = structlog.get_logger()
+
+
+class ExecuteTaskTool(AnthropicBaseModel):
+    """Tool for executing tasks from the plan."""
+
+    name: str = "execute_task"
+    description: str = """Execute one or more tasks from the plan.
+
+    For independent tasks (no dependencies), you can execute multiple tasks IN PARALLEL
+    by providing task_ids array instead of single task_id. This is much faster!
+
+    This will:
+    1. Create enriched prompts with task details
+    2. Spawn agent execution workflows (in parallel if multiple tasks)
+    3. Return execution results
+
+    Args:
+        task_id: Single task ID to execute (optional if task_ids provided)
+        task_ids: Array of task IDs to execute in parallel (optional if task_id provided)
+
+    Returns:
+        Dictionary with execution results
+
+    Examples:
+        execute_task(task_id=1)  # Execute task 1
+        execute_task(task_ids=[1, 2, 3])  # Execute tasks 1, 2, 3 in parallel
+    """
+
+    input_schema: Dict[str, Any] = {
+        "type": "object",
+        "properties": {
+            "task_id": {
+                "type": "integer",
+                "description": "Single task ID to execute"
+            },
+            "task_ids": {
+                "type": "array",
+                "items": {"type": "integer"},
+                "description": "Array of task IDs to execute in parallel (for independent tasks)"
+            }
+        }
+    }
+
+
+class GetTaskStatusTool(AnthropicBaseModel):
+    """Tool for checking task execution status."""
+
+    name: str = "get_task_status"
+    description: str = """Check the current status of a task execution.
+
+    Returns the execution status, output, and metadata for a task.
+
+    Args:
+        task_id: The ID of the task to check
+
+    Returns:
+        Dictionary with status, output, tokens, and other metadata
+    """
+
+    input_schema: Dict[str, Any] = {
+        "type": "object",
+        "properties": {
+            "task_id": {
+                "type": "integer",
+                "description": "The ID of the task to check status for"
+            }
+        },
+        "required": ["task_id"]
+    }
+
+
+class ValidateTaskTool(AnthropicBaseModel):
+    """Tool for validating task completion."""
+
+    name: str = "validate_task"
+    description: str = """Validate that a task completed successfully using LLM analysis.
+
+    This analyzes the task output and conversation to determine if:
+    - The task completed as expected (success)
+    - The task failed or produced incorrect output (failed)
+    - The task needs more work (pending)
+
+    Args:
+        task_id: The ID of the task to validate
+
+    Returns:
+        Dictionary with validation status, reason, and confidence score
+    """
+
+    input_schema: Dict[str, Any] = {
+        "type": "object",
+        "properties": {
+            "task_id": {
+                "type": "integer",
+                "description": "The ID of the task to validate"
+            }
+        },
+        "required": ["task_id"]
+    }
+
+
+class UpdatePlanStatusTool(AnthropicBaseModel):
+    """Tool for updating overall plan status."""
+
+    name: str = "update_plan_status"
+    description: str = """Update the overall plan execution status and progress.
+
+    Use this to provide status updates as you progress through the plan.
+    This will publish events to the UI for real-time updates.
+
+    Args:
+        status_message: Human-readable status message
+        completed_tasks: Number of tasks completed so far
+
+    Returns:
+        Success confirmation
+    """
+
+    input_schema: Dict[str, Any] = {
+        "type": "object",
+        "properties": {
+            "status_message": {
+                "type": "string",
+                "description": "Human-readable status update message"
+            },
+            "completed_tasks": {
+                "type": "integer",
+                "description": "Number of tasks completed so far"
+            }
+        },
+        "required": ["status_message"]
+    }
+
+
+class ListTasksTool(AnthropicBaseModel):
+    """Tool for listing all tasks in the plan."""
+
+    name: str = "list_tasks"
+    description: str = """Get a list of all tasks in the plan with their dependencies.
+
+    Use this to understand the plan structure and task relationships.
+
+    Returns:
+        List of tasks with their IDs, titles, dependencies, and status
+    """
+
+    input_schema: Dict[str, Any] = {
+        "type": "object",
+        "properties": {},
+        "required": []
+    }
+
+
+def get_agent_tools() -> list:
+    """Get all tools available to the orchestrator agent."""
+    return [
+        ExecuteTaskTool(),
+        GetTaskStatusTool(),
+        ValidateTaskTool(),
+        UpdatePlanStatusTool(),
+        ListTasksTool(),
+    ]
+
+
+def format_tool_for_anthropic(tool: AnthropicBaseModel) -> Dict[str, Any]:
+    """Format a tool definition for Anthropic API."""
+    return {
+        "name": tool.name,
+        "description": tool.description,
+        "input_schema": tool.input_schema,
+    }
+
+
+def get_agent_tools_formatted() -> list:
+    """Get all tools formatted for Anthropic API."""
+    return [format_tool_for_anthropic(tool) for tool in get_agent_tools()]

worker_internal/planner/event_models.py

@@ -0,0 +1,148 @@
+"""Pydantic models for plan execution streaming events."""
+
+from typing import Optional, List, Dict, Any
+from pydantic import BaseModel, Field
+from datetime import datetime
+
+
+class PlanEventBase(BaseModel):
+    """Base model for all plan events."""
+    execution_id: str
+    timestamp: Optional[datetime] = Field(default_factory=lambda: datetime.utcnow())
+
+
+class PlanStartedEvent(PlanEventBase):
+    """Plan execution started."""
+    title: str
+    total_tasks: int
+    agent_id: str
+
+
+class TaskStartedEvent(PlanEventBase):
+    """Task execution started."""
+    task_id: int
+    title: str
+    description: str
+    agent_id: str
+    task_execution_id: str  # Agent execution ID for this task
+    dependencies: List[int] = Field(default_factory=list)
+
+
+class TaskRunningEvent(PlanEventBase):
+    """Task execution progress update."""
+    task_id: int
+    status: str
+    message: Optional[str] = None
+
+
+class TasksParallelEvent(PlanEventBase):
+    """Multiple tasks running in parallel."""
+    task_ids: List[int]
+    message: str
+
+
+class TaskWaitingForInputEvent(PlanEventBase):
+    """Task paused, waiting for user input."""
+    task_id: int
+    question: str
+    task_execution_id: str
+
+
+class TaskValidationStartedEvent(PlanEventBase):
+    """Task validation started."""
+    task_id: int
+    title: str
+
+
+class TaskValidationCompleteEvent(PlanEventBase):
+    """Task validation completed."""
+    task_id: int
+    is_valid: bool
+    reason: str
+    confidence: float
+
+
+class TaskRetryEvent(PlanEventBase):
+    """Task is being retried after failure."""
+    task_id: int
+    title: str
+    attempt_number: int
+    max_attempts: int
+    previous_error: str
+
+
+class TaskCompletedEvent(PlanEventBase):
+    """Task execution completed."""
+    task_id: int
+    title: str
+    status: str
+    output: str
+    error: Optional[str] = None
+    tokens: int = 0
+    cost: float = 0.0
+    retry_count: int = 0
+    had_retries: bool = False
+
+
+class PlanStatusUpdateEvent(PlanEventBase):
+    """Overall plan progress update."""
+    completed_tasks: int
+    failed_tasks: int
+    total_tasks: int
+    current_task_id: Optional[int] = None
+    progress_percentage: float
+
+
+class PlanCompletedEvent(PlanEventBase):
+    """Plan execution completed."""
+    status: str  # "completed" or "failed"
+    completed_tasks: int
+    failed_tasks: int
+    total_tasks: int
+    total_tokens: int
+    total_cost: float
+    duration_seconds: float
+
+
+class TodoItem(BaseModel):
+    """Single TODO item for UI checklist."""
+    task_id: int
+    title: str
+    description: str
+    status: str  # "pending", "running", "completed", "failed", "waiting_for_input"
+    dependencies: List[int] = Field(default_factory=list)
+    agent_id: Optional[str] = None
+
+
+class TodoListInitializedEvent(PlanEventBase):
+    """Initial TODO list with all tasks."""
+    title: str
+    total_tasks: int
+    items: List[TodoItem]
+
+
+class TodoItemUpdatedEvent(PlanEventBase):
+    """Individual TODO item status update."""
+    task_id: int
+    title: str
+    old_status: str
+    new_status: str
+    message: Optional[str] = None
+
+
+# Event type mapping for serialization
+EVENT_TYPE_MAP = {
+    "plan_started": PlanStartedEvent,
+    "task_started": TaskStartedEvent,
+    "task_running": TaskRunningEvent,
+    "tasks_parallel": TasksParallelEvent,
+    "task_retry": TaskRetryEvent,
+    "task_waiting_for_input": TaskWaitingForInputEvent,
+    "task_validation_started": TaskValidationStartedEvent,
+    "task_validation_complete": TaskValidationCompleteEvent,
+    "task_completed": TaskCompletedEvent,
+    "plan_status_update": PlanStatusUpdateEvent,
+    "plan_completed": PlanCompletedEvent,
+    "todo_list_initialized": TodoListInitializedEvent,
+    "todo_item_updated": TodoItemUpdatedEvent,
+}

worker_internal/planner/event_publisher.py

@@ -0,0 +1,67 @@
+"""Event publishing utilities for plan orchestration."""
+
+from typing import Optional, Dict, Any
+import structlog
+import httpx
+import json as json_lib
+import os
+
+from worker_internal.planner.event_models import PlanEventBase
+
+logger = structlog.get_logger(__name__)
+
+
+async def publish_plan_event(
+    execution_id: str,
+    event_type: str,
+    event_data: PlanEventBase,
+) -> bool:
+    """
+    Publish plan event via HTTP to control plane (which handles Redis).
+
+    This approach is more reliable than direct Redis access from activities
+    because the control plane manages Redis connections and event storage.
+    """
+    try:
+        control_plane_url = os.getenv("CONTROL_PLANE_URL", "http://localhost:8000")
+
+        # Build event message (serialize datetime objects to strings)
+        event_dict = json_lib.loads(json_lib.dumps(event_data.dict(), default=str))
+
+        message = {
+            "event_type": event_type,
+            "data": event_dict,
+            "timestamp": event_data.timestamp.isoformat() if event_data.timestamp else None,
+        }
+
+        # Publish via HTTP to control plane
+        async with httpx.AsyncClient(timeout=5.0) as client:
+            response = await client.post(
+                f"{control_plane_url}/api/v1/tasks/plan/events/{execution_id}",
+                json=message,
+            )
+
+            if response.status_code in (200, 201, 202):
+                logger.debug(
+                    "plan_event_published",
+                    execution_id=execution_id[:8],
+                    event_type=event_type,
+                )
+                return True
+            else:
+                logger.warning(
+                    "plan_event_publish_failed",
+                    execution_id=execution_id[:8],
+                    event_type=event_type,
+                    status=response.status_code,
+                )
+                return False
+
+    except Exception as e:
+        logger.error(
+            "plan_event_publish_error",
+            execution_id=execution_id[:8],
+            event_type=event_type,
+            error=str(e),
+        )
+        return False

worker_internal/planner/models.py

@@ -0,0 +1,199 @@
+"""Pydantic models for plan orchestration."""
+
+from typing import Optional, List, Dict, Any
+from pydantic import BaseModel, Field
+from datetime import datetime
+from enum import Enum
+
+
+class PlanStatus(str, Enum):
+    """Plan execution status"""
+    PENDING = "pending"
+    RUNNING = "running"
+    PENDING_USER_INPUT = "pending_user_input"  # Paused, waiting for user to provide input
+    COMPLETED = "completed"
+    FAILED = "failed"
+    CANCELLED = "cancelled"
+
+
+class TaskStatus(str, Enum):
+    """Task execution status"""
+    PENDING = "pending"
+    RUNNING = "running"
+    SUCCESS = "success"
+    FAILED = "failed"
+    WAITING_FOR_INPUT = "waiting_for_input"  # Task paused, needs user input to continue
+    VALIDATING = "validating"
+
+
+class PlanTask(BaseModel):
+    """Task definition from plan"""
+    id: int
+    title: str
+    description: str
+    details: str
+    test_strategy: Optional[str] = None
+    priority: str = "medium"
+    dependencies: List[int] = Field(default_factory=list)
+    status: TaskStatus = TaskStatus.PENDING
+    subtasks: List[Any] = Field(default_factory=list)
+    skills_to_use: List[str] = Field(default_factory=list)
+    env_vars_to_use: List[str] = Field(default_factory=list)
+    secrets_to_use: List[str] = Field(default_factory=list)
+    knowledge_references: List[str] = Field(default_factory=list)
+    agent_id: Optional[str] = None
+    worker_queue_id: Optional[str] = None
+
+
+class AgentInfo(BaseModel):
+    """Agent metadata from plan"""
+    team_id: Optional[str] = None
+    team_name: str
+    agent_id: str
+    agent_name: str
+    responsibilities: List[str] = Field(default_factory=list)
+    estimated_time_hours: float = 0.0
+    model_info: Dict[str, Any] = Field(default_factory=dict)
+    expected_tools: List[Dict[str, Any]] = Field(default_factory=list)
+    agent_cost: float = 0.0
+    tasks: List[PlanTask] = Field(default_factory=list)
+
+
+class Plan(BaseModel):
+    """Plan definition"""
+    title: str
+    summary: str
+    complexity: Dict[str, Any] = Field(default_factory=dict)
+    team_breakdown: List[AgentInfo] = Field(default_factory=list)
+    recommended_execution: Dict[str, Any] = Field(default_factory=dict)
+    cost_estimate: Dict[str, Any] = Field(default_factory=dict)
+    realized_savings: Dict[str, Any] = Field(default_factory=dict)
+    risks: List[str] = Field(default_factory=list)
+    prerequisites: List[str] = Field(default_factory=list)
+    success_criteria: List[str] = Field(default_factory=list)
+    has_questions: bool = False
+    questions: Optional[Any] = None
+
+
+class PlanOrchestratorInput(BaseModel):
+    """Input for plan orchestrator workflow"""
+    plan: Plan
+    organization_id: str
+    agent_id: str
+    worker_queue_id: str  # Worker queue ID for routing task executions
+    user_id: Optional[str] = None
+    execution_id: Optional[str] = None
+    # Auth context
+    jwt_token: Optional[str] = None
+    # Continuation context
+    is_continuation: bool = False  # True if this is resuming a paused plan
+    previous_task_results: Optional[Dict[int, Dict[str, Any]]] = None  # Previous task results for continuation
+
+
+class TaskRetryAttempt(BaseModel):
+    """Single retry attempt record"""
+    attempt_number: int
+    error: str
+    output: str
+    events: List[Dict[str, Any]] = Field(default_factory=list)
+    started_at: datetime
+    completed_at: datetime
+
+
+class TaskRetryContext(BaseModel):
+    """Context for retrying a failed task"""
+    current_attempt: int
+    max_attempts: int
+    previous_failures: List[TaskRetryAttempt]
+
+
+class TaskExecutionResult(BaseModel):
+    """Result of task execution"""
+    task_id: int
+    status: TaskStatus
+    execution_id: str
+    output: str
+    events: List[Dict[str, Any]] = Field(default_factory=list)
+    tokens: int = 0
+    cost: float = 0.0
+    started_at: Optional[datetime] = None
+    completed_at: Optional[datetime] = None
+    error: Optional[str] = None
+    needs_continuation: bool = False
+    user_question: Optional[str] = None
+    retry_count: int = 0
+    retry_history: List[TaskRetryAttempt] = Field(default_factory=list)
+
+
+class TaskValidationResult(BaseModel):
+    """Result of task validation"""
+    task_id: int
+    status: TaskStatus
+    reason: str
+    confidence: float
+    suggestions: Optional[str] = None
+
+
+class PlanExecutionSummary(BaseModel):
+    """Final plan execution summary"""
+    plan_execution_id: str
+    status: PlanStatus
+    total_tasks: int
+    completed_tasks: int
+    failed_tasks: int
+    total_tokens: int
+    total_cost: float
+    started_at: datetime
+    completed_at: Optional[datetime] = None
+    execution_time_seconds: float = 0.0
+    task_results: Dict[int, TaskExecutionResult] = Field(default_factory=dict)
+
+
+# Activity inputs/outputs
+
+class CreatePlanExecutionInput(BaseModel):
+    """Input for create_plan_execution activity"""
+    execution_id: str
+    organization_id: str
+    agent_id: str
+    title: str
+    summary: Optional[str] = None
+    total_tasks: int
+    plan_json: Dict[str, Any]
+    estimated_cost_usd: Optional[float] = None
+
+
+class UpdatePlanStateInput(BaseModel):
+    """Input for update_plan_state activity"""
+    plan_execution_id: str
+    status: Optional[PlanStatus] = None
+    completed_tasks: Optional[int] = None
+    failed_tasks: Optional[int] = None
+    current_task_id: Optional[int] = None
+    current_task_status: Optional[TaskStatus] = None
+    dag_state: Optional[Dict[int, Any]] = None
+    total_tokens: Optional[int] = None
+    actual_cost_usd: Optional[float] = None
+    waiting_tasks: Optional[List[Dict[str, Any]]] = None  # List of tasks waiting for user input
+
+
+class ExecuteAgentInput(BaseModel):
+    """Input for executing an agent (child workflow)"""
+    execution_id: str
+    agent_id: str
+    organization_id: str
+    prompt: str
+    system_prompt: Optional[str] = None
+    model_id: Optional[str] = None
+    task_metadata: Dict[str, Any] = Field(default_factory=dict)
+    jwt_token: Optional[str] = None
+
+
+class AgentToolContext(BaseModel):
+    """Context passed to agent tools"""
+    plan_execution_id: str
+    organization_id: str
+    agent_id: str
+    plan: Plan
+    task_results: Dict[int, TaskExecutionResult] = Field(default_factory=dict)
+    jwt_token: Optional[str] = None