[Discussion] Pluggable model routing — API & design (PR #852 follow-up)

[purushah]

Following the review on #852 (thanks @xintongsong for the thorough look) — opening this to align on the design before going deeper on implementation. PR #852 is a working prototype/reference; below is a refined design and the considerations behind it. Feedback welcome.

1 · What & why

Let an agent route each request to the most appropriate chat model instead of pinning one — send cheap/simple asks to a small model and hard ones to a strong model (cost + quality). The prototype models this as a drop-in ChatModelRouter (a CHAT_MODEL an agent points at by name), with a pluggable selection strategy, optional fallback, and per-conversation caching.

2 · Core proposal: routing should select, not delegate

In the prototype the router is a ChatModelSetup that calls the chosen backend's chat() internally. That one fact causes the two issues raised in review: the backend call and the LLM-judge call happen inside one resource, so they're invisible to the EventLog and their metrics collapse onto the router.

Proposal: make routing a framework capability (a dedicated router resource) that returns a model name; the framework then runs the normal chat path against that model.

• Observability — the chosen backend call is an ordinary ChatRequestEvent/ChatResponseEvent in the EventLog; the decision is its own ModelRoutingEvent; an LLM judge's call goes through the standard chat path too. (This is also what keeps routing compatible with the tracing/evaluation direction in [Discussion] Agent observability: tracing & evaluation beyond metrics + event log #710 — the decision is data, not a black box.)
• Metrics — token usage attributes to the actual backend; the router only adds routing metrics (per-candidate routed count, fallback count, cache hit/miss).
• Bonus — tools/prompt/skills (cf. Agent Skills Integration Design #565) attach to the chosen model, not the router.

One detail to settle under this model: how the LLM-judge call is represented — a full ChatRequest/Response in the loop (most transparent, bigger change) vs. a lighter ModelRoutingEvent carrying its reply. I lean toward the lighter event first, with the full path as a follow-up.

3 · Public API: function-first, with built-ins

Every well-regarded routing framework I looked at makes a function the front door and ships a few built-ins — none require authoring rules in a nested map as the primary API:

Framework	Mechanism
Apache Camel (Content-Based Router)	Predicate.matches(Exchange) in .choice().when(...)
LangChain4j	QueryRouter.route(Query) + built-ins (LanguageModelQueryRouter)
Kafka Streams	split().branch(Predicate)
Spring Integration	@Router / AbstractMessageRouter
RouteLLM (LMSYS)	pluggable router types (bert / causal_llm / sw_ranking)

So the single public extension point is RoutingStrategy — a one-method interface returning a candidate name (or null to abstain). Everything else is config:

• Built-in strategies are RoutingStrategy objects handed in via small factories — Strategies.rules(...) (keyword/regex), Strategies.llm(...) (judge); later semantic/ML. Rule-based stays, but as just-a-built-in — no magic-string dispatch in code.
• Fallback and caching become config flags (fallback, cache/cache_size) — not interfaces users implement or wrap.
• Users implement only RoutingStrategy, and only for custom/ML.

4 · Framework vs. user — the target shape

Push everything mechanical into the framework; the user writes at most one function. Example shapes:

A · Common case — config only, no routing code

ModelRouter.of("small", "big")
    .strategy(Strategies.rules(Map.of("big", "\\b(code|sql|analyze)\\b")))  // built-in factory, no magic string
    .defaultModel("small")
    .fallback(true)
    .cache(true);

B · Custom — the only user code is one function (lambda)

ModelRouter.of("small", "big")
    .strategy(ctx -> ctx.firstUserMessage().length() > 200 ? "big" : "small")  // lambda
    .defaultModel("small");

C · LLM-as-router — a judge model decides (built-in)

ModelRouter.of("small", "big")
    .strategy(Strategies.llm("small"))  // judge model; candidates carry descriptions
    .defaultModel("small")
    .cache(true);   // judge runs ~once per conversation, not per tool-call round

D · ML / learned — a reusable strategy class (bring-your-own)

public class LearnedRouter implements RoutingStrategy {
    public String route(RoutingContext ctx) {
        double hardness = myModel.score(ctx.firstUserMessage());  // trained classifier / scorer
        return hardness >= 0.5 ? "big" : "small";                 // null = abstain -> default
    }
}
// wire it:
ModelRouter.of("small", "big").strategy(new LearnedRouter()).defaultModel("small");

Regex is just a rule (or a lambda)

// built-in rules with a regex (factory, not a name string):
.strategy(Strategies.rules(Map.of("big", "\\b(code|sql|prove|analyze)\\b")))
// or inline as a function, no config:
.strategy(ctx -> ctx.firstUserMessage().matches("(?is).*\\bselect\\b.*\\bfrom\\b.*") ? "big" : "small")

The agent side — identical in every case

@Action(listenEventTypes = {InputEvent.EVENT_TYPE})
public static void onInput(Event e, RunnerContext ctx) {
    String q = (String) InputEvent.fromEvent(e).getInput();
    ctx.sendEvent(new ChatRequestEvent("router", List.of(new ChatMessage(USER, q)), Map.of(), null));
}

Across all of these the agent's actions are unchanged — they just send ChatRequestEvent("router", …). Only the strategy slot differs: rules → LLM → ML is a one-line change.

Proposed API surface (names/semantics up for critique)

Member	Meaning
ModelRouter.of(String… candidates)	Start a router over these candidate model names (order matters for fallback).
.strategy(RoutingStrategy s)	The single strategy setter — takes a strategy object. Built-ins come from factories (Strategies.rules(...), Strategies.llm(...)); custom is a lambda or class instance. No name-string dispatch.
.defaultModel(String name)	Candidate used when the strategy abstains (null) or names a non-candidate. Must be one of the candidates.
.fallback(boolean)	On a chosen-model error, try the remaining candidates in declaration order (default false).
.cache(boolean) / .cacheSize(int)	Memoize the decision per conversation (sticky). Default on; bounded LRU.
RoutingStrategy#route(RoutingContext) → String	The SPI: return a candidate name, or null to abstain. (@FunctionalInterface, so a lambda works.)
RoutingContext	What the strategy sees: request messages, prompt args, candidates (name + description + metadata), and access to resources (e.g. an embedding model).

Points I'd like input on: (a) built-ins are factories that return a RoutingStrategy (Strategies.rules(...)) — the open bit is only whether we prefer that or typed builder methods (.rules(...), .llm(...)); either way, no hardcoded strategy keywords in Java. (b) should cache/fallback be simple flags as above, or richer policy objects? (c) does RoutingContext expose enough (candidate descriptions/metadata, resource access) for ML/semantic strategies?

5 · Design considerations — prior art & proposed direction

These are the decisions that shape the design. For each, I've laid out how established products already solve it and the direction I'm proposing, so we're converging on a concrete plan rather than starting from scratch. Refinements welcome.

Consideration 1 · Strategies are objects, not name strings — with built-in factories

Prior art — the split is about who dispatches on the name. In-code APIs pass objects: LangChain4j hands in a router instance (new LanguageModelQueryRouter(...)), Camel/Kafka Streams take a Predicate. Only LiteLLM selects a built-in by string (latency-based/cost-based/…) — because it's a config-driven proxy, the declarative exception, and even its built-ins are operational, not content-rules.

Proposal — no hardcoded strategy keywords in Java. One setter, strategy(RoutingStrategy); built-ins are factories (Strategies.rules(...), Strategies.llm(...)) and custom is a lambda. A name string is unavoidable only in declarative config (YAML — cf. #662 — or SQL), and there we use the fully-qualified class name (open/extensible), with a short alias as optional sugar, not a baked-in framework keyword. Keep a thin rules built-in (handy for the declarative/SQL path); operational strategies (latency/cost) are a natural future built-in set — added as more factories, still no dispatch strings.

Consideration 2 · ML/learned routing is bring-your-own, with embedding-similarity as the first built-in

Prior art — RouteLLM ships trained routers as in-process artifacts (BERT / causal-LLM / sw_ranking); vLLM Semantic Router uses an in-process embedding classifier. Learned/content routing is typically an in-process artifact or embedding classifier, not an external service.

Proposal — BYO for v1 via RoutingStrategy; make embedding-similarity (a registered embedding-model resource) the first built-in ML path (no training pipeline needed); trained-classifier-as-artifact as a follow-up.

Consideration 3 · Stay consistent with Flink SQL routing, but don't couple to it

Prior art — no LLM gateway is a SQL engine, so there's no direct analog; Flink's own ML_PREDICT (FLINK-39961) is the closest.

Proposal — keep semantics consistent with the SQL work, but don't couple or block the agents PR on it.

---

Cross-cutting — LiteLLM treats fallback / retries / timeout / caching as plain Router config, validating "fallback/cache are flags, not interfaces." (It also selects strategy by a name string, but only because it's config-driven — our in-code API passes strategy objects, matching LangChain4j/Camel.) The proposed API is squarely in line with the most-used LLM router in the wild.

PR #852 stays open as a prototype/reference. Happy to adjust it (or split it) once we align here. Thoughts?

[wenjin272]

Thank you for preparing this proposal. I will review it closely this week.