weiqingy commented on code in PR #1042:
URL: https://github.com/apache/flink-agents/pull/1042#discussion_r3911065729
##########
plan/src/main/java/org/apache/flink/agents/plan/actions/ModelRoutingResolver.java:
##########
@@ -87,20 +111,73 @@ public Class<RoutingDecision> getResultClass() {
@Override
public RoutingDecision call() throws Exception {
// Timed inside the durable call so the latency is
persisted with the
- // decision: a replayed run reports the original
strategy wall time.
+ // decision: a replayed run reports the original
strategy wall time — and
+ // the strategy is never re-executed on replay.
long start = System.nanoTime();
- RoutingDecision decision =
router.route(routingContext);
+ RoutingDecision decision = executePure(router,
strategy, routingContext);
return decision.withDecisionMs((System.nanoTime() -
start) / 1_000_000.0);
}
};
RoutingDecision decision = ctx.durableExecute(routeCallable);
- Double decisionMs = decision.getDecisionMs();
- FlinkAgentsMetricGroup actionMetrics = ctx.getActionMetricGroup();
- if (actionMetrics != null && decisionMs != null) {
-
actionMetrics.getHistogram("routingDecisionLatencyMs").update(Math.round(decisionMs));
+ recordDecisionLatency(ctx, decision);
+ return normalizeAndFinish(
+ requestId, model, router, decision,
ModelRoutingEvent.SOURCE_STRATEGY, ctx);
+ }
+
+ /** Executes the pure (engine-free) strategy types: built-in rules, or the
user's executor. */
+ private static RoutingDecision executePure(
+ ModelRouter router, RoutingStrategy strategy, RoutingContext
routingContext)
+ throws Exception {
+ switch (strategy.getType()) {
+ case RULE_BASED:
+ return executeRules(router, routingContext);
+ case CUSTOM:
+ return router.getCustomExecutor().route(strategy,
routingContext);
+ default:
+ throw new IllegalStateException(
+ "Unhandled routing strategy type: " +
strategy.getType());
}
+ }
+ /**
+ * Built-in keyword/regex rules: the first candidate whose pattern
(pre-compiled by the router)
+ * matches the most recent user message wins, in declaration order; no
match abstains so the
+ * router uses its default model.
+ */
+ private static RoutingDecision executeRules(ModelRouter router,
RoutingContext context) {
+ String text = context.lastUserMessage();
Review Comment:
The rule evaluator lost its multi-turn test in the move.
`RoutingTest.ruleMatchesLatestUserMessageNotFirst` (old `RoutingTest.java:64`)
matched the regex on the oldest USER turn and asserted abstain. Its replacement
(`RoutingTest.java:69-82`) only checks the `RoutingContext` accessors, so
nothing runs the rule evaluator now.
Every rule-routed request left in the suite has one user message, so
`firstUserMessage()` and `lastUserMessage()` return the same string. Nothing
pins the choice on this line. Would one multi-turn case be worth adding back
here?
##########
api/src/main/java/org/apache/flink/agents/api/chat/model/routing/ModelRouter.java:
##########
@@ -84,30 +96,113 @@ public ModelRouter(ResourceDescriptor descriptor,
ResourceContext resourceContex
}
this.fallbackEnabled =
Boolean.TRUE.equals(descriptor.getArgument("fallback",
Boolean.FALSE));
- String strategyClazz = descriptor.getArgument("strategy_clazz");
+ String typeTag = descriptor.getArgument(STRATEGY_TYPE_KEY);
+ if (typeTag == null || typeTag.isEmpty()) {
+ throw new IllegalArgumentException("ModelRouter requires a routing
strategy.");
+ }
Map<String, Object> strategyArgs =
- descriptor.getArgument("strategy_args",
Collections.emptyMap());
- this.strategy = instantiateStrategy(strategyClazz, strategyArgs);
+ descriptor.getArgument(STRATEGY_ARGS_KEY,
Collections.emptyMap());
+ String executorClass =
descriptor.getArgument(STRATEGY_EXECUTOR_CLASS_KEY);
+ // The declaration constructor owns the per-type argument rules, so a
structurally invalid
+ // configuration fails here (resource construction) with the same
message as at build().
+ this.strategy =
+ new RoutingStrategy(
+ RoutingStrategyType.fromTag(typeTag), strategyArgs,
executorClass);
+ this.compiledRules = compileRules(this.strategy);
+ this.customExecutor = instantiateCustomExecutor(this.strategy);
}
- @SuppressWarnings("unchecked")
- private static RoutingStrategy instantiateStrategy(String clazz,
Map<String, Object> args)
+ /**
+ * Instantiates the user's {@link CustomRoutingExecutor} once per router
instance (routers are
+ * cached per TaskManager), preserving executor instance state across
requests. The construction
Review Comment:
"cached per TaskManager" does not match the lifecycle. `ResourceCache` is an
instance field of `ActionExecutionOperator.java:101` on main, assigned at
`:187`, so the scope is one subtask. At parallelism 8 on one TaskManager a user
gets 8 executor instances, not 1.
This is the only place documenting the new extension point's lifetime, and
the same sentence invites executors to keep instance state. `:152` repeats it.
Is "cached per subtask" the accurate phrase for both?
##########
plan/src/main/java/org/apache/flink/agents/plan/actions/ModelRoutingResolver.java:
##########
@@ -111,12 +188,310 @@ 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 ({@link LlmJudgeRoutingExecutor}). 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.
Cancellation
+ * propagates and is never persisted as a routing outcome.
+ */
+ private static ResolvedModelRoute resolveViaJudge(
+ UUID requestId,
+ String model,
+ ModelRouter router,
+ RoutingStrategy strategy,
+ 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);
+ String judgeModel = LlmJudgeRoutingExecutor.judgeModel(strategy);
+
+ Map<String, Object> judgeMetadata = new LinkedHashMap<>();
+ judgeMetadata.put("judge_model", judgeModel);
+ String verdictModel = null;
+ String abstainReason = null;
+ // Runtime backstop to the plan-time check: plan validation only sees
descriptor-carried
+ // bindings, so a setup that binds a prompt/tools/skills at the
instance level (or via a
+ // non-Java provider) is caught here. Same policy as a failed judge
call: FAIL is loud
+ // (a config error should not hide), IGNORE abstains so the default
model keeps answering.
+ String misconfigured = judgeSetupMisconfiguration(judgeModel, ctx);
+ if (misconfigured != null && errorStrategy !=
Agent.ErrorHandlingStrategy.IGNORE) {
+ throw new IllegalStateException(misconfigured);
+ }
+ if (misconfigured != null) {
+ abstainReason = misconfigured;
+ } else
+ try {
+ boolean[] truncated = new boolean[1];
+ List<ChatMessage> effective = effectiveJudgeMessages(router,
routingContext, ctx);
+ List<ChatMessage> judgeInput =
+ LlmJudgeRoutingExecutor.buildJudgeMessages(
+ strategy,
+ routingContext,
+ effective,
+
pinnedRenderedIndices(routingContext.getMessages(), effective),
+ truncated);
+ if (truncated[0]) {
+
judgeMetadata.put(LlmJudgeRoutingExecutor.CONTEXT_TRUNCATED_KEY, true);
+ }
+ ChatModelInvoker.ChatAttemptResult judgeResult =
+ ChatModelInvoker.chatWithRetries(
+ requestId,
+ judgeModel,
+ "judge:" + model,
+ judgeInput,
+ 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);
+ }
+ if (completionTokens != null) {
+ judgeMetadata.put("judge_completion_tokens",
completionTokens);
+ }
+ verdictModel =
+ LlmJudgeRoutingExecutor.parseVerdict(
+ reply.getContent(),
router.getCandidateNames())
+ .orElse(null);
+ abstainReason =
+ verdictModel == null ? "judge verdict was not a
candidate name" : null;
+ } catch (InterruptedException cancellation) {
+ // Cancellation surfacing from the between-retries backoff
sleep.
+ Thread.currentThread().interrupt();
+ throw cancellation;
+ } catch (ChatModelInvoker.ChatAttemptFailed failure) {
+ ChatModelAction.recordAttemptRetryStats(
+ ctx,
+ requestId,
+ failure.chatModel,
+ failure.retryCount,
+ failure.totalRetryWaitSec);
+ // Cancellation surfacing from inside the judge attempt (the
invoker wraps every
+ // attempt exception): it must propagate, never persist as a
routing outcome.
+ if (isCancellation(failure)) {
+ Thread.currentThread().interrupt();
+ throw failure;
+ }
+ // A judge that exhausted its retries honors the request's
error-handling strategy,
+ // exactly like a throwing rule/custom strategy (see class
javadoc).
+ if (errorStrategy != Agent.ErrorHandlingStrategy.IGNORE) {
+ throw failure;
+ }
+ abstainReason = "judge call failed: " + failure.error;
+ }
+
+ RoutingDecision computed;
+ if (verdictModel != null) {
+ RoutingDecision.Builder builder =
+ RoutingDecision.builder(verdictModel).reason("llm judge
verdict");
+ for (Map.Entry<String, Object> entry : judgeMetadata.entrySet()) {
+ builder.metadata(entry.getKey(), entry.getValue());
+ }
+ computed = builder.build();
+ } else {
+ // Persisted as a real abstain: replay resolves to the router's
*current* default, so
+ // a candidate-set change across a restart degrades gracefully
(like the strategy
+ // path) instead of failing the non-candidate guard.
+ computed =
+ new RoutingDecision(
+ null, true, abstainReason, null, new
HashMap<>(judgeMetadata), null);
+ }
+ final RoutingDecision toStore =
+ computed.withDecisionMs((System.nanoTime() - start) /
1_000_000.0);
+
+ // Persist under the standard route id: on recovery the stored
decision (with its original
+ // judge-inclusive wall time) replays; the judge chat above replays
from its own durable
+ // record, so the recomputation feeding this call is deterministic.
+ RoutingDecision decision =
+ ctx.durableExecute(
Review Comment:
The two new replay tests (`ChatModelActionRoutingTest.java:1062`, `:1090`)
seed a custom-strategy router, so they take the durable call at `:122`. The
judge path stores its decision here instead, and the judge work runs before
this call. Nothing seeds `route:<router>` for a judge router. Does the judge
path want its own test, or do you read it as covered by the shared code?
##########
python/flink_agents/api/execution_environment.py:
##########
@@ -237,6 +237,14 @@ def add_resource(
AgentsExecutionEnvironment
The environment to register the resource.
"""
+ if resource_type == ResourceType.MODEL_ROUTER:
Review Comment:
nit: this guard is byte-identical to `agent.py:168-175`, message text
included, but only the `Agent` path has a test
(`test_model_router_not_supported.py:28-30`). The examples reach this copy
through `env.add_resource`, for example `rag_agent_example.py:42`. One shared
helper plus a parametrized test over both entry points would cover it. Would
you rather do that here, or leave it for the Python routing work?
##########
plan/src/main/java/org/apache/flink/agents/plan/AgentPlan.java:
##########
@@ -691,6 +698,156 @@ private void checkNoRouterModelNameClash(ResourceProvider
provider) {
}
}
+ /**
+ * Static routing-strategy constraints fail at plan construction — never
per record. The
+ * strategy travels as a language-neutral type tag plus arguments, so
validation reads
+ * declaration data directly: no reflective instantiation, whose failure
modes previously let a
+ * misconfigured strategy skip validation entirely.
+ *
+ * <p>{@code LLM_JUDGE}: the judge chat model must be registered (a typo'd
name would otherwise
+ * fail-and-abstain on every request, silently disabling routing — cf.
{@link
+ * #checkNoRouterModelNameClash}), and its descriptor must bind no prompt,
tools, or skills — a
+ * bound prompt would prepend an (unfilled) task prompt ahead of the
verdict contract, bound
+ * tools divert the reply into tool calls, and skills inject both a
discovery prompt and tools,
+ * each silently breaking verdict parsing on every request.
+ *
+ * <p>{@code CUSTOM}: the executor class must exist, implement {@link
CustomRoutingExecutor},
+ * and expose a supported constructor — checked without instantiation, so
plan construction
+ * never runs user constructors (or their static initializers).
+ */
+ 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;
+ }
+ String typeTag =
descriptor.getArgument(ModelRouter.STRATEGY_TYPE_KEY);
+ if (typeTag == null) {
+ // Fail here, not per record on the TaskManager: ModelRouter's
constructor
+ // unconditionally rejects a descriptor without a strategy,
and a throwing
+ // construction is never cached, so it would re-throw on every
routed request.
+ throw new IllegalArgumentException(
+ String.format(
+ "Model router '%s' declares no routing
strategy ('%s' missing"
+ + " from its descriptor).",
+ provider.getName(),
ModelRouter.STRATEGY_TYPE_KEY));
+ }
+ // The declaration constructor re-validates the per-type argument
rules, so a
+ // structurally invalid configuration (e.g. a judge without a
judge model) fails
+ // plan construction with the same message as build().
+ RoutingStrategy strategy =
+ new RoutingStrategy(
+ RoutingStrategyType.fromTag(typeTag),
+ descriptor.getArgument(
+ ModelRouter.STRATEGY_ARGS_KEY,
Collections.emptyMap()),
+
descriptor.getArgument(ModelRouter.STRATEGY_EXECUTOR_CLASS_KEY));
+ switch (strategy.getType()) {
+ case LLM_JUDGE:
+ validateJudge(provider.getName(), strategy, chatModels);
+ break;
+ case CUSTOM:
+ validateCustomExecutor(provider.getName(), strategy);
+ break;
+ default:
Review Comment:
`RULE_BASED` falls to `default: break` here, so rule keys are never checked
at plan time. The javadoc at `:701` says static routing constraints fail at
plan construction and never per record, and a rule key naming a non-candidate
is just as static.
Today it only surfaces at request time, at the candidate check in
`ModelRoutingResolver.java:153`, inside the durable call. `IGNORE` then drops
every matching record.
The note at `RoutingStrategy.java:111-112` says the builder validates rule
keys "where the candidates are in hand". That covers the fluent path, but a
descriptor read back from a plan skips the builder. Worth a check in this arm
too?
##########
plan/src/main/java/org/apache/flink/agents/plan/actions/ModelRoutingResolver.java:
##########
@@ -111,12 +188,310 @@ 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 ({@link LlmJudgeRoutingExecutor}). 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.
Cancellation
+ * propagates and is never persisted as a routing outcome.
+ */
+ private static ResolvedModelRoute resolveViaJudge(
+ UUID requestId,
+ String model,
+ ModelRouter router,
+ RoutingStrategy strategy,
+ 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);
+ String judgeModel = LlmJudgeRoutingExecutor.judgeModel(strategy);
+
+ Map<String, Object> judgeMetadata = new LinkedHashMap<>();
+ judgeMetadata.put("judge_model", judgeModel);
+ String verdictModel = null;
+ String abstainReason = null;
+ // Runtime backstop to the plan-time check: plan validation only sees
descriptor-carried
+ // bindings, so a setup that binds a prompt/tools/skills at the
instance level (or via a
+ // non-Java provider) is caught here. Same policy as a failed judge
call: FAIL is loud
+ // (a config error should not hide), IGNORE abstains so the default
model keeps answering.
+ String misconfigured = judgeSetupMisconfiguration(judgeModel, ctx);
+ if (misconfigured != null && errorStrategy !=
Agent.ErrorHandlingStrategy.IGNORE) {
+ throw new IllegalStateException(misconfigured);
+ }
+ if (misconfigured != null) {
+ abstainReason = misconfigured;
+ } else
+ try {
+ boolean[] truncated = new boolean[1];
+ List<ChatMessage> effective = effectiveJudgeMessages(router,
routingContext, ctx);
+ List<ChatMessage> judgeInput =
+ LlmJudgeRoutingExecutor.buildJudgeMessages(
+ strategy,
+ routingContext,
+ effective,
+
pinnedRenderedIndices(routingContext.getMessages(), effective),
+ truncated);
+ if (truncated[0]) {
+
judgeMetadata.put(LlmJudgeRoutingExecutor.CONTEXT_TRUNCATED_KEY, true);
+ }
+ ChatModelInvoker.ChatAttemptResult judgeResult =
+ ChatModelInvoker.chatWithRetries(
+ requestId,
+ judgeModel,
+ "judge:" + model,
+ judgeInput,
+ 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) {
Review Comment:
nit: these two guards check `!= null` only, so a non-`Number` value under
`promptTokens` or `completionTokens` goes into the durable metadata as-is. The
existing reader of the same two `extraArgs` keys checks the type first:
`ChatModelAction.java:225-232` uses `instanceof Number`. The new test passes
`Integer`, so it does not reach the gap. Should this line and `:293` match that
guard?
--
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]