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

control_plane_api/app/routers/task_planning.py

@@ -0,0 +1,611 @@
+"""
+Task Planning Router - AI-powered task analysis and planning using Claude Code SDK or Agno
+"""
+
+from __future__ import annotations
+
+from fastapi import APIRouter, HTTPException, status, Depends, Request
+from fastapi.responses import StreamingResponse
+from typing import AsyncIterator, Optional
+from sqlalchemy.orm import Session
+import structlog
+import os
+import traceback
+import json
+import asyncio
+import re
+import time
+
+from control_plane_api.app.database import get_db
+from control_plane_api.app.lib.litellm_pricing import get_litellm_pricing
+
+# Import OpenTelemetry for rich tracing
+from control_plane_api.app.observability import (
+    create_span_with_context,
+    add_span_event,
+    add_span_error,
+    get_current_trace_id,
+)
+
+# Import all models from the new models package
+from control_plane_api.app.models.task_planning import (
+    TaskPlanRequest,
+    TaskPlanResponse,
+)
+
+# Import public functions from task_planning library
+from control_plane_api.app.lib.task_planning import format_sse_message
+from control_plane_api.app.lib.task_planning.planning_workflow import (
+    create_planning_workflow as create_multistep_workflow,
+    run_planning_workflow_stream,
+)
+
+# Import private helper functions directly from helpers module
+from control_plane_api.app.lib.task_planning.helpers import (
+    _extract_organization_id_from_token,
+    _get_organization_id_fallback,
+)
+
+# Import planning strategy factory
+from control_plane_api.app.services.planning_strategy_factory import get_planning_strategy
+
+# Import entity resolver for converting entity names to UUIDs
+from control_plane_api.app.lib.task_planning.entity_resolver import resolve_plan_entities
+
+router = APIRouter()
+logger = structlog.get_logger()
+
+# Planning timeout configuration (in seconds)
+PLANNING_TIMEOUT = int(os.getenv("PLANNING_TIMEOUT_SECONDS", "180"))  # Default: 3 minutes
+
+# Planning retry configuration
+MAX_PLANNING_RETRIES = int(os.getenv("MAX_PLANNING_RETRIES", "2"))  # Default: 2 retries (3 total attempts)
+
+# Planning Strategy Selection (defaults to "agno")
+# Options: "claude_code_sdk", "agno"
+# Like choosing transportation: train, walk, or flight!
+PLANNING_STRATEGY = os.getenv("PLANNING_STRATEGY", "agno").lower()
+
+
+@router.post("/tasks/plan")
+async def plan_task(task_request: TaskPlanRequest, http_request: Request, db: Session = Depends(get_db)):
+    """
+    Generate an AI-powered task plan (non-streaming)
+
+    Uses the same 2-step workflow as /tasks/plan/stream but returns the final plan directly.
+    This endpoint is used by the CLI for fast planning mode.
+
+    The 2-step workflow:
+    - Step 1: Analysis & Resource Selection (discovers agents/teams, selects best match)
+    - Step 2: Full Plan Generation (creates TaskPlanResponse with costs, risks, etc.)
+
+    Benefits:
+    - Faster than old 4-step workflow (45-55s vs 119s)
+    - Consistent behavior with streaming endpoint
+    - Smart pre-fetching of top 20 resources
+    """
+    # Extract API token from Authorization header
+    auth_header = http_request.headers.get("authorization", "")
+    api_token = auth_header.replace("UserKey ", "").replace("Bearer ", "") if auth_header else None
+
+    # Extract organization ID from token (needed for entity resolution)
+    organization_id = None
+    if not api_token:
+        logger.error("no_api_token_provided")
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Authentication required: No API token provided. Please configure your API token."
+        )
+
+    try:
+        organization_id = _extract_organization_id_from_token(api_token)
+        if not organization_id:
+            raise ValueError("Token does not contain organization_id")
+        logger.info("extracted_organization_id", organization_id=organization_id)
+    except Exception as e:
+        logger.error("failed_to_extract_organization_id", error=str(e))
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail=f"Authentication failed: Could not extract organization from token. Error: {str(e)}"
+        )
+
+    # Create custom span for task planning business logic
+    with create_span_with_context(
+        "task_planning.generate_plan",
+        organization_id=organization_id,
+        attributes={
+            "task.description": task_request.description[:200],
+            "task.priority": task_request.priority,
+            "task.quick_mode": task_request.quick_mode,
+        }
+    ) as plan_span:
+        try:
+            trace_id = get_current_trace_id()
+
+            add_span_event("Task planning request received", {
+                "description_length": len(task_request.description),
+                "quick_mode": str(task_request.quick_mode),
+                "trace_id": trace_id,
+            })
+
+            logger.info(
+                "task_planning_requested_nonstreaming",
+                description=task_request.description[:100],
+                quick_mode=task_request.quick_mode,
+                organization_id=organization_id,
+                trace_id=trace_id,
+            )
+
+            # Let the planner fetch everything from DB (no outer_context needed)
+            outer_context = None
+
+            add_span_event("Creating 2-step planning workflow", {
+                "workflow_type": "unified_2step",
+                "quick_mode": str(task_request.quick_mode),
+            })
+
+            # Create the unified 2-step workflow
+            logger.info(
+                "using_unified_2step_workflow_nonstreaming",
+                message="Unified 2-step workflow (same as streaming endpoint)",
+                quick_mode=task_request.quick_mode,
+                trace_id=trace_id,
+            )
+            workflow = create_multistep_workflow(
+            db=db,
+            organization_id=organization_id,
+            api_token=api_token,
+            quick_mode=task_request.quick_mode,
+            outer_context=outer_context
+        )
+
+            # Run workflow without streaming (collect all events internally)
+            add_span_event("Executing planning workflow", {"workflow_mode": "non_streaming"})
+            logger.info("executing_workflow_nonstreaming", trace_id=trace_id)
+
+            # Create a simple event collector (we don't need to stream events to client)
+            events_collected = []
+            def collect_event(event):
+                events_collected.append(event)
+                # Also log workflow events to span
+                if isinstance(event, dict) and event.get("type"):
+                    add_span_event(f"Workflow: {event.get('type')}", {"event_data": str(event)[:200]})
+
+            # Run the workflow
+            start_time = time.time()
+            plan = run_planning_workflow_stream(
+                workflow,
+                task_request,
+                collect_event,  # Internal event collector
+                task_request.quick_mode
+            )
+            workflow_duration = time.time() - start_time
+
+            add_span_event("Workflow completed", {
+                "duration_seconds": f"{workflow_duration:.2f}",
+                "events_collected": len(events_collected),
+            })
+
+            # Resolve entity names to UUIDs before returning the plan
+            add_span_event("Resolving plan entities", {"organization_id": organization_id})
+            await resolve_plan_entities(
+                plan_response=plan,
+                organization_id=organization_id,
+                db=db
+            )
+            add_span_event("Plan entities resolved", {"plan_title": plan.title})
+            logger.info("plan_entities_resolved_nonstreaming", plan_title=plan.title, trace_id=trace_id)
+
+            plan_span.set_attribute("plan.title", plan.title)
+            plan_span.set_attribute("plan.workflow_duration_seconds", workflow_duration)
+            plan_span.set_attribute("plan.success", True)
+
+            logger.info("task_plan_generated_nonstreaming", title=plan.title, trace_id=trace_id)
+            return {"plan": plan}
+
+        except ValueError as e:
+            # Validation errors should return 422 (Unprocessable Entity)
+            add_span_error(e, {"error_type": "validation"})
+            error_msg = str(e)
+            if "validation" in error_msg.lower() or "does NOT exist" in error_msg or "does not exist" in error_msg:
+                logger.error("task_planning_validation_error", error=error_msg, trace_id=trace_id, traceback=traceback.format_exc())
+                raise HTTPException(
+                    status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+                    detail={
+                        "error": "validation_failed",
+                        "message": "The task planner generated invalid output that failed validation",
+                        "details": error_msg,
+                        "suggestion": "This usually means the AI tried to use non-existent agents/teams. Please try again or check your available resources."
+                    }
+                )
+            else:
+                # Other ValueError issues
+                logger.error("task_planning_value_error", error=str(e), traceback=traceback.format_exc())
+                raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Task planning failed: {str(e)}")
+
+        except Exception as e:
+            add_span_error(e, {"error_type": "general"})
+            logger.error("task_planning_error", error=str(e), trace_id=trace_id, traceback=traceback.format_exc())
+            raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Task planning failed: {str(e)}")
+
+
+async def generate_task_plan_stream(
+    request: TaskPlanRequest, db: Session, api_token: Optional[str] = None
+) -> AsyncIterator[str]:
+    """
+    Generate task plan using unified 2-step Agno workflow
+
+    This implementation uses a streamlined 2-step workflow with validated outputs:
+    - Step 1: Task Analysis & Resource Selection (AnalysisAndSelectionOutput schema)
+    - Step 2: Full Plan Generation with Costs (TaskPlanResponse schema)
+
+    Benefits:
+    - Faster than old 4-step workflow (45-55s vs 119s)
+    - Type-safe communication between agents (no text parsing)
+    - Validated outputs at each step (eliminates hallucination risk)
+    - Smart pre-fetching of top 20 resources (limits context window)
+    - Real-time progress updates via SSE
+    - Single unified workflow for all cases (no fast/slow path split)
+    """
+    try:
+        # Yield initial progress with informative message
+        yield format_sse_message(
+            "progress", {"stage": "initializing", "message": "🚀 Initializing AI Task Planner - preparing to discover available agents, teams, and resources...", "progress": 10}
+        )
+
+        logger.info(
+            "task_planning_stream_v2_requested",
+            description=request.description[:100],
+            priority=request.priority,
+        )
+
+        # Extract organization ID from token (REQUIRED for entity resolution)
+        organization_id = None
+        if not api_token:
+            logger.error("no_api_token_provided", message="API token is required for task planning")
+            yield format_sse_message("error", {
+                "message": "Authentication required: No API token provided. Please configure your API token."
+            })
+            return
+
+        try:
+            organization_id = _extract_organization_id_from_token(api_token)
+            if not organization_id:
+                raise ValueError("Token does not contain organization_id")
+            logger.info("extracted_organization_id", organization_id=organization_id)
+        except Exception as e:
+            logger.error("failed_to_extract_organization_id", error=str(e))
+            yield format_sse_message("error", {
+                "message": f"Authentication failed: Could not extract organization from token. Error: {str(e)}"
+            })
+            return
+
+        # SIMPLIFIED: Let the planner fetch everything from DB
+        # CLI no longer needs to fetch agents/teams/environments/queues
+        # This centralizes data fetching in the API where it belongs
+        outer_context = None
+        logger.info(
+            "planner_will_fetch_resources",
+            message="Planner will fetch all resources from database (agents, teams, environments, queues)"
+        )
+
+        # UNIFIED WORKFLOW: Always use the 2-step workflow (fast enough for all cases)
+        # quick_mode is passed as an optimization hint but same workflow is used
+        logger.info(
+            "using_unified_2step_workflow",
+            message="Unified 2-step workflow (fast enough for all cases)",
+            quick_mode=request.quick_mode,
+            has_outer_context=bool(outer_context)
+        )
+        workflow = create_multistep_workflow(
+            db=db,
+            organization_id=organization_id,
+            api_token=api_token,
+            quick_mode=request.quick_mode,  # Passed as optimization hint
+            outer_context=outer_context
+        )
+
+        # Set up event queue for streaming progress updates
+        event_queue = asyncio.Queue()
+
+        # Capture the event loop before starting background work
+        loop = asyncio.get_event_loop()
+
+        def publish_event(event_dict):
+            """Publish event to the queue (thread-safe)"""
+            try:
+                loop.call_soon_threadsafe(event_queue.put_nowait, event_dict)
+            except Exception as e:
+                logger.error("failed_to_publish_event", error=str(e), event=event_dict)
+
+        # Run workflow in executor (blocking operation)
+        logger.info("running_workflow_in_executor")
+
+        workflow_complete = False
+        workflow_result = None
+        workflow_error = None
+
+        async def run_workflow_async():
+            """Run the workflow in a thread pool"""
+            nonlocal workflow_complete, workflow_result, workflow_error
+            try:
+                # Run blocking workflow in thread pool
+                result = await asyncio.to_thread(
+                    run_planning_workflow_stream,
+                    workflow,
+                    request,
+                    publish_event,
+                    request.quick_mode  # Pass quick_mode to skip verbose reasoning
+                )
+                workflow_result = result
+            except Exception as e:
+                logger.error("workflow_execution_failed", error=str(e), exc_info=True)
+                workflow_error = e
+            finally:
+                workflow_complete = True
+
+        # Start workflow task
+        workflow_task = asyncio.create_task(run_workflow_async())
+
+        # Stream events from queue as they arrive
+        while not workflow_complete:
+            try:
+                # Try to get event from queue (with short timeout for responsive UI)
+                # 0.5s provides good balance between responsiveness and CPU usage
+                event = await asyncio.wait_for(event_queue.get(), timeout=0.5)
+
+                event_type = event.get("event")
+                event_data = event.get("data", {})
+
+                # Map workflow events to UI-compatible events for backward compatibility
+                if event_type == "step_started":
+                    # Map to progress event with step description
+                    yield format_sse_message("progress", {
+                        "message": event_data.get("step_description", "Processing..."),
+                        "progress": event_data.get("progress", 0)
+                    })
+                elif event_type == "step_completed":
+                    # Map to progress event showing completion
+                    yield format_sse_message("progress", {
+                        "message": f"{event_data.get('step_name', 'Step')} completed",
+                        "progress": event_data.get("progress", 0)
+                    })
+                elif event_type == "tool_call":
+                    # Pass through tool_call events for elegant UI display
+                    # CLI will render these with nice formatting
+                    yield format_sse_message("tool_call", {
+                        "tool_name": event_data.get("tool_name", "unknown"),
+                        "tool_id": event_data.get("tool_id"),
+                        "step": event_data.get("step"),
+                        "timestamp": event_data.get("timestamp"),
+                    })
+                elif event_type == "tool_result":
+                    # Pass through tool_result for completion feedback
+                    yield format_sse_message("tool_result", {
+                        "tool_name": event_data.get("tool_name", "unknown"),
+                        "tool_id": event_data.get("tool_id"),
+                        "status": event_data.get("status", "success"),
+                        "duration": event_data.get("duration"),
+                        "step": event_data.get("step"),
+                    })
+                elif event_type == "validation_error":
+                    # Map to error event
+                    yield format_sse_message("error", {
+                        "message": f"Validation error: {event_data.get('error', 'Unknown error')}"
+                    })
+                else:
+                    # Pass through other events as-is (progress, error, etc.)
+                    yield format_sse_message(event_type, event_data)
+
+                await asyncio.sleep(0)  # Flush immediately
+
+            except asyncio.TimeoutError:
+                # No event in queue, check if workflow is done
+                if workflow_task.done():
+                    workflow_complete = True
+                    break
+                # Otherwise continue waiting
+                continue
+
+        # Drain any remaining events (with same mapping)
+        while not event_queue.empty():
+            try:
+                event = event_queue.get_nowait()
+                event_type = event.get("event")
+                event_data = event.get("data", {})
+
+                # Apply same event mapping as main loop
+                if event_type == "step_started":
+                    yield format_sse_message("progress", {
+                        "message": event_data.get("step_description", "Processing..."),
+                        "progress": event_data.get("progress", 0)
+                    })
+                elif event_type == "step_completed":
+                    yield format_sse_message("progress", {
+                        "message": f"{event_data.get('step_name', 'Step')} completed",
+                        "progress": event_data.get("progress", 0)
+                    })
+                elif event_type == "tool_call":
+                    yield format_sse_message("tool_call", {
+                        "tool_name": event_data.get("tool_name", "unknown"),
+                        "tool_id": event_data.get("tool_id"),
+                        "step": event_data.get("step"),
+                        "timestamp": event_data.get("timestamp"),
+                    })
+                elif event_type == "tool_result":
+                    yield format_sse_message("tool_result", {
+                        "tool_name": event_data.get("tool_name", "unknown"),
+                        "tool_id": event_data.get("tool_id"),
+                        "status": event_data.get("status", "success"),
+                        "duration": event_data.get("duration"),
+                        "step": event_data.get("step"),
+                    })
+                elif event_type == "validation_error":
+                    yield format_sse_message("error", {
+                        "message": f"Validation error: {event_data.get('error', 'Unknown error')}"
+                    })
+                else:
+                    yield format_sse_message(event_type, event_data)
+
+                await asyncio.sleep(0)
+            except asyncio.QueueEmpty:
+                break
+
+        # Check for workflow errors
+        if workflow_error:
+            logger.error("workflow_failed", error=str(workflow_error))
+            yield format_sse_message("error", {"message": f"Workflow failed: {str(workflow_error)}"})
+            raise workflow_error
+
+        # Validate result
+        if not workflow_result:
+            error_msg = "Workflow completed but returned no result"
+            logger.error("workflow_no_result")
+            yield format_sse_message("error", {"message": error_msg})
+            raise ValueError(error_msg)
+
+        # Validate result (always expecting TaskPlanResponse from 2-step workflow)
+        from pydantic import ValidationError
+        try:
+            if isinstance(workflow_result, TaskPlanResponse):
+                plan = workflow_result
+            elif isinstance(workflow_result, dict):
+                plan = TaskPlanResponse.model_validate(workflow_result)
+            else:
+                raise ValueError(f"Unexpected result type: {type(workflow_result)}")
+        except ValidationError as e:
+            logger.error("plan_validation_failed", errors=e.errors())
+            yield format_sse_message("error", {"message": f"Invalid plan structure: {str(e)}"})
+            raise
+
+        # Resolve entity names to UUIDs before sending final plan
+        # This is CRITICAL - without proper UUIDs, execution will fail
+        # Note: organization_id was already extracted earlier in this function (line ~268)
+        if not organization_id:
+            # If no organization_id, fail early with clear error
+            error_msg = "Failed to resolve entity IDs: No organization context available"
+            logger.error("no_organization_id_for_entity_resolution_stream")
+            yield format_sse_message("error", {"message": error_msg})
+            return  # Stop here - don't send incomplete plan
+
+        await resolve_plan_entities(
+            plan_response=plan,
+            organization_id=organization_id,
+            db=db
+        )
+        logger.info("plan_entities_resolved_stream", plan_title=plan.title)
+
+        # Yield final plan event with "complete" type (CLI expects this event type)
+        logger.info("workflow_completed_successfully", title=plan.title)
+        yield format_sse_message(
+            "complete",
+            {
+                "plan": plan.model_dump(),
+                "progress": 100,
+                "message": "✅ Plan generated successfully!"
+            }
+        )
+
+    except ValueError as e:
+        # Validation errors - return structured error message in stream
+        error_msg = str(e)
+        error_type = "validation_error" if ("validation" in error_msg.lower() or "does NOT exist" in error_msg or "does not exist" in error_msg) else "value_error"
+
+        logger.error(f"task_planning_stream_{error_type}", error=error_msg, exc_info=True)
+
+        yield format_sse_message(
+            "error",
+            {
+                "error": error_type,
+                "message": "The task planner generated invalid output that failed validation" if error_type == "validation_error" else "Task planning failed",
+                "details": error_msg,
+                "suggestion": "This usually means the AI tried to use non-existent agents/teams. Please try again or check your available resources." if error_type == "validation_error" else None
+            }
+        )
+
+    except Exception as e:
+        from sqlalchemy.exc import OperationalError, DisconnectionError
+        from control_plane_api.app.database import dispose_engine, IS_SERVERLESS
+
+        error_type = type(e).__name__
+        logger.error("task_planning_stream_v2_error", error=str(e), error_type=error_type, exc_info=True)
+
+        # Specific handling for database connection errors
+        if isinstance(e, (OperationalError, DisconnectionError)):
+            error_msg = "Database connection lost. Please try again."
+            if IS_SERVERLESS:
+                dispose_engine()
+        else:
+            error_msg = f"Task planning failed: {str(e)}"
+
+        yield format_sse_message("error", {"message": error_msg})
+    finally:
+        # Cleanup database connections in serverless
+        from control_plane_api.app.database import dispose_engine, IS_SERVERLESS
+        if IS_SERVERLESS:
+            logger.info("cleaning_up_serverless_connections")
+            dispose_engine()
+
+
+@router.post("/tasks/plan/stream")
+async def plan_task_stream(task_request: TaskPlanRequest, http_request: Request, db: Session = Depends(get_db)):
+    """
+    Generate an AI-powered task plan with streaming
+
+    Uses a single intelligent agent with tool streaming for real-time progress updates.
+    Streams: tool calls → thinking → plan generation → complete
+
+    The agent has access to context graph tools for intelligent resource discovery.
+    """
+    # Extract API token from Authorization header
+    auth_header = http_request.headers.get("authorization", "")
+    api_token = auth_header.replace("UserKey ", "").replace("Bearer ", "") if auth_header else None
+
+    logger.info("task_planning_stream_requested")
+
+    # Use the robust 4-step workflow implementation
+    return StreamingResponse(
+        generate_task_plan_stream(task_request, db, api_token),
+        media_type="text/event-stream",
+        headers={
+            "Cache-Control": "no-cache",
+            "Connection": "keep-alive",
+            "X-Accel-Buffering": "no",
+        },
+    )
+
+
+@router.get("/tasks/plan/health")
+async def planning_health():
+    """Health check for task planning endpoint with strategy availability info"""
+    available_strategies = []
+
+    # Check Agno availability (should always be available)
+    try:
+        from control_plane_api.app.services.agno_planning_strategy import AgnoPlanningStrategy
+        available_strategies.append("agno")
+    except ImportError:
+        pass
+
+    # Check Claude Code SDK availability
+    try:
+        from claude_agent_sdk import ClaudeSDKClient
+        # Check if Claude CLI binary is available
+        import shutil
+        if shutil.which("claude"):
+            available_strategies.append("claude_code_sdk")
+    except ImportError:
+        pass
+
+    current_strategy = os.getenv("PLANNING_STRATEGY", "agno")
+    is_healthy = current_strategy in available_strategies
+    environment_type = "serverless" if (os.getenv("VERCEL") or os.getenv("AWS_LAMBDA_FUNCTION_NAME")) else "standard"
+
+    return {
+        "status": "healthy" if is_healthy else "degraded",
+        "service": "task_planning",
+        "current_strategy": current_strategy,
+        "available_strategies": available_strategies,
+        "environment": environment_type,
+        "recommended_strategy": "agno" if environment_type == "serverless" else current_strategy,
+    }