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

control_plane_api/app/lib/litellm_pricing.py

@@ -0,0 +1,166 @@
+"""
+LiteLLM Pricing Integration - Fetch and cache model pricing data
+"""
+
+import httpx
+import structlog
+from typing import Dict, Optional
+from datetime import datetime, timedelta
+import asyncio
+
+logger = structlog.get_logger()
+
+# Cache for pricing data
+_pricing_cache: Optional[Dict] = None
+_cache_timestamp: Optional[datetime] = None
+_cache_lock = asyncio.Lock()
+
+PRICING_URL = "https://raw.githubusercontent.com/BerriAI/litellm/refs/heads/main/model_prices_and_context_window.json"
+CACHE_TTL_HOURS = 24  # Refresh pricing data daily
+
+
+async def get_litellm_pricing() -> Dict:
+    """
+    Fetch LiteLLM pricing data with caching
+
+    Returns:
+        Dict containing model pricing information
+    """
+    global _pricing_cache, _cache_timestamp
+
+    async with _cache_lock:
+        # Check if cache is valid
+        if _pricing_cache and _cache_timestamp:
+            age = datetime.utcnow() - _cache_timestamp
+            if age < timedelta(hours=CACHE_TTL_HOURS):
+                logger.debug("litellm_pricing_cache_hit", age_hours=age.total_seconds() / 3600)
+                return _pricing_cache
+
+        # Fetch fresh pricing data
+        logger.info("fetching_litellm_pricing", url=PRICING_URL)
+        try:
+            async with httpx.AsyncClient(timeout=30.0) as client:
+                response = await client.get(PRICING_URL)
+                response.raise_for_status()
+                pricing_data = response.json()
+
+                # Update cache
+                _pricing_cache = pricing_data
+                _cache_timestamp = datetime.utcnow()
+
+                logger.info(
+                    "litellm_pricing_fetched_successfully",
+                    models_count=len(pricing_data),
+                    cached_until=(_cache_timestamp + timedelta(hours=CACHE_TTL_HOURS)).isoformat()
+                )
+
+                return pricing_data
+        except Exception as e:
+            logger.error("litellm_pricing_fetch_failed", error=str(e), error_type=type(e).__name__)
+            # Return empty dict on failure
+            return {}
+
+
+def get_model_pricing(model_id: str, pricing_data: Dict) -> Optional[Dict]:
+    """
+    Get pricing information for a specific model
+
+    Args:
+        model_id: Model identifier (e.g., "claude-sonnet-4", "gpt-4o")
+        pricing_data: Full pricing data from LiteLLM
+
+    Returns:
+        Dict with pricing info or None if not found
+    """
+    # Try exact match first
+    if model_id in pricing_data:
+        return pricing_data[model_id]
+
+    # Try common variations
+    variations = [
+        model_id,
+        f"openai/{model_id}",
+        f"anthropic/{model_id}",
+        f"anthropic.{model_id}",
+        f"bedrock/{model_id}",
+    ]
+
+    for variation in variations:
+        if variation in pricing_data:
+            return pricing_data[variation]
+
+    logger.warning("model_pricing_not_found", model_id=model_id, tried_variations=variations)
+    return None
+
+
+def calculate_llm_cost(
+    model_id: str,
+    estimated_input_tokens: int,
+    estimated_output_tokens: int,
+    pricing_data: Dict
+) -> tuple[float, float, float]:
+    """
+    Calculate LLM cost for a model
+
+    Args:
+        model_id: Model identifier
+        estimated_input_tokens: Expected input tokens
+        estimated_output_tokens: Expected output tokens
+        pricing_data: Full pricing data from LiteLLM
+
+    Returns:
+        Tuple of (cost_per_1k_input, cost_per_1k_output, total_cost)
+    """
+    model_pricing = get_model_pricing(model_id, pricing_data)
+
+    if not model_pricing:
+        # Fallback defaults
+        logger.warning("using_default_pricing", model_id=model_id)
+        return (0.003, 0.015, 0.0)
+
+    input_cost_per_token = model_pricing.get("input_cost_per_token", 0.000003)
+    output_cost_per_token = model_pricing.get("output_cost_per_token", 0.000015)
+
+    # Convert to per-1k pricing for display
+    input_cost_per_1k = input_cost_per_token * 1000
+    output_cost_per_1k = output_cost_per_token * 1000
+
+    # Calculate total cost
+    total_cost = (
+        (estimated_input_tokens * input_cost_per_token) +
+        (estimated_output_tokens * output_cost_per_token)
+    )
+
+    return (input_cost_per_1k, output_cost_per_1k, total_cost)
+
+
+def get_model_display_name(model_id: str) -> str:
+    """
+    Get human-readable model name
+
+    Args:
+        model_id: Model identifier
+
+    Returns:
+        Display name
+    """
+    # Map of common model IDs to display names
+    display_names = {
+        "claude-sonnet-4": "Claude Sonnet 4",
+        "claude-3-5-sonnet-20241022": "Claude 3.5 Sonnet",
+        "claude-3-5-sonnet-20240620": "Claude 3.5 Sonnet (Legacy)",
+        "claude-3-opus-20240229": "Claude 3 Opus",
+        "claude-3-haiku-20240307": "Claude 3 Haiku",
+        "gpt-4o": "GPT-4o",
+        "gpt-4o-mini": "GPT-4o Mini",
+        "gpt-4-turbo": "GPT-4 Turbo",
+        "gpt-4": "GPT-4",
+        "gpt-3.5-turbo": "GPT-3.5 Turbo",
+        "o1": "OpenAI o1",
+        "o1-mini": "OpenAI o1-mini",
+        "gemini-2.0-flash-exp": "Gemini 2.0 Flash",
+        "gemini-1.5-pro": "Gemini 1.5 Pro",
+        "gemini-1.5-flash": "Gemini 1.5 Flash",
+    }
+
+    return display_names.get(model_id, model_id.replace("-", " ").title())

control_plane_api/app/lib/mcp_validation.py

@@ -0,0 +1,163 @@
+"""
+MCP configuration validation helpers.
+
+Validates MCP server configurations and template syntax.
+"""
+
+import structlog
+from typing import Dict, List, Any, Optional
+from pydantic import ValidationError
+
+from control_plane_api.app.schemas.mcp_schemas import MCPServerConfig
+from control_plane_api.app.lib.templating import extract_all_variables, TemplateValidator, TemplateContext, get_default_engine
+
+logger = structlog.get_logger()
+
+
+class MCPValidationError(Exception):
+    """Raised when MCP configuration validation fails."""
+    pass
+
+
+def validate_mcp_server_config(
+    mcp_servers: Dict[str, Any],
+    available_secrets: Optional[List[str]] = None,
+    available_env_vars: Optional[List[str]] = None,
+    strict: bool = False
+) -> Dict[str, Any]:
+    """
+    Validate MCP server configuration.
+
+    Validates:
+    1. MCP server configuration schema (Pydantic validation)
+    2. Template syntax in all string fields
+    3. Referenced secrets exist (if available_secrets provided)
+    4. Referenced env vars exist (if available_env_vars provided)
+
+    Args:
+        mcp_servers: Dict of MCP server configurations
+        available_secrets: List of available secret names for validation
+        available_env_vars: List of available env var names for validation
+        strict: If True, raise error for missing secrets/env vars
+
+    Returns:
+        Dict with validation results:
+        {
+            "valid": bool,
+            "errors": List[str],
+            "warnings": List[str],
+            "required_secrets": List[str],
+            "required_env_vars": List[str]
+        }
+
+    Raises:
+        MCPValidationError: If strict=True and validation fails
+    """
+    errors = []
+    warnings = []
+    all_required_secrets = set()
+    all_required_env_vars = set()
+
+    if not mcp_servers:
+        return {
+            "valid": True,
+            "errors": [],
+            "warnings": [],
+            "required_secrets": [],
+            "required_env_vars": []
+        }
+
+    # Validate each server configuration
+    for server_name, server_config in mcp_servers.items():
+        try:
+            # Validate using Pydantic schema
+            validated_config = MCPServerConfig(**server_config)
+
+            # Extract template variables
+            variables = extract_all_variables(server_config)
+            required_secrets = variables.get("secrets", [])
+            required_env_vars = variables.get("env_vars", [])
+
+            all_required_secrets.update(required_secrets)
+            all_required_env_vars.update(required_env_vars)
+
+            # Check if required secrets are available
+            if available_secrets is not None:
+                missing_secrets = [s for s in required_secrets if s not in available_secrets]
+                if missing_secrets:
+                    error_msg = f"Server '{server_name}': Missing secrets: {', '.join(missing_secrets)}"
+                    if strict:
+                        errors.append(error_msg)
+                    else:
+                        warnings.append(error_msg)
+
+            # Check if required env vars are available
+            if available_env_vars is not None:
+                missing_env_vars = [v for v in required_env_vars if v not in available_env_vars]
+                if missing_env_vars:
+                    error_msg = f"Server '{server_name}': Missing environment variables: {', '.join(missing_env_vars)}"
+                    if strict:
+                        errors.append(error_msg)
+                    else:
+                        warnings.append(error_msg)
+
+            logger.debug(
+                "mcp_server_validated",
+                server_name=server_name,
+                transport_type=validated_config.get_transport_type().value,
+                required_secrets_count=len(required_secrets),
+                required_env_vars_count=len(required_env_vars)
+            )
+
+        except ValidationError as e:
+            error_msg = f"Server '{server_name}': Invalid configuration: {str(e)}"
+            errors.append(error_msg)
+            logger.warning("mcp_server_validation_error", server_name=server_name, error=str(e))
+
+        except Exception as e:
+            error_msg = f"Server '{server_name}': Validation error: {str(e)}"
+            errors.append(error_msg)
+            logger.error("mcp_server_validation_exception", server_name=server_name, error=str(e))
+
+    result = {
+        "valid": len(errors) == 0,
+        "errors": errors,
+        "warnings": warnings,
+        "required_secrets": sorted(list(all_required_secrets)),
+        "required_env_vars": sorted(list(all_required_env_vars))
+    }
+
+    if strict and not result["valid"]:
+        error_summary = "; ".join(errors)
+        raise MCPValidationError(f"MCP configuration validation failed: {error_summary}")
+
+    return result
+
+
+def validate_execution_environment_mcp(
+    execution_environment: Dict[str, Any],
+    strict: bool = False
+) -> Dict[str, Any]:
+    """
+    Validate MCP servers in an execution environment configuration.
+
+    Args:
+        execution_environment: Execution environment dict with mcp_servers
+        strict: If True, raise error on validation failure
+
+    Returns:
+        Validation result dict
+
+    Raises:
+        MCPValidationError: If strict=True and validation fails
+    """
+    mcp_servers = execution_environment.get("mcp_servers", {})
+    available_secrets = execution_environment.get("secrets", [])
+    available_env_vars = list(execution_environment.get("env_vars", {}).keys())
+
+    return validate_mcp_server_config(
+        mcp_servers,
+        available_secrets=available_secrets,
+        available_env_vars=available_env_vars,
+        strict=strict
+    )

control_plane_api/app/lib/nats/__init__.py

@@ -0,0 +1,13 @@
+"""NATS integration components for control plane."""
+
+from control_plane_api.app.lib.nats.credentials_manager import (
+    NATSCredentialsManager,
+    WorkerCredentials,
+)
+from control_plane_api.app.lib.nats.listener import NATSEventListener
+
+__all__ = [
+    "NATSCredentialsManager",
+    "WorkerCredentials",
+    "NATSEventListener",
+]

control_plane_api/app/lib/nats/credentials_manager.py

@@ -0,0 +1,288 @@
+"""NATS credentials manager for generating temporary worker credentials."""
+
+import os
+import json
+from datetime import datetime, timedelta, timezone
+from typing import Optional
+from pydantic import BaseModel, Field
+import structlog
+
+logger = structlog.get_logger(__name__)
+
+
+class WorkerCredentials(BaseModel):
+    """NATS credentials for a worker."""
+
+    jwt: str = Field(..., description="User JWT token")
+    seed: str = Field(..., description="NKey seed (private key)")
+    public_key: str = Field(..., description="NKey public key")
+    subject_prefix: str = Field(..., description="Subject prefix for publishing")
+    expires_at: datetime = Field(..., description="Credential expiration time")
+
+
+class NATSCredentialsManager:
+    """Generate temporary NATS credentials for workers."""
+
+    def __init__(
+        self,
+        operator_jwt: str,
+        operator_seed: str,
+        account_public_key: Optional[str] = None,
+    ):
+        """
+        Initialize NATS credentials manager.
+
+        Args:
+            operator_jwt: Operator JWT (not used for signing, just for reference)
+            operator_seed: Operator seed for signing JWTs
+            account_public_key: Optional account public key (extracted from operator JWT if not provided)
+        """
+        try:
+            import nkeys
+            import jwt as jwt_lib
+        except ImportError:
+            raise ImportError(
+                "nkeys and PyJWT are required for NATS credentials. "
+                "Install with: pip install nkeys pyjwt"
+            )
+
+        self.operator_jwt = operator_jwt
+        self.operator_seed = operator_seed
+        self.account_public_key = account_public_key
+
+        # Create signing key from operator seed
+        try:
+            self.signing_key = nkeys.from_seed(operator_seed.encode())
+            logger.info(
+                "nats_credentials_manager_initialized",
+                operator_key=self.signing_key.public_key.decode()[:10] + "...",
+            )
+        except Exception as e:
+            logger.error("nats_signing_key_init_failed", error=str(e))
+            raise
+
+    def create_worker_credentials(
+        self,
+        worker_id: str,
+        organization_id: str,
+        ttl_hours: int = 24,
+    ) -> WorkerCredentials:
+        """
+        Create temporary NATS user credentials (JWT + seed) for worker.
+
+        Permissions:
+        - Publish: events.{organization_id}.{worker_id}.>
+        - Subscribe: None (workers only publish)
+
+        Args:
+            worker_id: Worker UUID
+            organization_id: Organization ID
+            ttl_hours: Credential time-to-live in hours
+
+        Returns:
+            WorkerCredentials with JWT and seed
+
+        Raises:
+            Exception: If credential generation fails
+        """
+        try:
+            import nkeys
+            import jwt as jwt_lib
+
+            # Generate user nkey
+            user_key = nkeys.create_user()
+            user_seed = user_key.seed.decode()
+            user_public_key = user_key.public_key.decode()
+
+            # Calculate expiration
+            now = datetime.now(timezone.utc)
+            expires_at = now + timedelta(hours=ttl_hours)
+
+            # Subject prefix for this worker
+            subject_prefix = f"events.{organization_id}.{worker_id}"
+
+            # Create user JWT claims
+            # Reference: https://docs.nats.io/running-a-nats-service/configuration/securing_nats/jwt
+            claims = {
+                "jti": f"worker-{worker_id}",  # JWT ID
+                "iat": int(now.timestamp()),  # Issued at
+                "iss": self.signing_key.public_key.decode(),  # Issuer (operator key)
+                "sub": user_public_key,  # Subject (user public key)
+                "exp": int(expires_at.timestamp()),  # Expiration
+                "nats": {
+                    "pub": {
+                        # Allow publishing to worker-specific subjects
+                        "allow": [f"{subject_prefix}.>"]
+                    },
+                    "sub": {
+                        # Workers don't subscribe
+                        "allow": []
+                    },
+                    "subs": -1,  # Unlimited subscriptions (not used)
+                    "data": -1,  # Unlimited data
+                    "payload": -1,  # Unlimited payload size
+                    "type": "user",
+                    "version": 2,
+                },
+            }
+
+            # Sign JWT with operator key using Ed25519
+            # nkeys uses Ed25519, so we sign manually
+            signing_input = self._encode_jwt_parts(claims)
+            signature = self.signing_key.sign(signing_input)
+
+            # Build complete JWT
+            import base64
+
+            jwt_parts = [
+                base64.urlsafe_b64encode(
+                    json.dumps({"typ": "JWT", "alg": "ed25519-nkey"}).encode()
+                )
+                .decode()
+                .rstrip("="),
+                base64.urlsafe_b64encode(json.dumps(claims).encode())
+                .decode()
+                .rstrip("="),
+                base64.urlsafe_b64encode(signature).decode().rstrip("="),
+            ]
+
+            user_jwt = ".".join(jwt_parts)
+
+            logger.info(
+                "nats_worker_credentials_created",
+                worker_id=worker_id[:8],
+                organization_id=organization_id,
+                subject_prefix=subject_prefix,
+                expires_at=expires_at.isoformat(),
+                ttl_hours=ttl_hours,
+            )
+
+            return WorkerCredentials(
+                jwt=user_jwt,
+                seed=user_seed,
+                public_key=user_public_key,
+                subject_prefix=subject_prefix,
+                expires_at=expires_at,
+            )
+
+        except ImportError as e:
+            logger.error("nats_credentials_dependency_missing", error=str(e))
+            raise
+
+        except Exception as e:
+            logger.error(
+                "nats_worker_credentials_creation_failed",
+                error=str(e),
+                worker_id=worker_id[:8],
+                organization_id=organization_id,
+            )
+            raise
+
+    def _encode_jwt_parts(self, claims: dict) -> bytes:
+        """
+        Encode JWT header and payload for signing.
+
+        Args:
+            claims: JWT claims
+
+        Returns:
+            Signing input (header.payload as bytes)
+        """
+        import base64
+        import json
+
+        header = {"typ": "JWT", "alg": "ed25519-nkey"}
+
+        header_b64 = (
+            base64.urlsafe_b64encode(json.dumps(header).encode())
+            .decode()
+            .rstrip("=")
+        )
+        payload_b64 = (
+            base64.urlsafe_b64encode(json.dumps(claims).encode())
+            .decode()
+            .rstrip("=")
+        )
+
+        signing_input = f"{header_b64}.{payload_b64}"
+        return signing_input.encode()
+
+    def save_credentials_file(
+        self, credentials: WorkerCredentials, file_path: str
+    ) -> None:
+        """
+        Save credentials to .creds file format.
+
+        The .creds file format contains both the JWT and seed in a specific format
+        that NATS clients can read.
+
+        Args:
+            credentials: Worker credentials
+            file_path: Path to save .creds file
+
+        Raises:
+            Exception: If file write fails
+        """
+        try:
+            # Standard NATS .creds file format
+            creds_content = f"""-----BEGIN NATS USER JWT-----
+{credentials.jwt}
+------END NATS USER JWT------
+
+************************* IMPORTANT *************************
+NKEY Seed printed below can be used to sign and prove identity.
+NKEYs are sensitive and should be treated as secrets.
+
+-----BEGIN USER NKEY SEED-----
+{credentials.seed}
+------END USER NKEY SEED------
+
+*************************************************************
+"""
+
+            # Write file
+            with open(file_path, "w") as f:
+                f.write(creds_content)
+
+            # Secure file permissions (owner read/write only)
+            os.chmod(file_path, 0o600)
+
+            logger.info(
+                "nats_credentials_file_saved",
+                file_path=file_path,
+                permissions="0600",
+            )
+
+        except Exception as e:
+            logger.error(
+                "nats_credentials_file_save_failed",
+                error=str(e),
+                file_path=file_path,
+            )
+            raise
+
+    @staticmethod
+    def get_credentials_string(credentials: WorkerCredentials) -> str:
+        """
+        Get credentials as a string (for passing to workers without file).
+
+        Args:
+            credentials: Worker credentials
+
+        Returns:
+            Credentials in .creds file format as string
+        """
+        return f"""-----BEGIN NATS USER JWT-----
+{credentials.jwt}
+------END NATS USER JWT------
+
+************************* IMPORTANT *************************
+NKEY Seed printed below can be used to sign and prove identity.
+NKEYs are sensitive and should be treated as secrets.
+
+-----BEGIN USER NKEY SEED-----
+{credentials.seed}
+------END USER NKEY SEED------
+
+*************************************************************
+"""