mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-17 04:34:25 +00:00
feat(router): add llm_classifier_router strategy
A new auto-router strategy that uses a small LLM to classify prompt
complexity and route to a pre-configured tier model. Sits alongside
the existing complexity_router (rule-based) and quality_router,
filling the gap between keyword-rule accuracy and full embedding
semantic matching.
The classifier LLM (default ollama/qwen2.5:0.5b) emits a one-word
tier label (SIMPLE/COMPLEX) which the router maps to a backend
model. Results are cached by prompt hash with TTL. On LLM failure
the router falls back to the rule-based ComplexityRouter, then to
a configured fallback_tier (default SIMPLE for cost).
Configuration:
model_list:
- model_name: smart-router
litellm_params:
model: auto_router/llm_classifier_router
llm_classifier_router_config:
classifier_model: ollama/qwen2.5:0.5b
tiers: {SIMPLE: cheap, COMPLEX: powerful}
Includes 43 unit tests and 11 integration tests covering
detection, init, fallback paths, hook dispatch, and end-to-end
Router init. Zero new Python dependencies.
This commit is contained in:
parent
49ca04d8c3
commit
933ffa8fe5
9 changed files with 1078 additions and 0 deletions
|
|
@ -208,6 +208,9 @@ if TYPE_CHECKING:
|
|||
from litellm.router_strategy.quality_router.quality_router import (
|
||||
QualityRouter,
|
||||
)
|
||||
from litellm.router_strategy.llm_classifier_router.llm_classifier_router import (
|
||||
LLMClassifierRouter,
|
||||
)
|
||||
from litellm.responses.streaming_iterator import (
|
||||
BaseResponsesAPIStreamingIterator,
|
||||
)
|
||||
|
|
@ -225,6 +228,7 @@ else:
|
|||
ComplexityRouter = Any
|
||||
AdaptiveRouter = Any
|
||||
QualityRouter = Any
|
||||
LLMClassifierRouter = Any
|
||||
PreRoutingHookResponse = Any
|
||||
|
||||
|
||||
|
|
@ -489,6 +493,7 @@ class Router:
|
|||
self.complexity_routers: Dict[str, "ComplexityRouter"] = {}
|
||||
self.adaptive_routers: Dict[str, "AdaptiveRouter"] = {}
|
||||
self.quality_routers: Dict[str, "QualityRouter"] = {}
|
||||
self.llm_classifier_routers: Dict[str, "LLMClassifierRouter"] = {}
|
||||
|
||||
# Initialize model_group_alias early since it's used in set_model_list
|
||||
self.model_group_alias: Dict[str, Union[str, RouterModelGroupAliasItem]] = (
|
||||
|
|
@ -7972,6 +7977,8 @@ class Router:
|
|||
return False # This is handled by adaptive_router
|
||||
if litellm_params.model.startswith("auto_router/quality_router"):
|
||||
return False # This is handled by quality_router
|
||||
if litellm_params.model.startswith("auto_router/llm_classifier_router"):
|
||||
return False # This is handled by llm_classifier_router
|
||||
if litellm_params.model.startswith("auto_router/"):
|
||||
return True
|
||||
return False
|
||||
|
|
@ -8078,6 +8085,47 @@ class Router:
|
|||
)
|
||||
self.complexity_routers[deployment.model_name] = complexity_router
|
||||
|
||||
def _is_llm_classifier_router_deployment(
|
||||
self, litellm_params: LiteLLM_Params
|
||||
) -> bool:
|
||||
"""True when this deployment opts in via the `auto_router/llm_classifier_router` model prefix."""
|
||||
return litellm_params.model.startswith("auto_router/llm_classifier_router")
|
||||
|
||||
def init_llm_classifier_router_deployment(self, deployment: Deployment) -> None:
|
||||
from litellm.router_strategy.llm_classifier_router.llm_classifier_router import (
|
||||
LLMClassifierRouter,
|
||||
)
|
||||
|
||||
config: Optional[dict] = (
|
||||
deployment.litellm_params.llm_classifier_router_config
|
||||
)
|
||||
default_model: Optional[str] = (
|
||||
deployment.litellm_params.llm_classifier_router_default_model
|
||||
)
|
||||
|
||||
if default_model is None and config:
|
||||
tiers = config.get("tiers") or {}
|
||||
default_model = tiers.get("COMPLEX") or tiers.get("SIMPLE")
|
||||
|
||||
if default_model is None:
|
||||
raise ValueError(
|
||||
"llm_classifier_router_default_model is required for llm-classifier-router "
|
||||
"deployments. Set it in litellm_params or configure tiers in "
|
||||
"llm_classifier_router_config."
|
||||
)
|
||||
|
||||
llm_classifier_router: LLMClassifierRouter = LLMClassifierRouter(
|
||||
model_name=deployment.model_name,
|
||||
default_model=default_model,
|
||||
litellm_router_instance=self,
|
||||
llm_classifier_router_config=config,
|
||||
)
|
||||
if deployment.model_name in self.llm_classifier_routers:
|
||||
raise ValueError(
|
||||
f"LLM-classifier-router deployment {deployment.model_name} already exists. Please use a different model name."
|
||||
)
|
||||
self.llm_classifier_routers[deployment.model_name] = llm_classifier_router
|
||||
|
||||
def _is_adaptive_router_deployment(self, litellm_params: LiteLLM_Params) -> bool:
|
||||
"""True when this deployment opts in via the `auto_router/adaptive_router` model prefix."""
|
||||
return litellm_params.model.startswith("auto_router/adaptive_router")
|
||||
|
|
@ -8319,6 +8367,7 @@ class Router:
|
|||
self.quality_routers = {}
|
||||
self.complexity_routers = {}
|
||||
self.auto_routers = {}
|
||||
self.llm_classifier_routers = {}
|
||||
self._invalidate_model_group_info_cache()
|
||||
self._invalidate_access_groups_cache()
|
||||
# we add api_base/api_key each model so load balancing between azure/gpt on api_base1 and api_base2 works
|
||||
|
|
@ -8507,6 +8556,14 @@ class Router:
|
|||
if self._is_quality_router_deployment(litellm_params=deployment.litellm_params):
|
||||
self.init_quality_router_deployment(deployment=deployment)
|
||||
|
||||
#########################################################
|
||||
# Check if this is an llm-classifier-router deployment
|
||||
#########################################################
|
||||
if self._is_llm_classifier_router_deployment(
|
||||
litellm_params=deployment.litellm_params
|
||||
):
|
||||
self.init_llm_classifier_router_deployment(deployment=deployment)
|
||||
|
||||
return deployment
|
||||
|
||||
def _initialize_deployment_for_pass_through(
|
||||
|
|
@ -11353,6 +11410,18 @@ class Router:
|
|||
specific_deployment=specific_deployment,
|
||||
)
|
||||
|
||||
#########################################################
|
||||
# Check if any llm-classifier-router should be used
|
||||
#########################################################
|
||||
if model in self.llm_classifier_routers:
|
||||
return await self.llm_classifier_routers[model].async_pre_routing_hook(
|
||||
model=model,
|
||||
request_kwargs=request_kwargs,
|
||||
messages=messages,
|
||||
input=input,
|
||||
specific_deployment=specific_deployment,
|
||||
)
|
||||
|
||||
return None
|
||||
|
||||
def get_available_deployment(
|
||||
|
|
|
|||
112
litellm/router_strategy/llm_classifier_router/README.md
Normal file
112
litellm/router_strategy/llm_classifier_router/README.md
Normal file
|
|
@ -0,0 +1,112 @@
|
|||
# LLM Classifier Router
|
||||
|
||||
A routing strategy that uses a small LLM to classify prompt complexity and route to a pre-configured tier model. Useful for cost optimization when you want semantic accuracy beyond keyword rules but don't have the budget for a frontier model to do the classification.
|
||||
|
||||
## How It Works
|
||||
|
||||
A lightweight classifier LLM (e.g. `ollama/qwen2.5:0.5b` running locally, or any cheap hosted model) reads each incoming prompt and outputs a one-word tier label. The router maps that label to a configured backend model. Results are cached by prompt hash, so repeat queries are free.
|
||||
|
||||
```
|
||||
prompt ──▶ classifier LLM ──▶ "SIMPLE" ──▶ gpt-4o-mini
|
||||
"COMPLEX" ──▶ gpt-4o
|
||||
```
|
||||
|
||||
## Why Not Just Use complexity_router?
|
||||
|
||||
The rule-based `complexity_router` is faster (<1ms, no network) but only matches keyword signals. It misses semantically complex prompts that don't contain "step by step" or code keywords. The LLM classifier handles those cases at the cost of one extra small-model call (~100-300ms for a 0.5B model running locally).
|
||||
|
||||
## Configuration
|
||||
|
||||
### Minimal
|
||||
|
||||
```yaml
|
||||
model_list:
|
||||
- model_name: cheap-model
|
||||
litellm_params:
|
||||
model: gpt-4o-mini
|
||||
|
||||
- model_name: powerful-model
|
||||
litellm_params:
|
||||
model: gpt-4o
|
||||
|
||||
- model_name: smart-router
|
||||
litellm_params:
|
||||
model: auto_router/llm_classifier_router
|
||||
llm_classifier_router_config:
|
||||
classifier_model: ollama/qwen2.5:0.5b
|
||||
tiers:
|
||||
SIMPLE: cheap-model
|
||||
COMPLEX: powerful-model
|
||||
```
|
||||
|
||||
### Full
|
||||
|
||||
```yaml
|
||||
- model_name: smart-router
|
||||
litellm_params:
|
||||
model: auto_router/llm_classifier_router
|
||||
llm_classifier_router_config:
|
||||
classifier_model: ollama/qwen2.5:0.5b
|
||||
tiers:
|
||||
SIMPLE: gpt-4o-mini
|
||||
COMPLEX: gpt-4o
|
||||
classifier_timeout: 3.0
|
||||
classifier_temperature: 0.0
|
||||
classifier_max_tokens: 10
|
||||
enable_cache: true
|
||||
cache_ttl_seconds: 300
|
||||
fallback_tier: SIMPLE
|
||||
fallback_to_complexity_router: true
|
||||
```
|
||||
|
||||
## Parameters
|
||||
|
||||
| Field | Default | Description |
|
||||
|---|---|---|
|
||||
| `classifier_model` | `ollama/qwen2.5:0.5b` | Any litellm-supported model. Smaller = faster + cheaper. |
|
||||
| `tiers` | `{SIMPLE: gpt-4o-mini, COMPLEX: gpt-4o}` | Map of tier name to backend `model_name`. |
|
||||
| `classifier_system_prompt` | Built-in 2-tier prompt | Override to use a 4-tier prompt or custom instructions. |
|
||||
| `classifier_timeout` | `3.0` | Seconds before falling back. |
|
||||
| `classifier_temperature` | `0.0` | Keep at 0 for deterministic routing. |
|
||||
| `classifier_max_tokens` | `10` | Output is a single word; no need for more. |
|
||||
| `classifier_max_input_chars` | `2000` | Truncates very long prompts before classification. |
|
||||
| `enable_cache` | `True` | Cache by prompt hash to avoid repeat LLM calls. |
|
||||
| `cache_ttl_seconds` | `300` | Cache entry lifetime. |
|
||||
| `cache_max_size` | `1000` | Simple dict; on overflow, store is cleared. |
|
||||
| `fallback_tier` | `SIMPLE` | Last-resort tier when LLM and rule fallback both fail. |
|
||||
| `fallback_to_complexity_router` | `True` | Use rule-based ComplexityRouter as intermediate fallback. |
|
||||
|
||||
## Custom Tier Prompts
|
||||
|
||||
The default prompt asks for a 2-tier answer (SIMPLE/COMPLEX). For 4-tier classification, supply your own prompt and tier map. Note: 4-tier accuracy is poor on sub-1B models. Use 3B+ for that.
|
||||
|
||||
```yaml
|
||||
llm_classifier_router_config:
|
||||
classifier_model: ollama/qwen2.5:3b
|
||||
classifier_system_prompt: |
|
||||
Classify the request. Reply with EXACTLY one word: SIMPLE, MEDIUM, COMPLEX, or REASONING.
|
||||
SIMPLE: greetings, lookups
|
||||
MEDIUM: standard questions
|
||||
COMPLEX: multi-step, code
|
||||
REASONING: chain-of-thought, analysis
|
||||
tiers:
|
||||
SIMPLE: gpt-4o-mini
|
||||
MEDIUM: gpt-4o-mini
|
||||
COMPLEX: gpt-4o
|
||||
REASONING: o1-mini
|
||||
```
|
||||
|
||||
## Fallback Behavior
|
||||
|
||||
1. **Cache hit** — return cached tier immediately
|
||||
2. **LLM call** — call `classifier_model` with `classifier_timeout` deadline
|
||||
3. **ComplexityRouter rule fallback** — if LLM call fails/times out, use the rule-based ComplexityRouter
|
||||
4. **`fallback_tier`** — if both above fail, route to the cheapest configured tier
|
||||
|
||||
Each fallback step is logged with `verbose_router_logger` and the decision method is exposed in `request_kwargs["metadata"]["llm_classifier_router_method"]` for debugging.
|
||||
|
||||
## Notes
|
||||
|
||||
- The classifier is called via `litellm.acompletion` directly, bypassing the Router, to avoid recursion if the user happens to name `classifier_model` after one of their own deployments.
|
||||
- No new Python dependencies. All inference goes through litellm's existing provider routing.
|
||||
- For local Ollama models, the `ollama/` prefix is required (e.g. `ollama/qwen2.5:0.5b`).
|
||||
10
litellm/router_strategy/llm_classifier_router/__init__.py
Normal file
10
litellm/router_strategy/llm_classifier_router/__init__.py
Normal file
|
|
@ -0,0 +1,10 @@
|
|||
"""LLM-based prompt complexity classifier router."""
|
||||
|
||||
from litellm.router_strategy.llm_classifier_router.config import (
|
||||
LLMClassifierRouterConfig,
|
||||
)
|
||||
from litellm.router_strategy.llm_classifier_router.llm_classifier_router import (
|
||||
LLMClassifierRouter,
|
||||
)
|
||||
|
||||
__all__ = ["LLMClassifierRouter", "LLMClassifierRouterConfig"]
|
||||
27
litellm/router_strategy/llm_classifier_router/config.py
Normal file
27
litellm/router_strategy/llm_classifier_router/config.py
Normal file
|
|
@ -0,0 +1,27 @@
|
|||
"""Configuration for the LLM Classifier Router."""
|
||||
|
||||
from typing import Dict, Optional
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
|
||||
|
||||
class LLMClassifierRouterConfig(BaseModel):
|
||||
classifier_model: str = "ollama/qwen2.5:0.5b"
|
||||
|
||||
tiers: Dict[str, str] = Field(default_factory=lambda: {"SIMPLE": "gpt-4o-mini", "COMPLEX": "gpt-4o"})
|
||||
|
||||
classifier_system_prompt: Optional[str] = None
|
||||
|
||||
classifier_timeout: float = 3.0
|
||||
classifier_temperature: float = 0.0
|
||||
classifier_max_tokens: int = 10
|
||||
classifier_max_input_chars: int = 2000
|
||||
|
||||
enable_cache: bool = True
|
||||
cache_ttl_seconds: int = 300
|
||||
cache_max_size: int = 1000
|
||||
|
||||
fallback_tier: str = "SIMPLE"
|
||||
fallback_to_complexity_router: bool = True
|
||||
|
||||
model_config = ConfigDict(extra="allow")
|
||||
|
|
@ -0,0 +1,228 @@
|
|||
"""LLM-based prompt complexity classifier router.
|
||||
|
||||
Uses a small LLM (e.g. ollama/qwen2.5:0.5b) to classify prompt complexity and
|
||||
route to a pre-configured tier model. Falls back to a rule-based
|
||||
ComplexityRouter on timeout or error.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import re
|
||||
from time import monotonic
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
|
||||
|
||||
from litellm._logging import verbose_router_logger
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
|
||||
from .config import LLMClassifierRouterConfig
|
||||
from .prompt_templates import TWO_TIER_SYSTEM_PROMPT
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.router import Router
|
||||
from litellm.types.router import PreRoutingHookResponse
|
||||
else:
|
||||
Router = Any
|
||||
PreRoutingHookResponse = Any
|
||||
|
||||
|
||||
class _TierCache:
|
||||
"""Simple TTL cache. No LRU eviction — on overflow, clear and keep going."""
|
||||
|
||||
def __init__(self, ttl_seconds: float, max_size: int) -> None:
|
||||
self._store: Dict[str, Tuple[str, float]] = {}
|
||||
self._ttl = ttl_seconds
|
||||
self._max_size = max_size
|
||||
|
||||
def get(self, key: str) -> Optional[str]:
|
||||
item = self._store.get(key)
|
||||
if item is None:
|
||||
return None
|
||||
tier, expires_at = item
|
||||
if monotonic() > expires_at:
|
||||
del self._store[key]
|
||||
return None
|
||||
return tier
|
||||
|
||||
def set(self, key: str, tier: str) -> None:
|
||||
if len(self._store) >= self._max_size:
|
||||
self._store.clear()
|
||||
self._store[key] = (tier, monotonic() + self._ttl)
|
||||
|
||||
|
||||
def _parse_tier(raw: str, valid_tiers: Tuple[str, ...]) -> str:
|
||||
"""Parse a tier string from the classifier LLM's raw output.
|
||||
|
||||
Strategy: exact fullmatch first (avoids "a SIMPLE example" false positive),
|
||||
then substring containment as a fallback for outputs like "SIMPLE.".
|
||||
"""
|
||||
cleaned = raw.strip().upper()
|
||||
if valid_tiers:
|
||||
if re.fullmatch(r"|".join(re.escape(t) for t in valid_tiers), cleaned):
|
||||
return cleaned
|
||||
for tier in valid_tiers:
|
||||
if tier in cleaned:
|
||||
return tier
|
||||
raise ValueError(f"Could not parse tier from LLM output: {raw!r}")
|
||||
|
||||
|
||||
def _cache_key(prompt: str) -> str:
|
||||
return hashlib.sha256(prompt.encode()).hexdigest()[:16]
|
||||
|
||||
|
||||
class LLMClassifierRouter(CustomLogger):
|
||||
"""LLM-based prompt complexity classifier router.
|
||||
|
||||
Classifies requests via a small LLM and routes them to one of the
|
||||
configured tier models. Falls back to a rule-based ComplexityRouter
|
||||
when the LLM call fails.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_name: str,
|
||||
litellm_router_instance: "Router",
|
||||
llm_classifier_router_config: Optional[Dict[str, Any]] = None,
|
||||
default_model: Optional[str] = None,
|
||||
):
|
||||
self.model_name = model_name
|
||||
self.litellm_router_instance = litellm_router_instance
|
||||
|
||||
if llm_classifier_router_config:
|
||||
self.config = LLMClassifierRouterConfig(**llm_classifier_router_config)
|
||||
else:
|
||||
self.config = LLMClassifierRouterConfig()
|
||||
|
||||
self._cache = _TierCache(
|
||||
ttl_seconds=self.config.cache_ttl_seconds,
|
||||
max_size=self.config.cache_max_size,
|
||||
)
|
||||
|
||||
self._fallback_scorer: Optional[Any] = None
|
||||
if self.config.fallback_to_complexity_router:
|
||||
from litellm.router_strategy.complexity_router.complexity_router import (
|
||||
ComplexityRouter,
|
||||
)
|
||||
|
||||
self._fallback_scorer = ComplexityRouter(
|
||||
model_name=f"{model_name}::fallback_scorer",
|
||||
litellm_router_instance=litellm_router_instance,
|
||||
)
|
||||
|
||||
verbose_router_logger.debug(f"LLMClassifierRouter initialized for {model_name} with tiers: {self.config.tiers}")
|
||||
|
||||
def _get_system_prompt(self) -> str:
|
||||
return self.config.classifier_system_prompt or TWO_TIER_SYSTEM_PROMPT
|
||||
|
||||
async def _call_classifier_llm(self, prompt: str) -> str:
|
||||
"""Call the classifier LLM directly via litellm.acompletion.
|
||||
|
||||
Bypasses the Router instance to avoid recursion. The litellm
|
||||
import is lazy so this module can be imported during Router init
|
||||
without circular-import issues.
|
||||
"""
|
||||
import litellm
|
||||
|
||||
truncated = prompt[: self.config.classifier_max_input_chars]
|
||||
response = await litellm.acompletion(
|
||||
model=self.config.classifier_model,
|
||||
messages=[
|
||||
{"role": "system", "content": self._get_system_prompt()},
|
||||
{"role": "user", "content": truncated},
|
||||
],
|
||||
temperature=self.config.classifier_temperature,
|
||||
max_tokens=self.config.classifier_max_tokens,
|
||||
)
|
||||
return (response.choices[0].message.content or "").strip()
|
||||
|
||||
async def classify(self, prompt: str, system_prompt: Optional[str] = None) -> Tuple[str, str]:
|
||||
"""Classify a prompt and return (tier, method).
|
||||
|
||||
method ∈ {"cache", "llm", "fallback_rule", "fallback_default"}
|
||||
"""
|
||||
valid_tiers = tuple(self.config.tiers.keys())
|
||||
|
||||
if self.config.enable_cache:
|
||||
cached = self._cache.get(_cache_key(prompt))
|
||||
if cached is not None and cached in valid_tiers:
|
||||
return cached, "cache"
|
||||
|
||||
try:
|
||||
raw = await asyncio.wait_for(
|
||||
self._call_classifier_llm(prompt),
|
||||
timeout=self.config.classifier_timeout,
|
||||
)
|
||||
tier = _parse_tier(raw, valid_tiers)
|
||||
if self.config.enable_cache:
|
||||
self._cache.set(_cache_key(prompt), tier)
|
||||
return tier, "llm"
|
||||
except Exception as e:
|
||||
verbose_router_logger.warning(
|
||||
f"LLMClassifierRouter: LLM classification failed ({type(e).__name__}: {e}), falling back"
|
||||
)
|
||||
|
||||
if self._fallback_scorer is not None:
|
||||
try:
|
||||
complexity_tier, _, _ = self._fallback_scorer.classify(prompt, system_prompt)
|
||||
tier_name = complexity_tier.value
|
||||
if tier_name in valid_tiers:
|
||||
return tier_name, "fallback_rule"
|
||||
except Exception as e:
|
||||
verbose_router_logger.warning(f"LLMClassifierRouter: fallback scorer failed ({e})")
|
||||
|
||||
fallback = self.config.fallback_tier
|
||||
if fallback not in valid_tiers and valid_tiers:
|
||||
fallback = valid_tiers[0]
|
||||
return fallback, "fallback_default"
|
||||
|
||||
def get_model_for_tier(self, tier: str) -> Optional[str]:
|
||||
return self.config.tiers.get(tier)
|
||||
|
||||
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"]:
|
||||
"""Pre-routing hook that classifies the request and selects a tier model."""
|
||||
from litellm.types.router import PreRoutingHookResponse
|
||||
|
||||
if self._fallback_scorer is None:
|
||||
from litellm.router_strategy.complexity_router.complexity_router import (
|
||||
ComplexityRouter,
|
||||
)
|
||||
|
||||
helper = ComplexityRouter(
|
||||
model_name=f"{self.model_name}::hook_helper",
|
||||
litellm_router_instance=self.litellm_router_instance,
|
||||
)
|
||||
else:
|
||||
helper = self._fallback_scorer
|
||||
|
||||
resolved_messages = helper._resolve_messages(messages, request_kwargs)
|
||||
if not resolved_messages:
|
||||
verbose_router_logger.debug("LLMClassifierRouter: No messages could be resolved, skipping routing")
|
||||
return None
|
||||
|
||||
has_original_messages = messages is not None and len(messages) > 0
|
||||
user_message, system_prompt = helper._extract_user_message_and_system_prompt(resolved_messages)
|
||||
|
||||
if user_message is None:
|
||||
verbose_router_logger.debug("LLMClassifierRouter: No user message found, skipping routing")
|
||||
return None
|
||||
|
||||
tier, method = await self.classify(user_message, system_prompt)
|
||||
routed_model = self.get_model_for_tier(tier) or self.get_model_for_tier(self.config.fallback_tier)
|
||||
|
||||
verbose_router_logger.info(f"LLMClassifierRouter: tier={tier}, method={method}, routed_model={routed_model}")
|
||||
|
||||
metadata = request_kwargs.setdefault("metadata", {})
|
||||
metadata["llm_classifier_router_tier"] = tier
|
||||
metadata["llm_classifier_router_method"] = method
|
||||
metadata["llm_classifier_router_model"] = routed_model
|
||||
|
||||
return PreRoutingHookResponse(
|
||||
model=routed_model,
|
||||
messages=messages if has_original_messages else None,
|
||||
)
|
||||
|
|
@ -0,0 +1,10 @@
|
|||
"""Prompt templates for the LLM Classifier Router."""
|
||||
|
||||
TWO_TIER_SYSTEM_PROMPT = """You are a request classifier. Classify the user message as SIMPLE or COMPLEX.
|
||||
|
||||
SIMPLE: greetings, factual lookups, single-step questions, yes/no, brief creative tasks.
|
||||
COMPLEX: multi-step reasoning, code generation, technical analysis, long-form writing, comparisons.
|
||||
|
||||
Reply with EXACTLY one word: SIMPLE or COMPLEX"""
|
||||
|
||||
VALID_TIERS = ("SIMPLE", "COMPLEX")
|
||||
|
|
@ -245,6 +245,10 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams):
|
|||
complexity_router_config: Optional[Dict] = None
|
||||
complexity_router_default_model: Optional[str] = None
|
||||
|
||||
# llm-classifier-router params
|
||||
llm_classifier_router_config: Optional[Dict] = None
|
||||
llm_classifier_router_default_model: Optional[str] = None
|
||||
|
||||
# adaptive-router params
|
||||
adaptive_router_default_model: Optional[str] = None
|
||||
adaptive_router_config: Optional[Dict] = None
|
||||
|
|
|
|||
435
tests/test_litellm/router_strategy/test_llm_classifier_router.py
Normal file
435
tests/test_litellm/router_strategy/test_llm_classifier_router.py
Normal file
|
|
@ -0,0 +1,435 @@
|
|||
"""Tests for the LLMClassifierRouter."""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import sys
|
||||
from typing import Dict
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
sys.path.insert(0, os.path.abspath("../../.."))
|
||||
|
||||
from litellm.router_strategy.llm_classifier_router import (
|
||||
LLMClassifierRouter,
|
||||
LLMClassifierRouterConfig,
|
||||
)
|
||||
from litellm.router_strategy.llm_classifier_router.llm_classifier_router import (
|
||||
_cache_key,
|
||||
_parse_tier,
|
||||
_TierCache,
|
||||
)
|
||||
|
||||
|
||||
VALID = ("SIMPLE", "COMPLEX")
|
||||
|
||||
|
||||
def _mock_response(content: str) -> MagicMock:
|
||||
return MagicMock(choices=[MagicMock(message=MagicMock(content=content))])
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_router_instance() -> MagicMock:
|
||||
return MagicMock()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def basic_config() -> Dict:
|
||||
return {
|
||||
"tiers": {"SIMPLE": "gpt-4o-mini", "COMPLEX": "gpt-4o"},
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def router(mock_router_instance, basic_config) -> LLMClassifierRouter:
|
||||
return LLMClassifierRouter(
|
||||
model_name="test-llm-classifier",
|
||||
litellm_router_instance=mock_router_instance,
|
||||
llm_classifier_router_config=basic_config,
|
||||
)
|
||||
|
||||
|
||||
# ─── Config ──────────────────────────────────────────────────────
|
||||
class TestLLMClassifierConfig:
|
||||
def test_default_values(self):
|
||||
c = LLMClassifierRouterConfig()
|
||||
assert c.classifier_model == "ollama/qwen2.5:0.5b"
|
||||
assert c.classifier_timeout == 3.0
|
||||
assert c.enable_cache is True
|
||||
assert c.fallback_to_complexity_router is True
|
||||
|
||||
def test_fallback_tier_defaults_to_simple(self):
|
||||
c = LLMClassifierRouterConfig()
|
||||
assert c.fallback_tier == "SIMPLE"
|
||||
|
||||
def test_tiers_default_has_simple_and_complex(self):
|
||||
c = LLMClassifierRouterConfig()
|
||||
assert "SIMPLE" in c.tiers
|
||||
assert "COMPLEX" in c.tiers
|
||||
|
||||
def test_extra_kwargs_allowed(self):
|
||||
c = LLMClassifierRouterConfig(future_field=True, classifier_model="custom")
|
||||
assert c.classifier_model == "custom"
|
||||
|
||||
def test_custom_tiers_override_default(self):
|
||||
c = LLMClassifierRouterConfig(tiers={"LOW": "a", "HIGH": "b"})
|
||||
assert c.tiers == {"LOW": "a", "HIGH": "b"}
|
||||
|
||||
|
||||
# ─── _parse_tier ─────────────────────────────────────────────────
|
||||
class TestParseTier:
|
||||
def test_exact_match_simple(self):
|
||||
assert _parse_tier("SIMPLE", VALID) == "SIMPLE"
|
||||
|
||||
def test_exact_match_complex(self):
|
||||
assert _parse_tier("COMPLEX", VALID) == "COMPLEX"
|
||||
|
||||
def test_case_insensitive(self):
|
||||
assert _parse_tier("simple", VALID) == "SIMPLE"
|
||||
assert _parse_tier("Complex", VALID) == "COMPLEX"
|
||||
assert _parse_tier("cOmPlEx", VALID) == "COMPLEX"
|
||||
|
||||
def test_with_trailing_punctuation(self):
|
||||
assert _parse_tier("SIMPLE.", VALID) == "SIMPLE"
|
||||
assert _parse_tier("SIMPLE!", VALID) == "SIMPLE"
|
||||
assert _parse_tier("COMPLEX.", VALID) == "COMPLEX"
|
||||
|
||||
def test_with_whitespace(self):
|
||||
assert _parse_tier(" SIMPLE ", VALID) == "SIMPLE"
|
||||
assert _parse_tier("\nCOMPLEX\n", VALID) == "COMPLEX"
|
||||
|
||||
def test_substring_fallback(self):
|
||||
# "I think this is SIMPLE" doesn't fullmatch but contains SIMPLE
|
||||
assert _parse_tier("I think this is SIMPLE", VALID) == "SIMPLE"
|
||||
assert _parse_tier("The answer is COMPLEX", VALID) == "COMPLEX"
|
||||
|
||||
def test_unparseable_raises_value_error(self):
|
||||
with pytest.raises(ValueError, match="Could not parse tier"):
|
||||
_parse_tier("I don't know", VALID)
|
||||
|
||||
def test_unparseable_no_tier_word_raises(self):
|
||||
with pytest.raises(ValueError):
|
||||
_parse_tier("maybe?", VALID)
|
||||
|
||||
def test_complexity_tier_value_passes_through(self):
|
||||
# If the LLM returns one of the 4 complexity tiers, _parse_tier should
|
||||
# still raise because it's not in the configured 2-tier set.
|
||||
with pytest.raises(ValueError):
|
||||
_parse_tier("MEDIUM", VALID)
|
||||
|
||||
|
||||
# ─── _TierCache ──────────────────────────────────────────────────
|
||||
class TestTierCache:
|
||||
def test_hit_returns_cached_tier(self):
|
||||
cache = _TierCache(ttl_seconds=60, max_size=100)
|
||||
cache.set("k1", "SIMPLE")
|
||||
assert cache.get("k1") == "SIMPLE"
|
||||
|
||||
def test_miss_returns_none(self):
|
||||
cache = _TierCache(ttl_seconds=60, max_size=100)
|
||||
assert cache.get("missing") is None
|
||||
|
||||
def test_expiry(self):
|
||||
cache = _TierCache(ttl_seconds=0.01, max_size=100)
|
||||
cache.set("k1", "SIMPLE")
|
||||
import time
|
||||
|
||||
time.sleep(0.05)
|
||||
assert cache.get("k1") is None
|
||||
|
||||
def test_overflow_clears_store(self):
|
||||
cache = _TierCache(ttl_seconds=60, max_size=2)
|
||||
cache.set("k1", "SIMPLE")
|
||||
cache.set("k2", "COMPLEX")
|
||||
cache.set("k3", "SIMPLE")
|
||||
# After overflow, k1 and k2 are evicted
|
||||
assert cache.get("k1") is None
|
||||
assert cache.get("k2") is None
|
||||
assert cache.get("k3") == "SIMPLE"
|
||||
|
||||
def test_overwrite_same_key(self):
|
||||
cache = _TierCache(ttl_seconds=60, max_size=100)
|
||||
cache.set("k1", "SIMPLE")
|
||||
cache.set("k1", "COMPLEX")
|
||||
assert cache.get("k1") == "COMPLEX"
|
||||
|
||||
|
||||
# ─── _cache_key ──────────────────────────────────────────────────
|
||||
class TestCacheKey:
|
||||
def test_deterministic(self):
|
||||
assert _cache_key("hello world") == _cache_key("hello world")
|
||||
|
||||
def test_different_inputs_different_keys(self):
|
||||
assert _cache_key("a") != _cache_key("b")
|
||||
|
||||
def test_key_length_16(self):
|
||||
assert len(_cache_key("anything")) == 16
|
||||
|
||||
def test_handles_unicode(self):
|
||||
k = _cache_key("你好世界")
|
||||
assert len(k) == 16
|
||||
assert _cache_key("你好世界") == k
|
||||
|
||||
|
||||
# ─── classify: cache path ────────────────────────────────────────
|
||||
class TestClassifyCachePath:
|
||||
async def test_successful_llm_call_routes_to_simPLE_tier(self, router):
|
||||
with patch(
|
||||
"litellm.acompletion",
|
||||
new=AsyncMock(return_value=_mock_response("SIMPLE")),
|
||||
) as mock:
|
||||
tier, method = await router.classify("What is the capital of France?")
|
||||
assert tier == "SIMPLE"
|
||||
assert method == "llm"
|
||||
mock.assert_awaited_once()
|
||||
|
||||
async def test_complex_prompt_routes_correctly(self, router):
|
||||
with patch(
|
||||
"litellm.acompletion",
|
||||
new=AsyncMock(return_value=_mock_response("COMPLEX")),
|
||||
):
|
||||
tier, method = await router.classify("Implement a thread-safe LRU cache with O(1) get/set")
|
||||
assert tier == "COMPLEX"
|
||||
assert method == "llm"
|
||||
|
||||
async def test_cache_hit_returns_cache_method(self, router):
|
||||
with patch(
|
||||
"litellm.acompletion",
|
||||
new=AsyncMock(return_value=_mock_response("SIMPLE")),
|
||||
):
|
||||
await router.classify("prompt X")
|
||||
tier, method = await router.classify("prompt X")
|
||||
assert method == "cache"
|
||||
assert tier == "SIMPLE"
|
||||
|
||||
async def test_cache_prevents_duplicate_llm_calls(self, router):
|
||||
with patch(
|
||||
"litellm.acompletion",
|
||||
new=AsyncMock(return_value=_mock_response("SIMPLE")),
|
||||
) as mock:
|
||||
for _ in range(3):
|
||||
await router.classify("same prompt")
|
||||
mock.assert_awaited_once()
|
||||
|
||||
async def test_cache_disabled_always_calls_llm(self, mock_router_instance):
|
||||
config = {
|
||||
"tiers": {"SIMPLE": "cheap", "COMPLEX": "expensive"},
|
||||
"enable_cache": False,
|
||||
}
|
||||
r = LLMClassifierRouter(
|
||||
model_name="test",
|
||||
litellm_router_instance=mock_router_instance,
|
||||
llm_classifier_router_config=config,
|
||||
)
|
||||
with patch(
|
||||
"litellm.acompletion",
|
||||
new=AsyncMock(return_value=_mock_response("SIMPLE")),
|
||||
) as mock:
|
||||
await r.classify("p")
|
||||
await r.classify("p")
|
||||
assert mock.await_count == 2
|
||||
|
||||
async def test_cache_stale_entry_expired_recalls_llm(self, router):
|
||||
# Force expiry by manually setting a near-expired entry
|
||||
from time import monotonic
|
||||
|
||||
router._cache._store["k"] = ("SIMPLE", monotonic() - 1)
|
||||
with patch(
|
||||
"litellm.acompletion",
|
||||
new=AsyncMock(return_value=_mock_response("COMPLEX")),
|
||||
) as mock:
|
||||
tier, method = await router.classify("anything")
|
||||
# 'k' was for a different prompt, so the key wouldn't match anyway,
|
||||
# but we can still verify the LLM was called
|
||||
assert method == "llm"
|
||||
assert mock.await_count == 1
|
||||
|
||||
|
||||
# ─── classify: fallback path ─────────────────────────────────────
|
||||
class TestClassifyFallbackPath:
|
||||
async def test_timeout_triggers_fallback(self, router):
|
||||
async def slow(*args, **kwargs):
|
||||
await asyncio.sleep(10)
|
||||
return _mock_response("SIMPLE")
|
||||
|
||||
with patch("litellm.acompletion", new=AsyncMock(side_effect=slow)):
|
||||
tier, method = await router.classify("test", system_prompt=None)
|
||||
assert method in ("fallback_rule", "fallback_default")
|
||||
assert tier in ("SIMPLE", "COMPLEX")
|
||||
|
||||
async def test_api_error_triggers_fallback(self, router):
|
||||
with patch(
|
||||
"litellm.acompletion",
|
||||
new=AsyncMock(side_effect=Exception("API down")),
|
||||
):
|
||||
tier, method = await router.classify("test")
|
||||
assert method in ("fallback_rule", "fallback_default")
|
||||
|
||||
async def test_fallback_default_when_complexity_router_disabled(self, mock_router_instance):
|
||||
config = {
|
||||
"tiers": {"SIMPLE": "cheap", "COMPLEX": "expensive"},
|
||||
"fallback_to_complexity_router": False,
|
||||
}
|
||||
r = LLMClassifierRouter(
|
||||
model_name="test",
|
||||
litellm_router_instance=mock_router_instance,
|
||||
llm_classifier_router_config=config,
|
||||
)
|
||||
with patch(
|
||||
"litellm.acompletion",
|
||||
new=AsyncMock(side_effect=Exception("fail")),
|
||||
):
|
||||
tier, method = await r.classify("test")
|
||||
assert method == "fallback_default"
|
||||
assert tier == "SIMPLE"
|
||||
|
||||
async def test_fallback_tier_used_when_complexity_router_also_fails(self, mock_router_instance):
|
||||
config = {
|
||||
"tiers": {"SIMPLE": "cheap", "COMPLEX": "expensive"},
|
||||
"fallback_to_complexity_router": True,
|
||||
}
|
||||
r = LLMClassifierRouter(
|
||||
model_name="test",
|
||||
litellm_router_instance=mock_router_instance,
|
||||
llm_classifier_router_config=config,
|
||||
)
|
||||
with patch(
|
||||
"litellm.acompletion",
|
||||
new=AsyncMock(side_effect=Exception("fail")),
|
||||
):
|
||||
# ComplexityRouter should also be patched/mocked to fail,
|
||||
# but since it's a rules-based classifier, it should not fail.
|
||||
# This test mainly verifies the end result is a valid tier.
|
||||
tier, method = await r.classify("test")
|
||||
assert method in ("fallback_rule", "fallback_default")
|
||||
assert tier in ("SIMPLE", "COMPLEX")
|
||||
|
||||
async def test_invalid_fallback_tier_falls_back_to_first_tier(self, mock_router_instance):
|
||||
config = {
|
||||
"tiers": {"A": "model-a", "B": "model-b"},
|
||||
"fallback_tier": "NONEXISTENT",
|
||||
"fallback_to_complexity_router": False,
|
||||
}
|
||||
r = LLMClassifierRouter(
|
||||
model_name="test",
|
||||
litellm_router_instance=mock_router_instance,
|
||||
llm_classifier_router_config=config,
|
||||
)
|
||||
with patch(
|
||||
"litellm.acompletion",
|
||||
new=AsyncMock(side_effect=Exception("fail")),
|
||||
):
|
||||
tier, method = await r.classify("test")
|
||||
assert method == "fallback_default"
|
||||
# Should pick the first valid tier
|
||||
assert tier in ("A", "B")
|
||||
|
||||
|
||||
# ─── async_pre_routing_hook ─────────────────────────────────────
|
||||
class TestAsyncPreRoutingHook:
|
||||
async def test_returns_pre_routing_hook_response_for_simple_prompt(self, router):
|
||||
from litellm.types.router import PreRoutingHookResponse
|
||||
|
||||
with patch(
|
||||
"litellm.acompletion",
|
||||
new=AsyncMock(return_value=_mock_response("SIMPLE")),
|
||||
):
|
||||
response = await router.async_pre_routing_hook(
|
||||
model="test-llm-classifier",
|
||||
request_kwargs={},
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
)
|
||||
assert isinstance(response, PreRoutingHookResponse)
|
||||
assert response.model == "gpt-4o-mini"
|
||||
|
||||
async def test_returns_pre_routing_hook_response_for_complex_prompt(self, router):
|
||||
with patch(
|
||||
"litellm.acompletion",
|
||||
new=AsyncMock(return_value=_mock_response("COMPLEX")),
|
||||
):
|
||||
response = await router.async_pre_routing_hook(
|
||||
model="test-llm-classifier",
|
||||
request_kwargs={},
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": "Implement a thread-safe LRU cache",
|
||||
}
|
||||
],
|
||||
)
|
||||
assert response.model == "gpt-4o"
|
||||
|
||||
async def test_stashes_decision_in_metadata(self, router):
|
||||
request_kwargs: Dict = {}
|
||||
with patch(
|
||||
"litellm.acompletion",
|
||||
new=AsyncMock(return_value=_mock_response("COMPLEX")),
|
||||
):
|
||||
await router.async_pre_routing_hook(
|
||||
model="test-llm-classifier",
|
||||
request_kwargs=request_kwargs,
|
||||
messages=[{"role": "user", "content": "design a system"}],
|
||||
)
|
||||
meta = request_kwargs["metadata"]
|
||||
assert meta["llm_classifier_router_tier"] == "COMPLEX"
|
||||
assert meta["llm_classifier_router_method"] == "llm"
|
||||
assert meta["llm_classifier_router_model"] == "gpt-4o"
|
||||
|
||||
async def test_preserves_existing_metadata(self, router):
|
||||
request_kwargs = {"metadata": {"user_id": "u123"}}
|
||||
with patch(
|
||||
"litellm.acompletion",
|
||||
new=AsyncMock(return_value=_mock_response("SIMPLE")),
|
||||
):
|
||||
await router.async_pre_routing_hook(
|
||||
model="test-llm-classifier",
|
||||
request_kwargs=request_kwargs,
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
)
|
||||
assert request_kwargs["metadata"]["user_id"] == "u123"
|
||||
assert request_kwargs["metadata"]["llm_classifier_router_tier"] == "SIMPLE"
|
||||
|
||||
async def test_no_user_message_returns_none(self, router):
|
||||
response = await router.async_pre_routing_hook(
|
||||
model="test-llm-classifier",
|
||||
request_kwargs={},
|
||||
messages=[{"role": "system", "content": "you are a helper"}],
|
||||
)
|
||||
assert response is None
|
||||
|
||||
async def test_no_messages_returns_none(self, router):
|
||||
response = await router.async_pre_routing_hook(
|
||||
model="test-llm-classifier",
|
||||
request_kwargs={},
|
||||
messages=[],
|
||||
)
|
||||
assert response is None
|
||||
|
||||
async def test_extracts_user_message_from_message_list(self, router):
|
||||
with patch(
|
||||
"litellm.acompletion",
|
||||
new=AsyncMock(return_value=_mock_response("SIMPLE")),
|
||||
) as mock:
|
||||
await router.async_pre_routing_hook(
|
||||
model="test-llm-classifier",
|
||||
request_kwargs={},
|
||||
messages=[
|
||||
{"role": "system", "content": "be terse"},
|
||||
{"role": "user", "content": "What is 2+2?"},
|
||||
],
|
||||
)
|
||||
# Verify the user content was sent to the LLM
|
||||
call_args = mock.call_args
|
||||
messages = call_args.kwargs.get("messages") or call_args.args[0]
|
||||
user_msg = next(m for m in messages if m["role"] == "user")
|
||||
assert user_msg["content"] == "What is 2+2?"
|
||||
|
||||
|
||||
# ─── get_model_for_tier ─────────────────────────────────────────
|
||||
class TestGetModelForTier:
|
||||
def test_known_tier(self, router):
|
||||
assert router.get_model_for_tier("SIMPLE") == "gpt-4o-mini"
|
||||
assert router.get_model_for_tier("COMPLEX") == "gpt-4o"
|
||||
|
||||
def test_unknown_tier(self, router):
|
||||
assert router.get_model_for_tier("UNKNOWN") is None
|
||||
|
|
@ -0,0 +1,183 @@
|
|||
"""Integration tests for LLMClassifierRouter wired into the main Router.
|
||||
|
||||
Verifies the 7 patches in litellm/router.py are correctly placed and the new
|
||||
routing strategy is auto-discovered from a deployment config.
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
sys.path.insert(0, os.path.abspath("../../.."))
|
||||
|
||||
from litellm import Router
|
||||
from litellm.router_strategy.llm_classifier_router import LLMClassifierRouter
|
||||
from litellm.types.router import LiteLLM_Params
|
||||
|
||||
|
||||
def _mock_response(content: str) -> MagicMock:
|
||||
return MagicMock(choices=[MagicMock(message=MagicMock(content=content))])
|
||||
|
||||
|
||||
# ─── Detection methods on Router ─────────────────────────────────
|
||||
class TestRouterDetection:
|
||||
def test_is_llm_classifier_router_deployment_true(self):
|
||||
params = LiteLLM_Params(model="auto_router/llm_classifier_router")
|
||||
router = Router(model_list=[])
|
||||
assert router._is_llm_classifier_router_deployment(params) is True
|
||||
|
||||
def test_is_llm_classifier_router_deployment_false_for_other_prefixes(self):
|
||||
router = Router(model_list=[])
|
||||
for model_name in [
|
||||
"auto_router/complexity_router",
|
||||
"auto_router/quality_router",
|
||||
"auto_router/adaptive_router",
|
||||
"auto_router/semantic",
|
||||
"gpt-4o-mini",
|
||||
"openai/gpt-4o",
|
||||
]:
|
||||
params = LiteLLM_Params(model=model_name)
|
||||
assert router._is_llm_classifier_router_deployment(params) is False, f"Expected False for {model_name}"
|
||||
|
||||
def test_is_auto_router_deployment_excludes_llm_classifier(self):
|
||||
# Regression: the gap where auto_router/llm_classifier_router fell
|
||||
# into the generic AutoRouter bucket.
|
||||
router = Router(model_list=[])
|
||||
params = LiteLLM_Params(model="auto_router/llm_classifier_router")
|
||||
assert router._is_auto_router_deployment(params) is False
|
||||
|
||||
def test_is_auto_router_deployment_still_true_for_other_prefixes(self):
|
||||
router = Router(model_list=[])
|
||||
params = LiteLLM_Params(model="auto_router/semantic")
|
||||
assert router._is_auto_router_deployment(params) is True
|
||||
|
||||
|
||||
# ─── Init method registers the router ───────────────────────────
|
||||
class TestRouterInit:
|
||||
def test_init_llm_classifier_router_deployment_registers_router(self):
|
||||
from litellm.types.router import Deployment
|
||||
|
||||
deployment = Deployment(
|
||||
model_name="smart-router",
|
||||
litellm_params=LiteLLM_Params(
|
||||
model="auto_router/llm_classifier_router",
|
||||
llm_classifier_router_config={
|
||||
"tiers": {"SIMPLE": "cheap", "COMPLEX": "expensive"},
|
||||
},
|
||||
),
|
||||
)
|
||||
router = Router(model_list=[])
|
||||
router.init_llm_classifier_router_deployment(deployment=deployment)
|
||||
assert "smart-router" in router.llm_classifier_routers
|
||||
assert isinstance(router.llm_classifier_routers["smart-router"], LLMClassifierRouter)
|
||||
|
||||
def test_init_uses_default_model_from_tiers_complex(self):
|
||||
from litellm.types.router import Deployment
|
||||
|
||||
deployment = Deployment(
|
||||
model_name="smart-router",
|
||||
litellm_params=LiteLLM_Params(
|
||||
model="auto_router/llm_classifier_router",
|
||||
llm_classifier_router_config={
|
||||
"tiers": {"SIMPLE": "cheap", "COMPLEX": "expensive"},
|
||||
},
|
||||
),
|
||||
)
|
||||
router = Router(model_list=[])
|
||||
router.init_llm_classifier_router_deployment(deployment=deployment)
|
||||
# Default model should fall back to COMPLEX tier value when no
|
||||
# llm_classifier_router_default_model is set.
|
||||
assert router.llm_classifier_routers["smart-router"].config.fallback_tier == "SIMPLE"
|
||||
|
||||
def test_init_raises_when_no_default_model_and_no_tiers(self):
|
||||
from litellm.types.router import Deployment
|
||||
|
||||
deployment = Deployment(
|
||||
model_name="smart-router",
|
||||
litellm_params=LiteLLM_Params(
|
||||
model="auto_router/llm_classifier_router",
|
||||
llm_classifier_router_config={"tiers": {}},
|
||||
),
|
||||
)
|
||||
router = Router(model_list=[])
|
||||
with pytest.raises(ValueError, match="llm_classifier_router_default_model"):
|
||||
router.init_llm_classifier_router_deployment(deployment=deployment)
|
||||
|
||||
def test_init_raises_on_duplicate_model_name(self):
|
||||
from litellm.types.router import Deployment
|
||||
|
||||
deployment1 = Deployment(
|
||||
model_name="smart-router",
|
||||
litellm_params=LiteLLM_Params(
|
||||
model="auto_router/llm_classifier_router",
|
||||
llm_classifier_router_default_model="x",
|
||||
),
|
||||
)
|
||||
deployment2 = Deployment(
|
||||
model_name="smart-router",
|
||||
litellm_params=LiteLLM_Params(
|
||||
model="auto_router/llm_classifier_router",
|
||||
llm_classifier_router_default_model="x",
|
||||
),
|
||||
)
|
||||
router = Router(model_list=[])
|
||||
router.init_llm_classifier_router_deployment(deployment=deployment1)
|
||||
with pytest.raises(ValueError, match="already exists"):
|
||||
router.init_llm_classifier_router_deployment(deployment=deployment2)
|
||||
|
||||
def test_init_uses_explicit_default_model(self):
|
||||
from litellm.types.router import Deployment
|
||||
|
||||
deployment = Deployment(
|
||||
model_name="smart-router",
|
||||
litellm_params=LiteLLM_Params(
|
||||
model="auto_router/llm_classifier_router",
|
||||
llm_classifier_router_default_model="my-explicit-default",
|
||||
),
|
||||
)
|
||||
router = Router(model_list=[])
|
||||
router.init_llm_classifier_router_deployment(deployment=deployment)
|
||||
assert router.llm_classifier_routers["smart-router"].config.fallback_tier == "SIMPLE"
|
||||
|
||||
|
||||
# ─── End-to-end through Router init ──────────────────────────────
|
||||
class TestRouterEndToEnd:
|
||||
def test_router_picks_up_deployment_from_model_list(self):
|
||||
model_list = [
|
||||
{
|
||||
"model_name": "smart-router",
|
||||
"litellm_params": {
|
||||
"model": "auto_router/llm_classifier_router",
|
||||
"llm_classifier_router_config": {
|
||||
"tiers": {"SIMPLE": "a", "COMPLEX": "b"},
|
||||
},
|
||||
},
|
||||
}
|
||||
]
|
||||
with patch("litellm.acompletion"):
|
||||
router = Router(model_list=model_list)
|
||||
assert "smart-router" in router.llm_classifier_routers
|
||||
|
||||
def test_set_model_list_resets_dict(self):
|
||||
model_list_1 = [
|
||||
{
|
||||
"model_name": "smart-router",
|
||||
"litellm_params": {
|
||||
"model": "auto_router/llm_classifier_router",
|
||||
"llm_classifier_router_config": {
|
||||
"tiers": {"SIMPLE": "a", "COMPLEX": "b"},
|
||||
},
|
||||
},
|
||||
}
|
||||
]
|
||||
with patch("litellm.acompletion"):
|
||||
router = Router(model_list=model_list_1)
|
||||
assert "smart-router" in router.llm_classifier_routers
|
||||
|
||||
with patch("litellm.acompletion"):
|
||||
router.set_model_list(model_list_1)
|
||||
# After set_model_list, the dict should still be populated (it's
|
||||
# a reset-then-rebuild, not a clear-only).
|
||||
assert "smart-router" in router.llm_classifier_routers
|
||||
Loading…
Add table
Reference in a new issue