Model Routing › Code Implementation
Language AI HandbookAll chapters Model Routing › Code ImplementationOne section · Sections in this chapter7 / 11 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Let's implement the core components of a model routing system: a difficulty-based router, a confidence-threshold cascade, and an A/B test evaluation framework.
We start by installing required packages and setting up the simulation environment.
Simulating Requests and Models
We simulate a production environment with requests of varying difficulty and two model tiers with different quality and cost profiles.
@dataclass
class Request:
"""A single inference request with a ground-truth difficulty score."""
text: str
difficulty: float # 0.0 (easy) to 1.0 (hard)
tokens: int
@dataclass
class ModelResponse:
"""Response from a model tier."""
output: str
quality: float # 0.0 to 1.0
confidence: float # model's self-reported confidence
cost: float # cost in arbitrary units
latency_ms: float
def generate_requests(n: int) -> list[Request]:
"""Generate requests with a realistic difficulty distribution.
Most requests are easy (right-skewed), with a long tail of hard ones.
"""
difficulties = np.random.beta(a=1.5, b=4.0, size=n) # skewed toward easy
token_counts = (20 + difficulties * 200 + np.random.randn(n) * 20).astype(
int
)
token_counts = np.clip(token_counts, 10, 500)
requests = [
Request(text=f"request_{i}", difficulty=float(d), tokens=int(t))
for i, (d, t) in enumerate(zip(difficulties, token_counts))
]
return requests
def small_model_response(req: Request) -> ModelResponse:
"""Simulate a small (cheap) model. Struggles on hard requests."""
# Quality degrades on harder requests; small model is noisy
base_quality = 0.9 - 0.7 * req.difficulty
quality = float(np.clip(base_quality + np.random.randn() * 0.1, 0, 1))
# Confidence is correlated with actual quality, but noisy
confidence = float(np.clip(quality + np.random.randn() * 0.15, 0, 1))
return ModelResponse(
output="small_answer",
quality=quality,
confidence=confidence,
cost=0.001 * req.tokens,
latency_ms=10 + req.tokens * 0.2,
)
def large_model_response(req: Request) -> ModelResponse:
"""Simulate a large (expensive) model. High quality across all difficulties."""
base_quality = 0.97 - 0.15 * req.difficulty
quality = float(np.clip(base_quality + np.random.randn() * 0.05, 0, 1))
confidence = float(np.clip(quality + np.random.randn() * 0.05, 0, 1))
return ModelResponse(
output="large_answer",
quality=quality,
confidence=confidence,
cost=0.01 * req.tokens,
latency_ms=100 + req.tokens * 2.0,
)Generated 1000 requests Difficulty: mean=0.272, std=0.170 Tokens: mean=74.0
The simulated requests follow a realistic distribution: most are relatively easy (mean difficulty around 0.27), with a long tail of harder requests. The small model's quality degrades as difficulty increases, while the large model maintains high quality throughout. This asymmetry is the key structural fact that makes routing useful.
Rule-Based Router
Our first router uses a simple difficulty threshold based on request length as a proxy.
def rule_based_router(req: Request, token_threshold: int = 100) -> str:
"""Route based on token count as a proxy for difficulty."""
if req.tokens <= token_threshold:
return "small"
return "large"
def evaluate_router(
requests: list[Request],
router: Callable[[Request], str],
label: str,
) -> dict:
"""Evaluate a routing strategy and collect metrics."""
results = []
for req in requests:
decision = router(req)
if decision == "small":
resp = small_model_response(req)
else:
resp = large_model_response(req)
results.append(
{
"decision": decision,
"quality": resp.quality,
"cost": resp.cost,
"latency_ms": resp.latency_ms,
}
)
qualities = [r["quality"] for r in results]
costs = [r["cost"] for r in results]
latencies = [r["latency_ms"] for r in results]
small_frac = sum(1 for r in results if r["decision"] == "small") / len(
results
)
return {
"label": label,
"mean_quality": float(np.mean(qualities)),
"mean_cost": float(np.mean(costs)),
"mean_latency_ms": float(np.mean(latencies)),
"small_model_fraction": float(small_frac),
"total_cost": float(np.sum(costs)),
}Always Large: Mean quality: 0.9237 Mean cost: 0.74011 Mean latency: 248.0 ms Small model %: 0.0% Total cost: 740.11 Rule-Based (token): Mean quality: 0.8255 Mean cost: 0.47525 Mean latency: 140.6 ms Small model %: 60.5% Total cost: 475.25
The rule-based router routes a substantial fraction of requests to the small model, reducing total cost significantly. The quality drop is small because most short requests are easy. This illustrates the core value proposition: you pay less without sacrificing much quality. The token threshold of 80 was chosen somewhat arbitrarily; in a real system, you would calibrate this by sweeping thresholds and evaluating quality on a labeled validation set.
Cascade Router
Next, we implement a cascade that uses confidence thresholds to decide whether to escalate.
def cascade_router(req: Request, confidence_threshold: float = 0.75) -> dict:
"""Two-tier cascade: try small model first, escalate if confidence is low."""
small_resp = small_model_response(req)
if small_resp.confidence >= confidence_threshold:
return {
"decision": "small",
"quality": small_resp.quality,
"cost": small_resp.cost,
"latency_ms": small_resp.latency_ms,
"escalated": False,
}
# Escalate to large model
large_resp = large_model_response(req)
return {
"decision": "large",
"quality": large_resp.quality,
"cost": small_resp.cost + large_resp.cost, # paid for both
"latency_ms": small_resp.latency_ms + large_resp.latency_ms,
"escalated": True,
}
def evaluate_cascade(
requests: list[Request],
confidence_threshold: float,
label: str,
) -> dict:
results = [cascade_router(req, confidence_threshold) for req in requests]
qualities = [r["quality"] for r in results]
costs = [r["cost"] for r in results]
latencies = [r["latency_ms"] for r in results]
escalation_rate = float(np.mean([r["escalated"] for r in results]))
return {
"label": label,
"mean_quality": float(np.mean(qualities)),
"mean_cost": float(np.mean(costs)),
"mean_latency_ms": float(np.mean(latencies)),
"escalation_rate": escalation_rate,
"total_cost": float(np.sum(costs)),
}tau=0.50: quality=0.7674, cost=0.24381, escalation=16.0% tau=0.65: quality=0.8295, cost=0.41921, escalation=36.8% tau=0.75: quality=0.8734, cost=0.56904, escalation=56.0% tau=0.85: quality=0.9009, cost=0.67690, escalation=74.4% tau=0.95: quality=0.9148, cost=0.74571, escalation=86.5%
The cascade results show a clear tradeoff: higher confidence thresholds (requiring more certainty before returning the small model's answer) increase escalation rates and costs, but also improve quality. The right threshold depends on your quality requirements and cost budget. Notice that even at the highest threshold (), the cascade still costs less than always routing to the large model, because it correctly handles the clearly easy requests at the small model.
A/B Testing Framework
Now let's build a simple A/B testing framework to evaluate routing strategies with statistical rigor.
from scipy import stats
def run_ab_test(
requests: list[Request],
router_a: Callable[[Request], str],
router_b: Callable[[Request], str],
split_ratio: float = 0.5,
seed: int = 0,
) -> dict:
"""Run an A/B test comparing two routing strategies."""
rng = np.random.default_rng(seed)
# Randomly assign each request to arm A or arm B
assignments = rng.random(len(requests)) < split_ratio # True = arm A
arm_a_results = []
arm_b_results = []
for req, in_a in zip(requests, assignments):
if in_a:
decision = router_a(req)
resp = (
small_model_response(req)
if decision == "small"
else large_model_response(req)
)
arm_a_results.append(
{
"quality": resp.quality,
"cost": resp.cost,
"latency_ms": resp.latency_ms,
}
)
else:
decision = router_b(req)
resp = (
small_model_response(req)
if decision == "small"
else large_model_response(req)
)
arm_b_results.append(
{
"quality": resp.quality,
"cost": resp.cost,
"latency_ms": resp.latency_ms,
}
)
a_quality = [r["quality"] for r in arm_a_results]
b_quality = [r["quality"] for r in arm_b_results]
# Two-sample t-test for quality difference
t_stat, p_value = stats.ttest_ind(a_quality, b_quality)
quality_lift = float(np.mean(b_quality) - np.mean(a_quality))
return {
"n_a": len(arm_a_results),
"n_b": len(arm_b_results),
"quality_a": float(np.mean(a_quality)),
"quality_b": float(np.mean(b_quality)),
"quality_lift": quality_lift,
"t_statistic": float(t_stat),
"p_value": float(p_value),
"significant_at_05": bool(p_value < 0.05),
"cost_a": float(np.mean([r["cost"] for r in arm_a_results])),
"cost_b": float(np.mean([r["cost"] for r in arm_b_results])),
"latency_a": float(np.mean([r["latency_ms"] for r in arm_a_results])),
"latency_b": float(np.mean([r["latency_ms"] for r in arm_b_results])),
}A/B Test: Always Large (Control) vs Cascade tau=0.75 (Treatment) Sample sizes: n_A=473, n_B=527 Quality: A=0.9248, B=0.8497, lift=-0.0751 Cost: A=0.72780, B=0.52633 Latency: A=245.6ms, B=166.7ms t-statistic: 12.865 p-value: 0.0000 Significant at alpha=0.05: True
The A/B test quantifies exactly what we gain and lose by switching from always routing to the large model to using a cascade. A small quality drop at lower cost, with statistical significance determined by the p-value. Reading these results correctly requires looking at all metrics together: a statistically significant quality drop that is also economically meaningful requires a different decision than a statistically significant but practically negligible drop.
Visualizing the Cost-Quality Tradeoff
The escalation-quality curve captures an important practical insight: the marginal benefit of escalation changes across the operating range. Low-confidence requests tend to provide the largest early gains, while further escalation gradually approaches the always-large quality ceiling and continues to add cost. This means the optimal operating point is rarely "escalate everything uncertain" but rather "escalate until marginal quality gain equals marginal cost increase."
The quality comparison by difficulty bin reveals the core insight behind cascade routing: the models are closest on the easiest requests, while the small model's quality drops much faster as difficulty increases. This means the marginal value of escalation is concentrated in the high-difficulty tail. A good cascade exploits this by escalating only when it detects difficulty signals that predict poor small-model quality. Requests in the easy region can safely stay at the small model; the routing overhead on those requests is the price you pay for savings on the difficult ones.
Key Parameters
The key parameters for a cascade routing system are:
- confidence_threshold: The minimum confidence the small model must report to return its answer without escalation. Lower values route more to the small model (cheaper but lower quality); higher values escalate more frequently (better quality but higher cost).
- cost_ratio: The ratio of large model cost to small model cost per token. Determines how aggressively to optimize for the small model.
- escalation_cap: An optional maximum escalation rate, useful when you have a hard budget constraint. If your budget allows at most 20% of requests to reach the large model, set this cap explicitly rather than hoping the threshold happens to produce 20% escalation.
- routing_overhead_ms: The latency added by the routing decision itself. Must be small relative to the small model's response time to avoid negating the latency benefits of routing easy requests to the fast model.