Smart routing: pluggable LLM router with kNN routing memory + preference collection

[njbrake]

Note: this issue was drafted by Claude via back and forth with @njbrake. The reasoning and technical decisions are his; the prose is Claude's.

Summary

Add a router to otari that sends each request to the cheapest or fastest model
that is still good enough for the task, and gets better at that decision over
time by learning from the user's own preferences. Ships as a pluggable interface
with a default reference implementation, so the routing algorithm is swappable.

Tracks: mozilla-ai/otari-ai#1039

Motivation

The gateway is the only component that sees every request, model, dollar, and
latency number, so it is the right place to decide which model should handle a
given prompt. Routing simple prompts to a smaller model can cut cost
dramatically (RouteLLM reports about 95% of GPT-4 quality at about 14% of the
GPT-4 calls) without the caller changing anything.

Prior art and how this differs

• LiteLLM Router (OSS) load-balances across deployments of the same model
  group via latency-based, least-busy, and cost-based strategies.
  cost-based-routing picks the cheapest deployment and is not quality-aware: it
  does not decide whether a weaker model would suffice for a prompt.
  https://docs.litellm.ai/docs/routing
• LiteLLM Auto-Router (built on semantic-router) is content-aware, but it
  matches prompts against hand-authored utterance routes with a similarity
  threshold. The routes are configured by the operator, not learned, and there is
  no cost-quality objective. https://docs.litellm.ai/docs/proxy/auto_routing
  (An open request for preference-aligned routing exists:
  Auto-router - Content-Aware Preference-Aligned Routing BerriAI/litellm#25703)
• Closed products (Martian, NotDiamond) train per-user routers from eval data
  but are not open and not embeddable in a self-hosted gateway.

otari's proposal: a per-tenant router that learns from the user's own ranked
preferences and live traffic, with no hand-authored routes, an explicit
cost-versus-quality dial, a cascade fallback, and a pluggable backend so other
routers (including a trained one) can be dropped in.

Approach

1. RouterBackend protocol (same pattern as SandboxBackend /
   WebSearchBackend): route(ctx) -> RoutingDecision returns an ordered
   candidate model list plus a confidence; record_outcome(...) for online
   learning. The routing algorithm lives behind this seam.
2. Default impl: kNN "routing memory." No training pipeline. Embed the task
   signal, keep a per-tenant store of (embedding, model, quality, cost), and route
   by a cache-aware, cost-biased vote over the nearest neighbors:
   score(m) = mean_quality(m) - alpha * effective_cost(m) - switch_penalty(incumbent -> m).
   alpha is the single cost-versus-quality dial. Every preference a user submits is
   one more neighbor, so it improves online.
3. Preference collection. POST /router/preferences/compare fans a prompt out
   to N models; the user ranks the responses; each ranking becomes routing-memory
   records (rank-to-scalar). This is the training signal.
4. Cascade safety net. A low-confidence route should fall through to the strong
   model. Note this composes with otari's existing fallback only in platform mode
   (see Mode strategy below); in standalone, v1 either picks a single model or we
   add a standalone multi-attempt executor.

Routing granularity and prompt-cache economics

Per-request routing on the last turn is wrong for agent traces, and comparing models
on list price can make routing net-negative. Two linked issues, one answer.

• Single completion vs agent trace. A single completion is a trace of length one;
  route it every time. An agent trace is many acompletion calls in a loop where the
  last message is often a tool result, not the task. otari's mcp_tool_loop already
  locks in one provider and re-calls with the same model each round, so routing should
  decide once at the start of the trace and stay sticky (granularity=trace_sticky,
  the default). Per-step switching is an advanced, opt-in mode.
• Prompt-cache economics. The cross-request saving lever is provider prompt/prefix
  caching (Anthropic cache_control, OpenAI automatic prefix caching, Gemini context
  caching), which is per-model. Switching models forfeits the cached prefix and
  reprocesses it at full price on the new model. A cached read on an expensive model
  (around 0.1x its input price for Anthropic) can be cheaper than uncached input on a
  "cheaper" model, so the cost term must be effective cost (cache-aware), with a
  switch penalty for leaving a warm incumbent. Routing is therefore decided at
  cache-cold boundaries (trace start, single completions), where switching is free.

otari does not manage prompt caching today (it only passes Anthropic cache_control
through), so estimating cache state per provider is net-new work.

Why a kNN default instead of a trained classifier

Recent work ("Rethinking Predictive Modeling for LLM Routing: When Simple kNN Beats
Complex Learned Routers", arXiv:2505.12601) shows a non-parametric kNN over query
embeddings matches or beats trained matrix-factorization and BERT routers, with no
training pipeline and online updates for free. A trained router (RouteLLM or
NotDiamond style) becomes an optional plugin behind the same interface, not core
scope.

References: RouteLLM (arXiv:2406.18665), routing and cascading survey
(arXiv:2603.04445), RouterBench (arXiv:2403.12031).

Decisions to make first

• Mode strategy (this one is load-bearing). otari's multi-model fallback runs
  only in platform mode. Verified in code: chat.py forks on is_platform_mode,
  which is true only when a platform token (OTARI_AI_TOKEN) is set; the platform
  branch iterates a route.attempts list that is built by the otari.ai platform
  service, while the standalone branch is explicitly single-attempt with no fallback
  and no config to enable one. So in self-hosted standalone (the main OSS case) there
  is no attempt-walker for the router's ordered candidate list to feed into. v1 either
  targets platform mode first (walker exists) or adds a standalone multi-attempt
  executor so routing can cascade. Pick before coding.
• Embedding placement. Hosted embeddings add roughly 50 to 200 ms on the request
  hot path; sub-50 ms routing needs a local embedding model (heavier deploy).
• Routing granularity. trace_sticky (recommended default) vs per-step switching.

Validation gate (do this before building)

Confirm kNN routing memory actually predicts the "cheap model suffices" set on
otari-shaped traffic: sample representative prompts, fan them out to 2 or 3 models,
score the responses (LLM judge to start), and check whether nearest-neighbor lookup
over prompt embeddings predicts that set on held-out prompts. If yes, build with
confidence; if no, we learned it for the cost of a script.

Scope

v1

• RouterBackend protocol + RoutingContext / RoutingDecision / RoutingOutcome
  (including trace and cache fields), wired into dispatch behind config (default off)
• Trace-aware sticky routing (trace_sticky) with a cache-aware effective_cost
  term and a switch penalty, so routing is never net-negative on cached agent traces
• RouterPreference / RoutingMemory ORM entities + Alembic migration, with a
  per-tenant vector cap, eviction, and an embedding_model tag
• KnnRoutingMemory default backend (embed via any-llm, cosine kNN, the score above)
• POST /router/preferences/compare + /rank, docs, OpenAPI update
• Config: OTARI_ROUTER_BACKEND, OTARI_ROUTER_ALPHA, OTARI_ROUTER_K,
  OTARI_ROUTER_EMBEDDING_MODEL, OTARI_ROUTER_CONFIDENCE_FLOOR,
  OTARI_ROUTER_SEED_COUNT, OTARI_ROUTER_GRANULARITY, OTARI_ROUTER_SWITCH_PENALTY

Fast-follow

• Standalone multi-attempt execution (so standalone routing can cascade, not just
  pick one model)
• Step-level (per-call) routing with the cache-aware cost model
• Passive learning from live traffic plus an implicit quality signal
• Savings-analysis command over UsageLog ("a cheaper model would have sufficed on X%")
• Cost-quality eval harness (RouterBench style)
• pgvector or ANN backend for scale; trained-router plugin (RouteLLM or NotDiamond style)

Acceptance criteria

• A user runs the compare flow, ranks responses, and on prompts similar to ones
  already ranked sees otari route to a cheaper model
• An agent trace routes once and stays on one model for the loop; routing does not
  increase cost on a cached multi-step trace
• A second RouterBackend can be registered without touching the chat route
• Cost and latency for a routed request are tracked and reportable
• On a held-out traffic slice, the router reaches a target fraction of
  strong-model quality at a fraction of the cost, reported as a cost-quality curve

Notes

A full design doc exists with the architecture, interfaces, and open questions; can
attach or paste on request. A local harness for driving the gateway and collecting
preferences (fan a prompt out to N models and rank) is in progress.

[njbrake]

Note: drafted by Claude via discussion with Nathan; the reasoning and decisions are his, the prose is Claude's.

Design-discussion update, folding in feedback from besaleli and Davide. Grouped by what
it changes.

Strategy backends, not only the learned router

The pluggable RouterBackend should ship as a small menu, not just the learned kNN backend:

• none (passthrough), cheapest (lowest-price model in the pool; free since pricing
  already exists), fastest (lowest observed latency), and knn (the learned default).
• The deterministic strategies have no cold start, so they serve users who don't want to
  collect preferences and they backstop kNN while its store warms up.
• Caveat: fastest needs per-call latency, which otari does not persist today (UsageLog has
  tokens / cost / status, no duration). The RoutingMemory record already carries latency,
  so the learned path gets it; a standalone fastest strategy needs a latency column added.
• This is distinct from LiteLLM-style replica load-balancing (distributing across N
  deployments of the same model). otari routes provider:model with one endpoint and has no
  deployment-group concept, so that is a separate, larger feature and out of scope here.

Task metadata to scope retrieval and scoring (v1 schema)

A 0..1 preference score only means something within a task: "good enough" for drafting an
email is a different bar than for summarizing a PDF. kNN over embeddings already separates
tasks implicitly, so explicit task metadata earns its keep where embedding similarity diverges
from task similarity (same surface / different task, or same task / scattered surface). Plan:
optional task_id / tags on RoutingContext, stored per record, used as a soft filter
(boost same-task neighbors) with fallback to global kNN, plus within-task score scoping.
Untagged requests behave as today. Avoid hard partitioning: it loses valid cross-task
transfer and multiplies cold start.

User-defined tasks are the manual precursor to learned clustering (fast-follow)

Storing a task id per record is forward-compatible whether the id is user-assigned now or
discovered by clustering the embedding space later. That is the convergence point with the
federated / clustering routing work: same data model, manual now, learned later.

Judge-assisted labeling + label provenance (provenance field in v1)

The compare / rank flow is human-in-the-loop. An LLM judge can pre-fill a best-guess ranking
the user confirms or edits, and later backfill sparse data with judge labels (the RouteLLM
judge augmentation, which improved every router in their paper). Keep it honest: store whether
each label came from a human or a judge, keep human rankings as ground truth, and down-weight
judge labels. Judge-assisted labeling is fast-follow; the provenance field lands in v1.

Embeddings on long conversations (trace_sticky)

trace_sticky embeds the task signal (the first user turn, bounded), not the growing
transcript, so it sidesteps long-context embedding degradation by construction: routing keys
off the short early intent, not a 50k-token rolling context (which would truncate and
mean-pool-wash-out, i.e. be worse). Cap the embedding input, and treat a mid-trace task drift
as a re-route boundary rather than re-embedding every turn.

kNN internals evolve behind the interface

Confirming Davide's point: exact cosine now, with the swap to approximate nearest neighbors
(already in the fast-follow list) and any embedding-model change isolated behind
RouterBackend, so neither touches the interface.

Out of scope: agent orchestration (planner + executors)

This design does not, by itself, solve the agentic use case where one large model plans and
smaller models execute. That is orchestration (decomposing a task into a planning step and
executor subagents), which is a separate concern from routing: routing picks a model per call,
orchestration picks the structure and what each subagent's task is. The two compose cleanly,
each subagent is its own task and routes independently via the task metadata above, but the
orchestration layer itself is out of scope for this issue. For reference, what otari has today
is the single-model tool-loop (mcp_loop); there is no multi-agent orchestration layer.

Clarification: a None / abstain prediction

Raised relative to the federated version, where None means "no aggregate data for this model
in this region." Not needed here. The useful version already exists: on low confidence, route
to the requested / default model (confidence_floor), i.e. abstain from routing rather than
"call no model."

---

Net v1 additions: strategy-backend menu (cheapest now, fastest after a latency column),
optional task_id / tags in the store schema plus soft filtering, and a label-provenance
field. Fast-follow: learned clustering, judge-augmented labeling, ANN / pgvector.
Out of scope: replica load-balancing, agent orchestration.