purushah commented on code in PR #964:
URL: https://github.com/apache/flink-agents/pull/964#discussion_r3790092020
##########
plan/src/main/java/org/apache/flink/agents/plan/actions/ChatModelAction.java:
##########
@@ -400,56 +669,199 @@ public ChatMessage call() throws Exception {
Thread.sleep(currentWaitSec * 1000L);
totalWaitTimeSec += currentWaitSec;
}
- } else {
- LOG.debug(
- "Chat request {} failed, the input chat messages
are {}.",
- initialRequestId,
- messages);
- throw e;
+ continue;
}
+ throw new ChatAttemptFailed(
+ model, chatModel, e, actualRetryCount,
totalWaitTimeSec);
}
}
+ throw new IllegalStateException("Unreachable chat retry state.");
+ }
- if (actualRetryCount > 0) {
- accumulateRetryStats(
- ctx.getSensoryMemory(), initialRequestId,
actualRetryCount, totalWaitTimeSec);
+ 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.getConnectionName();
+ recordRetryMetrics(
+ ctx,
+ metricModel == null || metricModel.isEmpty() ? "unknown" :
metricModel,
+ retryCount,
+ retryWaitSec);
+ }
- 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();
+ private static List<String> candidateAttemptOrder(RoutingSelection
selection) {
+ List<String> order = new ArrayList<>();
+ order.add(selection.selectedModel);
+ if (selection.isRouter && selection.fallbackEnabled) {
+ for (String candidate : selection.candidates) {
+ if (!candidate.equals(selection.selectedModel)) {
+ order.add(candidate);
+ }
+ }
+ }
+ return order;
+ }
- recordRetryMetrics(
- ctx, chatModel.getConnectionName(), totalRetryCount,
totalRetryWaitSec);
+ private static String durableChatCallId(RoutingSelection selection, String
candidate) {
+ if (!selection.isRouter) {
+ return "chat";
+ }
+ return "chat:" + selection.requestedModel + ":" + candidate;
+ }
- ctx.sendEvent(
- new ChatResponseEvent(
- initialRequestId, response, totalRetryCount,
totalRetryWaitSec));
+ private static void attachRoutingMetadata(
+ ChatMessage response,
+ RoutingSelection selection,
+ String finalModel,
+ List<String> triedModels) {
+ boolean fallbackAttempted =
!finalModel.equals(selection.selectedModel);
+ List<String> fallbackModelsTried = new ArrayList<>();
+ for (int i = 1; i < triedModels.size(); i++) {
+ fallbackModelsTried.add(triedModels.get(i));
+ }
+ Map<String, Object> routing = new LinkedHashMap<>();
+ routing.put("router", selection.requestedModel);
+ routing.put("selected_model", selection.selectedModel);
+ routing.put("initial_selected_model", selection.selectedModel);
+ routing.put("final_model", finalModel);
+ routing.put("candidates", new ArrayList<>(selection.candidates));
+ routing.put(
+ "decision_source",
+ fallbackAttempted ? ModelRoutingEvent.SOURCE_FALLBACK :
selection.decisionSource);
+ routing.put("fallback_enabled", selection.fallbackEnabled);
+ routing.put("fallback_attempted", fallbackAttempted);
+ routing.put("fallback_models_tried", fallbackModelsTried);
+ routing.put("metadata", new LinkedHashMap<>(selection.metadata));
+ if (selection.reason != null) {
+ routing.put("reason", selection.reason);
+ }
+ if (selection.score != null) {
+ routing.put("score", selection.score);
}
+ response.getExtraArgs().put("model_routing", routing);
}
private static void processChatRequest(ChatRequestEvent event,
RunnerContext ctx)
throws Exception {
+ RoutingSelection selection =
+ resolveRouter(
+ event.getId(),
+ event.getModel(),
+ event.getMessages(),
+ event.getPromptArgs(),
+ ctx);
chat(
event.getId(),
- event.getModel(),
+ selection,
event.getMessages(),
event.getPromptArgs(),
event.getOutputSchema(),
ctx);
}
+ /**
+ * 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
{@link
+ * #handleToolCalls}), so this method is only reached with a router name
on the initial request.
+ */
+ private static RoutingSelection resolveRouter(
+ UUID requestId,
+ String model,
+ List<ChatMessage> messages,
+ Map<String, Object> promptArgs,
+ RunnerContext ctx)
+ throws Exception {
+ if (!ctx.hasResource(model, ResourceType.MODEL_ROUTER)) {
+ return RoutingSelection.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() {
+ return "route:" + requestId + ":" + model;
+ }
+
+ @Override
+ public Class<RoutingDecision> 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);
Review Comment:
Great catch — and you found *two* holes in one comment (the per-record throw
AND the policy bypass). Fixed both: `build()` validates rule keys so the typo
fails at registration, and `resolveRouter` failures now honor
`error-handling-strategy` — deliberately without retries, since v1 strategies
are CPU-only. Tests: `builderRejectsRuleKeyThatIsNotACandidate`,
`strategyFailureIsIgnoredUnderIgnorePolicy`,
`strategyFailurePropagatesUnderDefaultPolicy`.
##########
plan/src/main/java/org/apache/flink/agents/plan/actions/ChatModelAction.java:
##########
@@ -344,15 +505,129 @@ public static void chat(
: 0;
}
- ChatMessage response = null;
+ List<String> triedModels = new ArrayList<>();
+ Exception lastError = null;
+ for (String candidate : candidateAttemptOrder(selection)) {
+ triedModels.add(candidate);
+ try {
+ ChatAttemptResult result =
+ chatWithRetries(
+ initialRequestId,
+ candidate,
+ durableChatCallId(selection, candidate),
+ messages,
+ promptArgs,
+ outputSchema,
+ ctx,
+ strategy,
+ numRetries,
+ retryWaitIntervalSec);
+ recordAttemptRetryStats(
+ ctx,
+ initialRequestId,
+ result.chatModel,
+ result.retryCount,
+ result.totalRetryWaitSec);
+ if (selection.isRouter) {
+ attachRoutingMetadata(result.response, selection,
result.model, triedModels);
+ 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));
+ }
+ } else if (selection.carriedRouting != null) {
+ result.response
+ .getExtraArgs()
+ .put("model_routing", new
LinkedHashMap<>(selection.carriedRouting));
+ }
+
+ if
(!Objects.requireNonNull(result.response).getToolCalls().isEmpty()) {
+ handleToolCalls(
+ result.response,
+ initialRequestId,
+ result.model,
+ result.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();
+
+ ctx.sendEvent(
+ new ChatResponseEvent(
+ initialRequestId,
+ result.response,
+ totalRetryCount,
+ totalRetryWaitSec));
+ }
+ return;
+ } catch (ChatAttemptFailed e) {
+ recordAttemptRetryStats(
+ ctx, initialRequestId, e.chatModel, e.retryCount,
e.totalRetryWaitSec);
+ lastError = e.error;
Review Comment:
Really appreciate you noticing the test itself couldn't detect this (both
scripted errors shared "down") — that's a level of review rigor we rarely get.
Applied your suggestion: errors chain via `addSuppressed`, exhaustion WARNs
with the tried list, and the test now uses distinct markers so it fails against
the old behavior.
##########
plan/src/main/java/org/apache/flink/agents/plan/actions/ChatModelAction.java:
##########
@@ -344,15 +505,129 @@ public static void chat(
: 0;
}
- ChatMessage response = null;
+ List<String> triedModels = new ArrayList<>();
+ Exception lastError = null;
+ for (String candidate : candidateAttemptOrder(selection)) {
+ triedModels.add(candidate);
+ try {
+ ChatAttemptResult result =
+ chatWithRetries(
+ initialRequestId,
+ candidate,
+ durableChatCallId(selection, candidate),
+ messages,
+ promptArgs,
+ outputSchema,
+ ctx,
+ strategy,
+ numRetries,
+ retryWaitIntervalSec);
+ recordAttemptRetryStats(
+ ctx,
+ initialRequestId,
+ result.chatModel,
+ result.retryCount,
+ result.totalRetryWaitSec);
+ if (selection.isRouter) {
+ attachRoutingMetadata(result.response, selection,
result.model, triedModels);
+ 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));
+ }
+ } else if (selection.carriedRouting != null) {
+ result.response
+ .getExtraArgs()
+ .put("model_routing", new
LinkedHashMap<>(selection.carriedRouting));
+ }
+
+ if
(!Objects.requireNonNull(result.response).getToolCalls().isEmpty()) {
+ handleToolCalls(
+ result.response,
+ initialRequestId,
+ result.model,
+ result.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();
+
+ ctx.sendEvent(
+ new ChatResponseEvent(
+ initialRequestId,
+ result.response,
+ totalRetryCount,
+ totalRetryWaitSec));
+ }
+ return;
+ } catch (ChatAttemptFailed e) {
+ recordAttemptRetryStats(
+ ctx, initialRequestId, e.chatModel, e.retryCount,
e.totalRetryWaitSec);
+ lastError = e.error;
+ LOG.debug(
+ "Chat request {} failed for model {}, the input chat
messages are {}.",
+ initialRequestId,
+ e.model,
+ messages);
+ }
+ }
+
+ if (strategy == Agent.ErrorHandlingStrategy.IGNORE) {
+ LOG.warn(
+ "Chat request {} failed with error: {}, ignored.",
initialRequestId, lastError);
+ return;
+ }
+ throw Objects.requireNonNull(lastError);
+ }
+
+ private static ChatAttemptResult chatWithRetries(
+ UUID initialRequestId,
+ String model,
+ String durableCallId,
+ List<ChatMessage> messages,
+ Map<String, Object> promptArgs,
+ @Nullable Object outputSchema,
+ RunnerContext ctx,
+ Agent.ErrorHandlingStrategy strategy,
+ int numRetries,
+ int retryWaitIntervalSec)
+ throws ChatAttemptFailed, Exception {
+ BaseChatModelSetup chatModel =
+ (BaseChatModelSetup) ctx.getResource(model,
ResourceType.CHAT_MODEL);
Review Comment:
This was the sharpest catch of the review — the failure mode (real error
discarded, confusing "resource not found" for a model the user never picked,
both fallback and IGNORE bypassed) is exactly the worst debugging experience.
Fixed: the lookup moved inside the attempt conversion, so it counts as a failed
attempt and flows through fallback/IGNORE with prior errors preserved. (Side
effect: a plain chat with a missing model now honors IGNORE too — we think
that's more consistent.) Test:
`unresolvableCandidateCountsAsFailedAttemptAndFallsBack`.
##########
api/src/main/java/org/apache/flink/agents/api/chat/model/routing/RoutingContext.java:
##########
@@ -0,0 +1,123 @@
+/*
+ * 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.
Review Comment:
Good catch — fixed with a true deep copy, including the toolCalls list and
inner maps (`ChatMessage`'s ctor stores toolCalls by reference; re-verifying
your comment caught that second hole too, so this one paid off twice). Tests:
`routingContextMessagesAreDeepCopied`,
`routingContextToolCallsAreDeepCopiedToo`.
##########
api/src/main/java/org/apache/flink/agents/api/EventType.java:
##########
@@ -39,6 +39,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 =
Review Comment:
Right — will register `ModelRoutingEvent` (plus a condition-expression test)
during the rebase onto current main.
##########
api/src/main/java/org/apache/flink/agents/api/chat/model/routing/RuleBasedRoutingStrategy.java:
##########
@@ -0,0 +1,84 @@
+/*
+ * 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.
+ *
+ * <p>Constructed reflectively from a {@link RoutingStrategyDescriptor} via
the {@code
+ * (Map<String,Object>)} constructor; use {@link Strategies#rules(Map)} to
build one.
+ */
+public class RuleBasedRoutingStrategy implements RoutingStrategy {
+
+ private static final long serialVersionUID = 1L;
+
+ private final Map<String, Pattern> rules;
+
+ @SuppressWarnings("unchecked")
+ public RuleBasedRoutingStrategy(Map<String, Object> args) {
+ this.rules = new LinkedHashMap<>();
+ Object raw = args == null ? null : args.get("rules");
+ if (raw instanceof Map) {
+ for (Map.Entry<String, ?> entry : ((Map<String, ?>)
raw).entrySet()) {
+ String candidate = entry.getKey();
+ String regex = String.valueOf(entry.getValue());
Review Comment:
Nice spot on a classic `String.valueOf(null)` trap. Fixed: null/non-String
rule values are rejected with a descriptive error. Tests:
`ruleStrategyRejectsNullRuleValue`, `ruleStrategyRejectsNonStringRuleValue`.
##########
plan/src/main/java/org/apache/flink/agents/plan/actions/ChatModelAction.java:
##########
@@ -51,14 +55,43 @@
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.
+ *
+ * <h2>Model routing overview</h2>
+ *
+ * <p>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:
+ *
+ * <ol>
+ * <li><b>Decide</b> — {@code resolveRouter} runs the router's strategy and
normalizes the result
+ * (abstain → default model; non-candidate → fail).
+ * <li><b>Durably</b> — the strategy runs inside a durable call ({@code
"route:<requestId>:
Review Comment:
Fair — the guarantee is conditional and now reads that way: replay requires
a configured action-state store (`actionStateStoreBackend`); without one the
decision re-executes. Doc-only.
##########
plan/src/main/java/org/apache/flink/agents/plan/actions/ChatModelAction.java:
##########
@@ -344,15 +505,129 @@ public static void chat(
: 0;
}
- ChatMessage response = null;
+ List<String> triedModels = new ArrayList<>();
+ Exception lastError = null;
+ for (String candidate : candidateAttemptOrder(selection)) {
+ triedModels.add(candidate);
+ try {
+ ChatAttemptResult result =
+ chatWithRetries(
+ initialRequestId,
+ candidate,
+ durableChatCallId(selection, candidate),
+ messages,
+ promptArgs,
+ outputSchema,
+ ctx,
+ strategy,
+ numRetries,
+ retryWaitIntervalSec);
+ recordAttemptRetryStats(
+ ctx,
+ initialRequestId,
+ result.chatModel,
+ result.retryCount,
+ result.totalRetryWaitSec);
+ if (selection.isRouter) {
+ attachRoutingMetadata(result.response, selection,
result.model, triedModels);
+ 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));
+ }
+ } else if (selection.carriedRouting != null) {
+ result.response
+ .getExtraArgs()
+ .put("model_routing", new
LinkedHashMap<>(selection.carriedRouting));
+ }
+
+ if
(!Objects.requireNonNull(result.response).getToolCalls().isEmpty()) {
+ handleToolCalls(
+ result.response,
+ initialRequestId,
+ result.model,
+ result.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();
+
+ ctx.sendEvent(
+ new ChatResponseEvent(
+ initialRequestId,
+ result.response,
+ totalRetryCount,
+ totalRetryWaitSec));
+ }
+ return;
+ } catch (ChatAttemptFailed e) {
+ recordAttemptRetryStats(
Review Comment:
Good eye for catching the semantic shift — it's intentional (per-attempt
recording makes fallback attribution land on the model that served the
attempt), and it's now documented in a compat note on `recordAttemptRetryStats`
plus the PR description.
--
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]
