weiqingy commented on code in PR #964:
URL: https://github.com/apache/flink-agents/pull/964#discussion_r3793315732
##########
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");
Review Comment:
Java gets `MODEL_ROUTER` here, but the Python enum doesn't have it
(`python/flink_agents/api/resource.py:38-46`), and no Python files changed in
this PR. The Java plan is serialized over to Python and parsed against
`Dict[ResourceType, ...]` (`python/flink_agents/plan/agent_plan.py:77`), so a
Java agent using a router plus any Python action would fail at operator open
with a pydantic `ValidationError`. Full Python routing is a follow-up per the
description. Would adding just this one enum member be enough to keep mixed
jobs starting, or is there a reason to hold it for the follow-up?
##########
plan/src/main/java/org/apache/flink/agents/plan/AgentPlan.java:
##########
@@ -617,11 +617,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) {
Review Comment:
`checkNoRouterModelNameClash` doesn't have a test yet.
`RoutingResourceValidationTest.java:44-79` exercises the two `addResource` call
sites, but each of those only sees its own map. Registering the same name in
both places — `env.addResource("smart", CHAT_MODEL, ...)` and
`agent.addResource("smart", MODEL_ROUTER, ...)` — passes both checks, and
`Agent.addResourcesIfAbsent` (`Agent.java:104`) merges them without one, so
this is what actually catches it. That's the case `Agent.java:139` calls out as
the backstop. Should that path get a case of its own?
##########
plan/src/main/java/org/apache/flink/agents/plan/actions/ChatModelAction.java:
##########
@@ -352,106 +390,220 @@ public static void chat(
: 0;
}
- ChatMessage response = null;
- int actualRetryCount = 0;
- int totalWaitTimeSec = 0;
-
- DurableCallable<ChatMessage> callable =
- new DurableCallable<>() {
- @Override
- public String getId() {
- return "chat";
- }
-
- @Override
- public Class<ChatMessage> getResultClass() {
- return ChatMessage.class;
- }
-
- @Override
- public ChatMessage call() throws Exception {
- return chatModel.chat(messages, promptArgs, Map.of());
- }
- };
-
- for (int attempt = 0; attempt < numRetries + 1; attempt++) {
+ List<String> triedModels = new ArrayList<>();
+ Exception lastError = null;
+ for (String candidate : selection.attemptOrder()) {
+ triedModels.add(candidate);
try {
- response =
- chatAsync
- ? ctx.durableExecuteAsync(callable)
- : ctx.durableExecute(callable);
- recordChatTokenMetrics(chatModel, response,
requestMetricGroup);
- // only generate structured output for final response.
- if (outputSchema != null && response.getToolCalls().isEmpty())
{
- response = generateStructuredOutput(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));
+ }
}
- break;
- } 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<String, Object> 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;
- }
+ 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<String, Long> 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);
}
}
- 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<String, Long> 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
Review Comment:
The description still says the opposite of this — line 58 reads "Existing
agents are unaffected." Three things do change for a plain-chat agent, all only
when `error-handling-strategy` isn't the default: with `retry`, retry metrics
now count per attempt including failures, and the retry WARN moved to
`ChatModelInvoker`'s logger; with `ignore`, a request naming a missing chat
model is warned and dropped instead of failing the job. Default `FAIL` really
is unaffected, so the blast radius is small. `AGENTS.md` asks for compatibility
impact in the description — is line 58 still the message you want, or should it
name those three?
##########
api/src/main/java/org/apache/flink/agents/api/chat/model/routing/RoutingContext.java:
##########
@@ -0,0 +1,152 @@
+/*
+ * 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.
+ *
+ * <p>v1 exposes the request id, the request messages, prompt args, and the
router's candidates
+ * (name + description). It intentionally does <b>not</b> expose a
chat-invocation API, so a
+ * strategy cannot make a hidden synchronous model call; observable
LLM-as-router is a
+ * framework-managed follow-up.
+ */
+public final class RoutingContext {
+
+ private final UUID requestId;
+ private final String router;
+ private final List<ChatMessage> messages;
+ private final Map<String, Object> promptArgs;
+ private final List<RoutingCandidate> candidates;
+
+ public RoutingContext(
+ UUID requestId,
+ String router,
+ List<ChatMessage> messages,
+ Map<String, Object> promptArgs,
+ List<RoutingCandidate> 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 =
Review Comment:
nit: the javadoc calls this a read-only view (`:32`), but each copy only
goes one level — `promptArgs` here, and in `deepCopy` both `extraArgs` (via
`ChatMessage`'s constructor, `ChatMessage.java:72`) and each tool-call map
(`:90`). Nested values stay shared with what gets sent
(`ChatModelAction.java:604-610`). No in-tree strategy mutates one, so this is
about the shape the API locks in. Is the top-level copy deliberate, or would
deep-copying these fit better?
##########
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:<router>"} 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.
+ *
+ * <p>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<ChatMessage> messages,
+ Map<String, Object> 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<RoutingDecision> 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;
Review Comment:
nit: the description hasn't caught up with this line — body line 26 still
describes the id as `"route:<requestId>:<router>"`, the form this commit
removed. Nothing is broken, since this line and the javadoc at `:36` are
already right. It's just that durability is the claim people read closest here,
so worth syncing line 26 too?
##########
plan/src/test/java/org/apache/flink/agents/plan/actions/ChatModelActionRoutingTest.java:
##########
@@ -0,0 +1,725 @@
+/*
+ * 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<Object> outcomes = new ArrayDeque<>();
+
+ FakeChatModel(Object... outcomes) {
+ super(new ResourceDescriptor("fake", Map.of()), null);
+ Collections.addAll(this.outcomes, outcomes);
+ }
+
+ @Override
+ public Map<String, Object> getParameters() {
+ return Map.of();
+ }
+
+ @Override
+ public ChatMessage chat(
+ List<ChatMessage> messages,
+ Map<String, Object> promptArgs,
+ Map<String, Object> 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<Event> sentEvents = new ArrayList<>();
+ final List<String> resolvedChatModels = new ArrayList<>();
+ final List<String> durableCallIds = new ArrayList<>();
+ final Map<String, BaseChatModelSetup> models = new HashMap<>();
+ final Set<String> 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;
+ }
+
+ @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<String, Object> getActionConfig() {
+ return Map.of();
+ }
+
+ @Override
+ public Object getActionConfigValue(String key) {
+ return null;
+ }
+
+ @Override
+ public <T> T durableExecute(DurableCallable<T> callable) throws
Exception {
+ durableCallIds.add(callable.getId());
+ return callable.call();
+ }
+
+ @Override
+ public <T> T durableExecuteAsync(DurableCallable<T> 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 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");
Review Comment:
This is the only place `durableCallIds` gets asserted, and it pins the
legacy `"chat"` id. The routed ones, `route:<router>` and
`chat:<router>:<candidate>`, go through the same list (`:207`, `:213`) but
nothing checks them — including
`fallsBackToRemainingCandidateWhenSelectedModelFails` (`:367`), where the two
candidates' journal entries need to stay distinct. Since `6ef088fae` had to
change that format once already, could the routed ids get an assertion here too?
##########
plan/src/test/java/org/apache/flink/agents/plan/actions/ChatModelActionRoutingTest.java:
##########
@@ -0,0 +1,725 @@
+/*
+ * 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<Object> outcomes = new ArrayDeque<>();
+
+ FakeChatModel(Object... outcomes) {
+ super(new ResourceDescriptor("fake", Map.of()), null);
+ Collections.addAll(this.outcomes, outcomes);
+ }
+
+ @Override
+ public Map<String, Object> getParameters() {
+ return Map.of();
+ }
+
+ @Override
+ public ChatMessage chat(
+ List<ChatMessage> messages,
+ Map<String, Object> promptArgs,
+ Map<String, Object> 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<Event> sentEvents = new ArrayList<>();
+ final List<String> resolvedChatModels = new ArrayList<>();
+ final List<String> durableCallIds = new ArrayList<>();
+ final Map<String, BaseChatModelSetup> models = new HashMap<>();
+ final Set<String> 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) {
Review Comment:
`withErrorHandling` only ever gets `IGNORE` (`:511`, `:634`), and
`MAX_RETRIES` isn't in this file at all, so every routing test runs with
`numRetries == 0`. Nothing here would notice if retry and fallback swapped
order, even though the class javadoc calls that ordering a guarantee — selected
model with its full retry budget first, then the rest
(`ChatModelAction.java:73-75`). `FakeChatModel` already takes a scripted
sequence of outcomes. Is one `RETRY` case with two failures scripted enough to
pin it?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
