From feb3eb9cfc09f8fde0bf976ffca4d797faf9cbd9 Mon Sep 17 00:00:00 2001 From: purshotam shah Date: Tue, 27 Oct 2026 00:50:47 -0700 Subject: [PATCH 1/5] [api][plan][java] Add a decision-model routing strategy for MODEL_ROUTER Strategies.decision(setup) routes with a registered DECISION_MODEL: one choice question over the candidates, their describe(...) texts as the criteria, the returned choice as the selected model and its probability as the routing score. withMinProbability(p) abstains to the default model when the chosen candidate's probability is below p; a refusal or a non-choice answer abstains as well. The call runs under the durable id decide: with the configured retry budget, execution reports and model_resource metrics; exhausting the retries fails the request like a failed judge call. Every candidate must carry a description, enforced by the router, build() and plan validation. withMaxContextChars now applies to decision declarations too. Adds DecisionModelRoutingExecutor, the decision_model routing event source, plan-time validation, docs, an OpenAI decision routing example with a native Java decision connection, and the cross-language plan snapshot for the new strategy type. Closes #1325 Generated-by: Claude Code 2.1.309 (Claude Fable 5.1) --- .../api/chat/model/routing/ModelRouter.java | 41 +- .../routing/RoutingCandidateValidator.java | 62 ++ .../chat/model/routing/RoutingStrategy.java | 106 +++- .../model/routing/RoutingStrategyType.java | 7 + .../api/chat/model/routing/Strategies.java | 37 ++ .../agents/api/event/ModelRoutingEvent.java | 3 + .../api/chat/model/routing/RoutingTest.java | 137 +++++ .../development/calling-llms/model_routing.md | 77 ++- .../docs/development/decision_models.md | 4 +- .../java/agent_plan_with_decision_router.json | 88 +++ .../openai/OpenAiDecisionRoutingExample.java | 382 ++++++++++++ .../apache/flink/agents/plan/AgentPlan.java | 118 +++- .../plan/actions/DecisionModelAction.java | 203 +------ .../plan/actions/DecisionModelInvoker.java | 239 ++++++++ .../routing/DecisionModelRoutingExecutor.java | 435 ++++++++++++++ .../plan/routing/LlmJudgeRoutingExecutor.java | 14 +- .../plan/routing/ModelRoutingResolver.java | 7 +- .../agents/plan/routing/RoutingExecutors.java | 1 + .../agents/plan/utils/CancellationUtils.java | 14 + .../plan/AgentPlanCrossLanguageTest.java | 67 +++ .../AgentPlanDecisionModelValidationTest.java | 239 ++++++++ .../actions/ChatModelActionRoutingTest.java | 441 ++++++++++++++ .../DecisionModelRoutingExecutorTest.java | 549 ++++++++++++++++++ .../tests/test_agent_plan_cross_language.py | 25 + .../context/DecisionRoutingRecoveryTest.java | 300 ++++++++++ 25 files changed, 3360 insertions(+), 236 deletions(-) create mode 100644 e2e-test/cross-language-agent-plan-snapshots/java/agent_plan_with_decision_router.json create mode 100644 examples/src/main/java/org/apache/flink/agents/examples/openai/OpenAiDecisionRoutingExample.java create mode 100644 plan/src/main/java/org/apache/flink/agents/plan/actions/DecisionModelInvoker.java create mode 100644 plan/src/main/java/org/apache/flink/agents/plan/routing/DecisionModelRoutingExecutor.java create mode 100644 plan/src/test/java/org/apache/flink/agents/plan/AgentPlanDecisionModelValidationTest.java create mode 100644 plan/src/test/java/org/apache/flink/agents/plan/routing/DecisionModelRoutingExecutorTest.java create mode 100644 runtime/src/test/java/org/apache/flink/agents/runtime/context/DecisionRoutingRecoveryTest.java diff --git a/api/src/main/java/org/apache/flink/agents/api/chat/model/routing/ModelRouter.java b/api/src/main/java/org/apache/flink/agents/api/chat/model/routing/ModelRouter.java index bdab0fad96..341a8cce1e 100644 --- a/api/src/main/java/org/apache/flink/agents/api/chat/model/routing/ModelRouter.java +++ b/api/src/main/java/org/apache/flink/agents/api/chat/model/routing/ModelRouter.java @@ -65,6 +65,9 @@ public class ModelRouter extends Resource { /** Descriptor argument naming the candidate the router lands on when the strategy abstains. */ public static final String DEFAULT_MODEL_KEY = "default_model"; + /** Descriptor argument mapping candidate names to their {@code describe(...)} descriptions. */ + public static final String CANDIDATE_DESCRIPTIONS_KEY = "candidate_descriptions"; + private final List candidates; private final String defaultModel; private final boolean fallbackEnabled; @@ -78,16 +81,6 @@ public ModelRouter(ResourceDescriptor descriptor, ResourceContext resourceContex // with the candidate message instead of failing the cast here. Object defaultModel = descriptor.getArgument(DEFAULT_MODEL_KEY); RoutingCandidateValidator.validate("ModelRouter", names, defaultModel); - Map descriptions = - descriptor.getArgument("candidate_descriptions", Collections.emptyMap()); - List parsed = new ArrayList<>(); - for (String name : names) { - parsed.add(new RoutingCandidate(name, descriptions.get(name))); - } - this.candidates = Collections.unmodifiableList(parsed); - this.defaultModel = (String) defaultModel; - this.fallbackEnabled = - Boolean.TRUE.equals(descriptor.getArgument("fallback", Boolean.FALSE)); String typeTag = descriptor.getArgument(STRATEGY_TYPE_KEY); if (typeTag == null || typeTag.isEmpty()) { throw new IllegalArgumentException("ModelRouter requires a routing strategy."); @@ -100,6 +93,30 @@ public ModelRouter(ResourceDescriptor descriptor, ResourceContext resourceContex this.strategy = new RoutingStrategy( RoutingStrategyType.fromTag(typeTag), strategyArgs, executorClass); + Object rawDescriptions = descriptor.getArgument(CANDIDATE_DESCRIPTIONS_KEY); + if (rawDescriptions != null && !(rawDescriptions instanceof Map)) { + // A hand-built or deserialized descriptor may carry any JSON type here; diagnose it + // instead of surfacing a raw ClassCastException per routed request. + throw new IllegalArgumentException( + String.format( + "ModelRouter declares '%s' as %s; expected a map of candidate name to" + + " description.", + CANDIDATE_DESCRIPTIONS_KEY, + rawDescriptions.getClass().getSimpleName())); + } + Map descriptions = + rawDescriptions == null ? Collections.emptyMap() : (Map) rawDescriptions; + RoutingCandidateValidator.requireStringDescriptions("ModelRouter", names, descriptions); + List parsed = new ArrayList<>(); + for (String name : names) { + parsed.add(new RoutingCandidate(name, (String) descriptions.get(name))); + } + RoutingCandidateValidator.requireDescriptionsForDecisionStrategy( + "ModelRouter", this.strategy, names, descriptions); + this.candidates = Collections.unmodifiableList(parsed); + this.defaultModel = (String) defaultModel; + this.fallbackEnabled = + Boolean.TRUE.equals(descriptor.getArgument("fallback", Boolean.FALSE)); } /** The configured strategy declaration (type + arguments). */ @@ -234,10 +251,12 @@ public ResourceDescriptor build() { } } } + RoutingCandidateValidator.requireDescriptionsForDecisionStrategy( + "ModelRouter", strategy, candidates, descriptions); Map args = new HashMap<>(); args.put(CANDIDATES_KEY, new ArrayList<>(candidates)); if (!descriptions.isEmpty()) { - args.put("candidate_descriptions", new HashMap<>(descriptions)); + args.put(CANDIDATE_DESCRIPTIONS_KEY, new HashMap<>(descriptions)); } if (defaultModel != null) { args.put(DEFAULT_MODEL_KEY, defaultModel); diff --git a/api/src/main/java/org/apache/flink/agents/api/chat/model/routing/RoutingCandidateValidator.java b/api/src/main/java/org/apache/flink/agents/api/chat/model/routing/RoutingCandidateValidator.java index aa21b775df..77abac9620 100644 --- a/api/src/main/java/org/apache/flink/agents/api/chat/model/routing/RoutingCandidateValidator.java +++ b/api/src/main/java/org/apache/flink/agents/api/chat/model/routing/RoutingCandidateValidator.java @@ -22,8 +22,10 @@ import javax.annotation.Nullable; +import java.util.ArrayList; import java.util.LinkedHashSet; import java.util.List; +import java.util.Map; import java.util.Set; /** @@ -78,4 +80,64 @@ public static void validate( subject, defaultModel, candidates)); } } + + /** + * A decision-model strategy asks one choice question over the candidates, so it needs at least + * two of them, and it routes on their descriptions alone — they are the question's criteria — + * so a candidate without one is invisible to the decision model. Applied by the router + * constructor, by {@code ModelRouter.Builder#build()} and by plan validation, so every + * declaration path fails with the same diagnostic; {@code subject} names the router. + */ + /** + * Every description declared for a candidate must be a String. Applied by the router + * constructor and by plan validation, so a descriptor that skipped the builder fails with the + * same diagnostic at plan construction rather than with a ClassCastException per routed + * request; {@code subject} names the router. Entries for names that are not candidates are + * ignored, as the constructor ignores them. + */ + public static void requireStringDescriptions( + String subject, List candidates, Map descriptions) { + for (String name : candidates) { + Object description = descriptions.get(name); + if (description != null && !(description instanceof String)) { + throw new IllegalArgumentException( + String.format( + "%s candidate '%s' has a %s description; expected a String.", + subject, name, description.getClass().getSimpleName())); + } + } + } + + public static void requireDescriptionsForDecisionStrategy( + String subject, + RoutingStrategy strategy, + List candidates, + Map descriptions) { + if (strategy.getType() != RoutingStrategyType.DECISION_MODEL) { + return; + } + // A choice question needs something to choose between: decision providers reject a + // single option, so a one-candidate router would fail every request on the TaskManager. + if (candidates.size() < 2) { + throw new IllegalArgumentException( + String.format( + "%s uses Strategies.decision, which requires at least two candidates" + + " (got %s).", + subject, candidates)); + } + List missing = new ArrayList<>(); + for (String candidate : candidates) { + Object description = descriptions == null ? null : descriptions.get(candidate); + if (!(description instanceof String) || ((String) description).isBlank()) { + missing.add(candidate); + } + } + if (!missing.isEmpty()) { + throw new IllegalArgumentException( + String.format( + "%s uses Strategies.decision, which requires a description for every" + + " candidate (describe(name, description)); missing for %s.", + subject, missing)); + } + } } diff --git a/api/src/main/java/org/apache/flink/agents/api/chat/model/routing/RoutingStrategy.java b/api/src/main/java/org/apache/flink/agents/api/chat/model/routing/RoutingStrategy.java index b85be31e09..92b18662af 100644 --- a/api/src/main/java/org/apache/flink/agents/api/chat/model/routing/RoutingStrategy.java +++ b/api/src/main/java/org/apache/flink/agents/api/chat/model/routing/RoutingStrategy.java @@ -49,10 +49,26 @@ public final class RoutingStrategy implements Serializable { /** * Argument key: optional cap (in characters) on the conversation context included in the judge - * input ({@code LLM_JUDGE}). Unset means the judge sees the complete message list. + * or decision-model input ({@code LLM_JUDGE}, {@code DECISION_MODEL}). Unset means the model + * sees the complete message list. */ public static final String ARG_MAX_CONTEXT_CHARS = "max_context_chars"; + /** Argument key: the decision-model setup name ({@code DECISION_MODEL}). */ + public static final String ARG_DECISION_MODEL = "decision_model"; + + /** + * Argument key: optional instructions for the decision model's choice question ({@code + * DECISION_MODEL}). Unset means the built-in instruction is used. + */ + public static final String ARG_INSTRUCTIONS = "instructions"; + + /** + * Argument key: optional minimum probability of the chosen candidate ({@code DECISION_MODEL}). + * A choice whose probability is below it abstains to the router's default model. + */ + public static final String ARG_MIN_PROBABILITY = "min_probability"; + /** Argument key: the candidate-to-regex rule map ({@code RULE_BASED}). */ public static final String ARG_RULES = "rules"; @@ -101,20 +117,33 @@ private void validate() { + ARG_PROMPT_TEMPLATE + "' must be a non-empty String when provided."); } - Object cap = arguments.get(ARG_MAX_CONTEXT_CHARS); - // Compare as long: intValue() would wrap an out-of-range Long into a tiny (or - // negative) cap, and saturate a large Double to Integer.MAX_VALUE ("no cap"). - if (cap != null - && (!(cap instanceof Number) - || ((Number) cap).longValue() <= 0 - || ((Number) cap).longValue() > Integer.MAX_VALUE)) { + validateMaxContextChars(arguments.get(ARG_MAX_CONTEXT_CHARS)); + break; + case DECISION_MODEL: + Object decisionModel = arguments.get(ARG_DECISION_MODEL); + if (!(decisionModel instanceof String) || ((String) decisionModel).isEmpty()) { + throw new IllegalArgumentException( + "Strategies.decision requires a non-empty '" + + ARG_DECISION_MODEL + + "'."); + } + Object instructions = arguments.get(ARG_INSTRUCTIONS); + if (instructions != null + && (!(instructions instanceof String) + || ((String) instructions).isBlank())) { + throw new IllegalArgumentException( + "'" + ARG_INSTRUCTIONS + "' must be a non-empty String when provided."); + } + Object minProbability = arguments.get(ARG_MIN_PROBABILITY); + if (minProbability != null + && (!(minProbability instanceof Number) + || !isProbability(((Number) minProbability).doubleValue()))) { throw new IllegalArgumentException( "'" - + ARG_MAX_CONTEXT_CHARS - + "' must be a positive integer (at most " - + Integer.MAX_VALUE - + ") when provided."); + + ARG_MIN_PROBABILITY + + "' must be a number between 0 and 1 when provided."); } + validateMaxContextChars(arguments.get(ARG_MAX_CONTEXT_CHARS)); break; case CUSTOM: if (executorClass == null || executorClass.isEmpty()) { @@ -141,23 +170,62 @@ private void validate() { } } + private static void validateMaxContextChars(@Nullable Object cap) { + // Compare as long: intValue() would wrap an out-of-range Long into a tiny (or + // negative) cap, and saturate a large Double to Integer.MAX_VALUE ("no cap"). + if (cap != null + && (!(cap instanceof Number) + || ((Number) cap).longValue() <= 0 + || ((Number) cap).longValue() > Integer.MAX_VALUE)) { + throw new IllegalArgumentException( + "'" + + ARG_MAX_CONTEXT_CHARS + + "' must be a positive integer (at most " + + Integer.MAX_VALUE + + ") when provided."); + } + } + + private static boolean isProbability(double value) { + return !Double.isNaN(value) && value >= 0.0 && value <= 1.0; + } + /** - * Returns a copy of this {@code LLM_JUDGE} declaration with a cap (in characters) on the - * conversation context included in the judge input. Purely opt-in: without it the judge - * receives the complete message list. When set, the rendered request and the SYSTEM message are - * always kept, remaining messages fill newest-first within the budget, and any truncation is - * flagged in the decision metadata ({@code judge_context_truncated}). + * Returns a copy of this {@code LLM_JUDGE} or {@code DECISION_MODEL} declaration with a cap (in + * characters) on the conversation context included in the judge or decision-model input. Purely + * opt-in: without it the model receives the complete message list. When set, the rendered + * request and the SYSTEM message are always kept, remaining messages fill newest-first within + * the budget, and any truncation is flagged in the decision metadata ({@code + * judge_context_truncated} or {@code decision_context_truncated}). */ public RoutingStrategy withMaxContextChars(int maxContextChars) { - if (type != RoutingStrategyType.LLM_JUDGE) { + if (type != RoutingStrategyType.LLM_JUDGE && type != RoutingStrategyType.DECISION_MODEL) { throw new IllegalStateException( - "withMaxContextChars applies to Strategies.llm(...) declarations only."); + "withMaxContextChars applies to Strategies.llm(...) and" + + " Strategies.decision(...) declarations only."); } Map withCap = new LinkedHashMap<>(arguments); withCap.put(ARG_MAX_CONTEXT_CHARS, maxContextChars); return new RoutingStrategy(type, withCap, null); } + /** + * Returns a copy of this {@code DECISION_MODEL} declaration with a minimum probability for the + * chosen candidate. When the decision model's probability for its choice is below this value, + * the strategy abstains and the router's default model answers; the choice and the full + * probability map stay in the decision metadata. The probability is the provider's reported + * value for the chosen candidate; the framework does not normalize the distribution. + */ + public RoutingStrategy withMinProbability(double minProbability) { + if (type != RoutingStrategyType.DECISION_MODEL) { + throw new IllegalStateException( + "withMinProbability applies to Strategies.decision(...) declarations only."); + } + Map withMin = new LinkedHashMap<>(arguments); + withMin.put(ARG_MIN_PROBABILITY, minProbability); + return new RoutingStrategy(type, withMin, null); + } + /** * A mis-shaped 'rules' value must fail loudly at declaration time: silently compiling zero * rules would disable routing (every request abstains to the default) with no diagnostic. A diff --git a/api/src/main/java/org/apache/flink/agents/api/chat/model/routing/RoutingStrategyType.java b/api/src/main/java/org/apache/flink/agents/api/chat/model/routing/RoutingStrategyType.java index 8d9bdc059d..19525ec1c5 100644 --- a/api/src/main/java/org/apache/flink/agents/api/chat/model/routing/RoutingStrategyType.java +++ b/api/src/main/java/org/apache/flink/agents/api/chat/model/routing/RoutingStrategyType.java @@ -34,6 +34,13 @@ public enum RoutingStrategyType { */ LLM_JUDGE, + /** + * Framework-managed decision-model routing: the engine asks a {@code DECISION_MODEL} resource + * one choice question over the candidates and routes on the returned choice and its + * probability. + */ + DECISION_MODEL, + /** A user-provided {@link CustomRoutingExecutor} referenced by class name. */ CUSTOM; diff --git a/api/src/main/java/org/apache/flink/agents/api/chat/model/routing/Strategies.java b/api/src/main/java/org/apache/flink/agents/api/chat/model/routing/Strategies.java index 818a9d6153..6e2bbe8af9 100644 --- a/api/src/main/java/org/apache/flink/agents/api/chat/model/routing/Strategies.java +++ b/api/src/main/java/org/apache/flink/agents/api/chat/model/routing/Strategies.java @@ -90,6 +90,43 @@ public static RoutingStrategy llm(String judgeModel, String promptTemplate) { return new RoutingStrategy(RoutingStrategyType.LLM_JUDGE, args, null); } + /** + * Decision-model routing (framework-managed): a decision model — registered like any other + * {@code DECISION_MODEL} resource — answers one choice question whose options are the candidate + * names and whose criteria are the candidates' {@code describe(...)} descriptions. The engine + * executes the decision call on its durable, metered, observable path. The returned choice + * becomes the selected model and its reported probability the routing score. A refusal answer, + * or a choice whose probability is below {@link RoutingStrategy#withMinProbability(double)}, + * abstains to the router's default model; a decision call that exhausts its retries fails the + * request. + * + *

Every candidate must carry a description: the descriptions are all the decision model sees + * of the candidates. The decision model receives the conversation text (tool calls and results + * included, media and reasoning excluded); use {@link RoutingStrategy#withMaxContextChars(int)} + * to cap the context for cost. + */ + public static RoutingStrategy decision(String decisionModel) { + Map args = new HashMap<>(); + args.put(RoutingStrategy.ARG_DECISION_MODEL, decisionModel); + return new RoutingStrategy(RoutingStrategyType.DECISION_MODEL, args, null); + } + + /** + * Like {@link #decision(String)}, with custom instructions for the choice question in place of + * the built-in instruction. + */ + public static RoutingStrategy decision(String decisionModel, String instructions) { + Map args = new HashMap<>(); + args.put(RoutingStrategy.ARG_DECISION_MODEL, decisionModel); + if (instructions == null) { + throw new IllegalArgumentException( + "Strategies.decision instructions must be non-null; use decision(decisionModel)" + + " for the built-in instructions."); + } + args.put(RoutingStrategy.ARG_INSTRUCTIONS, instructions); + return new RoutingStrategy(RoutingStrategyType.DECISION_MODEL, args, null); + } + /** * A custom executor referenced by class. The class must implement {@link CustomRoutingExecutor} * with either a {@code (Map)} constructor or a no-arg constructor. This is the diff --git a/api/src/main/java/org/apache/flink/agents/api/event/ModelRoutingEvent.java b/api/src/main/java/org/apache/flink/agents/api/event/ModelRoutingEvent.java index 4e23cdd357..52b78685a8 100644 --- a/api/src/main/java/org/apache/flink/agents/api/event/ModelRoutingEvent.java +++ b/api/src/main/java/org/apache/flink/agents/api/event/ModelRoutingEvent.java @@ -51,6 +51,9 @@ public class ModelRoutingEvent extends Event { /** The decision came from a framework-managed LLM judge verdict. */ public static final String SOURCE_LLM_JUDGE = "llm_judge"; + /** The decision came from a framework-managed decision model choice. */ + public static final String SOURCE_DECISION_MODEL = "decision_model"; + /** Attribute/metadata key carrying the decision source (events, payloads, durable records). */ public static final String DECISION_SOURCE_KEY = "decision_source"; diff --git a/api/src/test/java/org/apache/flink/agents/api/chat/model/routing/RoutingTest.java b/api/src/test/java/org/apache/flink/agents/api/chat/model/routing/RoutingTest.java index d37ef56467..5bca1588c6 100644 --- a/api/src/test/java/org/apache/flink/agents/api/chat/model/routing/RoutingTest.java +++ b/api/src/test/java/org/apache/flink/agents/api/chat/model/routing/RoutingTest.java @@ -528,4 +528,141 @@ void routerConstructionRejectsBlankCandidateName() { new ModelRouter( new ResourceDescriptor(ModelRouter.class.getName(), args), null)); } + + @Test + void strategiesDecisionRequiresDecisionModelAndNonEmptyInstructions() { + assertThrows(IllegalArgumentException.class, () -> Strategies.decision("")); + assertThrows(IllegalArgumentException.class, () -> Strategies.decision("dm", "")); + assertThrows(IllegalArgumentException.class, () -> Strategies.decision("dm", null)); + RoutingStrategy strategy = Strategies.decision("dm", "Which model should answer?"); + assertEquals(RoutingStrategyType.DECISION_MODEL, strategy.getType()); + assertEquals("dm", strategy.getArguments().get(RoutingStrategy.ARG_DECISION_MODEL)); + assertEquals( + "Which model should answer?", + strategy.getArguments().get(RoutingStrategy.ARG_INSTRUCTIONS)); + assertNull(strategy.getExecutorClass()); + } + + @Test + void withMinProbabilityRequiresProbabilityAndDecisionType() { + assertThrows( + IllegalArgumentException.class, + () -> Strategies.decision("dm").withMinProbability(1.5)); + assertThrows( + IllegalArgumentException.class, + () -> Strategies.decision("dm").withMinProbability(-0.1)); + assertThrows( + IllegalArgumentException.class, + () -> Strategies.decision("dm").withMinProbability(Double.NaN)); + assertThrows( + IllegalStateException.class, () -> Strategies.llm("judge").withMinProbability(0.5)); + assertEquals( + 0.6, + Strategies.decision("dm") + .withMinProbability(0.6) + .getArguments() + .get(RoutingStrategy.ARG_MIN_PROBABILITY)); + // the bounds are inclusive + Strategies.decision("dm").withMinProbability(0.0); + Strategies.decision("dm").withMinProbability(1.0); + // a deserialized non-numeric value is rejected + Map args = new HashMap<>(); + args.put(RoutingStrategy.ARG_DECISION_MODEL, "dm"); + args.put(RoutingStrategy.ARG_MIN_PROBABILITY, "high"); + assertThrows( + IllegalArgumentException.class, + () -> new RoutingStrategy(RoutingStrategyType.DECISION_MODEL, args, null)); + } + + @Test + void withMaxContextCharsAppliesToDecisionDeclarations() { + assertEquals( + 4096, + Strategies.decision("dm") + .withMaxContextChars(4096) + .getArguments() + .get(RoutingStrategy.ARG_MAX_CONTEXT_CHARS)); + assertThrows( + IllegalArgumentException.class, + () -> Strategies.decision("dm").withMaxContextChars(0)); + } + + @Test + void decisionStrategyTypeTagRoundTrips() { + assertEquals("decision_model", RoutingStrategyType.DECISION_MODEL.tag()); + assertEquals( + RoutingStrategyType.DECISION_MODEL, RoutingStrategyType.fromTag("decision_model")); + } + + /** The descriptions are all the decision model sees of the candidates, so each needs one. */ + @Test + void decisionStrategyRequiresEveryCandidateDescribed() throws Exception { + assertThrows( + IllegalArgumentException.class, + () -> + ModelRouter.of("small", "big") + .describe("big", "code and sql") + .strategy(Strategies.decision("dm")) + .build()); + ResourceDescriptor descriptor = + ModelRouter.of("small", "big") + .describe("small", "chit-chat") + .describe("big", "code and sql") + .strategy(Strategies.decision("dm").withMinProbability(0.6)) + .defaultModel("small") + .build(); + ModelRouter router = new ModelRouter(descriptor, null); + assertEquals(RoutingStrategyType.DECISION_MODEL, router.getStrategy().getType()); + assertEquals( + 0.6, router.getStrategy().getArguments().get(RoutingStrategy.ARG_MIN_PROBABILITY)); + assertEquals("code and sql", router.getCandidates().get(1).getDescription()); + // a descriptor that skipped the builder is held to the same rule + Map args = new HashMap<>(descriptor.getInitialArguments()); + args.put(ModelRouter.CANDIDATE_DESCRIPTIONS_KEY, Map.of("big", "code and sql")); + assertThrows( + IllegalArgumentException.class, + () -> + new ModelRouter( + new ResourceDescriptor(ModelRouter.class.getName(), args), null)); + // other strategies keep descriptions optional + new ModelRouter( + ModelRouter.of("small", "big").strategy(Strategies.llm("judge")).build(), null); + // a blank description is as good as none + assertThrows( + IllegalArgumentException.class, + () -> + ModelRouter.of("small", "big") + .describe("small", "chit-chat") + .describe("big", " ") + .strategy(Strategies.decision("dm")) + .build()); + // a choice question needs at least two options + assertThrows( + IllegalArgumentException.class, + () -> + ModelRouter.of("only") + .describe("only", "everything") + .strategy(Strategies.decision("dm")) + .build()); + // a non-String description is diagnosed by candidate name and type, for every strategy + Map typed = new HashMap<>(descriptor.getInitialArguments()); + typed.put(ModelRouter.CANDIDATE_DESCRIPTIONS_KEY, Map.of("small", "chit-chat", "big", 42)); + IllegalArgumentException typeError = + assertThrows( + IllegalArgumentException.class, + () -> + new ModelRouter( + new ResourceDescriptor(ModelRouter.class.getName(), typed), + null)); + assertTrue(typeError.getMessage().contains("'big' has a Integer description")); + // and a descriptions value that is not a map at all + Map notAMap = new HashMap<>(descriptor.getInitialArguments()); + notAMap.put(ModelRouter.CANDIDATE_DESCRIPTIONS_KEY, List.of("small", "big")); + assertThrows( + IllegalArgumentException.class, + () -> + new ModelRouter( + new ResourceDescriptor(ModelRouter.class.getName(), notAMap), + null)); + } } diff --git a/docs/content/docs/development/calling-llms/model_routing.md b/docs/content/docs/development/calling-llms/model_routing.md index d9a21eb60d..f0a03f29ed 100644 --- a/docs/content/docs/development/calling-llms/model_routing.md +++ b/docs/content/docs/development/calling-llms/model_routing.md @@ -97,9 +97,9 @@ candidate or judge, or as a chat model connection. ## Routing Strategies -Choose rules for text matching, an LLM judge for semantic selection, or a custom -executor for application-specific logic. All strategies use the same -[selection and failure rules](#default-model-and-fallback). +Choose rules for text matching, an LLM judge or a decision model for semantic +selection, or a custom executor for application-specific logic. All strategies +use the same [selection and failure rules](#default-model-and-fallback). ### Rules @@ -184,6 +184,60 @@ Two optional settings customize the judge: For example, `.strategy(Strategies.llm("judge").withMaxContextChars(8000))` sets a text budget. This is not a token limit or a hard cap on total input. +### Decision Model + +`Strategies.decision(decisionModel)` asks a registered +[decision model]({{< ref "docs/development/decision_models" >}}) one choice +question: the options are the candidate names and the criteria are their +descriptions. The answer carries a probability for every candidate, so the +router gets a confidence for its decision and can decline to route on a weak +one. Register the decision model setup and its connection first; every candidate +must have a description. + +```java +ModelRouter.of("small", "big") + .describe("small", "short factual questions and everyday requests") + .describe("big", "code, SQL, math, and multi-step reasoning") + .strategy(Strategies.decision("dm").withMinProbability(0.6)) + .defaultModel("small") + .fallback(true) + .build(); +``` + +The decision model receives the conversation as text, including tool calls and +results, rendered as for the LLM judge. Media and reasoning blocks are left out, +and the number of media blocks dropped is recorded in the decision metadata. +When the request carries no text at all, for example a media-only message, the +prompt arguments are appended to the input as the request, so this fallback can +apply where the judge, which also accepts media, sends the message itself. With +no prompt arguments either, the strategy abstains to the default model without +calling the decision model. The chosen candidate becomes the selected +model, and its reported probability becomes the routing `score`. The probability +is the provider's value for the chosen candidate; the framework does not +normalize the distribution. + +The strategy abstains, and the default model answers, when the answer is a +refusal, when it is not a choice, when the result carries no answer for the +question, or when the chosen candidate's probability is below +`withMinProbability`. For a choice below the minimum, the choice and the +full probability map stay in the metadata. A router using this strategy needs at +least two candidates. A decision call that fails after the configured retries is a request +failure, like a failed judge call; it does not select the default or try +candidate fallback. + +| Setting | Description | +|---------|-------------| +| `Strategies.decision(decisionModel, instructions)` | Replace the built-in instruction for the choice question. | +| `withMinProbability(double)` | Abstain to the default model when the chosen candidate's probability is below this value, between 0 and 1 inclusive. | +| `withMaxContextChars(int)` | Limit the conversation text with the same rules as for the LLM judge. | + +The decision model's `provider_confidence`, when the provider reports one, is +recorded in the metadata but is not used for the threshold, since not every +provider returns it. Built-in decision providers are Python, so a Java job that +uses one needs the Python environment on the TaskManagers, as for any Java +decision call. Like the other strategies, decision-model routing is available +from the Java Programmatic API only. + ### Custom Executor `Strategies.custom(...)` runs your own selection logic. Implement `CustomRoutingExecutor` as a public class with a public `(Map)` constructor that receives the declaration's arguments, or a public no-arg constructor. The class must be on the classpath of the client and the TaskManagers. @@ -229,7 +283,7 @@ ModelRouter.of("small", "big") `RoutingContext` is a read-only view of the request: the messages, prompt arguments, last and first user message, the candidates with their descriptions, the default model, the router name, and the request ID. Put tenant or workload attributes in the prompt arguments when an executor should route on them. `RoutingDecision.of(name)` selects a candidate, `RoutingDecision.builder(name)` adds a reason, score, and metadata that appear on the routing event and the response, and `RoutingDecision.abstain()` defers to the default. An exception thrown by `route()` fails the request; the strategy itself is not retried. -Custom executors must only select. The context exposes no chat API; when the decision should come from a model, use `Strategies.llm(...)` so the call gets the framework's retries, metrics, and events. Keep constructors free of side effects: an executor may be constructed more than once, for example after recovery. +Custom executors must only select. The context exposes no chat API; when the decision should come from a model, use `Strategies.llm(...)` or `Strategies.decision(...)` so the call gets the framework's retries, metrics, and events. Keep constructors free of side effects: an executor may be constructed more than once, for example after recovery. ## Call a Router @@ -274,7 +328,8 @@ behaviors: | Strategy selects a candidate | Call that candidate. | | Strategy abstains, or returns a name outside the candidates | Call `defaultModel`, or the first candidate if no default is configured. | | Declaration has an invalid pattern, non-candidate rule target or default, or duplicate/empty candidates | Reject the configuration before processing requests. | -| Custom strategy throws, or the judge call exhausts retries | Fail the request; do not treat it as abstention or try candidate fallback. | +| Custom strategy throws, or the judge or decision call exhausts retries | Fail the request; do not treat it as abstention or try candidate fallback. | +| Decision model refuses, answers something other than a choice, or its chosen candidate is below `withMinProbability` | Call `defaultModel`, or the first candidate if no default is configured. | | Candidate call fails with fallback disabled | Fail the request after applicable retries. | | Candidate call fails with fallback enabled | Try the remaining candidates in declaration order, stopping at the first successful response. | | All attempted candidates fail | Fail the request with the last candidate's error; earlier errors remain in diagnostics. | @@ -312,15 +367,16 @@ In Trace Log, these Events have `entityType: event` and |-------------|-------------| | `request_id`, `router`, `candidates` | Request identity, router resource name, and candidate names. | | `selected_model` | Selected chat model resource name. On a fallback event, the candidate that successfully answered. | -| `decision_source` | `strategy` for rules/custom selection, `llm_judge` for a judge selection, `default` for abstention, or `fallback`. | +| `decision_source` | `strategy` for rules/custom selection, `llm_judge` for a judge selection, `decision_model` for a decision model selection, `default` for abstention, or `fallback`. | | `fallback_enabled` | Whether fallback is configured, not whether it happened. | -| `reason`, `score`, `metadata` | Strategy diagnostics. Judge metadata can include `judge_model`, token counts, and `judge_context_truncated`. | -| `decision_ms` | Decision latency, including judge calls and retries. Absent on the fallback event. | +| `reason`, `score`, `metadata` | Strategy diagnostics. Judge metadata can include `judge_model`, token counts, and `judge_context_truncated`. Decision model metadata includes `decision_model`, and can include `decision_probabilities` and `provider_confidence` when the answer was a choice, token counts, `decision_context_truncated`, `decision_media_dropped`, and, when the strategy abstained on a choice, `decision_choice`. For a decision model selection, `score` is the chosen candidate's probability. | +| `decision_ms` | Decision latency, including judge or decision model calls and retries. Absent on the fallback event. | When a custom or replayed decision selects a non-candidate, normalization records `rejected_model` and, if present, `rejected_reason` in metadata and drops the score. An unrecognized fresh judge verdict instead reports the reason -`judge verdict was not a candidate name`. +`judge verdict was not a candidate name`. A decision model abstention reports +the refusal, the non-choice answer, or the probability below the minimum. A **successful** final `ChatResponseEvent` for a routed request carries a `model_routing` attribute with `selected_model`, `final_model`, `fallback_attempted`, and @@ -362,5 +418,6 @@ for the persistence and replay guarantees. - {{< gh_link file="examples/src/main/java/org/apache/flink/agents/examples/ModelRoutingExample.java" name="`ModelRoutingExample`" >}}: rule-based routing between two Ollama models. - {{< gh_link file="examples/src/main/java/org/apache/flink/agents/examples/ModelRoutingJudgeExample.java" name="`ModelRoutingJudgeExample`" >}}: the same two models with an LLM judge and candidate descriptions. - {{< gh_link file="examples/src/main/java/org/apache/flink/agents/examples/openai/OpenAiModelRoutingExample.java" name="`OpenAiModelRoutingExample`" >}}: rule-based routing between `gpt-4o-mini` and `gpt-4o` over one OpenAI connection. +- {{< gh_link file="examples/src/main/java/org/apache/flink/agents/examples/openai/OpenAiDecisionRoutingExample.java" name="`OpenAiDecisionRoutingExample`" >}}: the same two OpenAI models routed by a decision model, with a native Java connection to the OpenAI Decisions API and an agent that reports the selected model, decision source and score. -All three register their resources on the execution environment and use the {{< gh_link file="examples/src/main/java/org/apache/flink/agents/examples/agents/ModelRoutingAgent.java" name="`ModelRoutingAgent`" >}} to send requests and handle responses. +All of them register their resources on the execution environment. The first three use the {{< gh_link file="examples/src/main/java/org/apache/flink/agents/examples/agents/ModelRoutingAgent.java" name="`ModelRoutingAgent`" >}} to send requests and handle responses. diff --git a/docs/content/docs/development/decision_models.md b/docs/content/docs/development/decision_models.md index bd798c025b..96ef0191c0 100644 --- a/docs/content/docs/development/decision_models.md +++ b/docs/content/docs/development/decision_models.md @@ -37,7 +37,9 @@ registered connection by name and a model. Call a model directly when the same Action will use its answers. Use [event-driven calls](#event-driven-calls) when you want the built-in Action to -manage durable execution, retries, and execution reports. +manage durable execution, retries, and execution reports. A decision model can also +select the chat model for a routed request; see +[Model Routing]({{< ref "docs/development/calling-llms/model_routing#decision-model" >}}). ## Direct Python call diff --git a/e2e-test/cross-language-agent-plan-snapshots/java/agent_plan_with_decision_router.json b/e2e-test/cross-language-agent-plan-snapshots/java/agent_plan_with_decision_router.json new file mode 100644 index 0000000000..8b6a88168b --- /dev/null +++ b/e2e-test/cross-language-agent-plan-snapshots/java/agent_plan_with_decision_router.json @@ -0,0 +1,88 @@ +{ + "agent_name" : null, + "actions" : { }, + "resource_providers" : { + "model_router" : { + "router" : { + "name" : "router", + "type" : "model_router", + "descriptor" : { + "language" : "java", + "target_module" : "", + "target_clazz" : "org.apache.flink.agents.api.chat.model.routing.ModelRouter", + "arguments" : { + "candidates" : [ "small", "big" ], + "candidate_descriptions" : { + "small" : "short factual questions", + "big" : "code and sql" + }, + "strategy_type" : "decision_model", + "default_model" : "small", + "fallback" : true, + "strategy_args" : { + "decision_model" : "dm", + "min_probability" : 0.6 + } + } + }, + "__resource_provider_type__" : "JavaResourceProvider" + } + }, + "chat_model" : { + "small" : { + "name" : "small", + "type" : "chat_model", + "descriptor" : { + "language" : "java", + "target_module" : "", + "target_clazz" : "some.Clazz", + "arguments" : { } + }, + "__resource_provider_type__" : "JavaResourceProvider" + }, + "big" : { + "name" : "big", + "type" : "chat_model", + "descriptor" : { + "language" : "java", + "target_module" : "", + "target_clazz" : "some.Clazz", + "arguments" : { } + }, + "__resource_provider_type__" : "JavaResourceProvider" + } + }, + "decision_model" : { + "dm" : { + "name" : "dm", + "type" : "decision_model", + "descriptor" : { + "language" : "java", + "target_module" : "", + "target_clazz" : "some.Setup", + "arguments" : { + "model" : "m", + "connection" : "conn" + } + }, + "__resource_provider_type__" : "JavaResourceProvider" + } + }, + "decision_model_connection" : { + "conn" : { + "name" : "conn", + "type" : "decision_model_connection", + "descriptor" : { + "language" : "java", + "target_module" : "", + "target_clazz" : "some.Clazz", + "arguments" : { } + }, + "__resource_provider_type__" : "JavaResourceProvider" + } + } + }, + "config" : { + "conf_data" : { } + } +} diff --git a/examples/src/main/java/org/apache/flink/agents/examples/openai/OpenAiDecisionRoutingExample.java b/examples/src/main/java/org/apache/flink/agents/examples/openai/OpenAiDecisionRoutingExample.java new file mode 100644 index 0000000000..36e6d25cfb --- /dev/null +++ b/examples/src/main/java/org/apache/flink/agents/examples/openai/OpenAiDecisionRoutingExample.java @@ -0,0 +1,382 @@ +/* + * 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.examples.openai; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.ArrayNode; +import com.fasterxml.jackson.databind.node.ObjectNode; +import org.apache.flink.agents.api.AgentsExecutionEnvironment; +import org.apache.flink.agents.api.EventType; +import org.apache.flink.agents.api.InputEvent; +import org.apache.flink.agents.api.OutputEvent; +import org.apache.flink.agents.api.agents.Agent; +import org.apache.flink.agents.api.agents.AgentExecutionOptions; +import org.apache.flink.agents.api.annotation.Action; +import org.apache.flink.agents.api.chat.messages.ChatMessage; +import org.apache.flink.agents.api.chat.model.routing.ModelRouter; +import org.apache.flink.agents.api.chat.model.routing.Strategies; +import org.apache.flink.agents.api.context.RunnerContext; +import org.apache.flink.agents.api.decision.BaseDecisionModelConnection; +import org.apache.flink.agents.api.decision.BaseDecisionModelSetup; +import org.apache.flink.agents.api.decision.DecisionRequest; +import org.apache.flink.agents.api.decision.DecisionResult; +import org.apache.flink.agents.api.event.ChatRequestEvent; +import org.apache.flink.agents.api.event.ChatResponseEvent; +import org.apache.flink.agents.api.resource.ResourceContext; +import org.apache.flink.agents.api.resource.ResourceDescriptor; +import org.apache.flink.agents.api.resource.ResourceName; +import org.apache.flink.agents.api.resource.ResourceType; +import org.apache.flink.streaming.api.datastream.DataStream; +import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment; + +import java.io.IOException; +import java.net.URI; +import java.net.http.HttpClient; +import java.net.http.HttpRequest; +import java.net.http.HttpResponse; +import java.time.Duration; +import java.util.Iterator; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; + +/** + * Routes each request with a decision model instead of an LLM judge: the OpenAI Decisions API + * answers one choice question over the candidates and reports a probability per candidate. The + * chosen candidate answers unless its probability is below the configured minimum, in which case + * the default model does. Requires {@code OPENAI_API_KEY} in the environment. + * + *

The decision model is a native Java provider: {@link OpenAiDecisionsConnection} implements + * {@link BaseDecisionModelConnection} over the Decisions endpoint, so the job needs no Python + * environment on the TaskManagers. The output lines carry the model that answered, the decision + * source and the routing score, read from the {@code model_routing} attribute of the response. + */ +public class OpenAiDecisionRoutingExample { + + private static ResourceDescriptor openAiModel(String model) { + return ResourceDescriptor.Builder.newBuilder( + ResourceName.ChatModel.OPENAI_COMPLETIONS_SETUP) + .addInitialArgument("connection", "openaiConnection") + .addInitialArgument("model", model) + .build(); + } + + /** Runs the example pipeline. */ + public static void main(String[] args) throws Exception { + String apiKey = System.getenv("OPENAI_API_KEY"); + if (apiKey == null || apiKey.isBlank()) { + throw new IllegalStateException("Set OPENAI_API_KEY in the environment to run."); + } + StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment(); + env.setParallelism(1); + AgentsExecutionEnvironment agentsEnv = + AgentsExecutionEnvironment.getExecutionEnvironment(env); + agentsEnv.getConfig().set(AgentExecutionOptions.NUM_ASYNC_THREADS, 2); + + // OpenAI connection shared by both candidate chat models. + agentsEnv.addResource( + "openaiConnection", + ResourceType.CHAT_MODEL_CONNECTION, + ResourceDescriptor.Builder.newBuilder( + ResourceName.ChatModel.OPENAI_COMPLETIONS_CONNECTION) + .addInitialArgument("api_key", apiKey) + .build()); + agentsEnv + .addResource("small", ResourceType.CHAT_MODEL, openAiModel("gpt-4o-mini")) + .addResource("big", ResourceType.CHAT_MODEL, openAiModel("gpt-4o")); + + // The decision model: a connection to the OpenAI Decisions API and a setup naming the + // decision model. Registered like any other decision model resource. + agentsEnv.addResource( + "openaiDecisions", + ResourceType.DECISION_MODEL_CONNECTION, + ResourceDescriptor.Builder.newBuilder(OpenAiDecisionsConnection.class.getName()) + .addInitialArgument("api_key", apiKey) + .build()); + agentsEnv.addResource( + "dm", + ResourceType.DECISION_MODEL, + ResourceDescriptor.Builder.newBuilder(BaseDecisionModelSetup.class.getName()) + .addInitialArgument("connection", "openaiDecisions") + .addInitialArgument("model", "gpt-6-luna") + .build()); + + // A router whose decision comes from the decision model; the descriptions are the choice + // criteria. A chosen candidate below 0.6 probability abstains to the default model. + agentsEnv.addResource( + "router", + ResourceType.MODEL_ROUTER, + ModelRouter.of("small", "big") + .describe("small", "fast and cheap: chit-chat, lookups, short factual asks") + .describe("big", "expensive: code, SQL, math, multi-step reasoning") + .strategy(Strategies.decision("dm").withMinProbability(0.6)) + .defaultModel("small") + .fallback(true) + .build()); + + // An easy request (-> small) and a coding request (-> big). + DataStream requestStream = + env.fromData( + "Hi, how are you today?", + "Write SQL to select all active users from the users table."); + DataStream resultStream = + agentsEnv + .fromDataStream(requestStream, value -> value) + .apply(new DecisionRoutingAgent()) + .toDataStream(); + resultStream.print(); + agentsEnv.execute("OpenAI Decision Routing Example"); + } + + /** Sends each input to the router and reports which model answered and why. */ + public static class DecisionRoutingAgent extends Agent { + + @Action(EventType.InputEvent) + public static void processInput(InputEvent event, RunnerContext ctx) { + ctx.sendEvent( + new ChatRequestEvent( + "router", List.of(ChatMessage.user((String) event.getInput())))); + } + + @Action(EventType.ChatResponseEvent) + public static void processChatResponse(ChatResponseEvent event, RunnerContext ctx) { + if (event.isFailed()) { + ctx.sendEvent(new OutputEvent("FAILED: " + event.getError())); + return; + } + Map routing = (Map) event.getAttr("model_routing"); + String answer = event.getResponse().getText().replace('\n', ' '); + ctx.sendEvent( + new OutputEvent( + String.format( + "final_model=%s decision_source=%s score=%s :: %s", + routing.get("final_model"), + routing.get("decision_source"), + routing.get("score"), + answer.length() > 80 + ? answer.substring(0, 80) + "..." + : answer))); + } + } + + /** + * A native Java decision model connection over the OpenAI Decisions API ({@code POST + * /v1/decisions}). Supports choice and predicate questions, which is what routing and most + * gating use; score questions are rejected. Descriptor arguments: {@code api_key} (required) + * and {@code base_url} (default {@code https://api.openai.com/v1}). + */ + public static class OpenAiDecisionsConnection extends BaseDecisionModelConnection { + + private static final ObjectMapper MAPPER = new ObjectMapper(); + private static final java.util.Set RESERVED_FIELDS = + java.util.Set.of("model", "input", "questions"); + + private final String apiKey; + private final String baseUrl; + private transient volatile HttpClient client; + + public OpenAiDecisionsConnection( + ResourceDescriptor descriptor, ResourceContext resourceContext) { + super(descriptor, resourceContext); + Object key = descriptor.getArgument("api_key"); + if (!(key instanceof String) || ((String) key).isBlank()) { + throw new IllegalArgumentException( + "OpenAiDecisionsConnection requires a non-blank 'api_key'."); + } + this.apiKey = (String) key; + this.baseUrl = descriptor.getArgument("base_url", "https://api.openai.com/v1"); + } + + @Override + public DecisionResult decide( + DecisionRequest request, String model, Map requestParams) { + ObjectNode body = MAPPER.createObjectNode(); + for (Map.Entry param : requestParams.entrySet()) { + if (RESERVED_FIELDS.contains(param.getKey())) { + throw new IllegalArgumentException( + "Request parameter '" + + param.getKey() + + "' cannot override the model, evidence or questions."); + } + body.set(param.getKey(), MAPPER.valueToTree(param.getValue())); + } + body.put("model", model); + body.set("input", MAPPER.valueToTree(request.getInput())); + ArrayNode questions = body.putArray("questions"); + for (Map.Entry entry : + request.getQuestions().entrySet()) { + questions.add(toQuestion(entry.getKey(), entry.getValue())); + } + JsonNode response = post(body); + Map answers = new LinkedHashMap<>(); + for (JsonNode answer : response.path("answers")) { + String name = answer.path("name").asText(); + DecisionRequest.Question question = request.getQuestions().get(name); + if (question == null) { + throw new IllegalStateException( + "OpenAI Decisions API answered an unknown question '" + name + "'."); + } + answers.put(name, toAnswer(answer, question)); + } + return new DecisionResult( + answers, + response.path("model").isTextual() ? response.get("model").asText() : null, + toUsage(response.path("usage")), + Map.of()); + } + + private static ObjectNode toQuestion(String name, DecisionRequest.Question question) { + ObjectNode node = MAPPER.createObjectNode(); + node.put("name", name); + if (question instanceof DecisionRequest.ChoiceQuestion) { + node.put("type", "choice"); + node.set("instructions", MAPPER.valueToTree(question.getInstructions())); + ArrayNode choices = node.putArray("choices"); + for (Map.Entry option : + ((DecisionRequest.ChoiceQuestion) question).getCriteria().entrySet()) { + ObjectNode choice = choices.addObject(); + choice.put("value", option.getKey()); + if (option.getValue() != null) { + choice.set("description", MAPPER.valueToTree(option.getValue())); + } + } + return node; + } + if (question instanceof DecisionRequest.PredicateQuestion) { + node.put("type", "predicate"); + node.set("instructions", MAPPER.valueToTree(question.getInstructions())); + return node; + } + throw new IllegalArgumentException( + "OpenAiDecisionsConnection supports choice and predicate questions only."); + } + + private static DecisionResult.Answer toAnswer( + JsonNode answer, DecisionRequest.Question question) { + String type = answer.path("type").asText(); + if ("refusal".equals(type)) { + return new DecisionResult.RefusalAnswer( + answer.path("reason").isTextual() ? answer.get("reason").asText() : null); + } + if (question instanceof DecisionRequest.ChoiceQuestion) { + if (!answer.path("choice").isTextual()) { + throw new IllegalStateException( + "OpenAI Decisions API returned a choice answer without a 'choice'."); + } + Map probabilities = new LinkedHashMap<>(); + for (JsonNode entry : answer.path("probabilities")) { + if (!entry.path("value").isTextual() || !entry.path("probability").isNumber()) { + throw new IllegalStateException( + "OpenAI Decisions API returned a malformed probability entry: " + + entry); + } + if (probabilities.put( + entry.path("value").asText(), + entry.path("probability").asDouble()) + != null) { + throw new IllegalStateException( + "OpenAI Decisions API returned duplicate choice '" + + entry.path("value").asText() + + "'."); + } + } + if (probabilities.isEmpty()) { + throw new IllegalStateException( + "OpenAI Decisions API returned a choice answer without" + + " probabilities."); + } + return new DecisionResult.ChoiceAnswer( + answer.path("choice").asText(), + probabilities, + answer.path("confidence").isNumber() + ? answer.get("confidence").asDouble() + : null); + } + if (!"predicate".equals(type) || !answer.path("probability").isNumber()) { + throw new IllegalStateException( + "OpenAI Decisions API returned an unexpected answer of type '" + + type + + "' for a predicate question."); + } + return new DecisionResult.PredicateAnswer(answer.path("probability").asDouble()); + } + + private static DecisionResult.TokenUsage toUsage(JsonNode usage) { + if (usage.isMissingNode() || usage.isNull()) { + return null; + } + return new DecisionResult.TokenUsage( + longOrNull(usage, "input_tokens", "prompt_tokens"), + longOrNull(usage, "output_tokens", "completion_tokens"), + longOrNull(usage, "total_tokens")); + } + + private static Long longOrNull(JsonNode node, String... names) { + for (String name : names) { + if (node.path(name).isNumber()) { + return node.get(name).asLong(); + } + } + return null; + } + + private JsonNode post(ObjectNode body) { + if (client == null) { + client = HttpClient.newBuilder().connectTimeout(Duration.ofSeconds(30)).build(); + } + try { + HttpRequest request = + HttpRequest.newBuilder(URI.create(baseUrl + "/decisions")) + .timeout(Duration.ofSeconds(60)) + .header("Authorization", "Bearer " + apiKey) + .header("Content-Type", "application/json") + .POST( + HttpRequest.BodyPublishers.ofString( + MAPPER.writeValueAsString(body))) + .build(); + HttpResponse response = + client.send(request, HttpResponse.BodyHandlers.ofString()); + if (response.statusCode() / 100 != 2) { + throw new IllegalStateException( + "OpenAI Decisions API returned HTTP " + + response.statusCode() + + ": " + + errorMessage(response.body())); + } + return MAPPER.readTree(response.body()); + } catch (IOException e) { + throw new IllegalStateException("OpenAI Decisions API call failed", e); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IllegalStateException("OpenAI Decisions API call interrupted", e); + } + } + + /** The provider's error message without echoing the whole body into logs. */ + private static String errorMessage(String body) { + try { + JsonNode error = MAPPER.readTree(body).path("error"); + Iterator fields = error.fieldNames(); + return fields.hasNext() ? error.path("message").asText(body) : body; + } catch (IOException e) { + return body; + } + } + } +} diff --git a/plan/src/main/java/org/apache/flink/agents/plan/AgentPlan.java b/plan/src/main/java/org/apache/flink/agents/plan/AgentPlan.java index 4dc831571a..b9ebf09097 100644 --- a/plan/src/main/java/org/apache/flink/agents/plan/AgentPlan.java +++ b/plan/src/main/java/org/apache/flink/agents/plan/AgentPlan.java @@ -829,6 +829,11 @@ private void checkNoRouterModelNameClash(ResourceProvider provider) { * 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. + * + *

{@code DECISION_MODEL}: the decision-model setup must be registered (as for the judge, a + * typo'd name would otherwise fail every request), and every candidate must carry a description + * — the descriptions are the choice question's criteria, so a candidate without one is + * invisible to the decision model. */ private void validateRoutingStrategies() { if (resourceProviders == null) { @@ -840,6 +845,11 @@ private void validateRoutingStrategies() { } Map chatModels = resourceProviders.getOrDefault(ResourceType.CHAT_MODEL, Collections.emptyMap()); + Map decisionModels = + resourceProviders.getOrDefault(ResourceType.DECISION_MODEL, Collections.emptyMap()); + Map decisionConnections = + resourceProviders.getOrDefault( + ResourceType.DECISION_MODEL_CONNECTION, Collections.emptyMap()); for (ResourceProvider provider : routers.values()) { // Java- and Python-provided declarations carry the same keys; a router without an // introspectable descriptor is left to its constructor. @@ -883,13 +893,32 @@ private void validateRoutingStrategies() { // shape guard applies to all of them — not only where the rule keys are checked. Object candidates = descriptor.getArgument(ModelRouter.CANDIDATES_KEY); validateCandidatesShape(provider.getName(), candidates); + Object candidateDescriptions = + descriptor.getArgument(ModelRouter.CANDIDATE_DESCRIPTIONS_KEY); + if (candidateDescriptions != null && !(candidateDescriptions instanceof Map)) { + throw new IllegalArgumentException( + String.format( + "Model router '%s' declares '%s' as %s; expected a map of" + + " candidate name to description.", + provider.getName(), + ModelRouter.CANDIDATE_DESCRIPTIONS_KEY, + candidateDescriptions.getClass().getSimpleName())); + } // The constructor's own candidate rules (non-empty, no duplicates, default model is a - // candidate), applied here so a descriptor that skipped the builder fails at plan - // construction rather than per routed request on the TaskManager. + // candidate, String descriptions), applied here so a descriptor that skipped the + // builder fails at plan construction rather than per routed request on the TaskManager. RoutingCandidateValidator.validate( String.format("Model router '%s'", provider.getName()), (List) candidates, descriptor.getArgument(ModelRouter.DEFAULT_MODEL_KEY)); + if (candidateDescriptions != null) { + @SuppressWarnings("unchecked") + List candidateNames = (List) candidates; + RoutingCandidateValidator.requireStringDescriptions( + String.format("Model router '%s'", provider.getName()), + candidateNames, + (Map) candidateDescriptions); + } switch (strategy.getType()) { case LLM_JUDGE: validateJudge(provider.getName(), strategy, chatModels); @@ -900,12 +929,97 @@ private void validateRoutingStrategies() { case RULE_BASED: validateRuleKeys(provider.getName(), strategy, candidates); break; + case DECISION_MODEL: + validateDecisionModel( + provider.getName(), + strategy, + decisionModels, + decisionConnections, + candidates, + descriptor); + break; default: break; } } } + private static void validateDecisionModel( + String routerName, + RoutingStrategy strategy, + Map decisionModels, + Map decisionConnections, + Object candidates, + ResourceDescriptor descriptor) { + String decisionModel = + (String) strategy.getArguments().get(RoutingStrategy.ARG_DECISION_MODEL); + ResourceProvider setupProvider = decisionModels.get(decisionModel); + if (setupProvider == null) { + throw new IllegalArgumentException( + String.format( + "Model router '%s' uses Strategies.decision with decision model '%s'," + + " but no DECISION_MODEL resource with that name is" + + " registered.", + routerName, decisionModel)); + } + // The setup opens its named connection on first use; a typo there would fail every + // routed request on the TaskManager, so an explicitly declared connection name is checked + // here. A setup that declares none may resolve it from provider defaults (Python-backed + // setups do), so absence is not an error. + ResourceDescriptor setupDescriptor = descriptorOf(setupProvider); + Object connection = + setupDescriptor == null || setupDescriptor.getInitialArguments() == null + ? null + : setupDescriptor.getArgument("connection"); + if (connection != null && !(connection instanceof String)) { + throw new IllegalArgumentException( + String.format( + "Decision model '%s' (router '%s') declares '%s' as its connection;" + + " expected a connection name.", + decisionModel, routerName, connection)); + } + if (connection instanceof String) { + if (((String) connection).isBlank()) { + throw new IllegalArgumentException( + String.format( + "Decision model '%s' (router '%s') declares a blank 'connection'.", + decisionModel, routerName)); + } + if (!decisionConnections.containsKey(connection)) { + throw new IllegalArgumentException( + String.format( + "Decision model '%s' (router '%s') names connection '%s', but no" + + " DECISION_MODEL_CONNECTION resource with that name is" + + " registered.", + decisionModel, routerName, connection)); + } + } + Object model = + setupDescriptor == null || setupDescriptor.getInitialArguments() == null + ? null + : setupDescriptor.getArgument("model"); + if (model != null && (!(model instanceof String) || ((String) model).isBlank())) { + throw new IllegalArgumentException( + String.format( + "Decision model '%s' (router '%s') declares a blank or non-String" + + " 'model'.", + decisionModel, routerName)); + } + // Same rule and message as the router constructor, applied here so a descriptor that + // skipped the builder fails at plan construction rather than at resource construction on + // the TaskManager. + Object descriptions = descriptor.getArgument(ModelRouter.CANDIDATE_DESCRIPTIONS_KEY); + List names = new ArrayList<>(); + for (Object candidate : (List) candidates) { + names.add(String.valueOf(candidate)); + } + RoutingCandidateValidator.requireDescriptionsForDecisionStrategy( + String.format("Model router '%s'", routerName), + strategy, + names, + descriptions == null ? Collections.emptyMap() : (Map) descriptions); + } + /** * 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 diff --git a/plan/src/main/java/org/apache/flink/agents/plan/actions/DecisionModelAction.java b/plan/src/main/java/org/apache/flink/agents/plan/actions/DecisionModelAction.java index e7f79ae86b..a9be4b80f6 100644 --- a/plan/src/main/java/org/apache/flink/agents/plan/actions/DecisionModelAction.java +++ b/plan/src/main/java/org/apache/flink/agents/plan/actions/DecisionModelAction.java @@ -17,37 +17,18 @@ */ package org.apache.flink.agents.plan.actions; -import com.fasterxml.jackson.annotation.JsonCreator; -import com.fasterxml.jackson.annotation.JsonProperty; import org.apache.flink.agents.api.Event; -import org.apache.flink.agents.api.agents.AgentExecutionOptions; -import org.apache.flink.agents.api.context.DurableCallable; import org.apache.flink.agents.api.context.RunnerContext; import org.apache.flink.agents.api.decision.BaseDecisionModelSetup; -import org.apache.flink.agents.api.decision.DecisionResult; import org.apache.flink.agents.api.event.DecisionRequestEvent; import org.apache.flink.agents.api.event.DecisionResponseEvent; -import org.apache.flink.agents.api.metrics.FlinkAgentsMetricGroup; import org.apache.flink.agents.api.resource.ResourceType; -import org.apache.flink.agents.api.trace.ExecutionReporter; -import org.apache.flink.agents.api.trace.ExecutionReporters; import org.apache.flink.agents.plan.JavaFunction; -import org.apache.flink.agents.plan.utils.CancellationUtils; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; -import javax.annotation.Nullable; - -import java.lang.reflect.InvocationTargetException; import java.util.List; -import java.util.Map; -import java.util.Objects; -import java.util.concurrent.CompletionException; -import java.util.concurrent.ExecutionException; /** Executes decision requests durably and emits one correlated success or failure response. */ public final class DecisionModelAction { - private static final Logger LOG = LoggerFactory.getLogger(DecisionModelAction.class); private DecisionModelAction() {} @@ -64,176 +45,36 @@ public static Action getDecisionModelAction() throws Exception { /** * Handles a decision request with the configured retry budget. Resource lookup failures, * cancellation and durable-state failures propagate instead of becoming failed decision - * responses. + * responses. The call itself runs through {@link DecisionModelInvoker} under the flat durable + * id {@code "decide"}. */ public static void processDecisionRequest(Event event, RunnerContext ctx) throws Exception { DecisionRequestEvent request = (DecisionRequestEvent) event; BaseDecisionModelSetup model = (BaseDecisionModelSetup) ctx.getResource(request.getModel(), ResourceType.DECISION_MODEL); - - boolean async = - ctx.getConfig().get(AgentExecutionOptions.DECISION_ASYNC) - && model.supportsAsyncExecution(); - int maxRetries = Math.max(0, ctx.getConfig().get(AgentExecutionOptions.MAX_RETRIES)); - int waitInterval = - Math.max(0, ctx.getConfig().get(AgentExecutionOptions.RETRY_WAIT_INTERVAL)); - FlinkAgentsMetricGroup actionMetrics = ctx.getActionMetricGroup(); - FlinkAgentsMetricGroup resourceMetrics = - actionMetrics == null - ? null - : actionMetrics.getSubGroup("model_resource", request.getModel()); - - DurableCallable call = - new DurableCallable<>() { - @Override - public String getId() { - return "decide"; - } - - @Override - public Class getResultClass() { - return InvocationOutcome.class; - } - - @Override - public InvocationOutcome call() throws Exception { - try { - return new InvocationOutcome( - Objects.requireNonNull( - model.decide( - request.getRequest(), - request.getRequestParams()), - "DecisionModel returned a null response."), - null); - } catch (Exception error) { - propagateCancellation(error); - LOG.debug("Decision provider failed", error); - return new InvocationOutcome(null, errorText(error)); - } - } - }; - - int retryCount = 0; - int totalWaitSec = 0; - DecisionResponseEvent terminal; - try { - while (true) { - String executionId = - ExecutionReporters.started( - ctx, - ExecutionReporter.EntityTypes.LLM, - request.getModel(), - Map.of()); - // Only the provider callback converts failures. Persistence/recovery errors escape. - InvocationOutcome outcome = - async ? ctx.durableExecuteAsync(call).await() : ctx.durableExecute(call); - if (outcome.error == null) { - ExecutionReporters.succeeded( - ctx, - executionId, - ExecutionReporter.EntityTypes.LLM, - request.getModel(), - Map.of()); - recordTokenMetrics(resourceMetrics, outcome.response); - terminal = - DecisionResponseEvent.success( - request.getId(), outcome.response, retryCount, totalWaitSec); - break; - } - - ExecutionReporters.failed( + DecisionModelInvoker.DecisionAttemptResult result = + DecisionModelInvoker.decideWithRetries( ctx, - executionId, - ExecutionReporter.EntityTypes.LLM, request.getModel(), - Map.of(), - new Exception(outcome.error), - ExecutionReporter.ProblemCategories.MODEL_CALL_FAILED); - if (retryCount == maxRetries) { - terminal = - DecisionResponseEvent.failed( - request.getId(), outcome.error, retryCount, totalWaitSec); - break; - } - int waitSec = waitInterval * (1 << retryCount); - retryCount++; - LOG.warn( - "Decision request {} failed with error: {}, retrying {} / {}, waiting {} s.", - request.getId(), - outcome.error, - retryCount, - maxRetries, - waitSec); - if (waitSec > 0) { - try { - Thread.sleep(waitSec * 1000L); - } catch (InterruptedException error) { - Thread.currentThread().interrupt(); - throw error; - } - totalWaitSec += waitSec; - } - } - } finally { - if (resourceMetrics != null && retryCount > 0) { - resourceMetrics.getCounter("retryCount").inc(retryCount); - resourceMetrics.getCounter("retryWaitSec").inc(totalWaitSec); - } - } + model, + request.getRequest(), + request.getRequestParams(), + "decide", + String.valueOf(request.getId()), + true); + DecisionResponseEvent terminal = + result.isSuccess() + ? DecisionResponseEvent.success( + request.getId(), + result.response, + result.retryCount, + result.totalWaitSec) + : DecisionResponseEvent.failed( + request.getId(), + result.error, + result.retryCount, + result.totalWaitSec); ctx.sendEvent(terminal); } - - private static void propagateCancellation(Exception error) throws Exception { - if (Thread.currentThread().isInterrupted() || CancellationUtils.isCancellation(error)) { - if (CancellationUtils.isInterruption(error)) { - Thread.currentThread().interrupt(); - } - throw error; - } - } - - private static String errorText(Exception error) { - while ((error instanceof InvocationTargetException - || error instanceof ExecutionException - || error instanceof CompletionException) - && error.getCause() instanceof Exception) { - error = (Exception) error.getCause(); - } - return error.getClass().getName() + ": " + error.getMessage(); - } - - private static void recordTokenMetrics( - @Nullable FlinkAgentsMetricGroup resourceMetrics, DecisionResult result) { - if (resourceMetrics == null || result.getUsage() == null) { - return; - } - FlinkAgentsMetricGroup metrics = resourceMetrics; - if (result.getModel() != null && !result.getModel().isEmpty()) { - metrics = metrics.getSubGroup("model", result.getModel()); - } - recordTokenCount(metrics, "promptTokens", result.getUsage().getInputTokens()); - recordTokenCount(metrics, "completionTokens", result.getUsage().getOutputTokens()); - recordTokenCount(metrics, "totalTokens", result.getUsage().getTotalTokens()); - } - - private static void recordTokenCount( - FlinkAgentsMetricGroup metrics, String name, @Nullable Long count) { - if (count != null) { - metrics.getCounter(name).inc(count); - } - } - - private static final class InvocationOutcome { - @JsonProperty private final DecisionResult response; - @JsonProperty private final String error; - - @JsonCreator - private InvocationOutcome( - @JsonProperty("response") DecisionResult response, - @JsonProperty("error") String error) { - this.response = response; - this.error = error; - } - } } diff --git a/plan/src/main/java/org/apache/flink/agents/plan/actions/DecisionModelInvoker.java b/plan/src/main/java/org/apache/flink/agents/plan/actions/DecisionModelInvoker.java new file mode 100644 index 0000000000..417208909d --- /dev/null +++ b/plan/src/main/java/org/apache/flink/agents/plan/actions/DecisionModelInvoker.java @@ -0,0 +1,239 @@ +/* + * 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 com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; +import org.apache.flink.agents.api.agents.AgentExecutionOptions; +import org.apache.flink.agents.api.context.DurableCallable; +import org.apache.flink.agents.api.context.RunnerContext; +import org.apache.flink.agents.api.decision.BaseDecisionModelSetup; +import org.apache.flink.agents.api.decision.DecisionRequest; +import org.apache.flink.agents.api.decision.DecisionResult; +import org.apache.flink.agents.api.metrics.FlinkAgentsMetricGroup; +import org.apache.flink.agents.api.trace.ExecutionReporter; +import org.apache.flink.agents.api.trace.ExecutionReporters; +import org.apache.flink.agents.plan.utils.CancellationUtils; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import javax.annotation.Nullable; + +import java.util.Map; +import java.util.Objects; + +/** + * The one implementation of a durable, retried, reported and metered decision-model call, shared by + * the built-in {@link DecisionModelAction} and the decision-model routing strategy. Callers supply + * the durable id and decide what to do with the terminal outcome; this class owns the retry loop, + * the execution reports, the token metrics and the durable record shape, so the two call sites + * cannot drift. + */ +public final class DecisionModelInvoker { + + private static final Logger LOG = LoggerFactory.getLogger(DecisionModelInvoker.class); + + private DecisionModelInvoker() {} + + /** The terminal outcome of the retry loop: a result, or the last error after the budget. */ + public static final class DecisionAttemptResult { + @Nullable public final DecisionResult response; + @Nullable public final String error; + public final int retryCount; + public final int totalWaitSec; + + DecisionAttemptResult( + @Nullable DecisionResult response, + @Nullable String error, + int retryCount, + int totalWaitSec) { + this.response = response; + this.error = error; + this.retryCount = retryCount; + this.totalWaitSec = totalWaitSec; + } + + public boolean isSuccess() { + return error == null; + } + } + + /** The durable record of one decision attempt: the provider's result or its error text. */ + public static final class DecisionOutcome { + @JsonProperty private final DecisionResult response; + @JsonProperty private final String error; + + @JsonCreator + public DecisionOutcome( + @JsonProperty("response") @Nullable DecisionResult response, + @JsonProperty("error") @Nullable String error) { + this.response = response; + this.error = error; + } + } + + /** + * Runs {@code setup.decide(request, requestParams)} under {@code durableCallId} with the + * configured retry budget. Only the provider callback converts failures into a retryable + * outcome; persistence and recovery errors escape, and cancellation propagates immediately. + * Token metrics for a successful result land on {@code model_resource/}; when {@code + * recordRetryMetrics} is set, the retry counters do too (a caller that accumulates retry stats + * through {@link ChatModelAction#recordAttemptRetryStats} passes {@code false}). + * + * @param label what is being decided, for the retry log line + */ + public static DecisionAttemptResult decideWithRetries( + RunnerContext ctx, + String modelName, + BaseDecisionModelSetup setup, + DecisionRequest request, + Map requestParams, + String durableCallId, + String label, + boolean recordRetryMetrics) + throws Exception { + boolean async = + ctx.getConfig().get(AgentExecutionOptions.DECISION_ASYNC) + && setup.supportsAsyncExecution(); + int maxRetries = Math.max(0, ctx.getConfig().get(AgentExecutionOptions.MAX_RETRIES)); + int waitInterval = + Math.max(0, ctx.getConfig().get(AgentExecutionOptions.RETRY_WAIT_INTERVAL)); + FlinkAgentsMetricGroup actionMetrics = ctx.getActionMetricGroup(); + FlinkAgentsMetricGroup resourceMetrics = + actionMetrics == null + ? null + : actionMetrics.getSubGroup("model_resource", modelName); + + DurableCallable call = + new DurableCallable<>() { + @Override + public String getId() { + return durableCallId; + } + + @Override + public Class getResultClass() { + return DecisionOutcome.class; + } + + @Override + public DecisionOutcome call() throws Exception { + try { + return new DecisionOutcome( + Objects.requireNonNull( + setup.decide(request, requestParams), + "DecisionModel returned a null response."), + null); + } catch (Exception error) { + CancellationUtils.rethrowIfCancellation(error); + LOG.debug("Decision provider failed", error); + return new DecisionOutcome(null, ChatModelInvoker.errorText(error)); + } + } + }; + + int retryCount = 0; + int totalWaitSec = 0; + try { + while (true) { + String executionId = + ExecutionReporters.started( + ctx, ExecutionReporter.EntityTypes.LLM, modelName, Map.of()); + // Only the provider callback converts failures. Persistence/recovery errors escape. + DecisionOutcome outcome; + try { + outcome = + async + ? ctx.durableExecuteAsync(call).await() + : ctx.durableExecute(call); + } catch (InterruptedException e) { + // A cancellation signal, not a model failure: restore the interrupt status + // and propagate immediately. + Thread.currentThread().interrupt(); + throw e; + } + if (outcome.error == null) { + ExecutionReporters.succeeded( + ctx, + executionId, + ExecutionReporter.EntityTypes.LLM, + modelName, + Map.of()); + recordTokenMetrics(resourceMetrics, outcome.response); + return new DecisionAttemptResult( + outcome.response, null, retryCount, totalWaitSec); + } + ExecutionReporters.failed( + ctx, + executionId, + ExecutionReporter.EntityTypes.LLM, + modelName, + Map.of(), + new Exception(outcome.error), + ExecutionReporter.ProblemCategories.MODEL_CALL_FAILED); + if (retryCount == maxRetries) { + return new DecisionAttemptResult(null, outcome.error, retryCount, totalWaitSec); + } + int waitSec = waitInterval * (1 << retryCount); + retryCount++; + LOG.warn( + "Decision request {} failed with error: {}, retrying {} / {}, waiting {} s.", + label, + outcome.error, + retryCount, + maxRetries, + waitSec); + if (waitSec > 0) { + try { + Thread.sleep(waitSec * 1000L); + } catch (InterruptedException error) { + Thread.currentThread().interrupt(); + throw error; + } + totalWaitSec += waitSec; + } + } + } finally { + if (recordRetryMetrics && resourceMetrics != null && retryCount > 0) { + resourceMetrics.getCounter("retryCount").inc(retryCount); + resourceMetrics.getCounter("retryWaitSec").inc(totalWaitSec); + } + } + } + + private static void recordTokenMetrics( + @Nullable FlinkAgentsMetricGroup resourceMetrics, DecisionResult result) { + if (resourceMetrics == null || result.getUsage() == null) { + return; + } + FlinkAgentsMetricGroup metrics = resourceMetrics; + if (result.getModel() != null && !result.getModel().isEmpty()) { + metrics = metrics.getSubGroup("model", result.getModel()); + } + recordTokenCount(metrics, "promptTokens", result.getUsage().getInputTokens()); + recordTokenCount(metrics, "completionTokens", result.getUsage().getOutputTokens()); + recordTokenCount(metrics, "totalTokens", result.getUsage().getTotalTokens()); + } + + private static void recordTokenCount( + FlinkAgentsMetricGroup metrics, String name, @Nullable Long count) { + if (count != null) { + metrics.getCounter(name).inc(count); + } + } +} diff --git a/plan/src/main/java/org/apache/flink/agents/plan/routing/DecisionModelRoutingExecutor.java b/plan/src/main/java/org/apache/flink/agents/plan/routing/DecisionModelRoutingExecutor.java new file mode 100644 index 0000000000..4884346348 --- /dev/null +++ b/plan/src/main/java/org/apache/flink/agents/plan/routing/DecisionModelRoutingExecutor.java @@ -0,0 +1,435 @@ +/* + * 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.routing; + +import org.apache.flink.agents.api.chat.messages.ChatMessage; +import org.apache.flink.agents.api.chat.messages.DataContentBlock; +import org.apache.flink.agents.api.chat.messages.MessageRole; +import org.apache.flink.agents.api.chat.messages.TextBlock; +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.decision.BaseDecisionModelSetup; +import org.apache.flink.agents.api.decision.DecisionRequest; +import org.apache.flink.agents.api.decision.DecisionResult; +import org.apache.flink.agents.api.event.ModelRoutingEvent; +import org.apache.flink.agents.api.resource.ResourceType; +import org.apache.flink.agents.plan.actions.ChatModelAction; +import org.apache.flink.agents.plan.actions.ChatModelInvoker; +import org.apache.flink.agents.plan.actions.DecisionModelInvoker; +import org.apache.flink.agents.plan.utils.CancellationUtils; + +import javax.annotation.Nullable; + +import java.util.HashMap; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; + +/** + * Executes the framework-managed decision-model strategy ({@code Strategies.decision(...)}): the + * engine — not the strategy — asks a registered {@code DECISION_MODEL} resource one choice question + * whose options are the router's candidates and whose criteria are their {@code describe(...)} + * descriptions, through a flat durable call with the id {@code "decide:"} (issued + * before the resolver's decision record — see the sequencing contract on {@link + * RoutingExecutor}), then derives the decision from the answer as a pure function. + * + *

The answer is constrained by construction: the provider must pick one of the requested + * choices, so the decision model cannot steer routing outside the declared candidates. The choice + * becomes the selected model and its reported probability the routing score. A refusal, a + * non-choice answer, a result without an answer for the question, or a probability below the + * declared {@code min_probability} abstains to the router's default model with the cause recorded. + * The decision model sees the conversation as text only (tool calls and results included, media and + * reasoning excluded), rendered like the LLM judge's input and subject to the same optional context + * cap. + * + *

Failure policy: a decision call that exhausts its retries fails the request, like a throwing + * rule/custom strategy or a failed judge call. Cancellation propagates and is never persisted as a + * routing outcome. The call shares {@link DecisionModelInvoker} with the built-in decision action + * but runs under its own durable id, so it does not collide with the action's flat {@code "decide"} + * record; retries are accumulated on the routed request like the judge's. + */ +final class DecisionModelRoutingExecutor implements RoutingExecutor { + + /** The name of the single choice question in the decision request and its answer. */ + static final String QUESTION_NAME = "model"; + + /** The built-in choice instructions, used unless the declaration carries its own. */ + static final String DEFAULT_INSTRUCTIONS = + "Choose which ONE candidate model should answer the conversation below. Each option" + + " describes the requests its model should handle; pick the option whose" + + " description fits the conversation best."; + + /** Decision-metadata key: the decision-model setup that produced the choice. */ + static final String DECISION_MODEL_KEY = "decision_model"; + + /** Decision-metadata key: the provider's probability per candidate, as reported. */ + static final String PROBABILITIES_KEY = "decision_probabilities"; + + /** Decision-metadata key: the provider's separate confidence measure, when it reports one. */ + static final String PROVIDER_CONFIDENCE_KEY = "provider_confidence"; + + /** Decision-metadata key: the chosen candidate when the choice abstained to the default. */ + static final String CHOICE_KEY = "decision_choice"; + + /** Decision-metadata flag set when the context cap dropped part of the conversation. */ + static final String CONTEXT_TRUNCATED_KEY = "decision_context_truncated"; + + /** Decision-metadata key: media blocks left out of the text-only decision input. */ + static final String MEDIA_DROPPED_KEY = "decision_media_dropped"; + + /** Decision-metadata keys: the decision call's token usage, when the provider reports both. */ + static final String INPUT_TOKENS_KEY = "decision_input_tokens"; + + static final String OUTPUT_TOKENS_KEY = "decision_output_tokens"; + + @Override + public boolean usesDurableExecutionInternally() { + return true; + } + + @Override + public String decisionSource() { + return ModelRoutingEvent.SOURCE_DECISION_MODEL; + } + + @Override + public RoutingDecision route( + RoutingStrategy strategy, RoutingContext context, RunnerContext ctx) throws Exception { + String decisionModel = decisionModel(strategy); + Map metadata = new LinkedHashMap<>(); + metadata.put(DECISION_MODEL_KEY, decisionModel); + DecisionRequest request; + try { + request = buildRequest(strategy, context, ctx, metadata); + } catch (Exception e) { + throw routingFailure(context, "could not build the decision request", e); + } + if (String.valueOf(request.getInput()).isEmpty()) { + // Nothing to decide on (a media-only or empty request with no prompt arguments): + // abstain to the default model instead of paying for a call the provider rejects. + // Resolved before the setup, so an unreachable provider cannot fail such a request. + return abstain("request carries no text for the decision model", metadata); + } + BaseDecisionModelSetup setup; + try { + setup = + (BaseDecisionModelSetup) + ctx.getResource(decisionModel, ResourceType.DECISION_MODEL); + } catch (Exception e) { + CancellationUtils.rethrowIfCancellation(e); + // An unresolvable decision model is an ordinary routing failure (failed response), + // not a runner failure: the request cannot be routed, the job keeps running. + throw new ModelRoutingResolver.RoutingFailure( + "Decision model '" + + decisionModel + + "' could not be resolved: " + + ChatModelInvoker.errorText(e)); + } + DecisionModelInvoker.DecisionAttemptResult result = + DecisionModelInvoker.decideWithRetries( + ctx, + decisionModel, + setup, + request, + Map.of(), + "decide:" + context.getRouter(), + "for router '" + context.getRouter() + "'", + false); + // Retries count against the routed request, like the judge's: response-level totals and + // the model_resource counters both come from this one call. + ChatModelAction.recordAttemptRetryStats( + ctx, context.getRequestId(), decisionModel, result.retryCount, result.totalWaitSec); + if (!result.isSuccess()) { + throw new ModelRoutingResolver.RoutingFailure( + "Decision model '" + decisionModel + "' failed: " + result.error); + } + try { + return toDecision(strategy, context, result.response, metadata); + } catch (Exception e) { + throw routingFailure(context, "could not read the decision answer", e); + } + } + + /** + * An ordinary routing failure (failed chat response, no candidate fallback) for a request that + * could not be prepared or whose answer could not be read. Runtime failures of the durable + * substrate never pass through here: they surface from the invoker and propagate to the runner, + * as the {@link RoutingExecutor} contract requires. Cancellation is rethrown as is. + */ + private static ModelRoutingResolver.RoutingFailure routingFailure( + RoutingContext context, String what, Exception cause) throws Exception { + CancellationUtils.rethrowIfCancellation(cause); + return new ModelRoutingResolver.RoutingFailure( + "Decision routing for router '" + + context.getRouter() + + "' " + + what + + ": " + + ChatModelInvoker.errorText(cause)); + } + + /** + * Builds the decision request: the conversation rendered as text (the judge's view, so the + * decision model routes on what the selected model will actually receive, including a bound + * prompt rendered by the anchor candidate) plus one choice question over the candidates. + * Truncation and dropped media are recorded in {@code metadata}. + */ + static DecisionRequest buildRequest( + RoutingStrategy strategy, + RoutingContext context, + RunnerContext ctx, + Map metadata) + throws Exception { + List effective = LlmJudgeRoutingExecutor.effectiveJudgeMessages(context, ctx); + Set pinned = + LlmJudgeRoutingExecutor.pinnedRenderedIndices(context.getMessages(), effective); + boolean[] truncated = new boolean[1]; + String input = renderInput(strategy, context, effective, pinned, truncated, metadata); + if (truncated[0]) { + metadata.put(CONTEXT_TRUNCATED_KEY, true); + } + return new DecisionRequest( + input, + Map.of( + QUESTION_NAME, + DecisionRequest.Question.choice( + criteria(context), instructions(strategy)))); + } + + /** + * The choice criteria: candidate name to description. Descriptions are required by the router + * declaration ({@code RoutingCandidateValidator#requireDescriptionsForDecisionStrategy}), so + * every option carries one. + */ + static Map criteria(RoutingContext context) { + Map criteria = new LinkedHashMap<>(); + for (RoutingCandidate candidate : context.getCandidates()) { + criteria.put(candidate.getName(), candidate.getDescription()); + } + return criteria; + } + + /** + * Renders the (optionally capped) conversation to text for the decision input, one line per + * message prefixed with its role, with tool interactions described as data. Media blocks are + * counted and left out: decision models take text and JSON evidence only. When the request + * itself contributes no text (no USER or rendered-template text: a media-only, empty or + * SYSTEM-only conversation), the prompt arguments are the only signal and are appended; with + * none, the input is empty and the strategy abstains. + */ + static String renderInput( + RoutingStrategy strategy, + RoutingContext context, + List effectiveMessages, + Set pinnedIndices, + boolean[] truncatedOut, + Map metadata) { + List selected = + LlmJudgeRoutingExecutor.selectWithinBudget( + effectiveMessages, + LlmJudgeRoutingExecutor.maxContextChars(strategy), + pinnedIndices, + truncatedOut); + // Whether the request carries text is judged on the complete effective list, as the judge + // does, so a context cap that drops the USER message does not by itself trigger the + // prompt-argument fallback. Messages generated by the anchor's request shaping (the + // rendered template) carry the request like USER text does. + boolean requestHasText = false; + for (int i = 0; i < effectiveMessages.size(); i++) { + ChatMessage message = effectiveMessages.get(i); + if (message.getRole() != MessageRole.USER && !pinnedIndices.contains(i)) { + continue; + } + for (DataContentBlock block : LlmJudgeRoutingExecutor.judgeContent(message)) { + if (block instanceof TextBlock && !((TextBlock) block).getText().isEmpty()) { + requestHasText = true; + } + } + } + StringBuilder text = new StringBuilder(); + int mediaDropped = 0; + for (ChatMessage message : selected) { + List content = LlmJudgeRoutingExecutor.judgeContent(message); + if (content.isEmpty()) { + continue; + } + appendLine(text, message.getRole().name() + ": "); + boolean firstBlock = true; + for (DataContentBlock block : content) { + if (block instanceof TextBlock) { + // Blocks of one message are separate pieces of text; keep them apart so a + // message split around media does not read as one run-on word. + String blockText = ((TextBlock) block).getText(); + if (firstBlock) { + text.append(blockText); + } else { + appendLine(text, blockText); + } + firstBlock = false; + } else { + mediaDropped++; + } + } + } + if (mediaDropped > 0) { + metadata.put(MEDIA_DROPPED_KEY, mediaDropped); + } + if (!requestHasText && context.getPromptArgs().isEmpty()) { + // No request text and no arguments: there is nothing to decide on, and role framing + // alone is not evidence. An empty input makes the executor abstain. + return ""; + } + if (!requestHasText && !context.getPromptArgs().isEmpty()) { + text.append(text.length() == 0 ? "" : "\n").append("Request arguments:"); + for (Map.Entry arg : context.getPromptArgs().entrySet()) { + text.append('\n').append(arg.getKey()).append(": ").append(arg.getValue()); + } + } + return text.toString(); + } + + /** + * Turns the decision result into a routing decision. A pure function of the answer and the + * declaration: the choice selects its candidate with the reported probability as score; + * anything else abstains with the cause in the reason and the provider's answer in the + * metadata. + */ + static RoutingDecision toDecision( + RoutingStrategy strategy, + RoutingContext context, + DecisionResult result, + Map metadata) { + DecisionResult.TokenUsage usage = result.getUsage(); + if (usage != null && usage.getInputTokens() != null && usage.getOutputTokens() != null) { + metadata.put(INPUT_TOKENS_KEY, jsonNumber(usage.getInputTokens())); + metadata.put(OUTPUT_TOKENS_KEY, jsonNumber(usage.getOutputTokens())); + } + DecisionResult.Answer answer = result.getAnswers().get(QUESTION_NAME); + if (answer == null) { + return abstain( + "decision model returned no answer for question '" + QUESTION_NAME + "'", + metadata); + } + if (answer instanceof DecisionResult.RefusalAnswer) { + String reason = ((DecisionResult.RefusalAnswer) answer).getReason(); + return abstain( + "decision model refused" + (reason == null ? "" : ": " + reason), metadata); + } + if (!(answer instanceof DecisionResult.ChoiceAnswer)) { + return abstain("decision answer was not a choice", metadata); + } + DecisionResult.ChoiceAnswer choice = (DecisionResult.ChoiceAnswer) answer; + metadata.put(PROBABILITIES_KEY, new LinkedHashMap<>(choice.getProbabilities())); + if (choice.getProviderConfidence() != null) { + metadata.put(PROVIDER_CONFIDENCE_KEY, choice.getProviderConfidence()); + } + String selected = choice.getChoice(); + if (!isCandidate(context, selected)) { + metadata.put(CHOICE_KEY, selected); + return abstain("decision choice '" + selected + "' is not a candidate", metadata); + } + Double score = choice.getProbabilities().get(selected); + Double minProbability = minProbability(strategy); + if (minProbability != null && score == null) { + metadata.put(CHOICE_KEY, selected); + return abstain( + "decision model reported no probability for '" + + selected + + "', required by min_probability " + + minProbability, + metadata); + } + if (minProbability != null && score < minProbability) { + metadata.put(CHOICE_KEY, selected); + return abstain( + String.format( + "decision probability %s for '%s' is below min_probability %s", + score, selected, minProbability), + metadata); + } + RoutingDecision.Builder builder = + RoutingDecision.builder(selected).reason("decision model choice"); + if (score != null) { + builder.score(score); + } + for (Map.Entry entry : metadata.entrySet()) { + builder.metadata(entry.getKey(), entry.getValue()); + } + return builder.build(); + } + + /** + * Persisted as a real abstain: replay resolves to the router's current default, and the + * provider's answer stays in the metadata for the routing event. + */ + /** + * Appends {@code piece} on a new line, without doubling a line break that either side already + * carries (tool calls and results are rendered with their own framing newlines). + */ + private static void appendLine(StringBuilder text, String piece) { + if (text.length() > 0 + && text.charAt(text.length() - 1) != '\n' + && !piece.startsWith("\n")) { + text.append('\n'); + } + text.append(piece); + } + + /** + * Metadata travels as JSON in the route record, where a whole number that fits an int comes + * back as Integer. Store it in that shape up front so a fresh decision and a replayed one carry + * the same types. + */ + private static Object jsonNumber(long value) { + return value >= Integer.MIN_VALUE && value <= Integer.MAX_VALUE + ? (Object) (int) value + : (Object) value; + } + + private static RoutingDecision abstain(String reason, Map metadata) { + return new RoutingDecision(null, true, reason, null, new HashMap<>(metadata), null); + } + + private static boolean isCandidate(RoutingContext context, String model) { + for (RoutingCandidate candidate : context.getCandidates()) { + if (candidate.getName().equals(model)) { + return true; + } + } + return false; + } + + static String decisionModel(RoutingStrategy strategy) { + return (String) strategy.getArguments().get(RoutingStrategy.ARG_DECISION_MODEL); + } + + static String instructions(RoutingStrategy strategy) { + Object instructions = strategy.getArguments().get(RoutingStrategy.ARG_INSTRUCTIONS); + return instructions instanceof String ? (String) instructions : DEFAULT_INSTRUCTIONS; + } + + @Nullable + static Double minProbability(RoutingStrategy strategy) { + Object min = strategy.getArguments().get(RoutingStrategy.ARG_MIN_PROBABILITY); + return min instanceof Number ? ((Number) min).doubleValue() : null; + } +} diff --git a/plan/src/main/java/org/apache/flink/agents/plan/routing/LlmJudgeRoutingExecutor.java b/plan/src/main/java/org/apache/flink/agents/plan/routing/LlmJudgeRoutingExecutor.java index e947902e8c..1e18a7e8b5 100644 --- a/plan/src/main/java/org/apache/flink/agents/plan/routing/LlmJudgeRoutingExecutor.java +++ b/plan/src/main/java/org/apache/flink/agents/plan/routing/LlmJudgeRoutingExecutor.java @@ -195,9 +195,10 @@ private static List candidateNames(RoutingContext context) { * 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. + * Cancellation and interruption are not failures of the anchor: they propagate. */ - private static List effectiveJudgeMessages( - RoutingContext context, RunnerContext ctx) { + static List effectiveJudgeMessages(RoutingContext context, RunnerContext ctx) + throws Exception { List messages = context.getMessages(); String anchor = context.getDefaultModel().orElseGet(() -> context.getCandidates().get(0).getName()); @@ -210,6 +211,7 @@ private static List effectiveJudgeMessages( // candidate, where abstains resolve) is a documented approximation. return setup.prepareRequestMessages(messages, context.getPromptArgs()); } catch (Exception unresolvable) { + CancellationUtils.rethrowIfCancellation(unresolvable); // An unresolvable candidate surfaces on the real chat path with its normal policy. return messages; } @@ -221,7 +223,7 @@ private static List effectiveJudgeMessages( * 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 pinnedRenderedIndices( + static Set pinnedRenderedIndices( List original, List effective) { if (effective == original) { return Set.of(); @@ -246,7 +248,7 @@ private static String promptTemplate(RoutingStrategy strategy) { return (String) strategy.getArguments().get(RoutingStrategy.ARG_PROMPT_TEMPLATE); } - private static int maxContextChars(RoutingStrategy strategy) { + 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; } @@ -338,7 +340,7 @@ private static boolean carriesRequestContent( * remaining messages fill newest-first, and dropped messages set {@code truncatedOut[0]}. * Without a cap this returns the input unchanged. */ - private static List selectWithinBudget( + static List selectWithinBudget( List messages, int maxChars, Set pinnedIndices, @@ -421,7 +423,7 @@ private static List renderConversation(List messa * executable calls to the judge. Reasoning and provider continuation metadata are not * forwarded. */ - private static List judgeContent(ChatMessage message) { + static List judgeContent(ChatMessage message) { List content = new ArrayList<>(); for (ContentBlock block : message.getBlocks()) { if (block instanceof TextBlock) { diff --git a/plan/src/main/java/org/apache/flink/agents/plan/routing/ModelRoutingResolver.java b/plan/src/main/java/org/apache/flink/agents/plan/routing/ModelRoutingResolver.java index 149eb2fdc9..9a9968ff86 100644 --- a/plan/src/main/java/org/apache/flink/agents/plan/routing/ModelRoutingResolver.java +++ b/plan/src/main/java/org/apache/flink/agents/plan/routing/ModelRoutingResolver.java @@ -90,12 +90,7 @@ RoutingDecision unwrap() { } private static RoutingFailure failure(Exception e) throws Exception { - if (Thread.currentThread().isInterrupted() || CancellationUtils.isCancellation(e)) { - if (CancellationUtils.isInterruption(e)) { - Thread.currentThread().interrupt(); - } - throw e; - } + CancellationUtils.rethrowIfCancellation(e); return new RoutingFailure(ChatModelInvoker.errorText(e)); } diff --git a/plan/src/main/java/org/apache/flink/agents/plan/routing/RoutingExecutors.java b/plan/src/main/java/org/apache/flink/agents/plan/routing/RoutingExecutors.java index f37bae89a6..436be6ddd5 100644 --- a/plan/src/main/java/org/apache/flink/agents/plan/routing/RoutingExecutors.java +++ b/plan/src/main/java/org/apache/flink/agents/plan/routing/RoutingExecutors.java @@ -36,6 +36,7 @@ final class RoutingExecutors { static { EXECUTORS.put(RoutingStrategyType.RULE_BASED, new RuleBasedRoutingExecutor()); EXECUTORS.put(RoutingStrategyType.LLM_JUDGE, new LlmJudgeRoutingExecutor()); + EXECUTORS.put(RoutingStrategyType.DECISION_MODEL, new DecisionModelRoutingExecutor()); EXECUTORS.put(RoutingStrategyType.CUSTOM, new CustomRoutingExecutorAdapter()); } diff --git a/plan/src/main/java/org/apache/flink/agents/plan/utils/CancellationUtils.java b/plan/src/main/java/org/apache/flink/agents/plan/utils/CancellationUtils.java index 7a3dfe12f5..0443bd6245 100644 --- a/plan/src/main/java/org/apache/flink/agents/plan/utils/CancellationUtils.java +++ b/plan/src/main/java/org/apache/flink/agents/plan/utils/CancellationUtils.java @@ -57,4 +57,18 @@ private static boolean containsCancellation(Throwable failure, boolean includeCa } return false; } + + /** + * Rethrows {@code failure} when it is a cancellation signal (or the current thread is already + * interrupted), restoring the interrupt status for a wrapped interruption; returns normally for + * an ordinary failure so the caller can convert it into a retryable or failed outcome. + */ + public static void rethrowIfCancellation(Exception failure) throws Exception { + if (Thread.currentThread().isInterrupted() || isCancellation(failure)) { + if (isInterruption(failure)) { + Thread.currentThread().interrupt(); + } + throw failure; + } + } } diff --git a/plan/src/test/java/org/apache/flink/agents/plan/AgentPlanCrossLanguageTest.java b/plan/src/test/java/org/apache/flink/agents/plan/AgentPlanCrossLanguageTest.java index f01a99944c..9069b13884 100644 --- a/plan/src/test/java/org/apache/flink/agents/plan/AgentPlanCrossLanguageTest.java +++ b/plan/src/test/java/org/apache/flink/agents/plan/AgentPlanCrossLanguageTest.java @@ -317,6 +317,73 @@ void javaPlanWithJudgeRouterSnapshotIsStable() throws Exception { assertThat(actual).isEqualTo(expected); } + /** + * A plan whose only resources are a decision-model-strategy router, its chat models and the + * decision model. + */ + private static AgentPlan compileWithDecisionRouter() throws Exception { + Map> providers = new LinkedHashMap<>(); + ResourceDescriptor routerDescriptor = + ModelRouter.of("small", "big") + .describe("small", "short factual questions") + .describe("big", "code and sql") + .strategy(Strategies.decision("dm").withMinProbability(0.6)) + .defaultModel("small") + .fallback(true) + .build(); + Map routers = new LinkedHashMap<>(); + routers.put( + "router", + new JavaResourceProvider("router", ResourceType.MODEL_ROUTER, routerDescriptor)); + providers.put(ResourceType.MODEL_ROUTER, routers); + Map chatModels = new LinkedHashMap<>(); + ResourceDescriptor model = new ResourceDescriptor("some.Clazz", Map.of()); + chatModels.put("small", new JavaResourceProvider("small", ResourceType.CHAT_MODEL, model)); + chatModels.put("big", new JavaResourceProvider("big", ResourceType.CHAT_MODEL, model)); + providers.put(ResourceType.CHAT_MODEL, chatModels); + Map decisionModels = new LinkedHashMap<>(); + decisionModels.put( + "dm", + new JavaResourceProvider( + "dm", + ResourceType.DECISION_MODEL, + new ResourceDescriptor( + "some.Setup", Map.of("connection", "conn", "model", "m")))); + providers.put(ResourceType.DECISION_MODEL, decisionModels); + Map connections = new LinkedHashMap<>(); + connections.put( + "conn", + new JavaResourceProvider("conn", ResourceType.DECISION_MODEL_CONNECTION, model)); + providers.put(ResourceType.DECISION_MODEL_CONNECTION, connections); + return new AgentPlan(Map.of(), providers); + } + + @Test + void regenerateJavaPlanWithDecisionRouterSnapshot() throws Exception { + assumeTrue(regenerateRequested(), "Set -Dregenerate.snapshots=true to refresh."); + AgentPlan plan = compileWithDecisionRouter(); + String json = MAPPER.writerWithDefaultPrettyPrinter().writeValueAsString(plan); + + Path target = snapshotDir.resolve("java/agent_plan_with_decision_router.json"); + Files.createDirectories(target.getParent()); + Files.writeString(target, json + "\n"); + } + + @Test + void javaPlanWithDecisionRouterSnapshotIsStable() throws Exception { + Path committed = snapshotDir.resolve("java/agent_plan_with_decision_router.json"); + assertTrue( + Files.exists(committed), + "Decision-router plan snapshot missing from " + + committed + + ". Regenerate with -Dregenerate.snapshots=true and commit alongside the test."); + + AgentPlan plan = compileWithDecisionRouter(); + JsonNode actual = MAPPER.readTree(MAPPER.writeValueAsString(plan)); + JsonNode expected = MAPPER.readTree(Files.readString(committed)); + assertThat(actual).isEqualTo(expected); + } + /** A plan that nests a child agent carried by an ordinary resource descriptor. */ private static AgentPlan compileWithInternalSubagent() throws Exception { Agent child = new Agent(); diff --git a/plan/src/test/java/org/apache/flink/agents/plan/AgentPlanDecisionModelValidationTest.java b/plan/src/test/java/org/apache/flink/agents/plan/AgentPlanDecisionModelValidationTest.java new file mode 100644 index 0000000000..5f633e1d02 --- /dev/null +++ b/plan/src/test/java/org/apache/flink/agents/plan/AgentPlanDecisionModelValidationTest.java @@ -0,0 +1,239 @@ +/* + * 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; + +import org.apache.flink.agents.api.chat.model.routing.ModelRouter; +import org.apache.flink.agents.api.chat.model.routing.RoutingStrategy; +import org.apache.flink.agents.api.chat.model.routing.RoutingStrategyType; +import org.apache.flink.agents.api.chat.model.routing.Strategies; +import org.apache.flink.agents.api.resource.ResourceDescriptor; +import org.apache.flink.agents.api.resource.ResourceType; +import org.apache.flink.agents.plan.resourceprovider.JavaResourceProvider; +import org.apache.flink.agents.plan.resourceprovider.ResourceProvider; +import org.junit.jupiter.api.Test; + +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +import static org.assertj.core.api.Assertions.assertThatCode; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** Plan-construction validation of decision-model routing declarations. */ +public class AgentPlanDecisionModelValidationTest { + + private static Map> providers( + boolean withDecisionModel) { + Map> providers = new HashMap<>(); + setRouter( + providers, + ModelRouter.of("small", "big") + .describe("small", "short factual questions") + .describe("big", "code and sql") + .strategy(Strategies.decision("dm")) + .defaultModel("small") + .build()); + Map chatModels = new HashMap<>(); + ResourceDescriptor model = new ResourceDescriptor("some.Clazz", Map.of()); + chatModels.put("small", new JavaResourceProvider("small", ResourceType.CHAT_MODEL, model)); + chatModels.put("big", new JavaResourceProvider("big", ResourceType.CHAT_MODEL, model)); + providers.put(ResourceType.CHAT_MODEL, chatModels); + if (withDecisionModel) { + Map decisionModels = new HashMap<>(); + decisionModels.put( + "dm", + new JavaResourceProvider( + "dm", + ResourceType.DECISION_MODEL, + new ResourceDescriptor( + "some.Setup", Map.of("connection", "conn", "model", "m")))); + providers.put(ResourceType.DECISION_MODEL, decisionModels); + Map connections = new HashMap<>(); + connections.put( + "conn", + new JavaResourceProvider( + "conn", ResourceType.DECISION_MODEL_CONNECTION, model)); + providers.put(ResourceType.DECISION_MODEL_CONNECTION, connections); + } + return providers; + } + + private static void setRouter( + Map> providers, + ResourceDescriptor routerDescriptor) { + providers + .computeIfAbsent(ResourceType.MODEL_ROUTER, k -> new HashMap<>()) + .put( + "router", + new JavaResourceProvider( + "router", ResourceType.MODEL_ROUTER, routerDescriptor)); + } + + @Test + void registeredDecisionModelWithDescribedCandidatesPassesPlanConstruction() { + assertThatCode(() -> new AgentPlan(Map.of(), providers(true))).doesNotThrowAnyException(); + } + + @Test + void typoedDecisionModelFailsAtPlanConstruction() { + // Without this check every routed request would fail at resource lookup on the + // TaskManager; a static constraint belongs at plan construction. + assertThatThrownBy(() -> new AgentPlan(Map.of(), providers(false))) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("decision model 'dm'") + .hasMessageContaining("router"); + } + + /** + * The setup's connection is opened per request, so a typo'd name fails here, not per record. + */ + @Test + void typoedDecisionConnectionFailsAtPlanConstruction() { + Map> providers = providers(true); + providers.remove(ResourceType.DECISION_MODEL_CONNECTION); + assertThatThrownBy(() -> new AgentPlan(Map.of(), providers)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("connection 'conn'") + .hasMessageContaining("router"); + } + + /** A setup that declares no connection may resolve one from provider defaults; not checked. */ + @Test + void decisionSetupWithoutExplicitConnectionPassesPlanConstruction() { + Map> providers = providers(true); + providers.remove(ResourceType.DECISION_MODEL_CONNECTION); + providers + .get(ResourceType.DECISION_MODEL) + .put( + "dm", + new JavaResourceProvider( + "dm", + ResourceType.DECISION_MODEL, + new ResourceDescriptor("some.Setup", Map.of("model", "m")))); + assertThatCode(() -> new AgentPlan(Map.of(), providers)).doesNotThrowAnyException(); + } + + /** A non-String connection can never open; it fails at plan time with the router named. */ + @Test + void nonStringDecisionConnectionFailsAtPlanConstruction() { + Map> providers = providers(true); + providers + .get(ResourceType.DECISION_MODEL) + .put( + "dm", + new JavaResourceProvider( + "dm", + ResourceType.DECISION_MODEL, + new ResourceDescriptor( + "some.Setup", Map.of("connection", 5, "model", "m")))); + assertThatThrownBy(() -> new AgentPlan(Map.of(), providers)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("expected a connection name"); + } + + /** A non-String candidate description is diagnosed at plan time, for any strategy. */ + @Test + void nonStringCandidateDescriptionFailsAtPlanConstruction() { + Map> providers = providers(true); + Map args = new HashMap<>(); + args.put(ModelRouter.CANDIDATES_KEY, List.of("small", "big")); + args.put( + ModelRouter.CANDIDATE_DESCRIPTIONS_KEY, + Map.of("small", "short factual questions", "big", 42)); + args.put(ModelRouter.STRATEGY_TYPE_KEY, RoutingStrategyType.LLM_JUDGE.tag()); + args.put(ModelRouter.STRATEGY_ARGS_KEY, Map.of(RoutingStrategy.ARG_JUDGE_MODEL, "small")); + setRouter(providers, new ResourceDescriptor(ModelRouter.class.getName(), args)); + assertThatThrownBy(() -> new AgentPlan(Map.of(), providers)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("candidate 'big' has a Integer description"); + } + + /** A declared but blank model can never open; it fails at plan time. */ + @Test + void blankDecisionModelNameFailsAtPlanConstruction() { + Map> providers = providers(true); + providers + .get(ResourceType.DECISION_MODEL) + .put( + "dm", + new JavaResourceProvider( + "dm", + ResourceType.DECISION_MODEL, + new ResourceDescriptor( + "some.Setup", Map.of("connection", "conn", "model", " ")))); + assertThatThrownBy(() -> new AgentPlan(Map.of(), providers)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("blank or non-String 'model'"); + } + + /** An explicitly blank connection can never open; it fails at plan time. */ + @Test + void blankDecisionConnectionFailsAtPlanConstruction() { + Map> providers = providers(true); + providers + .get(ResourceType.DECISION_MODEL) + .put( + "dm", + new JavaResourceProvider( + "dm", + ResourceType.DECISION_MODEL, + new ResourceDescriptor( + "some.Setup", Map.of("connection", " ", "model", "m")))); + assertThatThrownBy(() -> new AgentPlan(Map.of(), providers)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("blank 'connection'"); + } + + /** + * The builder already rejects an undescribed candidate; a descriptor built by hand (e.g. + * deserialized) never ran it, so the plan applies the same rule. + */ + @Test + void undescribedCandidateFailsAtPlanConstruction() { + Map> providers = providers(true); + Map args = new HashMap<>(); + args.put(ModelRouter.CANDIDATES_KEY, List.of("small", "big")); + args.put(ModelRouter.CANDIDATE_DESCRIPTIONS_KEY, Map.of("big", "code and sql")); + args.put(ModelRouter.DEFAULT_MODEL_KEY, "small"); + args.put("fallback", false); + args.put(ModelRouter.STRATEGY_TYPE_KEY, RoutingStrategyType.DECISION_MODEL.tag()); + args.put(ModelRouter.STRATEGY_ARGS_KEY, Map.of(RoutingStrategy.ARG_DECISION_MODEL, "dm")); + setRouter(providers, new ResourceDescriptor(ModelRouter.class.getName(), args)); + assertThatThrownBy(() -> new AgentPlan(Map.of(), providers)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("requires a description for every candidate") + .hasMessageContaining("small"); + } + + /** A structurally invalid declaration (no decision model) fails like a judge without one. */ + @Test + void missingDecisionModelArgumentFailsAtPlanConstruction() { + Map> providers = providers(true); + Map args = new HashMap<>(); + args.put(ModelRouter.CANDIDATES_KEY, List.of("small", "big")); + args.put( + ModelRouter.CANDIDATE_DESCRIPTIONS_KEY, + Map.of("small", "short factual questions", "big", "code and sql")); + args.put(ModelRouter.STRATEGY_TYPE_KEY, RoutingStrategyType.DECISION_MODEL.tag()); + args.put(ModelRouter.STRATEGY_ARGS_KEY, Map.of()); + setRouter(providers, new ResourceDescriptor(ModelRouter.class.getName(), args)); + assertThatThrownBy(() -> new AgentPlan(Map.of(), providers)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining(RoutingStrategy.ARG_DECISION_MODEL); + } +} diff --git a/plan/src/test/java/org/apache/flink/agents/plan/actions/ChatModelActionRoutingTest.java b/plan/src/test/java/org/apache/flink/agents/plan/actions/ChatModelActionRoutingTest.java index 29aa97b8d4..ae59ea870a 100644 --- a/plan/src/test/java/org/apache/flink/agents/plan/actions/ChatModelActionRoutingTest.java +++ b/plan/src/test/java/org/apache/flink/agents/plan/actions/ChatModelActionRoutingTest.java @@ -42,6 +42,9 @@ import org.apache.flink.agents.api.context.MemoryRef; import org.apache.flink.agents.api.context.Outcome; import org.apache.flink.agents.api.context.RunnerContext; +import org.apache.flink.agents.api.decision.BaseDecisionModelSetup; +import org.apache.flink.agents.api.decision.DecisionRequest; +import org.apache.flink.agents.api.decision.DecisionResult; import org.apache.flink.agents.api.event.ChatRequestEvent; import org.apache.flink.agents.api.event.ChatResponseEvent; import org.apache.flink.agents.api.event.ModelRoutingEvent; @@ -70,9 +73,12 @@ import java.util.List; import java.util.Map; import java.util.Set; +import java.util.concurrent.CancellationException; import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.atLeastOnce; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -104,6 +110,7 @@ static class FakeChatModel extends BaseChatModelSetup { private int callCount; private boolean nativeStructuredOutput; private RuntimeException gateFailure; + private RuntimeException prepareFailure; final List> structuredRequests = new ArrayList<>(); FakeChatModel(Object... outcomes) { @@ -122,6 +129,21 @@ FakeChatModel withGateFailure(RuntimeException gateFailure) { return this; } + /** Makes request shaping (prompt rendering) throw, e.g. a cancellation signal. */ + FakeChatModel withPrepareFailure(RuntimeException prepareFailure) { + this.prepareFailure = prepareFailure; + return this; + } + + @Override + public List prepareRequestMessages( + List messages, Map promptArgs) { + if (prepareFailure != null) { + throw prepareFailure; + } + return super.prepareRequestMessages(messages, promptArgs); + } + // This fake binds no connection, so it answers the gate itself: a test chooses the native // answer or injects a gate failure. @Override @@ -164,12 +186,69 @@ public ChatResult chat( } } + /** + * A decision model stub: answers from a queue of outcomes, throwing on RuntimeException + * entries, and records every request it received. + */ + static class FakeDecisionModel extends BaseDecisionModelSetup { + private final Deque outcomes = new ArrayDeque<>(); + int callCount; + boolean async; + final List requests = new ArrayList<>(); + + FakeDecisionModel(Object... outcomes) { + super( + new ResourceDescriptor( + "fake-decision", Map.of("connection", "c", "model", "m")), + null); + Collections.addAll(this.outcomes, outcomes); + } + + static DecisionResult choice(String choice, Map probabilities) { + return new DecisionResult( + Map.of("model", new DecisionResult.ChoiceAnswer(choice, probabilities, null)), + "fake-model", + new DecisionResult.TokenUsage(10L, 2L, 12L), + Map.of()); + } + + FakeDecisionModel supportingAsync() { + this.async = true; + return this; + } + + @Override + public boolean supportsAsyncExecution() { + return async; + } + + @Override + public DecisionResult decide(DecisionRequest request) { + return decide(request, Map.of()); + } + + @Override + public DecisionResult decide(DecisionRequest request, Map requestParams) { + callCount++; + requests.add(request); + Object next = outcomes.isEmpty() ? null : outcomes.poll(); + if (next instanceof RuntimeException) { + throw (RuntimeException) next; + } + if (next instanceof DecisionResult) { + return (DecisionResult) next; + } + return choice("small", Map.of("small", 1.0, "big", 0.0)); + } + } + static class FakeRunnerContext implements RunnerContext { final List sentEvents = new ArrayList<>(); final List resolvedChatModels = new ArrayList<>(); final List durableCallIds = new ArrayList<>(); final List durableAsyncCallIds = new ArrayList<>(); final Map models = new HashMap<>(); + final Map decisionModels = new HashMap<>(); final Set unresolvable = new HashSet<>(); private final ModelRouter router; private final MemoryObject sensoryMemory = new FakeMemoryObject(new HashMap<>()); @@ -188,12 +267,22 @@ FakeRunnerContext register(String name, BaseChatModelSetup model) { return this; } + FakeRunnerContext registerDecision(String name, BaseDecisionModelSetup model) { + decisionModels.put(name, model); + return this; + } + /** Marks a chat-model name whose resource lookup fails (e.g. a typo'd candidate). */ FakeRunnerContext unresolvable(String name) { unresolvable.add(name); return this; } + FakeRunnerContext withDecisionAsync() { + config.set(AgentExecutionOptions.DECISION_ASYNC, true); + return this; + } + FakeRunnerContext withRetryBudget(int maxRetries, int waitIntervalSec) { config.set(AgentExecutionOptions.MAX_RETRIES, maxRetries); config.set(AgentExecutionOptions.RETRY_WAIT_INTERVAL, waitIntervalSec); @@ -222,6 +311,13 @@ public Resource getResource(String name, ResourceType type) { resolvedChatModels.add(name); return models.getOrDefault(name, new FakeChatModel()); } + if (type == ResourceType.DECISION_MODEL) { + BaseDecisionModelSetup model = decisionModels.get(name); + if (model == null) { + throw new IllegalArgumentException("resource not found: " + name); + } + return model; + } throw new IllegalArgumentException("unexpected resource " + name + " " + type); } @@ -280,6 +376,8 @@ FakeRunnerContext seedDurable(String id, Object value) { value = outcome; } else if (value instanceof ChatResult) { value = new ChatModelInvoker.InvocationOutcome((ChatResult) value, null); + } else if (value instanceof DecisionResult) { + value = new DecisionModelInvoker.DecisionOutcome((DecisionResult) value, null); } durableStore.put(id, value); return this; @@ -1826,4 +1924,347 @@ void llmJudgeWithUndeclaredInstanceLevelPromptAbstainsToDefault() throws Excepti assertThat(event.getReason()).contains("not a candidate name"); assertThat(ctx.hasChatResponse()).isTrue(); } + + // ── Decision-model strategy ───────────────────────────────────────────── + + private static ModelRouter decisionRouter(RoutingStrategy strategy) throws Exception { + return new ModelRouter( + ModelRouter.of("small", "big") + .describe("small", "cheap chit-chat") + .describe("big", "code and sql") + .strategy(strategy) + .defaultModel("small") + .build(), + null); + } + + @Test + void decisionModelChoiceRoutesToNamedCandidate() throws Exception { + FakeDecisionModel dm = + new FakeDecisionModel( + FakeDecisionModel.choice("big", Map.of("small", 0.18, "big", 0.82))); + FakeRunnerContext ctx = + new FakeRunnerContext(decisionRouter(Strategies.decision("dm"))) + .registerDecision("dm", dm) + .register("big", new FakeChatModel()); + ChatModelAction.processChatRequestOrToolResponse( + new ChatRequestEvent( + "router", List.of(new ChatMessage(MessageRole.USER, "write some sql"))), + ctx); + + ModelRoutingEvent event = ctx.routingEvent(); + assertThat(event).isNotNull(); + assertThat(event.getSelectedModel()).isEqualTo("big"); + assertThat(event.getDecisionSource()).isEqualTo(ModelRoutingEvent.SOURCE_DECISION_MODEL); + assertThat(event.getScore()).isEqualTo(0.82); + assertThat(event.getMetadata()) + .containsEntry("decision_model", "dm") + .containsKey("decision_probabilities"); + assertThat(event.getDecisionMs()).isNotNull(); + // the decision call is durable under its own id; the decision persists under the route id + assertThat(ctx.durableCallIds).contains("decide:router", "route:router"); + // resolution order: the anchor candidate (bound-prompt rendering), then the winner answers + assertThat(ctx.resolvedChatModels).containsExactly("small", "big"); + assertThat(ctx.hasChatResponse()).isTrue(); + assertThat(ctx.chatResponse().isFailed()).isFalse(); + // one decision call, carrying the conversation and both candidates as choices + assertThat(dm.callCount).isEqualTo(1); + DecisionRequest request = dm.requests.get(0); + assertThat(String.valueOf(request.getInput())).contains("write some sql"); + DecisionRequest.ChoiceQuestion question = + (DecisionRequest.ChoiceQuestion) request.getQuestions().get("model"); + assertThat(question.getCriteria()) + .containsEntry("small", "cheap chit-chat") + .containsEntry("big", "code and sql"); + } + + @Test + void decisionBelowMinProbabilityAbstainsToDefault() throws Exception { + FakeDecisionModel dm = + new FakeDecisionModel( + FakeDecisionModel.choice("big", Map.of("small", 0.45, "big", 0.55))); + FakeRunnerContext ctx = + new FakeRunnerContext( + decisionRouter(Strategies.decision("dm").withMinProbability(0.6))) + .registerDecision("dm", dm) + .register("small", new FakeChatModel()); + ChatModelAction.processChatRequestOrToolResponse( + new ChatRequestEvent("router", List.of(new ChatMessage(MessageRole.USER, "hmm"))), + ctx); + + ModelRoutingEvent event = ctx.routingEvent(); + assertThat(event.getSelectedModel()).isEqualTo("small"); + assertThat(event.getDecisionSource()).isEqualTo(ModelRoutingEvent.SOURCE_DEFAULT); + assertThat(event.getReason()).contains("below min_probability"); + assertThat(event.getScore()).isNull(); + assertThat(event.getMetadata()) + .containsEntry("decision_choice", "big") + .containsKey("decision_probabilities") + // the strategy abstained on its own; non-candidate normalization is not involved + .doesNotContainKey(ModelRoutingResolver.REJECTED_MODEL_KEY); + assertThat(ctx.resolvedChatModels).containsExactly("small", "small"); + assertThat(ctx.hasChatResponse()).isTrue(); + } + + @Test + void decisionRefusalAbstainsToDefault() throws Exception { + DecisionResult refused = + new DecisionResult( + Map.of("model", new DecisionResult.RefusalAnswer("insufficient evidence")), + null, + null, + Map.of()); + FakeRunnerContext ctx = + new FakeRunnerContext(decisionRouter(Strategies.decision("dm"))) + .registerDecision("dm", new FakeDecisionModel(refused)) + .register("small", new FakeChatModel()); + ChatModelAction.processChatRequestOrToolResponse( + new ChatRequestEvent("router", List.of(new ChatMessage(MessageRole.USER, "hi"))), + ctx); + + ModelRoutingEvent event = ctx.routingEvent(); + assertThat(event.getSelectedModel()).isEqualTo("small"); + assertThat(event.getDecisionSource()).isEqualTo(ModelRoutingEvent.SOURCE_DEFAULT); + assertThat(event.getReason()).isEqualTo("decision model refused: insufficient evidence"); + assertThat(ctx.hasChatResponse()).isTrue(); + } + + @Test + void decisionModelUsesTheRetryBudgetAndRecordsMetrics() throws Exception { + FlinkAgentsMetricGroup actionMetricGroup = mock(FlinkAgentsMetricGroup.class); + FlinkAgentsMetricGroup dmMetricGroup = mock(FlinkAgentsMetricGroup.class); + FlinkAgentsMetricGroup tokenMetricGroup = mock(FlinkAgentsMetricGroup.class); + Counter retryCount = mock(Counter.class); + Counter retryWaitSec = mock(Counter.class); + Counter promptTokens = mock(Counter.class); + when(actionMetricGroup.getHistogram("routingDecisionLatencyMs")) + .thenReturn(mock(Histogram.class)); + when(actionMetricGroup.getSubGroup("model_resource", "dm")).thenReturn(dmMetricGroup); + when(dmMetricGroup.getCounter("retryCount")).thenReturn(retryCount); + when(dmMetricGroup.getCounter("retryWaitSec")).thenReturn(retryWaitSec); + when(dmMetricGroup.getSubGroup("model", "fake-model")).thenReturn(tokenMetricGroup); + when(tokenMetricGroup.getCounter(anyString())).thenReturn(mock(Counter.class)); + when(tokenMetricGroup.getCounter("promptTokens")).thenReturn(promptTokens); + FakeDecisionModel dm = + new FakeDecisionModel( + new RuntimeException("transient"), + FakeDecisionModel.choice("big", Map.of("small", 0.1, "big", 0.9))); + FakeRunnerContext ctx = + new FakeRunnerContext(decisionRouter(Strategies.decision("dm"))) + .withRetryBudget(1, 0) + .withActionMetricGroup(actionMetricGroup) + .registerDecision("dm", dm) + .register("big", new FakeChatModel()); + ChatModelAction.processChatRequestOrToolResponse( + new ChatRequestEvent( + "router", List.of(new ChatMessage(MessageRole.USER, "hard one"))), + ctx); + + ModelRoutingEvent event = ctx.routingEvent(); + assertThat(event.getSelectedModel()).isEqualTo("big"); + assertThat(event.getDecisionSource()).isEqualTo(ModelRoutingEvent.SOURCE_DECISION_MODEL); + assertThat(dm.callCount).isEqualTo(2); + // the decision retry counts against the routed request, like a judge retry + assertThat(ctx.chatResponse().getRetryCount()).isEqualTo(1); + // token metrics and retry counters both land on the decision model's resource group + verify(actionMetricGroup, atLeastOnce()).getSubGroup("model_resource", "dm"); + verify(retryCount).inc(1); + verify(retryWaitSec).inc(0); + verify(promptTokens).inc(10L); + } + + @Test + void decisionModelFailureFailsTheRequestWithoutCandidateFallback() throws Exception { + // The default zero retry budget produces a failed terminal response, like a failed judge. + FakeDecisionModel dm = new FakeDecisionModel(new RuntimeException("decision down")); + FakeRunnerContext ctx = + new FakeRunnerContext( + new ModelRouter( + ModelRouter.of("small", "big") + .describe("small", "cheap chit-chat") + .describe("big", "code and sql") + .strategy(Strategies.decision("dm")) + .defaultModel("small") + .fallback(true) + .build(), + null)) + .registerDecision("dm", dm) + .register("small", new FakeChatModel()) + .register("big", new FakeChatModel()); + ChatModelAction.processChatRequestOrToolResponse( + new ChatRequestEvent("router", List.of(new ChatMessage(MessageRole.USER, "hi"))), + ctx); + + assertThat(ctx.chatResponse().isFailed()).isTrue(); + assertThat(ctx.chatResponse().getError()) + .contains("Decision model 'dm' failed") + .contains("decision down"); + // no routing event and no candidate was called: only the anchor was resolved for + // rendering the decision input + assertThat(ctx.routingEventCount()).isZero(); + assertThat(ctx.resolvedChatModels).containsExactly("small"); + } + + @Test + void unresolvableDecisionModelFailsTheRequest() throws Exception { + FakeRunnerContext ctx = + new FakeRunnerContext(decisionRouter(Strategies.decision("dm"))) + .register("small", new FakeChatModel()); + ChatModelAction.processChatRequestOrToolResponse( + new ChatRequestEvent("router", List.of(new ChatMessage(MessageRole.USER, "hi"))), + ctx); + + assertThat(ctx.chatResponse().isFailed()).isTrue(); + assertThat(ctx.chatResponse().getError()).contains("Decision model 'dm'"); + assertThat(ctx.routingEventCount()).isZero(); + } + + @Test + void decisionReplayReusesTheStoredAnswerWithoutCallingTheProvider() throws Exception { + // Recovery: the decide: record is already stored, so the provider is not called + // again even though it would fail now; the stored choice routes the request. + FakeDecisionModel dm = new FakeDecisionModel(new RuntimeException("would fail now")); + FakeRunnerContext ctx = + new FakeRunnerContext(decisionRouter(Strategies.decision("dm"))) + .registerDecision("dm", dm) + .register("big", new FakeChatModel()) + .seedDurable( + "decide:router", + FakeDecisionModel.choice("big", Map.of("small", 0.2, "big", 0.8))); + ChatModelAction.processChatRequestOrToolResponse( + new ChatRequestEvent( + "router", List.of(new ChatMessage(MessageRole.USER, "write some sql"))), + ctx); + + assertThat(dm.callCount).isZero(); + ModelRoutingEvent event = ctx.routingEvent(); + assertThat(event.getSelectedModel()).isEqualTo("big"); + assertThat(event.getDecisionSource()).isEqualTo(ModelRoutingEvent.SOURCE_DECISION_MODEL); + assertThat(ctx.resolvedChatModels).containsExactly("small", "big"); + assertThat(ctx.chatResponse().isFailed()).isFalse(); + } + + /** + * Cancellation raised while the anchor renders the decision input is a cancellation signal, not + * an unresolvable anchor: it propagates before the provider or any candidate is called. + */ + @Test + void cancellationDuringDecisionInputPreparationPropagatesBeforeTheProviderCall() + throws Exception { + FakeDecisionModel dm = + new FakeDecisionModel( + FakeDecisionModel.choice("big", Map.of("small", 0.1, "big", 0.9))); + FakeChatModel big = new FakeChatModel(); + FakeRunnerContext ctx = + new FakeRunnerContext(decisionRouter(Strategies.decision("dm"))) + .registerDecision("dm", dm) + .register( + "small", + new FakeChatModel() + .withPrepareFailure(new CancellationException("cancelled"))) + .register("big", big); + assertThatThrownBy( + () -> + ChatModelAction.processChatRequestOrToolResponse( + new ChatRequestEvent( + "router", + List.of( + new ChatMessage( + MessageRole.USER, + "write some sql"))), + ctx)) + .isInstanceOf(CancellationException.class) + .hasMessage("cancelled"); + assertThat(dm.callCount).isZero(); + assertThat(big.callCount).isZero(); + assertThat(ctx.durableCallIds).isEmpty(); + assertThat(ctx.hasChatResponse()).isFalse(); + } + + /** The judge shares the rendering helper, so it propagates the same signal. */ + @Test + void cancellationDuringJudgeInputPreparationPropagatesBeforeTheJudgeCall() throws Exception { + FakeChatModel judge = new FakeChatModel(ChatFixtures.response("{\"model\": \"big\"}")); + FakeRunnerContext ctx = + new FakeRunnerContext( + new ModelRouter( + ModelRouter.of("small", "big") + .strategy(Strategies.llm("judge")) + .defaultModel("small") + .build(), + null)) + .register("judge", judge) + .register( + "small", + new FakeChatModel() + .withPrepareFailure(new CancellationException("cancelled"))) + .register("big", new FakeChatModel()); + assertThatThrownBy( + () -> + ChatModelAction.processChatRequestOrToolResponse( + new ChatRequestEvent( + "router", + List.of( + new ChatMessage( + MessageRole.USER, + "write some sql"))), + ctx)) + .isInstanceOf(CancellationException.class); + assertThat(judge.callCount).isZero(); + assertThat(ctx.durableCallIds).isEmpty(); + assertThat(ctx.hasChatResponse()).isFalse(); + } + + @Test + void asyncDecisionCallRunsAsFlatSiblingBeforeTheRouteRecord() throws Exception { + // The sequencing contract: the executor's durable call lands ahead of the resolver's + // route: record, on the async path as on the sync one. + FakeDecisionModel dm = + new FakeDecisionModel( + FakeDecisionModel.choice("big", Map.of("small", 0.1, "big", 0.9))) + .supportingAsync(); + FakeRunnerContext ctx = + new FakeRunnerContext(decisionRouter(Strategies.decision("dm"))) + .withDecisionAsync() + .registerDecision("dm", dm) + .register("big", new FakeChatModel()); + ChatModelAction.processChatRequestOrToolResponse( + new ChatRequestEvent( + "router", List.of(new ChatMessage(MessageRole.USER, "write some sql"))), + ctx); + + assertThat(ctx.durableAsyncCallIds).contains("decide:router"); + assertThat(ctx.durableCallIds.indexOf("decide:router")) + .isLessThan(ctx.durableCallIds.indexOf("route:router")); + assertThat(ctx.routingEvent().getSelectedModel()).isEqualTo("big"); + assertThat(ctx.chatResponse().isFailed()).isFalse(); + } + + @Test + void durableStateFailureDuringTheDecisionCallPropagatesToTheRunner() throws Exception { + // Persistence and recovery failures are not routing failures: they escape the executor + // and the resolver unchanged, with no terminal chat response. + FakeRunnerContext ctx = + new FakeRunnerContext(decisionRouter(Strategies.decision("dm"))) { + @Override + public T durableExecute(DurableCallable callable) throws Exception { + if ("decide:router".equals(callable.getId())) { + throw new IllegalStateException("state backend unavailable"); + } + return super.durableExecute(callable); + } + }; + ctx.registerDecision("dm", new FakeDecisionModel()).register("big", new FakeChatModel()); + assertThatThrownBy( + () -> + ChatModelAction.processChatRequestOrToolResponse( + new ChatRequestEvent( + "router", + List.of(new ChatMessage(MessageRole.USER, "hi"))), + ctx)) + .isInstanceOf(IllegalStateException.class) + .hasMessage("state backend unavailable"); + assertThat(ctx.hasChatResponse()).isFalse(); + assertThat(ctx.routingEventCount()).isZero(); + } } diff --git a/plan/src/test/java/org/apache/flink/agents/plan/routing/DecisionModelRoutingExecutorTest.java b/plan/src/test/java/org/apache/flink/agents/plan/routing/DecisionModelRoutingExecutorTest.java new file mode 100644 index 0000000000..fb6a0286a2 --- /dev/null +++ b/plan/src/test/java/org/apache/flink/agents/plan/routing/DecisionModelRoutingExecutorTest.java @@ -0,0 +1,549 @@ +/* + * 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.routing; + +import org.apache.flink.agents.api.chat.messages.ChatMessage; +import org.apache.flink.agents.api.chat.messages.ImageBlock; +import org.apache.flink.agents.api.chat.messages.MessageRole; +import org.apache.flink.agents.api.chat.messages.TextBlock; +import org.apache.flink.agents.api.chat.messages.ToolCallBlock; +import org.apache.flink.agents.api.chat.messages.ToolResultBlock; +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.chat.model.routing.Strategies; +import org.apache.flink.agents.api.context.RunnerContext; +import org.apache.flink.agents.api.decision.BaseDecisionModelSetup; +import org.apache.flink.agents.api.decision.DecisionRequest; +import org.apache.flink.agents.api.decision.DecisionResult; +import org.apache.flink.agents.api.resource.ResourceType; +import org.junit.jupiter.api.Test; + +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.UUID; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +/** + * Unit tests for the decision-model executor's pure functions (request building, answer mapping). + */ +class DecisionModelRoutingExecutorTest { + + private static RoutingContext ctx(List messages, Map promptArgs) { + return new RoutingContext( + UUID.randomUUID(), + "router", + messages, + promptArgs, + List.of( + new RoutingCandidate("small", "cheap chit-chat"), + new RoutingCandidate("big", "code and sql")), + "small"); + } + + /** A runner context whose anchor candidate cannot be resolved: the raw messages are used. */ + private static RunnerContext unresolvableAnchor() { + RunnerContext ctx = mock(RunnerContext.class); + try { + when(ctx.getResource(any(), any())) + .thenThrow(new IllegalArgumentException("no resource")); + } catch (Exception e) { + throw new AssertionError(e); + } + return ctx; + } + + private static DecisionResult choice( + String choice, Map probabilities, Double providerConfidence) { + return new DecisionResult( + Map.of( + DecisionModelRoutingExecutor.QUESTION_NAME, + new DecisionResult.ChoiceAnswer(choice, probabilities, providerConfidence)), + "dm-v1", + new DecisionResult.TokenUsage(120L, 4L, 124L), + Map.of()); + } + + private static DecisionRequest.ChoiceQuestion question(DecisionRequest request) { + return (DecisionRequest.ChoiceQuestion) + request.getQuestions().get(DecisionModelRoutingExecutor.QUESTION_NAME); + } + + @Test + void requestCarriesCandidatesAsChoiceCriteriaAndConversationAsText() throws Exception { + RoutingContext context = + ctx(List.of(new ChatMessage(MessageRole.USER, "write some sql")), Map.of()); + Map metadata = new LinkedHashMap<>(); + DecisionRequest request = + DecisionModelRoutingExecutor.buildRequest( + Strategies.decision("dm"), context, unresolvableAnchor(), metadata); + assertThat(request.getInput()).isEqualTo("USER: write some sql"); + DecisionRequest.ChoiceQuestion question = question(request); + assertThat(question.getCriteria()) + .containsExactly( + Map.entry("small", "cheap chit-chat"), Map.entry("big", "code and sql")); + assertThat(question.getInstructions()) + .isEqualTo(DecisionModelRoutingExecutor.DEFAULT_INSTRUCTIONS); + assertThat(metadata) + .doesNotContainKeys( + DecisionModelRoutingExecutor.CONTEXT_TRUNCATED_KEY, + DecisionModelRoutingExecutor.MEDIA_DROPPED_KEY); + } + + @Test + void customInstructionsReplaceTheDefault() throws Exception { + RoutingContext context = ctx(List.of(new ChatMessage(MessageRole.USER, "hi")), Map.of()); + DecisionRequest request = + DecisionModelRoutingExecutor.buildRequest( + Strategies.decision("dm", "Which model should answer?"), + context, + unresolvableAnchor(), + new LinkedHashMap<>()); + assertThat(question(request).getInstructions()).isEqualTo("Which model should answer?"); + } + + @Test + void requestIncludesFullConversationWithRoles() throws Exception { + RoutingContext context = + ctx( + List.of( + new ChatMessage(MessageRole.SYSTEM, "You review concurrency code"), + new ChatMessage(MessageRole.USER, "Focus on race conditions"), + new ChatMessage(MessageRole.ASSISTANT, "Understood"), + new ChatMessage(MessageRole.USER, "synchronized void transfer()")), + Map.of()); + String input = + (String) + DecisionModelRoutingExecutor.buildRequest( + Strategies.decision("dm"), + context, + unresolvableAnchor(), + new LinkedHashMap<>()) + .getInput(); + assertThat(input) + .contains("SYSTEM: You review concurrency code") + .contains("USER: Focus on race conditions") + .contains("ASSISTANT: Understood") + .contains("USER: synchronized void transfer()"); + } + + @Test + void contextCapPinsSystemAndNewestAndFlagsTruncation() throws Exception { + String oldTurn = "y".repeat(500); + RoutingContext context = + ctx( + List.of( + new ChatMessage(MessageRole.SYSTEM, "task framing"), + new ChatMessage(MessageRole.USER, oldTurn), + new ChatMessage(MessageRole.USER, "current question")), + Map.of()); + Map metadata = new LinkedHashMap<>(); + String input = + (String) + DecisionModelRoutingExecutor.buildRequest( + Strategies.decision("dm").withMaxContextChars(80), + context, + unresolvableAnchor(), + metadata) + .getInput(); + assertThat(input).contains("SYSTEM: task framing").contains("USER: current question"); + assertThat(input).doesNotContain(oldTurn); + assertThat(metadata) + .containsEntry(DecisionModelRoutingExecutor.CONTEXT_TRUNCATED_KEY, true); + } + + @Test + void mediaIsLeftOutOfTheTextInputAndCounted() throws Exception { + ChatMessage message = + ChatMessage.user( + List.of( + new TextBlock("before"), + ImageBlock.fromBase64("image/png", "aGk="), + new TextBlock("after"))); + Map metadata = new LinkedHashMap<>(); + String input = + (String) + DecisionModelRoutingExecutor.buildRequest( + Strategies.decision("dm"), + ctx(List.of(message), Map.of()), + unresolvableAnchor(), + metadata) + .getInput(); + assertThat(input).isEqualTo("USER: before\nafter"); + assertThat(metadata).containsEntry(DecisionModelRoutingExecutor.MEDIA_DROPPED_KEY, 1); + } + + /** Tool rounds reach the decision model as data, framed the way the judge frames them. */ + @Test + void toolCallsAndResultsAreRenderedAsText() throws Exception { + RoutingContext context = + ctx( + List.of( + ChatMessage.user("check the weather"), + ChatMessage.assistant( + List.of( + new ToolCallBlock( + "call-1", + "lookup", + Map.of("q", "weather")))), + ChatMessage.tool( + new ToolResultBlock( + "call-1", List.of(new TextBlock("sunny")), false))), + Map.of()); + String input = + (String) + DecisionModelRoutingExecutor.buildRequest( + Strategies.decision("dm"), + context, + unresolvableAnchor(), + new LinkedHashMap<>()) + .getInput(); + // One line per rendered piece, no blank lines: the tool framing's own newlines are not + // doubled by the block or message separators. JSON key order is the mapper's, so the + // payloads are checked by content. + List lines = List.of(input.split("\n")); + assertThat(lines).hasSize(5); + assertThat(lines.get(0)).isEqualTo("USER: check the weather"); + assertThat(lines.get(1)).isEqualTo("ASSISTANT: "); + assertThat(lines.get(2)) + .startsWith("Tool call {") + .contains("\"name\":\"lookup\"") + .contains("\"call_id\":\"call-1\""); + assertThat(lines.get(3)) + .startsWith("TOOL: Tool result {") + .contains("\"is_error\":false") + .endsWith("}:"); + assertThat(lines.get(4)).isEqualTo("sunny"); + } + + /** A result that answers other questions but not ours is a missing answer, not a wrong one. */ + @Test + void missingAnswerForTheQuestionAbstains() { + DecisionResult other = + new DecisionResult( + Map.of("other", new DecisionResult.PredicateAnswer(0.4)), + null, + null, + Map.of()); + RoutingDecision decision = + DecisionModelRoutingExecutor.toDecision( + Strategies.decision("dm"), + ctx(List.of(), Map.of()), + other, + new LinkedHashMap<>()); + assertThat(decision.isAbstain()).isTrue(); + assertThat(decision.getReason()) + .isEqualTo("decision model returned no answer for question 'model'"); + } + + /** The cap may drop the USER text from the input, but that is not a request without text. */ + @Test + void contextCapDoesNotTriggerThePromptArgsFallback() throws Exception { + RoutingContext context = + ctx( + List.of( + new ChatMessage(MessageRole.USER, "x".repeat(200)), + new ChatMessage(MessageRole.ASSISTANT, "ok")), + Map.of("input", "unused-args")); + Map metadata = new LinkedHashMap<>(); + String input = + (String) + DecisionModelRoutingExecutor.buildRequest( + Strategies.decision("dm").withMaxContextChars(60), + context, + unresolvableAnchor(), + metadata) + .getInput(); + assertThat(input).isEqualTo("ASSISTANT: ok"); + assertThat(metadata) + .containsEntry(DecisionModelRoutingExecutor.CONTEXT_TRUNCATED_KEY, true); + } + + /** A media-only request carries no text the decision model can read, so the args step in. */ + @Test + void mediaOnlyRequestFallsBackToPromptArgs() throws Exception { + Map metadata = new LinkedHashMap<>(); + String input = + (String) + DecisionModelRoutingExecutor.buildRequest( + Strategies.decision("dm"), + ctx( + List.of( + ChatMessage.user( + List.of( + ImageBlock.fromBase64( + "image/png", + "aGk=")))), + Map.of("input", "describe the picture")), + unresolvableAnchor(), + metadata) + .getInput(); + assertThat(input).contains("Request arguments:").contains("input: describe the picture"); + assertThat(metadata).containsEntry(DecisionModelRoutingExecutor.MEDIA_DROPPED_KEY, 1); + } + + @Test + void promptArgsAreTheInputWhenTheConversationCarriesNoRequestText() throws Exception { + RoutingContext context = + ctx( + List.of( + new ChatMessage(MessageRole.SYSTEM, "You are a helpful assistant"), + new ChatMessage(MessageRole.USER, "")), + Map.of("input", "write some sql for active users")); + String input = + (String) + DecisionModelRoutingExecutor.buildRequest( + Strategies.decision("dm"), + context, + unresolvableAnchor(), + new LinkedHashMap<>()) + .getInput(); + assertThat(input) + .contains("SYSTEM: You are a helpful assistant") + .contains("Request arguments:") + .contains("input: write some sql for active users"); + } + + @Test + void choiceSelectsItsCandidateWithTheReportedProbabilityAsScore() { + Map metadata = new LinkedHashMap<>(); + metadata.put(DecisionModelRoutingExecutor.DECISION_MODEL_KEY, "dm"); + RoutingDecision decision = + DecisionModelRoutingExecutor.toDecision( + Strategies.decision("dm"), + ctx(List.of(), Map.of()), + choice("big", Map.of("small", 0.18, "big", 0.82), 0.9), + metadata); + assertThat(decision.isAbstain()).isFalse(); + assertThat(decision.getSelectedModel()).isEqualTo("big"); + assertThat(decision.getScore()).isEqualTo(0.82); + assertThat(decision.getReason()).isEqualTo("decision model choice"); + assertThat(decision.getMetadata()) + .containsEntry(DecisionModelRoutingExecutor.DECISION_MODEL_KEY, "dm") + .containsEntry(DecisionModelRoutingExecutor.PROVIDER_CONFIDENCE_KEY, 0.9) + // JSON shape: whole numbers that fit an int are Integers, fresh or replayed. + .containsEntry(DecisionModelRoutingExecutor.INPUT_TOKENS_KEY, 120) + .containsEntry(DecisionModelRoutingExecutor.OUTPUT_TOKENS_KEY, 4); + @SuppressWarnings("unchecked") + Map probabilities = + (Map) + decision.getMetadata().get(DecisionModelRoutingExecutor.PROBABILITIES_KEY); + assertThat(probabilities).containsEntry("big", 0.82).containsEntry("small", 0.18); + } + + /** The provider's choice is trusted even when it is not the highest probability. */ + @Test + void scoreIsTheChosenCandidateProbabilityNotTheMaximum() { + RoutingDecision decision = + DecisionModelRoutingExecutor.toDecision( + Strategies.decision("dm"), + ctx(List.of(), Map.of()), + choice("small", Map.of("small", 0.4, "big", 0.6), null), + new LinkedHashMap<>()); + assertThat(decision.getSelectedModel()).isEqualTo("small"); + assertThat(decision.getScore()).isEqualTo(0.4); + assertThat(decision.getMetadata()) + .doesNotContainKey(DecisionModelRoutingExecutor.PROVIDER_CONFIDENCE_KEY); + } + + @Test + void probabilityBelowMinimumAbstainsAndRecordsTheChoice() { + RoutingDecision decision = + DecisionModelRoutingExecutor.toDecision( + Strategies.decision("dm").withMinProbability(0.9), + ctx(List.of(), Map.of()), + choice("big", Map.of("small", 0.18, "big", 0.82), null), + new LinkedHashMap<>()); + assertThat(decision.isAbstain()).isTrue(); + assertThat(decision.getSelectedModel()).isNull(); + assertThat(decision.getScore()).isNull(); + assertThat(decision.getReason()).contains("0.82").contains("below min_probability 0.9"); + assertThat(decision.getMetadata()) + .containsEntry(DecisionModelRoutingExecutor.CHOICE_KEY, "big") + .containsKey(DecisionModelRoutingExecutor.PROBABILITIES_KEY); + } + + /** A provider may omit the chosen candidate from the map; with a minimum set that abstains. */ + @Test + void missingProbabilityForTheChoiceAbstainsWhenAMinimumIsSet() { + RoutingDecision decision = + DecisionModelRoutingExecutor.toDecision( + Strategies.decision("dm").withMinProbability(0.6), + ctx(List.of(), Map.of()), + choice("big", Map.of("small", 1.0), null), + new LinkedHashMap<>()); + assertThat(decision.isAbstain()).isTrue(); + assertThat(decision.getReason()) + .isEqualTo( + "decision model reported no probability for 'big', required by" + + " min_probability 0.6"); + assertThat(decision.getMetadata()) + .containsEntry(DecisionModelRoutingExecutor.CHOICE_KEY, "big"); + } + + /** Nothing to decide on: no text and no prompt arguments abstains without a provider call. */ + @Test + void emptyRequestAbstainsWithoutCallingTheProvider() throws Exception { + BaseDecisionModelSetup setup = mock(BaseDecisionModelSetup.class); + RunnerContext runner = mock(RunnerContext.class); + when(runner.getResource(any(), any())) + .thenAnswer( + invocation -> { + if (invocation.getArgument(1) == ResourceType.DECISION_MODEL) { + return setup; + } + throw new IllegalArgumentException("no chat model"); + }); + RoutingDecision decision = + new DecisionModelRoutingExecutor() + .route( + Strategies.decision("dm"), + ctx( + List.of( + ChatMessage.user( + List.of( + ImageBlock.fromBase64( + "image/png", "aGk=")))), + Map.of()), + runner); + assertThat(decision.isAbstain()).isTrue(); + assertThat(decision.getReason()).contains("no text"); + org.mockito.Mockito.verifyNoInteractions(setup); + } + + /** + * A request that cannot be rendered (here: a tool call whose input Jackson cannot serialize) is + * an ordinary routing failure for this request, not a runner failure. + */ + @Test + void unrenderableRequestIsARoutingFailure() throws Exception { + RunnerContext runner = mock(RunnerContext.class); + when(runner.getResource(any(), any())) + .thenAnswer( + invocation -> { + if (invocation.getArgument(1) == ResourceType.DECISION_MODEL) { + return mock(BaseDecisionModelSetup.class); + } + throw new IllegalArgumentException("no chat model"); + }); + RoutingContext context = + ctx( + List.of( + ChatMessage.assistant( + List.of( + new ToolCallBlock( + "call-1", + "lookup", + Map.of("query", new Object()))))), + Map.of()); + assertThatThrownBy( + () -> + new DecisionModelRoutingExecutor() + .route(Strategies.decision("dm"), context, runner)) + .isInstanceOf(ModelRoutingResolver.RoutingFailure.class) + .hasMessageContaining("could not build the decision request"); + } + + @Test + void probabilityAtTheMinimumIsAccepted() { + RoutingDecision decision = + DecisionModelRoutingExecutor.toDecision( + Strategies.decision("dm").withMinProbability(0.82), + ctx(List.of(), Map.of()), + choice("big", Map.of("small", 0.18, "big", 0.82), null), + new LinkedHashMap<>()); + assertThat(decision.isAbstain()).isFalse(); + assertThat(decision.getSelectedModel()).isEqualTo("big"); + } + + @Test + void refusalAbstainsWithTheProviderReason() { + DecisionResult refused = + new DecisionResult( + Map.of( + DecisionModelRoutingExecutor.QUESTION_NAME, + new DecisionResult.RefusalAnswer("not enough evidence")), + null, + null, + Map.of()); + RoutingDecision decision = + DecisionModelRoutingExecutor.toDecision( + Strategies.decision("dm"), + ctx(List.of(), Map.of()), + refused, + new LinkedHashMap<>()); + assertThat(decision.isAbstain()).isTrue(); + assertThat(decision.getReason()).isEqualTo("decision model refused: not enough evidence"); + assertThat(decision.getMetadata()) + .doesNotContainKeys( + DecisionModelRoutingExecutor.INPUT_TOKENS_KEY, + DecisionModelRoutingExecutor.OUTPUT_TOKENS_KEY); + } + + @Test + void nonChoiceAnswerAbstains() { + DecisionResult predicate = + new DecisionResult( + Map.of( + DecisionModelRoutingExecutor.QUESTION_NAME, + new DecisionResult.PredicateAnswer(0.4)), + null, + null, + Map.of()); + RoutingDecision decision = + DecisionModelRoutingExecutor.toDecision( + Strategies.decision("dm"), + ctx(List.of(), Map.of()), + predicate, + new LinkedHashMap<>()); + assertThat(decision.isAbstain()).isTrue(); + assertThat(decision.getReason()).isEqualTo("decision answer was not a choice"); + } + + /** Defensive: the provider contract forbids it, but a non-candidate choice never routes. */ + @Test + void nonCandidateChoiceAbstains() { + RoutingDecision decision = + DecisionModelRoutingExecutor.toDecision( + Strategies.decision("dm"), + ctx(List.of(), Map.of()), + choice("gpt-attacker", Map.of("gpt-attacker", 1.0), null), + new LinkedHashMap<>()); + assertThat(decision.isAbstain()).isTrue(); + assertThat(decision.getReason()).contains("gpt-attacker").contains("not a candidate"); + assertThat(decision.getMetadata()) + .containsEntry(DecisionModelRoutingExecutor.CHOICE_KEY, "gpt-attacker"); + } + + @Test + void declarationAccessorsReadTheArguments() { + RoutingStrategy strategy = Strategies.decision("dm").withMinProbability(0.6); + assertThat(DecisionModelRoutingExecutor.decisionModel(strategy)).isEqualTo("dm"); + assertThat(DecisionModelRoutingExecutor.minProbability(strategy)).isEqualTo(0.6); + assertThat(DecisionModelRoutingExecutor.minProbability(Strategies.decision("dm"))).isNull(); + assertThat(DecisionModelRoutingExecutor.instructions(Strategies.decision("dm"))) + .isEqualTo(DecisionModelRoutingExecutor.DEFAULT_INSTRUCTIONS); + } +} diff --git a/python/flink_agents/plan/tests/test_agent_plan_cross_language.py b/python/flink_agents/plan/tests/test_agent_plan_cross_language.py index 50418b0b58..06671fb911 100644 --- a/python/flink_agents/plan/tests/test_agent_plan_cross_language.py +++ b/python/flink_agents/plan/tests/test_agent_plan_cross_language.py @@ -497,6 +497,31 @@ def test_python_can_deserialize_java_plan_with_judge_router() -> None: assert "strategy_clazz" not in args # identity is the tag, not a Java class +def test_python_can_deserialize_java_plan_with_decision_router() -> None: + """A Java-produced plan carrying the decision-model router strategy must deserialize + on the Python side: the strategy travels as the language-neutral tag + ``decision_model`` with its arguments, and the decision model itself is an ordinary + DECISION_MODEL resource provider. + """ + snapshot = _SNAPSHOT_DIR / "java" / "agent_plan_with_decision_router.json" + assert snapshot.exists(), ( + f"Java decision-router plan snapshot missing from {snapshot}; regenerate on the" + " Java side with -Dregenerate.snapshots=true and commit it." + ) + restored = AgentPlan.model_validate_json(snapshot.read_text()) + assert ResourceType.MODEL_ROUTER in restored.resource_providers + assert "dm" in restored.resource_providers[ResourceType.DECISION_MODEL] + assert "conn" in restored.resource_providers[ResourceType.DECISION_MODEL_CONNECTION] + args = restored.resource_providers[ResourceType.MODEL_ROUTER][ + "router" + ].descriptor.arguments + assert args["strategy_type"] == "decision_model" + assert args["strategy_args"]["decision_model"] == "dm" + assert args["strategy_args"]["min_probability"] == 0.6 + assert args["candidate_descriptions"]["big"] == "code and sql" + assert "strategy_clazz" not in args + + def test_python_can_deserialize_plan_with_java_model_router() -> None: """A Java agent may declare a MODEL_ROUTER resource (Java-side in-chat routing). diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/context/DecisionRoutingRecoveryTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/context/DecisionRoutingRecoveryTest.java new file mode 100644 index 0000000000..52e4553555 --- /dev/null +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/context/DecisionRoutingRecoveryTest.java @@ -0,0 +1,300 @@ +/* + * 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.runtime.context; + +import org.apache.flink.agents.api.Event; +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.ChatResult; +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.ModelRouter; +import org.apache.flink.agents.api.chat.model.routing.Strategies; +import org.apache.flink.agents.api.decision.BaseDecisionModelSetup; +import org.apache.flink.agents.api.decision.DecisionResult; +import org.apache.flink.agents.api.event.ChatRequestEvent; +import org.apache.flink.agents.api.event.ChatResponseEvent; +import org.apache.flink.agents.api.event.ModelRoutingEvent; +import org.apache.flink.agents.api.resource.ResourceDescriptor; +import org.apache.flink.agents.api.resource.ResourceType; +import org.apache.flink.agents.plan.AgentPlan; +import org.apache.flink.agents.plan.actions.ChatModelAction; +import org.apache.flink.agents.runtime.ResourceCache; +import org.apache.flink.agents.runtime.actionstate.ActionState; +import org.apache.flink.agents.runtime.actionstate.ActionStateSerde; +import org.apache.flink.agents.runtime.async.ContinuationActionExecutor; +import org.apache.flink.agents.runtime.async.ContinuationContext; +import org.apache.flink.agents.runtime.memory.CachedMemoryStore; +import org.apache.flink.agents.runtime.memory.ForTestMemoryMapState; +import org.apache.flink.agents.runtime.metrics.FlinkAgentsMetricGroupImpl; +import org.apache.flink.runtime.metrics.groups.UnregisteredMetricGroups; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; + +import java.util.List; +import java.util.Map; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Function; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.junit.jupiter.api.Assumptions.assumeTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyMap; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +/** + * Verifies decision-model routing against the runtime's durable journal and recovery serializer: + * the stored {@code decide:} outcomes (a failed attempt and the choice) survive + * serialization and are replayed positionally when the worker dies before the {@code + * route:} record exists. + */ +class DecisionRoutingRecoveryTest { + private static final String FAILURE = "simulated worker failure after persistence"; + + private final AtomicReference persisted = new AtomicReference<>(); + private final AtomicInteger decisionCalls = new AtomicInteger(); + private final AtomicInteger chatCalls = new AtomicInteger(); + private ContinuationActionExecutor executor; + private final ChatRequestEvent request = + new ChatRequestEvent( + "router", List.of(new ChatMessage(MessageRole.USER, "write some sql"))); + private final DecisionResult choice = + new DecisionResult( + Map.of( + "model", + new DecisionResult.ChoiceAnswer( + "big", Map.of("small", 0.2, "big", 0.8), 0.7)), + "fake-model", + new DecisionResult.TokenUsage(10L, 2L, 12L), + Map.of()); + + @AfterEach + void closeExecutor() { + // A routed request must not leak an interrupt onto the calling thread. + assertThat(Thread.interrupted()).as("interrupt leaked by the action").isFalse(); + if (executor != null) { + executor.close(); + } + } + + /** + * The worker dies right after the second decision attempt is journaled (first attempt failed, + * retry succeeded) and before the route record. Recovery replays both stored outcomes without + * calling the provider, routes on the stored choice and answers with the selected model. + */ + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void restoreBetweenTheDecisionCallAndTheRouteRecordReplaysTheStoredChoice(boolean async) + throws Exception { + RunnerContextImpl first = context(new ActionState(request), async, 2); + assertThatThrownBy(() -> runAction(request, first)).hasMessageContaining(FAILURE); + assertThat(decisionCalls).hasValue(2); + assertThat(chatCalls).hasValue(0); + assertThat(first.getPendingEvents()).isEmpty(); + + ActionState restored = ActionStateSerde.deserialize(persisted.get()); + assertThat(restored.getCallResultCount()).isEqualTo(2); + RunnerContextImpl replay = context(restored, async, Integer.MAX_VALUE); + runAction(restored.getTaskEvent(), replay); + + assertThat(decisionCalls).hasValue(2); + assertThat(chatCalls).hasValue(1); + assertRouted(replay, 1); + } + + /** A completed routed request replays end to end without a provider or chat call. */ + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void replaysACompletedRoutedRequestWithoutAnotherCall(boolean async) throws Exception { + RunnerContextImpl first = context(new ActionState(request), async, Integer.MAX_VALUE); + runAction(request, first); + assertThat(decisionCalls).hasValue(2); + assertThat(chatCalls).hasValue(1); + assertRouted(first, 1); + + ActionState restored = ActionStateSerde.deserialize(persisted.get()); + assertThat(restored.getCallResultCount()).isEqualTo(4); + RunnerContextImpl replay = context(restored, async, Integer.MAX_VALUE); + runAction(restored.getTaskEvent(), replay); + + assertThat(decisionCalls).hasValue(2); + assertThat(chatCalls).hasValue(1); + assertRouted(replay, 1); + } + + private void assertRouted(RunnerContextImpl context, int retryCount) { + ModelRoutingEvent routing = + event(context, ModelRoutingEvent.EVENT_TYPE, ModelRoutingEvent::fromEvent); + assertThat(routing.getRequestId()).isEqualTo(request.getId()); + assertThat(routing.getSelectedModel()).isEqualTo("big"); + assertThat(routing.getDecisionSource()).isEqualTo(ModelRoutingEvent.SOURCE_DECISION_MODEL); + assertThat(routing.getScore()).isEqualTo(0.8); + assertThat(routing.getMetadata()) + .containsEntry("decision_probabilities", Map.of("small", 0.2, "big", 0.8)) + .containsEntry("provider_confidence", 0.7) + // Same types fresh and after the JSON round trip of the route record. + .containsEntry("decision_input_tokens", 10) + .containsEntry("decision_output_tokens", 2); + ChatResponseEvent response = + event(context, ChatResponseEvent.EVENT_TYPE, ChatResponseEvent::fromEvent); + assertThat(response.getRequestId()).isEqualTo(request.getId()); + assertThat(response.isFailed()).isFalse(); + assertThat(response.getResponse().getText()).isEqualTo("answer"); + assertThat(response.getRetryCount()).isEqualTo(retryCount); + } + + private static T event(RunnerContextImpl context, String type, Function read) { + return context.getPendingEvents().stream() + .filter(e -> type.equals(e.getType())) + .map(read) + .findFirst() + .orElseThrow(() -> new AssertionError("no event of type " + type)); + } + + private BaseDecisionModelSetup decisionModel() { + BaseDecisionModelSetup model = mock(BaseDecisionModelSetup.class); + when(model.supportsAsyncExecution()).thenReturn(true); + when(model.decide(any(), anyMap())) + .thenAnswer( + invocation -> { + if (decisionCalls.incrementAndGet() == 1) { + throw new IllegalArgumentException("provider unavailable"); + } + return choice; + }); + return model; + } + + private BaseChatModelSetup chatModel() { + return new BaseChatModelSetup(new ResourceDescriptor("fake", Map.of()), null) { + @Override + public Map getParameters() { + return Map.of(); + } + + @Override + public boolean willApplyNativeStructuredOutput(Object outputSchema) { + return false; + } + + @Override + public ChatResult chat( + List messages, + Map promptArgs, + Map modelParams) { + chatCalls.incrementAndGet(); + return new ChatResult(new ChatMessage(MessageRole.ASSISTANT, "answer")); + } + }; + } + + private RunnerContextImpl context(ActionState state, boolean async, int failAfterCallResults) + throws Exception { + AgentPlan plan = new AgentPlan(Map.of()); + plan.getConfig().set(AgentExecutionOptions.DECISION_ASYNC, async); + plan.getConfig().set(AgentExecutionOptions.MAX_RETRIES, 1); + plan.getConfig().set(AgentExecutionOptions.RETRY_WAIT_INTERVAL, 0); + ModelRouter router = + new ModelRouter( + ModelRouter.of("small", "big") + .describe("small", "short factual questions") + .describe("big", "code and sql") + .strategy(Strategies.decision("dm")) + .defaultModel("small") + .build(), + null); + // Build the stubs first: stubbing one mock while another stubbing is open is a misuse. + BaseDecisionModelSetup decisionModel = decisionModel(); + BaseChatModelSetup small = chatModel(); + BaseChatModelSetup big = chatModel(); + ResourceCache resources = mock(ResourceCache.class); + when(resources.hasResource("router", ResourceType.MODEL_ROUTER)).thenReturn(true); + when(resources.getResource("router", ResourceType.MODEL_ROUTER)).thenReturn(router); + when(resources.getResource("dm", ResourceType.DECISION_MODEL)).thenReturn(decisionModel); + when(resources.getResource("small", ResourceType.CHAT_MODEL)).thenReturn(small); + when(resources.getResource("big", ResourceType.CHAT_MODEL)).thenReturn(big); + FlinkAgentsMetricGroupImpl metrics = + new FlinkAgentsMetricGroupImpl( + UnregisteredMetricGroups.createUnregisteredOperatorMetricGroup()); + RunnerContextImpl context; + if (async) { + assumeTrue(Runtime.version().feature() >= 21, "async recovery needs JDK 21+"); + assertThat(ContinuationActionExecutor.isContinuationSupported()).isTrue(); + if (executor == null) { + executor = new ContinuationActionExecutor(1); + } + JavaRunnerContextImpl asyncContext = + new JavaRunnerContextImpl( + metrics, () -> {}, plan, resources, "test-job", executor); + asyncContext.setContinuationContext(new ContinuationContext()); + context = asyncContext; + } else { + context = new RunnerContextImpl(metrics, () -> {}, plan, resources, "test-job"); + } + context.actionName = "chat_model_action"; + // The chat action reads and writes sensory memory; back it with in-memory stores. + context.memoryContext = + new RunnerContextImpl.MemoryContext( + new CachedMemoryStore(new ForTestMemoryMapState<>()), + new CachedMemoryStore(new ForTestMemoryMapState<>())); + context.durableExecutionContext = + new RunnerContextImpl.DurableExecutionContext( + "test-key", + 1L, + ChatModelAction.getChatModelAction(), + state.getTaskEvent(), + state, + (key, sequence, action, event, current) -> { + persisted.set(ActionStateSerde.serialize(current)); + if (current.getCallResultCount() == failAfterCallResults) { + throw new IllegalStateException(FAILURE); + } + }); + return context; + } + + private void runAction(Event event, RunnerContextImpl context) throws Exception { + if (!(context instanceof JavaRunnerContextImpl)) { + ChatModelAction.processChatRequestOrToolResponse(event, context); + return; + } + JavaRunnerContextImpl asyncContext = (JavaRunnerContextImpl) context; + AtomicReference failure = new AtomicReference<>(); + Runnable action = + () -> { + try { + ChatModelAction.processChatRequestOrToolResponse(event, context); + } catch (Exception error) { + failure.set(error); + } + }; + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(5); + while (!executor.executeAction(asyncContext.getContinuationContext(), action)) { + assertThat(System.nanoTime()).isLessThan(deadline); + Thread.yield(); + } + if (failure.get() != null) { + throw failure.get(); + } + } +} From 283190948c9a9726e9d2627b0b19acf9756c8eee Mon Sep 17 00:00:00 2001 From: purshotam shah Date: Mon, 2 Nov 2026 08:50:44 -0800 Subject: [PATCH 2/5] [plan][java] Extract the shared routing input helpers into RoutingInput Address a review comment on the decision-model routing strategy: the decision executor used the judge executor's package-private methods for message preparation, context budgeting and content projection. Move those helpers into RoutingInput, a package-private class both strategies use, so neither depends on the other's internals. Behaviour is unchanged and stays pinned by LlmJudgeRoutingExecutorTest. Generated-by: Claude Code 2.1.309 (Claude Fable 5.1) --- .../routing/DecisionModelRoutingExecutor.java | 21 +- .../plan/routing/LlmJudgeRoutingExecutor.java | 214 +-------------- .../agents/plan/routing/RoutingInput.java | 244 ++++++++++++++++++ 3 files changed, 261 insertions(+), 218 deletions(-) create mode 100644 plan/src/main/java/org/apache/flink/agents/plan/routing/RoutingInput.java diff --git a/plan/src/main/java/org/apache/flink/agents/plan/routing/DecisionModelRoutingExecutor.java b/plan/src/main/java/org/apache/flink/agents/plan/routing/DecisionModelRoutingExecutor.java index 4884346348..bf9d7a23bf 100644 --- a/plan/src/main/java/org/apache/flink/agents/plan/routing/DecisionModelRoutingExecutor.java +++ b/plan/src/main/java/org/apache/flink/agents/plan/routing/DecisionModelRoutingExecutor.java @@ -199,9 +199,8 @@ static DecisionRequest buildRequest( RunnerContext ctx, Map metadata) throws Exception { - List effective = LlmJudgeRoutingExecutor.effectiveJudgeMessages(context, ctx); - Set pinned = - LlmJudgeRoutingExecutor.pinnedRenderedIndices(context.getMessages(), effective); + List effective = RoutingInput.effectiveMessages(context, ctx); + Set pinned = RoutingInput.pinnedRenderedIndices(context.getMessages(), effective); boolean[] truncated = new boolean[1]; String input = renderInput(strategy, context, effective, pinned, truncated, metadata); if (truncated[0]) { @@ -244,9 +243,9 @@ static String renderInput( boolean[] truncatedOut, Map metadata) { List selected = - LlmJudgeRoutingExecutor.selectWithinBudget( + RoutingInput.selectWithinBudget( effectiveMessages, - LlmJudgeRoutingExecutor.maxContextChars(strategy), + RoutingInput.maxContextChars(strategy), pinnedIndices, truncatedOut); // Whether the request carries text is judged on the complete effective list, as the judge @@ -259,7 +258,7 @@ static String renderInput( if (message.getRole() != MessageRole.USER && !pinnedIndices.contains(i)) { continue; } - for (DataContentBlock block : LlmJudgeRoutingExecutor.judgeContent(message)) { + for (DataContentBlock block : RoutingInput.content(message)) { if (block instanceof TextBlock && !((TextBlock) block).getText().isEmpty()) { requestHasText = true; } @@ -268,7 +267,7 @@ static String renderInput( StringBuilder text = new StringBuilder(); int mediaDropped = 0; for (ChatMessage message : selected) { - List content = LlmJudgeRoutingExecutor.judgeContent(message); + List content = RoutingInput.content(message); if (content.isEmpty()) { continue; } @@ -377,10 +376,6 @@ static RoutingDecision toDecision( return builder.build(); } - /** - * Persisted as a real abstain: replay resolves to the router's current default, and the - * provider's answer stays in the metadata for the routing event. - */ /** * Appends {@code piece} on a new line, without doubling a line break that either side already * carries (tool calls and results are rendered with their own framing newlines). @@ -405,6 +400,10 @@ private static Object jsonNumber(long value) { : (Object) value; } + /** + * Persisted as a real abstain: replay resolves to the router's current default, and the + * provider's answer stays in the metadata for the routing event. + */ private static RoutingDecision abstain(String reason, Map metadata) { return new RoutingDecision(null, true, reason, null, new HashMap<>(metadata), null); } diff --git a/plan/src/main/java/org/apache/flink/agents/plan/routing/LlmJudgeRoutingExecutor.java b/plan/src/main/java/org/apache/flink/agents/plan/routing/LlmJudgeRoutingExecutor.java index 1e18a7e8b5..8200be52c2 100644 --- a/plan/src/main/java/org/apache/flink/agents/plan/routing/LlmJudgeRoutingExecutor.java +++ b/plan/src/main/java/org/apache/flink/agents/plan/routing/LlmJudgeRoutingExecutor.java @@ -17,25 +17,17 @@ */ package org.apache.flink.agents.plan.routing; -import com.fasterxml.jackson.core.JsonProcessingException; -import com.fasterxml.jackson.databind.ObjectMapper; import org.apache.flink.agents.api.chat.messages.ChatMessage; import org.apache.flink.agents.api.chat.messages.ChatResult; -import org.apache.flink.agents.api.chat.messages.ContentBlock; import org.apache.flink.agents.api.chat.messages.DataContentBlock; import org.apache.flink.agents.api.chat.messages.MessageRole; import org.apache.flink.agents.api.chat.messages.TextBlock; -import org.apache.flink.agents.api.chat.messages.ToolCallBlock; -import org.apache.flink.agents.api.chat.messages.ToolResultBlock; -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 org.apache.flink.agents.plan.actions.ChatModelAction; import org.apache.flink.agents.plan.actions.ChatModelInvoker; import org.apache.flink.agents.plan.utils.CancellationUtils; @@ -74,8 +66,6 @@ */ final class LlmJudgeRoutingExecutor implements RoutingExecutor { - private static final ObjectMapper MAPPER = new ObjectMapper(); - /** Matches {@code "model": ""} in the judge's JSON verdict. */ private static final Pattern VERDICT_JSON = Pattern.compile("\"model\"\\s*:\\s*\"([^\"]+)\""); @@ -106,13 +96,13 @@ public RoutingDecision route( String abstainReason = null; try { boolean[] truncated = new boolean[1]; - List effective = effectiveJudgeMessages(context, ctx); + List effective = RoutingInput.effectiveMessages(context, ctx); List judgeInput = buildJudgeMessages( strategy, context, effective, - pinnedRenderedIndices(context.getMessages(), effective), + RoutingInput.pinnedRenderedIndices(context.getMessages(), effective), truncated); if (truncated[0]) { judgeMetadata.put(CONTEXT_TRUNCATED_KEY, true); @@ -188,58 +178,6 @@ private static List candidateNames(RoutingContext context) { 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. - * Cancellation and interruption are not failures of the anchor: they propagate. - */ - static List effectiveJudgeMessages(RoutingContext context, RunnerContext ctx) - throws Exception { - List 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) { - CancellationUtils.rethrowIfCancellation(unresolvable); - // An unresolvable candidate surfaces on the real chat path with its normal policy. - return messages; - } - } - - /** - * Indices of effective messages that were generated 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. - */ - static Set pinnedRenderedIndices( - List original, List effective) { - if (effective == original) { - return Set.of(); - } - Set originals = - java.util.Collections.newSetFromMap(new java.util.IdentityHashMap<>()); - originals.addAll(original); - Set 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); } @@ -248,18 +186,13 @@ private static String promptTemplate(RoutingStrategy strategy) { return (String) strategy.getArguments().get(RoutingStrategy.ARG_PROMPT_TEMPLATE); } - 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. * *

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 {@link #effectiveJudgeMessages}). With the opt-in {@code + * rendered when one exists (see {@link RoutingInput#effectiveMessages}). With the opt-in {@code * max_context_chars} cap, the newest message and the SYSTEM message are always kept, remaining * messages fill newest-first within the budget, and {@code truncatedOut[0]} is set so the * decision metadata records that the judge saw a trimmed view. Media blocks are preserved in @@ -296,10 +229,10 @@ static List buildJudgeMessages( + " only task is to pick the model."; } List conversation = - renderConversation( - selectWithinBudget( + RoutingInput.renderConversation( + RoutingInput.selectWithinBudget( effectiveMessages, - maxContextChars(strategy), + RoutingInput.maxContextChars(strategy), pinnedIndices, truncatedOut)); // Fall back to prompt args when no rendered/user text or media is present @@ -328,146 +261,13 @@ private static boolean carriesRequestContent( for (int i = 0; i < messages.size(); i++) { ChatMessage message = messages.get(i); if ((message.getRole() == MessageRole.USER || pinnedIndices.contains(i)) - && !judgeContent(message).isEmpty()) { + && !RoutingInput.content(message).isEmpty()) { return true; } } return false; } - /** - * Applies the opt-in context budget: SYSTEM messages and the newest message are pinned, the - * remaining messages fill newest-first, and dropped messages set {@code truncatedOut[0]}. - * Without a cap this returns the input unchanged. - */ - static List selectWithinBudget( - List messages, - int maxChars, - Set pinnedIndices, - boolean[] truncatedOut) { - if (messages == null) { - return List.of(); - } - if (maxChars == Integer.MAX_VALUE) { - return messages; - } - long total = 0; - for (ChatMessage message : messages) { - total += length(message); - } - if (total <= maxChars) { - return messages; - } - Set kept = new LinkedHashSet<>(); - long budget = maxChars; - // Pinned first — even if they alone exceed the budget: SYSTEM messages (the task - // framing), messages generated by the anchor's request shaping (the rendered template - // carries the task definition), and the newest message (the request being routed). - for (int i = 0; i < messages.size(); i++) { - if (messages.get(i).getRole() == MessageRole.SYSTEM || pinnedIndices.contains(i)) { - kept.add(i); - budget -= length(messages.get(i)); - } - } - int newest = messages.size() - 1; - if (!kept.contains(newest)) { - kept.add(newest); - budget -= length(messages.get(newest)); - } - // Then newest-first within what remains. - for (int i = newest - 1; i >= 0; i--) { - if (kept.contains(i)) { - continue; - } - long len = length(messages.get(i)); - if (len <= budget) { - kept.add(i); - budget -= len; - } - } - truncatedOut[0] = kept.size() < messages.size(); - List selected = new ArrayList<>(); - for (int i = 0; i < messages.size(); i++) { - if (kept.contains(i)) { - selected.add(messages.get(i)); - } - } - return selected; - } - - /** Counts rendered content text; media stays attached to its message, without a byte budget. */ - private static long length(ChatMessage message) { - return judgeContent(message).stream() - .filter(block -> block instanceof TextBlock) - .mapToLong(block -> ((TextBlock) block).getText().length()) - .sum(); - } - - /** Frames conversation roles as data while preserving the order and payload of media blocks. */ - private static List renderConversation(List messages) { - List rendered = new ArrayList<>(); - for (ChatMessage message : messages) { - List content = judgeContent(message); - if (content.isEmpty()) { - continue; - } - String prefix = rendered.isEmpty() ? "" : "\n"; - rendered.add(new TextBlock(prefix + message.getRole().name() + ": ")); - rendered.addAll(content); - } - return rendered; - } - - /** - * Projects one message into judge input. Tool interactions are conversation data, not - * executable calls to the judge. Reasoning and provider continuation metadata are not - * forwarded. - */ - static List judgeContent(ChatMessage message) { - List content = new ArrayList<>(); - for (ContentBlock block : message.getBlocks()) { - if (block instanceof TextBlock) { - if (!((TextBlock) block).getText().isEmpty()) { - content.add((TextBlock) block); - } - } else if (block instanceof DataContentBlock) { - content.add((DataContentBlock) block); - } else if (block instanceof ToolCallBlock) { - ToolCallBlock call = (ToolCallBlock) block; - content.add( - new TextBlock( - "\nTool call " - + toolDescription( - Map.of( - "call_id", call.getCallId(), - "name", call.getName(), - "input", call.getInput())) - + "\n")); - } else if (block instanceof ToolResultBlock) { - ToolResultBlock result = (ToolResultBlock) block; - content.add( - new TextBlock( - "Tool result " - + toolDescription( - Map.of( - "call_id", result.getCallId(), - "is_error", result.isError())) - + ":\n")); - content.addAll(result.getBlocks()); - } - } - return content; - } - - private static String toolDescription(Map fields) { - try { - return MAPPER.writeValueAsString(fields); - } catch (JsonProcessingException e) { - throw new IllegalArgumentException( - "Cannot render tool interaction for routing judge", e); - } - } - /** * Extracts a candidate name from the judge's reply. Accepts the documented JSON verdict or a * reply that is exactly a candidate name; anything else — including a well-formed verdict diff --git a/plan/src/main/java/org/apache/flink/agents/plan/routing/RoutingInput.java b/plan/src/main/java/org/apache/flink/agents/plan/routing/RoutingInput.java new file mode 100644 index 0000000000..c4c6b95787 --- /dev/null +++ b/plan/src/main/java/org/apache/flink/agents/plan/routing/RoutingInput.java @@ -0,0 +1,244 @@ +/* + * 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.routing; + +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.ObjectMapper; +import org.apache.flink.agents.api.chat.messages.ChatMessage; +import org.apache.flink.agents.api.chat.messages.ContentBlock; +import org.apache.flink.agents.api.chat.messages.DataContentBlock; +import org.apache.flink.agents.api.chat.messages.MessageRole; +import org.apache.flink.agents.api.chat.messages.TextBlock; +import org.apache.flink.agents.api.chat.messages.ToolCallBlock; +import org.apache.flink.agents.api.chat.messages.ToolResultBlock; +import org.apache.flink.agents.api.chat.model.BaseChatModelSetup; +import org.apache.flink.agents.api.chat.model.routing.RoutingContext; +import org.apache.flink.agents.api.chat.model.routing.RoutingStrategy; +import org.apache.flink.agents.api.context.RunnerContext; +import org.apache.flink.agents.api.resource.ResourceType; +import org.apache.flink.agents.plan.utils.CancellationUtils; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.IdentityHashMap; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; + +/** + * The request view a model-backed routing strategy works from, shared by the LLM judge and the + * decision-model strategy so neither depends on the other's internals: the effective message list + * (the anchor candidate's bound prompt rendered, as the selected model will receive it), the opt-in + * context budget, and the projection of a message into text-and-media content with tool + * interactions framed as data. + */ +final class RoutingInput { + + private static final ObjectMapper MAPPER = new ObjectMapper(); + + private RoutingInput() {} + + /** + * A routing strategy judges what the selected model will actually receive. When the target + * setup binds a {@code 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. + * Cancellation and interruption are not failures of the anchor: they propagate. + */ + static List effectiveMessages(RoutingContext context, RunnerContext ctx) + throws Exception { + List 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 + // strategy's view cannot drift from what the selected model receives. Candidates that + // bind DIFFERENT prompts see their own rendering only at answer time; the anchor (the + // default candidate, where abstains resolve) is a documented approximation. + return setup.prepareRequestMessages(messages, context.getPromptArgs()); + } catch (Exception unresolvable) { + CancellationUtils.rethrowIfCancellation(unresolvable); + // An unresolvable candidate surfaces on the real chat path with its normal policy. + return messages; + } + } + + /** + * Indices of effective messages that were generated 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. + */ + static Set pinnedRenderedIndices( + List original, List effective) { + if (effective == original) { + return Set.of(); + } + Set originals = Collections.newSetFromMap(new IdentityHashMap<>()); + originals.addAll(original); + Set pinned = new LinkedHashSet<>(); + for (int i = 0; i < effective.size(); i++) { + if (!originals.contains(effective.get(i))) { + pinned.add(i); + } + } + return pinned; + } + + 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; + } + + /** + * Applies the opt-in context budget: SYSTEM messages and the newest message are pinned, the + * remaining messages fill newest-first, and dropped messages set {@code truncatedOut[0]}. + * Without a cap this returns the input unchanged. + */ + static List selectWithinBudget( + List messages, + int maxChars, + Set pinnedIndices, + boolean[] truncatedOut) { + if (messages == null) { + return List.of(); + } + if (maxChars == Integer.MAX_VALUE) { + return messages; + } + long total = 0; + for (ChatMessage message : messages) { + total += length(message); + } + if (total <= maxChars) { + return messages; + } + Set kept = new LinkedHashSet<>(); + long budget = maxChars; + // Pinned first — even if they alone exceed the budget: SYSTEM messages (the task + // framing), messages generated by the anchor's request shaping (the rendered template + // carries the task definition), and the newest message (the request being routed). + for (int i = 0; i < messages.size(); i++) { + if (messages.get(i).getRole() == MessageRole.SYSTEM || pinnedIndices.contains(i)) { + kept.add(i); + budget -= length(messages.get(i)); + } + } + int newest = messages.size() - 1; + if (!kept.contains(newest)) { + kept.add(newest); + budget -= length(messages.get(newest)); + } + // Then newest-first within what remains. + for (int i = newest - 1; i >= 0; i--) { + if (kept.contains(i)) { + continue; + } + long len = length(messages.get(i)); + if (len <= budget) { + kept.add(i); + budget -= len; + } + } + truncatedOut[0] = kept.size() < messages.size(); + List selected = new ArrayList<>(); + for (int i = 0; i < messages.size(); i++) { + if (kept.contains(i)) { + selected.add(messages.get(i)); + } + } + return selected; + } + + /** Counts rendered content text; media stays attached to its message, without a byte budget. */ + private static long length(ChatMessage message) { + return content(message).stream() + .filter(block -> block instanceof TextBlock) + .mapToLong(block -> ((TextBlock) block).getText().length()) + .sum(); + } + + /** Frames conversation roles as data while preserving the order and payload of media blocks. */ + static List renderConversation(List messages) { + List rendered = new ArrayList<>(); + for (ChatMessage message : messages) { + List content = content(message); + if (content.isEmpty()) { + continue; + } + String prefix = rendered.isEmpty() ? "" : "\n"; + rendered.add(new TextBlock(prefix + message.getRole().name() + ": ")); + rendered.addAll(content); + } + return rendered; + } + + /** + * Projects one message into strategy input (judge prompt or decision-model text). Tool + * interactions are conversation data, not executable calls to the strategy's model. Reasoning + * and provider continuation metadata are not forwarded. + */ + static List content(ChatMessage message) { + List content = new ArrayList<>(); + for (ContentBlock block : message.getBlocks()) { + if (block instanceof TextBlock) { + if (!((TextBlock) block).getText().isEmpty()) { + content.add((TextBlock) block); + } + } else if (block instanceof DataContentBlock) { + content.add((DataContentBlock) block); + } else if (block instanceof ToolCallBlock) { + ToolCallBlock call = (ToolCallBlock) block; + content.add( + new TextBlock( + "\nTool call " + + toolDescription( + Map.of( + "call_id", call.getCallId(), + "name", call.getName(), + "input", call.getInput())) + + "\n")); + } else if (block instanceof ToolResultBlock) { + ToolResultBlock result = (ToolResultBlock) block; + content.add( + new TextBlock( + "Tool result " + + toolDescription( + Map.of( + "call_id", result.getCallId(), + "is_error", result.isError())) + + ":\n")); + content.addAll(result.getBlocks()); + } + } + return content; + } + + private static String toolDescription(Map fields) { + try { + return MAPPER.writeValueAsString(fields); + } catch (JsonProcessingException e) { + throw new IllegalArgumentException("Cannot render tool interaction for routing", e); + } + } +} From f6cfdcf4a837c8c733aa5ae5ffb0d032e4442457 Mon Sep 17 00:00:00 2001 From: purshotam shah Date: Mon, 2 Nov 2026 09:39:15 -0800 Subject: [PATCH 3/5] [e2e][java] Add an end-to-end test for decision-model routing Address a review comment on the decision-model routing strategy: run the strategy through a real Flink job with a deterministic decision model and two mock candidate chat models, all Java resources registered on the agent, so no credentials or Python bridge are needed. One job sends four requests and checks a confident selection of each candidate, a choice below the configured minimum probability abstaining to the default model, and a provider that keeps failing until the retry budget is exhausted, which fails the request without candidate fallback. The test also checks the provider call count, including the retry, and the rendered decision input. Generated-by: Claude Code 2.1.309 (Claude Fable 5.1) --- .../test/DecisionModelRoutingAgent.java | 237 ++++++++++++++++++ .../test/DecisionModelRoutingE2ETest.java | 140 +++++++++++ 2 files changed, 377 insertions(+) create mode 100644 e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/DecisionModelRoutingAgent.java create mode 100644 e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/DecisionModelRoutingE2ETest.java diff --git a/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/DecisionModelRoutingAgent.java b/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/DecisionModelRoutingAgent.java new file mode 100644 index 0000000000..c6950c9c65 --- /dev/null +++ b/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/DecisionModelRoutingAgent.java @@ -0,0 +1,237 @@ +/* + * 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.integration.test; + +import org.apache.flink.agents.api.EventType; +import org.apache.flink.agents.api.InputEvent; +import org.apache.flink.agents.api.OutputEvent; +import org.apache.flink.agents.api.agents.Agent; +import org.apache.flink.agents.api.annotation.Action; +import org.apache.flink.agents.api.annotation.ChatModelConnection; +import org.apache.flink.agents.api.annotation.ChatModelSetup; +import org.apache.flink.agents.api.annotation.DecisionModelConnection; +import org.apache.flink.agents.api.annotation.DecisionModelSetup; +import org.apache.flink.agents.api.chat.messages.ChatMessage; +import org.apache.flink.agents.api.chat.messages.ChatResult; +import org.apache.flink.agents.api.chat.model.BaseChatModelConnection; +import org.apache.flink.agents.api.chat.model.BaseChatModelSetup; +import org.apache.flink.agents.api.chat.model.routing.ModelRouter; +import org.apache.flink.agents.api.chat.model.routing.Strategies; +import org.apache.flink.agents.api.context.RunnerContext; +import org.apache.flink.agents.api.decision.BaseDecisionModelConnection; +import org.apache.flink.agents.api.decision.BaseDecisionModelSetup; +import org.apache.flink.agents.api.decision.DecisionRequest; +import org.apache.flink.agents.api.decision.DecisionResult; +import org.apache.flink.agents.api.event.ChatRequestEvent; +import org.apache.flink.agents.api.event.ChatResponseEvent; +import org.apache.flink.agents.api.resource.ResourceContext; +import org.apache.flink.agents.api.resource.ResourceDescriptor; +import org.apache.flink.agents.api.resource.ResourceType; +import org.apache.flink.agents.api.tools.Tool; + +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.atomic.AtomicInteger; + +/** + * An agent whose chat requests go through a {@code MODEL_ROUTER} driven by a deterministic decision + * model, with two mock candidate chat models that echo their own name. No credentials and no Python + * bridge: everything the routed request touches is a Java resource registered on this agent, so the + * test exercises registration, initialization, routing, candidate invocation and the final response + * of a real job. + * + *

The decision connection scripts its answer from the request text: {@code SQL} picks {@code + * big} at 0.9, {@code unsure} picks {@code big} at 0.55 (below the router's 0.6 minimum, so the + * request abstains to the default model), {@code crash} throws (the retry budget is exhausted and + * the request fails), and anything else picks {@code small} at 0.95. + */ +public class DecisionModelRoutingAgent extends Agent { + + // Shared with the test through accessors; the job runs in this JVM. The test resets them + // before and after each run so no state leaks across tests in the same surefire fork. + private static final AtomicInteger DECISION_CALLS = new AtomicInteger(); + private static final List DECISION_INPUTS = new CopyOnWriteArrayList<>(); + + /** Provider calls across the job, including retried attempts. */ + public static int decisionCalls() { + return DECISION_CALLS.get(); + } + + /** The input text of every decision request, in call order. */ + public static List decisionInputs() { + return List.copyOf(DECISION_INPUTS); + } + + public static void reset() { + DECISION_CALLS.set(0); + DECISION_INPUTS.clear(); + } + + public DecisionModelRoutingAgent() { + addResource( + "router", + ResourceType.MODEL_ROUTER, + ModelRouter.of("small", "big") + .describe("small", "greetings and short factual questions") + .describe("big", "code, SQL and multi-step reasoning") + .strategy(Strategies.decision("dm").withMinProbability(0.6)) + .defaultModel("small") + .build()); + } + + /** + * A chat connection that answers with its tag and the last message, so the test sees who ran. + */ + public static class EchoChatConnection extends BaseChatModelConnection { + private final String tag; + + public EchoChatConnection(ResourceDescriptor descriptor, ResourceContext resourceContext) { + super(descriptor, resourceContext); + this.tag = descriptor.getArgument("tag"); + } + + @Override + public ChatResult chat( + List messages, List tools, Map modelParams) { + ChatMessage last = messages.get(messages.size() - 1); + return new ChatResult(ChatMessage.assistant(tag + ":" + last.getText())); + } + } + + /** Chat model setup bound to an {@link EchoChatConnection}. */ + public static class EchoChatModel extends BaseChatModelSetup { + public EchoChatModel(ResourceDescriptor descriptor, ResourceContext resourceContext) { + super(descriptor, resourceContext); + } + + @Override + public Map getParameters() { + return new HashMap<>(); + } + } + + @ChatModelConnection + public static ResourceDescriptor smallConnection() { + return ResourceDescriptor.Builder.newBuilder(EchoChatConnection.class.getName()) + .addInitialArgument("tag", "small") + .build(); + } + + @ChatModelConnection + public static ResourceDescriptor bigConnection() { + return ResourceDescriptor.Builder.newBuilder(EchoChatConnection.class.getName()) + .addInitialArgument("tag", "big") + .build(); + } + + @ChatModelSetup + public static ResourceDescriptor small() { + return ResourceDescriptor.Builder.newBuilder(EchoChatModel.class.getName()) + .addInitialArgument("connection", "smallConnection") + .addInitialArgument("model", "echo-small") + .build(); + } + + @ChatModelSetup + public static ResourceDescriptor big() { + return ResourceDescriptor.Builder.newBuilder(EchoChatModel.class.getName()) + .addInitialArgument("connection", "bigConnection") + .addInitialArgument("model", "echo-big") + .build(); + } + + /** The scripted decision provider (see the class comment). */ + public static class ScriptedDecisionConnection extends BaseDecisionModelConnection { + public ScriptedDecisionConnection( + ResourceDescriptor descriptor, ResourceContext resourceContext) { + super(descriptor, resourceContext); + } + + @Override + public DecisionResult decide( + DecisionRequest request, String model, Map requestParams) { + DECISION_CALLS.incrementAndGet(); + String input = String.valueOf(request.getInput()); + DECISION_INPUTS.add(input); + if (input.contains("crash")) { + throw new IllegalStateException("decision provider unavailable"); + } + String choice; + Map probabilities; + if (input.contains("SQL")) { + choice = "big"; + probabilities = Map.of("small", 0.1, "big", 0.9); + } else if (input.contains("unsure")) { + choice = "big"; + probabilities = Map.of("small", 0.45, "big", 0.55); + } else { + choice = "small"; + probabilities = Map.of("small", 0.95, "big", 0.05); + } + return new DecisionResult( + Map.of("model", new DecisionResult.ChoiceAnswer(choice, probabilities, null)), + model, + new DecisionResult.TokenUsage(5L, 1L, 6L), + Map.of()); + } + } + + @DecisionModelConnection + public static ResourceDescriptor decisionConnection() { + return ResourceDescriptor.Builder.newBuilder(ScriptedDecisionConnection.class.getName()) + .build(); + } + + @DecisionModelSetup + public static ResourceDescriptor dm() { + return ResourceDescriptor.Builder.newBuilder(BaseDecisionModelSetup.class.getName()) + .addInitialArgument("connection", "decisionConnection") + .addInitialArgument("model", "scripted") + .build(); + } + + @Action(EventType.InputEvent) + public static void processInput(InputEvent event, RunnerContext ctx) { + ctx.sendEvent( + new ChatRequestEvent( + "router", List.of(ChatMessage.user((String) event.getInput())))); + } + + @Action(EventType.ChatResponseEvent) + public static void processChatResponse(ChatResponseEvent event, RunnerContext ctx) { + if (event.isFailed()) { + ctx.sendEvent(new OutputEvent("FAILED: " + event.getError())); + return; + } + Map routing = (Map) event.getAttr("model_routing"); + Map metadata = (Map) routing.get("metadata"); + // decision_choice is recorded only when the strategy abstained on a choice. + Object choice = metadata == null ? null : metadata.get("decision_choice"); + ctx.sendEvent( + new OutputEvent( + String.format( + "final_model=%s decision_source=%s score=%s%s :: %s", + routing.get("final_model"), + routing.get("decision_source"), + routing.get("score"), + choice == null ? "" : " decision_choice=" + choice, + event.getResponse().getText()))); + } +} diff --git a/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/DecisionModelRoutingE2ETest.java b/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/DecisionModelRoutingE2ETest.java new file mode 100644 index 0000000000..c2d94e9c97 --- /dev/null +++ b/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/DecisionModelRoutingE2ETest.java @@ -0,0 +1,140 @@ +/* + * 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.integration.test; + +import org.apache.flink.agents.api.AgentsExecutionEnvironment; +import org.apache.flink.agents.api.agents.AgentExecutionOptions; +import org.apache.flink.api.java.functions.KeySelector; +import org.apache.flink.streaming.api.datastream.DataStream; +import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment; +import org.apache.flink.util.CloseableIterator; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; + +import java.util.ArrayList; +import java.util.List; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Decision-model routing through a running Flink job: the router, the decision model and both + * candidate chat models are Java resources registered on {@link DecisionModelRoutingAgent}, so the + * test covers resource registration and initialization, the routing decision, the candidate + * invocation and the final output, with no live credentials and no Python bridge. + */ +public class DecisionModelRoutingE2ETest { + + /** Routes by request text so every input is its own key. */ + public static class TextKeySelector implements KeySelector { + @Override + public String getKey(String value) { + return value; + } + } + + @BeforeEach + @AfterEach + void reset() { + DecisionModelRoutingAgent.reset(); + } + + /** + * One job, four requests: a confident choice of the cheap model, a confident choice of the + * expensive model, a choice below the configured minimum probability (abstains to the default + * model) and a provider that keeps failing (the retry budget is exhausted and the request + * fails, with no candidate fallback). + */ + @Test + @Timeout(120) + public void routesAbstainsAndFailsThroughARunningJob() throws Exception { + List outputs = + runJob( + agentsEnv -> { + agentsEnv.getConfig().set(AgentExecutionOptions.MAX_RETRIES, 1); + agentsEnv.getConfig().set(AgentExecutionOptions.RETRY_WAIT_INTERVAL, 0); + }, + "hello there", + "Write SQL to list active users", + "unsure which model fits", + "please crash"); + + assertThat(outputs).hasSize(4); + // Selection: the decision model's choice and probability drive the candidate. + assertThat(outputs) + .contains( + "final_model=small decision_source=decision_model score=0.95" + + " :: small:hello there", + "final_model=big decision_source=decision_model score=0.9" + + " :: big:Write SQL to list active users"); + // Abstention: 0.55 for "big" is below the 0.6 minimum, so the default model answers, the + // source records the default rather than a decision, no score is reported, and the + // provider's rejected choice stays in the metadata. + assertThat(outputs) + .contains( + "final_model=small decision_source=default score=null" + + " decision_choice=big :: small:unsure which model fits"); + // Exhausted retries: the request fails with the documented message naming the decision + // model and carrying the provider's error; no candidate answers on its behalf. + assertThat(outputs) + .anySatisfy( + output -> + assertThat(output) + .startsWith("FAILED: Decision model 'dm' failed") + .contains("decision provider unavailable") + .doesNotContain("small:please crash") + .doesNotContain("big:please crash")); + + // The provider saw three routed requests once each and the failing one twice (one retry). + assertThat(DecisionModelRoutingAgent.decisionCalls()).isEqualTo(5); + assertThat(DecisionModelRoutingAgent.decisionInputs()) + .filteredOn(input -> input.contains("please crash")) + .hasSize(2); + // The decision input is the rendered request, framed by role. + assertThat(DecisionModelRoutingAgent.decisionInputs()).contains("USER: hello there"); + } + + /** Runs a single-parallelism agent job and returns the collected outputs. */ + private static List runJob( + java.util.function.Consumer configure, String... inputs) + throws Exception { + StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment(); + env.setParallelism(1); + DataStream inputStream = env.fromData(inputs); + + AgentsExecutionEnvironment agentsEnv = + AgentsExecutionEnvironment.getExecutionEnvironment(env); + configure.accept(agentsEnv); + + DataStream outputStream = + agentsEnv + .fromDataStream(inputStream, new TextKeySelector()) + .apply(new DecisionModelRoutingAgent()) + .toDataStream(); + + List outputs = new ArrayList<>(); + try (CloseableIterator results = outputStream.collectAsync()) { + agentsEnv.execute(); + while (results.hasNext()) { + outputs.add(results.next().toString()); + } + } + return outputs; + } +} From c21651514b584f36482264fe3ab7686033f2d26e Mon Sep 17 00:00:00 2001 From: purshotam shah Date: Mon, 2 Nov 2026 09:39:15 -0800 Subject: [PATCH 4/5] [examples][docs] Use the Python OpenAI decision model in the routing example Address a review comment on the decision-model routing strategy: the example shipped its own Java HTTP adapter for the OpenAI Decisions API, which put provider-specific request and response handling into a routing example. Declare the built-in Python OpenAI decision model connection and setup through Python resource descriptors instead, as other cross-language examples do, and drop the adapter. Native Java support for the Decisions API, if wanted, belongs in a separate integration. The example now needs a Python environment with flink-agents on the TaskManagers, which its documentation states. Generated-by: Claude Code 2.1.309 (Claude Fable 5.1) --- .../development/calling-llms/model_routing.md | 2 +- .../openai/OpenAiDecisionRoutingExample.java | 243 +----------------- 2 files changed, 14 insertions(+), 231 deletions(-) diff --git a/docs/content/docs/development/calling-llms/model_routing.md b/docs/content/docs/development/calling-llms/model_routing.md index f0a03f29ed..22b3d0026f 100644 --- a/docs/content/docs/development/calling-llms/model_routing.md +++ b/docs/content/docs/development/calling-llms/model_routing.md @@ -418,6 +418,6 @@ for the persistence and replay guarantees. - {{< gh_link file="examples/src/main/java/org/apache/flink/agents/examples/ModelRoutingExample.java" name="`ModelRoutingExample`" >}}: rule-based routing between two Ollama models. - {{< gh_link file="examples/src/main/java/org/apache/flink/agents/examples/ModelRoutingJudgeExample.java" name="`ModelRoutingJudgeExample`" >}}: the same two models with an LLM judge and candidate descriptions. - {{< gh_link file="examples/src/main/java/org/apache/flink/agents/examples/openai/OpenAiModelRoutingExample.java" name="`OpenAiModelRoutingExample`" >}}: rule-based routing between `gpt-4o-mini` and `gpt-4o` over one OpenAI connection. -- {{< gh_link file="examples/src/main/java/org/apache/flink/agents/examples/openai/OpenAiDecisionRoutingExample.java" name="`OpenAiDecisionRoutingExample`" >}}: the same two OpenAI models routed by a decision model, with a native Java connection to the OpenAI Decisions API and an agent that reports the selected model, decision source and score. +- {{< gh_link file="examples/src/main/java/org/apache/flink/agents/examples/openai/OpenAiDecisionRoutingExample.java" name="`OpenAiDecisionRoutingExample`" >}}: the same two OpenAI models routed by a decision model (the built-in Python OpenAI decision model, used through the cross-language API), with an agent that reports the selected model, decision source and score. All of them register their resources on the execution environment. The first three use the {{< gh_link file="examples/src/main/java/org/apache/flink/agents/examples/agents/ModelRoutingAgent.java" name="`ModelRoutingAgent`" >}} to send requests and handle responses. diff --git a/examples/src/main/java/org/apache/flink/agents/examples/openai/OpenAiDecisionRoutingExample.java b/examples/src/main/java/org/apache/flink/agents/examples/openai/OpenAiDecisionRoutingExample.java index 36e6d25cfb..f3b2862036 100644 --- a/examples/src/main/java/org/apache/flink/agents/examples/openai/OpenAiDecisionRoutingExample.java +++ b/examples/src/main/java/org/apache/flink/agents/examples/openai/OpenAiDecisionRoutingExample.java @@ -17,10 +17,6 @@ */ package org.apache.flink.agents.examples.openai; -import com.fasterxml.jackson.databind.JsonNode; -import com.fasterxml.jackson.databind.ObjectMapper; -import com.fasterxml.jackson.databind.node.ArrayNode; -import com.fasterxml.jackson.databind.node.ObjectNode; import org.apache.flink.agents.api.AgentsExecutionEnvironment; import org.apache.flink.agents.api.EventType; import org.apache.flink.agents.api.InputEvent; @@ -32,27 +28,15 @@ import org.apache.flink.agents.api.chat.model.routing.ModelRouter; import org.apache.flink.agents.api.chat.model.routing.Strategies; import org.apache.flink.agents.api.context.RunnerContext; -import org.apache.flink.agents.api.decision.BaseDecisionModelConnection; -import org.apache.flink.agents.api.decision.BaseDecisionModelSetup; -import org.apache.flink.agents.api.decision.DecisionRequest; -import org.apache.flink.agents.api.decision.DecisionResult; import org.apache.flink.agents.api.event.ChatRequestEvent; import org.apache.flink.agents.api.event.ChatResponseEvent; -import org.apache.flink.agents.api.resource.ResourceContext; +import org.apache.flink.agents.api.resource.PythonResourceDescriptor; import org.apache.flink.agents.api.resource.ResourceDescriptor; import org.apache.flink.agents.api.resource.ResourceName; import org.apache.flink.agents.api.resource.ResourceType; import org.apache.flink.streaming.api.datastream.DataStream; import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment; -import java.io.IOException; -import java.net.URI; -import java.net.http.HttpClient; -import java.net.http.HttpRequest; -import java.net.http.HttpResponse; -import java.time.Duration; -import java.util.Iterator; -import java.util.LinkedHashMap; import java.util.List; import java.util.Map; @@ -62,10 +46,11 @@ * chosen candidate answers unless its probability is below the configured minimum, in which case * the default model does. Requires {@code OPENAI_API_KEY} in the environment. * - *

The decision model is a native Java provider: {@link OpenAiDecisionsConnection} implements - * {@link BaseDecisionModelConnection} over the Decisions endpoint, so the job needs no Python - * environment on the TaskManagers. The output lines carry the model that answered, the decision - * source and the routing score, read from the {@code model_routing} attribute of the response. + *

The decision model is the built-in Python OpenAI decision model provider, used from this Java + * job through the cross-language resource API, so the TaskManagers need a Python environment with + * {@code flink-agents} installed (see the cross-language resources documentation). The output lines + * carry the model that answered, the decision source and the routing score, read from the {@code + * model_routing} attribute of the response. */ public class OpenAiDecisionRoutingExample { @@ -101,18 +86,21 @@ public static void main(String[] args) throws Exception { .addResource("small", ResourceType.CHAT_MODEL, openAiModel("gpt-4o-mini")) .addResource("big", ResourceType.CHAT_MODEL, openAiModel("gpt-4o")); - // The decision model: a connection to the OpenAI Decisions API and a setup naming the - // decision model. Registered like any other decision model resource. + // The decision model: the Python OpenAI decision model provider (connection and setup), + // declared with Python resource descriptors and otherwise registered like any other + // decision model resource. agentsEnv.addResource( "openaiDecisions", ResourceType.DECISION_MODEL_CONNECTION, - ResourceDescriptor.Builder.newBuilder(OpenAiDecisionsConnection.class.getName()) + PythonResourceDescriptor.Builder.newBuilder( + ResourceName.DecisionModel.Python.OPENAI_CONNECTION) .addInitialArgument("api_key", apiKey) .build()); agentsEnv.addResource( "dm", ResourceType.DECISION_MODEL, - ResourceDescriptor.Builder.newBuilder(BaseDecisionModelSetup.class.getName()) + PythonResourceDescriptor.Builder.newBuilder( + ResourceName.DecisionModel.Python.OPENAI_SETUP) .addInitialArgument("connection", "openaiDecisions") .addInitialArgument("model", "gpt-6-luna") .build()); @@ -174,209 +162,4 @@ public static void processChatResponse(ChatResponseEvent event, RunnerContext ct : answer))); } } - - /** - * A native Java decision model connection over the OpenAI Decisions API ({@code POST - * /v1/decisions}). Supports choice and predicate questions, which is what routing and most - * gating use; score questions are rejected. Descriptor arguments: {@code api_key} (required) - * and {@code base_url} (default {@code https://api.openai.com/v1}). - */ - public static class OpenAiDecisionsConnection extends BaseDecisionModelConnection { - - private static final ObjectMapper MAPPER = new ObjectMapper(); - private static final java.util.Set RESERVED_FIELDS = - java.util.Set.of("model", "input", "questions"); - - private final String apiKey; - private final String baseUrl; - private transient volatile HttpClient client; - - public OpenAiDecisionsConnection( - ResourceDescriptor descriptor, ResourceContext resourceContext) { - super(descriptor, resourceContext); - Object key = descriptor.getArgument("api_key"); - if (!(key instanceof String) || ((String) key).isBlank()) { - throw new IllegalArgumentException( - "OpenAiDecisionsConnection requires a non-blank 'api_key'."); - } - this.apiKey = (String) key; - this.baseUrl = descriptor.getArgument("base_url", "https://api.openai.com/v1"); - } - - @Override - public DecisionResult decide( - DecisionRequest request, String model, Map requestParams) { - ObjectNode body = MAPPER.createObjectNode(); - for (Map.Entry param : requestParams.entrySet()) { - if (RESERVED_FIELDS.contains(param.getKey())) { - throw new IllegalArgumentException( - "Request parameter '" - + param.getKey() - + "' cannot override the model, evidence or questions."); - } - body.set(param.getKey(), MAPPER.valueToTree(param.getValue())); - } - body.put("model", model); - body.set("input", MAPPER.valueToTree(request.getInput())); - ArrayNode questions = body.putArray("questions"); - for (Map.Entry entry : - request.getQuestions().entrySet()) { - questions.add(toQuestion(entry.getKey(), entry.getValue())); - } - JsonNode response = post(body); - Map answers = new LinkedHashMap<>(); - for (JsonNode answer : response.path("answers")) { - String name = answer.path("name").asText(); - DecisionRequest.Question question = request.getQuestions().get(name); - if (question == null) { - throw new IllegalStateException( - "OpenAI Decisions API answered an unknown question '" + name + "'."); - } - answers.put(name, toAnswer(answer, question)); - } - return new DecisionResult( - answers, - response.path("model").isTextual() ? response.get("model").asText() : null, - toUsage(response.path("usage")), - Map.of()); - } - - private static ObjectNode toQuestion(String name, DecisionRequest.Question question) { - ObjectNode node = MAPPER.createObjectNode(); - node.put("name", name); - if (question instanceof DecisionRequest.ChoiceQuestion) { - node.put("type", "choice"); - node.set("instructions", MAPPER.valueToTree(question.getInstructions())); - ArrayNode choices = node.putArray("choices"); - for (Map.Entry option : - ((DecisionRequest.ChoiceQuestion) question).getCriteria().entrySet()) { - ObjectNode choice = choices.addObject(); - choice.put("value", option.getKey()); - if (option.getValue() != null) { - choice.set("description", MAPPER.valueToTree(option.getValue())); - } - } - return node; - } - if (question instanceof DecisionRequest.PredicateQuestion) { - node.put("type", "predicate"); - node.set("instructions", MAPPER.valueToTree(question.getInstructions())); - return node; - } - throw new IllegalArgumentException( - "OpenAiDecisionsConnection supports choice and predicate questions only."); - } - - private static DecisionResult.Answer toAnswer( - JsonNode answer, DecisionRequest.Question question) { - String type = answer.path("type").asText(); - if ("refusal".equals(type)) { - return new DecisionResult.RefusalAnswer( - answer.path("reason").isTextual() ? answer.get("reason").asText() : null); - } - if (question instanceof DecisionRequest.ChoiceQuestion) { - if (!answer.path("choice").isTextual()) { - throw new IllegalStateException( - "OpenAI Decisions API returned a choice answer without a 'choice'."); - } - Map probabilities = new LinkedHashMap<>(); - for (JsonNode entry : answer.path("probabilities")) { - if (!entry.path("value").isTextual() || !entry.path("probability").isNumber()) { - throw new IllegalStateException( - "OpenAI Decisions API returned a malformed probability entry: " - + entry); - } - if (probabilities.put( - entry.path("value").asText(), - entry.path("probability").asDouble()) - != null) { - throw new IllegalStateException( - "OpenAI Decisions API returned duplicate choice '" - + entry.path("value").asText() - + "'."); - } - } - if (probabilities.isEmpty()) { - throw new IllegalStateException( - "OpenAI Decisions API returned a choice answer without" - + " probabilities."); - } - return new DecisionResult.ChoiceAnswer( - answer.path("choice").asText(), - probabilities, - answer.path("confidence").isNumber() - ? answer.get("confidence").asDouble() - : null); - } - if (!"predicate".equals(type) || !answer.path("probability").isNumber()) { - throw new IllegalStateException( - "OpenAI Decisions API returned an unexpected answer of type '" - + type - + "' for a predicate question."); - } - return new DecisionResult.PredicateAnswer(answer.path("probability").asDouble()); - } - - private static DecisionResult.TokenUsage toUsage(JsonNode usage) { - if (usage.isMissingNode() || usage.isNull()) { - return null; - } - return new DecisionResult.TokenUsage( - longOrNull(usage, "input_tokens", "prompt_tokens"), - longOrNull(usage, "output_tokens", "completion_tokens"), - longOrNull(usage, "total_tokens")); - } - - private static Long longOrNull(JsonNode node, String... names) { - for (String name : names) { - if (node.path(name).isNumber()) { - return node.get(name).asLong(); - } - } - return null; - } - - private JsonNode post(ObjectNode body) { - if (client == null) { - client = HttpClient.newBuilder().connectTimeout(Duration.ofSeconds(30)).build(); - } - try { - HttpRequest request = - HttpRequest.newBuilder(URI.create(baseUrl + "/decisions")) - .timeout(Duration.ofSeconds(60)) - .header("Authorization", "Bearer " + apiKey) - .header("Content-Type", "application/json") - .POST( - HttpRequest.BodyPublishers.ofString( - MAPPER.writeValueAsString(body))) - .build(); - HttpResponse response = - client.send(request, HttpResponse.BodyHandlers.ofString()); - if (response.statusCode() / 100 != 2) { - throw new IllegalStateException( - "OpenAI Decisions API returned HTTP " - + response.statusCode() - + ": " - + errorMessage(response.body())); - } - return MAPPER.readTree(response.body()); - } catch (IOException e) { - throw new IllegalStateException("OpenAI Decisions API call failed", e); - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - throw new IllegalStateException("OpenAI Decisions API call interrupted", e); - } - } - - /** The provider's error message without echoing the whole body into logs. */ - private static String errorMessage(String body) { - try { - JsonNode error = MAPPER.readTree(body).path("error"); - Iterator fields = error.fieldNames(); - return fields.hasNext() ? error.path("message").asText(body) : body; - } catch (IOException e) { - return body; - } - } - } } From 598b068d4ac4628785b063e36ca88e3814c1143c Mon Sep 17 00:00:00 2001 From: purshotam shah Date: Tue, 3 Nov 2026 11:16:15 -0800 Subject: [PATCH 5/5] [plan][python][docs] Snapshot every decision-strategy argument and document a missing probability Address review comments on the decision-model routing strategy. The cross-language plan snapshot carried only the required strategy arguments; add a second router that sets instructions and max_context_chars as well, and have the Python test check each key's name and JSON type, so a later Python implementation of the strategy has the contract in front of it. Document what happens when the provider omits the chosen candidate from the probability map: with min_probability set the request abstains to the default model, since there is nothing to compare; without it the choice routes with no score. The second case gains a unit test next to the existing one, and the Javadoc of Strategies.decision and of the executor say the same. Generated-by: Claude Code 2.1.309 (Claude Fable 5.1) --- .../api/chat/model/routing/Strategies.java | 7 +++-- .../development/calling-llms/model_routing.md | 11 ++++--- .../java/agent_plan_with_decision_router.json | 26 +++++++++++++++++ .../routing/DecisionModelRoutingExecutor.java | 11 +++---- .../plan/AgentPlanCrossLanguageTest.java | 29 ++++++++++++++++--- .../DecisionModelRoutingExecutorTest.java | 17 +++++++++++ .../tests/test_agent_plan_cross_language.py | 17 ++++++++++- 7 files changed, 101 insertions(+), 17 deletions(-) diff --git a/api/src/main/java/org/apache/flink/agents/api/chat/model/routing/Strategies.java b/api/src/main/java/org/apache/flink/agents/api/chat/model/routing/Strategies.java index 6e2bbe8af9..ebe83ebe4d 100644 --- a/api/src/main/java/org/apache/flink/agents/api/chat/model/routing/Strategies.java +++ b/api/src/main/java/org/apache/flink/agents/api/chat/model/routing/Strategies.java @@ -96,9 +96,10 @@ public static RoutingStrategy llm(String judgeModel, String promptTemplate) { * names and whose criteria are the candidates' {@code describe(...)} descriptions. The engine * executes the decision call on its durable, metered, observable path. The returned choice * becomes the selected model and its reported probability the routing score. A refusal answer, - * or a choice whose probability is below {@link RoutingStrategy#withMinProbability(double)}, - * abstains to the router's default model; a decision call that exhausts its retries fails the - * request. + * a choice whose probability is below {@link RoutingStrategy#withMinProbability(double)}, or a + * choice with no reported probability while that minimum is set, abstains to the router's + * default model (without a minimum such a choice routes with no score); a decision call that + * exhausts its retries fails the request. * *

Every candidate must carry a description: the descriptions are all the decision model sees * of the candidates. The decision model receives the conversation text (tool calls and results diff --git a/docs/content/docs/development/calling-llms/model_routing.md b/docs/content/docs/development/calling-llms/model_routing.md index 22b3d0026f..b717e472e4 100644 --- a/docs/content/docs/development/calling-llms/model_routing.md +++ b/docs/content/docs/development/calling-llms/model_routing.md @@ -219,9 +219,11 @@ normalize the distribution. The strategy abstains, and the default model answers, when the answer is a refusal, when it is not a choice, when the result carries no answer for the question, or when the chosen candidate's probability is below -`withMinProbability`. For a choice below the minimum, the choice and the -full probability map stay in the metadata. A router using this strategy needs at -least two candidates. A decision call that fails after the configured retries is a request +`withMinProbability`. A provider may omit the chosen candidate from the +probability map: with `withMinProbability` set that is an abstention too, since +there is nothing to compare; without it the choice routes with no `score`. For +either abstention, the choice and the full probability map stay in the +metadata. A router using this strategy needs at least two candidates. A decision call that fails after the configured retries is a request failure, like a failed judge call; it does not select the default or try candidate fallback. @@ -329,7 +331,8 @@ behaviors: | Strategy abstains, or returns a name outside the candidates | Call `defaultModel`, or the first candidate if no default is configured. | | Declaration has an invalid pattern, non-candidate rule target or default, or duplicate/empty candidates | Reject the configuration before processing requests. | | Custom strategy throws, or the judge or decision call exhausts retries | Fail the request; do not treat it as abstention or try candidate fallback. | -| Decision model refuses, answers something other than a choice, or its chosen candidate is below `withMinProbability` | Call `defaultModel`, or the first candidate if no default is configured. | +| Decision model refuses, answers something other than a choice, or its chosen candidate is below `withMinProbability` or has no reported probability while `withMinProbability` is set | Call `defaultModel`, or the first candidate if no default is configured. | +| Decision model's chosen candidate has no reported probability and no `withMinProbability` is set | Call the chosen candidate; the routing event carries no `score`. | | Candidate call fails with fallback disabled | Fail the request after applicable retries. | | Candidate call fails with fallback enabled | Try the remaining candidates in declaration order, stopping at the first successful response. | | All attempted candidates fail | Fail the request with the last candidate's error; earlier errors remain in diagnostics. | diff --git a/e2e-test/cross-language-agent-plan-snapshots/java/agent_plan_with_decision_router.json b/e2e-test/cross-language-agent-plan-snapshots/java/agent_plan_with_decision_router.json index 8b6a88168b..8ddb79e1ed 100644 --- a/e2e-test/cross-language-agent-plan-snapshots/java/agent_plan_with_decision_router.json +++ b/e2e-test/cross-language-agent-plan-snapshots/java/agent_plan_with_decision_router.json @@ -26,6 +26,32 @@ } }, "__resource_provider_type__" : "JavaResourceProvider" + }, + "router_full" : { + "name" : "router_full", + "type" : "model_router", + "descriptor" : { + "language" : "java", + "target_module" : "", + "target_clazz" : "org.apache.flink.agents.api.chat.model.routing.ModelRouter", + "arguments" : { + "candidates" : [ "small", "big" ], + "candidate_descriptions" : { + "small" : "short factual questions", + "big" : "code and sql" + }, + "strategy_type" : "decision_model", + "default_model" : "small", + "fallback" : false, + "strategy_args" : { + "instructions" : "Pick the candidate that should answer.", + "decision_model" : "dm", + "min_probability" : 0.6, + "max_context_chars" : 4000 + } + } + }, + "__resource_provider_type__" : "JavaResourceProvider" } }, "chat_model" : { diff --git a/plan/src/main/java/org/apache/flink/agents/plan/routing/DecisionModelRoutingExecutor.java b/plan/src/main/java/org/apache/flink/agents/plan/routing/DecisionModelRoutingExecutor.java index bf9d7a23bf..118595babe 100644 --- a/plan/src/main/java/org/apache/flink/agents/plan/routing/DecisionModelRoutingExecutor.java +++ b/plan/src/main/java/org/apache/flink/agents/plan/routing/DecisionModelRoutingExecutor.java @@ -55,11 +55,12 @@ *

The answer is constrained by construction: the provider must pick one of the requested * choices, so the decision model cannot steer routing outside the declared candidates. The choice * becomes the selected model and its reported probability the routing score. A refusal, a - * non-choice answer, a result without an answer for the question, or a probability below the - * declared {@code min_probability} abstains to the router's default model with the cause recorded. - * The decision model sees the conversation as text only (tool calls and results included, media and - * reasoning excluded), rendered like the LLM judge's input and subject to the same optional context - * cap. + * non-choice answer, a result without an answer for the question, a probability below the declared + * {@code min_probability}, or no reported probability for the choice while a minimum is set, + * abstains to the router's default model with the cause recorded; without a minimum a choice with + * no reported probability routes with no score. The decision model sees the conversation as text + * only (tool calls and results included, media and reasoning excluded), rendered like the LLM + * judge's input and subject to the same optional context cap. * *

Failure policy: a decision call that exhausts its retries fails the request, like a throwing * rule/custom strategy or a failed judge call. Cancellation propagates and is never persisted as a diff --git a/plan/src/test/java/org/apache/flink/agents/plan/AgentPlanCrossLanguageTest.java b/plan/src/test/java/org/apache/flink/agents/plan/AgentPlanCrossLanguageTest.java index 9069b13884..e96a5385f2 100644 --- a/plan/src/test/java/org/apache/flink/agents/plan/AgentPlanCrossLanguageTest.java +++ b/plan/src/test/java/org/apache/flink/agents/plan/AgentPlanCrossLanguageTest.java @@ -318,9 +318,18 @@ void javaPlanWithJudgeRouterSnapshotIsStable() throws Exception { } /** - * A plan whose only resources are a decision-model-strategy router, its chat models and the - * decision model. + * A plan whose only resources are two decision-model-strategy routers, their chat models and + * the decision model: {@code router} declares the minimum (the common case) and {@code + * router_full} every optional argument, so the Python side has each key's name and JSON type. */ + /** Insertion-ordered so the regenerated snapshot does not reorder the keys per JVM run. */ + private static Map decisionSetupArguments() { + Map arguments = new LinkedHashMap<>(); + arguments.put("model", "m"); + arguments.put("connection", "conn"); + return arguments; + } + private static AgentPlan compileWithDecisionRouter() throws Exception { Map> providers = new LinkedHashMap<>(); ResourceDescriptor routerDescriptor = @@ -331,10 +340,23 @@ private static AgentPlan compileWithDecisionRouter() throws Exception { .defaultModel("small") .fallback(true) .build(); + ResourceDescriptor fullDescriptor = + ModelRouter.of("small", "big") + .describe("small", "short factual questions") + .describe("big", "code and sql") + .strategy( + Strategies.decision("dm", "Pick the candidate that should answer.") + .withMinProbability(0.6) + .withMaxContextChars(4000)) + .defaultModel("small") + .build(); Map routers = new LinkedHashMap<>(); routers.put( "router", new JavaResourceProvider("router", ResourceType.MODEL_ROUTER, routerDescriptor)); + routers.put( + "router_full", + new JavaResourceProvider("router_full", ResourceType.MODEL_ROUTER, fullDescriptor)); providers.put(ResourceType.MODEL_ROUTER, routers); Map chatModels = new LinkedHashMap<>(); ResourceDescriptor model = new ResourceDescriptor("some.Clazz", Map.of()); @@ -347,8 +369,7 @@ private static AgentPlan compileWithDecisionRouter() throws Exception { new JavaResourceProvider( "dm", ResourceType.DECISION_MODEL, - new ResourceDescriptor( - "some.Setup", Map.of("connection", "conn", "model", "m")))); + new ResourceDescriptor("some.Setup", decisionSetupArguments()))); providers.put(ResourceType.DECISION_MODEL, decisionModels); Map connections = new LinkedHashMap<>(); connections.put( diff --git a/plan/src/test/java/org/apache/flink/agents/plan/routing/DecisionModelRoutingExecutorTest.java b/plan/src/test/java/org/apache/flink/agents/plan/routing/DecisionModelRoutingExecutorTest.java index fb6a0286a2..56042cc289 100644 --- a/plan/src/test/java/org/apache/flink/agents/plan/routing/DecisionModelRoutingExecutorTest.java +++ b/plan/src/test/java/org/apache/flink/agents/plan/routing/DecisionModelRoutingExecutorTest.java @@ -385,6 +385,23 @@ void probabilityBelowMinimumAbstainsAndRecordsTheChoice() { .containsKey(DecisionModelRoutingExecutor.PROBABILITIES_KEY); } + /** Without a minimum there is nothing to compare, so the choice routes with no score. */ + @Test + void missingProbabilityForTheChoiceRoutesWithoutAScoreWhenNoMinimumIsSet() { + RoutingDecision decision = + DecisionModelRoutingExecutor.toDecision( + Strategies.decision("dm"), + ctx(List.of(), Map.of()), + choice("big", Map.of("small", 1.0), null), + new LinkedHashMap<>()); + assertThat(decision.isAbstain()).isFalse(); + assertThat(decision.getSelectedModel()).isEqualTo("big"); + assertThat(decision.getScore()).isNull(); + assertThat(decision.getMetadata()) + .containsEntry(DecisionModelRoutingExecutor.PROBABILITIES_KEY, Map.of("small", 1.0)) + .doesNotContainKey(DecisionModelRoutingExecutor.CHOICE_KEY); + } + /** A provider may omit the chosen candidate from the map; with a minimum set that abstains. */ @Test void missingProbabilityForTheChoiceAbstainsWhenAMinimumIsSet() { diff --git a/python/flink_agents/plan/tests/test_agent_plan_cross_language.py b/python/flink_agents/plan/tests/test_agent_plan_cross_language.py index 06671fb911..e0f9f140c8 100644 --- a/python/flink_agents/plan/tests/test_agent_plan_cross_language.py +++ b/python/flink_agents/plan/tests/test_agent_plan_cross_language.py @@ -501,7 +501,9 @@ def test_python_can_deserialize_java_plan_with_decision_router() -> None: """A Java-produced plan carrying the decision-model router strategy must deserialize on the Python side: the strategy travels as the language-neutral tag ``decision_model`` with its arguments, and the decision model itself is an ordinary - DECISION_MODEL resource provider. + DECISION_MODEL resource provider. The snapshot holds two routers: ``router`` with the + required arguments only, and ``router_full`` with every optional argument, which is + the contract a Python implementation of the strategy reads. """ snapshot = _SNAPSHOT_DIR / "java" / "agent_plan_with_decision_router.json" assert snapshot.exists(), ( @@ -520,6 +522,19 @@ def test_python_can_deserialize_java_plan_with_decision_router() -> None: assert args["strategy_args"]["min_probability"] == 0.6 assert args["candidate_descriptions"]["big"] == "code and sql" assert "strategy_clazz" not in args + # Every optional strategy argument, with the key names and JSON types a Python + # implementation of the strategy has to read. + full = restored.resource_providers[ResourceType.MODEL_ROUTER][ + "router_full" + ].descriptor.arguments["strategy_args"] + assert full == { + "decision_model": "dm", + "instructions": "Pick the candidate that should answer.", + "min_probability": 0.6, + "max_context_chars": 4000, + } + assert isinstance(full["min_probability"], float) + assert isinstance(full["max_context_chars"], int) def test_python_can_deserialize_plan_with_java_model_router() -> None: