weiqingy commented on code in PR #1042:
URL: https://github.com/apache/flink-agents/pull/1042#discussion_r3890813268
##########
plan/src/test/java/org/apache/flink/agents/plan/actions/ChatModelActionRoutingTest.java:
##########
@@ -345,6 +345,255 @@ void abstainWithoutDefaultUsesFirstCandidate() throws
Exception {
assertThat(ctx.resolvedChatModels).containsExactly("small");
}
+ @Test
+ void llmJudgeVerdictRoutesToNamedCandidate() throws Exception {
+ ModelRouter router =
+ new ModelRouter(
+ ModelRouter.of("small", "big")
+ .describe("big", "code and sql")
+ .strategy(Strategies.llm("judge"))
+ .defaultModel("small")
+ .build(),
+ null);
+ FakeRunnerContext ctx =
+ new FakeRunnerContext(router)
+ .register(
+ "judge",
+ new FakeChatModel(
+ new ChatMessage(
+ MessageRole.ASSISTANT,
"{\"model\": \"big\"}")))
+ .register("big", new FakeChatModel());
+ ChatModelAction.processChatRequestOrToolResponse(
+ new ChatRequestEvent(
+ "router", List.of(new ChatMessage(MessageRole.USER,
"write some sql"))),
+ ctx);
+
+ ModelRoutingEvent event = ctx.routingEvent();
+ assertThat(event).isNotNull();
+ assertThat(event.getSelectedModel()).isEqualTo("big");
+
assertThat(event.getDecisionSource()).isEqualTo(ModelRoutingEvent.SOURCE_LLM_JUDGE);
+ assertThat(event.getMetadata()).containsEntry("judge_model", "judge");
+ // every shipped payload key is read or asserted somewhere (v1 review
lesson)
+ assertThat(event.getMetadata()).containsKey("decision_source");
+ // judge call is durable under its own id; the decision persists under
the route id
+ assertThat(ctx.durableCallIds).contains("judge:router",
"route:router");
Review Comment:
This checks the durable ids. But `FakeRunnerContext` (`:212-221`) always
calls through in both `durableExecute` and `durableExecuteAsync`, and there is
no way to seed a stored result. So no test takes the replay path.
That leaves two things unproven: a stored abstain resolving to the router's
current default (`ModelRoutingResolver.java:279-281`), and `decision_ms`
surviving replay (`:290-292`), which is the reason the PR body gives for the
second durable write.
A small map in the fake keyed on `callable.getId()` would open that path.
Would that be worth adding here, or is replay meant to be covered end-to-end?
##########
plan/src/main/java/org/apache/flink/agents/plan/AgentPlan.java:
##########
@@ -691,6 +696,79 @@ private void checkNoRouterModelNameClash(ResourceProvider
provider) {
}
}
+ /**
+ * An LLM-judge router references its judge chat model by name, resolved
at request time. A
+ * typo'd judge name would not fail the job: every judge call would fail
and abstain to the
+ * default model, silently disabling routing. All resources are known
here, so fail at
+ * plan-construction time instead (cf. {@link
#checkNoRouterModelNameClash}).
+ */
+ private void validateLlmJudgeReferences() {
+ if (resourceProviders == null) {
+ return;
+ }
+ Map<String, ResourceProvider> routers =
resourceProviders.get(ResourceType.MODEL_ROUTER);
+ if (routers == null) {
+ return;
+ }
+ Map<String, ResourceProvider> chatModels =
+ resourceProviders.getOrDefault(ResourceType.CHAT_MODEL,
Collections.emptyMap());
+ for (ResourceProvider provider : routers.values()) {
+ if (!(provider instanceof JavaResourceProvider)) {
+ continue;
+ }
+ ResourceDescriptor descriptor = ((JavaResourceProvider)
provider).getDescriptor();
+ if (descriptor == null || descriptor.getInitialArguments() ==
null) {
+ continue;
+ }
+ // Runtime parity: the resolver dispatches on the *instantiated*
strategy
+ // (instanceof + getJudgeModel()), so validation instantiates the
same way —
+ // subclasses with their own constructors or overrides are judged
by what they
+ // actually return, not by raw descriptor args. Anything that
cannot be instantiated
+ // here is left for the runtime's own instantiation error.
+ LlmJudgeRoutingStrategy judge =
+ instantiateIfLlmJudge(
+
descriptor.getArgument(ModelRouter.STRATEGY_CLAZZ_KEY),
+ descriptor.getArgument(
+ ModelRouter.STRATEGY_ARGS_KEY,
Collections.emptyMap()));
+ if (judge == null) {
+ continue;
+ }
+ String judgeModel = judge.getJudgeModel();
+ if (judgeModel == null || !chatModels.containsKey(judgeModel)) {
+ throw new IllegalArgumentException(
+ String.format(
+ "Model router '%s' uses Strategies.llm with
judge model '%s', but no"
+ + " CHAT_MODEL resource with that name
is registered.",
+ provider.getName(), judgeModel));
+ }
+ }
+ }
+
+ private static LlmJudgeRoutingStrategy instantiateIfLlmJudge(
+ String strategyClazz, Map<String, Object> strategyArgs) {
+ if (strategyClazz == null) {
+ return null;
+ }
+ try {
+ // Gate BEFORE constructing: plan construction must not run
arbitrary custom-strategy
+ // constructors (or their static initializers — hence
initialize=false); only classes
+ // that opted into judge semantics are instantiated, and those via
the runtime's own
+ // contract (ModelRouter.instantiateStrategy) so validation judges
exactly the object
+ // the runtime will use. Anything that fails to construct here is
left to the
+ // runtime's own, louder error.
+ Class<?> clazz =
+ Class.forName(
+ strategyClazz, false,
Thread.currentThread().getContextClassLoader());
+ if (!LlmJudgeRoutingStrategy.class.isAssignableFrom(clazz)) {
+ return null;
+ }
+ return (LlmJudgeRoutingStrategy)
+ ModelRouter.instantiateStrategy(strategyClazz,
strategyArgs);
+ } catch (Exception | LinkageError notInstantiableHere) {
Review Comment:
This catch treats "not a judge" and "a judge with bad arguments" the same
way. A judge subclass without `judge_model` throws from `super(args)`, lands
here, and returns `null`. So `validateLlmJudgeReferences` skips that router,
including the judge-model check. `build()` misses it too, because the guard at
`ModelRouter.java:266` compares the exact class name.
`judgeSubclassIsValidatedByAssignability` passes valid args, so this case is
untested.
Under the default `FAIL` you get a loud error per record. Under `IGNORE` the
request is dropped with a warning, so routing is quietly off for the whole job.
That is what the javadoc at `:699-703` is trying to prevent.
Would it help for `instantiateIfLlmJudge` to tell the two apart, returning
`null` only when the class is absent or not a judge, and letting a bad-args
failure through?
##########
plan/src/main/java/org/apache/flink/agents/plan/actions/ResolvedModelRoute.java:
##########
@@ -116,7 +116,7 @@ Map<String, Object> buildResponseMetadata(String
finalModel, List<String> triedM
routing.put("final_model", finalModel);
routing.put("candidates", new ArrayList<>(this.candidates));
routing.put(
- "decision_source",
+
org.apache.flink.agents.api.event.ModelRoutingEvent.DECISION_SOURCE_KEY,
Review Comment:
`ModelRoutingResolver.java:272-273` and `:283-284` already put
`decision_source` into the decision metadata, and that map goes in under
`"metadata"` five lines below. So when a judge-routed request falls back, the
same block says `decision_source = "fallback"` here and
`metadata.decision_source = "llm_judge"` inside. The fallback event at
`ChatModelAction.java:413-424` splits the same way.
Nothing in production reads the nested one. Reads go to the event attribute
(`ModelRoutingEvent.java:153`) or to this top-level key
(`ChatModelActionRoutingTest.java:702`, `:909`). The only assertion on the
nested map is the `containsKey` at `:377`.
Since the PR body treats `decision_source` as consumer-visible, would
dropping the two resolver puts be enough? Or does the judge path want its own
key there, under a name that cannot clash?
##########
plan/src/main/java/org/apache/flink/agents/plan/actions/ModelRoutingResolver.java:
##########
@@ -111,12 +143,247 @@ public RoutingDecision call() throws Exception {
if (!router.isCandidate(selectedModel)) {
throw new IllegalStateException(
String.format(
- "Routing strategy for router '%s' returned
non-candidate model '%s'; candidates are %s.",
+ "Routing decision for router '%s' selected
non-candidate model '%s'; candidates are %s.",
model, selectedModel,
router.getCandidateNames()));
}
- decisionSource = ModelRoutingEvent.SOURCE_STRATEGY;
+ decisionSource = concreteSource;
+ }
+ return finish(requestId, model, router, decision, selectedModel,
decisionSource, ctx);
+ }
+
+ /** Records the decision latency histogram sample (also for decisions the
guards reject). */
+ private static void recordDecisionLatency(RunnerContext ctx,
RoutingDecision decision) {
+ Double decisionMs = decision.getDecisionMs();
+ FlinkAgentsMetricGroup actionMetrics = ctx.getActionMetricGroup();
+ if (actionMetrics != null && decisionMs != null) {
+
actionMetrics.getHistogram("routingDecisionLatencyMs").update(Math.round(decisionMs));
}
+ }
+
+ /**
+ * LLM-as-judge path (framework-managed, per discussion #897): the engine
runs the judge chat
+ * itself through the normal durable/metered/observable invoker path —
durable id {@code
+ * "judge:<router>"} so a recovered run replays the original verdict
instead of re-calling the
+ * judge (with a durable action-state store configured; without one the
judge re-runs on replay,
+ * like any non-deterministic strategy) — then derives the decision from
the verdict as a pure
+ * function. The decision (including its wall time, which covers the judge
call) is persisted
+ * under the standard {@code "route:<router>"} durable call, preserving
the replay-fingerprint
+ * property. Verdict abstains are persisted <i>as abstains</i>, so a
replay after a
+ * candidate-set change re-resolves to the current default exactly like
the strategy path.
+ *
+ * <p>Failure policy: an unparseable or non-candidate verdict always
abstains to the router's
+ * default model. A judge call that exhausts its retries honors the
request's error-handling
+ * strategy, exactly like a throwing rule/custom strategy: {@code FAIL}
surfaces the outage
+ * loudly, {@code IGNORE} degrades to the default with the cause recorded.
Interrupts
+ * (cancellation) propagate and are never persisted as routing outcomes.
+ */
+ private static ResolvedModelRoute resolveViaJudge(
+ UUID requestId,
+ String model,
+ ModelRouter router,
+ LlmJudgeRoutingStrategy judge,
+ RoutingContext routingContext,
+ RunnerContext ctx)
+ throws Exception {
+ long start = System.nanoTime();
+ Agent.ErrorHandlingStrategy errorStrategy =
+
ctx.getConfig().get(AgentExecutionOptions.ERROR_HANDLING_STRATEGY);
+ int numRetries = ChatModelInvoker.configuredRetries(ctx,
errorStrategy);
+ int retryWaitIntervalSec =
ChatModelInvoker.configuredRetryWaitSec(ctx, errorStrategy);
+ Map<String, Object> judgeMetadata = new LinkedHashMap<>();
+ judgeMetadata.put("judge_model", judge.getJudgeModel());
+ String verdictModel = null;
+ String abstainReason = null;
+ // A misconfigured judge follows the same policy as a failed judge
call: FAIL is loud
+ // (a config error should not hide), IGNORE abstains so the default
model keeps
+ // answering. Under IGNORE a replayed request is unaffected either way
— the stored
+ // decision below wins over the freshly computed abstain.
+ String misconfigured =
judgeSetupMisconfiguration(judge.getJudgeModel(), ctx);
+ if (misconfigured != null && errorStrategy !=
Agent.ErrorHandlingStrategy.IGNORE) {
+ throw new IllegalStateException(misconfigured);
+ }
+ if (misconfigured != null) {
+ abstainReason = misconfigured;
+ } else {
+ try {
+ ChatModelInvoker.ChatAttemptResult judgeResult =
+ ChatModelInvoker.chatWithRetries(
+ requestId,
+ judge.getJudgeModel(),
+ "judge:" + model,
+ judge.buildJudgeMessages(routingContext),
+ Map.of(),
+ null,
+ ctx,
+ errorStrategy,
+ numRetries,
+ retryWaitIntervalSec);
+ ChatModelAction.recordAttemptRetryStats(
+ ctx,
+ requestId,
+ judgeResult.chatModel,
+ judgeResult.retryCount,
+ judgeResult.totalRetryWaitSec);
+ ChatMessage reply = judgeResult.response;
+ Object promptTokens = reply.getExtraArgs().get("promptTokens");
+ Object completionTokens =
reply.getExtraArgs().get("completionTokens");
+ if (promptTokens != null) {
+ judgeMetadata.put("judge_prompt_tokens", promptTokens);
Review Comment:
`judge_prompt_tokens` and `judge_completion_tokens` are new here and appear
only at these two writes. `FakeChatModel.chat`
(`ChatModelActionRoutingTest.java:95-107`) never sets `extraArgs`, so both
branches are dead in every test and neither key has been produced once.
The PR body lists both under compatibility impact, and the comment at `:376`
says every shipped payload key is read or asserted somewhere.
One `FakeChatModel` outcome with `extraArgs = Map.of("promptTokens", 12,
"completionTokens", 3)` plus an assertion would close it. Would you rather do
that, or narrow what the comment claims?
##########
python/flink_agents/plan/tests/test_agent_plan_cross_language.py:
##########
@@ -411,6 +411,42 @@ def
test_python_preserves_conf_data_types_and_event_ordering() -> None:
assert list(restored.actions) == ["first", "second"]
+def test_python_can_deserialize_plan_with_java_llm_judge_router() -> None:
Review Comment:
This goes Python out and Python back in: `model_dump_json()` then
`model_validate_json()`. Java is never involved. If Java emitted a different
key or shape for `strategy_clazz` / `strategy_args`, this would still pass. So
it does not yet back the PR body's line about proving Java plans deserialize in
Python.
The snapshot helper in this file (`_SNAPSHOT_DIR` at `:52`, used by the test
at `:293`) could give that claim something real. Is pointing this at a
Java-produced plan practical, or would narrowing the claim be easier?
--
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]
