wenjin272 commented on code in PR #1042:
URL: https://github.com/apache/flink-agents/pull/1042#discussion_r3924001769
##########
api/src/main/java/org/apache/flink/agents/api/chat/model/routing/ModelRouter.java:
##########
@@ -84,30 +99,125 @@ 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 subtask — at parallelism N that is N router (and executor)
instances, so executor
+ * instance state spans the requests of one subtask, not the whole
TaskManager. The construction
+ * contract is a {@code (Map<String,Object>)} constructor fed the
declaration's arguments, then
+ * a no-arg constructor, via the thread context classloader — plan-time
validation checks the
+ * same contract without instantiating.
+ */
+ private static CustomRoutingExecutor
instantiateCustomExecutor(RoutingStrategy strategy)
throws Exception {
- if (clazz == null || clazz.isEmpty()) {
- throw new IllegalArgumentException("ModelRouter requires a routing
strategy.");
+ if (strategy.getType() != RoutingStrategyType.CUSTOM) {
+ return null;
+ }
+ Class<?> clazz =
+ Class.forName(
+ strategy.getExecutorClass(),
+ true,
+ Thread.currentThread().getContextClassLoader());
+ if (!CustomRoutingExecutor.class.isAssignableFrom(clazz)) {
+ throw new IllegalArgumentException(
+ String.format(
+ "Custom routing executor '%s' does not implement
%s.",
+ strategy.getExecutorClass(),
CustomRoutingExecutor.class.getName()));
}
- Class<?> c = Class.forName(clazz, true,
Thread.currentThread().getContextClassLoader());
try {
- Constructor<?> ctor = c.getConstructor(Map.class);
- return (RoutingStrategy) ctor.newInstance(args);
+ return (CustomRoutingExecutor)
+
clazz.getConstructor(Map.class).newInstance(strategy.getArguments());
} catch (NoSuchMethodException noMapCtor) {
- return (RoutingStrategy) c.getConstructor().newInstance();
+ return (CustomRoutingExecutor)
clazz.getConstructor().newInstance();
+ }
+ }
+
+ /** The user's custom executor instance ({@code null} unless the strategy
type is CUSTOM). */
+ public CustomRoutingExecutor getCustomExecutor() {
+ return customExecutor;
+ }
+
+ /**
+ * The single validation/compilation path for rule maps: null/empty keys,
non-String values and
+ * invalid regex all fail here with the same diagnostics everywhere it is
called — the builder
+ * ({@code build()}), the router constructor, and plan-time validation
({@code
+ * AgentPlan#validateRuleKeys}). Called once per router instance (routers
are cached per
+ * subtask), so rule evaluation stays regex-match-only per request.
+ */
+ public static Map<String, Pattern> compileRules(RoutingStrategy strategy) {
Review Comment:
Related to the execution-abstraction concern above, I do not think
`compileRules()` belongs in the API layer. It is currently there because
`ModelRouter` stores the compiled `Map<String, Pattern>`, and it was made
public so that `AgentPlan` could reuse the same validation path. This exposes
Java-specific execution details through an otherwise language-neutral API.
Ideally, `ModelRouter` should only carry candidates and the
`RoutingStrategy` declaration. A Plan-side `RuleBasedRoutingExecutor` should
validate, compile, cache, and execute the regex patterns, with a
package-private helper shared with `AgentPlan` if needed. Is there a specific
reason this compilation must remain in the API-layer `ModelRouter`?
##########
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);
+ }
+ if (completionTokens != null) {
+ judgeMetadata.put("judge_completion_tokens",
completionTokens);
+ }
+ verdictModel =
+ judge.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 (containsInterrupt(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");
+ builder.metadata(
+ ModelRoutingEvent.DECISION_SOURCE_KEY,
ModelRoutingEvent.SOURCE_LLM_JUDGE);
+ 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.
+ Map<String, Object> abstainMetadata = new
LinkedHashMap<>(judgeMetadata);
+ abstainMetadata.put(
+ ModelRoutingEvent.DECISION_SOURCE_KEY,
ModelRoutingEvent.SOURCE_DEFAULT);
+ computed = new RoutingDecision(null, true, abstainReason, null,
abstainMetadata, 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(
+ new DurableCallable<>() {
+ @Override
+ public String getId() {
+ return routeCallId(model);
+ }
+
+ @Override
+ public Class<RoutingDecision> getResultClass() {
+ return RoutingDecision.class;
+ }
+
+ @Override
+ public RoutingDecision call() {
+ return toStore;
+ }
+ });
+ recordDecisionLatency(ctx, decision);
+ return normalizeAndFinish(
+ requestId, model, router, decision,
ModelRoutingEvent.SOURCE_LLM_JUDGE, ctx);
+ }
+
+ /**
+ * The judge must be a plain chat model — nothing may rewrite the judge
conversation. 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 breaks verdict parsing on every request. Returns a diagnostic
when misconfigured,
+ * {@code null} when the setup is plain (or cannot be resolved — an
unresolvable judge takes the
+ * ChatAttemptFailed path with its normal policy).
+ */
+ private static String judgeSetupMisconfiguration(String judgeModel,
RunnerContext ctx) {
Review Comment:
Thanks for adding the plan-time validation. As I understand it, the runtime
backstop mainly protects against a custom `ChatModelSetup` that dynamically
adds prompt/tools/skills without declaring them in its descriptor. I think that
should be treated as an invalid user implementation and documented as part of
the judge model contract.
Keeping this check in the per-request path is not ideal: it delays a static
configuration error until runtime, makes the result depend on `FAIL`/`IGNORE`,
runs before durable replay, and duplicates the Plan validation. Could we remove
`judgeSetupMisconfiguration()` and keep this validation entirely in
`AgentPlan`? If any standard descriptor-backed providers are not covered yet,
we should extend the Plan validation for those providers instead.
##########
plan/src/main/java/org/apache/flink/agents/plan/actions/ModelRoutingResolver.java:
##########
@@ -66,6 +85,11 @@ static ResolvedModelRoute resolve(
ModelRouter router = (ModelRouter) ctx.getResource(model,
ResourceType.MODEL_ROUTER);
RoutingContext routingContext =
new RoutingContext(requestId, model, messages, promptArgs,
router.getCandidates());
+ RoutingStrategy strategy = router.getStrategy();
+
+ if (strategy.getType() == RoutingStrategyType.LLM_JUDGE) {
Review Comment:
Thanks for the substantial refactoring. I think the current implementation
unifies the strategy declarations, but not the execution path:
- `ModelRoutingResolver.resolve()` still special-cases `LLM_JUDGE`.
- Rule and custom strategies go through a separate `executePure()` switch.
- `LlmJudgeRoutingExecutor` is currently a collection of static helpers
rather than an executor implementing a common contract.
I expected a Plan-side abstraction along these lines:
```java
interface RoutingExecutor {
RoutingDecision route(
RoutingStrategy strategy,
RoutingContext routingContext,
RunnerContext runnerContext)
throws Exception;
}
```
`RuleBasedRoutingExecutor`, `LlmJudgeRoutingExecutor`, and the custom
executor could all implement this contract. `ModelRoutingResolver` would only
resolve the executor for the declared type, call `route()`, and perform the
common normalization/event handling. Any type-based dispatch could be confined
to an executor factory or registry.
The LLM implementation may need different internal durable/retry
orchestration, but that does not appear to require a separate top-level path in
the resolver. Could you clarify what concrete limitation prevents the three
strategy types from sharing a common `RoutingExecutor` abstraction?
--
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]
