Route each LLM request to the cheapest model that can actually answer it.
A classifier reads the incoming request, predicts whether a small model will do, and forwards accordingly. It runs as a Databricks Model Serving endpoint that speaks the standard chat API, so pointing your app at it is a one-line change.
Deciding costs well under a millisecond once warm. The whole loop — learning what is safe to route cheap, and deploying that — is automated.
Mosaic AI Gateway already gives you usage tracking, payload logging, guardrails, rate limits, and fallbacks. What it doesn't do is look at a prompt and pick a model — its traffic splitting is percentage-based A/B. That decision function is all this adds; everything around it stays Gateway's job.
┌─────────── uc-router endpoint ───────────┐
your app ────────► │ parse → features → gate → policy │
└───────┬─────────────────────────┬────────┘
▼ ▼
cheap endpoint good endpoint
(easy requests) (hard requests, tool calls)
The router decides in-process, then forwards. Two things drive the decision:
Policy rules, checked first. Requests carrying tools always go to the good model — small models lose function-calling accuracy long before they lose prose quality. No score can override that.
A trained gate for everything else: a LightGBM model over cheap structural features of the request (length, code blocks, tool count, reasoning cues). It predicts P(cheap model is good enough); above your threshold, the request goes cheap. No gate loaded means everything goes to the good model — a router that can't score should never guess cheap.
Five jobs, run as one pipeline:
| Step | What it does |
|---|---|
00_harvest |
Read logged prompts from your AI Gateway inference table |
01_replay |
Send each one to the cheap model |
02_judge |
Ask a judge model: was the cheap answer as good? |
03_train |
Train the gate on those labels |
04_promote |
Register it and deploy the routing endpoint |
05_frontier_report then shows what each threshold buys you:
tau small% qual% saved%
0.00 100.0 60.3 90.0 ← route everything cheap: 40% of answers get worse
0.25 67.2 89.7 39.2
0.40 60.3 100.0 27.6 ← pick a row like this one
0.75 43.1 100.0 1.6
1.00 6.9 100.0 0.0
Read the first row: routing everything cheap "saves 90%" while degrading 40% of answers. That is why there's a gate, and why you choose a threshold from this table instead of trusting a default. (Numbers above are from one small sample — yours will differ.)
Your savings depend entirely on your traffic. Routing pays off when easy and hard requests are mixed. If everything you send is a hard agentic task, the gate correctly routes it all to the good model and you've gained nothing. Run steps 1–3 and read the frontier before believing any number.
Prompt caching can invert the math. Moving a request to a different model
throws away its KV cache. With tens of thousands of cached tokens per request, the
"cheaper" model can cost more. eval/frontier.py prices cache reads, and
cache_read_guard_tokens can pin heavily-cached requests to whichever tier is
already warm.
Small samples give directional numbers. Training warns below 500 labeled examples and tags the run. Don't quote a percentage off a hundred rows.
The judge grades text, not tool use. It can't assess function-calling fidelity, which is why tool requests are excluded from judging and hard-routed instead.
Streaming isn't routed yet. predict_stream is unimplemented; streaming
callers should hit a tier directly.
The endpoint adds one network hop. If you own the calling code and that matters, run the same policy in your process instead:
from uc_router.client import RoutingClient
rc = RoutingClient.from_uc("main.uc_router.router@champion")
reply = rc.chat([{"role": "user", "content": "hello"}])
print(reply["_router"]["tier"])src/uc_router/
features.py structural features — no I/O, microseconds
parse.py normalizes the two Gateway request formats
policy.py hard rules + threshold → tier decision
client.py in-process routing (no extra hop)
serving/router_model.py the endpoint: classify, then forward
jobs/ the pipeline, plus smoke_test and seed_demo_data
eval/frontier.py cost/quality sweep, cache-aware
resources/pipeline.yml job definitions
dashboards/queries.sql cost, tier mix, and overhead reporting
pip install -e ".[dev]"
pytest
ruff check .These are the non-obvious constraints this code already handles. Worth knowing if you're writing similar jobs.
| Constraint | Symptom |
|---|---|
Serverless has no __file__ in spark_python_task |
path setup relative to the script fails |
| Serverless is Spark Connect | sparkContext unavailable; use spark.addArtifact |
| RDDs are barred on serverless | use mapInPandas |
| Executors have no ambient credentials | resolve auth on the driver, pass headers in |
Serving invocations take no /api/2.0 prefix |
the versioned path 404s |
spark_python_task has no default MLflow experiment |
start_run fails on experiment ID None |
ChatParams fills in defaults you didn't set |
forwarding stop=[] or strict=false gets a 400 |
| Reasoning blocks are signed by their origin model | replaying a logged conversation elsewhere 400s |
| Model instances holding SDK clients aren't picklable | log via model-from-code instead |
Apache 2.0