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