proxy/router.py

import time
import logging
import threading
import httpx
from typing import List, Dict, Any, Optional, Tuple
import tiktoken
import litellm
from litellm import Router
from litellm.integrations.custom_logger import CustomLogger

from .config import ProxyConfig, ModelEndpointConfig

logger = logging.getLogger("proxy.router")



class LiteLLMProxyRouter:
    """
    Core OOP Routing Engine wrapping LiteLLM's Router.
    Implements Token Per Request (TPR) checks, load balancing, fallback routing, and token usage optimization.
    """
    def __init__(self, config: ProxyConfig):
        self.config = config
        
        # Configure litellm global settings
        litellm.telemetry = False
        litellm.drop_params = True # Safely drop unsupported params per provider
        
        # Convert our configurations to the shape LiteLLM expects
        model_list = self.config.to_litellm_model_list()
        
        # Initialize LiteLLM's core Router
        logger.info(f"Initializing LiteLLM Router with strategy: {self.config.routing_strategy}")
        self.router = Router(
            model_list=model_list,
            routing_strategy=self.config.routing_strategy,
            num_retries=self.config.num_retries,
            timeout=self.config.timeout,
            fallbacks=self.config.general_fallbacks,
            context_window_fallbacks=self.config.context_window_fallbacks
        )
        
        # Local metrics tracking for the microservice
        self.metrics_lock = threading.Lock()
        self.metrics = {
            "total_requests": 0,
            "successful_requests": 0,
            "failed_requests": 0,
            "total_input_tokens": 0,
            "total_output_tokens": 0,
            "provider_calls": {}, # tracks calls per physical model
            "fallback_events": 0,
        }
        
        self.usage_history = []
        
        # Initialize Redis connection for centralized load balancing
        try:
            import os
            import redis
            redis_host = os.environ.get("REDIS_HOST", "localhost")
            redis_port = int(os.environ.get("REDIS_PORT", 6379))
            self.redis_client = redis.Redis(host=redis_host, port=redis_port, decode_responses=True, socket_timeout=1.0)
            self.redis_client.ping()
            logger.info(f"Successfully connected to Redis at {redis_host}:{redis_port} for load-balancing usage tracking.")
        except Exception as e:
            logger.warning(f"Could not connect to Redis: {e}. Falling back to dynamic in-memory Redis simulation.")
            self.redis_client = None

        self.routing_logs = []
        
        # Initialize Tiktoken for TPR (Tokens Per Request) estimation
        try:
            self.tokenizer = tiktoken.get_encoding("cl100k_base")
            logger.info("Tiktoken tokenizer initialized successfully.")
        except Exception as e:
            logger.warning(f"Failed to initialize tiktoken, falling back to character approximation: {e}")
            self.tokenizer = None

        # Team Bring-Your-Own Guardrails Registry
        self.team_guardrails = {}
        self.guardrail_instances = {}

        # PII Guardrail state parameters
        self.pii_enabled = False
        self.pii_action = "MASK"
        self.pii_policy = None
        self.pii_guardrail = None

        # Priority Preference Routing State Parameters
        self.preference_enabled = False
        self.preference_list = []
        self.credit_limits = {
            "groq/llama-3.1-8b-instant": 0.05,
            "cerebras/llama3.1-8b": 0.05,
            "groq/llama-3.3-70b-versatile": 0.05,
            "together_ai/meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo": 0.05,
            "together_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo": 0.05,
            "ollama/llama3.1": 0.05
        }
        self.accumulated_spend = {
            "groq/llama-3.1-8b-instant": 0.0,
            "cerebras/llama3.1-8b": 0.0,
            "groq/llama-3.3-70b-versatile": 0.0,
            "together_ai/meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo": 0.0,
            "together_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo": 0.0,
            "ollama/llama3.1": 0.0
        }

        self.log_event("Class-based LiteLLM Proxy online. Load balancer initialized.", "routing")

    def log_event(self, message: str, type: str = "routing"):
        """Logs an event and pushes it to a thread-safe rolling buffer for the telemetry UI."""
        logger.info(message)
        with self.metrics_lock:
            self.routing_logs.append({
                "timestamp": time.strftime("%H:%M:%S"),
                "message": message,
                "type": type
            })
            if len(self.routing_logs) > 50:
                self.routing_logs.pop(0)

    def _get_pii_guardrail(self):
        if self.pii_guardrail is None:
            from guardrails.deberta_pii_guardrail import DeBERTaPIIGuardrail
            self.pii_guardrail = DeBERTaPIIGuardrail()
        return self.pii_guardrail

    async def _sanitize_response(self, response: Dict[str, Any], guardrailed_query: Optional[str] = None) -> Dict[str, Any]:
        if guardrailed_query is not None:
            if not isinstance(response, dict):
                try:
                    response["guardrailed_query"] = guardrailed_query
                except TypeError:
                    if hasattr(response, "dict"):
                        response = response.dict()
                    elif hasattr(response, "model_dump"):
                        response = response.model_dump()
                    else:
                        response = dict(response)
                    response["guardrailed_query"] = guardrailed_query
            else:
                response["guardrailed_query"] = guardrailed_query
        if not self.pii_enabled:
            return response
        try:
            choices = response.get("choices")
            if not choices:
                return response
            
            content = choices[0].get("message", {}).get("content", "")
            if not content or not isinstance(content, str):
                return response
                
            guardrail = self._get_pii_guardrail()
            from guardrails.deberta_pii_guardrail import _Config, SERVER_DEFAULT_POLICY, SERVER_OLLAMA_MODEL
            policy = self.pii_policy if self.pii_policy is not None else {k: self.pii_action for k in SERVER_DEFAULT_POLICY.keys()}
            cfg = _Config(
                policy=policy,
                default_action=self.pii_action,
                threshold=0.4,
                rewrite_model=SERVER_OLLAMA_MODEL,
                apply_to="both",
                labels=list(policy.keys())
            )
            
            cleaned, was_modified, block_reason = await guardrail.process(content, cfg)
            if block_reason:
                choices[0]["message"]["content"] = "[Response suppressed: model output contained blocked PII]"
                self.log_event("[PII Guardrail] Output blocked due to PII violation.", "warning")
            elif was_modified:
                choices[0]["message"]["content"] = cleaned
                self.log_event("[PII Guardrail] Post-call sanitized response output.", "warning")
        except Exception as e:
            logger.warning(f"Error sanitizing output: {e}")
        return response

    def _prune_usage_history(self, now: float):
        """Removes usage records older than 60 seconds."""
        cutoff = now - 60.0
        self.usage_history = [r for r in self.usage_history if r["timestamp"] > cutoff]

    def track_request_usage(self, physical_model: str, tokens: int):
        """
        Logs a completed request, tracking input/output tokens and timestamp in Redis or local memory.
        """
        now = time.time()
        if self.redis_client:
            try:
                # Store dynamic rate indicators inside Redis sorted sets
                key_tpm = f"litellm:tpm:{physical_model}"
                key_rpm = f"litellm:rpm:{physical_model}"
                
                # ZADD payload using sliding score window matching the timestamp
                self.redis_client.zadd(key_tpm, {f"{now}:{tokens}": now})
                self.redis_client.zadd(key_rpm, {f"{now}": now})
                
                # Expire raw telemetry indexes after 65 seconds
                self.redis_client.expire(key_tpm, 65)
                self.redis_client.expire(key_rpm, 65)
                return
            except Exception as e:
                logger.error(f"[Redis Telemetry Error] Track request failure: {e}")
                
        # In-memory fallback
        with self.metrics_lock:
            self.usage_history.append({
                "timestamp": now,
                "model": physical_model,
                "tokens": tokens
            })

    def get_endpoint_usage(self, physical_model: str) -> Tuple[int, int]:
        """Returns the current (TPM, RPM) usage in the last 60 seconds for the given physical model."""
        now = time.time()
        cutoff = now - 60.0
        
        if self.redis_client:
            try:
                key_tpm = f"litellm:tpm:{physical_model}"
                key_rpm = f"litellm:rpm:{physical_model}"
                
                # Remove timestamps older than 60 seconds to prune sliding window
                self.redis_client.zremrangebyscore(key_tpm, 0, cutoff)
                self.redis_client.zremrangebyscore(key_rpm, 0, cutoff)
                
                # Compute TPM
                tpm = 0
                members = self.redis_client.zrange(key_tpm, 0, -1)
                for member in members:
                    if ":" in member:
                        try:
                            tpm += int(member.split(":")[1])
                        except ValueError:
                            pass
                            
                # Compute RPM
                rpm = self.redis_client.zcard(key_rpm)
                return tpm, rpm
            except Exception as e:
                logger.error(f"[Redis Telemetry Error] Get endpoint usage failure: {e}")
                
        # In-memory fallback
        with self.metrics_lock:
            self._prune_usage_history(now)
            tpm = 0
            rpm = 0
            for r in self.usage_history:
                if r["model"] == physical_model:
                    tpm += r["tokens"]
                    rpm += 1
            return tpm, rpm

    def get_fallbacks_for_model(self, virtual_model: str) -> List[str]:
        """Resolves target fallback clusters from the configuration."""
        fallbacks = []
        for fb_dict in self.config.general_fallbacks:
            if isinstance(fb_dict, dict) and virtual_model in fb_dict:
                fallbacks.extend(fb_dict[virtual_model])
        return fallbacks

    def estimate_tokens(self, text: str) -> int:
        """Estimates token count of a given string using tiktoken or robust approximation."""
        if not text:
            return 0
        if self.tokenizer:
            try:
                return len(self.tokenizer.encode(text))
            except Exception:
                pass
        # Fallback: Llama/Mistral models average ~4 characters per token
        return max(1, len(text) // 4)

    def estimate_request_tokens(self, messages: List[Dict[str, str]]) -> int:
        """Estimates the total token count of incoming chat messages."""
        total = 0
        for m in messages:
            content = m.get("content", "")
            role = m.get("role", "")
            total += self.estimate_tokens(content) + self.estimate_tokens(role) + 4
        return total + 2 # overhead



    def classify_prompt_complexity(self, messages: List[Dict[str, str]], required_context: int) -> str:
        """
        Classifies prompt complexity into 'low', 'medium', or 'high' based on:
        1. Context window requirement (required_context)
        2. Semantic features (presence of reasoning/coding keywords, structural code, etc.)
        """
        # Feature 1: Context requirement
        if required_context > 8192:
            return "high"
        
        # Feature 2: Semantic check on the latest user message
        user_msgs = [m.get("content", "") for m in messages if m.get("role") == "user"]
        latest_user_msg = user_msgs[-1] if user_msgs else ""
        
        # Keywords suggesting high complexity (reasoning, coding, architecture, deep math)
        high_complexity_keywords = [
            "code", "python", "javascript", "c++", "rust", "java", "html", "css", "sql", "git",
            "algorithm", "function", "refactor", "debug", "optimize", "regex", "database",
            "proof", "theorem", "math", "calculus", "derive", "solve", "equation",
            "analyze", "evaluate", "architecture", "design pattern", "system design",
            "compare and contrast", "step by step", "reasoning", "logical deduction"
        ]
        
        # Keywords suggesting medium complexity (formatting, summarizing, translations, drafting)
        medium_complexity_keywords = [
            "summar", "summary", "report", "translation", "translate", "synopsis", "outline", "draft", 
            "rewrite", "rephrase", "format", "extract", "list", "bullet points", "email",
            "explain", "what is", "how does"
        ]
        
        msg_lower = latest_user_msg.lower()
        
        # Count high-complexity indicators (keywords or code-like patterns)
        high_count = sum(1 for kw in high_complexity_keywords if kw in msg_lower)
        # Check for code blocks (```) or braces/indentation suggesting code
        if "```" in msg_lower or (msg_lower.count("{") > 2 and msg_lower.count("}") > 2) or "def " in msg_lower or "import " in msg_lower:
            high_count += 3
            
        medium_count = sum(1 for kw in medium_complexity_keywords if kw in msg_lower)
        
        if high_count >= 2 or (high_count >= 1 and required_context > 2048):
            return "high"
        elif medium_count >= 1 or required_context > 1024 or len(latest_user_msg) > 500:
            return "medium"
        else:
            return "low"

    async def execute_chat_completion(
        self, 
        model: str, 
        messages: List[Dict[str, str]], 
        **kwargs
    ) -> Dict[str, Any]:
        """
        Executes a chat completion request, applying local PII Guardrail shielding,
        TPM/TPR evaluation limits, and automatic multi-cluster routing.
        """
        with self.metrics_lock:
            self.metrics["total_requests"] += 1

        if self.pii_enabled:
            guardrail = self._get_pii_guardrail()
            from guardrails.deberta_pii_guardrail import _Config, SERVER_DEFAULT_POLICY, SERVER_OLLAMA_MODEL
            policy = self.pii_policy if self.pii_policy is not None else {k: self.pii_action for k in SERVER_DEFAULT_POLICY.keys()}
            cfg = _Config(
                policy=policy,
                default_action=self.pii_action,
                threshold=0.4,
                rewrite_model=SERVER_OLLAMA_MODEL,
                apply_to="both",
                labels=list(policy.keys())
            )
            
            import copy
            sanitized_messages = copy.deepcopy(messages)
            
            for msg in sanitized_messages:
                content = msg.get("content")
                if isinstance(content, str):
                    cleaned, was_modified, block_reason = await guardrail.process(content, cfg)
                    if block_reason:
                        raise ValueError(block_reason)
                    if was_modified:
                        msg["content"] = cleaned
                        self.log_event(f"[PII Guardrail] Pre-call sanitized input role={msg.get('role')}.", "warning")
        else:
            sanitized_messages = messages
        
        # Extract the post-guardrail user query
        user_msgs = [m for m in sanitized_messages if m.get("role") == "user"]
        guardrailed_query = user_msgs[-1]["content"] if user_msgs else None

        estimated_prompt_tokens = self.estimate_request_tokens(sanitized_messages)
        max_tokens = kwargs.get("max_tokens", 1000)
        required_context = estimated_prompt_tokens + max_tokens
        
        # Classify complexity
        complexity = self.classify_prompt_complexity(sanitized_messages, required_context)
        
        self.log_event(
            f"[Analysis] New request on '{model}'. Size: {estimated_prompt_tokens} prompt + {max_tokens} response = {required_context} required TPR. "
            f"Prompt Complexity classified as: '{complexity.upper()}'.", 
            "routing"
        )

        # 2. Select target endpoints and apply proactive TPR & TPM/RPM load-balancing/routing
        selected_cluster = model
        selected_endpoint = None
        is_fallback_triggered = False

        if self.preference_enabled and self.preference_list:
            for p_model in self.preference_list:
                # Find matching endpoint in config
                ep = next((e for e in self.config.endpoints if e.model == p_model), None)
                if not ep:
                    continue
                    
                # Check credit limit
                current_spend = self.accumulated_spend.get(p_model, 0.0)
                limit = self.credit_limits.get(p_model, 9999.0)
                if current_spend >= limit:
                    self.log_event(f"[Preference Route] Model '{p_model}' skipped: credit limit reached (spent ${current_spend:.6f} >= limit ${limit:.6f}).", "warning")
                    continue
                    
                # Check TPR
                if required_context > ep.tpr:
                    self.log_event(f"[Preference Route] Model '{p_model}' skipped: required context {required_context} exceeds limit {ep.tpr}.", "warning")
                    continue
                    
                # Check TPM / RPM
                current_tpm, current_rpm = self.get_endpoint_usage(ep.model)
                if current_tpm + required_context > ep.tpm:
                    self.log_event(f"[Preference Route] Model '{p_model}' skipped: current TPM {current_tpm} + required {required_context} exceeds limit {ep.tpm}.", "warning")
                    continue
                if current_rpm + 1 > ep.rpm:
                    self.log_event(f"[Preference Route] Model '{p_model}' skipped: current RPM {current_rpm} + 1 exceeds limit {ep.rpm}.", "warning")
                    continue
                    
                selected_endpoint = ep
                selected_cluster = ep.model_name
                self.log_event(f"[Preference Selection] Node '{p_model}' selected as highest priority active preferred LLM (spent: ${current_spend:.6f}/${limit:.6f}).", "success")
                break
                
            if not selected_endpoint:
                self.log_event("[Preference Route Exhaustion] All preferred models exceeded credit limits or TPR constraints! Falling back to the first preference to prevent hard outage.", "error")
                p_model = self.preference_list[0]
                ep = next((e for e in self.config.endpoints if e.model == p_model), None)
                if ep:
                    selected_endpoint = ep
                    selected_cluster = ep.model_name

        if not selected_endpoint:
            search_clusters = [model]
            search_clusters.extend(self.get_fallbacks_for_model(model))
            
            for cluster in search_clusters:
                endpoints = self.config.get_endpoints_for_model(cluster)
                if not endpoints:
                    continue
                    
                suitable_endpoints = []
                for ep in endpoints:
                    # Check TPR
                    if required_context > ep.tpr:
                        self.log_event(f"[TPR Limit] Node '{ep.model}' rejected: required context {required_context} exceeds limit {ep.tpr}.", "warning")
                        continue
                        
                    # Check TPM & RPM
                    current_tpm, current_rpm = self.get_endpoint_usage(ep.model)
                    if current_tpm + required_context > ep.tpm:
                        self.log_event(f"[TPM Limit] Node '{ep.model}' rejected: current TPM {current_tpm} + required {required_context} exceeds limit {ep.tpm}.", "warning")
                        continue
                    if current_rpm + 1 > ep.rpm:
                        self.log_event(f"[RPM Limit] Node '{ep.model}' rejected: current RPM {current_rpm} + 1 exceeds limit {ep.rpm}.", "warning")
                        continue
                        
                    suitable_endpoints.append((ep, current_tpm, current_rpm))
                if suitable_endpoints:
                    # Complexity-Aware and Cost-Aware Multi-Objective Load Balancing:
                    # 1. Primary Objective: Minimize tier mismatch penalty (align with classified complexity)
                    # 2. Secondary Objective: Minimize cost_per_million (least credit usage)
                    # 3. Tertiary Objective: Minimize utilization (balanced resource usage)
                    def get_suitability_score(item):
                        ep, t_used, r_used = item
                        
                        tier_map = {"low": 1, "medium": 2, "high": 3}
                        p_tier = tier_map.get(complexity, 2)
                        ep_tier = tier_map.get(ep.complexity_tier, 2)
                        
                        tier_mismatch = abs(p_tier - ep_tier)
                        
                        # Strong penalty if high complexity prompt is sent to a low reasoning node
                        if complexity == "high" and ep.complexity_tier == "low":
                            tier_mismatch += 5.0
                        # Strong penalty if low complexity prompt is sent to an expensive high-tier node
                        if complexity == "low" and ep.complexity_tier == "high":
                            tier_mismatch += 5.0
                            
                        util = max(t_used / ep.tpm, r_used / ep.rpm)
                        return (tier_mismatch, ep.cost_per_million, util)
                        
                    selected_endpoint, t_used, r_used = min(suitable_endpoints, key=get_suitability_score)
                    selected_cluster = cluster
                    if cluster != model:
                        is_fallback_triggered = True
                        with self.metrics_lock:
                            self.metrics["fallback_events"] += 1
                        self.log_event(f"[Fallback Route] Overload/limit on '{model}'. Cascading to cluster '{cluster}'.", "warning")
                    
                    self.log_event(
                        f"[Complexity-Aware Selection] Node '{selected_endpoint.model}' selected in cluster '{cluster}' "
                        f"(Complexity Tier: {selected_endpoint.complexity_tier.upper()}, Cost: ${selected_endpoint.cost_per_million}/M tokens, TPM usage: {t_used}/{selected_endpoint.tpm}).", 
                        "success"
                    )
                    break
                
        # Relax TPR constraint if prompt is extremely large and exceeds all limits
        if not selected_endpoint:
            all_eps = []
            for cluster in search_clusters:
                all_eps.extend(self.config.get_endpoints_for_model(cluster))
            if all_eps:
                max_tpr = max(e.tpr for e in all_eps)
                best_eps = [e for e in all_eps if e.tpr == max_tpr]
                backup_eps = [e for e in best_eps if e.model_name == "backup-cluster"]
                selected_endpoint = backup_eps[0] if backup_eps else best_eps[0]
                selected_cluster = selected_endpoint.model_name
                if selected_cluster != model:
                    is_fallback_triggered = True
                    with self.metrics_lock:
                        self.metrics["fallback_events"] += 1
                self.log_event(f"[TPR Overlimit] Required {required_context} exceeds all limits. Escalating to highest capacity node: '{selected_endpoint.model}'.", "warning")
            else:
                raise ValueError(f"No active endpoints configured for model routing group '{model}'.")

        # Determine if we execute in Mock Sandbox Mode
        is_mock = (
            kwargs.get("mock_sandbox", False) or
            (selected_endpoint and selected_endpoint.api_key and "mock" in selected_endpoint.api_key.lower())
        )
        
        if is_mock:
            response = self._execute_mock_sandbox_completion(
                virtual_model=selected_cluster,
                endpoint=selected_endpoint,
                estimated_prompt_tokens=estimated_prompt_tokens,
                max_tokens=max_tokens,
                messages=sanitized_messages
            )
            response["prompt_complexity"] = complexity
            response = await self._sanitize_response(response, guardrailed_query)
            return response

        # Real API Execution via LiteLLM Router - directly let LITELLM route the model!
        try:
            self.log_event(f"[API Dispatch] Sending request to backend '{selected_endpoint.model}' in cluster '{selected_cluster}'...", "routing")
            start_time = time.time()
            
            # Clean kwargs
            litellm_kwargs = kwargs.copy()
            litmm_kwargs_to_pop = ["mock_sandbox", "fallbacks"]
            for k in litmm_kwargs_to_pop:
                litellm_kwargs.pop(k, None)

            # We pass the selected target cluster to the LiteLLM Router so it uses its config
            # and automatically selects the active backend under that cluster!
            response = self.router.completion(
                model=selected_cluster,
                messages=sanitized_messages,
                **litellm_kwargs
            )
            
            latency = time.time() - start_time
            actual_routed_model = response.get("model", selected_endpoint.model)
            self.log_event(f"[Success] API call completed by '{actual_routed_model}' in {latency:.2f}s.", "success")
            
            # Update metrics and usage history
            usage = response.get("usage", {})
            input_tokens = usage.get("prompt_tokens", estimated_prompt_tokens)
            output_tokens = usage.get("completion_tokens", 0)
            total_tokens = input_tokens + output_tokens
            
            self.track_request_usage(actual_routed_model, total_tokens)
            
            self._update_success_metrics(actual_routed_model, input_tokens, output_tokens)
            try:
                response["prompt_complexity"] = complexity
            except Exception:
                pass
            
            response = await self._sanitize_response(response, guardrailed_query)
            return response
            
        except Exception as e:
            self.log_event(f"[API Error] Backend '{selected_endpoint.model}' failed: {e}. Trying fallback execution chain...", "error")
            
            # Try immediate physical model failover!
            # 1. First try other endpoints inside the CURRENT active cluster (Intra-Cluster Failover)
            current_cluster_eps = self.config.get_endpoints_for_model(selected_cluster)
            for alt_ep in current_cluster_eps:
                if alt_ep.model == selected_endpoint.model:
                    continue
                try:
                    self.log_event(f"[Intra-Cluster Failover] Cascading to alternate backend '{alt_ep.model}' in current cluster '{selected_cluster}'...", "warning")
                    start_time = time.time()
                    response = self.router.completion(
                        model=selected_cluster,
                        messages=sanitized_messages,
                        **litellm_kwargs
                    )
                    latency = time.time() - start_time
                    actual_routed_model = response.get("model", alt_ep.model)
                    
                    self.log_event(f"[Intra-Cluster Success] Succeeded on '{actual_routed_model}' in {latency:.2f}s.", "success")
                    
                    usage = response.get("usage", {})
                    input_tokens = usage.get("prompt_tokens", estimated_prompt_tokens)
                    output_tokens = usage.get("completion_tokens", 0)
                    total_tokens = input_tokens + output_tokens
                    
                    self.track_request_usage(actual_routed_model, total_tokens)
                    with self.metrics_lock:
                        self.metrics["fallback_events"] += 1
                        
                    self._update_success_metrics(actual_routed_model, input_tokens, output_tokens)
                    try:
                        response["prompt_complexity"] = complexity
                    except Exception:
                        pass
                    response = await self._sanitize_response(response, guardrailed_query)
                    return response
                except Exception as intra_ex:
                    self.log_event(f"[Intra-Cluster Failover Error] Alternate backend '{alt_ep.model}' also failed: {intra_ex}.", "error")

            # 2. Loop through other clusters
            for fallback_cluster in search_clusters:
                if fallback_cluster == selected_cluster:
                    continue
                fallback_eps = self.config.get_endpoints_for_model(fallback_cluster)
                if not fallback_eps:
                    continue
                alt_ep = fallback_eps[0]
                
                try:
                    self.log_event(f"[Failover API Dispatch] Cascading to alternate backend '{alt_ep.model}' in '{fallback_cluster}'...", "warning")
                    start_time = time.time()
                    
                    response = self.router.completion(
                        model=fallback_cluster,
                        messages=sanitized_messages,
                        **litellm_kwargs
                    )
                    latency = time.time() - start_time
                    actual_routed_model = response.get("model", alt_ep.model)
                    
                    self.log_event(f"[Failover Success] Alternate backend '{actual_routed_model}' succeeded in {latency:.2f}s.", "success")
                    
                    usage = response.get("usage", {})
                    input_tokens = usage.get("prompt_tokens", estimated_prompt_tokens)
                    output_tokens = usage.get("completion_tokens", 0)
                    total_tokens = input_tokens + output_tokens
                    
                    self.track_request_usage(actual_routed_model, total_tokens)
                    with self.metrics_lock:
                        self.metrics["fallback_events"] += 1
                        
                    self._update_success_metrics(actual_routed_model, input_tokens, output_tokens)
                    try:
                        response["prompt_complexity"] = complexity
                    except Exception:
                        pass
                    
                    response = await self._sanitize_response(response, guardrailed_query)
                    return response
                    
                except Exception as ex:
                    self.log_event(f"[Failover API Error] Alternate backend '{alt_ep.model}' also failed: {ex}.", "error")
            
            with self.metrics_lock:
                self.metrics["failed_requests"] += 1
            raise RuntimeError(f"All backends in the routing chain failed. Last error: {e}")

    def _update_success_metrics(self, model: str, input_tokens: int, output_tokens: int):
        """Updates metrics safely across threads."""
        cost = 0.0
        total_spent = 0.0
        with self.metrics_lock:
            self.metrics["successful_requests"] += 1
            self.metrics["total_input_tokens"] += input_tokens
            self.metrics["total_output_tokens"] += output_tokens
            self.metrics["provider_calls"][model] = self.metrics["provider_calls"].get(model, 0) + 1
            
            # Find the cost per million for this physical model
            cost_per_mil = 0.1
            for ep in self.config.endpoints:
                if ep.model == model:
                    cost_per_mil = ep.cost_per_million
                    break
            
            total_tokens = input_tokens + output_tokens
            cost = (total_tokens * cost_per_mil) / 1_000_000.0
            
            # Update accumulated spend
            self.accumulated_spend[model] = self.accumulated_spend.get(model, 0.0) + cost
            total_spent = self.accumulated_spend[model]
            
        self.log_event(f"[Cost Auditing] Spend on '{model}' increased by ${cost:.6f}. Total spent: ${total_spent:.6f}.", "success")

    def _execute_mock_sandbox_completion(
        self,
        virtual_model: str,
        endpoint: ModelEndpointConfig,
        estimated_prompt_tokens: int,
        max_tokens: int,
        messages: List[Dict[str, str]]
    ) -> Dict[str, Any]:
        """Simulates a model completion response instantly for local development and validation."""
        logger.info(f"[MOCK SANDBOX] Simulating completion for '{endpoint.model}'")
        time.sleep(0.1) # Simulate minimal network latency
        
        last_message = messages[-1].get("content", "") if messages else "Hello"
        mock_reply = (
            f"[LiteLLM Proxy Mock - {endpoint.model}]\n"
            f"Routing Group: {virtual_model}\n"
            f"Optimized TPR Context: {endpoint.tpr} max tokens.\n"
            f"Acknowledged request: '{last_message[:200]}...'"
        )
        
        output_tokens = self.estimate_tokens(mock_reply)
        total_tokens = estimated_prompt_tokens + output_tokens
        
        self.track_request_usage(endpoint.model, total_tokens)
            
        self._update_success_metrics(endpoint.model, estimated_prompt_tokens, output_tokens)
        self.log_event(f"[Mock Success] Simulated reply from '{endpoint.model}' in cluster '{virtual_model}' (Tokens: {total_tokens}).", "success")
        
        # Structure the dict like an OpenAI / LiteLLM chat.completion response
        return {
            "id": f"chatcmpl-mock-{int(time.time())}",
            "object": "chat.completion",
            "created": int(time.time()),
            "model": endpoint.model,
            "choices": [
                {
                    "index": 0,
                    "message": {
                        "role": "assistant",
                        "content": mock_reply
                    },
                    "finish_reason": "stop"
                }
            ],
            "usage": {
                "prompt_tokens": estimated_prompt_tokens,
                "completion_tokens": output_tokens,
                "total_tokens": estimated_prompt_tokens + output_tokens
            }
        }

    def get_metrics(self) -> Dict[str, Any]:
        """Thread-safe getter for router metrics."""
        with self.metrics_lock:
            m = self.metrics.copy()
            m["logs"] = list(self.routing_logs)
            return m