purushah commented on code in PR #1042:
URL: https://github.com/apache/flink-agents/pull/1042#discussion_r3917117767
##########
api/src/main/java/org/apache/flink/agents/api/chat/model/routing/ModelRouter.java:
##########
@@ -84,30 +96,113 @@ public ModelRouter(ResourceDescriptor descriptor,
ResourceContext resourceContex
}
this.fallbackEnabled =
Boolean.TRUE.equals(descriptor.getArgument("fallback",
Boolean.FALSE));
- String strategyClazz = descriptor.getArgument("strategy_clazz");
+ String typeTag = descriptor.getArgument(STRATEGY_TYPE_KEY);
+ if (typeTag == null || typeTag.isEmpty()) {
+ throw new IllegalArgumentException("ModelRouter requires a routing
strategy.");
+ }
Map<String, Object> strategyArgs =
- descriptor.getArgument("strategy_args",
Collections.emptyMap());
- this.strategy = instantiateStrategy(strategyClazz, strategyArgs);
+ descriptor.getArgument(STRATEGY_ARGS_KEY,
Collections.emptyMap());
+ String executorClass =
descriptor.getArgument(STRATEGY_EXECUTOR_CLASS_KEY);
+ // The declaration constructor owns the per-type argument rules, so a
structurally invalid
+ // configuration fails here (resource construction) with the same
message as at build().
+ this.strategy =
+ new RoutingStrategy(
+ RoutingStrategyType.fromTag(typeTag), strategyArgs,
executorClass);
+ this.compiledRules = compileRules(this.strategy);
+ this.customExecutor = instantiateCustomExecutor(this.strategy);
}
- @SuppressWarnings("unchecked")
- private static RoutingStrategy instantiateStrategy(String clazz,
Map<String, Object> args)
+ /**
+ * Instantiates the user's {@link CustomRoutingExecutor} once per router
instance (routers are
+ * cached per TaskManager), preserving executor instance state across
requests. The construction
Review Comment:
Good catch — "per subtask" is the accurate phrase. Will fix both spots;
since that sentence is what tells executor authors how instance state behaves,
it should be precise.
##########
plan/src/main/java/org/apache/flink/agents/plan/actions/ModelRoutingResolver.java:
##########
@@ -111,12 +188,310 @@ public RoutingDecision call() throws Exception {
if (!router.isCandidate(selectedModel)) {
throw new IllegalStateException(
String.format(
- "Routing strategy for router '%s' returned
non-candidate model '%s'; candidates are %s.",
+ "Routing decision for router '%s' selected
non-candidate model '%s'; candidates are %s.",
model, selectedModel,
router.getCandidateNames()));
}
- decisionSource = ModelRoutingEvent.SOURCE_STRATEGY;
+ decisionSource = concreteSource;
}
+ return finish(requestId, model, router, decision, selectedModel,
decisionSource, ctx);
+ }
+
+ /** Records the decision latency histogram sample (also for decisions the
guards reject). */
+ private static void recordDecisionLatency(RunnerContext ctx,
RoutingDecision decision) {
+ Double decisionMs = decision.getDecisionMs();
+ FlinkAgentsMetricGroup actionMetrics = ctx.getActionMetricGroup();
+ if (actionMetrics != null && decisionMs != null) {
+
actionMetrics.getHistogram("routingDecisionLatencyMs").update(Math.round(decisionMs));
+ }
+ }
+ /**
+ * LLM-as-judge path (framework-managed, per discussion #897): the engine
runs the judge chat
+ * itself through the normal durable/metered/observable invoker path —
durable id {@code
+ * "judge:<router>"} so a recovered run replays the original verdict
instead of re-calling the
+ * judge (with a durable action-state store configured; without one the
judge re-runs on replay,
+ * like any non-deterministic strategy) — then derives the decision from
the verdict as a pure
+ * function ({@link LlmJudgeRoutingExecutor}). The decision (including its
wall time, which
+ * covers the judge call) is persisted under the standard {@code
"route:<router>"} durable call,
+ * preserving the replay-fingerprint property. Verdict abstains are
persisted <i>as
+ * abstains</i>, so a replay after a candidate-set change re-resolves to
the current default
+ * exactly like the strategy path.
+ *
+ * <p>Failure policy: an unparseable or non-candidate verdict always
abstains to the router's
+ * default model. A judge call that exhausts its retries honors the
request's error-handling
+ * strategy, exactly like a throwing rule/custom strategy: {@code FAIL}
surfaces the outage
+ * loudly, {@code IGNORE} degrades to the default with the cause recorded.
Cancellation
+ * propagates and is never persisted as a routing outcome.
+ */
+ private static ResolvedModelRoute resolveViaJudge(
+ UUID requestId,
+ String model,
+ ModelRouter router,
+ RoutingStrategy strategy,
+ RoutingContext routingContext,
+ RunnerContext ctx)
+ throws Exception {
+ long start = System.nanoTime();
+ Agent.ErrorHandlingStrategy errorStrategy =
+
ctx.getConfig().get(AgentExecutionOptions.ERROR_HANDLING_STRATEGY);
+ int numRetries = ChatModelInvoker.configuredRetries(ctx,
errorStrategy);
+ int retryWaitIntervalSec =
ChatModelInvoker.configuredRetryWaitSec(ctx, errorStrategy);
+ String judgeModel = LlmJudgeRoutingExecutor.judgeModel(strategy);
+
+ Map<String, Object> judgeMetadata = new LinkedHashMap<>();
+ judgeMetadata.put("judge_model", judgeModel);
+ String verdictModel = null;
+ String abstainReason = null;
+ // Runtime backstop to the plan-time check: plan validation only sees
descriptor-carried
+ // bindings, so a setup that binds a prompt/tools/skills at the
instance level (or via a
+ // non-Java provider) is caught here. Same policy as a failed judge
call: FAIL is loud
+ // (a config error should not hide), IGNORE abstains so the default
model keeps answering.
+ String misconfigured = judgeSetupMisconfiguration(judgeModel, ctx);
+ if (misconfigured != null && errorStrategy !=
Agent.ErrorHandlingStrategy.IGNORE) {
+ throw new IllegalStateException(misconfigured);
+ }
+ if (misconfigured != null) {
+ abstainReason = misconfigured;
+ } else
+ try {
+ boolean[] truncated = new boolean[1];
+ List<ChatMessage> effective = effectiveJudgeMessages(router,
routingContext, ctx);
+ List<ChatMessage> judgeInput =
+ LlmJudgeRoutingExecutor.buildJudgeMessages(
+ strategy,
+ routingContext,
+ effective,
+
pinnedRenderedIndices(routingContext.getMessages(), effective),
+ truncated);
+ if (truncated[0]) {
+
judgeMetadata.put(LlmJudgeRoutingExecutor.CONTEXT_TRUNCATED_KEY, true);
+ }
+ ChatModelInvoker.ChatAttemptResult judgeResult =
+ ChatModelInvoker.chatWithRetries(
+ requestId,
+ judgeModel,
+ "judge:" + model,
+ judgeInput,
+ Map.of(),
+ null,
+ ctx,
+ errorStrategy,
+ numRetries,
+ retryWaitIntervalSec);
+ ChatModelAction.recordAttemptRetryStats(
+ ctx,
+ requestId,
+ judgeResult.chatModel,
+ judgeResult.retryCount,
+ judgeResult.totalRetryWaitSec);
+ ChatMessage reply = judgeResult.response;
+ Object promptTokens = reply.getExtraArgs().get("promptTokens");
+ Object completionTokens =
reply.getExtraArgs().get("completionTokens");
+ if (promptTokens != null) {
+ judgeMetadata.put("judge_prompt_tokens", promptTokens);
+ }
+ if (completionTokens != null) {
+ judgeMetadata.put("judge_completion_tokens",
completionTokens);
+ }
+ verdictModel =
+ LlmJudgeRoutingExecutor.parseVerdict(
+ reply.getContent(),
router.getCandidateNames())
+ .orElse(null);
+ abstainReason =
+ verdictModel == null ? "judge verdict was not a
candidate name" : null;
+ } catch (InterruptedException cancellation) {
+ // Cancellation surfacing from the between-retries backoff
sleep.
+ Thread.currentThread().interrupt();
+ throw cancellation;
+ } catch (ChatModelInvoker.ChatAttemptFailed failure) {
+ ChatModelAction.recordAttemptRetryStats(
+ ctx,
+ requestId,
+ failure.chatModel,
+ failure.retryCount,
+ failure.totalRetryWaitSec);
+ // Cancellation surfacing from inside the judge attempt (the
invoker wraps every
+ // attempt exception): it must propagate, never persist as a
routing outcome.
+ if (isCancellation(failure)) {
+ Thread.currentThread().interrupt();
+ throw failure;
+ }
+ // A judge that exhausted its retries honors the request's
error-handling strategy,
+ // exactly like a throwing rule/custom strategy (see class
javadoc).
+ if (errorStrategy != Agent.ErrorHandlingStrategy.IGNORE) {
+ throw failure;
+ }
+ abstainReason = "judge call failed: " + failure.error;
+ }
+
+ RoutingDecision computed;
+ if (verdictModel != null) {
+ RoutingDecision.Builder builder =
+ RoutingDecision.builder(verdictModel).reason("llm judge
verdict");
+ for (Map.Entry<String, Object> entry : judgeMetadata.entrySet()) {
+ builder.metadata(entry.getKey(), entry.getValue());
+ }
+ computed = builder.build();
+ } else {
+ // Persisted as a real abstain: replay resolves to the router's
*current* default, so
+ // a candidate-set change across a restart degrades gracefully
(like the strategy
+ // path) instead of failing the non-candidate guard.
+ computed =
+ new RoutingDecision(
+ null, true, abstainReason, null, new
HashMap<>(judgeMetadata), null);
+ }
+ final RoutingDecision toStore =
+ computed.withDecisionMs((System.nanoTime() - start) /
1_000_000.0);
+
+ // Persist under the standard route id: on recovery the stored
decision (with its original
+ // judge-inclusive wall time) replays; the judge chat above replays
from its own durable
+ // record, so the recomputation feeding this call is deterministic.
+ RoutingDecision decision =
+ ctx.durableExecute(
Review Comment:
Fair question — it deserves its own test. The judge path has enough of its
own orchestration before the shared store that "covered by shared code" is
exactly the kind of claim I shouldn't leave untested. Will add a replay test
that seeds route:<router> for a judge router and asserts the judge is never
invoked.
##########
plan/src/main/java/org/apache/flink/agents/plan/actions/ModelRoutingResolver.java:
##########
@@ -111,12 +188,310 @@ public RoutingDecision call() throws Exception {
if (!router.isCandidate(selectedModel)) {
throw new IllegalStateException(
String.format(
- "Routing strategy for router '%s' returned
non-candidate model '%s'; candidates are %s.",
+ "Routing decision for router '%s' selected
non-candidate model '%s'; candidates are %s.",
model, selectedModel,
router.getCandidateNames()));
}
- decisionSource = ModelRoutingEvent.SOURCE_STRATEGY;
+ decisionSource = concreteSource;
}
+ return finish(requestId, model, router, decision, selectedModel,
decisionSource, ctx);
+ }
+
+ /** Records the decision latency histogram sample (also for decisions the
guards reject). */
+ private static void recordDecisionLatency(RunnerContext ctx,
RoutingDecision decision) {
+ Double decisionMs = decision.getDecisionMs();
+ FlinkAgentsMetricGroup actionMetrics = ctx.getActionMetricGroup();
+ if (actionMetrics != null && decisionMs != null) {
+
actionMetrics.getHistogram("routingDecisionLatencyMs").update(Math.round(decisionMs));
+ }
+ }
+ /**
+ * LLM-as-judge path (framework-managed, per discussion #897): the engine
runs the judge chat
+ * itself through the normal durable/metered/observable invoker path —
durable id {@code
+ * "judge:<router>"} so a recovered run replays the original verdict
instead of re-calling the
+ * judge (with a durable action-state store configured; without one the
judge re-runs on replay,
+ * like any non-deterministic strategy) — then derives the decision from
the verdict as a pure
+ * function ({@link LlmJudgeRoutingExecutor}). The decision (including its
wall time, which
+ * covers the judge call) is persisted under the standard {@code
"route:<router>"} durable call,
+ * preserving the replay-fingerprint property. Verdict abstains are
persisted <i>as
+ * abstains</i>, so a replay after a candidate-set change re-resolves to
the current default
+ * exactly like the strategy path.
+ *
+ * <p>Failure policy: an unparseable or non-candidate verdict always
abstains to the router's
+ * default model. A judge call that exhausts its retries honors the
request's error-handling
+ * strategy, exactly like a throwing rule/custom strategy: {@code FAIL}
surfaces the outage
+ * loudly, {@code IGNORE} degrades to the default with the cause recorded.
Cancellation
+ * propagates and is never persisted as a routing outcome.
+ */
+ private static ResolvedModelRoute resolveViaJudge(
+ UUID requestId,
+ String model,
+ ModelRouter router,
+ RoutingStrategy strategy,
+ RoutingContext routingContext,
+ RunnerContext ctx)
+ throws Exception {
+ long start = System.nanoTime();
+ Agent.ErrorHandlingStrategy errorStrategy =
+
ctx.getConfig().get(AgentExecutionOptions.ERROR_HANDLING_STRATEGY);
+ int numRetries = ChatModelInvoker.configuredRetries(ctx,
errorStrategy);
+ int retryWaitIntervalSec =
ChatModelInvoker.configuredRetryWaitSec(ctx, errorStrategy);
+ String judgeModel = LlmJudgeRoutingExecutor.judgeModel(strategy);
+
+ Map<String, Object> judgeMetadata = new LinkedHashMap<>();
+ judgeMetadata.put("judge_model", judgeModel);
+ String verdictModel = null;
+ String abstainReason = null;
+ // Runtime backstop to the plan-time check: plan validation only sees
descriptor-carried
+ // bindings, so a setup that binds a prompt/tools/skills at the
instance level (or via a
+ // non-Java provider) is caught here. Same policy as a failed judge
call: FAIL is loud
+ // (a config error should not hide), IGNORE abstains so the default
model keeps answering.
+ String misconfigured = judgeSetupMisconfiguration(judgeModel, ctx);
+ if (misconfigured != null && errorStrategy !=
Agent.ErrorHandlingStrategy.IGNORE) {
+ throw new IllegalStateException(misconfigured);
+ }
+ if (misconfigured != null) {
+ abstainReason = misconfigured;
+ } else
+ try {
+ boolean[] truncated = new boolean[1];
+ List<ChatMessage> effective = effectiveJudgeMessages(router,
routingContext, ctx);
+ List<ChatMessage> judgeInput =
+ LlmJudgeRoutingExecutor.buildJudgeMessages(
+ strategy,
+ routingContext,
+ effective,
+
pinnedRenderedIndices(routingContext.getMessages(), effective),
+ truncated);
+ if (truncated[0]) {
+
judgeMetadata.put(LlmJudgeRoutingExecutor.CONTEXT_TRUNCATED_KEY, true);
+ }
+ ChatModelInvoker.ChatAttemptResult judgeResult =
+ ChatModelInvoker.chatWithRetries(
+ requestId,
+ judgeModel,
+ "judge:" + model,
+ judgeInput,
+ Map.of(),
+ null,
+ ctx,
+ errorStrategy,
+ numRetries,
+ retryWaitIntervalSec);
+ ChatModelAction.recordAttemptRetryStats(
+ ctx,
+ requestId,
+ judgeResult.chatModel,
+ judgeResult.retryCount,
+ judgeResult.totalRetryWaitSec);
+ ChatMessage reply = judgeResult.response;
+ Object promptTokens = reply.getExtraArgs().get("promptTokens");
+ Object completionTokens =
reply.getExtraArgs().get("completionTokens");
+ if (promptTokens != null) {
Review Comment:
Good catch — will match the instanceof Number guard the reader already uses,
and extend the test with a non-Number value.
##########
python/flink_agents/api/execution_environment.py:
##########
@@ -237,6 +237,14 @@ def add_resource(
AgentsExecutionEnvironment
The environment to register the resource.
"""
+ if resource_type == ResourceType.MODEL_ROUTER:
Review Comment:
Nice find. The shared helper is cheap, so I'll do it in this PR — one helper
plus a parametrized test over both entry points, rather than leaving a known
duplication for the Python routing work to trip over.
--
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]