diff --git a/api/src/main/java/org/apache/flink/agents/api/AgentsExecutionEnvironment.java b/api/src/main/java/org/apache/flink/agents/api/AgentsExecutionEnvironment.java index cb4a3c6b2..9095f23dd 100644 --- a/api/src/main/java/org/apache/flink/agents/api/AgentsExecutionEnvironment.java +++ b/api/src/main/java/org/apache/flink/agents/api/AgentsExecutionEnvironment.java @@ -209,6 +209,7 @@ public AgentsExecutionEnvironment addResource(String name, ResourceType type, Ob if (resources.get(type).containsKey(name)) { throw new IllegalArgumentException(String.format("%s %s already defined.", type, name)); } + Agent.checkNoChatModelRouterNameClash(name, type, resources); if (instance instanceof SerializableResource) { resources.get(type).put(name, instance); diff --git a/api/src/main/java/org/apache/flink/agents/api/EventType.java b/api/src/main/java/org/apache/flink/agents/api/EventType.java index 971134ec3..8dc923439 100644 --- a/api/src/main/java/org/apache/flink/agents/api/EventType.java +++ b/api/src/main/java/org/apache/flink/agents/api/EventType.java @@ -41,6 +41,8 @@ public final class EventType { org.apache.flink.agents.api.event.ContextRetrievalRequestEvent.EVENT_TYPE; public static final String ContextRetrievalResponseEvent = org.apache.flink.agents.api.event.ContextRetrievalResponseEvent.EVENT_TYPE; + public static final String ModelRoutingEvent = + org.apache.flink.agents.api.event.ModelRoutingEvent.EVENT_TYPE; public static final String ShortTermWriteEvent = org.apache.flink.agents.api.event.ShortTermWriteEvent.EVENT_TYPE; @@ -70,6 +72,7 @@ public final class EventType { Map.entry("ToolResponseEvent", ToolResponseEvent), Map.entry("ContextRetrievalRequestEvent", ContextRetrievalRequestEvent), Map.entry("ContextRetrievalResponseEvent", ContextRetrievalResponseEvent), + Map.entry("ModelRoutingEvent", ModelRoutingEvent), Map.entry("ShortTermWriteEvent", ShortTermWriteEvent), Map.entry("ShortTermReadEvent", ShortTermReadEvent), Map.entry("SensoryWriteEvent", SensoryWriteEvent), diff --git a/api/src/main/java/org/apache/flink/agents/api/agents/Agent.java b/api/src/main/java/org/apache/flink/agents/api/agents/Agent.java index 1f0377c5a..4894924a2 100644 --- a/api/src/main/java/org/apache/flink/agents/api/agents/Agent.java +++ b/api/src/main/java/org/apache/flink/agents/api/agents/Agent.java @@ -119,6 +119,7 @@ public Agent addResource(String name, ResourceType type, Object instance) { if (resources.get(type).containsKey(name)) { throw new IllegalArgumentException(String.format("%s %s already defined.", type, name)); } + checkNoChatModelRouterNameClash(name, type, resources); if (instance instanceof SerializableResource) { resources.get(type).put(name, instance); @@ -131,6 +132,29 @@ public Agent addResource(String name, ResourceType type, Object instance) { return this; } + /** + * Chat models and model routers share the chat request namespace ({@code ChatRequestEvent} + * names either), so one name must not be registered as both. Checked here, at the registration + * call site, so the failure points at the user's own {@code addResource} line; {@code + * AgentPlan} re-validates as a backstop. + */ + public static void checkNoChatModelRouterNameClash( + String name, ResourceType type, Map> resources) { + ResourceType clashing = + type == ResourceType.CHAT_MODEL + ? ResourceType.MODEL_ROUTER + : type == ResourceType.MODEL_ROUTER ? ResourceType.CHAT_MODEL : null; + if (clashing != null + && resources.containsKey(clashing) + && resources.get(clashing).containsKey(name)) { + throw new IllegalArgumentException( + String.format( + "'%s' is already registered as %s; chat models and model routers share the" + + " chat request namespace and must use distinct names.", + name, clashing)); + } + } + public enum ErrorHandlingStrategy { FAIL("fail"), RETRY("retry"), 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 new file mode 100644 index 000000000..4ba62d678 --- /dev/null +++ b/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 selects 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. + * + *

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. + * + *

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 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 names = descriptor.getArgument("candidates"); + if (names == null || names.isEmpty()) { + throw new IllegalArgumentException("ModelRouter requires at least one candidate."); + } + Map descriptions = + descriptor.getArgument("candidate_descriptions", Collections.emptyMap()); + List parsed = new ArrayList<>(); + Set 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 strategyArgs = + descriptor.getArgument("strategy_args", Collections.emptyMap()); + this.strategy = instantiateStrategy(strategyClazz, strategyArgs); + } + + @SuppressWarnings("unchecked") + private static RoutingStrategy instantiateStrategy(String clazz, Map 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 getCandidates() { + return candidates; + } + + public List getCandidateNames() { + List names = new ArrayList<>(); + for (RoutingCandidate candidate : candidates) { + names.add(candidate.getName()); + } + return names; + } + + public Optional 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 candidates; + private final Map descriptions = new HashMap<>(); + private RoutingStrategyDescriptor strategy; + private String defaultModel; + private boolean fallback = false; + + private Builder(List 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. + if (RuleBasedRoutingStrategy.class.getName().equals(strategy.getClazz())) { + Object rules = strategy.getArguments().get("rules"); + if (rules instanceof Map) { + for (Object ruleKey : ((Map) rules).keySet()) { + if (!candidates.contains(String.valueOf(ruleKey))) { + throw new IllegalArgumentException( + String.format( + "Routing rule key '%s' is not one of the candidates %s.", + ruleKey, candidates)); + } + } + } + } + Map args = new HashMap<>(); + args.put("candidates", 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()); + return new ResourceDescriptor(ModelRouter.class.getName(), args); + } + } +} diff --git a/api/src/main/java/org/apache/flink/agents/api/chat/model/routing/RoutingCandidate.java b/api/src/main/java/org/apache/flink/agents/api/chat/model/routing/RoutingCandidate.java new file mode 100644 index 000000000..b8f8006e6 --- /dev/null +++ b/api/src/main/java/org/apache/flink/agents/api/chat/model/routing/RoutingCandidate.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 java.io.Serializable; +import java.util.Objects; + +/** + * A candidate chat model a {@link ModelRouter} may select, as seen by a {@link RoutingStrategy}. + * + *

Carries the candidate's registered model name plus an optional human-readable description + * (declared via {@code ModelRouter.Builder#describe}) that semantic strategies — and future + * framework-managed LLM routing — can use to decide. The name must resolve to a registered {@code + * CHAT_MODEL}. Per-candidate metadata (e.g. cost or load hints) is deferred until a strategy can + * actually consume it. + */ +public final class RoutingCandidate implements Serializable { + + private static final long serialVersionUID = 1L; + + private final String name; + private final String description; + + public RoutingCandidate(String name, String description) { + if (name == null || name.isEmpty()) { + throw new IllegalArgumentException("Candidate name must be non-null and non-empty."); + } + this.name = name; + this.description = description == null ? "" : description; + } + + public RoutingCandidate(String name) { + this(name, ""); + } + + public String getName() { + return name; + } + + public String getDescription() { + return description; + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + RoutingCandidate that = (RoutingCandidate) o; + return name.equals(that.name) && description.equals(that.description); + } + + @Override + public int hashCode() { + return Objects.hash(name, description); + } + + @Override + public String toString() { + return "RoutingCandidate{name='" + name + "', description='" + description + "'}"; + } +} diff --git a/api/src/main/java/org/apache/flink/agents/api/chat/model/routing/RoutingContext.java b/api/src/main/java/org/apache/flink/agents/api/chat/model/routing/RoutingContext.java new file mode 100644 index 000000000..67a1688a3 --- /dev/null +++ b/api/src/main/java/org/apache/flink/agents/api/chat/model/routing/RoutingContext.java @@ -0,0 +1,159 @@ +/* + * 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.chat.messages.ChatMessage; +import org.apache.flink.agents.api.chat.messages.MessageRole; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.UUID; + +/** + * Read-only view a {@link RoutingStrategy} sees when deciding which model to route to. + * + *

v1 exposes the request id, the request messages, prompt args, and the router's candidates + * (name + description). It intentionally does not 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. + * + *

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 + * nested inside those maps are shared with the request that is actually sent. Strategies + * must treat the context as read-only; the copies exist to make accidental top-level mutation + * harmless, not to sandbox a hostile strategy (arbitrary-depth copies on every routing decision + * would tax the common case to guard a case the SPI already forbids). + */ +public final class RoutingContext { + + private final UUID requestId; + private final String router; + private final List messages; + private final Map promptArgs; + private final List candidates; + + public RoutingContext( + UUID requestId, + String router, + List messages, + Map promptArgs, + List candidates) { + this.requestId = requestId; + this.router = router; + // Deep copy: the wrapping list is unmodifiable, but ChatMessage is mutable and the + // caller passes the same instances that go to the model — a strategy calling + // setContent(...) on a shallow copy would silently rewrite the prompt actually sent. + this.messages = + messages == null + ? Collections.emptyList() + : Collections.unmodifiableList(deepCopy(messages)); + this.promptArgs = + promptArgs == null + ? Collections.emptyMap() + : Collections.unmodifiableMap(new HashMap<>(promptArgs)); + this.candidates = + candidates == null + ? Collections.emptyList() + : Collections.unmodifiableList(new ArrayList<>(candidates)); + } + + private static List deepCopy(List messages) { + List copy = new ArrayList<>(messages.size()); + for (ChatMessage m : messages) { + if (m == null) { + continue; + } + // ChatMessage's constructor copies extraArgs but stores toolCalls by reference, so + // copy the list AND each tool-call map — otherwise a strategy could still mutate the + // tool calls of the message actually sent. getToolCalls() can be null despite the + // constructor's default: the Jackson setter stores null as-is, so a message + // deserialized from JSON with an explicit "tool_calls": null carries null here. + List> source = m.getToolCalls(); + List> toolCalls; + if (source == null) { + toolCalls = null; + } else { + toolCalls = new ArrayList<>(source.size()); + for (Map call : source) { + toolCalls.add(call == null ? null : new HashMap<>(call)); + } + } + copy.add(new ChatMessage(m.getRole(), m.getContent(), toolCalls, m.getExtraArgs())); + } + return copy; + } + + /** + * Id of the initial chat request being routed. Lets strategies correlate their own logs with + * the framework's events, and enables deterministic per-request policies (e.g. hash-based A/B + * splits). + */ + public UUID getRequestId() { + return requestId; + } + + /** Name of the router resource handling this request. */ + public String getRouter() { + return router; + } + + public List getMessages() { + return messages; + } + + public Map getPromptArgs() { + return promptArgs; + } + + public List getCandidates() { + return candidates; + } + + /** + * Content of the first user message, or an empty string if there is none. Note that when the + * request carries conversation history this is the oldest user turn; strategies that + * should react to the current question want {@link #lastUserMessage()}. + */ + public String firstUserMessage() { + for (ChatMessage message : messages) { + if (message.getRole() == MessageRole.USER) { + return message.getContent() == null ? "" : message.getContent(); + } + } + return ""; + } + + /** + * Content of the most recent user message, or an empty string if there is none. This is the + * current question in a multi-turn conversation and the default input for rule/keyword + * strategies. + */ + public String lastUserMessage() { + for (int i = messages.size() - 1; i >= 0; i--) { + ChatMessage message = messages.get(i); + if (message.getRole() == MessageRole.USER) { + return message.getContent() == null ? "" : message.getContent(); + } + } + return ""; + } +} diff --git a/api/src/main/java/org/apache/flink/agents/api/chat/model/routing/RoutingDecision.java b/api/src/main/java/org/apache/flink/agents/api/chat/model/routing/RoutingDecision.java new file mode 100644 index 000000000..c711eea6a --- /dev/null +++ b/api/src/main/java/org/apache/flink/agents/api/chat/model/routing/RoutingDecision.java @@ -0,0 +1,179 @@ +/* + * 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 com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; + +import javax.annotation.Nullable; + +import java.io.Serializable; +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; + +/** + * The structured result of a {@link RoutingStrategy}. Carries the selected model name plus optional + * explanation (reason, score, metadata), or signals abstention so the router falls back to its + * default model. + * + *

Returning a name that is not one of the router's candidates is an invalid decision and + * is failed clearly by the router (not silently defaulted). Abstention ({@link #abstain()}) is the + * intended way to defer to the default model. + * + *

JSON-serializable so it can be persisted/replayed as a durable {@code "route"} call result. + */ +public final class RoutingDecision implements Serializable { + + private static final long serialVersionUID = 1L; + + private final String selectedModel; + private final boolean abstain; + private final String reason; + private final Double score; + private final Map metadata; + private final Double decisionMs; + + @JsonCreator + public RoutingDecision( + @JsonProperty("selected_model") String selectedModel, + @JsonProperty("abstain") boolean abstain, + @JsonProperty("reason") String reason, + @JsonProperty("score") Double score, + @JsonProperty("metadata") Map metadata, + @JsonProperty("decision_ms") Double decisionMs) { + // Invariants hold on every construction path — including JSON deserialization, which is + // the path durable replay takes. + if (abstain && selectedModel != null) { + throw new IllegalArgumentException( + "An abstaining decision must not carry a selected model."); + } + if (!abstain && (selectedModel == null || selectedModel.isEmpty())) { + throw new IllegalArgumentException( + "A non-abstain decision requires a selected model; use abstain() to defer."); + } + this.selectedModel = selectedModel; + this.abstain = abstain; + this.reason = reason; + this.score = score; + this.metadata = + metadata == null + ? Collections.emptyMap() + : Collections.unmodifiableMap(new HashMap<>(metadata)); + this.decisionMs = decisionMs; + } + + /** A decision selecting the given candidate model. */ + public static RoutingDecision of(String selectedModel) { + return new RoutingDecision(selectedModel, false, null, null, Collections.emptyMap(), null); + } + + /** A decision to abstain, deferring to the router's default model. */ + public static RoutingDecision abstain() { + return new RoutingDecision(null, true, null, null, Collections.emptyMap(), null); + } + + /** + * Copy of this decision stamped with the strategy's wall-clock time. Framework-recorded inside + * the durable {@code "route"} call, so — when an action-state store is configured — a replayed + * decision reports its original latency, not the replay's. Without a store (the + * default) the decision re-executes on recovery and reports fresh timing. + */ + public RoutingDecision withDecisionMs(double decisionMs) { + return new RoutingDecision(selectedModel, abstain, reason, score, metadata, decisionMs); + } + + /** Start building a rich decision (with reason/score/metadata) for the given model. */ + public static Builder builder(String selectedModel) { + return new Builder(selectedModel); + } + + @Nullable + @JsonProperty("selected_model") + public String getSelectedModel() { + return selectedModel; + } + + public boolean isAbstain() { + return abstain; + } + + @Nullable + public String getReason() { + return reason; + } + + @Nullable + public Double getScore() { + return score; + } + + public Map getMetadata() { + return metadata; + } + + /** Strategy wall-clock time in milliseconds, if recorded (see {@link #withDecisionMs}). */ + @Nullable + @JsonProperty("decision_ms") + public Double getDecisionMs() { + return decisionMs; + } + + @Override + public String toString() { + return abstain + ? "RoutingDecision{abstain}" + : "RoutingDecision{selectedModel='" + selectedModel + "', reason=" + reason + "}"; + } + + /** Builder for a rich {@link RoutingDecision}. */ + public static final class Builder { + private final String selectedModel; + private String reason; + private Double score; + private final Map metadata = new HashMap<>(); + + private Builder(String selectedModel) { + if (selectedModel == null || selectedModel.isEmpty()) { + throw new IllegalArgumentException( + "Selected model must be non-null and non-empty."); + } + this.selectedModel = selectedModel; + } + + public Builder reason(String reason) { + this.reason = reason; + return this; + } + + public Builder score(double score) { + this.score = score; + return this; + } + + public Builder metadata(String key, Object value) { + this.metadata.put(key, value); + return this; + } + + public RoutingDecision build() { + return new RoutingDecision(selectedModel, false, reason, score, metadata, null); + } + } +} 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 new file mode 100644 index 000000000..752aa5c02 --- /dev/null +++ b/api/src/main/java/org/apache/flink/agents/api/chat/model/routing/RoutingStrategy.java @@ -0,0 +1,47 @@ +/* + * 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 single extension point for model routing: given a {@link RoutingContext}, return a {@link + * RoutingDecision} (a chosen candidate, or {@link RoutingDecision#abstain()} to defer to the + * router's default model). + * + *

v1 strategies are pure selection logic: they must not invoke chat models or other + * external systems inside {@code route()}. LLM-as-router (a judge model call) is a follow-up that + * the framework will run on the observable, durable chat path — not hidden inside a strategy. + * + *

The deployable shape of a custom strategy is a named class (or descriptor) that serializes + * with the agent plan to the TaskManagers; a lambda is a local convenience and must be + * serializable. + */ +@FunctionalInterface +public interface RoutingStrategy extends Serializable { + + /** + * Select a model for the given routing context. + * + * @param context the request messages, prompt args, and candidates + * @return the routing decision (selected candidate or abstain) + * @throws Exception if the strategy fails + */ + RoutingDecision route(RoutingContext context) throws Exception; +} diff --git a/api/src/main/java/org/apache/flink/agents/api/chat/model/routing/RoutingStrategyDescriptor.java b/api/src/main/java/org/apache/flink/agents/api/chat/model/routing/RoutingStrategyDescriptor.java new file mode 100644 index 000000000..14c8f09e5 --- /dev/null +++ b/api/src/main/java/org/apache/flink/agents/api/chat/model/routing/RoutingStrategyDescriptor.java @@ -0,0 +1,55 @@ +/* + * 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; +import java.util.Collections; +import java.util.Map; + +/** + * Deployable description of a {@link RoutingStrategy}: the fully-qualified class name plus its + * construction arguments. Carried in the router's {@code ResourceDescriptor} so the strategy is + * plan-serializable and reconstructed on the TaskManagers by name — not shipped as a live closure. + * + *

Built-ins are produced by the {@link Strategies} factory (which fills in the class name); + * there is no magic-string strategy keyword. + */ +public final class RoutingStrategyDescriptor implements Serializable { + + private static final long serialVersionUID = 1L; + + private final String clazz; + private final Map arguments; + + public RoutingStrategyDescriptor(String clazz, Map arguments) { + if (clazz == null || clazz.isEmpty()) { + throw new IllegalArgumentException("Strategy class must be non-null and non-empty."); + } + this.clazz = clazz; + this.arguments = arguments == null ? Collections.emptyMap() : arguments; + } + + public String getClazz() { + return clazz; + } + + public Map getArguments() { + return arguments; + } +} diff --git a/api/src/main/java/org/apache/flink/agents/api/chat/model/routing/RuleBasedRoutingStrategy.java b/api/src/main/java/org/apache/flink/agents/api/chat/model/routing/RuleBasedRoutingStrategy.java new file mode 100644 index 000000000..4cc4dcafb --- /dev/null +++ b/api/src/main/java/org/apache/flink/agents/api/chat/model/routing/RuleBasedRoutingStrategy.java @@ -0,0 +1,96 @@ +/* + * 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.util.LinkedHashMap; +import java.util.Map; +import java.util.regex.Pattern; + +/** + * Built-in keyword/regex routing strategy. Configured with a map of {@code candidateModel -> + * regex}; the first candidate whose regex matches the most recent user message (case-insensitive + * find) wins, evaluated in the map's iteration order (pass a {@code LinkedHashMap} when precedence + * matters). If nothing matches, the strategy abstains so the router uses its default model. + * + *

Constructed reflectively from a {@link RoutingStrategyDescriptor} via the {@code + * (Map)} constructor; use {@link Strategies#rules(Map)} to build one. + */ +public class RuleBasedRoutingStrategy implements RoutingStrategy { + + private static final long serialVersionUID = 1L; + + private final Map rules; + + @SuppressWarnings("unchecked") + public RuleBasedRoutingStrategy(Map args) { + this.rules = new LinkedHashMap<>(); + Object raw = args == null ? null : args.get("rules"); + if (raw instanceof Map) { + for (Map.Entry entry : ((Map) raw).entrySet()) { + String candidate = entry.getKey(); + if (candidate == null || candidate.isEmpty()) { + throw new IllegalArgumentException( + "Routing rule has a null or empty candidate key."); + } + Object value = entry.getValue(); + // String.valueOf(null) would silently become the literal pattern "null" (and + // non-String values would coerce); reject both instead. + if (!(value instanceof String)) { + throw new IllegalArgumentException( + String.format( + "Routing rule for candidate '%s' must be a regex String, got %s.", + candidate, + value == null ? "null" : value.getClass().getSimpleName())); + } + this.rules.put( + candidate, Pattern.compile((String) value, Pattern.CASE_INSENSITIVE)); + } + } + } + + @Override + public RoutingDecision route(RoutingContext context) { + String text = context.lastUserMessage(); + if (text != null && !text.isEmpty()) { + for (Map.Entry entry : rules.entrySet()) { + if (entry.getValue().matcher(text).find()) { + if (!isCandidate(entry.getKey(), context)) { + throw new IllegalArgumentException( + "Routing rule selected non-candidate model '" + + entry.getKey() + + "'."); + } + return RoutingDecision.builder(entry.getKey()) + .reason("matched rule: " + entry.getValue().pattern()) + .build(); + } + } + } + return RoutingDecision.abstain(); + } + + private static boolean isCandidate(String model, RoutingContext context) { + for (RoutingCandidate candidate : context.getCandidates()) { + if (candidate.getName().equals(model)) { + return true; + } + } + return false; + } +} 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 new file mode 100644 index 000000000..0616ac78b --- /dev/null +++ b/api/src/main/java/org/apache/flink/agents/api/chat/model/routing/Strategies.java @@ -0,0 +1,62 @@ +/* + * 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.util.Collections; +import java.util.HashMap; +import java.util.Map; + +/** + * Factories for built-in routing strategies. Each returns a {@link RoutingStrategyDescriptor} + * (class name + args) rather than a live instance, so the strategy is plan-serializable. There is + * no magic-string strategy dispatch — the factory supplies the class name. + */ +public final class Strategies { + + private Strategies() {} + + /** + * Keyword/regex rules: a map of {@code candidateModel -> regex}. The first candidate whose + * regex matches the most recent user message wins; otherwise the strategy abstains (router + * falls back to its default model). + * + *

Rules are evaluated in the map's iteration order, so when precedence between overlapping + * patterns matters, pass a {@link java.util.LinkedHashMap} — {@code Map.of(...)} iteration + * order is unspecified. + */ + public static RoutingStrategyDescriptor rules(Map rules) { + Map args = new HashMap<>(); + args.put("rules", rules == null ? Collections.emptyMap() : rules); + return new RoutingStrategyDescriptor(RuleBasedRoutingStrategy.class.getName(), args); + } + + /** + * A custom strategy referenced by class. The class must be a {@link RoutingStrategy} with + * either a {@code (Map)} constructor or a no-arg constructor. This is the + * deployable shape for custom routing. + */ + public static RoutingStrategyDescriptor of(Class clazz) { + return new RoutingStrategyDescriptor(clazz.getName(), Collections.emptyMap()); + } + + /** A custom strategy referenced by class name plus construction arguments. */ + public static RoutingStrategyDescriptor of(String clazz, Map args) { + return new RoutingStrategyDescriptor(clazz, args); + } +} diff --git a/api/src/main/java/org/apache/flink/agents/api/context/RunnerContext.java b/api/src/main/java/org/apache/flink/agents/api/context/RunnerContext.java index c3e5d19ba..fb1f19305 100644 --- a/api/src/main/java/org/apache/flink/agents/api/context/RunnerContext.java +++ b/api/src/main/java/org/apache/flink/agents/api/context/RunnerContext.java @@ -96,6 +96,20 @@ public interface RunnerContext { */ Resource getResource(String name, ResourceType type) throws Exception; + /** + * Checks whether a resource of the given name and type is registered, without creating it. + * + *

Used, for example, to detect whether a requested chat-model name is actually a {@code + * MODEL_ROUTER}. Default returns {@code false} for contexts without resource support. + * + * @param name the resource name + * @param type the resource type + * @return true if such a resource is registered + */ + default boolean hasResource(String name, ResourceType type) { + return false; + } + /** * Gets the configuration for Flink Agents. * 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 new file mode 100644 index 000000000..5c9099dd6 --- /dev/null +++ b/api/src/main/java/org/apache/flink/agents/api/event/ModelRoutingEvent.java @@ -0,0 +1,201 @@ +/* + * 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.event; + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonIgnore; +import com.fasterxml.jackson.annotation.JsonProperty; +import org.apache.flink.agents.api.Event; + +import javax.annotation.Nullable; + +import java.util.ArrayList; +import java.util.Collection; +import java.util.HashMap; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.UUID; + +/** + * Observability-only event recording a model-routing decision. + * + *

Emitted by {@code ChatModelAction} after a {@code MODEL_ROUTER} selects a concrete chat model + * and before the selected model is invoked. It is a record for logging, tracing, and evaluation: it + * has no built-in consumer and does not drive dispatch. Removing every listener of + * this event does not change which model runs. + */ +public class ModelRoutingEvent extends Event { + + public static final String EVENT_TYPE = "_model_routing_event"; + + /** Where the final selected model came from. */ + public static final String SOURCE_STRATEGY = "strategy"; + + public static final String SOURCE_DEFAULT = "default"; + public static final String SOURCE_FALLBACK = "fallback"; + + public ModelRoutingEvent( + UUID requestId, + String router, + List candidates, + String selectedModel, + String decisionSource, + boolean fallbackEnabled, + @Nullable String reason, + @Nullable Double score, + @Nullable Map metadata, + @Nullable Double decisionMs) { + super(EVENT_TYPE); + setAttr("request_id", requestId); + setAttr("router", router); + setAttr("candidates", new ArrayList<>(candidates)); + setAttr("selected_model", selectedModel); + setAttr("decision_source", decisionSource); + setAttr("fallback_enabled", fallbackEnabled); + if (reason != null) { + setAttr("reason", reason); + } + if (score != null) { + setAttr("score", score); + } + setAttr("metadata", mutableRoutingValue(metadata == null ? new HashMap<>() : metadata)); + if (decisionMs != null) { + setAttr("decision_ms", decisionMs); + } + } + + @JsonCreator + public ModelRoutingEvent( + @JsonProperty("id") UUID id, + @JsonProperty("attributes") Map attributes) { + super(id, EVENT_TYPE, normalizeAttributes(attributes)); + } + + /** Convert the {@code request_id} back to a {@link UUID} after JSON deserialization. */ + private static Map normalizeAttributes(Map attributes) { + Map normalized = + attributes == null ? new HashMap<>() : new HashMap<>(attributes); + Object rawId = normalized.get("request_id"); + if (rawId instanceof String) { + normalized.put("request_id", UUID.fromString((String) rawId)); + } + Object candidates = normalized.get("candidates"); + if (candidates instanceof Collection) { + normalized.put("candidates", mutableRoutingValue(candidates)); + } + Object metadata = normalized.get("metadata"); + if (metadata instanceof Map) { + normalized.put("metadata", mutableRoutingValue(metadata)); + } + return normalized; + } + + /** Reconstructs a typed ModelRoutingEvent from a base Event. */ + public static ModelRoutingEvent fromEvent(Event event) { + ModelRoutingEvent result = + new ModelRoutingEvent(event.getId(), new HashMap<>(event.getAttributes())); + if (event.hasSourceTimestamp()) { + result.setSourceTimestamp(event.getSourceTimestamp()); + } + return result; + } + + @JsonIgnore + public UUID getRequestId() { + Object val = getAttr("request_id"); + if (val instanceof String) { + return UUID.fromString((String) val); + } + return (UUID) val; + } + + @JsonIgnore + public String getRouter() { + return (String) getAttr("router"); + } + + @JsonIgnore + @SuppressWarnings("unchecked") + public List getCandidates() { + return (List) getAttr("candidates"); + } + + @JsonIgnore + public String getSelectedModel() { + return (String) getAttr("selected_model"); + } + + @JsonIgnore + public String getDecisionSource() { + return (String) getAttr("decision_source"); + } + + /** Whether the router was configured with fallback (not whether fallback happened). */ + @JsonIgnore + public boolean isFallbackEnabled() { + Object value = getAttr("fallback_enabled"); + return value instanceof Boolean && (Boolean) value; + } + + @JsonIgnore + @Nullable + public String getReason() { + return (String) getAttr("reason"); + } + + @JsonIgnore + @Nullable + public Double getScore() { + Object value = getAttr("score"); + return value instanceof Number ? ((Number) value).doubleValue() : null; + } + + @JsonIgnore + @SuppressWarnings("unchecked") + public Map getMetadata() { + Map metadata = (Map) getAttr("metadata"); + return metadata == null ? new HashMap<>() : metadata; + } + + /** Wall-clock time spent resolving the routing decision, in milliseconds (if recorded). */ + @JsonIgnore + @Nullable + public Double getDecisionMs() { + Object value = getAttr("decision_ms"); + return value instanceof Number ? ((Number) value).doubleValue() : null; + } + + private static Object mutableRoutingValue(Object value) { + if (value instanceof Map) { + Map copy = new LinkedHashMap<>(); + ((Map) value) + .forEach((key, nestedValue) -> copy.put(key, mutableRoutingValue(nestedValue))); + return copy; + } + if (value instanceof Collection) { + List copy = new ArrayList<>(); + for (Object nestedValue : (Collection) value) { + copy.add(mutableRoutingValue(nestedValue)); + } + return copy; + } + return value; + } +} diff --git a/api/src/main/java/org/apache/flink/agents/api/resource/ResourceType.java b/api/src/main/java/org/apache/flink/agents/api/resource/ResourceType.java index 79e4ab52e..dc8b866b0 100644 --- a/api/src/main/java/org/apache/flink/agents/api/resource/ResourceType.java +++ b/api/src/main/java/org/apache/flink/agents/api/resource/ResourceType.java @@ -32,7 +32,8 @@ public enum ResourceType { PROMPT("prompt"), TOOL("tool"), MCP_SERVER("mcp_server"), - SKILLS("skills"); + SKILLS("skills"), + MODEL_ROUTER("model_router"); private final String value; diff --git a/api/src/test/java/org/apache/flink/agents/api/EventTypeTest.java b/api/src/test/java/org/apache/flink/agents/api/EventTypeTest.java index 8aca070de..de5d23e30 100644 --- a/api/src/test/java/org/apache/flink/agents/api/EventTypeTest.java +++ b/api/src/test/java/org/apache/flink/agents/api/EventTypeTest.java @@ -26,6 +26,7 @@ import org.apache.flink.agents.api.event.LongTermGetEvent; import org.apache.flink.agents.api.event.LongTermSearchEvent; import org.apache.flink.agents.api.event.LongTermUpdateEvent; +import org.apache.flink.agents.api.event.ModelRoutingEvent; import org.apache.flink.agents.api.event.SensoryReadEvent; import org.apache.flink.agents.api.event.SensoryWriteEvent; import org.apache.flink.agents.api.event.ShortTermReadEvent; @@ -58,6 +59,7 @@ void allConstantsProvidesAnUnmodifiableNameToValueMap() { Map.entry( "ContextRetrievalResponseEvent", ContextRetrievalResponseEvent.EVENT_TYPE), + Map.entry("ModelRoutingEvent", ModelRoutingEvent.EVENT_TYPE), Map.entry("ShortTermWriteEvent", ShortTermWriteEvent.EVENT_TYPE), Map.entry("ShortTermReadEvent", ShortTermReadEvent.EVENT_TYPE), Map.entry("SensoryWriteEvent", SensoryWriteEvent.EVENT_TYPE), diff --git a/api/src/test/java/org/apache/flink/agents/api/chat/model/routing/RoutingResourceValidationTest.java b/api/src/test/java/org/apache/flink/agents/api/chat/model/routing/RoutingResourceValidationTest.java new file mode 100644 index 000000000..05d289075 --- /dev/null +++ b/api/src/test/java/org/apache/flink/agents/api/chat/model/routing/RoutingResourceValidationTest.java @@ -0,0 +1,115 @@ +/* + * 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.AgentBuilder; +import org.apache.flink.agents.api.AgentsExecutionEnvironment; +import org.apache.flink.agents.api.agents.Agent; +import org.apache.flink.agents.api.resource.ResourceDescriptor; +import org.apache.flink.agents.api.resource.ResourceType; +import org.junit.jupiter.api.Test; + +import java.util.Map; + +import static org.junit.jupiter.api.Assertions.assertThrows; + +/** + * Chat models and model routers share the chat request namespace, so registering one name as both + * must fail at the registration call site (better failure locality than the {@code AgentPlan} + * backstop check). + */ +class RoutingResourceValidationTest { + + private static ResourceDescriptor descriptor() { + return new ResourceDescriptor("some.Clazz", Map.of()); + } + + @Test + void agentRejectsRouterNameAlreadyUsedByChatModel() { + Agent agent = new Agent(); + agent.addResource("shared", ResourceType.CHAT_MODEL, descriptor()); + IllegalArgumentException e = + assertThrows( + IllegalArgumentException.class, + () -> agent.addResource("shared", ResourceType.MODEL_ROUTER, descriptor())); + org.junit.jupiter.api.Assertions.assertTrue(e.getMessage().contains("CHAT_MODEL")); + } + + @Test + void agentRejectsChatModelNameAlreadyUsedByRouter() { + Agent agent = new Agent(); + agent.addResource("shared", ResourceType.MODEL_ROUTER, descriptor()); + assertThrows( + IllegalArgumentException.class, + () -> agent.addResource("shared", ResourceType.CHAT_MODEL, descriptor())); + } + + @Test + void distinctNamesAndUnrelatedTypesAreAllowed() { + Agent agent = new Agent(); + agent.addResource("router", ResourceType.MODEL_ROUTER, descriptor()); + agent.addResource("small", ResourceType.CHAT_MODEL, descriptor()); + // same name across unrelated types is not a clash + agent.addResource("router", ResourceType.PROMPT, descriptor()); + } + + @Test + void executionEnvironmentRejectsChatModelRouterNameClash() { + AgentsExecutionEnvironment env = stubEnvironment(); + env.addResource("shared", ResourceType.CHAT_MODEL, descriptor()); + assertThrows( + IllegalArgumentException.class, + () -> env.addResource("shared", ResourceType.MODEL_ROUTER, descriptor())); + } + + /** Minimal stub — resource registration lives in the abstract base class under test. */ + private static AgentsExecutionEnvironment stubEnvironment() { + return new AgentsExecutionEnvironment() { + @Override + public org.apache.flink.agents.api.configuration.Configuration getConfig() { + return null; + } + + @Override + public AgentBuilder fromList(java.util.List input) { + return null; + } + + @Override + public AgentBuilder fromDataStream( + org.apache.flink.streaming.api.datastream.DataStream input, + org.apache.flink.api.java.functions.KeySelector keySelector) { + return null; + } + + @Override + public AgentBuilder fromTable( + org.apache.flink.table.api.Table input, + org.apache.flink.api.java.functions.KeySelector keySelector) { + return null; + } + + @Override + public void execute() {} + + @Override + public void execute(String jobName) {} + }; + } +} 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 new file mode 100644 index 000000000..fd5a2fe30 --- /dev/null +++ b/api/src/test/java/org/apache/flink/agents/api/chat/model/routing/RoutingTest.java @@ -0,0 +1,337 @@ +/* + * 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 com.fasterxml.jackson.databind.ObjectMapper; +import org.apache.flink.agents.api.chat.messages.ChatMessage; +import org.apache.flink.agents.api.chat.messages.MessageRole; +import org.apache.flink.agents.api.event.ModelRoutingEvent; +import org.apache.flink.agents.api.resource.ResourceDescriptor; +import org.junit.jupiter.api.Test; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.UUID; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** Unit tests for the model-routing API layer (strategy, router, decision). */ +class RoutingTest { + + private static final UUID REQUEST_ID = UUID.fromString("00000000-0000-0000-0000-000000000001"); + + private static RoutingContext ctx(String userMessage) { + return new RoutingContext( + REQUEST_ID, + "router", + List.of(new ChatMessage(MessageRole.USER, userMessage)), + Map.of(), + List.of(new RoutingCandidate("small"), new RoutingCandidate("big"))); + } + + @Test + void ruleMatchSelectsCandidate() throws Exception { + RuleBasedRoutingStrategy strategy = + new RuleBasedRoutingStrategy(Map.of("rules", Map.of("big", "\\b(code|sql)\\b"))); + RoutingDecision decision = strategy.route(ctx("please write some SQL for me")); + assertFalse(decision.isAbstain()); + assertEquals("big", decision.getSelectedModel()); + } + + @Test + void ruleMatchesLatestUserMessageNotFirst() throws Exception { + // Multi-turn: turn 1 asked for SQL, but the current question is small talk. The rule + // strategy must route on the most recent user message, not the oldest. + RuleBasedRoutingStrategy strategy = + new RuleBasedRoutingStrategy(Map.of("rules", Map.of("big", "\\b(code|sql)\\b"))); + RoutingContext multiTurn = + new RoutingContext( + REQUEST_ID, + "router", + List.of( + new ChatMessage(MessageRole.USER, "please write some SQL for me"), + new ChatMessage(MessageRole.ASSISTANT, "SELECT 1;"), + new ChatMessage(MessageRole.USER, "thanks, how is the weather?")), + Map.of(), + List.of(new RoutingCandidate("small"), new RoutingCandidate("big"))); + assertTrue(strategy.route(multiTurn).isAbstain()); + assertEquals("thanks, how is the weather?", multiTurn.lastUserMessage()); + assertEquals("please write some SQL for me", multiTurn.firstUserMessage()); + } + + @Test + void ruleNoMatchAbstains() throws Exception { + RuleBasedRoutingStrategy strategy = + new RuleBasedRoutingStrategy(Map.of("rules", Map.of("big", "\\b(code|sql)\\b"))); + RoutingDecision decision = strategy.route(ctx("hello, how are you?")); + assertTrue(decision.isAbstain()); + assertNull(decision.getSelectedModel()); + } + + @Test + void modelRouterBuildsAndRoutes() throws Exception { + ResourceDescriptor descriptor = + ModelRouter.of("small", "big") + .strategy(Strategies.rules(Map.of("big", "\\b(code|sql)\\b"))) + .defaultModel("small") + .fallback(true) + .build(); + + ModelRouter router = new ModelRouter(descriptor, null); + + assertEquals(List.of("small", "big"), router.getCandidateNames()); + assertEquals("small", router.getDefaultModel().orElse(null)); + assertTrue(router.isFallbackEnabled()); + assertTrue(router.isCandidate("big")); + assertFalse(router.isCandidate("unknown")); + + RoutingContext context = + new RoutingContext( + REQUEST_ID, + "router", + List.of(new ChatMessage(MessageRole.USER, "write code")), + Map.of(), + router.getCandidates()); + assertEquals("big", router.route(context).getSelectedModel()); + assertEquals(REQUEST_ID, context.getRequestId()); + } + + @Test + void candidateDescriptionsReachStrategies() throws Exception { + // Descriptions declared on the router flow into the RoutingContext candidates, which is + // how semantic strategies (and future framework-managed LLM routing) learn what each + // candidate is for. + ModelRouter router = + new ModelRouter( + ModelRouter.of("small", "big") + .describe("small", "fast and cheap; chit-chat") + .describe("big", "strong; code and SQL") + .strategy(Strategies.rules(Map.of())) + .defaultModel("small") + .build(), + null); + assertEquals("fast and cheap; chit-chat", router.getCandidates().get(0).getDescription()); + assertEquals("strong; code and SQL", router.getCandidates().get(1).getDescription()); + } + + @Test + void describeRejectsUnknownCandidate() { + assertThrows( + IllegalArgumentException.class, + () -> ModelRouter.of("small", "big").describe("huge", "does not exist")); + } + + @Test + void defaultModelMustBeCandidate() { + assertThrows( + IllegalArgumentException.class, + () -> + new ModelRouter( + ModelRouter.of("small", "big") + .strategy(Strategies.rules(Map.of())) + .defaultModel("huge") + .build(), + null)); + } + + @Test + void modelRoutingEventRequestIdSurvivesStringForm() { + // Simulate an EventLog JSON round-trip where request_id came back as a String. + UUID id = UUID.fromString("00000000-0000-0000-0000-0000000000ab"); + HashMap attrs = new HashMap<>(); + attrs.put("request_id", id.toString()); + attrs.put("router", "router"); + attrs.put("candidates", List.of("small", "big")); + attrs.put("selected_model", "big"); + attrs.put("decision_source", "strategy"); + ModelRoutingEvent event = new ModelRoutingEvent(id, attrs); + assertEquals(id, event.getRequestId()); + assertEquals("big", event.getSelectedModel()); + } + + @Test + @SuppressWarnings("unchecked") + void modelRoutingEventCarriesMetadata() { + UUID id = UUID.fromString("00000000-0000-0000-0000-0000000000cd"); + ModelRoutingEvent event = + new ModelRoutingEvent( + id, + "router", + List.of("small", "big"), + "big", + "strategy", + true, + "matched sql", + 0.9, + Map.of("signals", List.of("sql")), + 1.5); + assertTrue(event.isFallbackEnabled()); + // metadata survives a JSON-shaped round-trip and is deeply mutable (serialization-safe). + ModelRoutingEvent restored = + new ModelRoutingEvent(id, new HashMap<>(event.getAttributes())); + assertEquals(List.of("sql"), restored.getMetadata().get("signals")); + restored.getMetadata().put("copy_check", true); + ((List) restored.getMetadata().get("signals")).add("probe"); + assertEquals(2, ((List) restored.getMetadata().get("signals")).size()); + } + + @Test + void builderRequiresStrategy() { + assertThrows( + IllegalStateException.class, + () -> ModelRouter.of("small", "big").defaultModel("small").build()); + } + + @Test + void candidateRejectsEmptyName() { + assertThrows(IllegalArgumentException.class, () -> new RoutingCandidate("")); + } + + @Test + void decisionRejectsEmptyModel() { + assertThrows(IllegalArgumentException.class, () -> RoutingDecision.of("")); + } + + @Test + void decisionJsonRoundTrips() throws Exception { + ObjectMapper mapper = new ObjectMapper(); + RoutingDecision original = + RoutingDecision.builder("big").reason("matched code").score(0.82).build(); + String json = mapper.writeValueAsString(original); + // wire shape is snake_case, matching the ModelRoutingEvent attributes + assertTrue(json.contains("\"selected_model\""), json); + RoutingDecision restored = mapper.readValue(json, RoutingDecision.class); + assertEquals("big", restored.getSelectedModel()); + assertFalse(restored.isAbstain()); + assertEquals("matched code", restored.getReason()); + assertEquals(0.82, restored.getScore()); + } + + @Test + void decisionRejectsInvalidStates() { + // abstain must not carry a model; non-abstain must carry one — on every construction + // path, including the JSON one used by durable replay. + assertThrows( + IllegalArgumentException.class, + () -> new RoutingDecision("big", true, null, null, null, null)); + assertThrows( + IllegalArgumentException.class, + () -> new RoutingDecision(null, false, null, null, null, null)); + } + + @Test + void decisionMsSurvivesJsonReplay() throws Exception { + // decision_ms is stamped inside the durable call; a replayed (deserialized) decision + // must report the original latency. + ObjectMapper mapper = new ObjectMapper(); + RoutingDecision timed = RoutingDecision.of("big").withDecisionMs(12.5); + RoutingDecision replayed = + mapper.readValue(mapper.writeValueAsString(timed), RoutingDecision.class); + assertEquals(12.5, replayed.getDecisionMs()); + } + + @Test + void abstainJsonRoundTrips() throws Exception { + ObjectMapper mapper = new ObjectMapper(); + RoutingDecision restored = + mapper.readValue( + mapper.writeValueAsString(RoutingDecision.abstain()), + RoutingDecision.class); + assertTrue(restored.isAbstain()); + assertNull(restored.getSelectedModel()); + } + + @Test + void builderRejectsRuleKeyThatIsNotACandidate() { + // A typo'd rule key must fail at the registration call site, not per record at runtime. + assertThrows( + IllegalArgumentException.class, + () -> + ModelRouter.of("small", "big") + .strategy(Strategies.rules(Map.of("bg", "\\bsql\\b"))) + .defaultModel("small") + .build()); + } + + @Test + void ruleStrategyRejectsNullRuleValue() { + Map rules = new HashMap<>(); + rules.put("big", null); + // String.valueOf(null) would otherwise compile the literal pattern "null". + assertThrows( + IllegalArgumentException.class, + () -> new RuleBasedRoutingStrategy(Map.of("rules", rules))); + } + + @Test + void ruleStrategyRejectsNonStringRuleValue() { + assertThrows( + IllegalArgumentException.class, + () -> new RuleBasedRoutingStrategy(Map.of("rules", Map.of("big", 42)))); + } + + @Test + void routingContextMessagesAreDeepCopied() { + ChatMessage original = new ChatMessage(MessageRole.USER, "original prompt"); + RoutingContext ctx = + new RoutingContext( + UUID.randomUUID(), "router", List.of(original), Map.of(), List.of()); + // A strategy mutating what it sees must not rewrite the message actually sent. + ctx.getMessages().get(0).setContent("REWRITTEN BY STRATEGY"); + assertEquals("original prompt", original.getContent()); + } + + @Test + void routingContextToolCallsAreDeepCopiedToo() { + // ChatMessage's constructor stores toolCalls by reference; the context must copy them. + List> toolCalls = new ArrayList<>(); + Map call = new HashMap<>(); + call.put("name", "originalTool"); + toolCalls.add(call); + ChatMessage original = new ChatMessage(MessageRole.ASSISTANT, "", toolCalls); + RoutingContext ctx = + new RoutingContext( + UUID.randomUUID(), "router", List.of(original), Map.of(), List.of()); + ctx.getMessages().get(0).getToolCalls().get(0).put("name", "HIJACKED"); + ctx.getMessages().get(0).getToolCalls().clear(); + assertEquals(1, original.getToolCalls().size()); + assertEquals("originalTool", original.getToolCalls().get(0).get("name")); + } + + @Test + void routingContextToleratesNullToolCallsFromJsonSetter() { + // The constructor defaults toolCalls to an empty list, but Jackson's setToolCalls stores + // null as-is — a message deserialized from JSON with "tool_calls": null carries null. + ChatMessage fromJson = new ChatMessage(MessageRole.USER, "hello"); + fromJson.setToolCalls(null); + RoutingContext ctx = + new RoutingContext( + UUID.randomUUID(), "router", List.of(fromJson), Map.of(), List.of()); + assertEquals(1, ctx.getMessages().size()); + assertEquals("hello", ctx.getMessages().get(0).getContent()); + // The copy re-normalizes through the constructor, so strategies see an empty list. + assertEquals(0, ctx.getMessages().get(0).getToolCalls().size()); + } +} diff --git a/examples/src/main/java/org/apache/flink/agents/examples/ModelRoutingExample.java b/examples/src/main/java/org/apache/flink/agents/examples/ModelRoutingExample.java new file mode 100644 index 000000000..c1716fc9a --- /dev/null +++ b/examples/src/main/java/org/apache/flink/agents/examples/ModelRoutingExample.java @@ -0,0 +1,105 @@ +/* + * 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; + +import org.apache.flink.agents.api.AgentsExecutionEnvironment; +import org.apache.flink.agents.api.agents.AgentExecutionOptions; +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.resource.ResourceDescriptor; +import org.apache.flink.agents.api.resource.ResourceName; +import org.apache.flink.agents.api.resource.ResourceType; +import org.apache.flink.agents.examples.agents.CustomTypesAndResources; +import org.apache.flink.agents.examples.agents.ModelRoutingAgent; +import org.apache.flink.streaming.api.datastream.DataStream; +import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment; + +import java.util.Map; + +/** + * Java example demonstrating rule-based in-chat model routing. + * + *

A stream of requests is processed by {@link ModelRoutingAgent}, which sends each to a {@code + * MODEL_ROUTER}. The router's rule strategy sends coding/SQL/analysis requests to a strong model + * ({@code big}) and everything else to a small model ({@code small}); when no rule matches it + * abstains and the router uses its default model. The selected model call is a first-class chat in + * the EventLog (tokens attributed to that model), and the decision itself is recorded as a {@code + * ModelRoutingEvent}. + * + *

Model names are illustrative — adjust them to models available on your Ollama server. + */ +public class ModelRoutingExample { + + private static ResourceDescriptor ollamaModel(String model) { + return ResourceDescriptor.Builder.newBuilder(ResourceName.ChatModel.OLLAMA_SETUP) + .addInitialArgument("connection", "ollamaChatModelConnection") + .addInitialArgument("model", model) + .build(); + } + + /** Runs the example pipeline. */ + public static void main(String[] args) throws Exception { + StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment(); + env.setParallelism(1); + AgentsExecutionEnvironment agentsEnv = + AgentsExecutionEnvironment.getExecutionEnvironment(env); + + // limit async request to avoid overwhelming ollama server + agentsEnv.getConfig().set(AgentExecutionOptions.NUM_ASYNC_THREADS, 2); + + // Ollama connection shared by both candidate models. + agentsEnv.addResource( + "ollamaChatModelConnection", + ResourceType.CHAT_MODEL_CONNECTION, + CustomTypesAndResources.OLLAMA_SERVER_DESCRIPTOR); + + // Two candidate chat models: a small (cheap) and a big (strong) one. + agentsEnv + .addResource("small", ResourceType.CHAT_MODEL, ollamaModel("qwen3:1.7b")) + .addResource("big", ResourceType.CHAT_MODEL, ollamaModel("qwen3:8b")); + + // A router over the two models: send code/SQL/analysis to "big", otherwise abstain -> + // "small". + agentsEnv.addResource( + "router", + ResourceType.MODEL_ROUTER, + ModelRouter.of("small", "big") + .strategy( + Strategies.rules( + Map.of("big", "\\b(code|sql|program|analyze|prove)\\b"))) + .defaultModel("small") + .fallback(true) + .build()); + + // A small stream of requests: an easy one (-> small) and a coding one (-> 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) + .apply(new ModelRoutingAgent()) + .toDataStream(); + + resultStream.print(); + + agentsEnv.execute("Model Routing Example Job"); + } +} diff --git a/examples/src/main/java/org/apache/flink/agents/examples/agents/ModelRoutingAgent.java b/examples/src/main/java/org/apache/flink/agents/examples/agents/ModelRoutingAgent.java new file mode 100644 index 000000000..701715f33 --- /dev/null +++ b/examples/src/main/java/org/apache/flink/agents/examples/agents/ModelRoutingAgent.java @@ -0,0 +1,59 @@ +/* + * 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.agents; + +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.chat.messages.ChatMessage; +import org.apache.flink.agents.api.chat.messages.MessageRole; +import org.apache.flink.agents.api.context.RunnerContext; +import org.apache.flink.agents.api.event.ChatRequestEvent; +import org.apache.flink.agents.api.event.ChatResponseEvent; + +import java.util.Collections; + +/** + * An agent that routes each request to a chat model chosen by a {@code MODEL_ROUTER}. + * + *

The agent simply sends a {@link ChatRequestEvent} naming the router {@code "router"}; the + * framework detects it as a router, selects a concrete chat model, emits a {@code + * ModelRoutingEvent} for observability, and runs the normal chat path against the selected model. + * The router, its candidate chat models, and the connection are registered in {@code + * ModelRoutingExample}. + */ +public class ModelRoutingAgent extends Agent { + + /** Send each input to the router, which selects the concrete model. */ + @Action(EventType.InputEvent) + public static void processInput(InputEvent event, RunnerContext ctx) { + ctx.sendEvent( + new ChatRequestEvent( + "router", + Collections.singletonList( + new ChatMessage(MessageRole.USER, (String) event.getInput())))); + } + + /** Emit the model's answer as output. */ + @Action(EventType.ChatResponseEvent) + public static void processChatResponse(ChatResponseEvent event, RunnerContext ctx) { + ctx.sendEvent(new OutputEvent(event.getResponse().getContent())); + } +} diff --git a/examples/src/main/java/org/apache/flink/agents/examples/openai/OpenAiModelRoutingExample.java b/examples/src/main/java/org/apache/flink/agents/examples/openai/OpenAiModelRoutingExample.java new file mode 100644 index 000000000..a16217e78 --- /dev/null +++ b/examples/src/main/java/org/apache/flink/agents/examples/openai/OpenAiModelRoutingExample.java @@ -0,0 +1,110 @@ +/* + * 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 org.apache.flink.agents.api.AgentsExecutionEnvironment; +import org.apache.flink.agents.api.agents.AgentExecutionOptions; +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.resource.ResourceDescriptor; +import org.apache.flink.agents.api.resource.ResourceName; +import org.apache.flink.agents.api.resource.ResourceType; +import org.apache.flink.agents.examples.agents.ModelRoutingAgent; +import org.apache.flink.streaming.api.datastream.DataStream; +import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment; + +import java.util.Map; + +/** + * Rule-based in-chat model routing against OpenAI, runnable as a local Flink job. + * + *

Routes coding/SQL/analysis requests to a strong model ({@code gpt-4o}) and everything else to + * a small model ({@code gpt-4o-mini}); on no match it abstains and the router uses its default + * ({@code small}). The selected model call is a first-class chat in the EventLog (tokens attributed + * to that model) and the decision is recorded as a {@code ModelRoutingEvent}. + * + *

Run with {@code OPENAI_API_KEY} set in the environment. + * + *

Lives in the {@code openai} subpackage (not directly under {@code examples}) so the + * submit-examples E2E job, which auto-submits every top-level example against a keyless local + * cluster, does not pick it up. + */ +public class OpenAiModelRoutingExample { + + 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 models. + agentsEnv.addResource( + "openaiConnection", + ResourceType.CHAT_MODEL_CONNECTION, + ResourceDescriptor.Builder.newBuilder( + ResourceName.ChatModel.OPENAI_COMPLETIONS_CONNECTION) + .addInitialArgument("api_key", apiKey) + .build()); + + // Two candidate chat models: a small (cheap) and a big (strong) one. + agentsEnv + .addResource("small", ResourceType.CHAT_MODEL, openAiModel("gpt-4o-mini")) + .addResource("big", ResourceType.CHAT_MODEL, openAiModel("gpt-4o")); + + // Router: code/SQL/analysis -> "big"; otherwise abstain -> default "small". + agentsEnv.addResource( + "router", + ResourceType.MODEL_ROUTER, + ModelRouter.of("small", "big") + .strategy( + Strategies.rules( + Map.of("big", "\\b(code|sql|program|analyze|prove)\\b"))) + .defaultModel("small") + .build()); + + 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) + .apply(new ModelRoutingAgent()) + .toDataStream(); + + resultStream.print(); + + agentsEnv.execute("OpenAI Model Routing Example"); + } +} 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 df6801132..c865c72ce 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 @@ -658,11 +658,39 @@ private ResourceProvider createResourceProvider( /** Adds a resource provider to the resourceProviders map. */ private void addResourceProvider(ResourceProvider provider) { + checkNoRouterModelNameClash(provider); resourceProviders .computeIfAbsent(provider.getType(), k -> new HashMap<>()) .put(provider.getName(), provider); } + /** + * A name must not be registered as both a {@link ResourceType#CHAT_MODEL} and a {@link + * ResourceType#MODEL_ROUTER}: an agent references either by putting it in {@code + * ChatRequestEvent.model}, so a name that is both would resolve ambiguously. Fail clearly at + * plan-construction time rather than at request time. + */ + private void checkNoRouterModelNameClash(ResourceProvider provider) { + final ResourceType conflicting; + if (provider.getType() == ResourceType.MODEL_ROUTER) { + conflicting = ResourceType.CHAT_MODEL; + } else if (provider.getType() == ResourceType.CHAT_MODEL) { + conflicting = ResourceType.MODEL_ROUTER; + } else { + return; + } + Map existing = resourceProviders.get(conflicting); + if (existing != null && existing.containsKey(provider.getName())) { + throw new IllegalArgumentException( + String.format( + "Resource name '%s' is registered as both %s and %s; a name must not be" + + " both a chat model and a model router.", + provider.getName(), + ResourceType.CHAT_MODEL, + ResourceType.MODEL_ROUTER)); + } + } + /** * Promote an api-layer {@link org.apache.flink.agents.api.function.Function} descriptor to its * plan-layer twin. Java parameter type strings are resolved to {@link Class} here; Python diff --git a/plan/src/main/java/org/apache/flink/agents/plan/actions/ChatModelAction.java b/plan/src/main/java/org/apache/flink/agents/plan/actions/ChatModelAction.java index 13f89d876..57c191c75 100644 --- a/plan/src/main/java/org/apache/flink/agents/plan/actions/ChatModelAction.java +++ b/plan/src/main/java/org/apache/flink/agents/plan/actions/ChatModelAction.java @@ -26,21 +26,18 @@ import org.apache.flink.agents.api.chat.messages.ChatMessage; import org.apache.flink.agents.api.chat.messages.MessageRole; import org.apache.flink.agents.api.chat.model.BaseChatModelSetup; -import org.apache.flink.agents.api.chat.model.python.PythonChatModelSetup; -import org.apache.flink.agents.api.context.DurableCallable; import org.apache.flink.agents.api.context.MemoryObject; import org.apache.flink.agents.api.context.RunnerContext; 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.event.ToolRequestEvent; import org.apache.flink.agents.api.event.ToolResponseEvent; import org.apache.flink.agents.api.metrics.FlinkAgentsMetricGroup; -import org.apache.flink.agents.api.resource.ResourceType; import org.apache.flink.agents.api.skills.Skills; import org.apache.flink.agents.api.tools.ToolResponse; import org.apache.flink.agents.api.trace.ExecutionReporter; import org.apache.flink.agents.api.trace.ExecutionReporters; -import org.apache.flink.agents.api.trace.LLMExecutionMetadataKeys; import org.apache.flink.agents.plan.JavaFunction; import org.apache.flink.api.java.typeutils.RowTypeInfo; import org.apache.flink.types.Row; @@ -52,9 +49,41 @@ import java.util.*; import static org.apache.flink.agents.api.agents.Agent.STRUCTURED_OUTPUT; -import static org.apache.flink.agents.plan.actions.Utils.supportAsync; -/** Built-in action for processing chat request and tool call result. */ +/** + * Built-in action for processing chat request and tool call result. + * + *

Model routing overview

+ * + *

When a {@link ChatRequestEvent} names a {@code MODEL_ROUTER} instead of a chat model, this + * action layers five jobs on top of the normal chat path; each is localized to one place: + * + *

    + *
  1. Decide — {@link ModelRoutingResolver} runs the router's strategy and normalizes the + * result (abstain → default model; non-candidate → fail). + *
  2. Durably — the strategy runs inside a durable call ({@code "route:"}; + * per-request uniqueness comes from the store's (key, sequence, event, action) scoping, and + * the id must stay deterministic across recovery re-processing), so recovery replays the + * persisted decision instead of re-running a possibly non-deterministic strategy. This replay + * guarantee requires an action-state store to be configured ({@code actionStateStoreBackend}, + * see {@code AgentConfigOptions#ACTION_STATE_STORE_BACKEND}); without one — the default — the + * decision and the chat call re-execute together on recovery, which is self-consistent but + * re-derives the decision. + *
  3. Once per reasoning loop — the selected concrete model is saved in the tool-request + * context and reused by tool rounds with no re-routing; the routing metadata block is parked + * once in an initial-request-keyed context and attached only to the loop's final response. + *
  4. Fallback over retries — {@link ResolvedModelRoute#attemptOrder} tries the selected + * model first (with its full retry budget, durable id {@code "chat::"}), + * then remaining candidates in declaration order if fallback is enabled. + *
  5. Observably — a {@link ModelRoutingEvent} records the decision (and a second one any + * fallback outcome); {@link ResolvedModelRoute#buildResponseMetadata} supplies the {@code + * model_routing} extra args stamped on the final response; decision latency feeds {@code + * decision_ms} and the {@code routingDecisionLatencyMs} histogram. + *
+ * + *

A request naming a plain chat model takes the pre-routing path unchanged, including the legacy + * durable call id {@code "chat"}. + */ public class ChatModelAction { private static final Logger LOG = LoggerFactory.getLogger(ChatModelAction.class); @@ -62,6 +91,7 @@ public class ChatModelAction { private static final String TOOL_REQUEST_EVENT_CONTEXT = "_TOOL_REQUEST_EVENT_CONTEXT"; private static final String INITIAL_REQUEST_ID = "initialRequestId"; private static final String MODEL = "model"; + private static final String ROUTING_METADATA_CONTEXT = "_ROUTING_METADATA_CONTEXT"; private static final String OUTPUT_SCHEMA = "outputSchema"; private static final String PROMPT_ARGS = "prompt_args"; private static final String RETRY_STATS_CONTEXT = "_RETRY_STATS_CONTEXT"; @@ -289,7 +319,7 @@ static String cleanLlmResponse(String rawResponse) { } @SuppressWarnings("unchecked") - private static ChatMessage generateStructuredOutput(ChatMessage response, Object outputSchema) + static ChatMessage generateStructuredOutput(ChatMessage response, Object outputSchema) throws JsonProcessingException { String output = response.getContent(); output = cleanLlmResponse(output); @@ -330,16 +360,23 @@ public static void chat( @Nullable Object outputSchema, RunnerContext ctx) throws Exception { - BaseChatModelSetup chatModel = - (BaseChatModelSetup) ctx.getResource(model, ResourceType.CHAT_MODEL); - FlinkAgentsMetricGroup requestMetricGroup = ctx.getActionMetricGroup(); - - boolean chatAsync = ctx.getConfig().get(AgentExecutionOptions.CHAT_ASYNC); - - if ((chatModel instanceof PythonChatModelSetup) && !supportAsync()) { - chatAsync = false; - } + chat( + initialRequestId, + ResolvedModelRoute.direct(model), + messages, + promptArgs, + outputSchema, + ctx); + } + private static void chat( + UUID initialRequestId, + ResolvedModelRoute selection, + List messages, + Map promptArgs, + @Nullable Object outputSchema, + RunnerContext ctx) + throws Exception { Agent.ErrorHandlingStrategy strategy = ctx.getConfig().get(AgentExecutionOptions.ERROR_HANDLING_STRATEGY); int numRetries = 0; @@ -355,121 +392,191 @@ public static void chat( : 0; } - ChatMessage response = null; - int actualRetryCount = 0; - int totalWaitTimeSec = 0; - - DurableCallable callable = - new DurableCallable<>() { - @Override - public String getId() { - return "chat"; - } - - @Override - public Class getResultClass() { - return ChatMessage.class; - } - - @Override - public ChatMessage call() throws Exception { - return chatModel.chat(messages, promptArgs, Map.of()); - } - }; - Map llmMetadata = - chatModel.getModel() == null - ? Map.of() - : Map.of(LLMExecutionMetadataKeys.MODEL, chatModel.getModel()); - - for (int attempt = 0; attempt < numRetries + 1; attempt++) { + List triedModels = new ArrayList<>(); + Exception lastError = null; + for (String candidate : selection.attemptOrder()) { + triedModels.add(candidate); try { - ExecutionReporters.started( - ctx, ExecutionReporter.EntityTypes.LLM, model, llmMetadata); - try { - response = - chatAsync - ? ctx.durableExecuteAsync(callable) - : ctx.durableExecute(callable); - Objects.requireNonNull(response, "ChatModel returned a null response."); - } catch (Throwable modelError) { - throw reportFailedAndPropagate( - ctx, - ExecutionReporter.EntityTypes.LLM, - model, - llmMetadata, - modelError, - ExecutionReporter.ProblemCategories.MODEL_CALL_FAILED); - } - ExecutionReporters.succeeded( - ctx, ExecutionReporter.EntityTypes.LLM, model, llmMetadata); - recordChatTokenMetrics(chatModel, response, requestMetricGroup); - if (outputSchema != null && response.getToolCalls().isEmpty()) { - response = generateStructuredOutputWithReport(ctx, response, outputSchema); + ChatModelInvoker.ChatAttemptResult result = + ChatModelInvoker.chatWithRetries( + initialRequestId, + candidate, + selection.durableChatCallId(candidate), + messages, + promptArgs, + outputSchema, + ctx, + strategy, + numRetries, + retryWaitIntervalSec); + recordAttemptRetryStats( + ctx, + initialRequestId, + result.chatModel, + result.retryCount, + result.totalRetryWaitSec); + if (selection.isRouter) { + if (!result.model.equals(selection.selectedModel)) { + // The strategy's pick failed and another candidate answered; record the + // outcome in the event log, not just on the response. + ctx.sendEvent( + new ModelRoutingEvent( + initialRequestId, + selection.requestedModel, + selection.candidates, + result.model, + ModelRoutingEvent.SOURCE_FALLBACK, + selection.fallbackEnabled, + String.format( + "fallback after selected model '%s' failed", + selection.selectedModel), + null, + selection.metadata, + null)); + } } - } catch (Exception e) { - if (strategy == Agent.ErrorHandlingStrategy.IGNORE) { - LOG.warn( - "Chat request {} failed with error: {}, ignored.", initialRequestId, e); - return; - } else if (strategy == Agent.ErrorHandlingStrategy.RETRY) { - if (attempt == numRetries) { - throw e; + + // Routing metadata is observability-only and needed exactly once, on the final + // response. If this response starts (or continues) a tool loop, park the block + // in an initial-request-keyed context instead of stamping intermediate messages + // and copying it through every tool round. + Map routingMetadata = + selection.isRouter + ? selection.buildResponseMetadata(result.model, triedModels) + : null; + if (!Objects.requireNonNull(result.response).getToolCalls().isEmpty()) { + if (routingMetadata != null) { + saveRoutingMetadata( + ctx.getSensoryMemory(), initialRequestId, routingMetadata); } - actualRetryCount = attempt + 1; - int currentWaitSec = retryWaitIntervalSec * (1 << (actualRetryCount - 1)); - LOG.warn( - "Chat request {} failed with error: {}, retrying {} / {}, waiting {} s.", + handleToolCalls( + result.response, initialRequestId, - e, - actualRetryCount, - numRetries, - currentWaitSec); - if (currentWaitSec > 0) { - Thread.sleep(currentWaitSec * 1000L); - totalWaitTimeSec += currentWaitSec; - } - continue; + result.model, + result.chatModel, + messages, + promptArgs, + outputSchema, + ctx); } else { - LOG.debug( - "Chat request {} failed, the input chat messages are {}.", - initialRequestId, - messages); - throw e; + if (routingMetadata == null) { + routingMetadata = + takeRoutingMetadata(ctx.getSensoryMemory(), initialRequestId); + } + if (routingMetadata != null) { + result.response.getExtraArgs().put("model_routing", routingMetadata); + } + Map retryStats = + getRetryStats(ctx.getSensoryMemory(), initialRequestId); + int totalRetryCount = retryStats.get(TOTAL_RETRY_COUNT).intValue(); + int totalRetryWaitSec = retryStats.get(TOTAL_RETRY_WAIT_SEC).intValue(); + + ctx.sendEvent( + new ChatResponseEvent( + initialRequestId, + result.response, + totalRetryCount, + totalRetryWaitSec)); } + return; + } catch (ChatModelInvoker.ChatAttemptFailed e) { + recordAttemptRetryStats( + ctx, initialRequestId, e.chatModel, e.retryCount, e.totalRetryWaitSec); + // Keep every candidate's failure: chain the previous error into the new one so + // exhaustion surfaces A's and B's errors as suppressed of C's, not just C's. + if (lastError != null && lastError != e.error) { + e.error.addSuppressed(lastError); + } + lastError = e.error; + LOG.debug( + "Chat request {} failed for model {} with error: {}. The input chat messages are {}.", + initialRequestId, + e.model, + e.error.toString(), + messages); } - break; } - if (actualRetryCount > 0) { - accumulateRetryStats( - ctx.getSensoryMemory(), initialRequestId, actualRetryCount, totalWaitTimeSec); + if (selection.isRouter && triedModels.size() > 1) { + LOG.warn( + "Chat request {} exhausted all candidates {} of router '{}'; last error: {}.", + initialRequestId, + triedModels, + selection.requestedModel, + lastError == null ? null : lastError.toString()); + } + // The reasoning loop is over; a routed loop that dies mid-way must not leak its + // parked metadata (matters under IGNORE, where the job keeps running). + takeRoutingMetadata(ctx.getSensoryMemory(), initialRequestId); + if (strategy == Agent.ErrorHandlingStrategy.IGNORE) { + LOG.warn( + "Chat request {} failed with error: {}, ignored.", initialRequestId, lastError); + return; } + throw Objects.requireNonNull(lastError); + } - if (!Objects.requireNonNull(response).getToolCalls().isEmpty()) { - handleToolCalls( - response, - initialRequestId, - model, - chatModel, - messages, - promptArgs, - outputSchema, - ctx); - } else { - Map retryStats = getRetryStats(ctx.getSensoryMemory(), initialRequestId); - int totalRetryCount = retryStats.get(TOTAL_RETRY_COUNT).intValue(); - int totalRetryWaitSec = retryStats.get(TOTAL_RETRY_WAIT_SEC).intValue(); + /** + * Compatibility note: retry metrics are recorded per attempt (including attempts on the failure + * path), where previously they were recorded once with cumulative totals on the final response. + * Totals over a completed request are unchanged; requests that ultimately fail now contribute + * their retry counts where they previously did not. + */ + private static void recordAttemptRetryStats( + RunnerContext ctx, + UUID initialRequestId, + BaseChatModelSetup chatModel, + int retryCount, + int retryWaitSec) + throws Exception { + if (retryCount <= 0) { + return; + } + accumulateRetryStats(ctx.getSensoryMemory(), initialRequestId, retryCount, retryWaitSec); + String metricModel = chatModel == null ? null : chatModel.getConnectionName(); + recordRetryMetrics( + ctx, + metricModel == null || metricModel.isEmpty() ? "unknown" : metricModel, + retryCount, + retryWaitSec); + } - recordRetryMetrics( - ctx, chatModel.getConnectionName(), totalRetryCount, totalRetryWaitSec); + /** + * Parks the routed request's {@code model_routing} block for the lifetime of its reasoning + * loop, keyed by the initial request id. Stored once when the loop starts; taken (removed) once + * when the final response is produced or the loop is abandoned. + */ + @SuppressWarnings("unchecked") + private static void saveRoutingMetadata( + MemoryObject sensoryMem, UUID initialRequestId, Map routing) + throws Exception { + Map context; + if (sensoryMem.isExist(ROUTING_METADATA_CONTEXT)) { + context = (Map) sensoryMem.get(ROUTING_METADATA_CONTEXT).getValue(); + } else { + context = new HashMap<>(); + } + context.put(initialRequestId, routing); + sensoryMem.set(ROUTING_METADATA_CONTEXT, context); + } - ctx.sendEvent( - new ChatResponseEvent( - initialRequestId, response, totalRetryCount, totalRetryWaitSec)); + @SuppressWarnings("unchecked") + @Nullable + private static Map takeRoutingMetadata( + MemoryObject sensoryMem, UUID initialRequestId) throws Exception { + if (!sensoryMem.isExist(ROUTING_METADATA_CONTEXT)) { + return null; + } + Map context = + (Map) sensoryMem.get(ROUTING_METADATA_CONTEXT).getValue(); + Map routing = (Map) context.remove(initialRequestId); + if (routing != null) { + sensoryMem.set(ROUTING_METADATA_CONTEXT, context); } + return routing; } - private static ChatMessage generateStructuredOutputWithReport( + static ChatMessage generateStructuredOutputWithReport( RunnerContext ctx, ChatMessage response, Object outputSchema) throws Exception { ExecutionReporters.started(ctx, ExecutionReporter.EntityTypes.PARSER, STRUCTURED_OUTPUT); try { @@ -490,9 +597,34 @@ private static ChatMessage generateStructuredOutputWithReport( private static void processChatRequest(ChatRequestEvent event, RunnerContext ctx) throws Exception { + ResolvedModelRoute selection; + try { + selection = + ModelRoutingResolver.resolve( + event.getId(), + event.getModel(), + event.getMessages(), + event.getPromptArgs(), + ctx); + } catch (Exception e) { + // A routing-strategy failure honors the same error-handling strategy as the chat + // call itself: under IGNORE the request is dropped with a warning instead of killing + // the job. (Retries are not applied to the decision; strategies that perform I/O are + // expected to absorb their own transient failures.) + if (ctx.getConfig().get(AgentExecutionOptions.ERROR_HANDLING_STRATEGY) + == Agent.ErrorHandlingStrategy.IGNORE) { + LOG.warn( + "Routing for chat request {} (model '{}') failed with error: {}, ignored.", + event.getId(), + event.getModel(), + e.toString()); + return; + } + throw e; + } chat( event.getId(), - event.getModel(), + selection, event.getMessages(), event.getPromptArgs(), event.getOutputSchema(), @@ -544,7 +676,16 @@ private static void processToolResponse(ToolResponseEvent event, RunnerContext c Collections.emptyList(), toolResponseMessages); - chat(initialRequestId, model, messages, promptArgs, outputSchema, ctx); + // Tool rounds reuse the already-selected concrete model (no re-routing). If the initial + // request was routed, its metadata block waits in ROUTING_METADATA_CONTEXT and is + // attached when this loop produces its final response. + chat( + initialRequestId, + ResolvedModelRoute.direct(model), + messages, + promptArgs, + outputSchema, + ctx); } /** @@ -574,7 +715,7 @@ public static void processChatRequestOrToolResponse(Event event, RunnerContext c * Reports a nested execution failure, then always throws the original failure. The Exception * return type exists so callers must {@code throw} the result and cannot fall through. */ - private static Exception reportFailedAndPropagate( + static Exception reportFailedAndPropagate( RunnerContext ctx, String entityType, String entityName, diff --git a/plan/src/main/java/org/apache/flink/agents/plan/actions/ChatModelInvoker.java b/plan/src/main/java/org/apache/flink/agents/plan/actions/ChatModelInvoker.java new file mode 100644 index 000000000..8dfacee4e --- /dev/null +++ b/plan/src/main/java/org/apache/flink/agents/plan/actions/ChatModelInvoker.java @@ -0,0 +1,208 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.flink.agents.plan.actions; + +import org.apache.flink.agents.api.agents.Agent; +import org.apache.flink.agents.api.agents.AgentExecutionOptions; +import org.apache.flink.agents.api.chat.messages.ChatMessage; +import org.apache.flink.agents.api.chat.model.BaseChatModelSetup; +import org.apache.flink.agents.api.chat.model.python.PythonChatModelSetup; +import org.apache.flink.agents.api.context.DurableCallable; +import org.apache.flink.agents.api.context.RunnerContext; +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.api.trace.LLMExecutionMetadataKeys; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import javax.annotation.Nullable; + +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.UUID; + +import static org.apache.flink.agents.plan.actions.Utils.supportAsync; + +/** + * Invokes one concrete chat model with the engine's durable-call and retry machinery. One call = + * one candidate attempt: success returns a {@link ChatAttemptResult}; failure (including an + * unresolvable model resource) surfaces as {@link ChatAttemptFailed} so the caller's fallback loop + * and error-handling strategy see every attempt uniformly. + */ +final class ChatModelInvoker { + + private static final Logger LOG = LoggerFactory.getLogger(ChatModelInvoker.class); + + private ChatModelInvoker() {} + + static final class ChatAttemptResult { + final String model; + final BaseChatModelSetup chatModel; + final ChatMessage response; + final int retryCount; + final int totalRetryWaitSec; + + ChatAttemptResult( + String model, + BaseChatModelSetup chatModel, + ChatMessage response, + int retryCount, + int totalRetryWaitSec) { + this.model = model; + this.chatModel = chatModel; + this.response = response; + this.retryCount = retryCount; + this.totalRetryWaitSec = totalRetryWaitSec; + } + } + + static final class ChatAttemptFailed extends Exception { + final String model; + final BaseChatModelSetup chatModel; + final Exception error; + final int retryCount; + final int totalRetryWaitSec; + + ChatAttemptFailed( + String model, + BaseChatModelSetup chatModel, + Exception error, + int retryCount, + int totalRetryWaitSec) { + super(error); + this.model = model; + this.chatModel = chatModel; + this.error = error; + this.retryCount = retryCount; + this.totalRetryWaitSec = totalRetryWaitSec; + } + } + + static ChatAttemptResult chatWithRetries( + UUID initialRequestId, + String model, + String durableCallId, + List messages, + Map promptArgs, + @Nullable Object outputSchema, + RunnerContext ctx, + Agent.ErrorHandlingStrategy strategy, + int numRetries, + int retryWaitIntervalSec) + throws ChatAttemptFailed, Exception { + BaseChatModelSetup chatModel; + try { + chatModel = (BaseChatModelSetup) ctx.getResource(model, ResourceType.CHAT_MODEL); + } catch (Exception e) { + // An unresolvable candidate (e.g. a typo in the router's candidate list) counts as + // that candidate failing, so the fallback loop and the error-handling strategy see + // it like any other attempt failure instead of it escaping chat() raw and discarding + // the previous candidate's real error. + throw new ChatAttemptFailed(model, null, e, 0, 0); + } + FlinkAgentsMetricGroup requestMetricGroup = ctx.getActionMetricGroup(); + + boolean chatAsync = ctx.getConfig().get(AgentExecutionOptions.CHAT_ASYNC); + + if ((chatModel instanceof PythonChatModelSetup) && !supportAsync()) { + chatAsync = false; + } + + int actualRetryCount = 0; + int totalWaitTimeSec = 0; + ChatMessage response; + + DurableCallable callable = + new DurableCallable<>() { + @Override + public String getId() { + return durableCallId; + } + + @Override + public Class getResultClass() { + return ChatMessage.class; + } + + @Override + public ChatMessage call() throws Exception { + return chatModel.chat(messages, promptArgs, Map.of()); + } + }; + Map llmMetadata = + chatModel.getModel() == null + ? Map.of() + : Map.of(LLMExecutionMetadataKeys.MODEL, chatModel.getModel()); + + for (int attempt = 0; attempt < numRetries + 1; attempt++) { + try { + ExecutionReporters.started( + ctx, ExecutionReporter.EntityTypes.LLM, model, llmMetadata); + try { + response = + chatAsync + ? ctx.durableExecuteAsync(callable) + : ctx.durableExecute(callable); + Objects.requireNonNull(response, "ChatModel returned a null response."); + } catch (Throwable modelError) { + throw ChatModelAction.reportFailedAndPropagate( + ctx, + ExecutionReporter.EntityTypes.LLM, + model, + llmMetadata, + modelError, + ExecutionReporter.ProblemCategories.MODEL_CALL_FAILED); + } + ExecutionReporters.succeeded( + ctx, ExecutionReporter.EntityTypes.LLM, model, llmMetadata); + ChatModelAction.recordChatTokenMetrics(chatModel, response, requestMetricGroup); + // only generate structured output for final response. + if (outputSchema != null && response.getToolCalls().isEmpty()) { + response = + ChatModelAction.generateStructuredOutputWithReport( + ctx, response, outputSchema); + } + return new ChatAttemptResult( + model, chatModel, response, actualRetryCount, totalWaitTimeSec); + } catch (Exception e) { + if (strategy == Agent.ErrorHandlingStrategy.RETRY && attempt < numRetries) { + actualRetryCount = attempt + 1; + int currentWaitSec = retryWaitIntervalSec * (1 << (actualRetryCount - 1)); + LOG.warn( + "Chat request {} failed with error: {}, retrying {} / {}, waiting {} s.", + initialRequestId, + e, + actualRetryCount, + numRetries, + currentWaitSec); + if (currentWaitSec > 0) { + Thread.sleep(currentWaitSec * 1000L); + totalWaitTimeSec += currentWaitSec; + } + continue; + } + throw new ChatAttemptFailed( + model, chatModel, e, actualRetryCount, totalWaitTimeSec); + } + } + throw new IllegalStateException("Unreachable chat retry state."); + } +} diff --git a/plan/src/main/java/org/apache/flink/agents/plan/actions/ModelRoutingResolver.java b/plan/src/main/java/org/apache/flink/agents/plan/actions/ModelRoutingResolver.java new file mode 100644 index 000000000..8ef0d346d --- /dev/null +++ b/plan/src/main/java/org/apache/flink/agents/plan/actions/ModelRoutingResolver.java @@ -0,0 +1,143 @@ +/* + * 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.chat.messages.ChatMessage; +import org.apache.flink.agents.api.chat.model.routing.ModelRouter; +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.context.DurableCallable; +import org.apache.flink.agents.api.context.RunnerContext; +import org.apache.flink.agents.api.event.ModelRoutingEvent; +import org.apache.flink.agents.api.metrics.FlinkAgentsMetricGroup; +import org.apache.flink.agents.api.resource.ResourceType; + +import java.util.List; +import java.util.Map; +import java.util.UUID; + +/** + * Resolves a chat request's target into a {@link ResolvedModelRoute}: if {@code model} names a + * {@link ModelRouter}, runs its strategy inside a durable {@code "route:"} call, emits the + * observability-only {@link ModelRoutingEvent}, and normalizes the decision; otherwise returns the + * direct route. + */ +final class ModelRoutingResolver { + + private ModelRoutingResolver() {} + + /** + * If {@code model} names a {@link ModelRouter}, run its strategy (as a durable {@code "route"} + * call so the decision replays deterministically on recovery), normalize the result (abstain -> + * default model, non-candidate -> fail clearly), emit an observability-only {@link + * ModelRoutingEvent}, and return the selected concrete model. Otherwise returns a direct + * selection. + * + *

Routing runs once for the initial chat request; tool-call rounds reuse the selected + * concrete model because it is saved in the tool-request context (see {@code + * ChatModelAction#handleToolCalls}), so this method is only reached with a router name on the + * initial request. + */ + static ResolvedModelRoute resolve( + UUID requestId, + String model, + List messages, + Map promptArgs, + RunnerContext ctx) + throws Exception { + if (!ctx.hasResource(model, ResourceType.MODEL_ROUTER)) { + return ResolvedModelRoute.direct(model); + } + ModelRouter router = (ModelRouter) ctx.getResource(model, ResourceType.MODEL_ROUTER); + RoutingContext routingContext = + new RoutingContext(requestId, model, messages, promptArgs, router.getCandidates()); + + DurableCallable routeCallable = + new DurableCallable<>() { + @Override + public String getId() { + // Deterministic across recovery re-processing: the durable store already + // scopes call results by (key, sequence number, event, action), so the id + // must NOT embed the request id — event ids are regenerated when Flink + // rolls back and re-processes, and a non-deterministic id turns every + // replay lookup into a miss (measured: 0/138 decisions replayed). + return "route:" + model; + } + + @Override + public Class getResultClass() { + return RoutingDecision.class; + } + + @Override + public RoutingDecision call() throws Exception { + // Timed inside the durable call so the latency is persisted with the + // decision: a replayed run reports the original strategy wall time. + long start = System.nanoTime(); + RoutingDecision decision = router.route(routingContext); + return decision.withDecisionMs((System.nanoTime() - start) / 1_000_000.0); + } + }; + + RoutingDecision decision = ctx.durableExecute(routeCallable); + Double decisionMs = decision.getDecisionMs(); + FlinkAgentsMetricGroup actionMetrics = ctx.getActionMetricGroup(); + if (actionMetrics != null && decisionMs != null) { + actionMetrics.getHistogram("routingDecisionLatencyMs").update(Math.round(decisionMs)); + } + + String selectedModel; + String decisionSource; + if (decision.isAbstain()) { + selectedModel = router.getDefaultModel().orElse(router.getCandidateNames().get(0)); + decisionSource = ModelRoutingEvent.SOURCE_DEFAULT; + } else { + selectedModel = decision.getSelectedModel(); + if (!router.isCandidate(selectedModel)) { + throw new IllegalStateException( + String.format( + "Routing strategy for router '%s' returned non-candidate model '%s'; candidates are %s.", + model, selectedModel, router.getCandidateNames())); + } + decisionSource = ModelRoutingEvent.SOURCE_STRATEGY; + } + + ctx.sendEvent( + new ModelRoutingEvent( + requestId, + model, + router.getCandidateNames(), + selectedModel, + decisionSource, + router.isFallbackEnabled(), + decision.getReason(), + decision.getScore(), + decision.getMetadata(), + decisionMs)); + return new ResolvedModelRoute( + model, + selectedModel, + router.getCandidateNames(), + true, + router.isFallbackEnabled(), + decisionSource, + decision.getReason(), + decision.getScore(), + decision.getMetadata()); + } +} diff --git a/plan/src/main/java/org/apache/flink/agents/plan/actions/ResolvedModelRoute.java b/plan/src/main/java/org/apache/flink/agents/plan/actions/ResolvedModelRoute.java new file mode 100644 index 000000000..96fb25382 --- /dev/null +++ b/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 candidates; + final boolean isRouter; + final boolean fallbackEnabled; + final String decisionSource; + @Nullable final String reason; + @Nullable final Double score; + final Map metadata; + + ResolvedModelRoute( + String requestedModel, + String selectedModel, + List candidates, + boolean isRouter, + boolean fallbackEnabled, + String decisionSource, + @Nullable String reason, + @Nullable Double score, + @Nullable Map 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 attemptOrder() { + List 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 buildResponseMetadata(String finalModel, List triedModels) { + boolean fallbackAttempted = !finalModel.equals(this.selectedModel); + List fallbackModelsTried = new ArrayList<>(); + for (int i = 1; i < triedModels.size(); i++) { + fallbackModelsTried.add(triedModels.get(i)); + } + Map routing = new LinkedHashMap<>(); + routing.put("router", this.requestedModel); + routing.put("selected_model", this.selectedModel); + routing.put("initial_selected_model", this.selectedModel); + routing.put("final_model", finalModel); + routing.put("candidates", new ArrayList<>(this.candidates)); + routing.put( + "decision_source", + fallbackAttempted ? ModelRoutingEvent.SOURCE_FALLBACK : this.decisionSource); + routing.put("fallback_enabled", this.fallbackEnabled); + routing.put("fallback_attempted", fallbackAttempted); + routing.put("fallback_models_tried", fallbackModelsTried); + routing.put("metadata", new LinkedHashMap<>(this.metadata)); + if (this.reason != null) { + routing.put("reason", this.reason); + } + if (this.score != null) { + routing.put("score", this.score); + } + return routing; + } +} 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 new file mode 100644 index 000000000..6afa94bd1 --- /dev/null +++ b/plan/src/test/java/org/apache/flink/agents/plan/AgentPlanRoutingBackstopTest.java @@ -0,0 +1,70 @@ +/* + * 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.agents.Agent; +import org.apache.flink.agents.api.resource.ResourceDescriptor; +import org.apache.flink.agents.api.resource.ResourceType; +import org.junit.jupiter.api.Test; + +import java.util.HashMap; +import java.util.Map; + +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * The chat-model/router name-clash has two per-call checks (agent and environment {@code + * addResource}), but each of those sees only its own registry. When the same name arrives from + * different registries — environment resources merged via {@code addResourcesIfAbsent}, + * agent resources added directly — only {@link AgentPlan}'s backstop check catches the clash. + */ +class AgentPlanRoutingBackstopTest { + + private static ResourceDescriptor descriptor(String clazz) { + return new ResourceDescriptor(clazz, Map.of()); + } + + @Test + void planBackstopRejectsClashArrivingFromDifferentRegistries() { + Agent agent = new Agent(); + agent.addResource("shared", ResourceType.MODEL_ROUTER, descriptor("some.Router")); + // Environment-level resources merge in via putIfAbsent, bypassing the per-call + // checks: neither addResource call ever saw both registrations. + agent.addResourcesIfAbsent( + Map.of( + ResourceType.CHAT_MODEL, + new HashMap<>(Map.of("shared", descriptor("some.ChatModel"))))); + + IllegalArgumentException e = + assertThrows(IllegalArgumentException.class, () -> new AgentPlan(agent)); + assertTrue(e.getMessage().contains("shared"), e.getMessage()); + } + + @Test + void planBackstopAllowsDistinctNamesAcrossRegistries() throws Exception { + Agent agent = new Agent(); + agent.addResource("router", ResourceType.MODEL_ROUTER, descriptor("some.Router")); + agent.addResourcesIfAbsent( + Map.of( + ResourceType.CHAT_MODEL, + new HashMap<>(Map.of("small", descriptor("some.ChatModel"))))); + new AgentPlan(agent); + } +} 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 new file mode 100644 index 000000000..46c9cc4b2 --- /dev/null +++ b/plan/src/test/java/org/apache/flink/agents/plan/actions/ChatModelActionRoutingTest.java @@ -0,0 +1,785 @@ +/* + * 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; +import org.apache.flink.agents.api.agents.Agent; +import org.apache.flink.agents.api.agents.AgentExecutionOptions; +import org.apache.flink.agents.api.chat.messages.ChatMessage; +import org.apache.flink.agents.api.chat.messages.MessageRole; +import org.apache.flink.agents.api.chat.model.BaseChatModelSetup; +import org.apache.flink.agents.api.chat.model.routing.ModelRouter; +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.configuration.ReadableConfiguration; +import org.apache.flink.agents.api.context.DurableCallable; +import org.apache.flink.agents.api.context.MemoryObject; +import org.apache.flink.agents.api.context.MemoryRef; +import org.apache.flink.agents.api.context.RunnerContext; +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.event.ToolRequestEvent; +import org.apache.flink.agents.api.event.ToolResponseEvent; +import org.apache.flink.agents.api.memory.BaseLongTermMemory; +import org.apache.flink.agents.api.metrics.FlinkAgentsMetricGroup; +import org.apache.flink.agents.api.resource.Resource; +import org.apache.flink.agents.api.resource.ResourceDescriptor; +import org.apache.flink.agents.api.resource.ResourceType; +import org.apache.flink.agents.api.tools.ToolResponse; +import org.apache.flink.agents.plan.AgentConfiguration; +import org.junit.jupiter.api.Test; + +import java.util.ArrayDeque; +import java.util.ArrayList; +import java.util.Collections; +import java.util.Deque; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** Integration tests for model routing inside {@link ChatModelAction}. */ +public class ChatModelActionRoutingTest { + + /** A strategy that returns a name that is not a candidate (to exercise the invalid path). */ + public static class SelectsUnknownStrategy implements RoutingStrategy { + public SelectsUnknownStrategy() {} + + @Override + public RoutingDecision route(RoutingContext context) { + return RoutingDecision.of("nonexistent"); + } + } + + /** + * A chat model returning scripted outcomes per call: a {@link ChatMessage} is returned, a + * {@link RuntimeException} is thrown. When the script is exhausted, returns a default assistant + * reply. + */ + static class FakeChatModel extends BaseChatModelSetup { + private final Deque outcomes = new ArrayDeque<>(); + + FakeChatModel(Object... outcomes) { + super(new ResourceDescriptor("fake", Map.of()), null); + Collections.addAll(this.outcomes, outcomes); + } + + @Override + public Map getParameters() { + return Map.of(); + } + + @Override + public ChatMessage chat( + List messages, + Map promptArgs, + Map modelParams) { + Object next = outcomes.isEmpty() ? null : outcomes.poll(); + if (next instanceof RuntimeException) { + throw (RuntimeException) next; + } + if (next instanceof ChatMessage) { + return (ChatMessage) next; + } + return new ChatMessage(MessageRole.ASSISTANT, "answer"); + } + } + + static class FakeRunnerContext implements RunnerContext { + final List sentEvents = new ArrayList<>(); + final List resolvedChatModels = new ArrayList<>(); + final List durableCallIds = new ArrayList<>(); + final Map models = new HashMap<>(); + final Set unresolvable = new HashSet<>(); + private final ModelRouter router; + private final MemoryObject sensoryMemory = new FakeMemoryObject(new HashMap<>()); + private final AgentConfiguration config = new AgentConfiguration(Map.of()); + + FakeRunnerContext(ModelRouter router) { + this.router = router; + } + + FakeRunnerContext register(String name, BaseChatModelSetup model) { + models.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 withErrorHandling(Agent.ErrorHandlingStrategy strategy) { + config.set(AgentExecutionOptions.ERROR_HANDLING_STRATEGY, strategy); + return this; + } + + FakeRunnerContext withRetryBudget(int maxRetries, int waitIntervalSec) { + config.set(AgentExecutionOptions.MAX_RETRIES, maxRetries); + config.set(AgentExecutionOptions.RETRY_WAIT_INTERVAL, waitIntervalSec); + return this; + } + + @Override + public boolean hasResource(String name, ResourceType type) { + return type == ResourceType.MODEL_ROUTER && "router".equals(name) && router != null; + } + + @Override + public Resource getResource(String name, ResourceType type) { + if (type == ResourceType.MODEL_ROUTER) { + return router; + } + if (type == ResourceType.CHAT_MODEL) { + if (unresolvable.contains(name)) { + throw new IllegalArgumentException("resource not found: " + name); + } + resolvedChatModels.add(name); + return models.getOrDefault(name, new FakeChatModel()); + } + throw new IllegalArgumentException("unexpected resource " + name + " " + type); + } + + @Override + public void sendEvent(Event event) { + sentEvents.add(event); + } + + @Override + public MemoryObject getSensoryMemory() { + return sensoryMemory; + } + + @Override + public MemoryObject getShortTermMemory() { + return null; + } + + @Override + public BaseLongTermMemory getLongTermMemory() { + return null; + } + + @Override + public FlinkAgentsMetricGroup getAgentMetricGroup() { + return null; + } + + @Override + public FlinkAgentsMetricGroup getActionMetricGroup() { + return null; + } + + @Override + public ReadableConfiguration getConfig() { + return config; + } + + @Override + public Map getActionConfig() { + return Map.of(); + } + + @Override + public Object getActionConfigValue(String key) { + return null; + } + + @Override + public T durableExecute(DurableCallable callable) throws Exception { + durableCallIds.add(callable.getId()); + return callable.call(); + } + + @Override + public T durableExecuteAsync(DurableCallable callable) throws Exception { + durableCallIds.add(callable.getId()); + return callable.call(); + } + + @Override + public void close() {} + + ModelRoutingEvent routingEvent() { + return sentEvents.stream() + .filter(e -> ModelRoutingEvent.EVENT_TYPE.equals(e.getType())) + .map(ModelRoutingEvent::fromEvent) + .findFirst() + .orElse(null); + } + + long routingEventCount() { + return sentEvents.stream() + .filter(e -> ModelRoutingEvent.EVENT_TYPE.equals(e.getType())) + .count(); + } + + ToolRequestEvent toolRequestEvent() { + return sentEvents.stream() + .filter(e -> ToolRequestEvent.EVENT_TYPE.equals(e.getType())) + .map(ToolRequestEvent::fromEvent) + .findFirst() + .orElse(null); + } + + ChatResponseEvent chatResponse() { + return sentEvents.stream() + .filter(e -> ChatResponseEvent.EVENT_TYPE.equals(e.getType())) + .map(ChatResponseEvent::fromEvent) + .findFirst() + .orElse(null); + } + + boolean hasChatResponse() { + return chatResponse() != null; + } + } + + private static ModelRouter router() throws Exception { + return new ModelRouter( + ModelRouter.of("small", "big") + .strategy(Strategies.rules(Map.of("big", "\\b(code|sql)\\b"))) + .defaultModel("small") + .build(), + null); + } + + @Test + void routesMatchingRequestToBigAndRunsNormalChat() throws Exception { + FakeRunnerContext ctx = new FakeRunnerContext(router()); + ChatModelAction.processChatRequestOrToolResponse( + new ChatRequestEvent( + "router", List.of(new ChatMessage(MessageRole.USER, "write some sql"))), + ctx); + + ModelRoutingEvent event = ctx.routingEvent(); + assertThat(event).isNotNull(); + assertThat(event.getRouter()).isEqualTo("router"); + assertThat(event.getSelectedModel()).isEqualTo("big"); + assertThat(event.getDecisionSource()).isEqualTo(ModelRoutingEvent.SOURCE_STRATEGY); + assertThat(event.getCandidates()).containsExactly("small", "big"); + // decision latency is stamped inside the durable route call + assertThat(event.getDecisionMs()).isNotNull(); + assertThat(event.isFallbackEnabled()).isFalse(); + // the selected concrete model was invoked via the normal chat path + assertThat(ctx.resolvedChatModels).containsExactly("big"); + assertThat(ctx.hasChatResponse()).isTrue(); + } + + @Test + void abstainRoutesToDefaultModel() throws Exception { + FakeRunnerContext ctx = new FakeRunnerContext(router()); + ChatModelAction.processChatRequestOrToolResponse( + new ChatRequestEvent( + "router", List.of(new ChatMessage(MessageRole.USER, "hello there"))), + ctx); + + ModelRoutingEvent event = ctx.routingEvent(); + assertThat(event).isNotNull(); + assertThat(event.getSelectedModel()).isEqualTo("small"); + assertThat(event.getDecisionSource()).isEqualTo(ModelRoutingEvent.SOURCE_DEFAULT); + assertThat(ctx.resolvedChatModels).containsExactly("small"); + } + + @Test + void invalidCandidateFailsClearly() throws Exception { + ModelRouter router = + new ModelRouter( + ModelRouter.of("small", "big") + .strategy(Strategies.of(SelectsUnknownStrategy.class)) + .defaultModel("small") + .build(), + null); + FakeRunnerContext ctx = new FakeRunnerContext(router); + assertThatThrownBy( + () -> + ChatModelAction.processChatRequestOrToolResponse( + new ChatRequestEvent( + "router", + List.of(new ChatMessage(MessageRole.USER, "hi"))), + ctx)) + .isInstanceOf(IllegalStateException.class) + .hasMessageContaining("non-candidate"); + } + + @Test + void abstainWithoutDefaultUsesFirstCandidate() throws Exception { + // No default model configured; on abstain the router falls back to the first candidate. + ModelRouter router = + new ModelRouter( + ModelRouter.of("small", "big") + .strategy(Strategies.rules(Map.of("big", "\\bsql\\b"))) + .build(), + null); + FakeRunnerContext ctx = new FakeRunnerContext(router); + ChatModelAction.processChatRequestOrToolResponse( + new ChatRequestEvent( + "router", List.of(new ChatMessage(MessageRole.USER, "hello there"))), + ctx); + + assertThat(ctx.routingEvent()).isNotNull(); + assertThat(ctx.routingEvent().getSelectedModel()).isEqualTo("small"); + assertThat(ctx.resolvedChatModels).containsExactly("small"); + } + + @Test + void nonRouterModelPassesThroughUnchanged() throws Exception { + FakeRunnerContext ctx = new FakeRunnerContext(null); + ChatModelAction.processChatRequestOrToolResponse( + new ChatRequestEvent( + "plainModel", List.of(new ChatMessage(MessageRole.USER, "write some sql"))), + ctx); + + assertThat(ctx.routingEvent()).isNull(); + assertThat(ctx.resolvedChatModels).containsExactly("plainModel"); + assertThat(ctx.hasChatResponse()).isTrue(); + } + + @Test + void routedRequestUsesRoutedDurableCallIds() throws Exception { + ModelRouter router = + new ModelRouter( + ModelRouter.of("small", "big") + .strategy(Strategies.rules(Map.of("big", "\\bsql\\b"))) + .defaultModel("small") + .build(), + null); + FakeRunnerContext ctx = + new FakeRunnerContext(router) + .register("big", new FakeChatModel(ChatMessage.assistant("ok"))); + ChatModelAction.processChatRequestOrToolResponse( + new ChatRequestEvent("router", List.of(ChatMessage.user("write sql"))), ctx); + // the decision and the chat attempt are distinct durable calls with routed ids + assertThat(ctx.durableCallIds).containsExactly("route:router", "chat:router:big"); + } + + @Test + void retryBudgetRunsBeforeFallback() throws Exception { + ModelRouter router = + new ModelRouter( + ModelRouter.of("small", "big") + .strategy(Strategies.rules(Map.of("big", "\\bsql\\b"))) + .defaultModel("small") + .fallback(true) + .build(), + null); + FakeRunnerContext ctx = + new FakeRunnerContext(router) + .withErrorHandling(Agent.ErrorHandlingStrategy.RETRY) + .withRetryBudget(1, 0) + .register( + "big", + new FakeChatModel( + new RuntimeException("transient"), + ChatMessage.assistant("recovered on retry"))) + .register( + "small", new FakeChatModel(ChatMessage.assistant("small answer"))); + + ChatModelAction.processChatRequestOrToolResponse( + new ChatRequestEvent("router", List.of(ChatMessage.user("write sql"))), ctx); + + // the selected model's retry budget is consumed BEFORE fallback: big's retry + // succeeds and small is never resolved — the ordering the class javadoc guarantees + assertThat(ctx.chatResponse().getResponse().getContent()).isEqualTo("recovered on retry"); + assertThat(ctx.resolvedChatModels).containsExactly("big"); + assertThat(ctx.routingEventCount()).isEqualTo(1L); + } + + @Test + void directModelKeepsLegacyDurableCallId() throws Exception { + FakeRunnerContext ctx = new FakeRunnerContext(null); + ChatModelAction.processChatRequestOrToolResponse( + new ChatRequestEvent("plain", List.of(ChatMessage.user("hi"))), ctx); + // a non-router request must keep the unchanged legacy durable chat-call id + assertThat(ctx.durableCallIds).containsExactly("chat"); + assertThat(ctx.routingEvent()).isNull(); + } + + @Test + @SuppressWarnings("unchecked") + void fallsBackToRemainingCandidateWhenSelectedModelFails() throws Exception { + ModelRouter router = + new ModelRouter( + ModelRouter.of("small", "big") + .strategy(Strategies.rules(Map.of("big", "\\bsql\\b"))) + .defaultModel("small") + .fallback(true) + .build(), + null); + FakeRunnerContext ctx = + new FakeRunnerContext(router) + .register("big", new FakeChatModel(new RuntimeException("big is down"))) + .register( + "small", new FakeChatModel(ChatMessage.assistant("ok from small"))); + + ChatModelAction.processChatRequestOrToolResponse( + new ChatRequestEvent("router", List.of(ChatMessage.user("write sql"))), ctx); + + // routed to big; big failed; fell back to small in declaration order + assertThat(ctx.resolvedChatModels).containsExactly("big", "small"); + // each stage has its own durable identity: the route decision, then one distinct + // chat call per candidate (the format recovery depends on, changed once already) + assertThat(ctx.durableCallIds) + .containsExactly("route:router", "chat:router:big", "chat:router:small"); + ChatResponseEvent response = ctx.chatResponse(); + assertThat(response).isNotNull(); + assertThat(response.getResponse().getContent()).isEqualTo("ok from small"); + Map routing = + (Map) response.getResponse().getExtraArgs().get("model_routing"); + assertThat(routing.get("final_model")).isEqualTo("small"); + assertThat(routing.get("decision_source")).isEqualTo(ModelRoutingEvent.SOURCE_FALLBACK); + List tried = (List) routing.get("fallback_models_tried"); + assertThat(tried).containsExactly("small"); + + // the fallback outcome is also in the event log: decision event + fallback event + assertThat(ctx.routingEventCount()).isEqualTo(2L); + ModelRoutingEvent fallbackEvent = + ctx.sentEvents.stream() + .filter(e -> ModelRoutingEvent.EVENT_TYPE.equals(e.getType())) + .map(ModelRoutingEvent::fromEvent) + .filter( + e -> + ModelRoutingEvent.SOURCE_FALLBACK.equals( + e.getDecisionSource())) + .findFirst() + .orElse(null); + assertThat(fallbackEvent).isNotNull(); + assertThat(fallbackEvent.getSelectedModel()).isEqualTo("small"); + assertThat(fallbackEvent.getReason()).contains("big"); + assertThat(fallbackEvent.isFallbackEnabled()).isTrue(); + } + + @Test + void fallbackExhaustedRethrows() throws Exception { + ModelRouter router = + new ModelRouter( + ModelRouter.of("small", "big") + .strategy(Strategies.rules(Map.of("big", "\\bsql\\b"))) + .defaultModel("small") + .fallback(true) + .build(), + null); + FakeRunnerContext ctx = + new FakeRunnerContext(router) + .register("big", new FakeChatModel(new RuntimeException("big-exploded"))) + .register( + "small", new FakeChatModel(new RuntimeException("small-exploded"))); + + // Distinct per-candidate markers: exhaustion must surface the LAST candidate's error + // with the earlier candidate's error chained as suppressed, not discarded. + assertThatThrownBy( + () -> + ChatModelAction.processChatRequestOrToolResponse( + new ChatRequestEvent( + "router", List.of(ChatMessage.user("write sql"))), + ctx)) + .isInstanceOf(RuntimeException.class) + .hasMessageContaining("small-exploded") + .satisfies( + t -> + assertThat(t.getSuppressed()) + .anySatisfy( + sup -> + assertThat(sup) + .hasMessageContaining( + "big-exploded"))); + assertThat(ctx.resolvedChatModels).containsExactly("big", "small"); + assertThat(ctx.hasChatResponse()).isFalse(); + } + + @Test + void unresolvableCandidateCountsAsFailedAttemptAndFallsBack() throws Exception { + // "big" is selected by the rule but its resource lookup fails (typo'd candidate); the + // failure must stay inside the fallback loop so "small" still answers. + ModelRouter router = + new ModelRouter( + ModelRouter.of("small", "big") + .strategy(Strategies.rules(Map.of("big", "\\bsql\\b"))) + .defaultModel("small") + .fallback(true) + .build(), + null); + FakeRunnerContext ctx = + new FakeRunnerContext(router) + .unresolvable("big") + .register("small", new FakeChatModel()); + + ChatModelAction.processChatRequestOrToolResponse( + new ChatRequestEvent("router", List.of(ChatMessage.user("write sql"))), ctx); + + assertThat(ctx.hasChatResponse()).isTrue(); + assertThat(ctx.resolvedChatModels).containsExactly("small"); + // the fallback outcome is recorded as a second routing event + long fallbackEvents = + ctx.sentEvents.stream() + .filter(e -> e instanceof ModelRoutingEvent) + .map(e -> (ModelRoutingEvent) e) + .filter( + e -> + ModelRoutingEvent.SOURCE_FALLBACK.equals( + e.getDecisionSource())) + .count(); + assertThat(fallbackEvents).isEqualTo(1); + } + + /** Strategy that always throws; must be public for reflective construction. */ + public static class ExplodingStrategy implements RoutingStrategy { + private static final long serialVersionUID = 1L; + + public ExplodingStrategy(Map args) {} + + @Override + public RoutingDecision route(RoutingContext context) { + throw new IllegalStateException("strategy exploded"); + } + } + + @Test + void strategyFailureIsIgnoredUnderIgnorePolicy() throws Exception { + ModelRouter router = + new ModelRouter( + ModelRouter.of("small", "big") + .strategy(Strategies.of(ExplodingStrategy.class)) + .defaultModel("small") + .build(), + null); + FakeRunnerContext ctx = + new FakeRunnerContext(router) + .withErrorHandling(Agent.ErrorHandlingStrategy.IGNORE) + .register("small", new FakeChatModel()); + + // Under IGNORE a strategy failure drops the request instead of killing the job. + ChatModelAction.processChatRequestOrToolResponse( + new ChatRequestEvent("router", List.of(ChatMessage.user("hello"))), ctx); + + assertThat(ctx.hasChatResponse()).isFalse(); + assertThat(ctx.resolvedChatModels).isEmpty(); + } + + @Test + void strategyFailurePropagatesUnderDefaultPolicy() throws Exception { + ModelRouter router = + new ModelRouter( + ModelRouter.of("small", "big") + .strategy(Strategies.of(ExplodingStrategy.class)) + .defaultModel("small") + .build(), + null); + FakeRunnerContext ctx = + new FakeRunnerContext(router).register("small", new FakeChatModel()); + + assertThatThrownBy( + () -> + ChatModelAction.processChatRequestOrToolResponse( + new ChatRequestEvent( + "router", List.of(ChatMessage.user("hello"))), + ctx)) + .isInstanceOf(IllegalStateException.class) + .hasMessageContaining("strategy exploded"); + } + + @Test + void routesOnceThenReusesSelectedModelAcrossToolRound() throws Exception { + ModelRouter router = + new ModelRouter( + ModelRouter.of("small", "big") + .strategy(Strategies.rules(Map.of("big", "\\bsql\\b"))) + .defaultModel("small") + .build(), + null); + List> toolCall = + List.of( + Map.of( + "id", + "call-1", + "type", + "function", + "function", + Map.of("name", "lookup", "arguments", Map.of()))); + ChatMessage intermediate = ChatMessage.assistant("", toolCall); + FakeRunnerContext ctx = + new FakeRunnerContext(router) + .register( + "big", + new FakeChatModel( + intermediate, ChatMessage.assistant("final answer"))); + + // initial routed request -> big -> tool call + ChatModelAction.processChatRequestOrToolResponse( + new ChatRequestEvent("router", List.of(ChatMessage.user("write sql"))), ctx); + assertThat(ctx.routingEventCount()).isEqualTo(1L); + ToolRequestEvent toolRequest = ctx.toolRequestEvent(); + assertThat(toolRequest).isNotNull(); + + // tool round: feed the tool response back + ChatModelAction.processChatRequestOrToolResponse( + new ToolResponseEvent( + toolRequest.getId(), + Map.of("call-1", ToolResponse.success("42")), + Map.of("call-1", true), + Map.of()), + ctx); + + // routing ran exactly once; the concrete model "big" was reused with no re-routing + assertThat(ctx.routingEventCount()).isEqualTo(1L); + assertThat(ctx.resolvedChatModels).containsExactly("big", "big"); + assertThat(ctx.chatResponse()).isNotNull(); + assertThat(ctx.chatResponse().getResponse().getContent()).isEqualTo("final answer"); + + // the routing metadata from the initial decision is carried onto the final response + @SuppressWarnings("unchecked") + Map routing = + (Map) + ctx.chatResponse().getResponse().getExtraArgs().get("model_routing"); + assertThat(routing).isNotNull(); + assertThat(routing.get("router")).isEqualTo("router"); + assertThat(routing.get("final_model")).isEqualTo("big"); + assertThat(routing.get("decision_source")).isEqualTo(ModelRoutingEvent.SOURCE_STRATEGY); + + // the intermediate tool-call message (which lives in the conversation history for the + // whole loop) is NOT stamped with observability metadata + assertThat(intermediate.getExtraArgs()).doesNotContainKey("model_routing"); + + // the parked metadata context was created for the loop and consumed by the final + // response (no leak) — asserted strictly so this fails if the context is never used + assertThat(ctx.getSensoryMemory().isExist("_ROUTING_METADATA_CONTEXT")).isTrue(); + Map parked = + (Map) ctx.getSensoryMemory().get("_ROUTING_METADATA_CONTEXT").getValue(); + assertThat(parked).isEmpty(); + } + + @Test + void routedLoopCleansParkedMetadataWhenToolRoundFailsUnderIgnore() throws Exception { + ModelRouter router = + new ModelRouter( + ModelRouter.of("small", "big") + .strategy(Strategies.rules(Map.of("big", "\\bsql\\b"))) + .defaultModel("small") + .build(), + null); + List> toolCall = + List.of( + Map.of( + "id", + "call-1", + "type", + "function", + "function", + Map.of("name", "lookup", "arguments", Map.of()))); + FakeRunnerContext ctx = + new FakeRunnerContext(router) + .withErrorHandling(Agent.ErrorHandlingStrategy.IGNORE) + .register( + "big", + new FakeChatModel( + ChatMessage.assistant("", toolCall), + new RuntimeException("tool round exploded"))); + + ChatModelAction.processChatRequestOrToolResponse( + new ChatRequestEvent("router", List.of(ChatMessage.user("write sql"))), ctx); + ToolRequestEvent toolRequest = ctx.toolRequestEvent(); + assertThat(toolRequest).isNotNull(); + // the routed round parked its metadata for the loop + Map parkedMidLoop = + (Map) ctx.getSensoryMemory().get("_ROUTING_METADATA_CONTEXT").getValue(); + assertThat(parkedMidLoop).hasSize(1); + + // the tool round's chat call fails and IGNORE drops the request... + ChatModelAction.processChatRequestOrToolResponse( + new ToolResponseEvent( + toolRequest.getId(), + Map.of("call-1", ToolResponse.success("42")), + Map.of("call-1", true), + Map.of()), + ctx); + assertThat(ctx.chatResponse()).isNull(); + + // ...and the abandoned loop's parked metadata was cleaned up, not leaked + Map parkedAfter = + (Map) ctx.getSensoryMemory().get("_ROUTING_METADATA_CONTEXT").getValue(); + assertThat(parkedAfter).isEmpty(); + } + + static class FakeMemoryObject implements MemoryObject { + private final Map values; + private final Object value; + + FakeMemoryObject(Map values) { + this(values, null); + } + + FakeMemoryObject(Map values, Object value) { + this.values = values; + this.value = value; + } + + @Override + public MemoryObject get(String path) { + return new FakeMemoryObject(values, values.get(path)); + } + + @Override + public MemoryObject get(MemoryRef ref) { + return get(ref.getPath()); + } + + @Override + public MemoryRef set(String path, Object value) { + values.put(path, value); + return null; + } + + @Override + public MemoryObject newObject(String path, boolean overwrite) { + throw new UnsupportedOperationException(); + } + + @Override + public boolean isExist(String path) { + return values.containsKey(path); + } + + @Override + public List getFieldNames() { + return new ArrayList<>(values.keySet()); + } + + @Override + public Map getFields() { + return Collections.unmodifiableMap(values); + } + + @Override + public Object getValue() { + return value; + } + + @Override + public boolean isNestedObject() { + return value == null; + } + } +} diff --git a/python/flink_agents/api/resource.py b/python/flink_agents/api/resource.py index 7187d3d6a..11542c5de 100644 --- a/python/flink_agents/api/resource.py +++ b/python/flink_agents/api/resource.py @@ -32,7 +32,7 @@ class ResourceType(Enum): """Type enum of resource. Currently, support chat_model, chat_model_server, tool, embedding_model, - vector_store, prompt, mcp_server, skills. + vector_store, prompt, mcp_server, skills, model_router. """ CHAT_MODEL = "chat_model" @@ -44,6 +44,11 @@ class ResourceType(Enum): PROMPT = "prompt" MCP_SERVER = "mcp_server" SKILLS = "skills" + # Java-side in-chat model routing (FLIP; full Python routing is a follow-up). + # Present so a Java plan containing a MODEL_ROUTER resource deserializes on the + # Python side: mixed jobs (Java router + Python actions) must not fail at + # operator open with a ValidationError. + MODEL_ROUTER = "model_router" class Resource(BaseModel, ABC): 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 f78b75750..5ed457474 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 @@ -35,6 +35,7 @@ from flink_agents.api.function import ( PythonFunction as ApiPythonFunction, ) +from flink_agents.api.resource import ResourceDescriptor, ResourceType from flink_agents.api.runner_context import RunnerContext from flink_agents.plan.agent_plan import AgentPlan from flink_agents.plan.configuration import AgentConfiguration @@ -44,6 +45,7 @@ from flink_agents.plan.function import ( PythonFunction as PlanPythonFunction, ) +from flink_agents.plan.resource_provider import JavaResourceProvider # python/flink_agents/plan/tests/test_*.py -> repo root is parents[4]. _REPO_ROOT = Path(__file__).resolve().parents[4] @@ -407,3 +409,34 @@ def test_python_preserves_conf_data_types_and_event_ordering() -> None: "k_str": "v1", } assert list(restored.actions) == ["first", "second"] + + +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). + + Python routing execution is a follow-up, but the Python side must still + deserialize such plans: a mixed job (Java router + any Python action) parses the + whole resource map against ResourceType, so a missing enum member fails the job + at operator open with a ValidationError. + """ + provider = JavaResourceProvider( + name="router", + type=ResourceType.MODEL_ROUTER, + descriptor=ResourceDescriptor( + target_module="java", + target_clazz=( + "org.apache.flink.agents.api.chat.model.routing.ModelRouter" + ), + arguments={"candidates": ["small", "big"]}, + ), + ) + plan = AgentPlan( + actions={}, + resource_providers={ResourceType.MODEL_ROUTER: {"router": provider}}, + ) + restored = AgentPlan.model_validate_json(_plan_dump_json(plan)) + assert ResourceType.MODEL_ROUTER in restored.resource_providers + assert isinstance( + restored.resource_providers[ResourceType.MODEL_ROUTER]["router"], + JavaResourceProvider, + ) diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/ResourceCache.java b/runtime/src/main/java/org/apache/flink/agents/runtime/ResourceCache.java index 01b920373..e3584ac26 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/ResourceCache.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/ResourceCache.java @@ -89,6 +89,24 @@ public ResourceContextImpl getResourceContext() { return resourceContext; } + /** + * Checks whether a resource of the given name and type is available, without creating it. + * Covers both registered providers and resources inserted directly into the cache via {@link + * #put} (which have no provider). + * + * @param name the resource name + * @param type the resource type + * @return true if such a resource has a registered provider or is already cached + */ + public boolean hasResource(String name, ResourceType type) { + Map cached = cache.get(type); + if (cached != null && cached.containsKey(name)) { + return true; + } + Map providers = resourceProviders.get(type); + return providers != null && providers.containsKey(name); + } + /** * Resolves a resource by name and type, creating it from its provider if not cached. * diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/context/RunnerContextImpl.java b/runtime/src/main/java/org/apache/flink/agents/runtime/context/RunnerContextImpl.java index 634ed2ae4..cbfe18636 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/context/RunnerContextImpl.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/context/RunnerContextImpl.java @@ -493,6 +493,11 @@ public Resource getResource(String name, ResourceType type) throws Exception { return resource; } + @Override + public boolean hasResource(String name, ResourceType type) { + return resourceCache != null && resourceCache.hasResource(name, type); + } + @Override public ReadableConfiguration getConfig() { return agentPlan.getConfig(); diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/ResourceCacheTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/ResourceCacheTest.java index 0d517e1cf..b0e5a685b 100644 --- a/runtime/src/test/java/org/apache/flink/agents/runtime/ResourceCacheTest.java +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/ResourceCacheTest.java @@ -218,6 +218,29 @@ public void testGetResourceNotFound() throws Exception { .hasMessageContaining("Resource not found: non-existent"); } + @Test + public void testHasResourceCoversProvidersAndCachedOnlyEntries() throws Exception { + TestAgentWithResources agent = new TestAgentWithResources(); + AgentPlan agentPlan = new AgentPlan(agent); + ResourceCache cache = new ResourceCache(agentPlan.getResourceProviders()); + + // registered provider, resource not yet created + assertThat(cache.hasResource("myTool", ResourceType.TOOL)).isTrue(); + assertThat(cache.hasResource("non-existent", ResourceType.TOOL)).isFalse(); + + // a resource inserted directly into the cache has no provider but must still be visible + Resource cachedOnly = + new Resource(new ResourceDescriptor("cachedOnly", Map.of()), null) { + @Override + public ResourceType getResourceType() { + return ResourceType.CHAT_MODEL; + } + }; + assertThat(cache.hasResource("cachedOnly", ResourceType.CHAT_MODEL)).isFalse(); + cache.put("cachedOnly", ResourceType.CHAT_MODEL, cachedOnly); + assertThat(cache.hasResource("cachedOnly", ResourceType.CHAT_MODEL)).isTrue(); + } + @Test public void testGetResourceFromResourceProvider() throws Exception { TestAgentWithResources agent = new TestAgentWithResources(); diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/condition/ConditionExpressionCompilerTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/condition/ConditionExpressionCompilerTest.java index 49fc4f8eb..494edb49d 100644 --- a/runtime/src/test/java/org/apache/flink/agents/runtime/condition/ConditionExpressionCompilerTest.java +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/condition/ConditionExpressionCompilerTest.java @@ -21,6 +21,7 @@ import dev.cel.runtime.CelEvaluationException; import dev.cel.runtime.CelRuntime; import org.apache.flink.agents.api.EventType; +import org.apache.flink.agents.api.event.ModelRoutingEvent; import org.apache.flink.agents.plan.condition.ConditionExpressionValidator; import org.apache.flink.agents.plan.condition.TriggerCondition; import org.apache.flink.agents.plan.condition.TriggerCondition.ExpressionCondition; @@ -78,6 +79,18 @@ void resolvesKnownEventTypeConstants() throws CelEvaluationException { assertThat(program.eval(activation)).isEqualTo(true); } + @Test + void resolvesModelRoutingEventConstant() throws CelEvaluationException { + CelRuntime.Program program = + compileValidated("type == EventType.ModelRoutingEvent").program(); + Map activation = new HashMap<>(); + activation.put("type", ModelRoutingEvent.EVENT_TYPE); + activation.put("attributes", new HashMap()); + activation.put("EventType", EventType.allConstants()); + + assertThat(program.eval(activation)).isEqualTo(true); + } + @Test void supportsDynamicAttributeTypes() throws CelEvaluationException { CelRuntime.Program program = compileValidated("score > 0").program();