Skip to content

[api][java][python] Record embedding token usage metrics #1047

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
hulincup wants to merge 2 commits into
base: main
Choose a base branch
from
Open

[api][java][python] Record embedding token usage metrics #1047

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
80869f0
[api][java][python] Record embedding token usage metrics
hulincup Aug 24, 2026
3dae4db
style: fix spotless violations in embedding token metrics
hulincup Aug 25, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
52 changes: 50 additions & 2 deletions api/src/main/java/org/apache/flink/agents/api/embedding/model/BaseEmbeddingModelSetup.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@

package org.apache.flink.agents.api.embedding.model;

import org.apache.flink.agents.api.metrics.FlinkAgentsMetricGroup;
import org.apache.flink.agents.api.resource.Resource;
import org.apache.flink.agents.api.resource.ResourceContext;
import org.apache.flink.agents.api.resource.ResourceDescriptor;
Expand Down Expand Up @@ -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),

@weiqingy weiqingy Aug 31, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

* 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();

@weiqingy weiqingy Aug 31, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@wenjin272 wenjin272 Sep 2, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

if (metricGroup == null) {
return;
}
FlinkAgentsMetricGroup modelGroup = metricGroup.getSubGroup("model", modelName);
modelGroup.getCounter("promptTokens").inc(promptTokens);
modelGroup.getCounter("totalTokens").inc(totalTokens);
}

/**
* Record the provider-reported embedding token usage, if any, onto this setup's bound metric
* group. Called from {@link #embedWithUsage} so direct calls and vector-store/RAG paths are
* both covered without each provider repeating the recording.
*/
protected void recordTokenUsage(@Nullable EmbeddingTokenUsage tokenUsage) {
if (tokenUsage == null || model == null || model.isBlank()) {
return;
}
recordTokenMetrics(model, tokenUsage.getPromptTokens(), tokenUsage.getTotalTokens());
}

/**
* Generate embeddings for the given text.
*
* <p>Token usage metrics are only recorded by {@link #embedWithUsage}; this method discards
* provider usage because it is not returned.
*
* @param text The input text to generate embeddings for
* @return An array of floating-point values representing the text embeddings
*/
Expand All @@ -117,7 +161,9 @@ public EmbeddingResult<float[]> embedWithUsage(String text, Map<String, Object>
Map<String, Object> params = this.getParameters();
params.putAll(parameters);
BaseEmbeddingModelConnection currentConnection = getConnection();
return currentConnection.embedWithUsage(text, params);
EmbeddingResult<float[]> result = currentConnection.embedWithUsage(text, params);
recordTokenUsage(result.getTokenUsage());
return result;
}

/**
Expand Down Expand Up @@ -146,6 +192,8 @@ public EmbeddingResult<List<float[]>> embedWithUsage(
Map<String, Object> params = this.getParameters();
params.putAll(parameters);
BaseEmbeddingModelConnection currentConnection = getConnection();
return currentConnection.embedWithUsage(texts, params);
EmbeddingResult<List<float[]>> result = currentConnection.embedWithUsage(texts, params);
recordTokenUsage(result.getTokenUsage());
return result;
}
}
10 changes: 8 additions & 2 deletions ...in/java/org/apache/flink/agents/api/embedding/model/python/PythonEmbeddingModelSetup.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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());

@weiqingy weiqingy Aug 31, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

return embeddingResult;
}

@Override
Expand All @@ -150,7 +153,10 @@ public EmbeddingResult<List<float[]>> embedWithUsage(
Map<String, Object> kwargs = new HashMap<>(parameters);
kwargs.put("text", texts);
Object result = adapter.invoke(CALL_EMBED_WITH_USAGE, embeddingModelSetup, kwargs);
return EmbeddingModelUtils.toBatchEmbeddingResult(result);
EmbeddingResult<List<float[]>> embeddingResult =
EmbeddingModelUtils.toBatchEmbeddingResult(result);
recordTokenUsage(embeddingResult.getTokenUsage());
return embeddingResult;
}

@Override
Expand Down
Loading
Loading