You are viewing a plain text version of this content. The canonical link for it is here.
Posted to issues@flink.apache.org by "hulincup (via GitHub)" <gi...@apache.org> on 2026/08/25 00:53:33 UTC
[PR] [api][java][python] Record embedding token usage metrics [flink-agents]
hulincup opened a new pull request, #1047:
URL: https://github.com/apache/flink-agents/pull/1047
Linked issue: #858
### Purpose of change
Embedding providers already report token usage (Bedrock on Java; OpenAI/Tongyi via the Python cross-language bridge), and the types to carry it (`EmbeddingTokenUsage` / `EmbeddingResult.tokenUsage`) already exist. But nothing reads it back to record metrics, so provider usage is dropped before it reaches the metric layer. The chat side already records `promptTokens`/`completionTokens`; embeddings had no equivalent, which made embedding cost/usage hard to validate — exactly the gap #858 calls out.
This mirrors the chat path:
- **Java** `BaseEmbeddingModelSetup.recordTokenMetrics(modelName, promptTokens, totalTokens)` records under the same `model` key-value group used by chat metrics, called from both `embedWithUsage` overloads (the chokepoint covering direct calls, vector stores, and RAG).
- **Python** `BaseEmbeddingModelSetup._record_token_metrics` / `_record_token_usage` does the same in `embed_with_usage`, keeping Java/Python parity (the chat side already records in Python via `_record_token_metrics`).
Two design notes:
- Embeddings record `totalTokens` in place of chat's `completionTokens`, since there is no completion.
- Embedding calls do not run inside a plan action (unlike chat, whose recording happens in `ChatModelAction` with a request-scoped group). Vector-store, RAG, and direct calls reach the setup directly, so the resource-bound metric group injected via `setMetricGroup` is used instead. This is the one deliberate asymmetry with the chat path and is documented on the method.
No change to `RowTypeInfo` handling, no payload change for a schema that renders, and `embed` (non-usage) methods are unchanged — they discard usage because it is not returned (now noted in their Javadoc).
### Tests
Mirrors `BaseChatModelSetupTokenMetricsTest` for both languages:
- Java: records under model group, batch path, no-op when no metric group bound, no-op when provider reports no usage, null/blank model name records nothing and is rejected by `recordTokenMetrics`, counters accumulate.
- Python: same matrix via a `_MockMetricGroup`, plus batch and accumulation.
Verification:
- Python: `pytest flink_agents/api/embedding_models/tests/` — 10 passed (4 existing + 6 new), `ruff check` and `ruff format --check` clean.
- Java: cannot build locally (Java 11 required, host is Java 8). Verification relies on CI (`./tools/ut.sh`); spotless/line-length checked manually against google-java-format AOSP 100-col rules.
### API
Yes. New public method `BaseEmbeddingModelSetup.recordTokenMetrics(String, long, long)` (Java) and `_record_token_metrics` / `_record_token_usage` (Python, protected). No existing signature changes. New counter names `promptTokens` / `totalTokens` under the existing `model` group, consistent with chat's `promptTokens` (chat also emits `completionTokens`; embeddings do not).
### Documentation
- [x] `doc-not-needed` — adds metric counters under the existing `model` group with names aligned to the chat side; no public API surface or config changes beyond the new recording method.
### Was this patch authored or co-authored using generative AI tooling?
- [x] Yes
Generated-by: Claude Code 2.1.220 (glm-5.2[1m])
--
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: issues-unsubscribe@flink.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [api][java][python] Record embedding token usage metrics [flink-agents]
Posted by "wenjin272 (via GitHub)" <gi...@apache.org>.
wenjin272 commented on code in PR #1047:
URL: https://github.com/apache/flink-agents/pull/1047#discussion_r3912075877
##########
api/src/main/java/org/apache/flink/agents/api/embedding/model/BaseEmbeddingModelSetup.java:
##########
@@ -93,9 +94,52 @@ public String getModel() {
return model;
}
+ /**
+ * Record embedding token usage metrics for the given model on this setup's bound metric group.
+ *
+ * <p>Mirrors {@code BaseChatModelSetup#recordTokenMetrics} but records input-side tokens only,
+ * since embeddings have no completion tokens. Counters are placed under the same {@code model}
+ * key-value group used by chat metrics, so embedding and chat usage for a model share one
+ * dimension.
+ *
+ * <p>Unlike the chat path, embedding calls do not run inside a plan action that hands in a
+ * request-scoped metric group (vector-store, RAG, and direct calls reach this setup directly),
+ * so the resource-bound metric group injected via {@link #setMetricGroup} is used instead.
+ *
+ * @param modelName the name of the model used
+ * @param promptTokens the number of prompt tokens
+ * @param totalTokens the total number of tokens reported by the provider
+ */
+ public void recordTokenMetrics(String modelName, long promptTokens, long totalTokens) {
+ Preconditions.checkArgument(
+ modelName != null && !modelName.isBlank(), "Model name must not be null or blank.");
+ FlinkAgentsMetricGroup metricGroup = getMetricGroup();
Review Comment:
For context, this is the same recording-boundary issue that led #870 to remove automatic metric recording and retain only usage transport. @joeyutong is working on the framework-level fix, so this PR may need to wait for that work and then adopt the resulting API.
--
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: issues-unsubscribe@flink.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [api][java][python] Record embedding token usage metrics [flink-agents]
Posted by "weiqingy (via GitHub)" <gi...@apache.org>.
weiqingy commented on code in PR #1047:
URL: https://github.com/apache/flink-agents/pull/1047#discussion_r3891277408
##########
api/src/main/java/org/apache/flink/agents/api/embedding/model/BaseEmbeddingModelSetup.java:
##########
@@ -93,9 +94,52 @@ public String getModel() {
return model;
}
+ /**
+ * Record embedding token usage metrics for the given model on this setup's bound metric group.
+ *
+ * <p>Mirrors {@code BaseChatModelSetup#recordTokenMetrics} but records input-side tokens only,
+ * since embeddings have no completion tokens. Counters are placed under the same {@code model}
+ * key-value group used by chat metrics, so embedding and chat usage for a model share one
+ * dimension.
+ *
+ * <p>Unlike the chat path, embedding calls do not run inside a plan action that hands in a
+ * request-scoped metric group (vector-store, RAG, and direct calls reach this setup directly),
Review Comment:
This line says vector-store and RAG calls reach this setup, but they do not reach `embedWithUsage`, which is where the recording happens. All four call sites use `embed`: `BaseVectorStore.java:179` (RAG query), `BaseVectorStore.java:337` (auto-embed on add and update), `vector_store.py:290`, and `vector_store.py:349`.
I grepped `embedWithUsage` and `embed_with_usage` on `main` and found no vector-store or RAG call site at all. That matches your own Javadoc on `embed(String)`, which says usage is discarded there. Those are the paths #858 asks for, so today they would still record nothing.
Is extending to the `embed` paths in scope for this PR? Or would you rather land the direct-call case first and reword this sentence, and the matching claim in the PR body, to match what it covers?
One more thing on the line above: it says embedding calls do not run inside a plan action, but RAG does. `ContextRetrievalAction` is registered as `context_retrieval_action` at `ContextRetrievalAction.java:44`. And while you are in the PR body, the sentence about `RowTypeInfo` and "a schema that renders" looks like it came from a different change. Worth dropping?
##########
api/src/test/java/org/apache/flink/agents/api/embedding/model/BaseEmbeddingModelSetupTokenMetricsTest.java:
##########
@@ -0,0 +1,318 @@
+/*
+ * 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 of
+ * limitations under the License.
+ */
+
+package org.apache.flink.agents.api.embedding.model;
+
+import org.apache.flink.agents.api.metrics.FlinkAgentsMetricGroup;
+import org.apache.flink.agents.api.metrics.UpdatableGauge;
+import org.apache.flink.agents.api.resource.ResourceContext;
+import org.apache.flink.agents.api.resource.ResourceDescriptor;
+import org.apache.flink.metrics.Counter;
+import org.apache.flink.metrics.Histogram;
+import org.apache.flink.metrics.Meter;
+import org.apache.flink.metrics.SimpleCounter;
+import org.junit.jupiter.api.DisplayName;
+import org.junit.jupiter.api.Test;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.mockito.Mockito.mock;
+
+/**
+ * Test cases for embedding token usage metrics recorded by {@link BaseEmbeddingModelSetup}. Mirrors
+ * {@code BaseChatModelSetupTokenMetricsTest}: embedding providers already populate {@link
+ * EmbeddingTokenUsage} on the returned {@link EmbeddingResult}, but nothing records it until this
+ * setup reads it back at the {@code embedWithUsage} chokepoint.
+ */
+class BaseEmbeddingModelSetupTokenMetricsTest {
+
+ /** Value-based metric group that mirrors the one in the chat token-metrics test. */
+ private static class TestMetricGroup implements FlinkAgentsMetricGroup {
+ final Map<String, TestMetricGroup> subGroups = new HashMap<>();
+ final Map<String, SimpleCounter> counters = new HashMap<>();
+
+ @Override
+ public FlinkAgentsMetricGroup getSubGroup(String name) {
+ return subGroups.computeIfAbsent(name, k -> new TestMetricGroup());
+ }
+
+ @Override
+ public FlinkAgentsMetricGroup getSubGroup(String key, String value) {
+ return subGroups.computeIfAbsent(key + "=" + value, k -> new TestMetricGroup());
+ }
+
+ @Override
+ public Counter getCounter(String name) {
+ return counters.computeIfAbsent(name, k -> new SimpleCounter());
+ }
+
+ @Override
+ public UpdatableGauge getGauge(String name) {
+ return null;
+ }
+
+ @Override
+ public Meter getMeter(String name) {
+ return null;
+ }
+
+ @Override
+ public Meter getMeter(String name, Counter counter) {
+ return null;
+ }
+
+ @Override
+ public Histogram getHistogram(String name) {
+ return null;
+ }
+
+ @Override
+ public Histogram getHistogram(String name, int windowSize) {
+ return null;
+ }
+ }
+
+ private static final float[] VEC = new float[] {0.1f, 0.2f};
+
+ /**
+ * Builds a setup bound to a connection that reports the given usage on single-text embed, with
+ * the given model name in its descriptor (may be {@code null} to exercise the guard).
+ */
+ private static BaseEmbeddingModelSetup setupWithSingleUsageAndModel(
+ EmbeddingTokenUsage usage, String model) {
+ BaseEmbeddingModelSetup setup =
+ new BaseEmbeddingModelSetup(
+ new ResourceDescriptor("test", descriptorArgs(model)),
+ mock(ResourceContext.class)) {
+ @Override
+ public Map<String, Object> getParameters() {
+ return new HashMap<>();
+ }
+ };
+ setup.connection =
+ new BaseEmbeddingModelConnection(
+ new ResourceDescriptor("conn", Collections.emptyMap()),
+ mock(ResourceContext.class)) {
+ @Override
+ public float[] embed(String text, Map<String, Object> parameters) {
+ return VEC;
+ }
+
+ @Override
+ public List<float[]> embed(List<String> texts, Map<String, Object> parameters) {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public EmbeddingResult<float[]> embedWithUsage(
+ String text, Map<String, Object> parameters) {
+ return new EmbeddingResult<>(VEC, usage);
+ }
+ };
+ return setup;
+ }
+
+ /** Builds a setup bound to a connection that reports the given usage on single-text embed. */
+ private static BaseEmbeddingModelSetup setupWithSingleUsage(EmbeddingTokenUsage usage) {
+ return setupWithSingleUsageAndModel(usage, "bedrock-text");
+ }
+
+ /** Descriptor args with an optional model (omitted when null/blank so it stays unset). */
+ private static Map<String, String> descriptorArgs(String model) {
Review Comment:
This helper returns `Map<String, String>`, but `ResourceDescriptor`'s constructor takes `Map<String, Object>`. Java generics are invariant, so the call at `:106` does not compile.
That is what all 18 red CI checks are hitting. Every one of them stops at `flink-agents-api` testCompile, including the Elasticsearch job, which builds this module through `-am`:
```
BaseEmbeddingModelSetupTokenMetricsTest.java:[106,70] incompatible types:
java.util.Map<java.lang.String,java.lang.String> cannot be converted to
java.util.Map<java.lang.String,java.lang.Object>
```
Would changing this helper and its local `HashMap` to `Map<String, Object>` be enough? I tried just that change locally, and the file compiles, all 9 new tests pass, and `spotless:check` stays green. The two sibling call sites already work because `Map.of(...)` at `:156` and `Collections.emptyMap()` at `:115` take their type from the target on the spot, so only the explicitly typed helper trips.
One thing worth knowing, since you mentioned leaning on CI: a green `Code Style Check` does not tell you the Java side builds. Spotless formats test sources without compiling them, so it stays green straight through a compile error.
##########
python/flink_agents/api/embedding_models/embedding_model.py:
##########
@@ -155,4 +155,50 @@ def embed_with_usage(
"""Generate embeddings and return provider token usage when available."""
merged_kwargs = self.model_kwargs.copy()
merged_kwargs.update(kwargs)
- return self._get_connection().embed_with_usage(text, **merged_kwargs)
+ result = self._get_connection().embed_with_usage(text, **merged_kwargs)
+ self._record_token_usage(result.token_usage)
+ return result
+
+ def _record_token_metrics(
+ self, model_name: str, prompt_tokens: int, total_tokens: int
+ ) -> None:
+ """Record embedding token usage metrics for the given model.
+
+ Mirrors ``BaseChatModelSetup._record_token_metrics`` but records input-side
+ tokens only, since embeddings have no completion tokens. Counters are placed
+ under the same ``model`` key-value group used by chat metrics, so embedding
+ and chat usage for a model share one dimension.
+
+ Unlike the chat path, embedding calls do not run inside a plan action that
+ hands in a request-scoped metric group (vector-store, RAG, and direct calls
+ reach this setup directly), so the resource-bound metric group injected via
+ ``set_metric_group`` is used instead.
+
+ Parameters
+ ----------
+ model_name : str
+ The name of the model used
+ prompt_tokens : int
+ The number of prompt tokens
+ total_tokens : int
+ The total number of tokens reported by the provider
+ """
+ metric_group = self.metric_group
+ if metric_group is None:
+ return
+
+ model_group = metric_group.get_sub_group("model", model_name)
+ model_group.get_counter("promptTokens").inc(prompt_tokens)
+ model_group.get_counter("totalTokens").inc(total_tokens)
+
+ def _record_token_usage(self, token_usage: EmbeddingTokenUsage | None) -> None:
+ """Record the provider-reported embedding token usage, if any.
+
+ Called from ``embed_with_usage`` so direct calls and vector-store/RAG paths
+ are both covered without each provider repeating the recording.
+ """
+ if token_usage is None or not self.model:
Review Comment:
`not self.model` and Java's `model == null || model.isBlank()` (`BaseEmbeddingModelSetup.java:131`) do not agree when the model name is only spaces. I ran the PR's own code to check: with `model=' '` Python records under a group literally named `model= `, while Java skips it. With `model=''` both skip.
`AGENTS.md` asks for the Java, Python and YAML APIs to stay semantically aligned.
There is no Python test on this guard right now. Removing `or not self.model` leaves all 10 tests passing, while Java has `testEmbedWithUsageNullModelRecordsNothing` covering the same thing.
Is the whitespace case worth folding in? Something like this, in case it helps:
```python
if token_usage is None or not (self.model or "").strip():
```
##########
api/src/main/java/org/apache/flink/agents/api/embedding/model/python/PythonEmbeddingModelSetup.java:
##########
@@ -137,7 +137,10 @@ public EmbeddingResult<float[]> embedWithUsage(String text, Map<String, Object>
Map<String, Object> kwargs = new HashMap<>(parameters);
kwargs.put("text", text);
Object result = adapter.invoke(CALL_EMBED_WITH_USAGE, embeddingModelSetup, kwargs);
- return EmbeddingModelUtils.toSingleEmbeddingResult(result);
+ EmbeddingResult<float[]> embeddingResult =
+ EmbeddingModelUtils.toSingleEmbeddingResult(result);
+ recordTokenUsage(embeddingResult.getTokenUsage());
Review Comment:
I think this records the same tokens twice, for a Java agent using a Python embedding model.
`adapter.invoke(CALL_EMBED_WITH_USAGE, ...)` at `:139` calls into the Python `BaseEmbeddingModelSetup.embed_with_usage`, and this PR makes that method record too, at `embedding_model.py:159`. Then this line records again.
Both writes land on the same Java `SimpleCounter`. `setMetricGroup` here (`:177-181`) binds the Java field and forwards the same `FlinkAgentsMetricGroup` to the Python side, and `FlinkMetricGroup` and `FlinkCounter` just pass through to that Java object. Neither guard stops it, because both `model` fields read the same descriptor argument (`EmbeddingCrossLanguageAgent.java:68` sets it).
The other direction stays at one write: `JavaEmbeddingModelSetupImpl.embed_with_usage` (`java_embedding_model.py:174-182`) calls the Java resource without chaining to `super()`.
Nothing would catch it today either. I deleted both `recordTokenUsage(...)` lines and the whole `api` module still passed 385/385. `PythonEmbeddingModelSetupTest` builds the setup from a `@Mock ResourceDescriptor` (`:50`), so `getArgument("model")` comes back null and the recording returns early, and no metric group is ever bound.
So: which side should own the write for a cross-language resource? Dropping these two lines and letting the Python setup own it would match what `JavaEmbeddingModelSetupImpl` already does, unless there is a reason the Java wrapper needs its own. Either way, would a test with a real descriptor carrying `model` plus a bound group be worth adding, so these lines are covered?
##########
api/src/main/java/org/apache/flink/agents/api/embedding/model/BaseEmbeddingModelSetup.java:
##########
@@ -93,9 +94,52 @@ public String getModel() {
return model;
}
+ /**
+ * Record embedding token usage metrics for the given model on this setup's bound metric group.
+ *
+ * <p>Mirrors {@code BaseChatModelSetup#recordTokenMetrics} but records input-side tokens only,
+ * since embeddings have no completion tokens. Counters are placed under the same {@code model}
+ * key-value group used by chat metrics, so embedding and chat usage for a model share one
+ * dimension.
+ *
+ * <p>Unlike the chat path, embedding calls do not run inside a plan action that hands in a
+ * request-scoped metric group (vector-store, RAG, and direct calls reach this setup directly),
+ * so the resource-bound metric group injected via {@link #setMetricGroup} is used instead.
+ *
+ * @param modelName the name of the model used
+ * @param promptTokens the number of prompt tokens
+ * @param totalTokens the total number of tokens reported by the provider
+ */
+ public void recordTokenMetrics(String modelName, long promptTokens, long totalTokens) {
+ Preconditions.checkArgument(
+ modelName != null && !modelName.isBlank(), "Model name must not be null or blank.");
+ FlinkAgentsMetricGroup metricGroup = getMetricGroup();
Review Comment:
Reading the group bound to the resource here, and at `embedding_model.py:186`, is the pattern #859 reported and #861 moved chat off.
Here is what worries me. `Resource.metricGroup` is a single mutable field on an object that is cached and shared across the whole subtask. `RunnerContextImpl.getResource:485-495` rewrites it to the current action's group on every fetch. Actions yield to other keys while they run. So another action can rebind that field in between a `ctx.getResource(...)` call and the later `embedWithUsage`.
Chat takes a different route: it captures the group up front and passes it in (`ChatModelInvoker.java:121`, used at `:176`). There is a test guarding exactly that, `BaseChatModelSetupTokenMetricsTest.java:95-111`, which asserts the bound group is not used.
The vector-store path looks worse than racy. A vector store resolves its model through `ResourceContext` (`BaseVectorStore.java:96-101`), and that path never sets `metricGroup` at all, so it stays null there.
Given RAG does run inside an action, what would make passing the group in the way chat does hard here? And whichever way it settles, could a test pin it? None of the new tests bind two groups, so nothing would fail today if the choice were reversed.
##########
api/src/test/java/org/apache/flink/agents/api/embedding/model/BaseEmbeddingModelSetupTokenMetricsTest.java:
##########
@@ -0,0 +1,318 @@
+/*
+ * 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 of
Review Comment:
nit: `permissions of` looks like it wants to be `permissions and`. RAT accepts the file either way, so nothing is broken. It is the only Java file in the repo with that wording though. Worth folding into the compile fix?
--
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: issues-unsubscribe@flink.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org