[api][plan][python][examples] Framework-managed LLM-as-judge routing (Strategies.llm)

api/src/main/java/org/apache/flink/agents/api/chat/model/BaseChatModelSetup.java

@@ -133,24 +133,23 @@ public void recordTokenMetrics(
         modelGroup.getCounter("completionTokens").inc(completionTokens);
     }
 
-    public ChatMessage chat(List<ChatMessage> messages) {
-        return this.chat(messages, Collections.emptyMap(), Collections.emptyMap());
-    }
-
-    public ChatMessage chat(
-            List<ChatMessage> messages,
-            Map<String, Object> promptArgs,
-            Map<String, Object> modelParams) {
-        Preconditions.checkNotNull(
-                connection,
-                "Connection is not initialized. Ensure open() is called before chat().");
-
-        // Format input messages if set prompt.
-        if (this.prompt != null) {
+    /**
+     * The setup's request-shaping step, shared by {@link #chat} and by the model-routing judge
+     * (which must route on exactly what the selected model will receive): renders the bound {@link
+     * Prompt} (if any) with the prompt args and prepends it to the non-empty conversation messages,
+     * then injects the skill-discovery prompt (if any). Returns the input unchanged when neither is
+     * configured.
+     */
+    public List<ChatMessage> prepareRequestMessages(
+            List<ChatMessage> messages, Map<String, Object> promptArgs) {
+        // Format input messages if set prompt. Read via the accessor so subclasses that override
+        // getPrompt() are honored — the same contract the routing layer inspects.
+        Object boundPrompt = getPrompt();
+        if (boundPrompt != null) {
             Preconditions.checkState(
-                    prompt instanceof Prompt,
+                    boundPrompt instanceof Prompt,
                     "Prompt is not initialized. Ensure open() is called before chat().");
-            Prompt prompt = (Prompt) this.prompt;
+            Prompt prompt = (Prompt) boundPrompt;
             Map<String, String> stringified = new HashMap<>();
             if (promptArgs != null) {
                 for (Map.Entry<String, Object> entry : promptArgs.entrySet()) {
@@ -177,6 +176,22 @@ public ChatMessage chat(
             mutated.add(idx + 1, new ChatMessage(MessageRole.SYSTEM, this.skillDiscoveryPrompt));
             messages = mutated;
         }
+        return messages;
+    }
+
+    public ChatMessage chat(List<ChatMessage> messages) {
+        return this.chat(messages, Collections.emptyMap(), Collections.emptyMap());
+    }
+
+    public ChatMessage chat(
+            List<ChatMessage> messages,
+            Map<String, Object> promptArgs,
+            Map<String, Object> modelParams) {
+        Preconditions.checkNotNull(
+                connection,
+                "Connection is not initialized. Ensure open() is called before chat().");
+
+        messages = prepareRequestMessages(messages, promptArgs);
 
         Map<String, Object> params = this.getParameters();
         if (modelParams != null) {

api/src/main/java/org/apache/flink/agents/api/chat/model/routing/CustomRoutingExecutor.java

@@ -0,0 +1,49 @@
+/*
+ * 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 java.io.Serializable;
+
+/**
+ * The extension point for user-defined routing: given the strategy declaration and a {@link
+ * RoutingContext}, return a {@link RoutingDecision} (a chosen candidate, or {@link
+ * RoutingDecision#abstain()} to defer to the router's default model).
+ *
+ * <p>Custom executors are <b>pure selection logic</b> over request data: the context is data-only
+ * by design, so a custom executor cannot invoke chat models or other engine services inside {@code
+ * route()} — every model call stays on the engine's durable, metered, observable path.
+ * LLM-as-router is available as the framework-managed {@code Strategies.llm(...)} strategy.
+ *
+ * <p>The deployable shape is a named class carried by the strategy declaration ({@code
+ * Strategies.custom(...)}): the class must expose a {@code (Map<String,Object>)} constructor (fed
+ * the declaration's arguments) or a no-arg constructor, and is reconstructed by name on the
+ * TaskManagers — not shipped as a live closure.
+ */
+public interface CustomRoutingExecutor extends Serializable {
+
+    /**
+     * Select a model for the given routing context.
+     *
+     * @param strategy the declaration this executor was configured with (type + arguments)
+     * @param context the request messages, prompt args, and candidates
+     * @return the routing decision (selected candidate or abstain)
+     * @throws Exception if the executor fails
+     */
+    RoutingDecision route(RoutingStrategy strategy, RoutingContext context) throws Exception;
+}

api/src/main/java/org/apache/flink/agents/api/chat/model/routing/ModelRouter.java

@@ -23,7 +23,6 @@
 import org.apache.flink.agents.api.resource.ResourceDescriptor;
 import org.apache.flink.agents.api.resource.ResourceType;
 
-import java.lang.reflect.Constructor;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
@@ -33,24 +32,37 @@
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
-import java.util.regex.Pattern;
-import java.util.regex.PatternSyntaxException;
 
 /**
  * A framework resource that <b>selects</b> a concrete chat model for a request. It does not call
- * the backend itself — {@code ChatModelAction} resolves the router, runs its {@link
- * RoutingStrategy} to get a {@link RoutingDecision}, and then runs the normal chat path against the
- * chosen model.
+ * the backend itself, nor does it execute routing logic: it carries the candidates plus a {@link
+ * RoutingStrategy} <i>declaration</i>, and the engine's plan layer resolves an executor for the
+ * declared type ({@code ChatModelAction} runs it, then runs the normal chat path against the chosen
+ * model).
  *
  * <p>Built with the fluent {@link #of(String...)} builder, which produces a {@link
- * ResourceDescriptor} the framework instantiates reflectively. The strategy is carried by class
- * name + args (see {@link RoutingStrategyDescriptor}) so it is plan-serializable.
+ * ResourceDescriptor} the framework instantiates reflectively. The router carries declaration data
+ * only — candidates plus the strategy as a language-neutral type tag + arguments — so it is
+ * plan-serializable across runtimes; compilation, caching and execution of strategies live in the
+ * engine's plan layer (the {@code RoutingExecutor} implementations).
  *
  * <p>Abstain ({@link RoutingDecision#abstain()}) → {@link #getDefaultModel()}. A returned name that
  * is not a candidate is an invalid decision and is failed clearly by the caller.
  */
 public class ModelRouter extends Resource {
 
+    /** Descriptor key carrying the candidate model names. */
+    public static final String CANDIDATES_KEY = "candidates";
+
+    /** Descriptor key carrying the strategy type tag ({@link RoutingStrategyType#tag()}). */
+    public static final String STRATEGY_TYPE_KEY = "strategy_type";
+
+    /** Descriptor key carrying the strategy arguments. */
+    public static final String STRATEGY_ARGS_KEY = "strategy_args";
+
+    /** Descriptor key carrying the custom executor class name ({@code CUSTOM} only). */
+    public static final String STRATEGY_EXECUTOR_CLASS_KEY = "strategy_executor_class";
+
     private final List<RoutingCandidate> candidates;
     private final String defaultModel;
     private final boolean fallbackEnabled;
@@ -59,7 +71,7 @@ public class ModelRouter extends Resource {
     public ModelRouter(ResourceDescriptor descriptor, ResourceContext resourceContext)
             throws Exception {
         super(descriptor, resourceContext);
-        List<String> names = descriptor.getArgument("candidates");
+        List<String> names = descriptor.getArgument(CANDIDATES_KEY);
         if (names == null || names.isEmpty()) {
             throw new IllegalArgumentException("ModelRouter requires at least one candidate.");
         }
@@ -84,30 +96,23 @@ public ModelRouter(ResourceDescriptor descriptor, ResourceContext resourceContex
         }
         this.fallbackEnabled =
                 Boolean.TRUE.equals(descriptor.getArgument("fallback", Boolean.FALSE));
-        String strategyClazz = descriptor.getArgument("strategy_clazz");
-        Map<String, Object> strategyArgs =
-                descriptor.getArgument("strategy_args", Collections.emptyMap());
-        this.strategy = instantiateStrategy(strategyClazz, strategyArgs);
-    }
-
-    @SuppressWarnings("unchecked")
-    private static RoutingStrategy instantiateStrategy(String clazz, Map<String, Object> args)
-            throws Exception {
-        if (clazz == null || clazz.isEmpty()) {
+        String typeTag = descriptor.getArgument(STRATEGY_TYPE_KEY);
+        if (typeTag == null || typeTag.isEmpty()) {
             throw new IllegalArgumentException("ModelRouter requires a routing strategy.");
         }
-        Class<?> c = Class.forName(clazz, true, Thread.currentThread().getContextClassLoader());
-        try {
-            Constructor<?> ctor = c.getConstructor(Map.class);
-            return (RoutingStrategy) ctor.newInstance(args);
-        } catch (NoSuchMethodException noMapCtor) {
-            return (RoutingStrategy) c.getConstructor().newInstance();
-        }
+        Map<String, Object> strategyArgs =
+                descriptor.getArgument(STRATEGY_ARGS_KEY, Collections.emptyMap());
+        String executorClass = descriptor.getArgument(STRATEGY_EXECUTOR_CLASS_KEY);
+        // The declaration constructor owns the per-type argument rules, so a structurally invalid
+        // configuration fails here (resource construction) with the same message as at build().
+        this.strategy =
+                new RoutingStrategy(
+                        RoutingStrategyType.fromTag(typeTag), strategyArgs, executorClass);
     }
 
-    /** Run the strategy for the given context. */
-    public RoutingDecision route(RoutingContext context) throws Exception {
-        return strategy.route(context);
+    /** The configured strategy declaration (type + arguments). */
+    public RoutingStrategy getStrategy() {
+        return strategy;
     }
 
     public List<RoutingCandidate> getCandidates() {
@@ -156,23 +161,23 @@ public static Builder of(String... candidates) {
     public static final class Builder {
         private final List<String> candidates;
         private final Map<String, String> descriptions = new HashMap<>();
-        private RoutingStrategyDescriptor strategy;
+        private RoutingStrategy strategy;
         private String defaultModel;
         private boolean fallback = false;
 
         private Builder(List<String> candidates) {
             this.candidates = candidates;
         }
 
-        public Builder strategy(RoutingStrategyDescriptor strategy) {
+        public Builder strategy(RoutingStrategy strategy) {
             this.strategy = strategy;
             return this;
         }
 
         /**
          * Attach a human-readable description to a candidate, surfaced to strategies via {@link
-         * RoutingCandidate#getDescription()}. Descriptions are how semantic strategies — and future
-         * framework-managed LLM routing — learn what each candidate is for, so declare them here
+         * RoutingCandidate#getDescription()}. Descriptions are how semantic strategies — including
+         * the framework-managed LLM judge — learn what each candidate is for, so declare them here
          * (once, on the router) rather than in per-strategy arguments.
          */
         public Builder describe(String candidate, String description) {
@@ -207,52 +212,39 @@ public ResourceDescriptor build() {
             if (strategy == null) {
                 throw new IllegalStateException("ModelRouter requires a strategy(...).");
             }
-            // Rule keys are candidate names and rule values are regex patterns; validate both
-            // here, where the full map is in hand, so a typo fails at the registration call site
-            // instead of throwing per record at runtime (an invalid pattern is never cached by
-            // the resource cache, so it would otherwise re-throw on every routed request).
-            if (RuleBasedRoutingStrategy.class.getName().equals(strategy.getClazz())) {
-                Object rules = strategy.getArguments().get("rules");
+            // 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
+            // the builder get the same check at plan construction (AgentPlan#validateRuleKeys,
+            // with a router-scoped message).
+            if (strategy.getType() == RoutingStrategyType.RULE_BASED) {
+                Object rules = strategy.getArguments().get(RoutingStrategy.ARG_RULES);
                 if (rules instanceof Map) {
-                    for (Map.Entry<?, ?> rule : ((Map<?, ?>) rules).entrySet()) {
-                        if (!candidates.contains(String.valueOf(rule.getKey()))) {
-                            throw new IllegalArgumentException(
-                                    String.format(
-                                            "Routing rule key '%s' is not one of the candidates %s.",
-                                            rule.getKey(), candidates));
-                        }
-                        if (!(rule.getValue() instanceof String)) {
+                    for (Object ruleKey : ((Map<?, ?>) rules).keySet()) {
+                        if (!candidates.contains(ruleKey)) {
                             throw new IllegalArgumentException(
                                     String.format(
-                                            "Routing rule pattern for candidate '%s' must be a non-null String, got %s.",
-                                            rule.getKey(),
-                                            rule.getValue() == null
-                                                    ? "null"
-                                                    : rule.getValue().getClass().getSimpleName()));
-                        }
-                        try {
-                            Pattern.compile((String) rule.getValue());
-                        } catch (PatternSyntaxException e) {
-                            throw new IllegalArgumentException(
-                                    String.format(
-                                            "Routing rule pattern '%s' for candidate '%s' is not a valid regex.",
-                                            rule.getValue(), rule.getKey()),
-                                    e);
+                                            "Routing rule key '%s' is not one of the candidates"
+                                                    + " %s.",
+                                            ruleKey, candidates));
                         }
                     }
                 }
             }
             Map<String, Object> args = new HashMap<>();
-            args.put("candidates", new ArrayList<>(candidates));
+            args.put(CANDIDATES_KEY, new ArrayList<>(candidates));
             if (!descriptions.isEmpty()) {
                 args.put("candidate_descriptions", new HashMap<>(descriptions));
             }
             if (defaultModel != null) {
                 args.put("default_model", defaultModel);
             }
             args.put("fallback", fallback);
-            args.put("strategy_clazz", strategy.getClazz());
-            args.put("strategy_args", strategy.getArguments());
+            args.put(STRATEGY_TYPE_KEY, strategy.getType().tag());
+            args.put(STRATEGY_ARGS_KEY, strategy.getArguments());
+            if (strategy.getExecutorClass() != null) {
+                args.put(STRATEGY_EXECUTOR_CLASS_KEY, strategy.getExecutorClass());
+            }
             return new ResourceDescriptor(ModelRouter.class.getName(), args);
         }
     }

api/src/main/java/org/apache/flink/agents/api/chat/model/routing/RoutingContext.java

@@ -26,15 +26,16 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.UUID;
 
 /**
- * Read-only view a {@link RoutingStrategy} sees when deciding which model to route to.
+ * Read-only view a routing executor sees when deciding which model to route to.
  *
- * <p>v1 exposes the request id, the request messages, prompt args, and the router's candidates
- * (name + description). It intentionally does <b>not</b> expose a chat-invocation API, so a
- * strategy cannot make a hidden synchronous model call; observable LLM-as-router is a
- * framework-managed follow-up.
+ * <p>It exposes the request id, the request messages, prompt args, and the router's candidates
+ * (name + description). It intentionally does <b>not</b> expose a chat-invocation API, so a {@link
+ * CustomRoutingExecutor} cannot make a hidden synchronous model call; observable LLM-as-router is
+ * provided by the framework-managed {@code Strategies.llm(...)} strategy instead.
  *
  * <p>The isolation boundary is deliberate and one level deep: the message list, each message's
  * tool-call maps and extra args, and the prompt-args map are defensive copies, but values
@@ -50,13 +51,24 @@ public final class RoutingContext {
     private final List<ChatMessage> messages;
     private final Map<String, Object> promptArgs;
     private final List<RoutingCandidate> candidates;
+    private final String defaultModel;
 
     public RoutingContext(
             UUID requestId,
             String router,
             List<ChatMessage> messages,
             Map<String, Object> promptArgs,
             List<RoutingCandidate> candidates) {
+        this(requestId, router, messages, promptArgs, candidates, null);
+    }
+
+    public RoutingContext(
+            UUID requestId,
+            String router,
+            List<ChatMessage> messages,
+            Map<String, Object> promptArgs,
+            List<RoutingCandidate> candidates,
+            String defaultModel) {
         this.requestId = requestId;
         this.router = router;
         // Deep copy: the wrapping list is unmodifiable, but ChatMessage is mutable and the
@@ -74,6 +86,7 @@ public RoutingContext(
                 candidates == null
                         ? Collections.emptyList()
                         : Collections.unmodifiableList(new ArrayList<>(candidates));
+        this.defaultModel = defaultModel;
     }
 
     private static List<ChatMessage> deepCopy(List<ChatMessage> messages) {
@@ -116,6 +129,11 @@ public String getRouter() {
         return router;
     }
 
+    /** The router's declared default model — where abstains resolve — if one is configured. */
+    public Optional<String> getDefaultModel() {
+        return Optional.ofNullable(defaultModel);
+    }
+
     public List<ChatMessage> getMessages() {
         return messages;
     }