purushah commented on code in PR #964: URL: https://github.com/apache/flink-agents/pull/964#discussion_r3837429713
########## api/src/main/java/org/apache/flink/agents/api/chat/model/routing/ModelRouter.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.api.chat.model.routing; + +import org.apache.flink.agents.api.resource.Resource; +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 java.lang.reflect.Constructor; +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 + * 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. + * + * <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. + * + * <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 { + + private final List<RoutingCandidate> candidates; + private final String defaultModel; + private final boolean fallbackEnabled; + private final RoutingStrategy strategy; + + public ModelRouter(ResourceDescriptor descriptor, ResourceContext resourceContext) + throws Exception { + super(descriptor, resourceContext); + List<String> names = descriptor.getArgument("candidates"); + if (names == null || names.isEmpty()) { + throw new IllegalArgumentException("ModelRouter requires at least one candidate."); + } + 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.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()) { + 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(); + } + } + + /** Run the strategy for the given context. */ + public RoutingDecision route(RoutingContext context) throws Exception { + return strategy.route(context); + } + + public List<RoutingCandidate> getCandidates() { + return candidates; + } + + public List<String> getCandidateNames() { + List<String> names = new ArrayList<>(); + for (RoutingCandidate candidate : candidates) { + names.add(candidate.getName()); + } + return names; + } + + public Optional<String> getDefaultModel() { + return Optional.ofNullable(defaultModel); + } + + public boolean isFallbackEnabled() { + return fallbackEnabled; + } + + /** Whether the given model name is one of this router's candidates. */ + public boolean isCandidate(String model) { + for (RoutingCandidate candidate : candidates) { + if (candidate.getName().equals(model)) { + return true; + } + } + return false; + } + + @Override + public ResourceType getResourceType() { + return ResourceType.MODEL_ROUTER; + } + + /** + * Start building a router over the given candidate model names (order matters for fallback). + */ + public static Builder of(String... candidates) { + return new Builder(Arrays.asList(candidates)); + } + + /** Fluent builder that produces a {@link ResourceDescriptor} for a {@link ModelRouter}. */ + public static final class Builder { + private final List<String> candidates; + private final Map<String, String> descriptions = new HashMap<>(); + private RoutingStrategyDescriptor strategy; + private String defaultModel; + private boolean fallback = false; + + private Builder(List<String> candidates) { + this.candidates = candidates; + } + + public Builder strategy(RoutingStrategyDescriptor 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 + * (once, on the router) rather than in per-strategy arguments. + */ + public Builder describe(String candidate, String description) { + if (!candidates.contains(candidate)) { + throw new IllegalArgumentException( + String.format( + "Cannot describe '%s': not one of the candidates %s.", + candidate, candidates)); + } + descriptions.put(candidate, description); + return this; + } + + public Builder defaultModel(String defaultModel) { + this.defaultModel = defaultModel; + return this; + } + + /** + * Whether to try remaining candidates (in declaration order) after the selected model has + * exhausted its own retry policy. Applies to the initial routed request only; tool-call + * rounds keep the already-selected model for conversation coherence. Fallback outcomes are + * recorded on the response ({@code model_routing} extra args) and as a second {@code + * ModelRoutingEvent} with source {@code fallback}. + */ + public Builder fallback(boolean fallback) { + this.fallback = fallback; + return this; + } + + public ResourceDescriptor build() { + if (strategy == null) { + throw new IllegalStateException("ModelRouter requires a strategy(...)."); + } + // Rule keys are candidate names; validate here, where both lists are in hand, so a + // typo fails at the registration call site instead of throwing per record at runtime. Review Comment: Good catch on the caching consequence — fixed: `build()` now compiles every rule pattern (and rejects null/non-String values) in the same block as the key check, so a malformed regex fails at the registration call site. Two new builder tests cover it. ########## plan/src/main/java/org/apache/flink/agents/plan/actions/ResolvedModelRoute.java: ########## @@ -0,0 +1,134 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.flink.agents.plan.actions; + +import org.apache.flink.agents.api.event.ModelRoutingEvent; + +import javax.annotation.Nullable; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; + +/** + * The outcome of route resolution for one chat request: the concrete model to call first, the + * candidate set and fallback policy, and the decision facts (source, reason, score, metadata) that + * feed the {@code model_routing} observability block. A plain (unrouted) request is the degenerate + * {@link #direct(String)} route. + */ +final class ResolvedModelRoute { + final String requestedModel; + final String selectedModel; + final List<String> candidates; + final boolean isRouter; + final boolean fallbackEnabled; + final String decisionSource; + @Nullable final String reason; + @Nullable final Double score; + final Map<String, Object> metadata; + + ResolvedModelRoute( + String requestedModel, + String selectedModel, + List<String> candidates, + boolean isRouter, + boolean fallbackEnabled, + String decisionSource, + @Nullable String reason, + @Nullable Double score, + @Nullable Map<String, Object> metadata) { + this.requestedModel = requestedModel; + this.selectedModel = selectedModel; + this.candidates = Collections.unmodifiableList(new ArrayList<>(candidates)); + this.isRouter = isRouter; + this.fallbackEnabled = fallbackEnabled; + this.decisionSource = decisionSource; + this.reason = reason; + this.score = score; + this.metadata = + metadata == null + ? Collections.emptyMap() + : Collections.unmodifiableMap(new HashMap<>(metadata)); + } + + static ResolvedModelRoute direct(String model) { + return new ResolvedModelRoute( + model, + model, + Collections.singletonList(model), + false, + false, + "direct", + null, + null, + null); + } + + /** Candidate order: the strategy's pick first, then declaration order if fallback is on. */ + List<String> attemptOrder() { + List<String> order = new ArrayList<>(); + order.add(this.selectedModel); + if (this.isRouter && this.fallbackEnabled) { + for (String candidate : this.candidates) { + if (!candidate.equals(this.selectedModel)) { + order.add(candidate); + } + } + } + return order; + } + + String durableChatCallId(String candidate) { + if (!this.isRouter) { + return "chat"; + } + return "chat:" + this.requestedModel + ":" + candidate; + } + + /** The {@code model_routing} extra-args block stamped on the loop's final response. */ + Map<String, Object> buildResponseMetadata(String finalModel, List<String> triedModels) { + boolean fallbackAttempted = !finalModel.equals(this.selectedModel); + List<String> fallbackModelsTried = new ArrayList<>(); + for (int i = 1; i < triedModels.size(); i++) { + fallbackModelsTried.add(triedModels.get(i)); + } + Map<String, Object> routing = new LinkedHashMap<>(); + routing.put("router", this.requestedModel); + routing.put("selected_model", this.selectedModel); + routing.put("initial_selected_model", this.selectedModel); Review Comment: You are right — it duplicated `selected_model` on every path (fallback winners are `final_model`). Removed before it ships and becomes a compat constraint. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
