purushah commented on code in PR #964:
URL: https://github.com/apache/flink-agents/pull/964#discussion_r3793792530
##########
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:
You're right, and the one enum member is exactly the right scope:
`MODEL_ROUTER = "model_router"` is added to the Python `ResourceType` so a Java
plan containing a router deserializes on the Python side — mixed jobs start;
Python never needs to execute the router (the Java `ChatModelAction` owns it).
A cross-language round-trip test covers it, and we verified the wire shape
against the Java serializers (field names, enum value keys). One adjacent gap
we noticed for the follow-up: Python's `Agent.add_resource` now accepts
`MODEL_ROUTER` but the compiled plan silently drops it — the Python-routing
follow-up should make that an explicit not-yet-supported error. Full Python
routing (a Python `RoutingStrategy` API and events) stays that follow-up.
##########
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:
Fair — line 58 was written before the review rounds changed those paths. The
description now has a "Compatibility impact" section naming exactly your three:
per-attempt retry metrics (+ the retry WARN's logger move to
`ChatModelInvoker`) under `retry`, and missing-model drop under `ignore`;
default `FAIL` unchanged.
##########
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:
Synced — description line 26 now shows `"route:<router>"` with the
determinism rationale.
##########
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:
Added: `routedRequestUsesRoutedDurableCallIds` asserts `route:<router>` +
`chat:<router>:<candidate>` on the happy path, and the fallback test now
asserts the two candidates' distinct ids in order.
##########
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:
Added `retryBudgetRunsBeforeFallback`: the selected model fails once and
recovers on its retry (MAX_RETRIES=1, wait 0), the fallback candidate is never
resolved, and no fallback event is emitted — pinning the retry-before-fallback
ordering. `FakeRunnerContext` gained `withRetryBudget`.
##########
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:
Added `AgentPlanRoutingBackstopTest`: the same name arrives as MODEL_ROUTER
via `agent.addResource` and as CHAT_MODEL via `addResourcesIfAbsent` (bypassing
both per-call checks), and plan construction throws; a distinct-names control
passes.
##########
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:
Deliberate, and now documented: the javadoc states the isolation boundary is
one level deep — the copies make accidental top-level mutation harmless, nested
values stay shared, and the SPI forbids mutation. Arbitrary-depth copies on
every routing decision would tax the common case to guard a case the API
already forbids; happy to revisit if a real strategy needs deeper isolation.
--
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]
