weiqingy commented on code in PR #1042: URL: https://github.com/apache/flink-agents/pull/1042#discussion_r4001798226
########## 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: nit: `effectiveJudgeMessages` moved into this class (`:199`), so this `ModelRoutingResolver#effectiveJudgeMessages` reference is out of date. ########## 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: This return also skips a judge from a `PythonResourceProvider`. A Java agent gets one by declaring a chat model with `pythonClazz` (`:340-341`, `:346-348`). The cross-language e2e does this with `tools` (`ChatModelCrossLanguageAgent.java:88-92` on main). Plan validation never checked these judges. The runtime `judgeSetupMisconfiguration` check did, and it is removed now. So nothing rejects a Python judge with a prompt, tools, or skills anymore. `PythonResourceProvider` also has `getDescriptor()`, and the Python setup uses the same three argument names. Could the check below cover it too? -- 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]
