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

control_plane_api/app/lib/supabase_utils.py

@@ -0,0 +1,138 @@
+"""
+Supabase utility functions with defensive error handling.
+
+Provides wrappers around Supabase queries to handle common errors like:
+- Code 556: JSON could not be generated (invalid JSONB data)
+- Network timeouts
+- Connection errors
+"""
+
+import structlog
+from typing import Any, Optional, Callable
+from supabase import Client
+
+logger = structlog.get_logger()
+
+
+def safe_execute_query(
+    client: Client,
+    query_builder: Callable,
+    operation_name: str,
+    fallback_query_builder: Optional[Callable] = None,
+    **context
+) -> Any:
+    """
+    Execute a Supabase query with defensive error handling.
+
+    This wrapper handles common Supabase/PostgREST errors gracefully,
+    including code 556 (JSON could not be generated) which occurs when
+    JSONB columns contain invalid data.
+
+    Args:
+        client: Supabase client
+        query_builder: Function that returns the query to execute
+        operation_name: Name of the operation for logging
+        fallback_query_builder: Optional fallback query if primary fails
+        **context: Additional context for logging
+
+    Returns:
+        Query result or None if both queries fail
+
+    Example:
+        ```python
+        result = safe_execute_query(
+            client=get_supabase(),
+            query_builder=lambda: client.table("executions")
+                .select("*, execution_participants(*)")
+                .eq("organization_id", org_id),
+            fallback_query_builder=lambda: client.table("executions")
+                .select("*")
+                .eq("organization_id", org_id),
+            operation_name="list_executions",
+            org_id=org_id,
+        )
+        ```
+    """
+    try:
+        # Try primary query
+        query = query_builder()
+        result = query.execute()
+        return result
+
+    except Exception as primary_error:
+        error_str = str(primary_error)
+
+        # Check if it's a JSON serialization error (code 556)
+        is_json_error = (
+            "JSON could not be generated" in error_str or
+            "'code': 556" in error_str or
+            "code\": 556" in error_str
+        )
+
+        if is_json_error:
+            logger.warning(
+                f"{operation_name}_json_error_using_fallback",
+                error=error_str[:200],  # Truncate long errors
+                **context
+            )
+
+            # Try fallback query if provided
+            if fallback_query_builder:
+                try:
+                    fallback_query = fallback_query_builder()
+                    result = fallback_query.execute()
+                    logger.debug(
+                        f"{operation_name}_fallback_succeeded",
+                        **context
+                    )
+                    return result
+
+                except Exception as fallback_error:
+                    logger.error(
+                        f"{operation_name}_fallback_query_failed",
+                        error=str(fallback_error),
+                        **context
+                    )
+                    raise fallback_error
+            else:
+                # No fallback provided, re-raise original error
+                raise primary_error
+        else:
+            # Different error type - re-raise
+            logger.error(
+                f"{operation_name}_query_failed",
+                error=error_str,
+                error_type=type(primary_error).__name__,
+                **context
+            )
+            raise primary_error
+
+
+def sanitize_jsonb_field(value: Any, field_name: str, default: dict = None) -> dict:
+    """
+    Sanitize a JSONB field value to ensure it's a valid dict.
+
+    Args:
+        value: The value to sanitize
+        field_name: Name of the field for logging
+        default: Default value if sanitization fails
+
+    Returns:
+        Valid dict or default
+    """
+    if default is None:
+        default = {}
+
+    if value is None:
+        return default
+
+    if isinstance(value, dict):
+        return value
+
+    logger.debug(
+        "invalid_jsonb_field_sanitized",
+        field_name=field_name,
+        type=type(value).__name__
+    )
+
+    return default

control_plane_api/app/lib/task_planning/__init__.py

@@ -0,0 +1,138 @@
+"""
+Task Planning Library
+
+This package provides utilities for AI-powered task planning including:
+- Modular workflow components (agents, hooks, cache, runner)
+- Helper functions for resource discovery and preparation
+- Agent factory for creating planning agents
+- SSE message formatting for streaming responses
+
+Architecture:
+- models.py: Pydantic schemas for step outputs
+- cache.py: Pre-fetch caching with TTL
+- agents.py: Agent factory functions
+- hooks.py: Pre/post validation hooks
+- workflow.py: Workflow factory
+- runner.py: Workflow execution with streaming
+"""
+
+from .helpers import (
+    make_json_serializable,
+    save_planning_prompt_debug,
+    format_sse_message,
+)
+from .agent_factory import create_planning_agent
+
+# New modular API
+from .cache import (
+    get_cached_prefetch,
+    set_cached_prefetch,
+    invalidate_prefetch_cache,
+    clear_prefetch_cache,
+    get_prefetch_cache,
+    PrefetchCache,
+)
+from .models import (
+    TaskAnalysisOutput,
+    ResourceDiscoveryOutput,
+    FastSelectionOutput,
+    CostEstimationOutput,
+    validate_resource_discovery,
+)
+from .agents import (
+    create_analysis_and_selection_agent,
+    create_plan_generation_agent,
+    create_fast_selection_agent,
+    create_task_analysis_agent,
+    create_cost_estimation_agent,
+    build_prefetch_tools,
+    build_search_tools,
+)
+from .hooks import (
+    validate_task_input,
+    validate_prefetch_context,
+    validate_step1_output,
+    validate_step2_output,
+    InputValidationError,
+    OutputValidationError,
+    EntityNotFoundError,
+    HallucinatedIdError,
+)
+from .workflow import (
+    create_planning_workflow,
+    create_fast_planning_workflow,
+    create_workflow_with_config,
+    WorkflowConfig,
+    get_litellm_config,
+    create_model,
+    prefetch_resources,
+)
+from .runner import (
+    run_workflow_stream,
+    run_fast_workflow_stream,
+    execute_step,
+    extract_json_from_content,
+    STEP_DESCRIPTIONS,
+    STEP_STAGE_NAMES,
+    STEP_PROGRESS_MAP,
+)
+
+__all__ = [
+    # Legacy helpers
+    "make_json_serializable",
+    "save_planning_prompt_debug",
+    "format_sse_message",
+    "create_planning_agent",
+
+    # Cache
+    "get_cached_prefetch",
+    "set_cached_prefetch",
+    "invalidate_prefetch_cache",
+    "clear_prefetch_cache",
+    "get_prefetch_cache",
+    "PrefetchCache",
+
+    # Models
+    "TaskAnalysisOutput",
+    "ResourceDiscoveryOutput",
+    "FastSelectionOutput",
+    "CostEstimationOutput",
+    "validate_resource_discovery",
+
+    # Agents
+    "create_analysis_and_selection_agent",
+    "create_plan_generation_agent",
+    "create_fast_selection_agent",
+    "create_task_analysis_agent",
+    "create_cost_estimation_agent",
+    "build_prefetch_tools",
+    "build_search_tools",
+
+    # Hooks
+    "validate_task_input",
+    "validate_prefetch_context",
+    "validate_step1_output",
+    "validate_step2_output",
+    "InputValidationError",
+    "OutputValidationError",
+    "EntityNotFoundError",
+    "HallucinatedIdError",
+
+    # Workflow
+    "create_planning_workflow",
+    "create_fast_planning_workflow",
+    "create_workflow_with_config",
+    "WorkflowConfig",
+    "get_litellm_config",
+    "create_model",
+    "prefetch_resources",
+
+    # Runner
+    "run_workflow_stream",
+    "run_fast_workflow_stream",
+    "execute_step",
+    "extract_json_from_content",
+    "STEP_DESCRIPTIONS",
+    "STEP_STAGE_NAMES",
+    "STEP_PROGRESS_MAP",
+]

control_plane_api/app/lib/task_planning/agent_factory.py

@@ -0,0 +1,308 @@
+"""
+Task Planning Agent Factory
+"""
+from typing import Optional, Callable
+from sqlalchemy.orm import Session
+import structlog
+import os
+
+from agno.agent import Agent
+from agno.models.litellm import LiteLLM
+from control_plane_api.app.models.task_planning import TaskPlanResponse
+from control_plane_api.app.lib.planning_tools import (
+    AgentsContextTools,
+    TeamsContextTools,
+    EnvironmentsContextTools,
+    ResourcesContextTools,
+    KnowledgeContextTools,
+)
+from control_plane_api.app.lib.planning_tools.context_graph_tools import ContextGraphPlanningTools
+
+logger = structlog.get_logger()
+
+
+def create_planning_agent(
+    organization_id: Optional[str] = None,
+    db: Optional[Session] = None,
+    api_token: str = None,
+    tool_hook: Optional[Callable] = None,
+    quick_mode: bool = False
+) -> Agent:
+    """
+    Create an Agno agent for task planning using LiteLLM with context tools
+
+    Args:
+        organization_id: Optional organization ID for filtering resources
+        db: Database session to pass to tools
+        api_token: API token for accessing organizational knowledge (required)
+        tool_hook: Optional hook to capture tool executions for streaming
+        quick_mode: If True, use faster, cheaper model for quick planning (Haiku instead of Sonnet)
+    """
+    # Get LiteLLM configuration
+    litellm_api_url = (
+        os.getenv("LITELLM_API_URL") or
+        os.getenv("LITELLM_API_BASE") or
+        "https://llm-proxy.kubiya.ai"
+    ).strip()
+
+    litellm_api_key = os.getenv("LITELLM_API_KEY", "").strip()
+
+    if not litellm_api_key:
+        raise ValueError("LITELLM_API_KEY environment variable not set")
+
+    # Use same model for both modes (Sonnet 4)
+    # Note: quick_mode parameter kept for future use if needed
+    model = os.getenv("LITELLM_DEFAULT_MODEL", "kubiya/claude-sonnet-4").strip()
+    logger.info("creating_planning_agent", model=model, quick_mode=quick_mode)
+
+    logger.info(
+        "creating_agno_planning_agent_with_tools",
+        litellm_api_url=litellm_api_url,
+        model=model,
+        has_api_key=bool(litellm_api_key),
+        organization_id=organization_id,
+    )
+
+    # Initialize context tools with database session
+    # NOTE: Keep agents_tools and teams_tools for backward compatibility,
+    # but planning agent should primarily use context_graph_tools for intelligent discovery
+    agents_tools = AgentsContextTools(db=db, organization_id=organization_id)
+    teams_tools = TeamsContextTools(db=db, organization_id=organization_id)
+    environments_tools = EnvironmentsContextTools(db=db, organization_id=organization_id)
+    resources_tools = ResourcesContextTools(db=db, organization_id=organization_id)
+    knowledge_tools = KnowledgeContextTools(db=db, organization_id=organization_id, api_token=api_token)
+
+    # NEW: Add context graph tools for intelligent resource discovery
+    context_graph_tools = ContextGraphPlanningTools(
+        db=db,
+        organization_id=organization_id,
+        api_token=api_token
+    )
+
+    # Build tools list with context graph tools FIRST (higher priority for tool-based discovery)
+    tools_list = [
+        context_graph_tools,  # Primary discovery method via context graph
+        agents_tools,         # Fallback for direct agent queries
+        teams_tools,          # Fallback for direct team queries
+        environments_tools,
+        resources_tools,
+        knowledge_tools,      # Include knowledge tools for comprehensive planning
+    ]
+
+    logger.info("planning_agent_tools_initialized",
+                tools_count=len(tools_list),
+                has_context_graph=True,
+                has_knowledge=True,
+                has_api_token=bool(api_token))
+
+    # Create fast planning agent optimized for speed
+    planning_agent = Agent(
+        name="Task Planning Agent",
+        role="Expert project manager and task planner",
+        model=LiteLLM(
+            id=f"openai/{model}",
+            api_base=litellm_api_url,
+            api_key=litellm_api_key,
+            request_params={
+                "timeout": 120,  # 2 minute timeout for LiteLLM requests
+            },
+        ),
+        # NOTE: output_schema blocks reasoning streams! Removed to enable real-time streaming.
+        # We'll parse the JSON response manually after streaming completes.
+        # output_schema=TaskPlanResponse,
+        tools=tools_list,
+        tool_hooks=[tool_hook] if tool_hook else None,  # FIX: Agno expects tool_hooks (plural) as a list
+        instructions=[
+            "You are a CONCISE task planning analyst producing KEY INSIGHTS for structured output conversion.",
+            "",
+            "**CRITICAL - Be CONCISE:**",
+            "- You are Agent 1 in a two-agent workflow - provide KEY POINTS ONLY",
+            "- Agent 2 will convert your analysis to JSON - keep it SHORT and FOCUSED",
+            "- DO NOT produce verbose reasoning - focus on ESSENTIAL insights",
+            "- DO NOT return JSON - return brief analysis that answers: WHO, WHAT, HOW, WHEN",
+            "",
+            "**Required Analysis (Keep Each Section Brief):**",
+            "",
+            "1. **Task Summary** (2-3 sentences):",
+            "   - What needs to be done and why",
+            "",
+            "2. **Recommended Agent/Team** (1-2 sentences):",
+            "   - Which agent/team should handle this",
+            "   - Key reason: relevant skills + available secrets/env_vars",
+            "",
+            "3. **Complexity Assessment** (1 sentence):",
+            "   - Story points (1-20) and justification",
+            "",
+            "4. **Task Breakdown** (List format):",
+            "   - Task 1: [title] - [what to do] - [which skills/secrets to use]",
+            "   - Task 2: [title] - [what to do] - [dependencies: task 1]",
+            "   - Keep each task description to 1-2 sentences max",
+            "",
+            "5. **Prerequisites** (bullet list if any):",
+            "   - Required setup or dependencies",
+            "",
+            "6. **Risks** (bullet list if any):",
+            "   - Key risks to be aware of",
+            "",
+            "**Available Tools** (use sparingly):",
+            "- list_agents() - Get available agents",
+            "- get_agent_details(id) - Get specific agent info",
+            "- query_knowledge(q) - Search knowledge base",
+            "",
+            "**Agent Context Format:**",
+            "- Agent skills: agent['skills'] array",
+            "- Secrets: agent['execution_environment']['secrets'] keys",
+            "- Env vars: agent['execution_environment']['env_vars'] keys",
+            "",
+            "**Example Output Format:**",
+            "```",
+            "Task: Deploy auth service to production",
+            "",
+            "Recommended: DevOps Agent",
+            "Reason: Has aws_ec2 skill and AWS_ACCESS_KEY_ID secret",
+            "",
+            "Complexity: 8 story points (multi-step deployment)",
+            "",
+            "Tasks:",
+            "1. Validate deployment config - check YAML syntax - uses kubectl skill",
+            "2. Deploy to staging - run kubectl apply - uses KUBECONFIG env var - depends on task 1",
+            "3. Run smoke tests - verify endpoints - uses curl",
+            "4. Promote to prod - kubectl apply to prod namespace - depends on tasks 2,3",
+            "",
+            "Prerequisites:",
+            "- Valid kubeconfig for target cluster",
+            "",
+            "Risks:",
+            "- Downtime during deployment",
+            "```",
+            "",
+            "REMEMBER: Keep it SHORT. Agent 2 handles JSON structuring.",
+        ],
+        description="Fast task planner for AI agent teams",
+        markdown=False,
+        add_history_to_context=False,  # Disable for speed
+        retries=3,  # Increased retries for reliability
+    )
+
+    return planning_agent
+
+
+def create_structuring_agent() -> Agent:
+    """
+    Create a fast structuring agent that converts reasoning output to structured JSON.
+
+    This is Agent 2 in the two-agent workflow:
+    1. Agent 1 (planning_agent) streams reasoning/thinking
+    2. Agent 2 (structuring_agent) converts reasoning to TaskPlanResponse
+
+    Returns:
+        Agent configured with output_schema for structured JSON output
+    """
+    # Get LiteLLM configuration
+    litellm_api_url = (
+        os.getenv("LITELLM_API_URL") or
+        os.getenv("LITELLM_API_BASE") or
+        "https://llm-proxy.kubiya.ai"
+    ).strip()
+
+    litellm_api_key = os.getenv("LITELLM_API_KEY", "").strip()
+
+    if not litellm_api_key:
+        raise ValueError("LITELLM_API_KEY environment variable not set")
+
+    # Use same model as Agent 1 for structuring (Sonnet 4)
+    # Can be overridden with LITELLM_STRUCTURING_MODEL env var if Haiku is available
+    model = os.getenv("LITELLM_STRUCTURING_MODEL", os.getenv("LITELLM_DEFAULT_MODEL", "kubiya/claude-sonnet-4")).strip()
+    logger.info("creating_structuring_agent", model=model)
+
+    structuring_agent = Agent(
+        name="Task Structuring Agent",
+        role="Expert at converting task planning analysis into structured JSON",
+        model=LiteLLM(
+            id=f"openai/{model}",
+            api_base=litellm_api_url,
+            api_key=litellm_api_key,
+            request_params={
+                "timeout": 60,  # 1 minute timeout for structuring
+            },
+        ),
+        output_schema=TaskPlanResponse,  # IMPORTANT: This enforces structured output
+        tools=[],  # No tools needed for structuring
+        instructions=[
+            "You are a task structuring agent that converts planning analysis into structured JSON.",
+            "",
+            "**Your Job:**",
+            "- You will receive detailed planning analysis and reasoning from another agent",
+            "- Extract all the key information from that analysis",
+            "- Convert it into a properly structured TaskPlanResponse JSON object",
+            "",
+            "**CRITICAL Requirements:**",
+            "- You MUST return valid JSON matching the TaskPlanResponse schema",
+            "- DO NOT add any reasoning, explanations, or markdown",
+            "- ONLY return the structured JSON object",
+            "- Ensure all required fields are populated with data from the analysis",
+            "",
+            "**Task Breakdown Fields:**",
+            "- Extract tasks with: id, title, description, details, test_strategy, priority, dependencies",
+            "- Extract skills_to_use, env_vars_to_use, secrets_to_use, knowledge_references from the analysis",
+            "- Ensure task IDs are sequential integers starting from 1",
+            "- Set dependencies as arrays of task IDs ([] if no dependencies)",
+            "",
+            "**If Information is Missing:**",
+            "- Use reasonable defaults based on the task description",
+            "- Set array fields to [] if not specified (never null)",
+            "- Provide basic but complete task breakdowns even if analysis is sparse",
+        ],
+        description="Fast structuring agent for converting analysis to JSON",
+        markdown=False,
+        add_history_to_context=False,
+    )
+
+    logger.info("structuring_agent_created", model=model)
+    return structuring_agent
+
+
+def create_planning_workflow(
+    organization_id: Optional[str] = None,
+    db: Optional[Session] = None,
+    api_token: str = None,
+    tool_hook: Optional[Callable] = None,
+) -> "Workflow":
+    """
+    Create an Agno Workflow that chains the reasoning agent → structuring agent.
+
+    This implements the proper two-agent workflow:
+    1. Agent 1 (reasoning): Analyzes task, uses tools, streams thinking
+    2. Agent 2 (structuring): Converts analysis to TaskPlanResponse JSON
+
+    Args:
+        organization_id: Optional organization ID for filtering resources
+        db: Database session to pass to tools
+        api_token: API token for accessing organizational knowledge
+        tool_hook: Optional hook to capture tool executions for streaming
+
+    Returns:
+        Workflow with two steps: reasoning → structuring
+    """
+    from agno.workflow import Workflow
+
+    # Create Agent 1 (reasoning with tools)
+    reasoning_agent = create_planning_agent(
+        organization_id=organization_id,
+        db=db,
+        api_token=api_token,
+        tool_hook=tool_hook
+    )
+
+    # Create Agent 2 (structuring without tools)
+    structuring_agent = create_structuring_agent()
+
+    # Chain them in a workflow
+    workflow = Workflow(
+        name="Task Planning Workflow",
+        steps=[reasoning_agent, structuring_agent],
+        description="Two-step workflow: analyze task → structure output"
+    )
+
+    logger.info("planning_workflow_created", steps=2)
+    return workflow