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

control_plane_api/worker/services/agent_executor.py

@@ -0,0 +1,485 @@
+"""Agent executor service - handles agent execution business logic"""
+
+from typing import Dict, Any, Optional, List
+from datetime import datetime, timezone
+import structlog
+import asyncio
+import os
+
+from agno.agent import Agent
+from agno.models.litellm import LiteLLM
+
+from control_plane_api.worker.control_plane_client import ControlPlaneClient
+from control_plane_api.worker.services.session_service import SessionService
+from control_plane_api.worker.services.cancellation_manager import CancellationManager
+from control_plane_api.worker.services.skill_factory import SkillFactory
+from control_plane_api.worker.utils.streaming_utils import StreamingHelper
+
+logger = structlog.get_logger()
+
+
+class AgentExecutorService:
+    """
+    Service for executing agents with full session management and cancellation support.
+
+    This service orchestrates:
+    - Session loading and restoration
+    - Agent creation with LiteLLM configuration
+    - Skill instantiation
+    - Streaming execution with real-time updates
+    - Session persistence
+    - Cancellation support via CancellationManager
+    """
+
+    def __init__(
+        self,
+        control_plane: ControlPlaneClient,
+        session_service: SessionService,
+        cancellation_manager: CancellationManager
+    ):
+        self.control_plane = control_plane
+        self.session_service = session_service
+        self.cancellation_manager = cancellation_manager
+
+    async def execute(self, input: Any) -> Dict[str, Any]:
+        """
+        Execute an agent with full session management and streaming.
+
+        Args:
+            input: AgentExecutionInput with execution details
+
+        Returns:
+            Dict with response, usage, success flag, etc.
+        """
+        execution_id = input.execution_id
+
+        logger.info(
+            "agent_workflow_start",
+            execution_id=execution_id[:8],
+            agent_id=input.agent_id,
+            session_id=input.session_id
+        )
+
+        try:
+            # STEP 1: Load session history
+            session_history = self.session_service.load_session(
+                execution_id=execution_id,
+                session_id=input.session_id
+            )
+
+            if session_history:
+                print(f"✅ Loaded {len(session_history)} messages from previous session\n")
+            else:
+                print("ℹ️  Starting new conversation\n")
+
+            # STEP 2: Build conversation context for Agno
+            conversation_context = self.session_service.build_conversation_context(session_history)
+
+            # STEP 3: Get LiteLLM configuration
+            litellm_api_base = os.getenv("LITELLM_API_BASE", "https://llm-proxy.kubiya.ai")
+            litellm_api_key = os.getenv("LITELLM_API_KEY")
+
+            if not litellm_api_key:
+                raise ValueError("LITELLM_API_KEY environment variable not set")
+
+            model = input.model_id or os.environ.get("LITELLM_DEFAULT_MODEL", "kubiya/claude-sonnet-4")
+
+            # STEP 4: Fetch and instantiate skills
+            skills = []
+            if input.agent_id:
+                print(f"🔧 Fetching skills from Control Plane...")
+                try:
+                    skill_configs = self.control_plane.get_skills(input.agent_id)
+                    if skill_configs:
+                        print(f"✅ Resolved {len(skill_configs)} skills")
+                        print(f"   Types: {[t.get('type') for t in skill_configs]}")
+                        print(f"   Names: {[t.get('name') for t in skill_configs]}\n")
+
+                        skills = SkillFactory.create_skills_from_list(skill_configs)
+
+                        if skills:
+                            print(f"✅ Instantiated {len(skills)} skill(s)\n")
+                    else:
+                        print(f"⚠️  No skills found\n")
+                except Exception as e:
+                    print(f"❌ Error fetching skills: {str(e)}\n")
+                    logger.error("skill_fetch_error", error=str(e))
+
+            # STEP 5: Create agent with streaming helper
+            print(f"\n🤖 Creating Agno Agent:")
+            print(f"   Model: {model}")
+            print(f"   Skills: {len(skills)}")
+
+            # Create streaming helper for this execution
+            streaming_helper = StreamingHelper(
+                control_plane_client=self.control_plane,
+                execution_id=execution_id
+            )
+
+            # Create tool hook for real-time updates
+            def tool_hook(name: str = None, function_name: str = None, function=None, arguments: dict = None, **kwargs):
+                """Hook to capture tool execution for real-time streaming"""
+                tool_name = name or function_name or "unknown"
+                tool_args = arguments or {}
+
+                # Generate unique tool execution ID
+                import time
+                tool_execution_id = f"{tool_name}_{int(time.time() * 1000000)}"
+
+                print(f"   🔧 Tool Starting: {tool_name} (ID: {tool_execution_id})")
+
+                # Publish tool start event
+                streaming_helper.publish_tool_start(
+                    tool_name=tool_name,
+                    tool_execution_id=tool_execution_id,
+                    tool_args=tool_args,
+                    source="agent"
+                )
+
+                # Execute the tool
+                result = None
+                error = None
+                try:
+                    if function and callable(function):
+                        result = function(**tool_args) if tool_args else function()
+                    else:
+                        raise ValueError(f"Function not callable: {function}")
+
+                    status = "success"
+                    print(f"   ✅ Tool Success: {tool_name}")
+
+                except Exception as e:
+                    error = e
+                    status = "failed"
+                    print(f"   ❌ Tool Failed: {tool_name} - {str(e)}")
+
+                # Publish tool completion event
+                streaming_helper.publish_tool_complete(
+                    tool_name=tool_name,
+                    tool_execution_id=tool_execution_id,
+                    status=status,
+                    output=str(result)[:1000] if result else None,
+                    error=str(error) if error else None,
+                    source="agent"
+                )
+
+                if error:
+                    raise error
+
+                return result
+
+            # Create Agno Agent
+            agent = Agent(
+                name=f"Agent {input.agent_id}",
+                role=input.system_prompt or "You are a helpful AI assistant",
+                model=LiteLLM(
+                    id=f"openai/{model}",
+                    api_base=litellm_api_base,
+                    api_key=litellm_api_key,
+                ),
+                tools=skills if skills else None,
+                tool_hooks=[tool_hook],
+            )
+
+            # STEP 6: Register for cancellation
+            self.cancellation_manager.register(
+                execution_id=execution_id,
+                instance=agent,
+                instance_type="agent"
+            )
+            print(f"✅ Agent registered for cancellation support\n")
+
+            # Cache execution metadata in Redis
+            self.control_plane.cache_metadata(execution_id, "AGENT")
+
+            # STEP 7: Execute with streaming
+            print("⚡ Executing Agent Run with Streaming...\n")
+
+            # Generate deterministic message IDs for this turn (matches V2 executor pattern)
+            # This ensures streaming and persisted messages have the SAME message_id
+            turn_number = len(session_history) // 2 + 1
+            user_message_id = f"{execution_id}_user_{turn_number}"
+            message_id = f"{execution_id}_assistant_{turn_number}"
+
+            # Publish user message to stream immediately so it appears in chronological order
+            from datetime import datetime, timezone
+            user_message_timestamp = datetime.now(timezone.utc).isoformat()
+            self.control_plane.publish_event(
+                execution_id=execution_id,
+                event_type="message",
+                data={
+                    "role": "user",
+                    "content": input.prompt,
+                    "timestamp": user_message_timestamp,
+                    "message_id": user_message_id,
+                    "user_id": input.user_id,
+                    "user_name": getattr(input, "user_name", None),
+                    "user_email": getattr(input, "user_email", None),
+                    "user_avatar": getattr(input, "user_avatar", None),
+                }
+            )
+            print(f"   📤 Published user message to stream (ID: {user_message_id})")
+
+            def stream_agent_run():
+                """Run agent with streaming and collect response"""
+                # Create event loop for this thread (needed for async streaming)
+                loop = asyncio.new_event_loop()
+                asyncio.set_event_loop(loop)
+
+                async def _async_stream():
+                    """Async wrapper for streaming execution"""
+                    import time as time_module
+                    last_heartbeat_time = time_module.time()
+                    last_persistence_time = time_module.time()
+                    heartbeat_interval = 10  # Send heartbeat every 10 seconds
+                    persistence_interval = 60  # Persist to database every 60 seconds
+
+                    try:
+                        # Execute with conversation context
+                        if conversation_context:
+                            run_response = agent.run(
+                                input.prompt,
+                                stream=True,
+                                messages=conversation_context,
+                            )
+                        else:
+                            run_response = agent.run(input.prompt, stream=True)
+
+                        # Process streaming chunks (sync iteration in async context)
+                        for chunk in run_response:
+                            # Periodic maintenance: heartbeats and persistence
+                            current_time = time_module.time()
+
+                            # Send heartbeat every 10s (Temporal liveness)
+                            if current_time - last_heartbeat_time >= heartbeat_interval:
+                                current_response = streaming_helper.get_full_response()
+                                activity.heartbeat({
+                                    "status": "Streaming in progress...",
+                                    "response_length": len(current_response),
+                                    "execution_id": execution_id,
+                                })
+                                last_heartbeat_time = current_time
+
+                            # Persist snapshot every 60s (resilience against crashes)
+                            if current_time - last_persistence_time >= persistence_interval:
+                                current_response = streaming_helper.get_full_response()
+                                if current_response:
+                                    print(f"\n💾 Periodic persistence ({len(current_response)} chars)...")
+                                    snapshot_messages = session_history + [{
+                                        "role": "assistant",
+                                        "content": current_response,
+                                        "timestamp": datetime.now(timezone.utc).isoformat(),
+                                        "message_id": message_id,  # Include deterministic message_id
+                                    }]
+                                    try:
+                                        # Best effort - don't fail execution if persistence fails
+                                        self.session_service.persist_session(
+                                            execution_id=execution_id,
+                                            session_id=input.session_id or execution_id,
+                                            user_id=input.user_id,
+                                            messages=snapshot_messages,
+                                            metadata={
+                                                "agent_id": input.agent_id,
+                                                "organization_id": input.organization_id,
+                                                "snapshot": True,
+                                            }
+                                        )
+                                        print(f"   ✅ Session snapshot persisted with message_id: {message_id}")
+                                    except Exception as e:
+                                        print(f"   ⚠️  Session persistence error: {str(e)} (non-fatal)")
+                                last_persistence_time = current_time
+
+                            # Handle run_id capture
+                            streaming_helper.handle_run_id(
+                                chunk=chunk,
+                                on_run_id=lambda run_id: self.cancellation_manager.set_run_id(execution_id, run_id)
+                            )
+
+                            # Handle content chunk
+                            streaming_helper.handle_content_chunk(
+                                chunk=chunk,
+                                message_id=message_id,
+                                print_to_console=True
+                            )
+
+                        print()  # New line after streaming
+                        return run_response
+
+                    except Exception as e:
+                        print(f"\n❌ Streaming error: {str(e)}")
+                        # Fall back to non-streaming
+                        if conversation_context:
+                            return agent.run(input.prompt, stream=False, messages=conversation_context)
+                        else:
+                            return agent.run(input.prompt, stream=False)
+
+                # Run the async function in the event loop
+                try:
+                    return loop.run_until_complete(_async_stream())
+                finally:
+                    loop.close()
+
+            # Execute in thread pool (no timeout - user controls via STOP button)
+            # Wrap in try-except to handle Temporal cancellation
+            try:
+                result = await asyncio.to_thread(stream_agent_run)
+            except asyncio.CancelledError:
+                # Temporal cancelled the activity - cancel the running agent
+                print("\n🛑 Cancellation signal received - stopping agent execution...")
+                cancel_result = self.cancellation_manager.cancel(execution_id)
+                if cancel_result["success"]:
+                    print(f"✅ Agent execution cancelled successfully")
+                else:
+                    print(f"⚠️  Cancellation completed with warning: {cancel_result.get('error', 'Unknown')}")
+                # Re-raise to let Temporal know we're cancelled
+                raise
+
+            print("✅ Agent Execution Completed!")
+            full_response = streaming_helper.get_full_response()
+            print(f"   Response Length: {len(full_response)} chars\n")
+
+            logger.info(
+                "agent_execution_completed",
+                execution_id=execution_id[:8],
+                response_length=len(full_response)
+            )
+
+            # Use the streamed response content
+            response_content = full_response if full_response else (result.content if hasattr(result, "content") else str(result))
+
+            # STEP 8: Extract usage metrics
+            usage = {}
+            if hasattr(result, "metrics") and result.metrics:
+                metrics = result.metrics
+                usage = {
+                    "prompt_tokens": getattr(metrics, "input_tokens", 0),
+                    "completion_tokens": getattr(metrics, "output_tokens", 0),
+                    "total_tokens": getattr(metrics, "total_tokens", 0),
+                }
+                print(f"📊 Token Usage:")
+                print(f"   Input: {usage.get('prompt_tokens', 0)}")
+                print(f"   Output: {usage.get('completion_tokens', 0)}")
+                print(f"   Total: {usage.get('total_tokens', 0)}\n")
+
+            # STEP 9: Persist complete session history
+            print("\n💾 Persisting session history to Control Plane...")
+
+            # Build message_ids dict to pass to extract_messages_from_result
+            # This ensures persisted messages use the SAME IDs as streaming messages
+            message_ids_map = {
+                len(session_history): user_message_id,      # User message index
+                len(session_history) + 1: message_id         # Assistant message index
+            }
+
+            # Extract messages from Agno result to get accurate timestamps
+            # Agno tracks message timestamps as they're created (msg.created_at)
+            extracted_messages = self.session_service.extract_messages_from_result(
+                result=run_response,
+                user_id=input.user_id,
+                execution_id=execution_id,
+                message_ids=message_ids_map  # Pass deterministic IDs
+            )
+            print(f"   📊 Extracted {len(extracted_messages)} messages with deterministic IDs")
+
+            # Use extracted messages which have proper timestamps from Agno
+            # These include both user and assistant messages with accurate created_at times
+            new_messages = extracted_messages
+
+            # Fallback: if no messages extracted (shouldn't happen), create them manually
+            if not extracted_messages:
+                from datetime import datetime, timezone
+                current_timestamp = datetime.now(timezone.utc).isoformat()
+                print("   ⚠️  No messages extracted from Agno result, creating manually")
+                new_messages = [
+                    {
+                        "role": "user",
+                        "content": input.prompt,
+                        "timestamp": current_timestamp,
+                        "message_id": f"{execution_id}_user_{turn_number}",
+                        "user_id": input.user_id,
+                        "user_name": getattr(input, "user_name", None),
+                        "user_email": getattr(input, "user_email", None),
+                    },
+                    {
+                        "role": "assistant",
+                        "content": response_content,
+                        "timestamp": current_timestamp,
+                        "message_id": message_id,  # Use the same message_id as streaming
+                    },
+                ]
+
+            # Extract tool messages from streaming helper
+            tool_messages = streaming_helper.get_tool_messages()
+            print(f"   📊 Collected {len(tool_messages)} tool messages during streaming")
+
+            # Combine with previous history: session_history + new_messages + tool_messages
+            complete_session = session_history + new_messages + tool_messages
+
+            # CRITICAL: Deduplicate messages by message_id AND content to prevent duplicates
+            # Use session_service.deduplicate_messages() which has enhanced two-level deduplication
+            original_count = len(complete_session)
+            complete_session = self.session_service.deduplicate_messages(complete_session)
+
+            # CRITICAL: Sort by timestamp to ensure chronological order
+            # Tool messages happen DURING streaming, so they need to be interleaved with user/assistant messages
+            complete_session.sort(key=lambda msg: msg.get("timestamp", ""))
+            print(f"   ✅ Messages deduplicated ({original_count} -> {len(complete_session)}) and sorted by timestamp")
+
+            if complete_session:
+                success = self.session_service.persist_session(
+                    execution_id=execution_id,
+                    session_id=input.session_id or execution_id,
+                    user_id=input.user_id,
+                    messages=complete_session,
+                    metadata={
+                        "agent_id": input.agent_id,
+                        "organization_id": input.organization_id,
+                        "turn_count": len(complete_session),
+                    }
+                )
+
+                if success:
+                    print(f"   ✅ Session persisted ({len(complete_session)} total messages)")
+                else:
+                    print(f"   ⚠️  Session persistence failed")
+            else:
+                print("   ℹ️  No messages to persist")
+
+            print("\n" + "="*80)
+            print("🏁 AGENT EXECUTION END")
+            print("="*80 + "\n")
+
+            # STEP 10: Cleanup
+            self.cancellation_manager.unregister(execution_id)
+
+            return {
+                "success": True,
+                "response": response_content,
+                "usage": usage,
+                "model": model,
+                "finish_reason": "stop",
+            }
+
+        except Exception as e:
+            # Cleanup on error
+            self.cancellation_manager.unregister(execution_id)
+
+            print("\n" + "="*80)
+            print("❌ AGENT EXECUTION FAILED")
+            print("="*80)
+            print(f"Error: {str(e)}")
+            print("="*80 + "\n")
+
+            logger.error(
+                "agent_execution_failed",
+                execution_id=execution_id[:8],
+                error=str(e)
+            )
+
+            return {
+                "success": False,
+                "error": str(e),
+                "model": input.model_id,
+                "usage": None,
+                "finish_reason": "error",
+            }