mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-12 23:01:41 +00:00
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:
parent
97f722f558
commit
8ccecaf9bd
19 changed files with 3188 additions and 1 deletions
|
|
@ -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)
|
||||
|
|
|
|||
193
litellm/proxy/management_endpoints/content_routing_endpoints.py
Normal file
193
litellm/proxy/management_endpoints/content_routing_endpoints.py
Normal 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,
|
||||
)
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
#########################################################
|
||||
|
|
|
|||
0
litellm/router_strategy/content_aware_router/__init__.py
Normal file
0
litellm/router_strategy/content_aware_router/__init__.py
Normal 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,
|
||||
)
|
||||
273
litellm/router_strategy/content_aware_router/utils.py
Normal file
273
litellm/router_strategy/content_aware_router/utils.py
Normal 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 (Robertson–Sparck 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
|
||||
|
|
@ -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}"
|
||||
)
|
||||
|
|
@ -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",
|
||||
),
|
||||
]
|
||||
|
|
|
|||
|
|
@ -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=())
|
||||
|
||||
|
|
|
|||
639
tests/test_litellm/router_strategy/test_content_aware_router.py
Normal file
639
tests/test_litellm/router_strategy/test_content_aware_router.py
Normal 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
|
||||
321
tests/test_litellm/router_strategy/test_model_affinity_router.py
Normal file
321
tests/test_litellm/router_strategy/test_model_affinity_router.py
Normal 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
|
||||
|
|
@ -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 "Add Preference" 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>
|
||||
);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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`;
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
});
|
||||
});
|
||||
|
|
@ -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.0–1.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;
|
||||
|
|
@ -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 })
|
||||
|
|
|
|||
|
|
@ -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>
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue