GitHub user purushah created a discussion: [Discussion] Pluggable model routing
— API & design (PR #852 follow-up)
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 #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. #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**
```java
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)**
```java
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)**
```java
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)**
```java
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)**
```java
// 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**
```java
@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?
GitHub link: https://github.com/apache/flink-agents/discussions/897
----
This is an automatically sent email for [email protected].
To unsubscribe, please send an email to: [email protected]
