purushah opened a new pull request, #964:
URL: https://github.com/apache/flink-agents/pull/964
Linked issue: #897 (design discussion). Supersedes #852.
### Purpose of change
Today an agent must name a concrete chat model in each `ChatRequestEvent`;
any routing logic lives in user code and is not first-class in the runtime.
Most requests are easy and a small model handles them fine, but a minority
really need the strong model — so a fixed choice either overpays on the easy
majority or under-serves the hard ones. This PR makes per-request routing a
framework capability: an agent sends its request to a router, and the runtime
picks one of several candidate models.
My first attempt at this (#852) made the router a special `ChatModelSetup`
that called the chosen backend itself. Review rightly pointed out that this
breaks the two things the framework is good at: token metrics collapse onto the
router instead of the real backend, and the actual model call disappears from
the EventLog. The redesign discussed in #897 fixes this by making the router
**select instead of delegate**: the router only returns a model *name*, and
`ChatModelAction` runs the normal chat path against that model. A routed
request produces the same events and the same per-model metrics as if the agent
had named the chosen model directly.
From the user side it stays simple:
```java
env.addResource("router", ResourceType.MODEL_ROUTER,
ModelRouter.of("small", "big")
.describe("small", "fast and cheap; chit-chat, simple facts")
.describe("big", "strong; code, SQL, analysis")
.strategy(Strategies.rules(Map.of("big",
"\\b(code|sql|analyze)\\b")))
.defaultModel("small")
.fallback(true)
.build());
// the agent just sends its ChatRequestEvent to "router"
```
A few behaviors were worth getting right, and they came out of the #897
discussion:
- **The decision is durable.** The strategy runs inside its own durable call
(`"route:<requestId>:<router>"`), so on recovery the persisted decision is
replayed instead of re-running the strategy. Decision latency is measured
inside that call, so a replayed run reports the original `decision_ms`.
- **Route once per ReAct loop.** Tool-call rounds reuse the already-selected
model and carry the routing metadata onto the final response — no re-routing
mid-conversation.
- **Fallback sits on top of retries, not instead of them.** The selected
model gets its full retry budget first; only then are the remaining candidates
tried in declaration order, and a fallback emits a second `ModelRoutingEvent`
so the event log shows which model actually answered.
- **Abstain is not an error, an invalid pick is.** A strategy that has no
opinion abstains and the router's default model handles the request
(`decision_source=default`); a strategy that returns a non-candidate name fails
loudly, because that is a bug, not a runtime condition.
- **`ModelRoutingEvent` is observability-only.** It has no built-in consumer
and does not drive dispatch — removing every listener does not change which
model runs.
Since routers and chat models share the `ChatRequestEvent` namespace, one
name must not be registered as both; this is validated at the `addResource`
call site (with an `AgentPlan` backstop) so the failure points at the user's
own line.
Per the discussion, LLM-as-judge routing is deliberately **not** in this PR.
`RoutingContext` exposes no chat API, so a strategy cannot make hidden model
calls; the framework-managed observable judge (the judge call running on the
normal durable/metered chat path) is the agreed follow-up. YAML support is
intentionally left out of v1 to keep the first Java API/runtime change small;
it and Python support are follow-ups.
### Tests
- 20 API unit tests (`RoutingTest`, `RoutingResourceValidationTest`): rule
matching and abstain over the latest user message, builder and registration
validation, candidate descriptions reaching strategies, `RoutingDecision`
invariants and snake_case JSON round-trips, `decision_ms` surviving replay
deserialization, event attribute normalization.
- 9 `ChatModelActionRoutingTest` integration tests against a scripted fake
chat model: routing to the matched candidate, abstain-to-default, invalid
candidate failing clearly, non-router requests keeping the legacy `"chat"`
durable id unchanged, fallback across candidates (including exhaustion), and
route-once semantics across a tool round with metadata carried to the final
response.
- A `ResourceCacheTest` case covering `hasResource` visibility of resources
inserted directly into the cache (no provider).
Verified locally (JDK 17):
- `mvn spotless:check -pl api,plan,runtime,examples`
- `mvn test -pl api,plan` — 312 + 196 tests, 0 failures
- `mvn test -pl runtime -Dtest=ResourceCacheTest`
- `mvn compile -pl examples`
### API
New public surface, all additive:
- `ResourceType.MODEL_ROUTER`
- `org.apache.flink.agents.api.chat.model.routing`: `ModelRouter` (+
builder), `RoutingStrategy`, `RoutingContext`, `RoutingDecision`,
`RoutingCandidate`, `RoutingStrategyDescriptor`, `Strategies`,
`RuleBasedRoutingStrategy`
- `ModelRoutingEvent`
- `RunnerContext.hasResource(name, type)` (default method, returns false)
Existing agents are unaffected: a request naming a plain chat model takes
the unchanged path, including the legacy durable call id.
### Documentation
- [x] `doc-needed`
- [ ] `doc-not-needed`
- [ ] `doc-included`
The Java examples show usage, but this adds public API, so proper docs are
warranted — happy to do a docs follow-up PR once the API settles in review.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
