This is an automated email from the ASF dual-hosted git repository.
wenjin272 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/flink-agents.git
The following commit(s) were added to refs/heads/main by this push:
new 57949309 [api][plan][java] Validate router candidates and default
model at declaration (#1151)
57949309 is described below
commit 57949309ab1ac392fab9c39a06b50b7c90e4710b
Author: Sangkyoon Nam <[email protected]>
AuthorDate: Mon Sep 28 15:08:54 2026 +0900
[api][plan][java] Validate router candidates and default model at
declaration (#1151)
Generated-by: Claude Code 2.1.280 (Claude Fable 5.1)
---
.../agents/api/chat/model/routing/ModelRouter.java | 38 ++++-----
.../model/routing/RoutingCandidateValidator.java | 81 +++++++++++++++++++
.../agents/api/chat/model/routing/RoutingTest.java | 90 +++++++++++++++++++++-
docs/content/docs/development/model_routing.md | 4 +-
.../org/apache/flink/agents/plan/AgentPlan.java | 17 ++--
.../plan/AgentPlanLlmJudgeValidationTest.java | 43 ++++++++++-
.../agents/plan/AgentPlanRoutingBackstopTest.java | 15 +++-
7 files changed, 254 insertions(+), 34 deletions(-)
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 9edf4349..a9de9b0c 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
@@ -27,11 +27,9 @@ import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
-import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Optional;
-import java.util.Set;
/**
* A framework resource that <b>selects</b> a concrete chat model for a
request. It does not call
@@ -63,6 +61,9 @@ public class ModelRouter extends Resource {
/** Descriptor key carrying the custom executor class name ({@code CUSTOM}
only). */
public static final String STRATEGY_EXECUTOR_CLASS_KEY =
"strategy_executor_class";
+ /** Descriptor argument naming the candidate the router lands on when the
strategy abstains. */
+ public static final String DEFAULT_MODEL_KEY = "default_model";
+
private final List<RoutingCandidate> candidates;
private final String defaultModel;
private final boolean fallbackEnabled;
@@ -72,28 +73,18 @@ public class ModelRouter extends Resource {
throws Exception {
super(descriptor, resourceContext);
List<String> names = descriptor.getArgument(CANDIDATES_KEY);
- if (names == null || names.isEmpty()) {
- throw new IllegalArgumentException("ModelRouter requires at least
one candidate.");
- }
+ // Read as Object so a deserialized non-String default model is
rejected by the validator
+ // with the candidate message instead of failing the cast here.
+ Object defaultModel = descriptor.getArgument(DEFAULT_MODEL_KEY);
+ RoutingCandidateValidator.validate("ModelRouter", names, defaultModel);
Map<String, String> descriptions =
descriptor.getArgument("candidate_descriptions",
Collections.emptyMap());
List<RoutingCandidate> parsed = new ArrayList<>();
- Set<String> uniqueNames = new LinkedHashSet<>();
for (String name : names) {
- if (!uniqueNames.add(name)) {
- throw new IllegalArgumentException(
- String.format("ModelRouter candidate '%s' is
duplicated.", name));
- }
parsed.add(new RoutingCandidate(name, descriptions.get(name)));
}
this.candidates = Collections.unmodifiableList(parsed);
- this.defaultModel = descriptor.getArgument("default_model");
- if (this.defaultModel != null && !isCandidate(this.defaultModel)) {
- throw new IllegalArgumentException(
- String.format(
- "ModelRouter default model '%s' is not one of the
candidates %s.",
- this.defaultModel, getCandidateNames()));
- }
+ this.defaultModel = (String) defaultModel;
this.fallbackEnabled =
Boolean.TRUE.equals(descriptor.getArgument("fallback",
Boolean.FALSE));
String typeTag = descriptor.getArgument(STRATEGY_TYPE_KEY);
@@ -208,10 +199,21 @@ public class ModelRouter extends Resource {
return this;
}
+ /**
+ * Builds the descriptor to register after checking the declaration.
+ *
+ * @throws IllegalStateException if no strategy was set
+ * @throws IllegalArgumentException if the candidate list is empty, a
candidate name is
+ * empty or duplicated, the default model is not a candidate, or a
rule key is not a
+ * candidate
+ */
public ResourceDescriptor build() {
if (strategy == null) {
throw new IllegalStateException("ModelRouter requires a
strategy(...).");
}
+ // Same check as the constructor: a duplicate or blank candidate,
or a default model
+ // that is not a candidate, fails at the registration call site.
+ RoutingCandidateValidator.validate("ModelRouter", candidates,
defaultModel);
// Rule shape/pattern validation ran in the RoutingStrategy
constructor (the single
// declaration-validation path). build() additionally checks rule
keys against the
// candidate set so a typo fails at the registration call site;
descriptors that skip
@@ -237,7 +239,7 @@ public class ModelRouter extends Resource {
args.put("candidate_descriptions", new
HashMap<>(descriptions));
}
if (defaultModel != null) {
- args.put("default_model", defaultModel);
+ args.put(DEFAULT_MODEL_KEY, defaultModel);
}
args.put("fallback", fallback);
args.put(STRATEGY_TYPE_KEY, strategy.getType().tag());
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
new file mode 100644
index 00000000..aa21b775
--- /dev/null
+++
b/api/src/main/java/org/apache/flink/agents/api/chat/model/routing/RoutingCandidateValidator.java
@@ -0,0 +1,81 @@
+/*
+ * 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.api.chat.model.routing;
+
+import org.apache.flink.annotation.Internal;
+
+import javax.annotation.Nullable;
+
+import java.util.LinkedHashSet;
+import java.util.List;
+import java.util.Set;
+
+/**
+ * Validates a model router's candidate declaration: at least one candidate,
every name a non-empty
+ * String, no duplicates, and a default model (when set) that is one of them.
+ *
+ * <p>This is the single check behind every declaration path: {@link
ModelRouter.Builder#build()} at
+ * the registration call site, plan construction for descriptors that never
went through the
+ * builder, and the {@link ModelRouter} constructor itself. A typo therefore
fails with one message
+ * wherever it is caught, rather than per routed request on the TaskManager,
when the router is
+ * resolved.
+ *
+ * <p>Internal contract shared with the plan module; not a stable public API.
+ */
+@Internal
+public final class RoutingCandidateValidator {
+
+ private RoutingCandidateValidator() {}
+
+ /**
+ * Validate the candidate declaration.
+ *
+ * @param subject how to name the router in the message, e.g. {@code
"ModelRouter"} or {@code
+ * "Model router 'name'"}
+ * @param candidates the declared candidate names; the elements are
checked, so a raw list from
+ * a deserialized descriptor may be passed
+ * @param defaultModel the declared default model, or null for none
+ * @throws IllegalArgumentException when the declaration is invalid
+ */
+ public static void validate(
+ String subject, @Nullable List<?> candidates, @Nullable Object
defaultModel) {
+ if (candidates == null || candidates.isEmpty()) {
+ throw new IllegalArgumentException(subject + " requires at least
one candidate.");
+ }
+ Set<Object> uniqueNames = new LinkedHashSet<>();
+ for (Object name : candidates) {
+ if (!(name instanceof String) || ((String) name).isEmpty()) {
+ throw new IllegalArgumentException(
+ String.format(
+ "%s candidate names must be non-empty strings,
got %s.",
+ subject, name == null ? "null" : "'" + name +
"'"));
+ }
+ if (!uniqueNames.add(name)) {
+ throw new IllegalArgumentException(
+ String.format("%s candidate '%s' is duplicated.",
subject, name));
+ }
+ }
+ if (defaultModel != null && !uniqueNames.contains(defaultModel)) {
+ throw new IllegalArgumentException(
+ String.format(
+ "%s default model '%s' is not one of the
candidates %s.",
+ subject, defaultModel, candidates));
+ }
+ }
+}
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 4dbfe234..cce0ed06 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
@@ -141,17 +141,39 @@ class RoutingTest {
@Test
void defaultModelMustBeCandidate() {
+ // A descriptor built outside the builder (e.g. deserialized) meets
the same default-model
+ // check in the constructor; going through build() would throw before
the constructor runs.
assertThrows(
IllegalArgumentException.class,
() ->
new ModelRouter(
- ModelRouter.of("small", "big")
- .strategy(Strategies.rules(Map.of()))
- .defaultModel("huge")
- .build(),
+ new ResourceDescriptor(
+ ModelRouter.class.getName(),
routerArgs("huge")),
null));
}
+ @Test
+ void routerConstructionRejectsNonStringDefaultModel() {
+ Map<String, Object> args = routerArgs(null);
+ args.put(ModelRouter.DEFAULT_MODEL_KEY, 123);
+ assertThrows(
+ IllegalArgumentException.class,
+ () ->
+ new ModelRouter(
+ new
ResourceDescriptor(ModelRouter.class.getName(), args), null));
+ }
+
+ private static Map<String, Object> routerArgs(String defaultModel) {
+ Map<String, Object> args = new HashMap<>();
+ args.put(ModelRouter.CANDIDATES_KEY, List.of("small", "big"));
+ args.put(ModelRouter.STRATEGY_TYPE_KEY, "rule_based");
+ args.put(ModelRouter.STRATEGY_ARGS_KEY,
Map.of(RoutingStrategy.ARG_RULES, Map.of()));
+ if (defaultModel != null) {
+ args.put(ModelRouter.DEFAULT_MODEL_KEY, defaultModel);
+ }
+ return args;
+ }
+
@Test
void modelRoutingEventRequestIdSurvivesStringForm() {
// Simulate an EventLog JSON round-trip where request_id came back as
a String.
@@ -457,4 +479,64 @@ class RoutingTest {
// The copy re-normalizes through the constructor, so strategies see
an empty list.
assertEquals(0, ctx.getMessages().get(0).getToolCalls().size());
}
+
+ /**
+ * The builder is the registration call site: a duplicate candidate, a
blank one, or a default
+ * model that is not a candidate must fail here, not in the router
constructor on the
+ * TaskManager (the constructor runs when the router is resolved, per
routed request, before the
+ * durable routing call).
+ */
+ @Test
+ void builderRejectsDuplicateCandidate() {
+ IllegalArgumentException e =
+ assertThrows(
+ IllegalArgumentException.class,
+ () ->
+ ModelRouter.of("small", "small")
+ .strategy(Strategies.rules(Map.of()))
+ .build());
+ assertTrue(e.getMessage().contains("duplicated"), e.getMessage());
+ }
+
+ @Test
+ void builderRejectsDefaultModelThatIsNotACandidate() {
+ IllegalArgumentException e =
+ assertThrows(
+ IllegalArgumentException.class,
+ () ->
+ ModelRouter.of("small", "big")
+ .strategy(Strategies.rules(Map.of()))
+ .defaultModel("huge")
+ .build());
+ assertTrue(e.getMessage().contains("huge"), e.getMessage());
+ }
+
+ @Test
+ void builderRequiresAtLeastOneCandidate() {
+ assertThrows(
+ IllegalArgumentException.class,
+ () ->
ModelRouter.of().strategy(Strategies.rules(Map.of())).build());
+ }
+
+ @Test
+ void builderRejectsBlankCandidateName() {
+ assertThrows(
+ IllegalArgumentException.class,
+ () -> ModelRouter.of("small",
"").strategy(Strategies.rules(Map.of())).build());
+ }
+
+ @Test
+ void routerConstructionRejectsBlankCandidateName() {
+ // A descriptor built outside the builder (e.g. deserialized) meets
the same candidate
+ // check in the constructor.
+ Map<String, Object> args = new HashMap<>();
+ args.put(ModelRouter.CANDIDATES_KEY, List.of("small", ""));
+ args.put(ModelRouter.STRATEGY_TYPE_KEY, "rule_based");
+ args.put(ModelRouter.STRATEGY_ARGS_KEY,
Map.of(RoutingStrategy.ARG_RULES, Map.of()));
+ assertThrows(
+ IllegalArgumentException.class,
+ () ->
+ new ModelRouter(
+ new
ResourceDescriptor(ModelRouter.class.getName(), args), null));
+ }
}
diff --git a/docs/content/docs/development/model_routing.md
b/docs/content/docs/development/model_routing.md
index b286c20f..145211a9 100644
--- a/docs/content/docs/development/model_routing.md
+++ b/docs/content/docs/development/model_routing.md
@@ -118,12 +118,12 @@ Resources can equally be registered on the execution
environment with `agentsEnv
| Method | Description |
|--------|-------------|
-| `of(String... candidates)` | Start a router over the given chat model names.
|
+| `of(String... candidates)` | Start a router over the given chat model names.
At least one, each non-empty and unique. |
| `strategy(RoutingStrategy)` | Required. One of the `Strategies` factories
below. |
| `describe(candidate, description)` | Describe a candidate. Descriptions are
the criteria the [LLM judge](#llm-judge) reads. The name must be a candidate. |
| `defaultModel(String)` | Where the router lands when the strategy abstains.
Optional; without it the first candidate is the default. Must be a candidate. |
| `fallback(boolean)` | Try the remaining candidates, in declaration order,
after the selected model fails. Off by default. |
-| `build()` | Produces the `ResourceDescriptor` to register. |
+| `build()` | Produces the `ResourceDescriptor` to register. Fails on a
duplicate or empty candidate, a default model that is not a candidate, or a
rule key that is not a candidate. |
## Routing Strategies
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 06f94fbe..54587434 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
@@ -32,6 +32,7 @@ import org.apache.flink.agents.api.annotation.Tool;
import org.apache.flink.agents.api.annotation.VectorStore;
import org.apache.flink.agents.api.chat.model.routing.CustomRoutingExecutor;
import org.apache.flink.agents.api.chat.model.routing.ModelRouter;
+import
org.apache.flink.agents.api.chat.model.routing.RoutingCandidateValidator;
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.function.JavaFunctionUtils;
@@ -816,6 +817,13 @@ public class AgentPlan implements Serializable {
// 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);
+ // 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.
+ RoutingCandidateValidator.validate(
+ String.format("Model router '%s'", provider.getName()),
+ (List<?>) candidates,
+ descriptor.getArgument(ModelRouter.DEFAULT_MODEL_KEY));
switch (strategy.getType()) {
case LLM_JUDGE:
validateJudge(provider.getName(), strategy, chatModels);
@@ -836,15 +844,15 @@ public class AgentPlan implements Serializable {
* Rule declarations are static constraints like the judge checks above:
the fluent builder
* rejects a bad one at build(), but a descriptor read back from a plan
(deserialized or
* hand-built) never went through the builder. Without this arm they would
surface only per
- * record at request time — inside the durable call — where the IGNORE
error policy silently
+ * record at request time — when the router is resolved — where the IGNORE
error policy silently
* drops every matching record. Rule shape, value types and pattern
validity were already
* enforced by the {@link RoutingStrategy} constructor (regardless of the
'candidates' shape);
* the key-vs-candidate check here mirrors build().
*/
/**
* Fail here, not per record: the router constructor's unchecked read
would turn a mis-shaped
- * 'candidates' value into a raw ClassCastException inside the durable
call. A missing argument
- * is left to the constructor's own message ("requires at least one
candidate").
+ * 'candidates' value into a raw ClassCastException when the router is
resolved. A missing or
+ * empty list is rejected next by {@link RoutingCandidateValidator}.
*/
private static void validateCandidatesShape(String routerName, Object
candidates) {
if (candidates != null && !(candidates instanceof List)) {
@@ -860,9 +868,6 @@ public class AgentPlan implements Serializable {
private static void validateRuleKeys(
String routerName, RoutingStrategy strategy, Object candidates) {
- if (candidates == null) {
- return;
- }
Object rules = strategy.getArguments().get(RoutingStrategy.ARG_RULES);
if (!(rules instanceof Map)) {
return;
diff --git
a/plan/src/test/java/org/apache/flink/agents/plan/AgentPlanLlmJudgeValidationTest.java
b/plan/src/test/java/org/apache/flink/agents/plan/AgentPlanLlmJudgeValidationTest.java
index 7dd1844e..25eb3f2d 100644
---
a/plan/src/test/java/org/apache/flink/agents/plan/AgentPlanLlmJudgeValidationTest.java
+++
b/plan/src/test/java/org/apache/flink/agents/plan/AgentPlanLlmJudgeValidationTest.java
@@ -356,7 +356,7 @@ public class AgentPlanLlmJudgeValidationTest {
/**
* Mis-shaped 'candidates' with VALID rules also fails at plan
construction (review): the router
* constructor's unchecked read would otherwise turn it into a raw
per-record ClassCastException
- * inside the durable call.
+ * when the router is resolved on the TaskManager.
*/
@Test
void misShapedCandidatesFailAtPlanConstruction() {
@@ -397,4 +397,45 @@ public class AgentPlanLlmJudgeValidationTest {
void registeredJudgeModelPassesValidation() {
assertThatCode(() -> new AgentPlan(Map.of(),
providers(true))).doesNotThrowAnyException();
}
+
+ /**
+ * A default model that is not a candidate is a static declaration error
like a typo'd rule key:
+ * the builder catches it, but a descriptor-built plan only met it in the
router constructor —
+ * when the router is resolved on the TaskManager, per routed request,
where IGNORE drops every
+ * record.
+ */
+ @Test
+ void defaultModelNamingNonCandidateFailsAtPlanConstruction() {
+ Map<ResourceType, Map<String, ResourceProvider>> providers =
providers(true);
+ Map<String, Object> args = ruleBasedRouterArgs(List.of("small",
"big"), Map.of());
+ args.put("default_model", "huge");
+ setRouter(providers, new
ResourceDescriptor(ModelRouter.class.getName(), args));
+ assertThatThrownBy(() -> new AgentPlan(Map.of(), providers))
+ .isInstanceOf(IllegalArgumentException.class)
+ .hasMessageContaining("default model")
+ .hasMessageContaining("huge")
+ .hasMessageContaining("router");
+ }
+
+ @Test
+ void duplicateCandidateFailsAtPlanConstruction() {
+ Map<ResourceType, Map<String, ResourceProvider>> providers =
+ ruleBasedProviders(List.of("small", "small"), Map.of());
+ assertThatThrownBy(() -> new AgentPlan(Map.of(), providers))
+ .isInstanceOf(IllegalArgumentException.class)
+ .hasMessageContaining("duplicated")
+ .hasMessageContaining("router");
+ }
+
+ @Test
+ void emptyCandidateListFailsAtPlanConstruction() {
+ Map<ResourceType, Map<String, ResourceProvider>> providers =
providers(true);
+ Map<String, Object> args = ruleBasedRouterArgs(List.of(), Map.of());
+ args.remove("default_model");
+ setRouter(providers, new
ResourceDescriptor(ModelRouter.class.getName(), args));
+ assertThatThrownBy(() -> new AgentPlan(Map.of(), providers))
+ .isInstanceOf(IllegalArgumentException.class)
+ .hasMessageContaining("candidate")
+ .hasMessageContaining("router");
+ }
}
diff --git
a/plan/src/test/java/org/apache/flink/agents/plan/AgentPlanRoutingBackstopTest.java
b/plan/src/test/java/org/apache/flink/agents/plan/AgentPlanRoutingBackstopTest.java
index c3a6744a..a2a07913 100644
---
a/plan/src/test/java/org/apache/flink/agents/plan/AgentPlanRoutingBackstopTest.java
+++
b/plan/src/test/java/org/apache/flink/agents/plan/AgentPlanRoutingBackstopTest.java
@@ -24,6 +24,7 @@ import org.apache.flink.agents.api.resource.ResourceType;
import org.junit.jupiter.api.Test;
import java.util.HashMap;
+import java.util.List;
import java.util.Map;
import static org.junit.jupiter.api.Assertions.assertThrows;
@@ -38,10 +39,18 @@ import static org.junit.jupiter.api.Assertions.assertTrue;
class AgentPlanRoutingBackstopTest {
private static ResourceDescriptor descriptor(String clazz) {
- // Router descriptors must declare a strategy (validated at plan
construction), so the
- // hand-built minimal descriptor carries an empty rules strategy.
+ // Router descriptors must declare a strategy and at least one
candidate (both validated
+ // at plan construction), so the hand-built minimal descriptor carries
an empty rules
+ // strategy over one candidate.
return new ResourceDescriptor(
- clazz, Map.of("strategy_type", "rule_based", "strategy_args",
Map.of()));
+ clazz,
+ Map.of(
+ "candidates",
+ List.of("small"),
+ "strategy_type",
+ "rule_based",
+ "strategy_args",
+ Map.of()));
}
@Test
