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

control_plane_api/worker/runtimes/agno/config.py

@@ -0,0 +1,248 @@
+"""
+Configuration builder for Agno runtime.
+
+This module handles:
+- LiteLLM API configuration
+- Model selection and setup
+- Agent configuration
+- Environment variable management
+"""
+
+import os
+import structlog
+from typing import Any, Optional, List
+
+from agno.agent import Agent
+from agno.models.litellm import LiteLLM
+from control_plane_api.worker.services.system_prompt_enhancement import create_default_prompt_builder
+from control_plane_api.worker.runtimes.model_utils import get_effective_model, is_model_override_active
+
+# Import LiteLLMLangfuse for proper Langfuse metadata support
+try:
+    from agno.models.litellm import LiteLLMLangfuse
+    LANGFUSE_SUPPORT = True
+except ImportError:
+    LANGFUSE_SUPPORT = False
+
+logger = structlog.get_logger(__name__)
+
+# Module-level singleton for system prompt enhancement
+# NOTE: This is now created per-execution to support dynamic skill context
+# _prompt_builder = create_default_prompt_builder()
+
+
+def build_agno_agent_config(
+    agent_id: str,
+    system_prompt: Optional[str] = None,
+    model_id: Optional[str] = None,
+    skills: Optional[List[Any]] = None,
+    mcp_tools: Optional[List[Any]] = None,
+    tool_hooks: Optional[List[Any]] = None,
+    user_id: Optional[str] = None,
+    session_id: Optional[str] = None,
+    organization_id: Optional[str] = None,
+    agent_name: Optional[str] = None,
+    skill_configs: Optional[List[Any]] = None,
+    user_metadata: Optional[dict] = None,
+    additional_context: Optional[dict] = None,
+) -> Agent:
+    """
+    Build Agno Agent configuration with LiteLLM.
+
+    Args:
+        agent_id: Unique identifier for the agent
+        system_prompt: System-level instructions
+        model_id: Model identifier (overrides default)
+        skills: List of skills/tools available to agent
+        mcp_tools: List of MCPTools instances
+        tool_hooks: List of tool execution hooks
+        user_id: User identifier (email) for Langfuse tracking
+        session_id: Session identifier for Langfuse tracking
+        organization_id: Organization identifier for Langfuse tracking
+        agent_name: Agent name for generation naming
+        skill_configs: Original skill configuration dictionaries for prompt enhancement
+        user_metadata: User metadata dict containing user_id, user_name, user_email, user_avatar
+        additional_context: Additional contextual information to inject (custom key-value pairs)
+
+    Returns:
+        Configured Agno Agent instance
+
+    Raises:
+        ValueError: If required environment variables are missing
+    """
+    # Get LiteLLM configuration from environment
+    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")
+
+    # Determine model to use with override support
+    # Priority: KUBIYA_MODEL_OVERRIDE > model_id > LITELLM_DEFAULT_MODEL > default
+    model = get_effective_model(
+        context_model_id=model_id,
+        log_context={"agent_id": agent_id},
+    )
+
+    logger.info(
+        "building_agno_agent_config",
+        agent_id=agent_id,
+        model=model,
+        has_skills=bool(skills),
+        has_mcp_tools=bool(mcp_tools),
+        mcp_tools_count=len(mcp_tools) if mcp_tools else 0,
+        has_tool_hooks=bool(tool_hooks),
+    )
+
+    # Build metadata for Langfuse tracking
+    # Format: trace_user_id = EMAIL-ORG_NAME, trace_name = "agent-chat" (simple), metadata = {details}
+    metadata = {}
+
+    # Always set a simple, consistent trace name
+    metadata["trace_name"] = "agent-chat"
+    metadata["generation_name"] = "agent-chat"
+
+    if user_id and organization_id:
+        # Format: EMAIL-ORG_NAME (EMAIL is the user_id)
+        metadata["trace_user_id"] = f"{user_id}-{organization_id}"
+        metadata["user_id"] = f"{user_id}-{organization_id}"
+
+    if session_id:
+        # Use session_id as trace_id to group all messages in same conversation
+        metadata["trace_id"] = session_id
+        metadata["session_id"] = session_id
+
+    # Add additional details directly to metadata for Langfuse "Metadata" column
+    # Any extra fields not in the standard spec will be saved as metadata
+    if agent_id:
+        metadata["agent_id"] = agent_id
+    if agent_name:
+        metadata["agent_name"] = agent_name
+    if user_id:
+        metadata["user_email"] = user_id
+    if organization_id:
+        metadata["organization_id"] = organization_id
+    if model_id:
+        metadata["model"] = model_id
+
+    # DEBUG: Log the complete metadata being sent
+    logger.warning(
+        "🔍 DEBUG: AGNO RUNTIME - LANGFUSE METADATA",
+        metadata=metadata,
+        user_id_input=user_id,
+        organization_id_input=organization_id,
+        agent_name_input=agent_name,
+        session_id_input=session_id,
+        metadata_json=str(metadata),
+    )
+
+    # Create LiteLLM model instance with metadata
+    # Use LiteLLMLangfuse for proper Langfuse integration if available
+    if LANGFUSE_SUPPORT and metadata:
+        logger.warning(
+            "🔍 DEBUG: USING LiteLLMLangfuse CLASS FOR LANGFUSE INTEGRATION",
+            metadata=metadata,
+        )
+        litellm_model = LiteLLMLangfuse(
+            id=f"openai/{model}",
+            api_base=litellm_api_base,
+            api_key=litellm_api_key,
+            metadata=metadata,
+        )
+    else:
+        logger.warning(
+            "🔍 DEBUG: USING STANDARD LiteLLM CLASS (No Langfuse integration)",
+            has_metadata=bool(metadata),
+            langfuse_available=LANGFUSE_SUPPORT,
+            note="Install agno with Langfuse support for proper metadata tracking"
+        )
+        litellm_model = LiteLLM(
+            id=f"openai/{model}",
+            api_base=litellm_api_base,
+            api_key=litellm_api_key,
+            metadata=metadata if metadata else None,
+        )
+
+    logger.warning(
+        "🔍 DEBUG: LITELLM MODEL CREATED",
+        model_class=litellm_model.__class__.__name__,
+        model_id=f"openai/{model}",
+        api_base=litellm_api_base,
+        has_metadata=bool(metadata),
+        metadata_keys=list(metadata.keys()) if metadata else [],
+    )
+
+    # Combine skills and MCP tools
+    all_tools = []
+    if skills:
+        all_tools.extend(skills)
+    if mcp_tools:
+        all_tools.extend(mcp_tools)
+
+    # Enhance system prompt with runtime-specific additions
+    # Create per-execution prompt builder to support dynamic skill context and user context
+    from control_plane_api.worker.services.skill_context_enhancement import (
+        SkillContextEnhancement,
+    )
+
+    # Create prompt builder with user context and additional context
+    prompt_builder = create_default_prompt_builder(
+        user_metadata=user_metadata,
+        additional_context=additional_context,
+    )
+
+    # Add skill context enhancement if enabled and skills are configured
+    skill_context_enabled = os.getenv("ENABLE_SKILL_CONTEXT_ENHANCEMENT", "true").lower() == "true"
+    if skill_context_enabled and skill_configs:
+        skill_context_enhancement = SkillContextEnhancement(skill_configs)
+        prompt_builder.add_enhancement(skill_context_enhancement)
+        logger.info(
+            "skill_context_enhancement_enabled",
+            skill_count=len(skill_configs),
+            agent_id=agent_id,
+        )
+
+    enhanced_system_prompt = prompt_builder.build(
+        base_prompt=system_prompt,
+        runtime_type="agno",
+    )
+
+    # Build agent configuration
+    agent = Agent(
+        name=f"Agent {agent_id}",
+        role=enhanced_system_prompt or "You are a helpful AI assistant",
+        model=litellm_model,
+        tools=all_tools if all_tools else None,
+        tool_hooks=tool_hooks if tool_hooks else None,
+    )
+
+    return agent
+
+
+def validate_litellm_config() -> bool:
+    """
+    Validate LiteLLM configuration is present.
+
+    Returns:
+        True if configuration is valid
+
+    Raises:
+        ValueError: If configuration is invalid
+    """
+    litellm_api_key = os.getenv("LITELLM_API_KEY")
+
+    if not litellm_api_key:
+        raise ValueError(
+            "LITELLM_API_KEY environment variable not set. "
+            "This is required for Agno runtime to function."
+        )
+
+    logger.debug(
+        "litellm_config_validated",
+        api_base=os.getenv("LITELLM_API_BASE", "https://llm-proxy.kubiya.ai"),
+        default_model=os.environ.get("LITELLM_DEFAULT_MODEL", "kubiya/claude-sonnet-4"),
+    )
+
+    return True

control_plane_api/worker/runtimes/agno/hooks.py

@@ -0,0 +1,385 @@
+"""
+Tool execution hooks for Agno runtime.
+
+This module provides:
+- Tool execution event hooks
+- Real-time event publishing to Control Plane
+- Event callback creation
+- Tool execution tracking
+- Policy enforcement for tool executions
+"""
+
+import uuid
+import asyncio
+import json
+import structlog
+from typing import Callable, Any, Optional, Dict, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from control_plane_client import ControlPlaneClient
+    from control_plane_api.worker.services.tool_enforcement import ToolEnforcementService
+
+logger = structlog.get_logger(__name__)
+
+
+def _parse_json_arguments(arguments: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Parse JSON strings in tool arguments to their proper types.
+
+    Agno sometimes passes dict/list arguments as JSON strings instead of actual objects.
+    This function detects and parses those strings into proper Python types.
+
+    Args:
+        arguments: Raw tool arguments from agno
+
+    Returns:
+        Parsed arguments with JSON strings converted to dicts/lists
+    """
+    if not arguments:
+        return arguments
+
+    parsed = {}
+    for key, value in arguments.items():
+        # If it's a string that looks like JSON, try to parse it
+        if isinstance(value, str) and value.strip().startswith(('{', '[')):
+            try:
+                parsed[key] = json.loads(value)
+            except (json.JSONDecodeError, ValueError):
+                # Not valid JSON, keep as string
+                parsed[key] = value
+        else:
+            parsed[key] = value
+
+    return parsed
+
+
+def create_tool_hook_for_streaming(
+    control_plane: "ControlPlaneClient",
+    execution_id: str,
+    message_id: str,
+    enforcement_context: Optional[Dict[str, Any]] = None,
+    enforcement_service: Optional["ToolEnforcementService"] = None,
+) -> Callable:
+    """
+    Create a tool hook for streaming execution that publishes directly to Control Plane.
+
+    This hook publishes tool events immediately (not batched) for real-time visibility.
+    Used in streaming execution mode. Includes policy enforcement checks.
+
+    IMPORTANT: Parameter names MUST match Agno's expected hook signature exactly:
+    - function_name: str - The name of the tool being called
+    - function_call: Callable - The actual function to call
+    - arguments: Dict[str, Any] - Arguments to pass to the function
+
+    Args:
+        control_plane: Control Plane client for publishing events
+        execution_id: Execution ID for this run
+        message_id: Message ID for the current assistant turn (links tool to conversation)
+        enforcement_context: Optional context for policy enforcement
+        enforcement_service: Optional enforcement service for policy checks
+
+    Returns:
+        Tool hook function for Agno agent
+    """
+    def tool_hook(
+        function_name: str,
+        function_call: Callable,
+        arguments: Dict[str, Any],
+        **kwargs,
+    ):
+        """
+        Hook to capture tool execution for real-time streaming.
+
+        Agno calls this hook with exact parameter names: function_name, function_call, arguments.
+        The hook MUST call function_call(**arguments) and return the result.
+        """
+        tool_name = function_name or "unknown"
+        # Parse JSON strings in arguments to proper types (dict/list)
+        tool_args = _parse_json_arguments(arguments or {})
+        # Use UUID for tool_execution_id to avoid collisions when tools run simultaneously
+        tool_execution_id = f"{tool_name}_{uuid.uuid4().hex[:12]}"
+
+        logger.info(
+            "tool_hook_invoked",
+            tool_name=tool_name,
+            tool_execution_id=tool_execution_id,
+            execution_id=execution_id[:8],
+            has_function_call=function_call is not None and callable(function_call),
+        )
+
+        # Enforcement check (non-blocking)
+        enforcement_allowed = True
+        enforcement_violation = None
+        enforcement_metadata = {}
+
+        if enforcement_service and enforcement_context:
+            try:
+                # Run enforcement check asynchronously using control_plane's thread-local loop
+                # This reuses the same event loop for all async operations in this thread,
+                # preventing resource leaks and "await wasn't used with future" errors
+                loop = control_plane._get_thread_event_loop()
+
+                enforcement_allowed, enforcement_violation, enforcement_metadata = (
+                    loop.run_until_complete(
+                        enforcement_service.enforce_tool_execution(
+                            tool_name=tool_name,
+                            tool_args=tool_args,
+                            enforcement_context={
+                                **enforcement_context,
+                                "execution_id": execution_id,
+                                "message_id": message_id,
+                                "tool_execution_id": tool_execution_id,
+                            },
+                        )
+                    )
+                )
+            except Exception as e:
+                logger.error(
+                    "enforcement_check_failed",
+                    tool_name=tool_name,
+                    error=str(e),
+                )
+                # Fail open - allow execution
+                enforcement_metadata = {"enforcer": "error", "error": str(e)}
+
+        # Publish tool start event (blocking call - OK in thread)
+        control_plane.publish_event(
+            execution_id=execution_id,
+            event_type="tool_started",
+            data={
+                "tool_name": tool_name,
+                "tool_execution_id": tool_execution_id,
+                "tool_arguments": tool_args,
+                "message_id": message_id,  # Link tool to assistant message
+                "message": f"Executing tool: {tool_name}",
+                "enforcement": enforcement_metadata,  # Add enforcement metadata
+            }
+        )
+
+        # Execute tool using Agno's function_call pattern
+        # IMPORTANT: Must call function_call(**arguments) as per Agno hook contract
+        result = None
+        error = None
+        try:
+            if function_call and callable(function_call):
+                result = function_call(**tool_args) if tool_args else function_call()
+            else:
+                raise ValueError(f"function_call not callable: {function_call}")
+            status = "success"
+
+            # Inject enforcement violation into result if blocked
+            if not enforcement_allowed and enforcement_violation:
+                # Prepend violation message to result
+                violation_prefix = f"\n{'='*60}\n"
+                violation_prefix += "⛔ POLICY VIOLATION DETECTED\n"
+                violation_prefix += f"{'='*60}\n"
+                violation_prefix += enforcement_violation
+                violation_prefix += f"\n{'='*60}\n\n"
+
+                if result:
+                    result = violation_prefix + str(result)
+                else:
+                    result = violation_prefix + "(Tool execution completed but blocked by policy)"
+
+                logger.warning(
+                    "tool_execution_policy_violation",
+                    tool_name=tool_name,
+                    tool_execution_id=tool_execution_id,
+                    enforcement_metadata=enforcement_metadata,
+                )
+
+        except Exception as e:
+            error = e
+            status = "failed"
+            logger.error(
+                "tool_execution_failed",
+                tool_name=tool_name,
+                error=str(e),
+            )
+
+        # Publish tool completion event with enforcement metadata
+        control_plane.publish_event(
+            execution_id=execution_id,
+            event_type="tool_completed",
+            data={
+                "tool_name": tool_name,
+                "tool_execution_id": tool_execution_id,
+                "status": status,
+                "tool_output": str(result)[:1000] if result else None,
+                "tool_error": str(error) if error else None,
+                "message_id": message_id,  # Link tool to assistant message
+                "message": f"Tool {status}: {tool_name}",
+                "enforcement": enforcement_metadata,  # Add enforcement metadata
+            }
+        )
+
+        if error:
+            raise error
+        return result
+
+    return tool_hook
+
+
+def create_tool_hook_with_callback(
+    execution_id: str,
+    message_id: str,
+    event_callback: Callable[[Dict[str, Any]], None],
+    enforcement_context: Optional[Dict[str, Any]] = None,
+    enforcement_service: Optional["ToolEnforcementService"] = None,
+) -> Callable:
+    """
+    Create a tool hook that uses a callback for event publishing.
+
+    This hook uses a callback function to publish events, allowing for flexible
+    event handling (batching, filtering, etc.). Used in non-streaming execution.
+    Includes policy enforcement checks.
+
+    IMPORTANT: Parameter names MUST match Agno's expected hook signature exactly:
+    - function_name: str - The name of the tool being called
+    - function_call: Callable - The actual function to call
+    - arguments: Dict[str, Any] - Arguments to pass to the function
+
+    Args:
+        execution_id: Execution ID for this run
+        message_id: Message ID for the current assistant turn
+        event_callback: Callback function for event publishing
+        enforcement_context: Optional context for policy enforcement
+        enforcement_service: Optional enforcement service for policy checks
+
+    Returns:
+        Tool hook function for Agno agent
+    """
+    def tool_hook(
+        function_name: str,
+        function_call: Callable,
+        arguments: Dict[str, Any],
+        **kwargs,
+    ):
+        """
+        Hook to capture tool execution for callback-based publishing.
+
+        Agno calls this hook with exact parameter names: function_name, function_call, arguments.
+        The hook MUST call function_call(**arguments) and return the result.
+        """
+        tool_name = function_name or "unknown"
+        # Parse JSON strings in arguments to proper types (dict/list)
+        tool_args = _parse_json_arguments(arguments or {})
+        # Use UUID for tool_execution_id to avoid collisions
+        tool_execution_id = f"{tool_name}_{uuid.uuid4().hex[:12]}"
+
+        # Enforcement check (non-blocking)
+        enforcement_allowed = True
+        enforcement_violation = None
+        enforcement_metadata = {}
+
+        if enforcement_service and enforcement_context:
+            try:
+                # Run enforcement check asynchronously
+                # For callback-based hooks (used in non-streaming/testing), create a temporary loop
+                try:
+                    loop = asyncio.get_event_loop()
+                    if loop.is_closed():
+                        raise RuntimeError("Event loop is closed")
+                except RuntimeError:
+                    loop = asyncio.new_event_loop()
+                    asyncio.set_event_loop(loop)
+
+                enforcement_allowed, enforcement_violation, enforcement_metadata = (
+                    loop.run_until_complete(
+                        enforcement_service.enforce_tool_execution(
+                            tool_name=tool_name,
+                            tool_args=tool_args,
+                            enforcement_context={
+                                **enforcement_context,
+                                "execution_id": execution_id,
+                                "message_id": message_id,
+                                "tool_execution_id": tool_execution_id,
+                            },
+                        )
+                    )
+                )
+            except Exception as e:
+                logger.error(
+                    "enforcement_check_failed",
+                    tool_name=tool_name,
+                    error=str(e),
+                )
+                # Fail open - allow execution
+                enforcement_metadata = {"enforcer": "error", "error": str(e)}
+
+        # Publish tool start event via callback
+        event_callback(
+            {
+                "type": "tool_start",
+                "tool_name": tool_name,
+                "tool_execution_id": tool_execution_id,
+                "tool_args": tool_args,
+                "message_id": message_id,
+                "execution_id": execution_id,
+                "enforcement": enforcement_metadata,  # Add enforcement metadata
+            }
+        )
+
+        # Execute tool using Agno's function_call pattern
+        # IMPORTANT: Must call function_call(**arguments) as per Agno hook contract
+        result = None
+        error = None
+        try:
+            if function_call and callable(function_call):
+                result = function_call(**tool_args) if tool_args else function_call()
+            else:
+                raise ValueError(f"function_call not callable: {function_call}")
+
+            status = "success"
+
+            # Inject enforcement violation into result if blocked
+            if not enforcement_allowed and enforcement_violation:
+                # Prepend violation message to result
+                violation_prefix = f"\n{'='*60}\n"
+                violation_prefix += "⛔ POLICY VIOLATION DETECTED\n"
+                violation_prefix += f"{'='*60}\n"
+                violation_prefix += enforcement_violation
+                violation_prefix += f"\n{'='*60}\n\n"
+
+                if result:
+                    result = violation_prefix + str(result)
+                else:
+                    result = violation_prefix + "(Tool execution completed but blocked by policy)"
+
+                logger.warning(
+                    "tool_execution_policy_violation",
+                    tool_name=tool_name,
+                    tool_execution_id=tool_execution_id,
+                    enforcement_metadata=enforcement_metadata,
+                )
+
+        except Exception as e:
+            error = e
+            status = "failed"
+            logger.error(
+                "tool_execution_failed",
+                tool_name=tool_name,
+                error=str(e),
+            )
+
+        # Publish tool completion event via callback with enforcement metadata
+        event_callback(
+            {
+                "type": "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,
+                "message_id": message_id,
+                "execution_id": execution_id,
+                "enforcement": enforcement_metadata,  # Add enforcement metadata
+            }
+        )
+
+        if error:
+            raise error
+
+        return result
+
+    return tool_hook