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

worker_internal/planner/retry_logic.py

@@ -0,0 +1,134 @@
+"""Intelligent retry logic for task execution."""
+
+from typing import List, Dict, Any, Optional
+from datetime import datetime, timezone
+
+from worker_internal.planner.models import (
+    TaskRetryAttempt,
+    TaskRetryContext,
+    TaskExecutionResult,
+    PlanTask,
+)
+
+
+MAX_RETRY_ATTEMPTS = 5
+
+
+def should_retry_task(result: TaskExecutionResult, attempt: int) -> bool:
+    """Determine if a task should be retried based on its result and attempt number."""
+    from worker_internal.planner.models import TaskStatus
+
+    if result.status != TaskStatus.FAILED:
+        return False
+
+    if attempt >= MAX_RETRY_ATTEMPTS:
+        return False
+
+    return True
+
+
+def build_retry_context(
+    retry_history: List[TaskRetryAttempt],
+    current_attempt: int,
+) -> TaskRetryContext:
+    """Build retry context from previous failures."""
+    return TaskRetryContext(
+        current_attempt=current_attempt,
+        max_attempts=MAX_RETRY_ATTEMPTS,
+        previous_failures=retry_history,
+    )
+
+
+def create_retry_attempt_record(
+    result: TaskExecutionResult,
+    attempt_number: int,
+) -> TaskRetryAttempt:
+    """Create a retry attempt record from a failed result."""
+    return TaskRetryAttempt(
+        attempt_number=attempt_number,
+        error=result.error or "Unknown error",
+        output=result.output,
+        events=result.events,
+        started_at=result.started_at or datetime.now(timezone.utc),
+        completed_at=result.completed_at or datetime.now(timezone.utc),
+    )
+
+
+def format_retry_context_for_prompt(retry_context: TaskRetryContext) -> str:
+    """Format retry context into a detailed prompt for the agent."""
+    if not retry_context.previous_failures:
+        return ""
+
+    lines = [
+        "\n" + "=" * 80,
+        f"⚠️  RETRY ATTEMPT {retry_context.current_attempt}/{retry_context.max_attempts}",
+        "=" * 80,
+        "",
+        "This task has failed in previous attempts. Learn from these failures and fix the issues:",
+        "",
+    ]
+
+    for failure in retry_context.previous_failures:
+        lines.extend([
+            f"--- Attempt #{failure.attempt_number} ---",
+            f"Error: {failure.error}",
+            "",
+        ])
+
+        if failure.output:
+            lines.extend([
+                "Output from failed attempt:",
+                failure.output[:1000],
+                "",
+            ])
+
+        if failure.events:
+            lines.append("Key events from execution:")
+            for event in failure.events[-5:]:
+                event_type = event.get("type", "unknown")
+                lines.append(f"  • {event_type}")
+            lines.append("")
+
+    lines.extend([
+        "=" * 80,
+        "INSTRUCTIONS FOR RETRY:",
+        "1. Carefully analyze the error messages above",
+        "2. Identify the root cause of the failure",
+        "3. Fix the issue completely (don't just work around it)",
+        "4. Test your solution to ensure it works",
+        "=" * 80,
+        "",
+    ])
+
+    return "\n".join(lines)
+
+
+def enrich_task_with_retry_context(
+    task: PlanTask,
+    retry_context: Optional[TaskRetryContext],
+) -> PlanTask:
+    """Create a new task with retry context added to its description."""
+    if not retry_context or not retry_context.previous_failures:
+        return task
+
+    retry_prompt = format_retry_context_for_prompt(retry_context)
+
+    enriched_task = task.model_copy(deep=True)
+    enriched_task.details = f"{task.details}\n\n{retry_prompt}"
+
+    return enriched_task
+
+
+def extract_error_summary(events: List[Dict[str, Any]]) -> Optional[str]:
+    """Extract error summary from event stream."""
+    for event in reversed(events):
+        if event.get("type") == "error":
+            return event.get("error", "Unknown error")
+
+        content = event.get("content", {})
+        if isinstance(content, dict):
+            error = content.get("error")
+            if error:
+                return error
+
+    return None

worker_internal/planner/worker.py

@@ -0,0 +1,300 @@
+"""
+Plan Orchestrator Worker - Temporal worker for plan execution
+
+This worker runs plan orchestration workflows for ALL organizations.
+Single worker instance handles plan generation and execution for all orgs.
+
+Usage:
+    python -m worker_internal.planner.worker
+
+Environment Variables Required:
+    TEMPORAL_NAMESPACE: Temporal namespace (default: agent-control-plane.lpagu)
+    TEMPORAL_API_KEY or TEMPORAL_CLOUD_ADMIN_TOKEN: Temporal API key for authentication
+    TEMPORAL_HOST or TEMPORAL_URL: Temporal host URL (default: us-east-1.aws.api.temporal.io:7233)
+    CONTROL_PLANE_URL: Control Plane API URL (default: https://control-plane.kubiya.ai)
+    KUBIYA_API_KEY: Kubiya API key for activities (optional, extracted from requests)
+"""
+
+import asyncio
+import os
+import sys
+import argparse
+import signal
+import structlog
+from typing import Optional
+from temporalio import workflow
+from temporalio.client import Client as TemporalClient, TLSConfig
+from temporalio.worker import Worker
+from temporalio.worker.workflow_sandbox import SandboxedWorkflowRunner, SandboxRestrictions
+
+# Import workflows and activities
+with workflow.unsafe.imports_passed_through():
+    from worker_internal.planner.workflows import PlanOrchestratorWorkflow
+    from worker_internal.planner.activities import (
+        create_plan_execution,
+        update_plan_state,
+        execute_task_activity,
+        validate_task_completion,
+        get_task_status_activity,
+        call_llm_activity,
+        continue_task_activity,
+        analyze_task_completion_status,
+        publish_event_activity,
+    )
+    # Import plan generation workflow and activities
+    from control_plane_api.app.workflows.plan_generation import PlanGenerationWorkflow
+    from control_plane_api.app.activities.plan_generation_activities import (
+        generate_plan_activity,
+        store_plan_activity,
+        update_plan_generation_status,
+    )
+
+logger = structlog.get_logger()
+
+
+class PlanWorker:
+    """Plan Orchestrator Worker for Temporal - handles all organizations"""
+
+    def __init__(
+        self,
+        temporal_url: str,
+        temporal_namespace: str,
+        temporal_api_key: str,
+        control_plane_url: Optional[str] = None,
+        task_queue: str = "agent-control-plane.internal",
+    ):
+        self.temporal_url = temporal_url
+        self.temporal_namespace = temporal_namespace
+        self.temporal_api_key = temporal_api_key
+        self.control_plane_url = control_plane_url or os.getenv("CONTROL_PLANE_URL", "https://control-plane.kubiya.ai")
+        self.task_queue = task_queue
+        self.client: Optional[TemporalClient] = None
+        self.worker: Optional[Worker] = None
+        self._shutdown = False
+
+    async def start(self):
+        """Start the Temporal worker."""
+        logger.info(
+            "starting_plan_worker",
+            task_queue=self.task_queue,
+            temporal_url=self.temporal_url,
+            namespace=self.temporal_namespace,
+        )
+
+        # Set environment variables for activities
+        os.environ["CONTROL_PLANE_URL"] = self.control_plane_url
+
+        # Connect to Temporal
+        try:
+            # Check if using Temporal Cloud (has API key)
+            is_cloud = "tmprl.cloud" in self.temporal_url or "api.temporal.io" in self.temporal_url
+
+            if is_cloud and self.temporal_api_key:
+                # Temporal Cloud with API key authentication
+                api_key_preview = f"{self.temporal_api_key[:20]}...{self.temporal_api_key[-20:]}" if len(self.temporal_api_key) > 40 else self.temporal_api_key
+                logger.info(
+                    "connecting_to_temporal_cloud",
+                    namespace=self.temporal_namespace,
+                    has_api_key=bool(self.temporal_api_key),
+                    api_key_preview=api_key_preview,
+                    host=self.temporal_url,
+                )
+
+                # Connect using rpc_metadata with Bearer token (same as control plane)
+                self.client = await TemporalClient.connect(
+                    self.temporal_url,
+                    namespace=self.temporal_namespace,
+                    tls=TLSConfig(),  # TLS without client cert
+                    rpc_metadata={"authorization": f"Bearer {self.temporal_api_key}"},
+                )
+            else:
+                # Local Temporal or non-cloud
+                logger.info("connecting_to_local_temporal", namespace=self.temporal_namespace)
+                self.client = await TemporalClient.connect(
+                    self.temporal_url,
+                    namespace=self.temporal_namespace,
+                )
+
+            logger.info(
+                "temporal_client_connected",
+                url=self.temporal_url,
+                namespace=self.temporal_namespace
+            )
+        except Exception as e:
+            logger.error("failed_to_connect_to_temporal", error=str(e))
+            raise
+
+        # Configure workflow sandbox with passthrough modules
+        # These modules use non-deterministic operations at import time but are safe
+        # because they're only used in activities, not workflow logic
+        sandbox_restrictions = SandboxRestrictions.default.with_passthrough_modules(
+            "structlog",
+            "structlog.dev",
+            "structlog.processors",
+            "structlog.tracebacks",
+            "rich",
+            "rich.traceback",
+            "httpx",
+            "worker_internal",
+            "worker_internal.planner",
+        )
+
+        # Create worker
+        self.worker = Worker(
+            self.client,
+            task_queue=self.task_queue,
+            workflows=[
+                PlanOrchestratorWorkflow,  # Plan execution
+                PlanGenerationWorkflow,     # Plan generation (async)
+            ],
+            activities=[
+                # Plan orchestrator activities
+                create_plan_execution,
+                update_plan_state,
+                execute_task_activity,
+                validate_task_completion,
+                get_task_status_activity,
+                call_llm_activity,
+                continue_task_activity,
+                analyze_task_completion_status,
+                publish_event_activity,
+                # Plan generation activities
+                generate_plan_activity,
+                store_plan_activity,
+                update_plan_generation_status,
+            ],
+            max_concurrent_workflow_tasks=10,
+            max_concurrent_activities=20,
+            workflow_runner=SandboxedWorkflowRunner(restrictions=sandbox_restrictions),
+            # Temporarily disabled build_id to avoid versioning mismatch issues
+            # build_id=f"plan-orchestrator-{self.organization_id}-v1",
+        )
+
+        logger.info(
+            "plan_worker_started",
+            task_queue=self.task_queue,
+            namespace=self.temporal_namespace,
+        )
+
+        # Run worker
+        await self.worker.run()
+
+    async def stop(self):
+        """Stop the worker gracefully."""
+        logger.info("stopping_plan_worker")
+        self._shutdown = True
+        # Worker.run() will exit when workflows/activities complete
+
+
+async def main():
+    """Main entry point for worker."""
+    # Load environment from .env.local
+    from dotenv import load_dotenv
+    env_path = os.path.join(os.path.dirname(__file__), '..', '..', '.env.local')
+    if os.path.exists(env_path):
+        load_dotenv(env_path)
+        logger.info("loaded_env_from_file", path=env_path)
+
+    parser = argparse.ArgumentParser(
+        description="Plan Orchestrator Worker - Handles all organizations",
+        epilog="""
+Examples:
+  # Basic usage (uses env vars)
+  python -m worker_internal.planner.worker
+
+  # Override defaults
+  python -m worker_internal.planner.worker --temporal-namespace my-namespace.lpagu
+        """
+    )
+    parser.add_argument(
+        "--temporal-url",
+        default=os.getenv("TEMPORAL_HOST") or os.getenv("TEMPORAL_URL", "us-east-1.aws.api.temporal.io:7233"),
+        help="Temporal server URL",
+    )
+    parser.add_argument(
+        "--temporal-namespace",
+        default=os.getenv("TEMPORAL_NAMESPACE", "agent-control-plane.lpagu"),
+        help="Temporal namespace",
+    )
+    parser.add_argument(
+        "--temporal-api-key",
+        default=os.getenv("TEMPORAL_API_KEY") or os.getenv("TEMPORAL_CLOUD_ADMIN_TOKEN"),
+        help="Temporal API key",
+    )
+    parser.add_argument(
+        "--task-queue",
+        default=os.getenv("TASK_QUEUE", "agent-control-plane.internal"),
+        help="Task queue name",
+    )
+    parser.add_argument(
+        "--control-plane-url",
+        default=os.getenv("CONTROL_PLANE_URL", "https://control-plane.kubiya.ai"),
+        help="Control Plane API URL",
+    )
+
+    args = parser.parse_args()
+
+    # Validate required parameters
+    if not args.temporal_api_key:
+        logger.error(
+            "TEMPORAL_API_KEY required",
+            message="Set TEMPORAL_API_KEY or TEMPORAL_CLOUD_ADMIN_TOKEN env var"
+        )
+        sys.exit(1)
+
+    logger.info(
+        "worker_configuration",
+        temporal_url=args.temporal_url,
+        temporal_namespace=args.temporal_namespace,
+        task_queue=args.task_queue,
+        control_plane_url=args.control_plane_url,
+    )
+
+    # Create worker with env-based credentials
+    worker = PlanWorker(
+        temporal_url=args.temporal_url,
+        temporal_namespace=args.temporal_namespace,
+        temporal_api_key=args.temporal_api_key,
+        control_plane_url=args.control_plane_url,
+        task_queue=args.task_queue,
+    )
+
+    # Handle shutdown signals
+    loop = asyncio.get_event_loop()
+
+    def handle_shutdown(sig):
+        logger.info("shutdown_signal_received", signal=sig)
+        asyncio.create_task(worker.stop())
+
+    for sig in (signal.SIGTERM, signal.SIGINT):
+        loop.add_signal_handler(sig, lambda s=sig: handle_shutdown(s))
+
+    # Start worker
+    try:
+        await worker.start()
+    except KeyboardInterrupt:
+        logger.info("keyboard_interrupt_received")
+    except Exception as e:
+        import traceback
+        logger.error(
+            "worker_error",
+            error=str(e),
+            traceback=traceback.format_exc()
+        )
+        sys.exit(1)
+    finally:
+        logger.info("worker_shutdown_complete")
+
+
+if __name__ == "__main__":
+    # Setup logging
+    structlog.configure(
+        processors=[
+            structlog.processors.TimeStamper(fmt="iso"),
+            structlog.processors.add_log_level,
+            structlog.processors.JSONRenderer(),
+        ]
+    )
+
+    # Run
+    asyncio.run(main())