The update introduces a content-aware routing system designed to classify prompts and route them to the most suitable model based on specific preferences before considering infrastructure metrics like cost or latency.

Key Components:
Backend: Features three classifier backends (rule-based TF-IDF, semantic embedding similarity, and external HTTP) and a model affinity router. It includes new configuration types and a dry-run endpoint for testing classification.
UI: Adds a new settings panel for configuring classifiers, default models, and confidence thresholds, along with an inline prompt testing tool.
Testing: Comprehensive unit tests cover the new routing logic and 26 frontend tests verify UI functionality, including API paths and security guards.
This commit is contained in:
Chetan Soni 2026-04-14 14:15:31 -07:00
parent 97f722f558
commit 8ccecaf9bd
19 changed files with 3188 additions and 1 deletions

View file

@ -561,6 +561,30 @@ class ProxyBaseLLMRequestProcessing:
if logging_caching_headers:
headers.update(logging_caching_headers)
# Content-aware routing decision header
content_routing_decision = (
request_data.get("metadata", {}) or {}
).get("content_routing_decision")
if content_routing_decision:
pref = content_routing_decision.get("matched_preference", "")
model = content_routing_decision.get("model", "")
if pref and model:
headers["x-litellm-content-route"] = f"{pref} -> {model}"
# Model affinity (session pinning) header
model_affinity_decision = (
request_data.get("metadata", {}) or {}
).get("model_affinity_decision")
if model_affinity_decision:
status = model_affinity_decision.get("status", "")
affinity_model = model_affinity_decision.get("model", "")
session_id = model_affinity_decision.get("session_id", "")
if status and affinity_model:
headers["x-litellm-model-affinity-status"] = status
headers["x-litellm-model-affinity-model"] = affinity_model
if session_id:
headers["x-litellm-model-affinity-session"] = session_id
try:
return {
key: str(value)

View file

@ -0,0 +1,193 @@
"""
CONTENT-AWARE ROUTING ENDPOINTS
POST /utils/content_route_test - Dry-run: classify a prompt without making an LLM call
GET /router/content_routing/preferences - List all models with their routing_preferences
"""
from typing import Any, Dict, List, Optional
from fastapi import APIRouter, Depends
from pydantic import BaseModel
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
router = APIRouter()
class ContentRouteTestRequest(BaseModel):
prompt: Optional[str] = None
messages: Optional[List[Dict[str, Any]]] = None
class ContentRouteTestResponse(BaseModel):
matched_preference: str
matched_model: str
confidence: float
classifier: str
all_scores: Optional[Dict[str, float]] = None
class ModelRoutingPreferences(BaseModel):
model_name: str
routing_preferences: List[Dict[str, str]]
class ContentRoutingPreferencesResponse(BaseModel):
models: List[ModelRoutingPreferences]
content_routing_enabled: bool
classifier: Optional[str] = None
default_model: Optional[str] = None
@router.post(
"/utils/content_route_test",
tags=["router"],
dependencies=[Depends(user_api_key_auth)],
response_model=ContentRouteTestResponse,
summary="Test content-aware routing for a prompt without making an LLM call",
)
async def content_route_test(
request: ContentRouteTestRequest,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
) -> ContentRouteTestResponse:
"""
Classify a prompt against the configured routing_preferences and return the
routing decision without actually calling any LLM.
Useful for validating routing_preferences config and debugging routing decisions.
"""
from litellm.proxy.proxy_server import llm_router
from litellm.router_strategy.content_aware_router.utils import (
build_tfidf_vectors,
extract_prompt_text,
tfidf_score,
tokenize,
)
if llm_router is None:
from fastapi import HTTPException
raise HTTPException(status_code=503, detail="Router not initialized")
content_aware_router = llm_router.content_aware_router
if content_aware_router is None:
from fastapi import HTTPException
raise HTTPException(
status_code=400,
detail="Content-aware routing is not enabled. Set router_settings.content_routing.enabled=true in your config.",
)
# Build messages from prompt if needed
messages: Optional[List[Dict[str, Any]]] = request.messages
if not messages and request.prompt:
messages = [{"role": "user", "content": request.prompt}]
if not messages:
from fastapi import HTTPException
raise HTTPException(
status_code=400, detail="Either 'prompt' or 'messages' must be provided"
)
user_text, system_text = extract_prompt_text(messages)
if not user_text:
from fastapi import HTTPException
raise HTTPException(status_code=400, detail="No user message content found")
classifier = content_aware_router.config.classifier
# Run classification and collect all scores for transparency
all_scores: Dict[str, float] = {}
if classifier == "rule_based":
prompt_tokens = tokenize(f"{system_text or ''} {user_text}".strip())
for i, (model_name, pref) in enumerate(content_aware_router._index):
score = tfidf_score(
prompt_tokens,
content_aware_router._idf_weights,
content_aware_router._tfidf_vectors[i],
)
all_scores[f"{model_name}/{pref.name}"] = round(score, 4)
matched_model, matched_pref, confidence = (
content_aware_router._classify_rule_based(user_text, system_text)
)
elif classifier == "embedding_similarity":
matched_model, matched_pref, confidence = (
await content_aware_router._classify_embedding_similarity(
user_text, system_text
)
)
else: # external_model
matched_model, matched_pref, confidence = (
await content_aware_router._classify_external_model(user_text, system_text)
)
verbose_proxy_logger.info(
f"content_route_test: classifier={classifier} matched={matched_pref} "
f"model={matched_model} confidence={confidence:.4f}"
)
return ContentRouteTestResponse(
matched_preference=matched_pref,
matched_model=matched_model,
confidence=round(confidence, 4),
classifier=classifier,
all_scores=all_scores if all_scores else None,
)
@router.get(
"/router/content_routing/preferences",
tags=["router"],
dependencies=[Depends(user_api_key_auth)],
response_model=ContentRoutingPreferencesResponse,
summary="List all models with their configured routing_preferences",
)
async def get_content_routing_preferences(
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
) -> ContentRoutingPreferencesResponse:
"""
Returns all models in the router that have routing_preferences configured,
along with the current content routing settings.
"""
from litellm.proxy.proxy_server import llm_router
if llm_router is None:
from fastapi import HTTPException
raise HTTPException(status_code=503, detail="Router not initialized")
models: List[ModelRoutingPreferences] = []
for deployment in llm_router.model_list or []:
if isinstance(deployment, dict):
prefs = deployment.get("routing_preferences")
model_name = deployment.get("model_name")
else:
prefs = getattr(deployment, "routing_preferences", None)
model_name = getattr(deployment, "model_name", None)
if prefs and model_name:
if isinstance(prefs, list):
pref_dicts = [
p if isinstance(p, dict) else p.model_dump() for p in prefs
]
else:
pref_dicts = []
models.append(
ModelRoutingPreferences(
model_name=model_name,
routing_preferences=pref_dicts,
)
)
config = llm_router._content_routing_config
return ContentRoutingPreferencesResponse(
models=models,
content_routing_enabled=config.enabled if config else False,
classifier=config.classifier if config else None,
default_model=config.default_model if config else None,
)

View file

@ -410,6 +410,9 @@ from litellm.proxy.management_endpoints.policy_endpoints import router as policy
from litellm.proxy.management_endpoints.project_endpoints import (
router as project_router,
)
from litellm.proxy.management_endpoints.content_routing_endpoints import (
router as content_routing_router,
)
from litellm.proxy.management_endpoints.router_settings_endpoints import (
router as router_settings_router,
)
@ -13938,6 +13941,7 @@ app.include_router(tag_management_router)
app.include_router(tool_management_router)
app.include_router(cost_tracking_settings_router)
app.include_router(router_settings_router)
app.include_router(content_routing_router)
app.include_router(fallback_management_router)
app.include_router(cache_settings_router)
app.include_router(config_override_router)

View file

@ -200,12 +200,23 @@ if TYPE_CHECKING:
from litellm.router_strategy.complexity_router.complexity_router import (
ComplexityRouter,
)
from litellm.router_strategy.content_aware_router.content_aware_router import (
ContentAwareRouter,
)
from litellm.router_strategy.model_affinity_router.model_affinity_router import (
ModelAffinityRouter,
)
from litellm.types.router import ContentRoutingConfig, ModelAffinityConfig
Span = Union[_Span, Any]
else:
Span = Any
AutoRouter = Any
ComplexityRouter = Any
ContentAwareRouter = Any
ModelAffinityRouter = Any
ContentRoutingConfig = Any
ModelAffinityConfig = Any
PreRoutingHookResponse = Any
@ -311,6 +322,8 @@ class Router:
enable_health_check_routing: bool = False,
health_check_staleness_threshold: Optional[int] = None,
health_check_ignore_transient_errors: bool = False,
content_routing: Optional[Union[Dict, "ContentRoutingConfig"]] = None,
model_affinity: Optional[Union[Dict, "ModelAffinityConfig"]] = None,
) -> None:
"""
Initialize the Router class with the given parameters for caching, reliability, and routing strategy.
@ -464,6 +477,9 @@ class Router:
) # {"TEAM_ID": PatternMatchRouter}
self.auto_routers: Dict[str, "AutoRouter"] = {}
self.complexity_routers: Dict[str, "ComplexityRouter"] = {}
self.content_aware_router: Optional["ContentAwareRouter"] = None
self._content_routing_config: Optional["ContentRoutingConfig"] = None
self.model_affinity_router: Optional["ModelAffinityRouter"] = None
# Initialize model_group_alias early since it's used in set_model_list
self.model_group_alias: Dict[str, Union[str, RouterModelGroupAliasItem]] = (
@ -706,6 +722,14 @@ class Router:
if self.alerting_config is not None:
self._initialize_alerting()
### CONTENT-AWARE ROUTING SETUP ###
if content_routing is not None:
self._init_content_aware_router(content_routing)
### MODEL AFFINITY (SESSION PINNING) SETUP ###
if model_affinity is not None:
self._init_model_affinity_router(model_affinity)
self.initialize_assistants_endpoint()
self.initialize_router_endpoints()
self.apply_default_settings()
@ -6914,6 +6938,96 @@ class Router:
)
self.complexity_routers[deployment.model_name] = complexity_router
def _init_content_aware_router(
self,
content_routing: Union[Dict, "ContentRoutingConfig"],
) -> None:
"""
Build a ContentAwareRouter from all deployments that have routing_preferences.
Called once at Router init (if content_routing config is present) and again
whenever deployments change (_refresh_content_aware_router).
"""
from litellm.router_strategy.content_aware_router.content_aware_router import (
ContentAwareRouter,
)
from litellm.types.router import ContentRoutingConfig
if isinstance(content_routing, dict):
config = ContentRoutingConfig(**content_routing)
else:
config = content_routing
if not config.enabled:
self.content_aware_router = None
self._content_routing_config = config
return
self._content_routing_config = config
# Collect routing_preferences from all deployed models.
# self.model_list stores plain dicts (deployment.to_json()), so
# routing_preferences entries are plain dicts and must be coerced.
from litellm.types.router import RoutingPreference
preferences_by_model: Dict[str, list] = {}
for deployment in self.model_list:
if isinstance(deployment, dict):
prefs_raw = deployment.get("routing_preferences")
model_name = deployment.get("model_name")
else:
prefs_raw = getattr(deployment, "routing_preferences", None)
model_name = getattr(deployment, "model_name", None)
if prefs_raw and model_name:
coerced = [
RoutingPreference(**p) if isinstance(p, dict) else p
for p in prefs_raw
]
preferences_by_model[model_name] = coerced
if not preferences_by_model:
verbose_router_logger.warning(
"ContentAwareRouter: content_routing.enabled=true but no deployments "
"have routing_preferences configured. Content routing will be skipped."
)
self.content_aware_router = ContentAwareRouter(
preferences_by_model=preferences_by_model,
config=config,
litellm_router_instance=self,
)
def _refresh_content_aware_router(self) -> None:
"""Rebuild the ContentAwareRouter index after deployments change."""
if self._content_routing_config is not None:
self._init_content_aware_router(self._content_routing_config)
def _init_model_affinity_router(
self,
model_affinity: Union[Dict, "ModelAffinityConfig"],
) -> None:
"""
Instantiate a ModelAffinityRouter from the model_affinity config.
Called once at Router init and can be re-called to update config.
"""
from litellm.router_strategy.model_affinity_router.model_affinity_router import (
ModelAffinityRouter,
)
from litellm.types.router import ModelAffinityConfig
if isinstance(model_affinity, dict):
config = ModelAffinityConfig(**model_affinity)
else:
config = model_affinity
if not config.enabled:
self.model_affinity_router = None
return
self.model_affinity_router = ModelAffinityRouter(config=config)
def deployment_is_active_for_environment(self, deployment: Deployment) -> bool:
"""
Function to check if a llm deployment is active for a given environment. Allows using the same config.yaml across multople environments
@ -7134,6 +7248,11 @@ class Router:
):
self.init_complexity_router_deployment(deployment=deployment)
#########################################################
# Refresh content-aware router index if needed
#########################################################
self._refresh_content_aware_router()
return deployment
def _initialize_deployment_for_pass_through(
@ -9621,6 +9740,61 @@ class Router:
Used for the litellm auto-router to modify the request before the routing decision is made.
"""
#########################################################
# Model Affinity — session pinning (runs before content routing)
#########################################################
from litellm.types.router import (
PreRoutingHookResponse as _PreRoutingHookResponse,
)
session_id: Optional[str] = None
if self.model_affinity_router is not None and not specific_deployment:
headers = (request_kwargs.get("metadata") or {}).get("headers") or {}
session_id = headers.get("x-model-affinity") or headers.get(
"X-Model-Affinity"
)
if session_id:
pinned_model = await self.model_affinity_router.get_pinned_model(
session_id
)
if pinned_model:
metadata = request_kwargs.setdefault("metadata", {})
metadata["model_affinity_decision"] = {
"session_id": session_id,
"model": pinned_model,
"status": "pinned",
}
return _PreRoutingHookResponse(
model=pinned_model, messages=messages
)
#########################################################
# Check if content-aware routing is enabled (global, runs second)
#########################################################
content_result: Optional[PreRoutingHookResponse] = None
if self.content_aware_router is not None and not specific_deployment:
content_result = await self.content_aware_router.async_pre_routing_hook(
model=model,
request_kwargs=request_kwargs,
messages=messages,
input=input,
specific_deployment=specific_deployment,
)
# Pin the selected model for new sessions before returning
if session_id and self.model_affinity_router is not None:
selected_model = content_result.model if content_result else model
await self.model_affinity_router.pin_model(session_id, selected_model)
metadata = request_kwargs.setdefault("metadata", {})
metadata["model_affinity_decision"] = {
"session_id": session_id,
"model": selected_model,
"status": "new",
}
if content_result is not None:
return content_result
#########################################################
# Check if any auto-router should be used
#########################################################

View file

@ -0,0 +1,324 @@
"""
Content-Aware Preference-Aligned Router
Routes incoming requests to the best-matching model by classifying the prompt
content against per-model routing_preference descriptions configured in the
LiteLLM YAML config.
Supports three classifiers:
- rule_based: TF-IDF cosine similarity, zero latency, no external deps
- embedding_similarity: uses litellm.aembedding() for semantic matching
- external_model: delegates to an external HTTP classification endpoint
"""
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
import httpx
from litellm._logging import verbose_router_logger
from litellm.integrations.custom_logger import CustomLogger
from .utils import (
bm25_score,
build_bm25_index,
cosine_similarity,
extract_prompt_text,
tokenize,
)
if TYPE_CHECKING:
from litellm.router import Router
from litellm.types.router import (
ContentRoutingConfig,
PreRoutingHookResponse,
RoutingPreference,
)
else:
Router = Any
ContentRoutingConfig = Any
PreRoutingHookResponse = Any
RoutingPreference = Any
class ContentAwareRouter(CustomLogger):
"""
Pre-routing hook that classifies prompt content and selects the best-matched
model based on routing_preferences declared in the model_list config.
Instantiated by the Router when router_settings.content_routing.enabled = true.
"""
def __init__(
self,
preferences_by_model: Dict[str, List["RoutingPreference"]],
config: "ContentRoutingConfig",
litellm_router_instance: "Router",
):
"""
Args:
preferences_by_model: {model_name: [RoutingPreference, ...]} for all
deployments that have routing_preferences set.
config: ContentRoutingConfig parsed from router_settings.content_routing.
litellm_router_instance: The Router instance (used for embedding calls).
"""
self.preferences_by_model = preferences_by_model
self.config = config
self.litellm_router_instance = litellm_router_instance
# Flat ordered list of (model_name, preference) for index alignment
self._index: List[Tuple[str, "RoutingPreference"]] = [
(model, pref)
for model, prefs in preferences_by_model.items()
for pref in prefs
]
# Rule-based: BM25 index
self._bm25_corpus: List[List[str]] = [] # stemmed tokens per preference
self._bm25_idf: Dict[str, float] = {}
self._bm25_avgdl: float = 0.0
# Embedding-similarity: pre-computed description embeddings
self._description_embeddings: List[List[float]] = []
if self._index:
# Always build the rule_based index — it is cheap and serves as a
# fallback when embedding or external classifiers fail at runtime.
self._build_rule_based_index()
# Embedding vectors are built lazily on first request to allow async init
verbose_router_logger.info(
f"ContentAwareRouter initialized with {len(self._index)} preferences "
f"across {len(preferences_by_model)} models, classifier={config.classifier}"
)
# ------------------------------------------------------------------
# Index builders
# ------------------------------------------------------------------
def _build_rule_based_index(self) -> None:
"""Pre-compute BM25 index for all preference descriptions."""
descriptions = [pref.description for _, pref in self._index]
self._bm25_corpus, self._bm25_idf, self._bm25_avgdl = build_bm25_index(
descriptions
)
async def _ensure_embedding_index(self) -> None:
"""Build embedding vectors for all preference descriptions (once)."""
if self._description_embeddings:
return # already built
import litellm
embedding_model = self.config.embedding_model or "text-embedding-3-small"
descriptions = [pref.description for _, pref in self._index]
try:
response = await litellm.aembedding(
model=embedding_model,
input=descriptions,
)
self._description_embeddings = [
item["embedding"] for item in response.data
]
verbose_router_logger.debug(
f"ContentAwareRouter: built {len(self._description_embeddings)} "
f"description embeddings with {embedding_model}"
)
except Exception as e:
verbose_router_logger.warning(
f"ContentAwareRouter: failed to build embedding index: {e}. "
"Falling back to rule_based classifier."
)
self._build_rule_based_index()
# ------------------------------------------------------------------
# Classifiers — all return (model_name, preference_name, confidence)
# ------------------------------------------------------------------
def _classify_rule_based(
self, user_text: str, system_text: Optional[str]
) -> Tuple[str, str, float]:
"""BM25 scoring against each preference description."""
# Use full text for scoring (system prompt provides deployment context)
combined = f"{system_text or ''} {user_text}".strip()
prompt_tokens = tokenize(combined)
best_score = -1.0
best_model = self.config.default_model or ""
best_pref = ""
for i, (model_name, pref) in enumerate(self._index):
score = bm25_score(
prompt_tokens, self._bm25_corpus[i], self._bm25_idf, self._bm25_avgdl
)
verbose_router_logger.debug(
f"ContentAwareRouter rule_based: model={model_name} "
f"pref={pref.name} score={score:.4f}"
)
if score > best_score:
best_score = score
best_model = model_name
best_pref = pref.name
return best_model, best_pref, best_score
async def _classify_embedding_similarity(
self, user_text: str, system_text: Optional[str]
) -> Tuple[str, str, float]:
"""Embed the prompt and find the nearest preference description."""
await self._ensure_embedding_index()
# If embedding index failed, fall back to rule-based
if not self._description_embeddings:
return self._classify_rule_based(user_text, system_text)
import litellm
embedding_model = self.config.embedding_model or "text-embedding-3-small"
try:
response = await litellm.aembedding(
model=embedding_model,
input=[user_text],
)
prompt_embedding: List[float] = response.data[0]["embedding"]
except Exception as e:
verbose_router_logger.warning(
f"ContentAwareRouter: embedding call failed ({e}), "
"falling back to rule_based"
)
return self._classify_rule_based(user_text, system_text)
best_score = -1.0
best_model = self.config.default_model or ""
best_pref = ""
for i, (model_name, pref) in enumerate(self._index):
score = cosine_similarity(prompt_embedding, self._description_embeddings[i])
verbose_router_logger.debug(
f"ContentAwareRouter embedding: model={model_name} "
f"pref={pref.name} score={score:.4f}"
)
if score > best_score:
best_score = score
best_model = model_name
best_pref = pref.name
return best_model, best_pref, best_score
async def _classify_external_model(
self, user_text: str, system_text: Optional[str]
) -> Tuple[str, str, float]:
"""
POST prompt to an external classifier endpoint.
Expected response JSON:
{"matched_preference": "code_generation", "model": "claude-sonnet", "confidence": 0.92}
"""
url = self.config.external_classifier_url
if not url:
verbose_router_logger.warning(
"ContentAwareRouter: external_classifier_url not set, "
"falling back to rule_based"
)
return self._classify_rule_based(user_text, system_text)
payload = {
"prompt": user_text,
"system_prompt": system_text,
"preferences": [
{"model": m, "name": p.name, "description": p.description}
for m, p in self._index
],
}
try:
async with httpx.AsyncClient(timeout=5.0) as client:
resp = await client.post(url, json=payload)
resp.raise_for_status()
data = resp.json()
model_name = data.get("model", self.config.default_model or "")
pref_name = data.get("matched_preference", "")
confidence = float(data.get("confidence", 0.0))
return model_name, pref_name, confidence
except Exception as e:
verbose_router_logger.warning(
f"ContentAwareRouter: external classifier call failed ({e}), "
"falling back to rule_based"
)
return self._classify_rule_based(user_text, system_text)
# ------------------------------------------------------------------
# Public pre-routing hook
# ------------------------------------------------------------------
async def async_pre_routing_hook(
self,
model: str,
request_kwargs: Dict,
messages: Optional[List[Dict[str, Any]]] = None,
input: Optional[Union[str, List]] = None,
specific_deployment: Optional[bool] = False,
) -> Optional["PreRoutingHookResponse"]:
"""
Called by Router.async_pre_routing_hook() before infrastructure routing.
Classifies the prompt content and returns the best-matched model.
Returns None when content routing should be skipped (no preferences, etc.).
"""
from litellm.types.router import PreRoutingHookResponse
if not self._index:
return None
user_text, system_text = extract_prompt_text(messages)
if not user_text:
verbose_router_logger.debug(
"ContentAwareRouter: no user message found, skipping"
)
return None
classifier = self.config.classifier
if classifier == "rule_based":
matched_model, matched_pref, confidence = self._classify_rule_based(
user_text, system_text
)
elif classifier == "embedding_similarity":
matched_model, matched_pref, confidence = (
await self._classify_embedding_similarity(user_text, system_text)
)
else: # external_model
matched_model, matched_pref, confidence = (
await self._classify_external_model(user_text, system_text)
)
threshold = self.config.confidence_threshold
if confidence < threshold:
fallback = self.config.default_model
verbose_router_logger.info(
f"ContentAwareRouter: confidence {confidence:.4f} below threshold "
f"{threshold}, using default_model={fallback}"
)
if not fallback:
return None
matched_model = fallback
matched_pref = "default"
verbose_router_logger.info(
f"ContentAwareRouter: classifier={classifier} "
f"matched_preference={matched_pref} model={matched_model} "
f"confidence={confidence:.4f}"
)
# Store decision in metadata for response header propagation
metadata = request_kwargs.setdefault("metadata", {})
metadata["content_routing_decision"] = {
"matched_preference": matched_pref,
"model": matched_model,
"confidence": confidence,
"classifier": classifier,
}
return PreRoutingHookResponse(
model=matched_model,
messages=messages,
)

View file

@ -0,0 +1,273 @@
"""
Utility helpers for the content-aware router.
All functions are pure (no I/O) and dependency-free so they can be unit-tested
without standing up any LLM infrastructure.
Matching pipeline
-----------------
tokenize (lowercase + punctuation removal + stop-word filter + light stemming)
build_bm25_index (precomputes per-term IDF and average document length)
bm25_score (Okapi BM25 handles term saturation and length normalisation)
Stemming ensures morphological variants ("reason" / "reasoning",
"function" / "functions", "debug" / "debugging") map to the same stem so they
match across descriptions and prompts even when written in different forms.
"""
import math
import re
import string
from typing import Dict, List, Optional, Tuple
# Okapi BM25 hyper-parameters (standard defaults)
_BM25_K1 = 1.5 # term-frequency saturation
_BM25_B = 0.75 # length normalisation
# ---------------------------------------------------------------------------
# Stop words
# ---------------------------------------------------------------------------
_STOP_WORDS = frozenset(
{
"a", "an", "the", "and", "or", "but", "in", "on", "at", "to", "for",
"of", "with", "by", "from", "is", "are", "was", "were", "be", "been",
"being", "have", "has", "had", "do", "does", "did", "will", "would",
"could", "should", "may", "might", "shall", "can", "i", "you", "he",
"she", "it", "we", "they", "this", "that", "these", "those", "as",
"if", "then", "than", "so", "not", "no", "nor", "yet", "both",
"either", "neither", "about", "into", "through", "during", "including",
"until", "while", "of", "about", "against", "between", "into",
"through", "such", "any", "more", "also", "use", "using", "used",
"help", "based", "related", "request", "task", "tasks",
}
)
# ---------------------------------------------------------------------------
# Light stemmer
# ---------------------------------------------------------------------------
def stem(word: str) -> str:
"""
Light suffix-stripping stemmer for English.
Reduces the most common inflected forms to a shared base so that
morphological variants produce the same index token:
reasoning reason (-ing, 6+ chars remaining)
debugging debug (-ing + de-double trailing consonant)
implementing implement (-ing)
functions function (-s)
algorithms algorithm (-s)
processes process (-es)
stories story (-ies y)
Deliberately conservative: only strips -ing, -ed, -ies, -es, and -s so
that common technical terms ("function", "computation", "section") are
not truncated to unrecognisable stems by over-aggressive -tion/-ation
rules.
Rules are applied in priority order; the first matching rule wins.
Minimum word-length guards prevent over-truncation on short words.
"""
n = len(word)
# -ing → strip, then de-double trailing consonant
# reasoning→reason, debugging→debugg→debug, implementing→implement
if n > 5 and word.endswith("ing"):
candidate = word[:-3]
if (
len(candidate) >= 2
and candidate[-1] == candidate[-2]
and candidate[-1] not in "aeiou"
):
candidate = candidate[:-1]
return candidate
# -ed → strip, then de-double (only for longer words)
if n > 5 and word.endswith("ed"):
candidate = word[:-2]
if (
len(candidate) >= 2
and candidate[-1] == candidate[-2]
and candidate[-1] not in "aeiou"
):
candidate = candidate[:-1]
return candidate
# -ies → y (entries→entry, stories→story)
if n > 4 and word.endswith("ies"):
return word[:-3] + "y"
# -es → strip (processes→process, classes→class, accesses→access)
if n > 4 and word.endswith("es"):
return word[:-2]
# -s → strip (functions→function, algorithms→algorithm)
# Skip words ending in -ss (class, process, etc.) to avoid truncation
if n > 3 and word.endswith("s") and not word.endswith("ss"):
return word[:-1]
return word
# ---------------------------------------------------------------------------
# Tokenizer
# ---------------------------------------------------------------------------
def tokenize(text: str) -> List[str]:
"""
Normalise *text* to a list of stemmed tokens.
Steps:
1. Lowercase
2. Replace punctuation with spaces
3. Split on whitespace
4. Drop stop words and single-character tokens
5. Apply light stemming so inflected forms share a common base
"""
text = text.lower()
text = re.sub(r"[" + re.escape(string.punctuation) + r"]", " ", text)
raw = text.split()
return [
stem(t)
for t in raw
if t and t not in _STOP_WORDS and len(t) > 1
]
# ---------------------------------------------------------------------------
# BM25 index
# ---------------------------------------------------------------------------
def build_bm25_index(
corpus: List[str],
) -> Tuple[List[List[str]], Dict[str, float], float]:
"""
Build an Okapi BM25 index for a list of documents.
Args:
corpus: list of raw text strings (preference descriptions).
Returns:
tokenized_corpus : stemmed token list per document
idf_weights : per-term IDF values (RobertsonSparck Jones, always 0)
avgdl : average document length in tokens
"""
tokenized: List[List[str]] = [tokenize(doc) for doc in corpus]
n = len(tokenized)
avgdl = sum(len(tokens) for tokens in tokenized) / max(n, 1)
# Document frequency
df: Dict[str, int] = {}
for tokens in tokenized:
for t in set(tokens):
df[t] = df.get(t, 0) + 1
# IDF: log((N - df + 0.5) / (df + 0.5) + 1) — always positive
idf: Dict[str, float] = {
t: math.log((n - count + 0.5) / (count + 0.5) + 1)
for t, count in df.items()
}
return tokenized, idf, avgdl
def bm25_score(
query_tokens: List[str],
doc_tokens: List[str],
idf: Dict[str, float],
avgdl: float,
) -> float:
"""
Okapi BM25 relevance score for a query against a single document.
Args:
query_tokens : stemmed prompt tokens.
doc_tokens : stemmed description tokens (from the BM25 index).
idf : shared IDF weights built by build_bm25_index().
avgdl : average document length from build_bm25_index().
Returns:
Non-negative float; higher means more relevant.
"""
if not query_tokens or not doc_tokens:
return 0.0
dl = len(doc_tokens)
doc_freq: Dict[str, int] = {}
for t in doc_tokens:
doc_freq[t] = doc_freq.get(t, 0) + 1
score = 0.0
for term in set(query_tokens):
f = doc_freq.get(term, 0)
if f == 0:
continue
term_idf = idf.get(term, 0.0)
if term_idf <= 0.0:
continue
# BM25 term-frequency component with length normalisation
length_norm = 1 - _BM25_B + _BM25_B * dl / max(avgdl, 1)
tf_component = f * (_BM25_K1 + 1) / (f + _BM25_K1 * length_norm)
score += term_idf * tf_component
return score
# ---------------------------------------------------------------------------
# Dense-vector cosine similarity (used by embedding_similarity classifier)
# ---------------------------------------------------------------------------
def cosine_similarity(a: List[float], b: List[float]) -> float:
"""Cosine similarity between two dense float vectors."""
if not a or not b or len(a) != len(b):
return 0.0
dot = sum(x * y for x, y in zip(a, b))
norm_a = math.sqrt(sum(x * x for x in a))
norm_b = math.sqrt(sum(y * y for y in b))
if norm_a == 0.0 or norm_b == 0.0:
return 0.0
return dot / (norm_a * norm_b)
# ---------------------------------------------------------------------------
# Message extraction
# ---------------------------------------------------------------------------
def extract_prompt_text(
messages: Optional[List[Dict]],
) -> Tuple[str, Optional[str]]:
"""
Extract the last user message and the last system prompt from a messages list.
Returns:
(user_text, system_text) system_text may be None if no system message found.
"""
if not messages:
return "", None
user_text: Optional[str] = None
system_text: Optional[str] = None
for msg in reversed(messages):
role = msg.get("role", "")
content = msg.get("content") or ""
if isinstance(content, list):
# Content-part format: [{"type": "text", "text": "..."}]
parts = [
p.get("text", "")
for p in content
if isinstance(p, dict) and p.get("type") == "text"
]
content = " ".join(parts).strip()
if isinstance(content, str) and content:
if role == "user" and user_text is None:
user_text = content
elif role == "system" and system_text is None:
system_text = content
if user_text is not None and system_text is not None:
break
return user_text or "", system_text

View file

@ -0,0 +1,210 @@
"""
Model Affinity Router Session Pinning
Prevents mid-session model switching in agentic loops where successive calls
may have different content characteristics that would otherwise cause
content-aware routing to select different models.
How it works
------------
1. Client sends ``X-Model-Affinity: <session-id>`` (any opaque string, typically
a UUID) in the request header.
2. First request for that session-id routes normally (content-aware routing runs,
or infrastructure routing picks the model). The selected **model name** is
stored in the affinity cache keyed by session-id.
3. All subsequent requests carrying the same session-id are pinned to the cached
model content-aware routing is bypassed entirely.
4. Pinning is at the **model group** level, not the deployment level, so
load-balancing and failover within the pinned model group still work normally.
5. Entries expire after a configurable TTL (default 10 min) and the cache uses
LRU eviction once its capacity is reached.
Storage backends
----------------
- ``local`` in-process LRU + TTL cache (default, zero extra deps)
- ``redis`` shared across multiple proxy replicas; requires the ``redis``
package and a ``redis_url`` in the config
"""
import asyncio
import time
from collections import OrderedDict
from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple
from litellm._logging import verbose_router_logger
if TYPE_CHECKING:
from litellm.types.router import ModelAffinityConfig
else:
ModelAffinityConfig = Any
# ---------------------------------------------------------------------------
# Local LRU cache
# ---------------------------------------------------------------------------
class _LocalAffinityCache:
"""
Thread-safe, async-friendly LRU cache with per-entry TTL.
Entries are stored as ``(model_name, expires_at)`` tuples.
``get()`` evicts expired entries on access so they cannot be returned.
The oldest entry is evicted when ``max_size`` is reached.
"""
def __init__(self, max_size: int, ttl: float) -> None:
self._max_size = max_size
self._ttl = ttl
# OrderedDict: oldest entries at the front (for LRU eviction)
self._store: OrderedDict[str, Tuple[str, float]] = OrderedDict()
self._lock = asyncio.Lock()
async def get(self, key: str) -> Optional[str]:
async with self._lock:
entry = self._store.get(key)
if entry is None:
return None
model, expires_at = entry
if time.monotonic() > expires_at:
del self._store[key]
return None
# Promote to most-recently-used
self._store.move_to_end(key)
return model
async def set(self, key: str, value: str) -> None:
async with self._lock:
if key in self._store:
self._store.move_to_end(key)
self._store[key] = (value, time.monotonic() + self._ttl)
if len(self._store) > self._max_size:
# Evict the least-recently-used entry
self._store.popitem(last=False)
async def delete(self, key: str) -> None:
async with self._lock:
self._store.pop(key, None)
def __len__(self) -> int:
return len(self._store)
# ---------------------------------------------------------------------------
# Redis cache wrapper
# ---------------------------------------------------------------------------
class _RedisAffinityCache:
"""
Redis-backed affinity cache. TTL is delegated to Redis SETEX so entries
expire server-side even if the Python process restarts.
"""
_KEY_PREFIX = "litellm:model_affinity:"
def __init__(self, redis_url: str, ttl: int) -> None:
self._ttl = ttl
try:
import redis.asyncio as aioredis
self._client = aioredis.from_url(redis_url, decode_responses=True)
except ImportError as exc:
raise ImportError(
"redis package is required for ModelAffinityRouter storage='redis'. "
"Install it with: pip install redis"
) from exc
def _key(self, session_id: str) -> str:
return f"{self._KEY_PREFIX}{session_id}"
async def get(self, session_id: str) -> Optional[str]:
try:
return await self._client.get(self._key(session_id))
except Exception as e:
verbose_router_logger.warning(
f"ModelAffinityRouter: Redis GET failed ({e}), treating as cache miss"
)
return None
async def set(self, session_id: str, model: str) -> None:
try:
await self._client.setex(self._key(session_id), self._ttl, model)
except Exception as e:
verbose_router_logger.warning(
f"ModelAffinityRouter: Redis SETEX failed ({e}), pin will not persist"
)
async def delete(self, session_id: str) -> None:
try:
await self._client.delete(self._key(session_id))
except Exception as e:
verbose_router_logger.warning(
f"ModelAffinityRouter: Redis DELETE failed ({e})"
)
# ---------------------------------------------------------------------------
# Public interface
# ---------------------------------------------------------------------------
class ModelAffinityRouter:
"""
Session-pin cache used by the Router to enforce model affinity.
Instantiated by the Router when ``router_settings.model_affinity.enabled=true``.
The routing logic itself lives in ``Router.async_pre_routing_hook`` this
class is only responsible for managing the underlying cache.
"""
def __init__(self, config: "ModelAffinityConfig") -> None:
self.config = config
storage = config.storage or "local"
if storage == "redis":
if not config.redis_url:
raise ValueError(
"ModelAffinityConfig.redis_url is required when storage='redis'"
)
self._cache: Any = _RedisAffinityCache(
redis_url=config.redis_url,
ttl=config.ttl,
)
else:
self._cache = _LocalAffinityCache(
max_size=config.max_sessions,
ttl=float(config.ttl),
)
verbose_router_logger.info(
f"ModelAffinityRouter initialized: storage={storage} "
f"ttl={config.ttl}s max_sessions={config.max_sessions}"
)
async def get_pinned_model(self, session_id: str) -> Optional[str]:
"""Return the pinned model name for *session_id*, or None if not pinned."""
model = await self._cache.get(session_id)
if model:
verbose_router_logger.debug(
f"ModelAffinityRouter: cache HIT session={session_id} -> model={model}"
)
else:
verbose_router_logger.debug(
f"ModelAffinityRouter: cache MISS session={session_id}"
)
return model
async def pin_model(self, session_id: str, model: str) -> None:
"""Pin *model* for *session_id* (creates or refreshes the TTL)."""
await self._cache.set(session_id, model)
verbose_router_logger.info(
f"ModelAffinityRouter: pinned session={session_id} -> model={model} "
f"(ttl={self.config.ttl}s)"
)
async def clear_session(self, session_id: str) -> None:
"""Explicitly remove a session pin (e.g., on logout or explicit reset)."""
await self._cache.delete(session_id)
verbose_router_logger.info(
f"ModelAffinityRouter: cleared session={session_id}"
)

View file

@ -265,4 +265,18 @@ ROUTER_SETTINGS_FIELDS: List[RouterSettingsField] = [
field_default=None,
ui_field_name="Disable Cooldowns",
),
RouterSettingsField(
field_name="content_routing",
field_type="Dictionary",
field_value=None,
field_description=(
"Content-aware routing configuration. When enabled, classifies incoming "
"prompt content against per-model routing_preferences and routes to the "
"best-matched model. Supports rule_based (TF-IDF), embedding_similarity, "
"and external_model classifiers."
),
field_default=None,
ui_field_name="Content Routing",
link="https://docs.litellm.ai/docs/routing",
),
]

View file

@ -359,6 +359,44 @@ class DeploymentTypedDict(TypedDict, total=False):
model_info: dict
class RoutingPreference(BaseModel):
"""A content-based routing preference declaration for a model deployment."""
name: str
description: str
class ContentRoutingConfig(BaseModel):
"""Global content-aware routing configuration set in router_settings.content_routing."""
enabled: bool = False
classifier: Literal["rule_based", "embedding_similarity", "external_model"] = (
"rule_based"
)
default_model: Optional[str] = None
confidence_threshold: float = 0.1
# embedding_similarity only
embedding_model: Optional[str] = None
# external_model only
external_classifier_url: Optional[str] = None
class ModelAffinityConfig(BaseModel):
"""
Session-pinning configuration set in router_settings.model_affinity.
When enabled, the first request for a given X-Model-Affinity session-id
selects a model normally; all subsequent requests reuse that model for the
lifetime of the pin (ttl seconds), regardless of content.
"""
enabled: bool = False
ttl: int = 600 # seconds before a session pin expires (default 10 min)
max_sessions: int = 10_000 # LRU capacity for the local cache
storage: Literal["local", "redis"] = "local"
redis_url: Optional[str] = None # required when storage="redis"
SPECIAL_MODEL_INFO_PARAMS = [
"input_cost_per_token",
"output_cost_per_token",
@ -371,6 +409,7 @@ class Deployment(BaseModel):
model_name: str
litellm_params: LiteLLM_Params
model_info: ModelInfo
routing_preferences: Optional[List[RoutingPreference]] = None
model_config = ConfigDict(extra="allow", protected_namespaces=())

View file

@ -0,0 +1,639 @@
"""
Unit tests for content-aware routing.
All tests are self-contained no LLM calls, no external services.
"""
import math
from typing import Dict, List, Optional
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from litellm.router_strategy.content_aware_router.utils import (
bm25_score,
build_bm25_index,
cosine_similarity,
extract_prompt_text,
stem,
tokenize,
)
# ---------------------------------------------------------------------------
# Utils tests
# ---------------------------------------------------------------------------
class TestStem:
def test_ing_suffix(self):
assert stem("reasoning") == "reason"
assert stem("debugging") == "debug" # de-doubles trailing g
assert stem("implementing") == "implement"
def test_s_suffix(self):
assert stem("functions") == "function"
assert stem("algorithms") == "algorithm"
def test_ed_suffix(self):
# "sorted" → strip -ed → "sort"
assert stem("sorted") == "sort"
def test_ies_suffix(self):
assert stem("stories") == "story"
assert stem("entries") == "entry"
def test_es_suffix(self):
assert stem("processes") == "process"
assert stem("classes") == "class"
def test_no_change_on_short_word(self):
assert stem("go") == "go"
assert stem("code") == "code"
def test_function_not_over_stemmed(self):
# "function" must NOT be truncated by a -tion rule to "func"
assert stem("function") == "function"
# "computation" must stay intact (no destructive -tion stripping)
assert stem("computation") == "computation"
def test_morphological_variants_share_stem(self):
# reasoning and reason must produce the same stem
assert stem("reasoning") == stem("reason") == "reason"
# functions and function must share stem
assert stem("functions") == stem("function") == "function"
class TestTokenize:
def test_lowercases(self):
# "hello" is stem("hello") == "hello"
assert "hello" in tokenize("Hello World")
def test_removes_stop_words(self):
tokens = tokenize("this is a test")
assert "this" not in tokens
assert "is" not in tokens
assert "a" not in tokens
assert stem("test") in tokens
def test_removes_punctuation(self):
tokens = tokenize("code, debugging, and explaining!")
assert stem("code") in tokens
assert stem("debugging") in tokens
assert stem("explaining") in tokens
def test_empty_string(self):
assert tokenize("") == []
def test_single_char_removed(self):
# Single-character tokens are filtered out
tokens = tokenize("a b c hello")
assert "a" not in tokens
assert "b" not in tokens
assert "hello" in tokens
def test_stemming_applied(self):
# tokenize must apply stemming so inflected and base forms share the same token
t_reasoning = tokenize("reasoning")
t_reason = tokenize("reason")
assert t_reasoning == t_reason, (
f"'reasoning' tokenised to {t_reasoning} but 'reason' to {t_reason}; "
"they should share the same stem"
)
class TestCosineSimilarity:
def test_identical_vectors(self):
v = [1.0, 2.0, 3.0]
assert abs(cosine_similarity(v, v) - 1.0) < 1e-9
def test_orthogonal_vectors(self):
a = [1.0, 0.0]
b = [0.0, 1.0]
assert abs(cosine_similarity(a, b)) < 1e-9
def test_zero_vector_returns_zero(self):
assert cosine_similarity([0.0, 0.0], [1.0, 2.0]) == 0.0
def test_empty_returns_zero(self):
assert cosine_similarity([], []) == 0.0
def test_length_mismatch_returns_zero(self):
assert cosine_similarity([1.0, 2.0], [1.0]) == 0.0
class TestBuildBM25Index:
def test_returns_correct_length(self):
corpus = ["code debugging programming", "creative writing storytelling"]
corpus_tokens, idf, avgdl = build_bm25_index(corpus)
assert len(corpus_tokens) == 2
assert len(idf) > 0
assert avgdl > 0
def test_unique_terms_get_positive_idf(self):
corpus = ["code programming", "creative writing"]
_, idf, _ = build_bm25_index(corpus)
# stem("code") appears only in doc 0 → idf should be positive
assert idf.get(stem("code"), 0) > 0
def test_single_document(self):
corpus = ["hello world testing"]
corpus_tokens, idf, avgdl = build_bm25_index(corpus)
assert len(corpus_tokens) == 1
def test_bm25_score_nonzero_for_matching_term(self):
corpus = ["code programming python", "storytelling creative fiction"]
corpus_tokens, idf, avgdl = build_bm25_index(corpus)
query = tokenize("write python code")
score_code = bm25_score(query, corpus_tokens[0], idf, avgdl)
score_story = bm25_score(query, corpus_tokens[1], idf, avgdl)
assert score_code > score_story, (
f"code description ({score_code:.4f}) should score higher than "
f"story description ({score_story:.4f}) for a code prompt"
)
def test_bm25_score_zero_for_no_overlap(self):
corpus = ["apple orange mango"]
corpus_tokens, idf, avgdl = build_bm25_index(corpus)
query = tokenize("quantum physics entanglement")
assert bm25_score(query, corpus_tokens[0], idf, avgdl) == 0.0
def test_stemmed_variants_match(self):
# 'reasoning' in description should match 'reason' in query after stemming
corpus = ["analyze explain reasoning logic mathematical math proof theorem equation"]
corpus_tokens, idf, avgdl = build_bm25_index(corpus)
query_base = tokenize("reason")
query_inflected = tokenize("reasoning")
# Both queries must produce identical scores (same stem)
score_base = bm25_score(query_base, corpus_tokens[0], idf, avgdl)
score_inflected = bm25_score(query_inflected, corpus_tokens[0], idf, avgdl)
assert score_base == score_inflected, (
f"'reason' ({score_base:.4f}) and 'reasoning' ({score_inflected:.4f}) "
"should produce the same BM25 score after stemming"
)
assert score_base > 0.0
class TestExtractPromptText:
def test_extracts_last_user_message(self):
messages = [
{"role": "user", "content": "first message"},
{"role": "assistant", "content": "response"},
{"role": "user", "content": "last user message"},
]
user_text, system_text = extract_prompt_text(messages)
assert user_text == "last user message"
assert system_text is None
def test_extracts_system_prompt(self):
messages = [
{"role": "system", "content": "You are a coding assistant"},
{"role": "user", "content": "write a function"},
]
user_text, system_text = extract_prompt_text(messages)
assert user_text == "write a function"
assert system_text == "You are a coding assistant"
def test_handles_content_parts(self):
messages = [
{
"role": "user",
"content": [{"type": "text", "text": "write python code"}],
}
]
user_text, _ = extract_prompt_text(messages)
assert user_text == "write python code"
def test_empty_messages_returns_empty(self):
user_text, system_text = extract_prompt_text([])
assert user_text == ""
assert system_text is None
def test_none_returns_empty(self):
user_text, system_text = extract_prompt_text(None)
assert user_text == ""
assert system_text is None
# ---------------------------------------------------------------------------
# ContentAwareRouter tests
# ---------------------------------------------------------------------------
def _make_router(preferences_by_model, classifier="rule_based", default_model="gpt-4o", threshold=0.01):
"""Create a ContentAwareRouter with minimal mocking."""
from litellm.router_strategy.content_aware_router.content_aware_router import (
ContentAwareRouter,
)
from litellm.types.router import ContentRoutingConfig, RoutingPreference
prefs = {
model: [RoutingPreference(**p) for p in pref_list]
for model, pref_list in preferences_by_model.items()
}
config = ContentRoutingConfig(
enabled=True,
classifier=classifier,
default_model=default_model,
confidence_threshold=threshold,
)
mock_router = MagicMock()
return ContentAwareRouter(
preferences_by_model=prefs,
config=config,
litellm_router_instance=mock_router,
)
class TestRuleBasedClassifier:
def test_routes_code_prompt(self):
"""Code-heavy prompts should route to the code model."""
router = _make_router(
{
"gpt-4o": [
{"name": "creative_writing", "description": "creative storytelling narrative fiction writing"},
],
"claude-sonnet": [
{"name": "code_generation", "description": "code programming debugging python javascript function"},
],
},
threshold=0.0,
)
model, pref, score = router._classify_rule_based(
"write a python function to sort a list", None
)
assert model == "claude-sonnet"
assert pref == "code_generation"
assert score > 0
def test_routes_creative_prompt(self):
"""Creative prompts should route to the creative writing model."""
router = _make_router(
{
"gpt-4o": [
{"name": "creative_writing", "description": "creative storytelling narrative fiction writing poetry"},
],
"claude-sonnet": [
{"name": "code_generation", "description": "code programming debugging function"},
],
},
threshold=0.0,
)
model, pref, score = router._classify_rule_based(
"write a short story about a lonely astronaut", None
)
assert model == "gpt-4o"
assert pref == "creative_writing"
def test_below_threshold_uses_default(self):
"""When confidence is below threshold, default_model is returned."""
from litellm.router_strategy.content_aware_router.content_aware_router import (
ContentAwareRouter,
)
from litellm.types.router import ContentRoutingConfig, RoutingPreference
prefs = {
"claude-sonnet": [
RoutingPreference(name="code_generation", description="very specific coding words only"),
],
}
# Set threshold absurdly high so nothing matches
config = ContentRoutingConfig(
enabled=True,
classifier="rule_based",
default_model="gpt-4o",
confidence_threshold=999.0,
)
mock_router = MagicMock()
car = ContentAwareRouter(prefs, config, mock_router)
import asyncio
messages = [{"role": "user", "content": "hello how are you"}]
result = asyncio.get_event_loop().run_until_complete(
car.async_pre_routing_hook(
model="any",
request_kwargs={},
messages=messages,
)
)
assert result is not None
assert result.model == "gpt-4o"
def test_no_preferences_returns_none(self):
"""Router with no routing_preferences should return None."""
from litellm.router_strategy.content_aware_router.content_aware_router import (
ContentAwareRouter,
)
from litellm.types.router import ContentRoutingConfig
config = ContentRoutingConfig(enabled=True, classifier="rule_based")
mock_router = MagicMock()
car = ContentAwareRouter({}, config, mock_router)
import asyncio
messages = [{"role": "user", "content": "write a python function"}]
result = asyncio.get_event_loop().run_until_complete(
car.async_pre_routing_hook(
model="any",
request_kwargs={},
messages=messages,
)
)
assert result is None
def test_no_user_message_returns_none(self):
"""Missing user message should return None."""
router = _make_router(
{"gpt-4o": [{"name": "general", "description": "general conversation chat"}]},
threshold=0.0,
)
import asyncio
result = asyncio.get_event_loop().run_until_complete(
router.async_pre_routing_hook(
model="any",
request_kwargs={},
messages=[{"role": "system", "content": "system only"}],
)
)
assert result is None
def test_routing_decision_stored_in_metadata(self):
"""Classification decision should be written into request_kwargs metadata."""
router = _make_router(
{
"claude-sonnet": [
{"name": "code_generation", "description": "code programming debugging function python"},
],
},
threshold=0.0,
)
import asyncio
request_kwargs: dict = {}
messages = [{"role": "user", "content": "write a python function to sort a list"}]
asyncio.get_event_loop().run_until_complete(
router.async_pre_routing_hook(
model="any",
request_kwargs=request_kwargs,
messages=messages,
)
)
decision = request_kwargs.get("metadata", {}).get("content_routing_decision")
assert decision is not None
assert decision["model"] == "claude-sonnet"
assert decision["matched_preference"] == "code_generation"
assert decision["classifier"] == "rule_based"
class TestEmbeddingSimilarityClassifier:
@pytest.mark.asyncio
async def test_embedding_similarity_uses_cosine(self):
"""Embedding classifier should pick the model with highest cosine similarity."""
from litellm.router_strategy.content_aware_router.content_aware_router import (
ContentAwareRouter,
)
from litellm.types.router import ContentRoutingConfig, RoutingPreference
prefs = {
"gpt-4o": [
RoutingPreference(name="creative_writing", description="storytelling"),
],
"claude-sonnet": [
RoutingPreference(name="code_generation", description="code programming"),
],
}
config = ContentRoutingConfig(
enabled=True,
classifier="embedding_similarity",
default_model="gpt-4o",
confidence_threshold=0.0,
)
mock_router = MagicMock()
car = ContentAwareRouter(prefs, config, mock_router)
# Pre-populate description embeddings (2 descriptions)
# Code description embedding points towards [1, 0], creative towards [0, 1]
car._description_embeddings = [
[0.0, 1.0], # gpt-4o / creative_writing
[1.0, 0.0], # claude-sonnet / code_generation
]
# Prompt embedding pointing towards code
mock_embedding_response = MagicMock()
mock_embedding_response.data = [{"embedding": [0.9, 0.1]}]
with patch("litellm.aembedding", new=AsyncMock(return_value=mock_embedding_response)):
model, pref, score = await car._classify_embedding_similarity(
"write python code", None
)
assert model == "claude-sonnet"
assert pref == "code_generation"
assert score > 0.5
@pytest.mark.asyncio
async def test_embedding_fallback_on_error(self):
"""Embedding failures should fall back to rule_based without raising."""
from litellm.router_strategy.content_aware_router.content_aware_router import (
ContentAwareRouter,
)
from litellm.types.router import ContentRoutingConfig, RoutingPreference
prefs = {
"claude-sonnet": [
RoutingPreference(name="code_generation", description="code programming python function"),
],
}
config = ContentRoutingConfig(
enabled=True,
classifier="embedding_similarity",
default_model="gpt-4o",
confidence_threshold=0.0,
embedding_model="text-embedding-3-small",
)
mock_router = MagicMock()
car = ContentAwareRouter(prefs, config, mock_router)
with patch("litellm.aembedding", new=AsyncMock(side_effect=Exception("API error"))):
# Should not raise; falls back to rule_based
model, pref, score = await car._classify_embedding_similarity(
"write a python function", None
)
# Fell back to rule_based — claude-sonnet should win for code prompt
assert model == "claude-sonnet"
class TestExternalModelClassifier:
@pytest.mark.asyncio
async def test_external_model_parses_response(self):
"""External classifier should parse JSON response correctly."""
from litellm.router_strategy.content_aware_router.content_aware_router import (
ContentAwareRouter,
)
from litellm.types.router import ContentRoutingConfig, RoutingPreference
prefs = {
"claude-sonnet": [
RoutingPreference(name="code_generation", description="code programming"),
],
}
config = ContentRoutingConfig(
enabled=True,
classifier="external_model",
default_model="gpt-4o",
confidence_threshold=0.0,
external_classifier_url="http://arch-router/classify",
)
mock_router = MagicMock()
car = ContentAwareRouter(prefs, config, mock_router)
mock_response = MagicMock()
mock_response.json.return_value = {
"matched_preference": "code_generation",
"model": "claude-sonnet",
"confidence": 0.95,
}
mock_response.raise_for_status = MagicMock()
mock_client = AsyncMock()
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
mock_client.__aexit__ = AsyncMock(return_value=False)
mock_client.post = AsyncMock(return_value=mock_response)
with patch("httpx.AsyncClient", return_value=mock_client):
model, pref, confidence = await car._classify_external_model(
"write python code", None
)
assert model == "claude-sonnet"
assert pref == "code_generation"
assert abs(confidence - 0.95) < 1e-9
@pytest.mark.asyncio
async def test_external_model_fallback_on_error(self):
"""HTTP errors should fall back to rule_based."""
from litellm.router_strategy.content_aware_router.content_aware_router import (
ContentAwareRouter,
)
from litellm.types.router import ContentRoutingConfig, RoutingPreference
prefs = {
"claude-sonnet": [
RoutingPreference(name="code_generation", description="code programming python function"),
],
}
config = ContentRoutingConfig(
enabled=True,
classifier="external_model",
default_model="gpt-4o",
confidence_threshold=0.0,
external_classifier_url="http://arch-router/classify",
)
mock_router = MagicMock()
car = ContentAwareRouter(prefs, config, mock_router)
mock_client = AsyncMock()
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
mock_client.__aexit__ = AsyncMock(return_value=False)
mock_client.post = AsyncMock(side_effect=Exception("connection refused"))
with patch("httpx.AsyncClient", return_value=mock_client):
model, pref, score = await car._classify_external_model(
"write a python function", None
)
# Fell back to rule_based
assert model == "claude-sonnet"
@pytest.mark.asyncio
async def test_external_model_no_url_falls_back(self):
"""Missing URL should fall back to rule_based."""
from litellm.router_strategy.content_aware_router.content_aware_router import (
ContentAwareRouter,
)
from litellm.types.router import ContentRoutingConfig, RoutingPreference
prefs = {
"claude-sonnet": [
RoutingPreference(name="code_generation", description="code programming python function"),
],
}
config = ContentRoutingConfig(
enabled=True,
classifier="external_model",
default_model="gpt-4o",
confidence_threshold=0.0,
external_classifier_url=None, # no URL set
)
mock_router = MagicMock()
car = ContentAwareRouter(prefs, config, mock_router)
model, pref, score = await car._classify_external_model("write python code", None)
assert model == "claude-sonnet"
class TestRouterIntegration:
def test_router_init_with_content_routing(self):
"""Router should initialize ContentAwareRouter when content_routing is set."""
import litellm
router = litellm.Router(
model_list=[
{
"model_name": "gpt-4o",
"litellm_params": {"model": "openai/gpt-4o", "api_key": "fake"},
"routing_preferences": [
{"name": "creative_writing", "description": "creative storytelling fiction"},
],
},
{
"model_name": "claude-sonnet",
"litellm_params": {"model": "anthropic/claude-sonnet-4-20250514", "api_key": "fake"},
"routing_preferences": [
{"name": "code_generation", "description": "code programming debugging"},
],
},
],
content_routing={
"enabled": True,
"classifier": "rule_based",
"default_model": "gpt-4o",
"confidence_threshold": 0.0,
},
)
assert router.content_aware_router is not None
assert len(router.content_aware_router._index) == 2
assert router._content_routing_config is not None
assert router._content_routing_config.enabled is True
def test_router_disabled_content_routing(self):
"""Router should not set content_aware_router when enabled=False."""
import litellm
router = litellm.Router(
model_list=[
{
"model_name": "gpt-4o",
"litellm_params": {"model": "openai/gpt-4o", "api_key": "fake"},
},
],
content_routing={"enabled": False},
)
assert router.content_aware_router is None
def test_router_no_content_routing(self):
"""Router without content_routing param should have None content_aware_router."""
import litellm
router = litellm.Router(
model_list=[
{
"model_name": "gpt-4o",
"litellm_params": {"model": "openai/gpt-4o", "api_key": "fake"},
},
],
)
assert router.content_aware_router is None

View file

@ -0,0 +1,321 @@
"""
Unit tests for Model Affinity (Session Pinning) router.
All tests are self-contained no LLM calls, no external services.
"""
import asyncio
import time
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from litellm.router_strategy.model_affinity_router.model_affinity_router import (
ModelAffinityRouter,
_LocalAffinityCache,
)
from litellm.types.router import ModelAffinityConfig
# ---------------------------------------------------------------------------
# _LocalAffinityCache tests
# ---------------------------------------------------------------------------
class TestLocalAffinityCache:
def _cache(self, max_size=100, ttl=60.0) -> _LocalAffinityCache:
return _LocalAffinityCache(max_size=max_size, ttl=ttl)
def _run(self, coro):
return asyncio.get_event_loop().run_until_complete(coro)
def test_set_and_get(self):
cache = self._cache()
self._run(cache.set("s1", "gpt-4o"))
assert self._run(cache.get("s1")) == "gpt-4o"
def test_miss_returns_none(self):
cache = self._cache()
assert self._run(cache.get("nonexistent")) is None
def test_delete_removes_entry(self):
cache = self._cache()
self._run(cache.set("s1", "gpt-4o"))
self._run(cache.delete("s1"))
assert self._run(cache.get("s1")) is None
def test_delete_missing_key_is_noop(self):
cache = self._cache()
self._run(cache.delete("nonexistent")) # should not raise
def test_ttl_expiry(self):
cache = self._cache(ttl=0.05) # 50 ms TTL
self._run(cache.set("s1", "gpt-4o"))
assert self._run(cache.get("s1")) == "gpt-4o"
time.sleep(0.1) # outlast the TTL
assert self._run(cache.get("s1")) is None # expired
def test_lru_eviction(self):
"""When max_size is exceeded the oldest entry is evicted."""
cache = self._cache(max_size=2)
self._run(cache.set("a", "m1"))
self._run(cache.set("b", "m2"))
self._run(cache.set("c", "m3")) # should evict "a"
assert self._run(cache.get("a")) is None # evicted
assert self._run(cache.get("b")) == "m2"
assert self._run(cache.get("c")) == "m3"
def test_get_promotes_to_mru(self):
"""Accessing an entry should move it to MRU so it is not evicted first."""
cache = self._cache(max_size=2)
self._run(cache.set("a", "m1"))
self._run(cache.set("b", "m2"))
# access "a" so it becomes MRU
self._run(cache.get("a"))
self._run(cache.set("c", "m3")) # should evict "b", not "a"
assert self._run(cache.get("a")) == "m1"
assert self._run(cache.get("b")) is None # evicted
def test_overwrite_refreshes_ttl(self):
"""Re-setting a key should reset its TTL."""
cache = self._cache(ttl=0.05)
self._run(cache.set("s1", "gpt-4o"))
time.sleep(0.03)
self._run(cache.set("s1", "gpt-4o")) # refresh
time.sleep(0.03)
# 60 ms have passed in total, but TTL reset after 30 ms so still alive
assert self._run(cache.get("s1")) == "gpt-4o"
# ---------------------------------------------------------------------------
# ModelAffinityRouter tests
# ---------------------------------------------------------------------------
def _make_affinity_router(ttl=600, max_sessions=1000) -> ModelAffinityRouter:
config = ModelAffinityConfig(
enabled=True, ttl=ttl, max_sessions=max_sessions, storage="local"
)
return ModelAffinityRouter(config=config)
class TestModelAffinityRouter:
def _run(self, coro):
return asyncio.get_event_loop().run_until_complete(coro)
def test_new_session_is_unpinned(self):
r = _make_affinity_router()
assert self._run(r.get_pinned_model("new-session")) is None
def test_pin_and_retrieve(self):
r = _make_affinity_router()
self._run(r.pin_model("s1", "claude-sonnet"))
assert self._run(r.get_pinned_model("s1")) == "claude-sonnet"
def test_clear_session(self):
r = _make_affinity_router()
self._run(r.pin_model("s1", "gpt-4o"))
self._run(r.clear_session("s1"))
assert self._run(r.get_pinned_model("s1")) is None
def test_multiple_sessions_are_independent(self):
r = _make_affinity_router()
self._run(r.pin_model("s1", "gpt-4o"))
self._run(r.pin_model("s2", "claude-sonnet"))
assert self._run(r.get_pinned_model("s1")) == "gpt-4o"
assert self._run(r.get_pinned_model("s2")) == "claude-sonnet"
def test_pin_can_be_updated(self):
r = _make_affinity_router()
self._run(r.pin_model("s1", "gpt-4o"))
self._run(r.pin_model("s1", "claude-sonnet")) # overwrite
assert self._run(r.get_pinned_model("s1")) == "claude-sonnet"
# ---------------------------------------------------------------------------
# Router integration tests
# ---------------------------------------------------------------------------
class TestRouterAffinityIntegration:
def _run(self, coro):
return asyncio.get_event_loop().run_until_complete(coro)
def _make_router(self, **affinity_kwargs):
import litellm
return litellm.Router(
model_list=[
{
"model_name": "gpt-4o",
"litellm_params": {"model": "openai/gpt-4o", "api_key": "fake"},
},
{
"model_name": "claude-sonnet",
"litellm_params": {
"model": "anthropic/claude-sonnet-4-20250514",
"api_key": "fake",
},
"routing_preferences": [
{
"name": "code_generation",
"description": "code programming debugging python function",
}
],
},
],
content_routing={
"enabled": True,
"classifier": "rule_based",
"default_model": "gpt-4o",
"confidence_threshold": 0.0,
},
model_affinity={
"enabled": True,
"ttl": 600,
"max_sessions": 1000,
"storage": "local",
**affinity_kwargs,
},
)
def test_router_initializes_affinity_router(self):
router = self._make_router()
assert router.model_affinity_router is not None
def test_first_request_pins_content_routed_model(self):
"""First request with session header should pin the content-routed model."""
router = self._make_router()
request_kwargs = {
"metadata": {
"headers": {"x-model-affinity": "session-abc"},
}
}
messages = [{"role": "user", "content": "write a python function to sort a list"}]
result = self._run(
router.async_pre_routing_hook(
model="gpt-4o",
request_kwargs=request_kwargs,
messages=messages,
)
)
# Content routing selected claude-sonnet; affinity should have pinned it
assert result is not None
assert result.model == "claude-sonnet"
decision = request_kwargs["metadata"].get("model_affinity_decision")
assert decision is not None
assert decision["status"] == "new"
assert decision["model"] == "claude-sonnet"
# Second request should be served from pin, bypassing content routing
request_kwargs2 = {
"metadata": {
"headers": {"x-model-affinity": "session-abc"},
}
}
messages2 = [{"role": "user", "content": "what is the weather today?"}]
result2 = self._run(
router.async_pre_routing_hook(
model="gpt-4o",
request_kwargs=request_kwargs2,
messages=messages2,
)
)
assert result2 is not None
assert result2.model == "claude-sonnet" # pinned, despite conversation prompt
decision2 = request_kwargs2["metadata"].get("model_affinity_decision")
assert decision2["status"] == "pinned"
def test_different_sessions_are_independent(self):
"""Two sessions with different IDs must not interfere."""
router = self._make_router()
# Session A: code prompt → claude-sonnet
kw_a = {"metadata": {"headers": {"x-model-affinity": "session-a"}}}
self._run(
router.async_pre_routing_hook(
model="gpt-4o",
request_kwargs=kw_a,
messages=[{"role": "user", "content": "write a python function"}],
)
)
# Session B: different ID → gets its own independent routing
kw_b = {"metadata": {"headers": {"x-model-affinity": "session-b"}}}
result_b = self._run(
router.async_pre_routing_hook(
model="gpt-4o",
request_kwargs=kw_b,
messages=[{"role": "user", "content": "write a python function"}],
)
)
# Both session-a and session-b exist independently
pinned_a = self._run(router.model_affinity_router.get_pinned_model("session-a"))
pinned_b = self._run(router.model_affinity_router.get_pinned_model("session-b"))
assert pinned_a is not None
assert pinned_b is not None
def test_no_session_header_bypasses_affinity(self):
"""Requests without X-Model-Affinity should route normally (no pin created)."""
router = self._make_router()
kw = {
"metadata": {
"headers": {}, # no affinity header
}
}
result = self._run(
router.async_pre_routing_hook(
model="gpt-4o",
request_kwargs=kw,
messages=[{"role": "user", "content": "write a python function"}],
)
)
# Content routing still runs normally
assert result is not None
assert result.model == "claude-sonnet"
# But no pin was created
assert "model_affinity_decision" not in kw["metadata"]
def test_affinity_disabled_leaves_router_none(self):
import litellm
router = litellm.Router(
model_list=[
{
"model_name": "gpt-4o",
"litellm_params": {"model": "openai/gpt-4o", "api_key": "fake"},
}
],
model_affinity={"enabled": False},
)
assert router.model_affinity_router is None
def test_affinity_without_config_leaves_router_none(self):
import litellm
router = litellm.Router(
model_list=[
{
"model_name": "gpt-4o",
"litellm_params": {"model": "openai/gpt-4o", "api_key": "fake"},
}
],
)
assert router.model_affinity_router is None
def test_specific_deployment_bypasses_affinity(self):
"""Affinity must not activate when specific_deployment=True."""
router = self._make_router()
# Pre-pin a session
self._run(router.model_affinity_router.pin_model("s1", "claude-sonnet"))
kw = {"metadata": {"headers": {"x-model-affinity": "s1"}}}
result = self._run(
router.async_pre_routing_hook(
model="gpt-4o",
request_kwargs=kw,
messages=[{"role": "user", "content": "hello"}],
specific_deployment=True,
)
)
# specific_deployment=True skips both affinity and content routing
assert result is None

View file

@ -493,6 +493,7 @@ export default function ModelInfoView({
<TabGroup>
<TabList className="mb-6">
<Tab>Overview</Tab>
<Tab>Routing Preferences</Tab>
<Tab>Raw JSON</Tab>
</TabList>
@ -1202,6 +1203,31 @@ export default function ModelInfoView({
</Card>
</TabPanel>
<TabPanel>
<RoutingPreferencesTab
modelData={modelData}
accessToken={accessToken}
onSave={async (preferences) => {
if (!accessToken) return;
setIsSaving(true);
try {
await modelPatchUpdateCall(accessToken, {
routing_preferences: preferences,
}, modelData.model_info?.id);
NotificationsManager.success("Routing preferences saved");
setLocalModelData((prev: any) => ({
...prev,
routing_preferences: preferences,
}));
} catch (e: any) {
NotificationsManager.fromBackend("Failed to save routing preferences: " + e?.message);
} finally {
setIsSaving(false);
}
}}
/>
</TabPanel>
<TabPanel>
<Card>
<pre className="bg-gray-100 p-4 rounded text-xs overflow-auto">{JSON.stringify(modelData, null, 2)}</pre>
@ -1269,3 +1295,110 @@ export default function ModelInfoView({
</div>
);
}
// ---------------------------------------------------------------------------
// RoutingPreferencesTab — inline sub-component
// ---------------------------------------------------------------------------
interface RoutingPreference {
name: string;
description: string;
}
interface RoutingPreferencesTabProps {
modelData: any;
accessToken: string | null;
onSave: (preferences: RoutingPreference[]) => Promise<void>;
}
function RoutingPreferencesTab({ modelData, accessToken, onSave }: RoutingPreferencesTabProps) {
const [preferences, setPreferences] = useState<RoutingPreference[]>(
modelData?.routing_preferences ?? []
);
const [saving, setSaving] = useState(false);
const add = () => setPreferences((prev) => [...prev, { name: "", description: "" }]);
const remove = (idx: number) =>
setPreferences((prev) => prev.filter((_, i) => i !== idx));
const update = (idx: number, field: keyof RoutingPreference, value: string) =>
setPreferences((prev) =>
prev.map((p, i) => (i === idx ? { ...p, [field]: value } : p))
);
const handleSave = async () => {
setSaving(true);
try {
await onSave(preferences);
} finally {
setSaving(false);
}
};
return (
<Card>
<div className="space-y-4">
<div>
<Title className="text-sm">Routing Preferences</Title>
<Text className="text-xs text-gray-500 mt-1">
Define content-based routing preferences for this model. When{" "}
<code className="bg-gray-100 px-1 rounded">content_routing</code> is enabled in
router settings, incoming prompts are classified against these descriptions.
</Text>
</div>
{preferences.length === 0 && (
<Text className="text-xs text-gray-400 italic">
No routing preferences configured. Click &quot;Add Preference&quot; to get started.
</Text>
)}
{preferences.map((pref, idx) => (
<div key={idx} className="border border-gray-200 rounded-lg p-3 space-y-2">
<div className="flex items-center justify-between">
<span className="text-xs font-medium text-gray-600">Preference {idx + 1}</span>
<Button
type="link"
danger
size="small"
onClick={() => remove(idx)}
>
Remove
</Button>
</div>
<div className="space-y-2">
<div>
<label className="block text-xs font-medium text-gray-700 mb-1">Name</label>
<TextInput
placeholder="e.g. code_generation"
value={pref.name}
onChange={(e) => update(idx, "name", e.target.value)}
/>
</div>
<div>
<label className="block text-xs font-medium text-gray-700 mb-1">
Description
</label>
<TextInput
placeholder="e.g. generating, debugging, and explaining code"
value={pref.description}
onChange={(e) => update(idx, "description", e.target.value)}
/>
</div>
</div>
</div>
))}
<div className="flex gap-2">
<TremorButton size="xs" variant="secondary" onClick={add}>
+ Add Preference
</TremorButton>
<TremorButton size="xs" loading={saving} onClick={handleSave}>
Save Preferences
</TremorButton>
</div>
</div>
</Card>
);
}

View file

@ -4361,6 +4361,68 @@ export const getRouterSettingsCall = async (accessToken: string) => {
}
};
export const testContentRoutingCall = async (
accessToken: string,
prompt: string,
messages?: Array<{ role: string; content: string }>
) => {
try {
const url = proxyBaseUrl
? `${proxyBaseUrl}/utils/content_route_test`
: `/utils/content_route_test`;
const body: Record<string, unknown> = messages ? { messages } : { prompt };
const response = await fetch(url, {
method: "POST",
headers: {
[globalLitellmHeaderName]: `Bearer ${accessToken}`,
"Content-Type": "application/json",
},
body: JSON.stringify(body),
});
if (!response.ok) {
const errorData = await response.json();
const errorMessage = deriveErrorMessage(errorData);
throw new Error(errorMessage);
}
return await response.json();
} catch (error) {
console.error("Failed to test content routing:", error);
throw error;
}
};
export const getContentRoutingPreferencesCall = async (accessToken: string) => {
try {
const url = proxyBaseUrl
? `${proxyBaseUrl}/router/content_routing/preferences`
: `/router/content_routing/preferences`;
const response = await fetch(url, {
method: "GET",
headers: {
[globalLitellmHeaderName]: `Bearer ${accessToken}`,
"Content-Type": "application/json",
},
});
if (!response.ok) {
const errorData = await response.json();
const errorMessage = deriveErrorMessage(errorData);
handleError(errorMessage);
throw new Error(errorMessage);
}
return await response.json();
} catch (error) {
console.error("Failed to get content routing preferences:", error);
throw error;
}
};
export const getCacheSettingsCall = async (accessToken: string) => {
try {
let url = proxyBaseUrl ? `${proxyBaseUrl}/cache/settings` : `/cache/settings`;

View file

@ -0,0 +1,463 @@
import { describe, it, expect, vi, beforeEach } from "vitest";
import { render, screen, waitFor, fireEvent } from "@testing-library/react";
import userEvent from "@testing-library/user-event";
import ContentRoutingConfiguration from "./ContentRoutingConfiguration";
import type { ContentRoutingConfig } from "./ContentRoutingConfiguration";
const defaultConfig: ContentRoutingConfig = {
enabled: false,
classifier: "rule_based",
};
const enabledConfig: ContentRoutingConfig = {
enabled: true,
classifier: "rule_based",
default_model: "",
confidence_threshold: 0.1,
};
const baseProps = {
config: defaultConfig,
onChange: vi.fn(),
accessToken: "test-token",
};
describe("ContentRoutingConfiguration", () => {
beforeEach(() => {
vi.clearAllMocks();
});
// ── Rendering ────────────────────────────────────────────────────────────────
it("renders the section heading", () => {
render(<ContentRoutingConfiguration {...baseProps} />);
expect(screen.getByText("Content-Aware Routing")).toBeInTheDocument();
});
it("renders the enable toggle", () => {
render(<ContentRoutingConfiguration {...baseProps} />);
expect(screen.getByRole("switch")).toBeInTheDocument();
});
it("does not render classifier options when disabled", () => {
render(<ContentRoutingConfiguration {...baseProps} />);
expect(screen.queryByText("Classifier")).not.toBeInTheDocument();
});
it("renders classifier and common fields when enabled", () => {
render(
<ContentRoutingConfiguration {...baseProps} config={enabledConfig} />
);
expect(screen.getByText("Classifier")).toBeInTheDocument();
expect(screen.getByPlaceholderText("e.g. gpt-4o")).toBeInTheDocument();
expect(screen.getByPlaceholderText("0.1")).toBeInTheDocument();
});
// ── Toggle ───────────────────────────────────────────────────────────────────
it("calls onChange with enabled=true when toggle is clicked", async () => {
const onChange = vi.fn();
const user = userEvent.setup();
render(
<ContentRoutingConfiguration
{...baseProps}
onChange={onChange}
config={defaultConfig}
/>
);
await user.click(screen.getByRole("switch"));
expect(onChange).toHaveBeenCalledWith(
expect.objectContaining({ enabled: true })
);
});
it("calls onChange with enabled=false when toggled off", async () => {
const onChange = vi.fn();
const user = userEvent.setup();
render(
<ContentRoutingConfiguration
{...baseProps}
onChange={onChange}
config={enabledConfig}
/>
);
await user.click(screen.getByRole("switch"));
expect(onChange).toHaveBeenCalledWith(
expect.objectContaining({ enabled: false })
);
});
// ── Classifier descriptions ───────────────────────────────────────────────────
it("shows rule_based description by default", () => {
render(
<ContentRoutingConfiguration {...baseProps} config={enabledConfig} />
);
expect(screen.getByText(/TF-IDF keyword matching/i)).toBeInTheDocument();
});
// ── Conditional fields: embedding_similarity ──────────────────────────────────
it("shows embedding model field when classifier is embedding_similarity", () => {
render(
<ContentRoutingConfiguration
{...baseProps}
config={{ ...enabledConfig, classifier: "embedding_similarity" }}
/>
);
expect(screen.getByText(/embedding model/i)).toBeInTheDocument();
expect(
screen.getByPlaceholderText("text-embedding-3-small")
).toBeInTheDocument();
});
it("does not show embedding model field for rule_based classifier", () => {
render(
<ContentRoutingConfiguration {...baseProps} config={enabledConfig} />
);
expect(screen.queryByText(/embedding model/i)).not.toBeInTheDocument();
});
// ── Conditional fields: external_model ────────────────────────────────────────
it("shows external classifier URL field when classifier is external_model", () => {
render(
<ContentRoutingConfiguration
{...baseProps}
config={{ ...enabledConfig, classifier: "external_model" }}
/>
);
expect(screen.getByText(/external classifier url/i)).toBeInTheDocument();
expect(
screen.getByPlaceholderText(/arch-router/i)
).toBeInTheDocument();
});
it("does not show external classifier URL for rule_based classifier", () => {
render(
<ContentRoutingConfiguration {...baseProps} config={enabledConfig} />
);
expect(
screen.queryByText(/external classifier url/i)
).not.toBeInTheDocument();
});
// ── Default model field ───────────────────────────────────────────────────────
it("calls onChange with updated default_model when edited", async () => {
const onChange = vi.fn();
const user = userEvent.setup();
render(
<ContentRoutingConfiguration
{...baseProps}
onChange={onChange}
config={enabledConfig}
/>
);
await user.type(screen.getByPlaceholderText("e.g. gpt-4o"), "g");
expect(onChange).toHaveBeenLastCalledWith(
expect.objectContaining({ default_model: "g" })
);
});
// ── Confidence threshold field ────────────────────────────────────────────────
it("renders confidence threshold with the configured value", () => {
render(
<ContentRoutingConfiguration
{...baseProps}
config={{ ...enabledConfig, confidence_threshold: 0.75 }}
/>
);
expect(screen.getByDisplayValue("0.75")).toBeInTheDocument();
});
it("calls onChange with updated confidence_threshold when valid number is entered", () => {
const onChange = vi.fn();
render(
<ContentRoutingConfiguration
{...baseProps}
onChange={onChange}
config={{ ...enabledConfig, confidence_threshold: 0.1 }}
/>
);
const input = screen.getByDisplayValue("0.1");
fireEvent.change(input, { target: { value: "0.5" } });
expect(onChange).toHaveBeenLastCalledWith(
expect.objectContaining({ confidence_threshold: 0.5 })
);
});
// ── Test panel ────────────────────────────────────────────────────────────────
it("shows 'Open tester' button when enabled", () => {
render(
<ContentRoutingConfiguration {...baseProps} config={enabledConfig} />
);
expect(
screen.getByRole("button", { name: /open tester/i })
).toBeInTheDocument();
});
it("toggles the test panel visibility when 'Open tester' is clicked", async () => {
const user = userEvent.setup();
render(
<ContentRoutingConfiguration {...baseProps} config={enabledConfig} />
);
const toggleBtn = screen.getByRole("button", { name: /open tester/i });
await user.click(toggleBtn);
expect(
screen.getByPlaceholderText(/write a python function/i)
).toBeInTheDocument();
expect(
screen.getByRole("button", { name: /hide/i })
).toBeInTheDocument();
});
it("hides the test panel again when 'Hide' is clicked", async () => {
const user = userEvent.setup();
render(
<ContentRoutingConfiguration {...baseProps} config={enabledConfig} />
);
await user.click(screen.getByRole("button", { name: /open tester/i }));
await user.click(screen.getByRole("button", { name: /hide/i }));
expect(
screen.queryByPlaceholderText(/write a python function/i)
).not.toBeInTheDocument();
});
it("disables 'Run classification' when the prompt is empty", async () => {
const user = userEvent.setup();
render(
<ContentRoutingConfiguration {...baseProps} config={enabledConfig} />
);
await user.click(screen.getByRole("button", { name: /open tester/i }));
expect(
screen.getByRole("button", { name: /run classification/i })
).toBeDisabled();
});
it("enables 'Run classification' after the user types a prompt", async () => {
const user = userEvent.setup();
render(
<ContentRoutingConfiguration {...baseProps} config={enabledConfig} />
);
await user.click(screen.getByRole("button", { name: /open tester/i }));
await user.type(
screen.getByPlaceholderText(/write a python function/i),
"hello"
);
expect(
screen.getByRole("button", { name: /run classification/i })
).not.toBeDisabled();
});
// ── runTest — successful response ─────────────────────────────────────────────
it("displays classification results after a successful API call", async () => {
const mockResult = {
matched_preference: "coding",
matched_model: "gpt-4o",
confidence: 0.9123,
classifier: "rule_based",
};
global.fetch = vi.fn().mockResolvedValueOnce({
ok: true,
json: async () => mockResult,
} as Response);
const user = userEvent.setup();
render(
<ContentRoutingConfiguration {...baseProps} config={enabledConfig} />
);
await user.click(screen.getByRole("button", { name: /open tester/i }));
await user.type(
screen.getByPlaceholderText(/write a python function/i),
"sort a list"
);
await user.click(screen.getByRole("button", { name: /run classification/i }));
await waitFor(() =>
expect(screen.getByText("coding")).toBeInTheDocument()
);
expect(screen.getByText("gpt-4o")).toBeInTheDocument();
expect(screen.getByText("0.9123")).toBeInTheDocument();
});
it("posts to /utils/content_route_test with the correct payload", async () => {
const mockFetch = vi.fn().mockResolvedValueOnce({
ok: true,
json: async () => ({
matched_preference: "coding",
matched_model: "gpt-4o",
confidence: 0.9,
classifier: "rule_based",
}),
} as Response);
global.fetch = mockFetch;
const user = userEvent.setup();
render(
<ContentRoutingConfiguration {...baseProps} config={enabledConfig} />
);
await user.click(screen.getByRole("button", { name: /open tester/i }));
await user.type(
screen.getByPlaceholderText(/write a python function/i),
"my prompt"
);
await user.click(screen.getByRole("button", { name: /run classification/i }));
await waitFor(() => expect(mockFetch).toHaveBeenCalledOnce());
const [url, init] = mockFetch.mock.calls[0] as [string, RequestInit];
expect(url).toBe("/utils/content_route_test");
expect(init.method).toBe("POST");
expect(JSON.parse(init.body as string)).toEqual({ prompt: "my prompt" });
expect((init.headers as Record<string, string>)["Authorization"]).toBe(
"Bearer test-token"
);
});
// ── runTest — error response ───────────────────────────────────────────────────
it("displays an error message when the API call fails", async () => {
global.fetch = vi.fn().mockResolvedValueOnce({
ok: false,
json: async () => ({ detail: "Router not configured" }),
} as Response);
const user = userEvent.setup();
render(
<ContentRoutingConfiguration {...baseProps} config={enabledConfig} />
);
await user.click(screen.getByRole("button", { name: /open tester/i }));
await user.type(
screen.getByPlaceholderText(/write a python function/i),
"something"
);
await user.click(screen.getByRole("button", { name: /run classification/i }));
await waitFor(() =>
expect(screen.getByText("Router not configured")).toBeInTheDocument()
);
});
it("displays a fallback error message when response has no detail field", async () => {
global.fetch = vi.fn().mockResolvedValueOnce({
ok: false,
status: 422,
json: async () => ({}),
} as Response);
const user = userEvent.setup();
render(
<ContentRoutingConfiguration {...baseProps} config={enabledConfig} />
);
await user.click(screen.getByRole("button", { name: /open tester/i }));
await user.type(
screen.getByPlaceholderText(/write a python function/i),
"something"
);
await user.click(screen.getByRole("button", { name: /run classification/i }));
await waitFor(() =>
expect(screen.getByText(/HTTP 422/i)).toBeInTheDocument()
);
});
it("displays an error message when fetch throws a network error", async () => {
global.fetch = vi.fn().mockRejectedValueOnce(new Error("Network failure"));
const user = userEvent.setup();
render(
<ContentRoutingConfiguration {...baseProps} config={enabledConfig} />
);
await user.click(screen.getByRole("button", { name: /open tester/i }));
await user.type(
screen.getByPlaceholderText(/write a python function/i),
"something"
);
await user.click(screen.getByRole("button", { name: /run classification/i }));
await waitFor(() =>
expect(screen.getByText("Network failure")).toBeInTheDocument()
);
});
// ── all_scores details block ──────────────────────────────────────────────────
it("renders all_scores when present in the result", async () => {
global.fetch = vi.fn().mockResolvedValueOnce({
ok: true,
json: async () => ({
matched_preference: "coding",
matched_model: "gpt-4o",
confidence: 0.9,
classifier: "rule_based",
all_scores: { coding: 0.9, writing: 0.1 },
}),
} as Response);
const user = userEvent.setup();
render(
<ContentRoutingConfiguration {...baseProps} config={enabledConfig} />
);
await user.click(screen.getByRole("button", { name: /open tester/i }));
await user.type(
screen.getByPlaceholderText(/write a python function/i),
"sort a list"
);
await user.click(screen.getByRole("button", { name: /run classification/i }));
await waitFor(() =>
expect(screen.getByText(/all scores/i)).toBeInTheDocument()
);
});
// ── accessToken guard ─────────────────────────────────────────────────────────
it("does not fetch when accessToken is null", async () => {
const mockFetch = vi.fn();
global.fetch = mockFetch;
const user = userEvent.setup();
render(
<ContentRoutingConfiguration
config={enabledConfig}
onChange={vi.fn()}
accessToken={null}
/>
);
await user.click(screen.getByRole("button", { name: /open tester/i }));
await user.type(
screen.getByPlaceholderText(/write a python function/i),
"hello"
);
await user.click(screen.getByRole("button", { name: /run classification/i }));
expect(mockFetch).not.toHaveBeenCalled();
});
});

View file

@ -0,0 +1,291 @@
import React, { useState } from "react";
import { Switch, TextInput, Select, SelectItem, Button, Badge } from "@tremor/react";
export interface ContentRoutingConfig {
enabled: boolean;
classifier: "rule_based" | "embedding_similarity" | "external_model";
default_model?: string;
confidence_threshold?: number;
embedding_model?: string;
external_classifier_url?: string;
}
interface ContentRouteTestResult {
matched_preference: string;
matched_model: string;
confidence: number;
classifier: string;
all_scores?: Record<string, number>;
}
interface ContentRoutingConfigurationProps {
config: ContentRoutingConfig;
onChange: (config: ContentRoutingConfig) => void;
accessToken: string | null;
}
const DEFAULT_CONFIG: ContentRoutingConfig = {
enabled: false,
classifier: "rule_based",
default_model: "",
confidence_threshold: 0.1,
};
const CLASSIFIER_DESCRIPTIONS: Record<string, string> = {
rule_based:
"TF-IDF keyword matching — zero latency, no extra API calls. Best starting point.",
embedding_similarity:
"Semantic embedding similarity using litellm.aembedding(). More flexible but adds one embedding call per request.",
external_model:
"Delegates to an external HTTP classifier (e.g. Arch-Router). Fully pluggable.",
};
const ContentRoutingConfiguration: React.FC<ContentRoutingConfigurationProps> = ({
config,
onChange,
accessToken,
}) => {
const [testPrompt, setTestPrompt] = useState("");
const [testResult, setTestResult] = useState<ContentRouteTestResult | null>(null);
const [testLoading, setTestLoading] = useState(false);
const [testError, setTestError] = useState<string | null>(null);
const [showTestPanel, setShowTestPanel] = useState(false);
const effective = { ...DEFAULT_CONFIG, ...config };
const update = (patch: Partial<ContentRoutingConfig>) =>
onChange({ ...effective, ...patch });
const runTest = async () => {
if (!testPrompt.trim() || !accessToken) return;
setTestLoading(true);
setTestError(null);
setTestResult(null);
try {
const resp = await fetch("/utils/content_route_test", {
method: "POST",
headers: {
"Content-Type": "application/json",
Authorization: `Bearer ${accessToken}`,
},
body: JSON.stringify({ prompt: testPrompt }),
});
if (!resp.ok) {
const data = await resp.json().catch(() => ({}));
throw new Error(data?.detail ?? `HTTP ${resp.status}`);
}
setTestResult(await resp.json());
} catch (e: any) {
setTestError(e?.message ?? "Unknown error");
} finally {
setTestLoading(false);
}
};
return (
<div className="space-y-6">
{/* Header */}
<div className="max-w-3xl">
<h3 className="text-sm font-medium text-gray-900">Content-Aware Routing</h3>
<p className="text-xs text-gray-500 mt-1">
Classify prompt content and route to the best-matched model based on{" "}
<code className="bg-gray-100 px-1 rounded">routing_preferences</code> configured
per model. Runs before infrastructure routing (latency, cost, etc.).
</p>
</div>
{/* Enable toggle */}
<div className="flex items-center gap-3">
<Switch
id="content-routing-enabled"
checked={effective.enabled}
onChange={(v) => update({ enabled: v })}
/>
<label htmlFor="content-routing-enabled" className="text-sm text-gray-700 cursor-pointer">
Enable content-aware routing
</label>
</div>
{effective.enabled && (
<div className="space-y-6 pl-0">
{/* Classifier */}
<div className="space-y-2 max-w-md">
<label className="block text-xs font-medium text-gray-700 uppercase tracking-wide">
Classifier
</label>
<Select
value={effective.classifier}
onValueChange={(v) =>
update({ classifier: v as ContentRoutingConfig["classifier"] })
}
>
<SelectItem value="rule_based">Rule-Based (TF-IDF)</SelectItem>
<SelectItem value="embedding_similarity">Embedding Similarity</SelectItem>
<SelectItem value="external_model">External Model</SelectItem>
</Select>
<p className="text-xs text-gray-500">
{CLASSIFIER_DESCRIPTIONS[effective.classifier]}
</p>
</div>
{/* Common fields */}
<div className="grid grid-cols-1 gap-6 lg:grid-cols-2">
<div className="space-y-2">
<label className="block text-xs font-medium text-gray-700 uppercase tracking-wide">
Default Model
</label>
<p className="text-xs text-gray-500">
Fallback when no preference matches above the confidence threshold.
</p>
<TextInput
placeholder="e.g. gpt-4o"
value={effective.default_model ?? ""}
onChange={(e) => update({ default_model: e.target.value })}
/>
</div>
<div className="space-y-2">
<label className="block text-xs font-medium text-gray-700 uppercase tracking-wide">
Confidence Threshold
</label>
<p className="text-xs text-gray-500">
Minimum score (0.01.0) required to route to a matched model.
</p>
<TextInput
placeholder="0.1"
value={(effective.confidence_threshold ?? 0.1).toString()}
onChange={(e) => {
const n = parseFloat(e.target.value);
if (!isNaN(n)) update({ confidence_threshold: n });
}}
/>
</div>
</div>
{/* Embedding-specific */}
{effective.classifier === "embedding_similarity" && (
<div className="space-y-2 max-w-md">
<label className="block text-xs font-medium text-gray-700 uppercase tracking-wide">
Embedding Model
</label>
<p className="text-xs text-gray-500">
Model used to embed preference descriptions and incoming prompts.
</p>
<TextInput
placeholder="text-embedding-3-small"
value={effective.embedding_model ?? ""}
onChange={(e) => update({ embedding_model: e.target.value })}
/>
</div>
)}
{/* External model-specific */}
{effective.classifier === "external_model" && (
<div className="space-y-2 max-w-lg">
<label className="block text-xs font-medium text-gray-700 uppercase tracking-wide">
External Classifier URL
</label>
<p className="text-xs text-gray-500">
POST endpoint that accepts{" "}
<code className="bg-gray-100 px-1 rounded">
{"{ prompt, preferences }"}
</code>{" "}
and returns{" "}
<code className="bg-gray-100 px-1 rounded">
{"{ matched_preference, model, confidence }"}
</code>
. Compatible with Arch-Router.
</p>
<TextInput
placeholder="http://arch-router-host/classify"
value={effective.external_classifier_url ?? ""}
onChange={(e) => update({ external_classifier_url: e.target.value })}
/>
</div>
)}
{/* Test routing panel */}
<div className="border border-gray-200 rounded-lg p-4 space-y-3 max-w-2xl">
<div className="flex items-center justify-between">
<div>
<p className="text-sm font-medium text-gray-900">Test Content Routing</p>
<p className="text-xs text-gray-500">
Classify a prompt without making an LLM call.
</p>
</div>
<Button
size="xs"
variant="secondary"
onClick={() => setShowTestPanel((v) => !v)}
>
{showTestPanel ? "Hide" : "Open tester"}
</Button>
</div>
{showTestPanel && (
<div className="space-y-3">
<TextInput
placeholder="e.g. write a Python function to sort a list"
value={testPrompt}
onChange={(e) => setTestPrompt(e.target.value)}
/>
<Button
size="xs"
loading={testLoading}
disabled={!testPrompt.trim()}
onClick={runTest}
>
Run classification
</Button>
{testError && (
<p className="text-xs text-red-600">{testError}</p>
)}
{testResult && (
<div className="bg-gray-50 rounded p-3 space-y-2 text-xs">
<div className="flex items-center gap-2">
<span className="text-gray-500">Matched preference:</span>
<Badge color="blue" size="xs">{testResult.matched_preference}</Badge>
</div>
<div className="flex items-center gap-2">
<span className="text-gray-500">Routed to:</span>
<span className="font-mono font-medium">{testResult.matched_model}</span>
</div>
<div className="flex items-center gap-2">
<span className="text-gray-500">Confidence:</span>
<span className="font-mono">{testResult.confidence.toFixed(4)}</span>
</div>
<div className="flex items-center gap-2">
<span className="text-gray-500">Classifier:</span>
<span className="font-mono">{testResult.classifier}</span>
</div>
{testResult.all_scores && Object.keys(testResult.all_scores).length > 0 && (
<details className="mt-1">
<summary className="cursor-pointer text-gray-500 hover:text-gray-700">
All scores
</summary>
<div className="mt-1 space-y-1 pl-2">
{Object.entries(testResult.all_scores)
.sort(([, a], [, b]) => b - a)
.map(([key, score]) => (
<div key={key} className="flex justify-between font-mono">
<span className="text-gray-600">{key}</span>
<span>{score.toFixed(4)}</span>
</div>
))}
</div>
</details>
)}
</div>
)}
</div>
)}
</div>
</div>
)}
</div>
);
};
export default ContentRoutingConfiguration;

View file

@ -104,7 +104,7 @@ describe("RouterSettingsForm", () => {
const user = userEvent.setup();
render(<RouterSettingsForm {...baseProps} onChange={onChange} />);
await user.click(screen.getByRole("switch"));
await user.click(screen.getAllByRole("switch")[0]);
expect(onChange).toHaveBeenCalledWith(
expect.objectContaining({ enableTagFiltering: true })

View file

@ -1,4 +1,5 @@
import React from "react";
import ContentRoutingConfiguration, { ContentRoutingConfig } from "./ContentRoutingConfiguration";
import LatencyBasedConfiguration from "./LatencyBasedConfiguration";
import ReliabilityRetriesSection from "./ReliabilityRetriesSection";
import RoutingStrategySelector from "./RoutingStrategySelector";
@ -16,6 +17,7 @@ interface RouterSettingsFormProps {
routerFieldsMetadata: { [key: string]: any };
availableRoutingStrategies: string[];
routingStrategyDescriptions: { [key: string]: string };
accessToken?: string | null;
}
const RouterSettingsForm: React.FC<RouterSettingsFormProps> = ({
@ -24,6 +26,7 @@ const RouterSettingsForm: React.FC<RouterSettingsFormProps> = ({
routerFieldsMetadata,
availableRoutingStrategies,
routingStrategyDescriptions,
accessToken,
}) => {
const handleStrategyChange = (strategy: string) => {
onChange({
@ -39,6 +42,16 @@ const RouterSettingsForm: React.FC<RouterSettingsFormProps> = ({
});
};
const handleContentRoutingChange = (config: ContentRoutingConfig) => {
onChange({
...value,
routerSettings: {
...value.routerSettings,
content_routing: config,
},
});
};
return (
<div className="w-full space-y-8 py-2">
{/* Routing Settings Section */}
@ -75,6 +88,16 @@ const RouterSettingsForm: React.FC<RouterSettingsFormProps> = ({
<LatencyBasedConfiguration routingStrategyArgs={value.routerSettings["routing_strategy_args"]} />
)}
{/* Content-Aware Routing */}
<ContentRoutingConfiguration
config={value.routerSettings.content_routing ?? { enabled: false, classifier: "rule_based" }}
onChange={handleContentRoutingChange}
accessToken={accessToken ?? null}
/>
{/* Divider */}
<div className="border-t border-gray-200" />
{/* Other Settings */}
<ReliabilityRetriesSection routerSettings={value.routerSettings} routerFieldsMetadata={routerFieldsMetadata} />
</div>