purushah commented on code in PR #1042:
URL: https://github.com/apache/flink-agents/pull/1042#discussion_r4009455680
##########
plan/src/main/java/org/apache/flink/agents/plan/AgentPlan.java:
##########
@@ -691,6 +698,209 @@ 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).
+ *
+ * <p>{@code RULE_BASED}: rule shape and pattern validity are enforced by
the {@link
+ * RoutingStrategy} constructor invoked below (the single
declaration-validation path, so
+ * diagnostics match the builder's); this arm additionally checks that
every rule key names a
+ * declared candidate.
+ */
+ private void validateRoutingStrategies() {
+ 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;
+ case RULE_BASED:
+ validateRuleKeys(
+ provider.getName(),
+ strategy,
+
descriptor.getArgument(ModelRouter.CANDIDATES_KEY));
+ break;
+ default:
+ break;
+ }
+ }
+ }
+
+ /**
+ * Rule declarations are static constraints like the judge checks above:
the fluent builder
+ * rejects a bad one at build(), but a descriptor read back from a plan
(deserialized or
+ * hand-built) never went through the builder. Without this arm they would
surface only per
+ * record at request time — inside the durable call — where the IGNORE
error policy silently
+ * drops every matching record. Rule shape, value types and pattern
validity were already
+ * enforced by the {@link RoutingStrategy} constructor (regardless of the
'candidates' shape);
+ * the key-vs-candidate check here mirrors build().
+ */
+ private static void validateRuleKeys(
+ String routerName, RoutingStrategy strategy, Object candidates) {
+ if (candidates == null) {
+ // A missing 'candidates' argument gets the router constructor's
own message
+ // ("requires at least one candidate").
+ return;
+ }
+ if (!(candidates instanceof List)) {
+ // Fail here, not per record: the router constructor's unchecked
read would turn a
+ // mis-shaped value into a raw ClassCastException inside the
durable call.
+ throw new IllegalArgumentException(
+ String.format(
+ "Model router '%s' declares '%s' as %s; expected a
list of model"
+ + " names.",
+ routerName,
+ ModelRouter.CANDIDATES_KEY,
+ candidates.getClass().getSimpleName()));
+ }
+ Object rules = strategy.getArguments().get(RoutingStrategy.ARG_RULES);
+ if (!(rules instanceof Map)) {
+ return;
+ }
+ for (Object ruleKey : ((Map<?, ?>) rules).keySet()) {
+ if (!((List<?>) candidates).contains(ruleKey)) {
+ throw new IllegalArgumentException(
+ String.format(
+ "Model router '%s' has routing rule key '%s'
which is not one of"
+ + " the candidates %s.",
+ routerName, ruleKey, candidates));
+ }
+ }
+ }
+
+ private static void validateJudge(
+ String routerName, RoutingStrategy strategy, Map<String,
ResourceProvider> chatModels) {
+ String judgeModel = (String)
strategy.getArguments().get(RoutingStrategy.ARG_JUDGE_MODEL);
+ ResourceProvider judgeProvider = chatModels.get(judgeModel);
+ if (judgeProvider == null) {
+ 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.",
+ routerName, judgeModel));
+ }
+ // The judge must be a plain chat model — nothing may rewrite the
judge conversation.
+ // Only descriptor-carried bindings are visible here; a setup that is
not introspectable
+ // at plan time surfaces its bindings on the judge's normal chat path.
+ if (!(judgeProvider instanceof JavaResourceProvider)) {
Review Comment:
Fixed. `validateJudge` now reads the descriptor from a
`JavaResourceProvider` or a `PythonResourceProvider` (`descriptorOf`), so a
judge declared with `pythonClazz` gets the same prompt/tools/skills check.
Test-first: `pythonBackedJudgeWithBoundToolsFailsAtPlanConstruction` failed
before the change, plus a positive case. `validateRoutingStrategies` uses the
same helper, so a Python-provided router is no longer skipped either.
##########
plan/src/main/java/org/apache/flink/agents/plan/actions/LlmJudgeRoutingExecutor.java:
##########
@@ -0,0 +1,450 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.flink.agents.plan.actions;
+
+import org.apache.flink.agents.api.agents.Agent;
+import org.apache.flink.agents.api.agents.AgentExecutionOptions;
+import org.apache.flink.agents.api.chat.messages.ChatMessage;
+import org.apache.flink.agents.api.chat.messages.MessageRole;
+import org.apache.flink.agents.api.chat.model.BaseChatModelSetup;
+import org.apache.flink.agents.api.chat.model.routing.RoutingCandidate;
+import org.apache.flink.agents.api.chat.model.routing.RoutingContext;
+import org.apache.flink.agents.api.chat.model.routing.RoutingDecision;
+import org.apache.flink.agents.api.chat.model.routing.RoutingStrategy;
+import org.apache.flink.agents.api.context.RunnerContext;
+import org.apache.flink.agents.api.event.ModelRoutingEvent;
+import org.apache.flink.agents.api.prompt.Prompt;
+import org.apache.flink.agents.api.resource.ResourceType;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.LinkedHashMap;
+import java.util.LinkedHashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+
+/**
+ * Executes the framework-managed LLM-as-judge strategy ({@code
Strategies.llm(...)}): the engine —
+ * not the strategy — runs the judge chat, through the normal
durable/metered/observable invoker
+ * path with the flat durable id {@code "judge:<router>"} (issued
<i>before</i> the resolver's
+ * decision record — see the sequencing contract on {@link RoutingExecutor}),
then derives the
+ * decision from the verdict as a pure function.
+ *
+ * <p>Failure policy: an unparseable or non-candidate verdict 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.
+ *
+ * <p>The judge must be a plain chat model (no prompt, tools, or skills) —
enforced at plan
+ * construction for descriptor-carried bindings; a setup that binds them
dynamically without
+ * declaring them is outside the {@code Strategies.llm} contract (its verdicts
fail to parse and
+ * every request abstains to the default, visible in the routing events).
+ *
+ * <p>The verdict is constrained by construction: only candidate names are
accepted, so a judge that
+ * gets hijacked by instructions inside the user's request (a measured failure
mode) cannot steer
+ * routing outside the declared candidates — an unparseable or non-candidate
reply abstains to the
+ * router's default model.
+ */
+final class LlmJudgeRoutingExecutor implements RoutingExecutor {
+
+ /** Matches {@code "model": "<name>"} in the judge's JSON verdict. */
+ private static final Pattern VERDICT_JSON =
Pattern.compile("\"model\"\\s*:\\s*\"([^\"]+)\"");
+
+ /** Decision-metadata flag set when the context cap dropped part of the
conversation. */
+ static final String CONTEXT_TRUNCATED_KEY = "judge_context_truncated";
+
+ @Override
+ public boolean usesDurableExecutionInternally() {
+ return true;
+ }
+
+ @Override
+ public String decisionSource() {
+ return ModelRoutingEvent.SOURCE_LLM_JUDGE;
+ }
+
+ @Override
+ public RoutingDecision route(
+ RoutingStrategy strategy, RoutingContext context, RunnerContext
ctx) throws Exception {
+ Agent.ErrorHandlingStrategy errorStrategy =
+
ctx.getConfig().get(AgentExecutionOptions.ERROR_HANDLING_STRATEGY);
+ int numRetries = ChatModelInvoker.configuredRetries(ctx,
errorStrategy);
+ int retryWaitIntervalSec =
ChatModelInvoker.configuredRetryWaitSec(ctx, errorStrategy);
+ String judgeModel = judgeModel(strategy);
+ List<String> candidateNames = candidateNames(context);
+
+ Map<String, Object> judgeMetadata = new LinkedHashMap<>();
+ judgeMetadata.put("judge_model", judgeModel);
+ String verdictModel = null;
+ String abstainReason = null;
+ try {
+ boolean[] truncated = new boolean[1];
+ List<ChatMessage> effective = effectiveJudgeMessages(context, ctx);
+ List<ChatMessage> judgeInput =
+ buildJudgeMessages(
+ strategy,
+ context,
+ effective,
+ pinnedRenderedIndices(context.getMessages(),
effective),
+ truncated);
+ if (truncated[0]) {
+ judgeMetadata.put(CONTEXT_TRUNCATED_KEY, true);
+ }
+ ChatModelInvoker.ChatAttemptResult judgeResult =
+ ChatModelInvoker.chatWithRetries(
+ context.getRequestId(),
+ judgeModel,
+ "judge:" + context.getRouter(),
+ judgeInput,
+ Map.of(),
+ null,
+ ctx,
+ errorStrategy,
+ numRetries,
+ retryWaitIntervalSec);
+ ChatModelAction.recordAttemptRetryStats(
+ ctx,
+ context.getRequestId(),
+ judgeResult.chatModel,
+ judgeResult.retryCount,
+ judgeResult.totalRetryWaitSec);
+ ChatMessage reply = judgeResult.response;
+ // Same both-or-neither type guard as the metrics reader of these
extraArgs
+ // keys (ChatModelAction#recordChatTokenMetrics): a half-populated
or non-Number
+ // pair must not leak into the durable decision metadata.
+ Object promptTokens = reply.getExtraArgs().get("promptTokens");
+ Object completionTokens =
reply.getExtraArgs().get("completionTokens");
+ if (promptTokens instanceof Number && completionTokens instanceof
Number) {
+ judgeMetadata.put("judge_prompt_tokens", promptTokens);
+ judgeMetadata.put("judge_completion_tokens", completionTokens);
+ }
+ verdictModel = parseVerdict(reply.getContent(),
candidateNames).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,
+ context.getRequestId(),
+ 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 (ModelRoutingResolver.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;
+ }
+
+ 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());
+ }
+ return builder.build();
+ }
+ // 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.
+ return new RoutingDecision(
+ null, true, abstainReason, null, new HashMap<>(judgeMetadata),
null);
+ }
+
+ private static List<String> candidateNames(RoutingContext context) {
+ List<String> names = new ArrayList<>();
+ for (RoutingCandidate candidate : context.getCandidates()) {
+ names.add(candidate.getName());
+ }
+ return names;
+ }
+
+ /**
+ * The judge routes on what the selected model will actually receive. When
the target setup
+ * binds a {@link Prompt}, this mirrors {@code BaseChatModelSetup#chat}:
the template is
+ * rendered with the request's prompt args and prepended to the non-empty
conversation messages.
+ * The rendering anchor is the router's default candidate (or the first
candidate) — where
+ * abstains resolve, and in practice the workload-level prompt shared by
the candidates. If the
+ * anchor can't be resolved or binds no prompt, the raw message list is
used unchanged.
+ */
+ private static List<ChatMessage> effectiveJudgeMessages(
+ RoutingContext context, RunnerContext ctx) {
+ List<ChatMessage> messages = context.getMessages();
+ String anchor =
+ context.getDefaultModel().orElseGet(() ->
context.getCandidates().get(0).getName());
+ try {
+ BaseChatModelSetup setup =
+ (BaseChatModelSetup) ctx.getResource(anchor,
ResourceType.CHAT_MODEL);
+ // One shared implementation with the chat path
(prepareRequestMessages), so the
+ // judge's view cannot drift from what the selected model
receives. Candidates binding
+ // DIFFERENT prompts see their own rendering only at answer time —
the anchor (default
+ // candidate, where abstains resolve) is a documented
approximation.
+ return setup.prepareRequestMessages(messages,
context.getPromptArgs());
+ } catch (Exception unresolvable) {
+ // An unresolvable candidate surfaces on the real chat path with
its normal policy.
+ return messages;
+ }
+ }
+
+ /**
+ * Indices of effective messages that were <i>generated</i> by the
anchor's request shaping
+ * (rendered template, skill-discovery prompt) rather than taken from the
conversation —
+ * identified by object identity, since {@code prepareRequestMessages}
appends the original
+ * message instances unchanged. They carry the task definition, so the
context cap pins them.
+ */
+ private static Set<Integer> pinnedRenderedIndices(
+ List<ChatMessage> original, List<ChatMessage> effective) {
+ if (effective == original) {
+ return Set.of();
+ }
+ Set<ChatMessage> originals =
+ java.util.Collections.newSetFromMap(new
java.util.IdentityHashMap<>());
+ originals.addAll(original);
+ Set<Integer> pinned = new LinkedHashSet<>();
+ for (int i = 0; i < effective.size(); i++) {
+ if (!originals.contains(effective.get(i))) {
+ pinned.add(i);
+ }
+ }
+ return pinned;
+ }
+
+ static String judgeModel(RoutingStrategy strategy) {
+ return (String)
strategy.getArguments().get(RoutingStrategy.ARG_JUDGE_MODEL);
+ }
+
+ private static String promptTemplate(RoutingStrategy strategy) {
+ return (String)
strategy.getArguments().get(RoutingStrategy.ARG_PROMPT_TEMPLATE);
+ }
+
+ private static int maxContextChars(RoutingStrategy strategy) {
+ Object cap =
strategy.getArguments().get(RoutingStrategy.ARG_MAX_CONTEXT_CHARS);
+ return cap instanceof Number ? ((Number) cap).intValue() :
Integer.MAX_VALUE;
+ }
+
+ /**
+ * Builds the judge conversation: a system message carrying the candidates
(with their {@code
+ * describe(...)} descriptions) and the verdict contract, plus the request
under judgment.
+ *
+ * <p>The judge routes on what the selected model will actually receive:
{@code
+ * effectiveMessages} is the complete message list, with the target
setup's bound prompt already
+ * rendered when one exists (see {@code
ModelRoutingResolver#effectiveJudgeMessages}). With the
Review Comment:
Fixed, now `{@link #effectiveJudgeMessages}`.
##########
plan/src/main/java/org/apache/flink/agents/plan/routing/RoutingExecutor.java:
##########
Review Comment:
Done: the resolver, executor interface, registry and the three executors
moved to `org.apache.flink.agents.plan.routing`, with
`LlmJudgeRoutingExecutorTest` alongside. `ChatModelAction` reads the route
through public getters on `ResolvedModelRoute`; the judge executor reuses
`ChatModelInvoker`, which became public. That leaves a `plan.routing` <->
`plan.actions` reference cycle. If you'd rather not have it, the judge's retry
helper could move into `plan.routing`. I used a flat `plan.routing` rather than
mirroring `api.chat.model.routing`; happy to rename.
--
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]
