[api][plan][python] Use request-scoped metric groups for chat token metrics by jlon · Pull Request #861 · apache/flink-agents

jlon · 2026-06-26T06:55:13Z

What

Capture the current action metric group when a chat request is initiated.
Record delayed chat token metrics against that captured group instead of the mutable metric group on cached chat model resources.
Add Java and Python regression coverage for resource metric-group rebinding.

Why

Cached resources can be rebound when another action retrieves the same resource before the original request records token usage. The delayed token metrics should stay under the action that initiated the request.

Closes #859.

Tests

mvn -pl api -DskipITs -Dspotless.check.skip=true -Dcheckstyle.skip=true -Dtest=BaseChatModelSetupTokenMetricsTest test
mvn -pl plan -am -DskipITs -Dspotless.check.skip=true -Dcheckstyle.skip=true -Dsurefire.failIfNoSpecifiedTests=false -Dtest=ChatModelActionTest,ChatModelActionRetryTest test
bash -lc 'export PATH="$HOME/.local/bin:$PATH"; cd python && uv run --no-sync pytest flink_agents/api/chat_models/tests/test_token_metrics.py -q'
bash -lc 'export PATH="$HOME/.local/bin:$PATH"; cd python && uv run --no-sync ruff format --check flink_agents/api/chat_models/chat_model.py flink_agents/api/chat_models/tests/test_token_metrics.py flink_agents/plan/actions/chat_model_action.py && uv run --no-sync ruff check flink_agents/api/chat_models/chat_model.py flink_agents/api/chat_models/tests/test_token_metrics.py flink_agents/plan/actions/chat_model_action.py'
bash -lc 'export JAVA_HOME=/usr/lib/jvm/java-17-openjdk-amd64; export PATH="$JAVA_HOME/bin:$PATH"; mvn -pl api,plan -am spotless:check -Dspotless.skip=false -DskipTests -DskipITs -Dcheckstyle.skip=true'

…etrics

weiqingy · 2026-06-30T22:20:34Z

+            this resource's currently bound metric group is used.
        """
-        metric_group = self.metric_group
+        if metric_group is None:


The two languages now document slightly different contracts for the absent-group case: Java's 4-arg recordTokenMetrics treats a null group as "skip recording" (BaseChatModelSetup.java:135), while here metric_group is None falls back to self.metric_group — the live, possibly-rebound group. Every real call site passes a captured group (chat_model_action.py:334), so there's no impact on today's flow. The part I'm less sure about: a future caller passing None expecting Java parity would record under whatever action is currently bound — the #859 scenario again on the Python side, where Java would no-op instead. Is the asymmetry intentional? If the = None default is only there to keep the param optional for the test helper, dropping it so callers must pass the captured group would match Java's no-fallback contract; alternatively a one-line javadoc note that Java's null means skip would at least make each side self-documenting.

Updated in ecd96c5. _record_token_metrics now requires the metric group argument and no longer falls back to the resource-bound self.metric_group. The action path also skips token metric recording when no request/action metric group is available.

wenjin272

Thanks for taking this on @jlon.

Since Resource instances are shared across multiple Actions, reporting metrics using the metricGroup bound at creation time can lead to incorrect metric matching. This PR addresses the issue by passing the current Action’s metricGroup at reporting time, ensuring metrics are reported to the correct Action scope.

In this case, I think Resource no longer need to hold a metricGroup field.

wenjin272 · 2026-07-06T08:34:35Z

@@ -116,7 +116,22 @@ public void open() throws Exception {
     * @param completionTokens the number of completion tokens
     */
    public void recordTokenMetrics(String modelName, long promptTokens, long completionTokens) {


It appears that only tests call this method, so I think it can be removed directly.

Updated in ecd96c5. Removed the 3-arg Java recordTokenMetrics(...) overload and updated the tests to pass an explicit request/action metric group.

wenjin272 · 2026-07-06T08:36:09Z

+     * @param completionTokens the number of completion tokens
+     */
+    public void recordTokenMetrics(
+            @Nullable FlinkAgentsMetricGroup metricGroup,


In practice, we now require that the metricGroup of the current action be passed in when reporting metrics. Therefore, I believe metricGroup should not be nullable.

Updated in ecd96c5. The Java recordTokenMetrics(...) metric group parameter is no longer nullable, and ChatModelAction now skips token metric recording before calling it when the action metric group is absent.

joeyutong · 2026-07-07T09:27:32Z

Thanks for taking this on @jlon.

Since Resource instances are shared across multiple Actions, reporting metrics using the metricGroup bound at creation time can lead to incorrect metric matching. This PR addresses the issue by passing the current Action’s metricGroup at reporting time, ensuring metrics are reported to the correct Action scope.

In this case, I think Resource no longer need to hold a metricGroup field.

I agree. This PR fixes the concrete bug locally for chat token metrics by passing the request/action metric group at reporting time, instead of reading the mutable metric group currently bound to the cached resource. But the underlying model issue is broader than this one recording path.

A cleaner abstraction would be to separate the stable resource instance from the per-action resource-use context:

Resource is a stable cached capability.
ResourceUseContext carries action-scoped state, including the current action metric group.
Metric reporting uses the current ResourceUseContext, not mutable state stored on the resource.
Cross-language bridges propagate this context when invoking resource methods or resolving nested resources.

One possible shape is for ctx.getResource(...) to return a bound resource handle internally, i.e. the cached Resource plus the current ResourceUseContext, even if the public API is migrated gradually.

Under that model, Resource would no longer need a mutable metricGroup field. It would also simplify the cross-language propagation work in #860: instead of making every Python wrapper participate in setMetricGroup(...), the bridge can pass the current use context through the existing callMethod / getResource path. CC @weiqingy

This is probably a larger refactor than this PR should take on, but I agree with the direction: action-scoped metric state should not be owned by the cached Resource instance.

jlon · 2026-07-08T06:08:31Z

Pushed ecd96c5 to address the review comments:

Removed the 3-arg Java recordTokenMetrics(...) overload so token metrics are recorded only with an explicit request/action metric group.
Made the Java metric group parameter non-null and added an action-layer skip when no metric group is available.
Removed the Python fallback from metric_group=None to self.metric_group; Python now skips instead of recording under a possibly rebound resource group.

I kept the broader Resource.metricGroup removal out of this PR because that crosses resource binding and cross-language bridge behavior. This patch keeps the issue #859 fix focused on chat token metrics while avoiding the fallback contract that could recreate the bug.

Local verification:

mvn -pl api -DskipITs -Dspotless.check.skip=true -Dcheckstyle.skip=true -Dtest=BaseChatModelSetupTokenMetricsTest test
mvn -pl plan -am -DskipITs -Dspotless.check.skip=true -Dcheckstyle.skip=true -Dsurefire.failIfNoSpecifiedTests=false -Dtest=ChatModelActionTest,ChatModelActionRetryTest test
cd python && uv run --no-sync pytest flink_agents/api/chat_models/tests/test_token_metrics.py -q
cd python && uv run --no-sync ruff format --check ... && uv run --no-sync ruff check ...
mvn -pl api,plan -am spotless:check -Dspotless.skip=false -DskipTests -DskipITs -Dcheckstyle.skip=true
git diff --check

wenjin272

I agree with excluding the removal of metricGroup bindings from this PR. We should create a separate issue to evaluate this holistically.

I’ve left one comment regarding this PR.

wenjin272 · 2026-07-10T07:58:12Z

+            metrics are skipped.
        """
-        metric_group = self.metric_group
        if metric_group is None:


One parity question: Python _record_token_metrics(..., metric_group=None) is a no-op, while Java recordTokenMetrics(null, ...) fails via checkNotNull. Should we align the lower-level helper contract?

# 影响范围 - Java/Python `BaseEmbeddingModelConnection` 与 `BaseEmbeddingModelSetup` 保留 `embedWithUsage`/`embed_with_usage` 的结果携带链路；普通 `embed` API 的返回类型不变。 - `BaseVectorStore.query`、自动补 embedding 以及 Python 对应实现恢复仅调用 `embed`，不再进入指标组或计数器。 - OpenAI、Tongyi、Bedrock 的 provider usage 提取及批量聚合能力保持不变；测试改为断言 usage 结果本身。 # 改动影响面 - 主链路：embedding provider -> `EmbeddingResult` -> 调用方显式消费 usage；本提交不再把 usage 写入 `MetricGroup`。 - 向量库扩展点：恢复 `BaseEmbeddingModelSetup.embed` 覆盖语义，并且仅在确有缺失 embedding 时延迟解析模型，避免 Mem0 无模型路径报错。 - 不受影响范围：apache#860/apache#861 的 resource metric-group 传播和 action-scoped metric 语义未改动；未引入新的 metric abstraction。 # 功能改进/开发/新增 - 类型: refactor - 触发条件: 普通资源 API 可并发调用，无法安全保证 `MetricGroup`/`SimpleCounter` 的线程模型；Python vector-store 还绕过了 `embed` 扩展点。 - 根因类别: 指标写入边界不明确、公共 API 扩展点回归、惰性资源解析边界遗漏。 - 行为变化: 有；PR 仅暴露 provider usage 返回 API，不再自动记录 embedding token metrics。 # 验证 - `mvn --batch-mode --no-transfer-progress -pl api -Dtest=BaseEmbeddingModelSetupEmbeddingResultTest test` PASS - `mvn --batch-mode --no-transfer-progress -pl integrations/embedding-models/bedrock -am -Dtest=BedrockEmbeddingModelTest -Dsurefire.failIfNoSpecifiedTests=false test` PASS - `uv run --python 3.12 --extra test pytest ...vector_stores... -q` PASS (49 passed, 4 skipped) - `./tools/lint.sh --check` PASS - `uv run --python 3.12 --extra lint ruff format --check ... && uv run --python 3.12 --extra lint ruff check ...` PASS - `git diff --check` PASS

[api][plan][python] Use request-scoped metric groups for chat token m…

a3d7330

…etrics

github-actions Bot added doc-label-missing The Bot applies this label either because none or multiple labels were provided. fixVersion/0.4.0 priority/major Default priority of the PR or issue. labels Jun 26, 2026

weiqingy reviewed Jun 30, 2026

View reviewed changes

zxs1633079383 mentioned this pull request Jul 4, 2026

[Feature] Track token usage metrics for embedding models #858

Open

2 tasks

wenjin272 reviewed Jul 6, 2026

View reviewed changes

[api][plan][python] Address request-scoped metrics review

ecd96c5

zxs1633079383 mentioned this pull request Jul 9, 2026

[api][python][java] Track embedding token usage metrics #870

Open

3 tasks

wenjin272 added doc-not-needed Your PR changes do not impact docs and removed doc-label-missing The Bot applies this label either because none or multiple labels were provided. labels Jul 10, 2026

wenjin272 reviewed Jul 10, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[api][plan][python] Use request-scoped metric groups for chat token metrics#861

[api][plan][python] Use request-scoped metric groups for chat token metrics#861
jlon wants to merge 2 commits into
apache:mainfrom
jlon:fix-action-scoped-token-metrics

jlon commented Jun 26, 2026

Uh oh!

weiqingy Jun 30, 2026

Uh oh!

jlon Jul 8, 2026

Uh oh!

wenjin272 left a comment

Uh oh!

wenjin272 Jul 6, 2026

Uh oh!

jlon Jul 8, 2026

Uh oh!

wenjin272 Jul 6, 2026

Uh oh!

jlon Jul 8, 2026

Uh oh!

joeyutong commented Jul 7, 2026

Uh oh!

jlon commented Jul 8, 2026

Uh oh!

wenjin272 left a comment

Uh oh!

wenjin272 Jul 10, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Uh oh!

Conversation

jlon commented Jun 26, 2026

What

Why

Tests

Uh oh!

weiqingy Jun 30, 2026

Choose a reason for hiding this comment

Uh oh!

jlon Jul 8, 2026

Choose a reason for hiding this comment

Uh oh!

wenjin272 left a comment

Choose a reason for hiding this comment

Uh oh!

wenjin272 Jul 6, 2026

Choose a reason for hiding this comment

Uh oh!

jlon Jul 8, 2026

Choose a reason for hiding this comment

Uh oh!

wenjin272 Jul 6, 2026

Choose a reason for hiding this comment

Uh oh!

jlon Jul 8, 2026

Choose a reason for hiding this comment

Uh oh!

joeyutong commented Jul 7, 2026

Uh oh!

jlon commented Jul 8, 2026

Uh oh!

wenjin272 left a comment

Choose a reason for hiding this comment

Uh oh!

wenjin272 Jul 10, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants