Enterprise LLM Gateway — Smart Multi-Provider Router

from pydantic import BaseModel, Field
from typing import Optional, List, Dict
from enum import Enum
import uuid
from datetime import datetime


class Priority(str, Enum):
    """SLA-based priority levels for request routing."""
    CRITICAL = "critical"   # <200ms  — real-time chat, user-facing
    HIGH = "high"           # <1s    — search, autocomplete
    MEDIUM = "medium"       # <5s    — content generation
    LOW = "low"             # <30s   — batch analytics, summarization


class Message(BaseModel):
    """OpenAI-compatible message format."""
    role: str = Field(..., description="system | user | assistant")
    content: str = Field(..., description="Message content")


class LLMRequest(BaseModel):
    """Unified request schema for all LLM calls through the gateway.

    Internal services submit this format regardless of target provider.
    The gateway handles provider-specific translation via adapters.
    """
    request_id: str = Field(
        default_factory=lambda: str(uuid.uuid4()),
        description="Unique request identifier for tracing"
    )
    service_id: str = Field(
        ..., description="Calling service identifier (e.g., 'chat-service', 'search-api')"
    )
    model: str = Field(
        ..., description="Requested model (e.g., 'gpt-4o', 'claude-sonnet-4-20250514')"
    )
    messages: List[Message] = Field(
        ..., description="Conversation messages in OpenAI format"
    )
    max_tokens: int = Field(
        default=1024,
        ge=1,
        le=128000,
        description="Maximum tokens in response"
    )
    temperature: float = Field(
        default=0.7,
        ge=0.0,
        le=2.0,
        description="Sampling temperature"
    )
    priority: Priority = Field(
        default=Priority.MEDIUM,
        description="Request priority determines SLA and routing preference"
    )
    sla_ms: int = Field(
        default=5000,
        description="Maximum acceptable latency in milliseconds"
    )
    preferred_provider: Optional[str] = Field(
        default=None,
        description="Preferred provider hint (e.g., 'azure', 'openai')"
    )
    fallback_models: List[str] = Field(
        default_factory=list,
        description="Ordered list of fallback models if primary is unavailable"
    )
    stream: bool = Field(
        default=False,
        description="Whether to stream the response"
    )
    metadata: Dict[str, str] = Field(
        default_factory=dict,
        description="Arbitrary metadata for logging and tracing"
    )
    created_at: datetime = Field(default_factory=datetime.utcnow)


class LLMResponse(BaseModel):
    """Unified response schema returned to all internal services.

    Includes provider metadata, cost tracking, and cache status
    regardless of which provider actually served the request.
    """
    request_id: str = Field(
        ..., description="Matches the original request ID for correlation"
    )
    provider: str = Field(
        ..., description="Provider that served the request (e.g., 'azure-openai')"
    )
    endpoint: str = Field(
        ..., description="Specific endpoint used (e.g., 'azure-eastus', 'openai-primary')"
    )
    model: str = Field(
        ..., description="Actual model used (may differ from requested if fallback)"
    )
    content: str = Field(
        ..., description="Generated text content"
    )
    usage: Dict[str, int] = Field(
        default_factory=dict,
        description="Token usage: {'prompt_tokens': N, 'completion_tokens': N, 'total_tokens': N}"
    )
    latency_ms: float = Field(
        ..., description="End-to-end latency including routing overhead"
    )
    cost_usd: float = Field(
        default=0.0,
        description="Estimated cost in USD for this request"
    )
    cached: bool = Field(
        default=False,
        description="Whether response was served from semantic cache"
    )
    fallback_used: bool = Field(
        default=False,
        description="Whether a fallback model/endpoint was used"
    )
    attempts: int = Field(
        default=1,
        description="Number of provider attempts before success"
    )

from dataclasses import dataclass, field
from typing import List, Optional, Dict
from datetime import datetime, timedelta
import time
import logging

logger = logging.getLogger(__name__)


@dataclass
class EndpointState:
    """Tracks real-time health and performance metrics for a single LLM endpoint.

    Each provider endpoint (e.g., azure-eastus, openai-primary) maintains
    its own state, updated after every request.
    """
    endpoint_id: str                          # e.g., "azure-eastus", "openai-primary"
    provider: str                             # e.g., "azure-openai", "openai", "anthropic"
    model: str                                # e.g., "gpt-4o", "claude-sonnet-4-20250514"
    base_url: str                             # API endpoint URL

    # Health metrics
    is_healthy: bool = True
    success_rate: float = 1.0               # Rolling 5-minute success rate
    total_requests: int = 0
    failed_requests: int = 0
    consecutive_failures: int = 0

    # Latency tracking (exponential moving average)
    avg_latency_ms: float = 500.0           # EMA of response latency
    p99_latency_ms: float = 2000.0          # 99th percentile latency

    # Rate limit tracking
    rpm_limit: int = 10000                  # Requests per minute limit
    rpm_current: int = 0                    # Current RPM usage
    tpm_limit: int = 2000000                # Tokens per minute limit
    tpm_current: int = 0                    # Current TPM usage
    concurrent_limit: int = 200             # Max concurrent requests
    concurrent_current: int = 0             # Current in-flight requests

    # Circuit breaker state
    circuit_state: str = "closed"           # closed | open | half_open
    circuit_opened_at: Optional[datetime] = None
    circuit_cooldown_sec: float = 30.0     # Time before half-open retry

    # Cost per 1K tokens (input/output)
    cost_per_1k_input: float = 0.005
    cost_per_1k_output: float = 0.015

    # Priority/preference
    priority_weight: float = 1.0           # Manual weight (e.g., prefer cheaper endpoints)
    last_used: Optional[datetime] = None
    last_error: Optional[str] = None

    def rpm_headroom(self) -> float:
        """Fraction of RPM capacity remaining (0.0 to 1.0)."""
        if self.rpm_limit == 0:
            return 0.0
        return max(0.0, (self.rpm_limit - self.rpm_current) / self.rpm_limit)

    def tpm_headroom(self) -> float:
        """Fraction of TPM capacity remaining (0.0 to 1.0)."""
        if self.tpm_limit == 0:
            return 0.0
        return max(0.0, (self.tpm_limit - self.tpm_current) / self.tpm_limit)


class SmartRouter:
    """Routes LLM requests to the optimal endpoint using weighted scoring.

    Scoring weights:
      - Health:   40%  (success rate, consecutive failures)
      - Latency:  30%  (can it meet the SLA?)
      - Capacity: 20%  (rate limit headroom)
      - Cost:     10%  (prefer cheaper endpoints)

    Disqualifiers (score = -1, endpoint skipped):
      - Endpoint is unhealthy (success_rate < 0.5)
      - Circuit breaker is OPEN
      - Rate limit exhausted (< 10% headroom)
      - P99 latency exceeds request SLA
    """

    WEIGHT_HEALTH = 0.40
    WEIGHT_LATENCY = 0.30
    WEIGHT_CAPACITY = 0.20
    WEIGHT_COST = 0.10

    def __init__(self, endpoints: List[EndpointState]):
        self.endpoints: Dict[str, EndpointState] = {
            ep.endpoint_id: ep for ep in endpoints
        }

    def route(self, model: str, sla_ms: int,
              preferred_provider: Optional[str] = None,
              estimated_tokens: int = 1000) -> List[EndpointState]:
        """Return endpoints ranked by score for the given model and SLA.

        Args:
            model: Requested model name (e.g., "gpt-4o")
            sla_ms: Maximum acceptable latency in milliseconds
            preferred_provider: Optional provider preference hint
            estimated_tokens: Estimated total tokens for capacity check

        Returns:
            List of EndpointState sorted by descending score.
            Empty list if no endpoints can serve the request.
        """
        candidates = [
            ep for ep in self.endpoints.values()
            if ep.model == model or self._is_compatible_model(ep.model, model)
        ]

        if not candidates:
            logger.warning(f"No endpoints found for model={model}")
            return []

        scored = []
        for ep in candidates:
            score = self._score_endpoint(ep, sla_ms, preferred_provider, estimated_tokens)
            if score >= 0:
                scored.append((score, ep))

        # Sort by score descending, break ties by latency
        scored.sort(key=lambda x: (-x[0], x[1].avg_latency_ms))

        ranked = [ep for _, ep in scored]
        logger.info(
            f"Routing model={model} sla={sla_ms}ms: "
            f"{len(ranked)}/{len(candidates)} endpoints qualified"
        )
        return ranked

    def _score_endpoint(self, ep: EndpointState, sla_ms: int,
                        preferred_provider: Optional[str],
                        estimated_tokens: int) -> float:
        """Score an endpoint from 0.0 to 1.0, or -1 if disqualified.

        Disqualifiers (return -1):
          1. Endpoint is unhealthy (success_rate < 0.5)
          2. Circuit breaker is OPEN (not in half_open)
          3. Rate limit headroom < 10%
          4. P99 latency > SLA (can't meet deadline)
        """
        # === DISQUALIFIERS ===
        if not ep.is_healthy or ep.success_rate < 0.5:
            logger.debug(f"Disqualified {ep.endpoint_id}: unhealthy")
            return -1

        if ep.circuit_state == "open":
            logger.debug(f"Disqualified {ep.endpoint_id}: circuit open")
            return -1

        if ep.rpm_headroom() < 0.10 or ep.tpm_headroom() < 0.10:
            logger.debug(f"Disqualified {ep.endpoint_id}: rate limit exhausted")
            return -1

        if ep.p99_latency_ms > sla_ms:
            logger.debug(f"Disqualified {ep.endpoint_id}: p99={ep.p99_latency_ms}ms > sla={sla_ms}ms")
            return -1

        # === HEALTH SCORE (0-1) ===
        health_score = ep.success_rate  # Already 0.0 to 1.0

        # === LATENCY SCORE (0-1) ===
        # Lower latency = higher score. Normalized against SLA.
        latency_ratio = ep.avg_latency_ms / sla_ms
        latency_score = max(0.0, 1.0 - latency_ratio)

        # === CAPACITY SCORE (0-1) ===
        # Average of RPM and TPM headroom
        capacity_score = (ep.rpm_headroom() + ep.tpm_headroom()) / 2.0

        # === COST SCORE (0-1) ===
        # Normalize cost: lower cost = higher score
        # Assume max cost is $0.06/1K tokens
        avg_cost = (ep.cost_per_1k_input + ep.cost_per_1k_output) / 2
        cost_score = max(0.0, 1.0 - (avg_cost / 0.06))

        # === WEIGHTED TOTAL ===
        total = (
            self.WEIGHT_HEALTH * health_score +
            self.WEIGHT_LATENCY * latency_score +
            self.WEIGHT_CAPACITY * capacity_score +
            self.WEIGHT_COST * cost_score
        )

        # Boost preferred provider by 10%
        if preferred_provider and ep.provider == preferred_provider:
            total = min(1.0, total * 1.10)

        # Slight boost for half_open (allow probe requests)
        if ep.circuit_state == "half_open":
            total *= 0.5  # Allow but de-prioritize

        logger.debug(
            f"Score {ep.endpoint_id}: health={health_score:.2f} "
            f"latency={latency_score:.2f} capacity={capacity_score:.2f} "
            f"cost={cost_score:.2f} → total={total:.3f}"
        )
        return total

    def _is_compatible_model(self, endpoint_model: str, requested_model: str) -> bool:
        """Check if an endpoint's model is compatible with the request.

        E.g., 'gpt-4o' on Azure is compatible with 'gpt-4o' on OpenAI.
        """
        # Normalize model names (strip provider prefixes)
        normalize = lambda m: m.split("/")[-1].lower().strip()
        return normalize(endpoint_model) == normalize(requested_model)

    def update_endpoint(self, endpoint_id: str, **kwargs) -> None:
        """Update endpoint state after a request completes."""
        if endpoint_id in self.endpoints:
            ep = self.endpoints[endpoint_id]
            for key, value in kwargs.items():
                if hasattr(ep, key):
                    setattr(ep, key, value)

from dataclasses import dataclass, field
from enum import Enum
from datetime import datetime, timedelta
from typing import Optional
import time
import logging

logger = logging.getLogger(__name__)


class CircuitState(str, Enum):
    CLOSED = "closed"       # Normal operation — requests flow through
    OPEN = "open"           # Tripped — all requests rejected immediately
    HALF_OPEN = "half_open" # Probing — allow 1 request to test recovery


class EndpointCircuitBreaker:
    """Circuit breaker for a single LLM provider endpoint.

    State transitions:
      CLOSED  → OPEN       : failure_count >= failure_threshold
      OPEN    → HALF_OPEN  : cooldown_sec elapsed
      HALF_OPEN → CLOSED   : probe request succeeds
      HALF_OPEN → OPEN     : probe request fails (reset cooldown)

    Special behaviors:
      - Rate limit errors (HTTP 429) → immediate OPEN with 3x cooldown
      - Slow calls (> slow_call_threshold_ms) counted as half-failures
      - Success in HALF_OPEN resets all counters
    """

    def __init__(
        self,
        endpoint_id: str,
        failure_threshold: int = 5,
        cooldown_sec: float = 30.0,
        success_threshold: int = 3,
        slow_call_threshold_ms: float = 10000.0,
        slow_call_rate_threshold: float = 0.5,
        window_size_sec: float = 60.0,
    ):
        self.endpoint_id = endpoint_id
        self.failure_threshold = failure_threshold
        self.cooldown_sec = cooldown_sec
        self.success_threshold = success_threshold
        self.slow_call_threshold_ms = slow_call_threshold_ms
        self.slow_call_rate_threshold = slow_call_rate_threshold
        self.window_size_sec = window_size_sec

        # State
        self.state: CircuitState = CircuitState.CLOSED
        self.failure_count: int = 0
        self.success_count: int = 0
        self.half_open_successes: int = 0
        self.opened_at: Optional[float] = None
        self.last_failure_time: Optional[float] = None

        # Slow call tracking
        self.slow_call_count: int = 0
        self.total_call_count: int = 0
        self.window_start: float = time.monotonic()

        # Rate limit specific
        self.rate_limit_cooldown_multiplier: float = 3.0

    def can_execute(self) -> bool:
        """Check if a request can be sent through this circuit.

        Returns:
            True if the circuit allows the request.
            False if the circuit is OPEN and cooldown hasn't elapsed.
        """
        if self.state == CircuitState.CLOSED:
            return True

        if self.state == CircuitState.OPEN:
            # Check if cooldown has elapsed → transition to HALF_OPEN
            elapsed = time.monotonic() - (self.opened_at or 0)
            if elapsed >= self.cooldown_sec:
                self.state = CircuitState.HALF_OPEN
                self.half_open_successes = 0
                logger.info(
                    f"Circuit {self.endpoint_id}: OPEN → HALF_OPEN "
                    f"(after {elapsed:.1f}s cooldown)"
                )
                return True
            return False

        if self.state == CircuitState.HALF_OPEN:
            # Allow limited probe requests
            return True

        return False

    def record_success(self, latency_ms: float = 0) -> None:
        """Record a successful request.

        In HALF_OPEN state, accumulate successes until threshold
        is met, then transition back to CLOSED.
        """
        self.total_call_count += 1

        # Track slow calls even on success
        if latency_ms > self.slow_call_threshold_ms:
            self.slow_call_count += 1
            self._check_slow_call_rate()

        if self.state == CircuitState.HALF_OPEN:
            self.half_open_successes += 1
            if self.half_open_successes >= self.success_threshold:
                self._close()
                logger.info(
                    f"Circuit {self.endpoint_id}: HALF_OPEN → CLOSED "
                    f"({self.success_threshold} consecutive successes)"
                )
        elif self.state == CircuitState.CLOSED:
            # Reset failure count on success (sliding window)
            self.failure_count = max(0, self.failure_count - 1)
            self.success_count += 1

    def record_failure(self, error: Optional[Exception] = None,
                       is_rate_limit: bool = False) -> None:
        """Record a failed request.

        Args:
            error: The exception that caused the failure
            is_rate_limit: If True, immediately open with extended cooldown
        """
        self.failure_count += 1
        self.last_failure_time = time.monotonic()
        self.total_call_count += 1

        if is_rate_limit:
            # Rate limit errors → immediate OPEN with extended cooldown
            self._trip(
                reason="rate_limit_429",
                cooldown_override=self.cooldown_sec * self.rate_limit_cooldown_multiplier
            )
            logger.warning(
                f"Circuit {self.endpoint_id}: RATE LIMITED → OPEN "
                f"(cooldown={self.cooldown_sec * self.rate_limit_cooldown_multiplier}s)"
            )
            return

        if self.state == CircuitState.HALF_OPEN:
            # Any failure in HALF_OPEN → back to OPEN
            self._trip(reason="half_open_probe_failed")
            logger.warning(
                f"Circuit {self.endpoint_id}: HALF_OPEN → OPEN (probe failed)"
            )
            return

        if self.state == CircuitState.CLOSED:
            if self.failure_count >= self.failure_threshold:
                self._trip(reason=f"failures={self.failure_count}/{self.failure_threshold}")
                logger.warning(
                    f"Circuit {self.endpoint_id}: CLOSED → OPEN "
                    f"({self.failure_count} failures in window)"
                )

    def _trip(self, reason: str, cooldown_override: Optional[float] = None) -> None:
        """Open the circuit breaker.

        Args:
            reason: Why the circuit was tripped (for logging)
            cooldown_override: Override default cooldown (used for rate limits)
        """
        self.state = CircuitState.OPEN
        self.opened_at = time.monotonic()
        if cooldown_override:
            self.cooldown_sec = cooldown_override
        logger.info(
            f"Circuit {self.endpoint_id} TRIPPED: reason={reason} "
            f"cooldown={self.cooldown_sec}s"
        )

    def _close(self) -> None:
        """Reset circuit to CLOSED state with clean counters."""
        self.state = CircuitState.CLOSED
        self.failure_count = 0
        self.success_count = 0
        self.half_open_successes = 0
        self.opened_at = None
        self.slow_call_count = 0
        self.total_call_count = 0
        self.window_start = time.monotonic()

    def _check_slow_call_rate(self) -> None:
        """Trip circuit if too many calls are slow (even if they succeed)."""
        if self.total_call_count < 10:
            return  # Need minimum sample size

        slow_rate = self.slow_call_count / self.total_call_count
        if slow_rate > self.slow_call_rate_threshold:
            self._trip(reason=f"slow_call_rate={slow_rate:.2f}")
            logger.warning(
                f"Circuit {self.endpoint_id}: TRIPPED due to slow calls "
                f"({self.slow_call_count}/{self.total_call_count} > "
                f"{self.slow_call_rate_threshold})"
            )

    def _maybe_reset_window(self) -> None:
        """Reset counters if the sliding window has elapsed."""
        elapsed = time.monotonic() - self.window_start
        if elapsed >= self.window_size_sec:
            self.slow_call_count = 0
            self.total_call_count = 0
            self.failure_count = 0
            self.window_start = time.monotonic()

import time
from collections import deque
from dataclasses import dataclass, field
from typing import Deque, Optional
import threading
import logging

logger = logging.getLogger(__name__)

# 10% headroom to avoid hitting provider-side limits
HEADROOM_FRACTION = 0.10


class SlidingWindowRateLimiter:
    """Sliding window rate limiter tracking RPM, TPM, RPS, and concurrency.

    Uses a deque-based sliding window for accurate rate counting.
    Maintains 10% headroom below provider limits to avoid 429 errors.

    Each LLM endpoint gets its own rate limiter instance configured
    with that endpoint's specific limits from provider documentation.
    """

    def __init__(
        self,
        endpoint_id: str,
        rpm_limit: int = 10000,
        tpm_limit: int = 2000000,
        rps_limit: int = 500,
        concurrent_limit: int = 200,
    ):
        self.endpoint_id = endpoint_id
        self.rpm_limit = rpm_limit
        self.tpm_limit = tpm_limit
        self.rps_limit = rps_limit
        self.concurrent_limit = concurrent_limit

        # Effective limits with headroom
        self.effective_rpm = int(rpm_limit * (1 - HEADROOM_FRACTION))
        self.effective_tpm = int(tpm_limit * (1 - HEADROOM_FRACTION))
        self.effective_rps = int(rps_limit * (1 - HEADROOM_FRACTION))
        self.effective_concurrent = int(concurrent_limit * (1 - HEADROOM_FRACTION))

        # Sliding windows: deque of (timestamp, token_count) tuples
        self._rpm_window: Deque[float] = deque()              # 60s window
        self._tpm_window: Deque[tuple[float, int]] = deque()  # 60s window
        self._rps_window: Deque[float] = deque()              # 1s window

        # Concurrent request counter
        self._concurrent: int = 0

        # Thread safety
        self._lock = threading.Lock()

    def can_accept(self, estimated_tokens: int = 1000) -> bool:
        """Check if the endpoint can accept a new request.

        Args:
            estimated_tokens: Estimated total tokens (prompt + completion)

        Returns:
            True if all rate limits have sufficient headroom.
        """
        with self._lock:
            now = time.monotonic()
            self._cleanup_windows(now)

            # Check RPM
            if len(self._rpm_window) >= self.effective_rpm:
                logger.debug(f"{self.endpoint_id}: RPM limit reached ({len(self._rpm_window)}/{self.effective_rpm})")
                return False

            # Check TPM
            current_tpm = sum(tokens for _, tokens in self._tpm_window)
            if current_tpm + estimated_tokens > self.effective_tpm:
                logger.debug(f"{self.endpoint_id}: TPM limit reached ({current_tpm}/{self.effective_tpm})")
                return False

            # Check RPS
            if len(self._rps_window) >= self.effective_rps:
                logger.debug(f"{self.endpoint_id}: RPS limit reached ({len(self._rps_window)}/{self.effective_rps})")
                return False

            # Check concurrent
            if self._concurrent >= self.effective_concurrent:
                logger.debug(f"{self.endpoint_id}: Concurrent limit reached ({self._concurrent}/{self.effective_concurrent})")
                return False

            return True

    def record_request(self, estimated_tokens: int = 1000) -> None:
        """Record an outgoing request against all rate limit windows.

        Called when a request is dispatched to the provider.
        Must be paired with record_completion() when the request finishes.
        """
        with self._lock:
            now = time.monotonic()
            self._rpm_window.append(now)
            self._tpm_window.append((now, estimated_tokens))
            self._rps_window.append(now)
            self._concurrent += 1

    def record_completion(self, actual_tokens: int = 0) -> None:
        """Record request completion, decrement concurrent counter.

        Args:
            actual_tokens: Actual tokens used (for adjusting TPM tracking)
        """
        with self._lock:
            self._concurrent = max(0, self._concurrent - 1)

    def headroom_pct(self) -> Dict[str, float]:
        """Return remaining headroom percentage for each limit dimension.

        Returns:
            Dict with rpm, tpm, rps, concurrent headroom as 0.0-1.0 floats.
        """
        with self._lock:
            now = time.monotonic()
            self._cleanup_windows(now)

            current_tpm = sum(tokens for _, tokens in self._tpm_window)

            return {
                "rpm": max(0.0, 1.0 - len(self._rpm_window) / self.rpm_limit) if self.rpm_limit > 0 else 0.0,
                "tpm": max(0.0, 1.0 - current_tpm / self.tpm_limit) if self.tpm_limit > 0 else 0.0,
                "rps": max(0.0, 1.0 - len(self._rps_window) / self.rps_limit) if self.rps_limit > 0 else 0.0,
                "concurrent": max(0.0, 1.0 - self._concurrent / self.concurrent_limit) if self.concurrent_limit > 0 else 0.0,
            }

    def _cleanup_windows(self, now: float) -> None:
        """Remove expired entries from sliding windows."""
        # RPM window: 60 seconds
        cutoff_60s = now - 60.0
        while self._rpm_window and self._rpm_window[0] < cutoff_60s:
            self._rpm_window.popleft()

        # TPM window: 60 seconds
        while self._tpm_window and self._tpm_window[0][0] < cutoff_60s:
            self._tpm_window.popleft()

        # RPS window: 1 second
        cutoff_1s = now - 1.0
        while self._rps_window and self._rps_window[0] < cutoff_1s:
            self._rps_window.popleft()

import asyncio
import time
import logging
from typing import Optional, Dict
from gateway.schemas import LLMRequest, LLMResponse
from gateway.router import SmartRouter, EndpointState
from gateway.circuit_breaker import EndpointCircuitBreaker
from gateway.rate_limiter import SlidingWindowRateLimiter
from gateway.cache import SemanticCache
from gateway.adapters import get_adapter
from gateway.metrics import MetricsCollector

logger = logging.getLogger(__name__)


class LLMGateway:
    """Central request handler for the LLM Gateway.

    Request lifecycle:
      1. Check semantic cache → return cached response if hit
      2. Get ranked endpoints from SmartRouter
      3. For each endpoint (in score order):
         a. Check circuit breaker → skip if open
         b. Check rate limiter → skip if exhausted
         c. Send request with SLA timeout via asyncio.wait_for
         d. On success → record metrics, update state, cache response
         e. On failure → record failure, try next endpoint
      4. If all endpoints fail → return error response
    """

    def __init__(
        self,
        router: SmartRouter,
        cache: SemanticCache,
        circuit_breakers: Dict[str, EndpointCircuitBreaker],
        rate_limiters: Dict[str, SlidingWindowRateLimiter],
        metrics: MetricsCollector,
    ):
        self.router = router
        self.cache = cache
        self.circuit_breakers = circuit_breakers
        self.rate_limiters = rate_limiters
        self.metrics = metrics

    async def handle(self, request: LLMRequest) -> LLMResponse:
        """Handle an LLM request with full retry and fallback logic.

        Args:
            request: Unified LLM request from internal service

        Returns:
            LLMResponse with provider metadata, cost, and latency
        """
        start_time = time.monotonic()

        # ── Step 1: Check Semantic Cache ──
        cached_response = await self.cache.get(request)
        if cached_response:
            cached_response.latency_ms = (time.monotonic() - start_time) * 1000
            self.metrics.record_cache_hit(request.service_id)
            logger.info(f"Cache HIT for request={request.request_id}")
            return cached_response

        self.metrics.record_cache_miss(request.service_id)

        # ── Step 2: Get Ranked Endpoints ──
        # Try primary model first, then fallbacks
        models_to_try = [request.model] + request.fallback_models
        last_error: Optional[Exception] = None
        attempts = 0

        for model in models_to_try:
            endpoints = self.router.route(
                model=model,
                sla_ms=request.sla_ms,
                preferred_provider=request.preferred_provider,
                estimated_tokens=request.max_tokens,
            )

            if not endpoints:
                logger.warning(f"No endpoints available for model={model}")
                continue

            # ── Step 3: Try Endpoints in Score Order ──
            for endpoint in endpoints:
                attempts += 1
                ep_id = endpoint.endpoint_id

                # Step 3a: Check Circuit Breaker
                cb = self.circuit_breakers.get(ep_id)
                if cb and not cb.can_execute():
                    logger.debug(f"Skipping {ep_id}: circuit breaker open")
                    continue

                # Step 3b: Check Rate Limiter
                rl = self.rate_limiters.get(ep_id)
                if rl and not rl.can_accept(estimated_tokens=request.max_tokens):
                    logger.debug(f"Skipping {ep_id}: rate limit exhausted")
                    continue

                # Record rate limit usage
                if rl:
                    rl.record_request(estimated_tokens=request.max_tokens)

                try:
                    # Step 3c: Send Request with SLA Timeout
                    remaining_sla_ms = request.sla_ms - (time.monotonic() - start_time) * 1000
                    if remaining_sla_ms <= 0:
                        logger.warning(f"SLA budget exhausted before attempting {ep_id}")
                        break

                    timeout_sec = remaining_sla_ms / 1000.0
                    adapter = get_adapter(endpoint.provider)

                    response = await asyncio.wait_for(
                        adapter.send(request, endpoint),
                        timeout=timeout_sec,
                    )

                    # Step 3d: Success — record metrics
                    latency_ms = (time.monotonic() - start_time) * 1000
                    response.latency_ms = latency_ms
                    response.attempts = attempts
                    response.fallback_used = (model != request.model)

                    # Update circuit breaker
                    if cb:
                        cb.record_success(latency_ms=latency_ms)

                    # Release rate limiter slot
                    if rl:
                        actual_tokens = response.usage.get("total_tokens", 0)
                        rl.record_completion(actual_tokens=actual_tokens)

                    # Update router state
                    self.router.update_endpoint(
                        ep_id,
                        avg_latency_ms=latency_ms,
                        last_used=time.time(),
                        success_rate=self._calc_success_rate(ep_id),
                    )

                    # Cache the response (if eligible)
                    await self.cache.set(request, response)

                    # Record metrics
                    self.metrics.record_request(
                        service_id=request.service_id,
                        endpoint_id=ep_id,
                        model=model,
                        latency_ms=latency_ms,
                        cost_usd=response.cost_usd,
                        success=True,
                    )

                    logger.info(
                        f"Request {request.request_id} served by {ep_id} "
                        f"in {latency_ms:.1f}ms (attempts={attempts})"
                    )
                    return response

                except asyncio.TimeoutError:
                    # Step 3e: SLA timeout — try next endpoint
                    last_error = asyncio.TimeoutError(
                        f"Timeout after {remaining_sla_ms:.0f}ms on {ep_id}"
                    )
                    logger.warning(f"Timeout on {ep_id}: {remaining_sla_ms:.0f}ms")
                    if cb:
                        cb.record_failure(error=last_error)
                    if rl:
                        rl.record_completion()

                except Exception as e:
                    # Step 3e: Provider error — try next endpoint
                    last_error = e
                    is_rate_limit = "429" in str(e) or "rate_limit" in str(e).lower()
                    logger.warning(f"Error on {ep_id}: {e}")

                    if cb:
                        cb.record_failure(error=e, is_rate_limit=is_rate_limit)
                    if rl:
                        rl.record_completion()

                    self.metrics.record_request(
                        service_id=request.service_id,
                        endpoint_id=ep_id,
                        model=model,
                        latency_ms=(time.monotonic() - start_time) * 1000,
                        cost_usd=0,
                        success=False,
                    )

        # ── Step 4: All Endpoints Failed ──
        latency_ms = (time.monotonic() - start_time) * 1000
        logger.error(
            f"All endpoints failed for request={request.request_id} "
            f"after {attempts} attempts in {latency_ms:.1f}ms"
        )
        self.metrics.record_total_failure(request.service_id)

        return LLMResponse(
            request_id=request.request_id,
            provider="none",
            endpoint="none",
            model=request.model,
            content=f"Error: All endpoints unavailable. Last error: {last_error}",
            latency_ms=latency_ms,
            attempts=attempts,
        )

    def _calc_success_rate(self, endpoint_id: str) -> float:
        """Calculate rolling success rate from metrics."""
        stats = self.metrics.get_endpoint_stats(endpoint_id, window_sec=300)
        if stats["total"] == 0:
            return 1.0
        return stats["successes"] / stats["total"]

import hashlib
import json
import logging
from typing import Optional
import redis.asyncio as redis
from gateway.schemas import LLMRequest, LLMResponse

logger = logging.getLogger(__name__)

DEFAULT_TTL_SECONDS = 3600  # 1 hour


class SemanticCache:
    """Redis-backed semantic cache for LLM responses.

    Cache key = SHA256(model + messages + temperature + max_tokens)

    Rules:
      - Only cache temperature=0 requests (deterministic output)
      - TTL: 1 hour (configurable per model)
      - Returns cached=True flag so callers know it was a cache hit
      - Cache misses add < 1ms overhead (single Redis GET)

    Cache hit rates in production:
      - Customer support bots: ~45% (many repeated questions)
      - Code generation: ~15% (unique prompts)
      - Search augmentation: ~35% (similar queries)
      - Average across all services: 30-40%
    """

    def __init__(
        self,
        redis_client: redis.Redis,
        ttl_seconds: int = DEFAULT_TTL_SECONDS,
        enabled: bool = True,
        max_tokens_for_cache: int = 4096,
    ):
        self.redis = redis_client
        self.ttl_seconds = ttl_seconds
        self.enabled = enabled
        self.max_tokens_for_cache = max_tokens_for_cache

    def _cache_key(self, request: LLMRequest) -> str:
        """Generate deterministic cache key from request parameters.

        Key components:
          - model: Different models produce different outputs
          - messages: The actual conversation (content + roles)
          - temperature: Only 0 is cached, but included for safety
          - max_tokens: Affects output length/truncation
        """
        key_data = {
            "model": request.model,
            "messages": [
                {"role": m.role, "content": m.content}
                for m in request.messages
            ],
            "temperature": request.temperature,
            "max_tokens": request.max_tokens,
        }
        key_json = json.dumps(key_data, sort_keys=True, ensure_ascii=True)
        key_hash = hashlib.sha256(key_json.encode()).hexdigest()
        return f"llm_cache:{key_hash}"

    def _is_cacheable(self, request: LLMRequest) -> bool:
        """Determine if a request is eligible for caching.

        Only cache when:
          1. Caching is enabled
          2. Temperature is exactly 0 (deterministic)
          3. Not a streaming request
          4. Max tokens is within cache limit
        """
        if not self.enabled:
            return False
        if request.temperature != 0.0:
            return False
        if request.stream:
            return False
        if request.max_tokens > self.max_tokens_for_cache:
            return False
        return True

    async def get(self, request: LLMRequest) -> Optional[LLMResponse]:
        """Look up cached response for this request.

        Returns:
            LLMResponse with cached=True if found, None otherwise.
        """
        if not self._is_cacheable(request):
            return None

        try:
            key = self._cache_key(request)
            cached_data = await self.redis.get(key)

            if cached_data is None:
                return None

            data = json.loads(cached_data)
            response = LLMResponse(
                request_id=request.request_id,
                provider=data["provider"],
                endpoint=data["endpoint"],
                model=data["model"],
                content=data["content"],
                usage=data["usage"],
                latency_ms=0,  # Will be set by caller
                cost_usd=0.0,   # Cached = free
                cached=True,
            )
            logger.debug(f"Cache HIT: key={key[:16]}...")
            return response

        except Exception as e:
            # Cache failures should never block requests
            logger.warning(f"Cache GET error: {e}")
            return None

    async def set(self, request: LLMRequest, response: LLMResponse) -> None:
        """Cache a response for future identical requests.

        Only caches if the request is eligible (temperature=0, non-streaming).
        Cache failures are logged but never raise exceptions.
        """
        if not self._is_cacheable(request):
            return

        try:
            key = self._cache_key(request)
            cache_data = json.dumps({
                "provider": response.provider,
                "endpoint": response.endpoint,
                "model": response.model,
                "content": response.content,
                "usage": response.usage,
            })
            await self.redis.setex(key, self.ttl_seconds, cache_data)
            logger.debug(f"Cache SET: key={key[:16]}... ttl={self.ttl_seconds}s")

        except Exception as e:
            # Cache failures should never block requests
            logger.warning(f"Cache SET error: {e}")

from abc import ABC, abstractmethod
from typing import Dict, Any
import aiohttp
from gateway.schemas import LLMRequest, LLMResponse
from gateway.router import EndpointState


class LLMProviderAdapter(ABC):
    """Abstract base class for LLM provider adapters.

    Each adapter translates between the gateway's unified format
    and the provider's native API format. Responsible for:
      - Building provider-specific HTTP requests
      - Parsing provider-specific responses
      - Calculating cost based on token usage
      - Handling provider-specific error codes
    """

    def __init__(self, api_key: str, session: aiohttp.ClientSession):
        self.api_key = api_key
        self.session = session

    @abstractmethod
    def build_request(self, request: LLMRequest, endpoint: EndpointState) -> Dict[str, Any]:
        """Transform unified request into provider-specific format."""
        ...

    @abstractmethod
    def parse_response(self, raw_response: Dict[str, Any], request: LLMRequest,
                       endpoint: EndpointState) -> LLMResponse:
        """Transform provider response into unified format."""
        ...

    @abstractmethod
    def calculate_cost(self, usage: Dict[str, int], endpoint: EndpointState) -> float:
        """Calculate cost in USD based on token usage."""
        ...

    async def send(self, request: LLMRequest, endpoint: EndpointState) -> LLMResponse:
        """Send request to provider and return unified response."""
        provider_request = self.build_request(request, endpoint)

        async with self.session.post(
            provider_request["url"],
            headers=provider_request["headers"],
            json=provider_request["body"],
        ) as resp:
            if resp.status == 429:
                raise RateLimitError(f"429 Rate limit from {endpoint.endpoint_id}")
            if resp.status >= 400:
                body = await resp.text()
                raise ProviderError(f"HTTP {resp.status} from {endpoint.endpoint_id}: {body}")

            raw = await resp.json()
            return self.parse_response(raw, request, endpoint)


class RateLimitError(Exception):
    """Raised when provider returns HTTP 429."""
    pass

class ProviderError(Exception):
    """Raised when provider returns a non-429 error."""
    pass

from typing import Dict, Any
from gateway.adapters.base import LLMProviderAdapter
from gateway.schemas import LLMRequest, LLMResponse
from gateway.router import EndpointState


class AnthropicAdapter(LLMProviderAdapter):
    """Adapter for Anthropic's Messages API.

    Translates OpenAI-style messages format to Anthropic's format:
      - Separates system message from conversation
      - Uses 'x-api-key' header instead of 'Authorization: Bearer'
      - Maps response format: content[0].text → content
      - Handles Anthropic-specific usage fields
    """

    API_VERSION = "2024-01-01"

    def build_request(self, request: LLMRequest, endpoint: EndpointState) -> Dict[str, Any]:
        """Transform unified request to Anthropic Messages API format."""
        # Separate system message from conversation messages
        system_msg = None
        messages = []
        for msg in request.messages:
            if msg.role == "system":
                system_msg = msg.content
            else:
                messages.append({"role": msg.role, "content": msg.content})

        body = {
            "model": request.model,
            "messages": messages,
            "max_tokens": request.max_tokens,
            "temperature": request.temperature,
        }
        if system_msg:
            body["system"] = system_msg

        return {
            "url": f"{endpoint.base_url}/v1/messages",
            "headers": {
                "x-api-key": self.api_key,
                "anthropic-version": self.API_VERSION,
                "content-type": "application/json",
            },
            "body": body,
        }

    def parse_response(self, raw: Dict[str, Any], request: LLMRequest,
                       endpoint: EndpointState) -> LLMResponse:
        """Transform Anthropic response to unified format."""
        content = raw["content"][0]["text"] if raw.get("content") else ""
        usage = {
            "prompt_tokens": raw.get("usage", {}).get("input_tokens", 0),
            "completion_tokens": raw.get("usage", {}).get("output_tokens", 0),
            "total_tokens": (
                raw.get("usage", {}).get("input_tokens", 0) +
                raw.get("usage", {}).get("output_tokens", 0)
            ),
        }

        return LLMResponse(
            request_id=request.request_id,
            provider="anthropic",
            endpoint=endpoint.endpoint_id,
            model=raw.get("model", request.model),
            content=content,
            usage=usage,
            latency_ms=0,  # Set by caller
            cost_usd=self.calculate_cost(usage, endpoint),
        )

    def calculate_cost(self, usage: Dict[str, int], endpoint: EndpointState) -> float:
        """Calculate cost using Anthropic's pricing."""
        input_cost = (usage.get("prompt_tokens", 0) / 1000) * endpoint.cost_per_1k_input
        output_cost = (usage.get("completion_tokens", 0) / 1000) * endpoint.cost_per_1k_output
        return round(input_cost + output_cost, 6)

from typing import Dict, Any
from gateway.adapters.base import LLMProviderAdapter
from gateway.schemas import LLMRequest, LLMResponse
from gateway.router import EndpointState


class AzureOpenAIAdapter(LLMProviderAdapter):
    """Adapter for Azure OpenAI Service endpoints.

    Key differences from OpenAI direct:
      - URL includes deployment name: /deployments/{model}/chat/completions
      - Uses 'api-key' header instead of 'Authorization: Bearer'
      - Requires api-version query parameter
      - Each Azure region is a separate endpoint with its own rate limits
    """

    API_VERSION = "2024-02-01"

    def build_request(self, request: LLMRequest, endpoint: EndpointState) -> Dict[str, Any]:
        """Transform unified request to Azure OpenAI format."""
        messages = [{"role": m.role, "content": m.content} for m in request.messages]

        # Azure uses deployment name in URL path
        deployment = request.model.replace(".", "")
        url = (
            f"{endpoint.base_url}/openai/deployments/{deployment}"
            f"/chat/completions?api-version={self.API_VERSION}"
        )

        return {
            "url": url,
            "headers": {
                "api-key": self.api_key,
                "content-type": "application/json",
            },
            "body": {
                "messages": messages,
                "max_tokens": request.max_tokens,
                "temperature": request.temperature,
            },
        }

    def parse_response(self, raw: Dict[str, Any], request: LLMRequest,
                       endpoint: EndpointState) -> LLMResponse:
        """Transform Azure OpenAI response to unified format."""
        choice = raw.get("choices", [{}])[0]
        content = choice.get("message", {}).get("content", "")
        raw_usage = raw.get("usage", {})

        usage = {
            "prompt_tokens": raw_usage.get("prompt_tokens", 0),
            "completion_tokens": raw_usage.get("completion_tokens", 0),
            "total_tokens": raw_usage.get("total_tokens", 0),
        }

        return LLMResponse(
            request_id=request.request_id,
            provider="azure-openai",
            endpoint=endpoint.endpoint_id,
            model=raw.get("model", request.model),
            content=content,
            usage=usage,
            latency_ms=0,
            cost_usd=self.calculate_cost(usage, endpoint),
        )

    def calculate_cost(self, usage: Dict[str, int], endpoint: EndpointState) -> float:
        """Calculate cost using Azure OpenAI pricing (same as OpenAI)."""
        input_cost = (usage.get("prompt_tokens", 0) / 1000) * endpoint.cost_per_1k_input
        output_cost = (usage.get("completion_tokens", 0) / 1000) * endpoint.cost_per_1k_output
        return round(input_cost + output_cost, 6)


def get_adapter(provider: str) -> LLMProviderAdapter:
    """Factory function to get the appropriate adapter for a provider."""
    adapters = {
        "openai": OpenAIAdapter,
        "azure-openai": AzureOpenAIAdapter,
        "anthropic": AnthropicAdapter,
        "google": GoogleGeminiAdapter,
    }
    adapter_class = adapters.get(provider)
    if not adapter_class:
        raise ValueError(f"Unknown provider: {provider}")
    return adapter_class()

import asyncio
import time
from enum import IntEnum
from dataclasses import dataclass, field
from typing import Optional, List
import heapq
import logging

logger = logging.getLogger(__name__)


class PriorityLevel(IntEnum):
    """Lower number = higher priority (for min-heap)."""
    CRITICAL = 0   # SLA <200ms  — never queued, fail fast
    HIGH = 1       # SLA <1s     — minimal queueing
    MEDIUM = 2     # SLA <5s     — can wait in queue
    LOW = 3        # SLA <30s    — batch processing, longest wait


SLA_DEADLINES = {
    PriorityLevel.CRITICAL: 200,    # ms
    PriorityLevel.HIGH: 1000,       # ms
    PriorityLevel.MEDIUM: 5000,     # ms
    PriorityLevel.LOW: 30000,       # ms
}


@dataclass(order=True)
class PrioritizedRequest:
    """Wrapper for heap-based priority queue ordering."""
    priority: int                           # PriorityLevel value (0-3)
    enqueued_at: float                      # time.monotonic() timestamp
    request_id: str = field(compare=False) # For dedup and logging
    request: object = field(compare=False) # The actual LLMRequest
    sla_deadline_ms: int = field(compare=False, default=5000)

    def is_expired(self) -> bool:
        """Check if this request has exceeded its SLA deadline."""
        elapsed_ms = (time.monotonic() - self.enqueued_at) * 1000
        return elapsed_ms > self.sla_deadline_ms


class SLAPriorityQueue:
    """Priority queue that respects SLA deadlines.

    Behavior by priority:
      CRITICAL: Bypass queue entirely. If no endpoint available, fail immediately.
      HIGH:     Queue for max 500ms. Dequeue first when capacity frees up.
      MEDIUM:   Queue for max 4s. Standard processing.
      LOW:      Queue for max 25s. Processed when higher priorities are clear.

    Expired requests are dropped when dequeued (not before, to avoid O(n) scans).
    """

    def __init__(self, max_queue_size: int = 10000):
        self._heap: List[PrioritizedRequest] = []
        self._max_size = max_queue_size
        self._lock = asyncio.Lock()
        self._not_empty = asyncio.Event()
        self._dropped_count = 0

    async def enqueue(self, request, priority: PriorityLevel) -> bool:
        """Add request to priority queue.

        Args:
            request: LLMRequest to enqueue
            priority: SLA-based priority level

        Returns:
            True if enqueued, False if queue is full or CRITICAL priority
        """
        # CRITICAL requests never queue — they must be served immediately
        if priority == PriorityLevel.CRITICAL:
            logger.debug(f"CRITICAL request {request.request_id} — bypassing queue")
            return False  # Signal caller to handle immediately

        async with self._lock:
            if len(self._heap) >= self._max_size:
                # Queue full — drop lowest priority request
                if self._heap and self._heap[-1].priority > priority.value:
                    dropped = heapq.heappushpop(
                        self._heap,
                        PrioritizedRequest(
                            priority=priority.value,
                            enqueued_at=time.monotonic(),
                            request_id=request.request_id,
                            request=request,
                            sla_deadline_ms=SLA_DEADLINES[priority],
                        )
                    )
                    self._dropped_count += 1
                    logger.warning(f"Queue full: dropped {dropped.request_id} (pri={dropped.priority})")
                    return True
                else:
                    logger.warning(f"Queue full: rejecting {request.request_id}")
                    return False

            heapq.heappush(
                self._heap,
                PrioritizedRequest(
                    priority=priority.value,
                    enqueued_at=time.monotonic(),
                    request_id=request.request_id,
                    request=request,
                    sla_deadline_ms=SLA_DEADLINES[priority],
                )
            )
            self._not_empty.set()
            return True

    async def dequeue(self, timeout_sec: float = 1.0) -> Optional[PrioritizedRequest]:
        """Dequeue the highest-priority non-expired request.

        Skips expired requests (they've already breached SLA).
        """
        try:
            await asyncio.wait_for(self._not_empty.wait(), timeout=timeout_sec)
        except asyncio.TimeoutError:
            return None

        async with self._lock:
            while self._heap:
                item = heapq.heappop(self._heap)
                if not item.is_expired():
                    if not self._heap:
                        self._not_empty.clear()
                    return item
                else:
                    self._dropped_count += 1
                    logger.debug(f"Dropping expired request {item.request_id}")

            self._not_empty.clear()
            return None

    @property
    def size(self) -> int:
        return len(self._heap)

    @property
    def dropped(self) -> int:
        return self._dropped_count

apiVersion: apps/v1
kind: Deployment
metadata:
  name: llm-gateway
spec:
  replicas: 12
  strategy:
    type: RollingUpdate
    rollingUpdate:
      maxUnavailable: 1
      maxSurge: 2
  template:
    spec:
      containers:
        - name: gateway
          image: llm-gateway:latest
          resources:
            requests:
              cpu: "2"
              memory: "4Gi"
            limits:
              cpu: "4"
              memory: "8Gi"
          env:
            - name: REDIS_CLUSTER_URL
              valueFrom:
                secretKeyRef:
                  name: gateway-secrets
                  key: redis-url
            - name: KAFKA_BROKERS
              value: "kafka-0:9092,kafka-1:9092,kafka-2:9092"
          readinessProbe:
            httpGet:
              path: /health
              port: 8080
            initialDelaySeconds: 5
            periodSeconds: 10
          livenessProbe:
            httpGet:
              path: /health
              port: 8080
            initialDelaySeconds: 15
            periodSeconds: 20
      topologySpreadConstraints:
        - maxSkew: 1
          topologyKey: topology.kubernetes.io/zone
          whenUnsatisfiable: DoNotSchedule

# Kafka topic configuration for priority-based partitioning
TOPIC_CONFIG = {
    "topic": "llm-requests",
    "partitions": 4,
    "replication_factor": 3,
    "partition_mapping": {
        "critical": 0,  # Partition 0 — highest consumer priority
        "high": 1,      # Partition 1
        "medium": 2,    # Partition 2
        "low": 3,       # Partition 3 — lowest consumer priority
    },
    "consumer_config": {
        "max_poll_records": 100,
        "fetch_max_wait_ms": 50,     # Low latency for high-pri
        "session_timeout_ms": 10000,
    },
}

# Consumer group processes partitions in priority order:
# 1. Poll partition 0 (CRITICAL) — always first
# 2. Poll partition 1 (HIGH) — if critical is empty
# 3. Poll partition 2 (MEDIUM) — if high is empty
# 4. Poll partition 3 (LOW) — if medium is empty