[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.

[wenjin272]

Thanks for the clear and comprehensive proposal @purushah. I have a few implementation details that I’d like to confirm:

1. My understanding is that users only need to send a ChatRequestEvent whose model name refers to a registered ModelRouter. ChatModelAction would then detect the router, invoke ModelRouter.route() to resolve the concrete chat model, emit a ModelRoutingEvent, and continue through the existing chat/tool-call loop using the selected model. In this design, is ModelRoutingEvent intended purely as an observability event for logging, tracing, and evaluation, with no built-in consumer and no second ChatRequestEvent being dispatched? Or do you expect it to drive a separate step that re-dispatches the request to the selected model?

2. In [api][routing] Pluggable in-chat LLM routing (ChatModelRouter + RoutingStrategy) #852, CachingStrategy keys the routing decision by the first user message because the router is itself a ChatModelSetup and is invoked again on every tool-call round. With ModelRouter separated from the chat model, would it be cleaner for ChatModelAction to associate the RoutingDecision with the initialRequestId and store the selected concrete model in the existing tool-call context? The router would then run once per ReAct loop, and all subsequent ToolResponseEvents would reuse that model until the final ChatResponseEvent. In that case, the routing decision is part of the loop execution state, and the cache introduced in [api][routing] Pluggable in-chat LLM routing (ChatModelRouter + RoutingStrategy) #852 may no longer be necessary. Is there another use case for caching routing decisions in the new design, or could cache and cacheSize be removed?

3. Is returning only a String from ModelRouter.route() sufficient? A ModelRoutingEvent may need to record additional information about the decision, such as the reason, score, strategy metadata, or the LLM judge’s response. Would it be more appropriate for route() to return a structured RoutingDecision instead of only the selected model name?

Regarding the three open questions in the proposal:

1. I prefer keeping RoutingStrategy as the single extension point, with built-in strategies provided through factories such as Strategies.rules(...) and Strategies.llm(...). This keeps built-in and custom strategies under the same abstraction and preserves extensibility without expanding the ModelRouter API for every new strategy.

2. I think fallback and similar options can remain simple configuration flags for now. We can introduce richer policy objects later when there are concrete requirements that cannot be expressed through configuration. For cache specifically, however, I think we should first clarify whether it is still necessary in the new design, as discussed above.

3. The proposed information exposed by RoutingContext seems sufficient for the initial scope. Request messages, prompt arguments, candidate descriptions and metadata, and access to resources should cover rule-based, LLM-based, and custom semantic or ML strategies. We can extend it later when a concrete strategy requires additional context.

[purushah]

Thanks @wenjin272, I agree with your understanding and suggestions.

1. ModelRoutingEvent should be observability-only in v1. It should be used for logging, tracing, evaluation, and metrics context. It should not have a built-in consumer, and it should not trigger a second ChatRequestEvent. The flow would be: ChatModelAction detects that the requested model is a router, invokes ModelRouter.route(), emits ModelRoutingEvent, and then directly continues the existing chat/tool-call loop with the selected concrete model.

2. I agree that cache / cacheSize can be removed from the public v1 API. With ModelRouter separated from ChatModelSetup, ChatModelAction can associate the selected model / RoutingDecision with initialRequestId and store it in the existing tool-call context. Subsequent ToolResponseEvents should reuse that model until the final ChatResponseEvent. That gives us per-ReAct-loop stickiness without a separate LRU cache keyed by first user message.

3. I agree that returning only a String is too limited. route() should return a structured RoutingDecision, so ModelRoutingEvent can record selected model, reason, score/confidence, strategy metadata, abstain/default behavior, and optional LLM judge output. We can still keep the common case simple with helpers like RoutingDecision.of("big") and RoutingDecision.abstain().

I also agree with keeping RoutingStrategy as the single extension point, built-ins via factories such as Strategies.rules(...) and Strategies.llm(...), fallback as a simple config flag for now, and the current RoutingContext scope for v1.

[wenjin272]

Thanks for the reply. No further comments from my side.

[weiqingy]

+1 to @wenjin272's questions, especially around the runtime/event contract and whether caching is still needed in the framework-owned design.

Thanks for writing this up. I like the direction of treating routing as a framework capability rather than a specialized ChatModelSetup. That seems like a promising way to address the observability and metrics concerns raised on #852, while keeping the agent-side usage simple.

A few additional questions came to mind while reading:

1. Could we make the event/control-flow contract more explicit?
   If a user sends ChatRequestEvent("router", ...), would ChatModelAction resolve the selected model and call it directly, or would the framework create/record a selected-model chat request as a separate step? If it is a direct call, I wonder whether the selected backend invocation would still be visible as an ordinary model invocation in the EventLog. If it is represented as a second/internal chat request, it may be helpful to define request IDs, parent/original request IDs, and whether user-defined actions listening to chat requests can observe it.

2. What role should ModelRoutingEvent play?
   Should it be purely an observability event for logging/tracing/evaluation, or is it expected to drive another framework step that dispatches the request to the selected model? I think either shape can work, but making this explicit would help clarify the runtime call path and avoid ambiguity for users who consume events.

3. Would a structured RoutingDecision be more useful than returning only a String?
   A selected model name may be enough for execution, but the routing event may want to record why the model was selected, whether the strategy abstained, whether default/fallback was used, score/confidence, strategy metadata, or an LLM judge response. Maybe the common custom case can still stay lightweight with helpers, but the SPI could return a structured decision so this information is not lost.

4. Is router-level caching still needed once routing is owned by the framework?
   In [api][routing] Pluggable in-chat LLM routing (ChatModelRouter + RoutingStrategy) #852, caching helped because the router was itself a chat model and could be invoked again on each tool-call round. If routing moves into ChatModelAction, could it route once for the initial request, associate the selected concrete model with initialRequestId, and reuse that model for subsequent tool-call rounds? That might avoid a separate LRU keyed by the first user message, which could be tricky when different conversations share the same first message but differ in prompt args, tools, tenant, or metadata.

5. How should LLM-as-router be represented?
   A rule-based strategy feels like a pure selection function, while an LLM judge is itself a model invocation with latency, retries, metrics, durability, and EventLog implications. In the current shape, a resource/strategy does not have RunnerContext.sendEvent(...), and a direct judge.chat(...) call would bypass the current action-level EventLog and token-recording path. Would it make sense for the judge call to go through a framework-managed observable path from the beginning, or should the first version focus on the basic routing contract and leave fully observable LLM routing as a follow-up?

6. How should routing decisions behave during recovery/replay?
   Routing makes the selected backend a runtime decision. If a job fails after the original decision but before the run finishes, should recovery replay the original selected model, or is it acceptable to re-run the strategy and possibly choose a different backend? This seems especially important because the current chat durable call identity is generic to the chat operation rather than to the selected model/decision. My intuition is that the decision should become replayable data, not re-derived data, so a recovered run reaches the same backend or can detect that it would not. Would RoutingDecision, ModelRoutingEvent, or the existing action/durable state be the right place to persist that selected model and related metadata?

7. Could fallback semantics be clarified relative to existing retry/error handling?
   ChatModelAction already has retry/ignore/fail behavior. It would be useful to know whether fallback happens before retries, after retries are exhausted, or only for certain exception categories. For observability, it may also help if the routing/response data records both the initially selected model and the model that ultimately produced the response.

8. What is the deployable shape for custom strategies?
   The lambda example is very ergonomic for local Java code, but I wonder how that function reaches TaskManagers once the agent plan is serialized. Existing function/resource paths tend to move named classes or descriptors rather than closures. Is the lambda form intended mainly as an embedded/local convenience, with a named RoutingStrategy class as the deployable shape? Or should custom routing follow a named-method hook similar to actions/tools?

9. What ResourceType should a framework router use?
   If ChatModelAction needs to detect the router, should the router still be registered as CHAT_MODEL, or should there be a distinct router resource type? Keeping it as CHAT_MODEL preserves the agent-side experience but may keep some of the type-level conflation that the framework design is trying to remove, and it also raises questions like whether routers can point to other routers. A distinct type would make the contract clearer, but it would need the usual Java/Python/YAML parity work.

10. Should abstain and invalid candidate be treated differently?
    Using defaultModel when a strategy returns null seems intuitive. If a strategy returns a non-candidate model name, though, could that indicate a typo or buggy strategy? I wonder whether failing clearly, or at least recording it as an invalid decision, would be safer than treating it the same as abstain.

11. Could we record the Python/YAML scope explicitly?
    On [api][routing] Pluggable in-chat LLM routing (ChatModelRouter + RoutingStrategy) #852, the Python story was intentionally narrow: Python users could reference the Java router/strategies by FQCN, while a native Python router/SPI was deferred. Does that remain the intended v1 scope for the framework-level design? If so, it may be worth recording as a deliberate non-goal, together with the trigger for when native Python routing should be added.

Overall, I'm supportive of the framework-level router direction and the single RoutingStrategy extension point with built-in factories. The main thing I'd love to understand better is the runtime/event contract: when routing runs, how the selected model is stored for tool-call rounds and recovery, what exactly ModelRoutingEvent represents, and how selected backend calls remain observable with clear IDs and metrics attribution.

[purushah]

Thanks @weiqingy for the detailed questions. I agree these are the right contracts to clarify before changing the implementation. My updated thinking is:

1. Event/control-flow contract

The user-facing API should stay simple: users send ChatRequestEvent("router", ...). ChatModelAction should detect that the name refers to a ModelRouter, resolve a RoutingDecision, emit ModelRoutingEvent, and then directly call the selected concrete chat model through the existing chat/tool-call loop.

For v1, I would avoid creating a second user-visible ChatRequestEvent. That keeps user-defined actions from observing an internal dispatch and avoids adding parent/original request ID semantics in the first version.

2. Role of ModelRoutingEvent

I propose ModelRoutingEvent as an observability event only. It should record the router, selected model, candidates, reason/metadata, default/fallback status, and timing if available. It should not be the event that drives dispatch to the selected model.

3. Structured RoutingDecision

I agree returning only String is too limited. route() should return RoutingDecision. The simple case can remain easy with helpers like RoutingDecision.of("big") and RoutingDecision.abstain(), but the structured result can preserve reason, score/confidence, metadata, invalid/abstain status, and optional LLM judge output.

4. Router-level caching

I would remove public cache/cacheSize from v1. In the old PR, caching was needed because the router was itself a ChatModelSetup and could be invoked again on each tool-call round. If routing is owned by ChatModelAction, routing should run once per initialRequestId / ReAct loop, and later ToolResponseEvent rounds should reuse the stored selected model.

5. LLM-as-router

I agree LLM routing is not just a pure function. A judge model call has latency, retries, token usage, metrics, durability, and EventLog implications. The best design is for the framework to manage this path so the judge decision is durable and observable. If that is too much for the first patch, I would keep v1 focused on the basic routing contract and document fully observable LLM routing as a follow-up.

6. Recovery/replay

RoutingDecision should be replayable data, not re-derived data. ChatModelAction should store the decision with the initialRequestId / ReAct loop state before invoking the selected backend. With actionStateStoreBackend enabled, this should be persisted through the existing ActionStateStore and coordinated by Flink checkpoints through recovery markers and sequence numbers. On recovery, if the decision exists, reuse it and do not call route() again. If the stored selected model is no longer valid, fail clearly or report a recovery mismatch.

7. Fallback vs retry

My proposed v1 behavior is: first route to the selected model, then apply the existing retry policy for that model. If retries are exhausted and router fallback is enabled, try the remaining candidates in declaration order. The metadata should record both the initially selected model and the final model that produced the response.

8. Deployable custom strategies

The lambda examples are useful for API readability and local examples, but the deployable shape should be a named RoutingStrategy class or descriptor that can be serialized/shipped to TaskManagers. That fits better with the existing resource/action deployment model.

9. ResourceType

I lean toward a distinct MODEL_ROUTER resource type, while keeping the user-facing API unchanged. Users can still put the router name in ChatRequestEvent.model. ChatModelAction can resolve it as a router before resolving a normal CHAT_MODEL, or we can reject duplicate names to avoid ambiguity.

10. Abstain vs invalid candidate

I agree these should be different. null/abstain should use defaultModel. But returning a non-candidate model name likely means a typo or buggy strategy, so it should be recorded as an invalid routing decision and fail clearly by default.

11. Python/YAML scope

For v1, I would record this explicitly as Java-first framework routing. Python/YAML can reference the Java router/strategy descriptors if supported by the existing resource path, but native Python router SPI and YAML aliases can be follow-ups unless maintainers want full parity in the first patch.

---

Overall, the revised target is: framework-owned routing in ChatModelAction, structured RoutingDecision, observability-only ModelRoutingEvent, route-once-per-ReAct-loop state keyed by initialRequestId, deterministic recovery, no public router cache in v1, and explicit fallback/retry/deployment scope.

[purushah]

Hi @weiqingy — thanks again for the detailed review. I've replied to all your points above. Whenever you get a chance, could you take a look and let me know if the direction works or if anything still looks off? Happy to adjust. Thanks!

[weiqingy]

Thanks for the detailed reply. The updated direction works for me at the design-discussion level.

I'm aligned with the main v1 shape: framework-owned routing in ChatModelAction, structured RoutingDecision, observability-only ModelRoutingEvent, route once per initialRequestId / ReAct loop, reuse the selected concrete model across tool rounds, and remove public cache / cacheSize from v1.

I also agree that the remaining questions are mostly implementation-contract details rather than blockers for closing this discussion. A few that seem worth attention during implementation:

1. how the selected model / routing decision is persisted and replayed, especially for mid-action recovery and when no action state store is configured;
2. whether the deployable custom-strategy shape is descriptor/class-based rather than relying on live lambda/object instances that cannot cross the plan boundary;
3. whether MODEL_ROUTER resource registration has Java/Python/YAML parity and clear duplicate-name semantics;
4. how fallback composes with the existing fail/retry/ignore behavior in ChatModelAction;
5. how routing lines up with the in-flight observability work, especially model-resource health/retry metrics and trace metadata for the selected backend.

One related point seems worth settling before implementation, since it shapes the SPI rather than only the implementation. If Strategies.llm(...) ships in v1 while its observability is deferred, the judge call likely has to sit inside a synchronous route(). Moving that judge call onto a framework-managed observable path later may require changing the extension point rather than just extending the implementation. So I think we should decide whether Strategies.llm(...) is part of v1 or explicitly a follow-up before the interface is finalized.

Given the direction is now clear and the remaining items are implementation-specific, no further comments from my side on this discussion. I'm fine with closing it and continuing the detailed review on the PR.

[purushah]

Thanks @weiqingy — good catch that this is an interface question, not just an implementation detail.

I agree that if Strategies.llm(...) shipped under a plain synchronous route(), the judge call would likely sit inside route() as a hidden judge.chat(...). Moving that later onto a framework-managed observable path could require changing the extension point, not just changing the implementation.

The direction I want for observable LLM routing is: the framework runs the judge, not the strategy. Concretely, LLM routing should be expressed as a framework-managed routing step: the routing code identifies that a judge model needs to be invoked, and ChatModelAction executes that judge call through the normal durable/metered/observable chat path before continuing with the selected backend model. That keeps the judge from being a black box, and keeps the durable execution and observability hooks in the framework where they already live.

Given that, I think Strategies.llm(...) should be an explicit follow-up rather than part of the first implementation PR. It pulls in the two-phase routing shape, judge orchestration inside ChatModelAction, and a second model call with its own durability/observability semantics. Folding that into v1 would make the PR larger and harder to review.

So for v1, I'll keep the scope to framework-owned routing in ChatModelAction, structured RoutingDecision, observability-only ModelRoutingEvent, route-once-per-initialRequestId behavior across tool rounds, fallback semantics, and rule/custom pure-function strategies through the simple route() contract. Strategies.llm(...) will be documented as a follow-up, implemented through the framework-managed observable path rather than as a hidden synchronous model call inside route().

Agreed that the rest are implementation-contract details for the PR. Happy to close this discussion and continue there. Thanks again for the careful review.