purushah commented on code in PR #1042:
URL: https://github.com/apache/flink-agents/pull/1042#discussion_r3929058328
##########
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:
You're right — the round-1 restructure unified the declarations but left the
execution path uneven. Adopted as you sketched it: a plan-side RoutingExecutor
implemented by rule/judge/custom executors, all type dispatch confined to a
registry, and the resolver reduced to executor lookup, the single
route:<router> persistence boundary, normalization and events.
On your question — the only hard constraint I found is that durable calls
replay as a flat, order-matched sequence, so they can't nest: the judge's
durable chat is issued as a flat sibling before the decision record rather than
inside it. That fits inside the interface as one contract point
(issuesDurableCalls()), not a separate resolver path — happy to rename/reshape
that flag if you'd prefer a different form.
##########
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:
Agreed — with the plan-time checks in place the backstop is duplicate
machinery. Removed judgeSetupMisconfiguration(); the plain-judge contract is
now documented on Strategies.llm, and a setup that binds a prompt/tools
dynamically without declaring them abstains to the default (visible in the
routing events).
##########
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:
You're right — making compileRules public was the wrong direction. Moved
compilation/caching/execution into the plan-side RuleBasedRoutingExecutor,
moved the custom-executor instantiation out of the API layer on the same
grounds, and ModelRouter now carries only candidates + the declaration. Rule
shape/pattern validation lives in the RoutingStrategy constructor, so the
builder, the router and plan validation share one path.
--
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]
