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

[zxs1633079383]

Linked issue: #858

Purpose of change

This adds provider-neutral token usage metrics for embedding model calls.

Embedding providers can now report input-side usage through the embedding
connection, while embedding setup records the usage under the current model
metric group after the request completes. Providers without token usage remain
no-op.

Covered paths in this PR:

• Java embedding setup token metrics.
• Python embedding setup token metrics.
• OpenAI-compatible Python embedding usage extraction.
• DashScope/Tongyi Python embedding usage extraction.
• Bedrock Titan Java embedding inputTextTokenCount extraction when present.

Tests

• cd python && uv run --python 3.12 --extra test pytest flink_agents/api/embedding_models/tests/test_token_metrics.py flink_agents/integrations/embedding_models/tests/test_openai_embedding_model.py flink_agents/integrations/embedding_models/tests/test_tongyi_embedding_model.py -q
• cd python && uv run --python 3.12 --extra lint ruff format --check flink_agents/api/embedding_models/embedding_model.py flink_agents/api/embedding_models/tests/test_token_metrics.py flink_agents/integrations/embedding_models/openai_embedding_model.py flink_agents/integrations/embedding_models/tests/test_openai_embedding_model.py flink_agents/integrations/embedding_models/tongyi_embedding_model.py flink_agents/integrations/embedding_models/tests/test_tongyi_embedding_model.py && uv run --python 3.12 --extra lint ruff check flink_agents/api/embedding_models/embedding_model.py flink_agents/api/embedding_models/tests/test_token_metrics.py flink_agents/integrations/embedding_models/openai_embedding_model.py flink_agents/integrations/embedding_models/tests/test_openai_embedding_model.py flink_agents/integrations/embedding_models/tongyi_embedding_model.py flink_agents/integrations/embedding_models/tests/test_tongyi_embedding_model.py
• export JAVA_HOME=/usr/local/opt/openjdk@17/libexec/openjdk.jdk/Contents/Home; export PATH="$JAVA_HOME/bin:$PATH"; mvn --batch-mode --no-transfer-progress -pl api -Dtest=BaseEmbeddingModelSetupTokenMetricsTest test
• export JAVA_HOME=/usr/local/opt/openjdk@17/libexec/openjdk.jdk/Contents/Home; export PATH="$JAVA_HOME/bin:$PATH"; mvn --batch-mode --no-transfer-progress -pl api,integrations/embedding-models/bedrock -am spotless:check -Dspotless.skip=false
• export JAVA_HOME=/usr/local/opt/openjdk@17/libexec/openjdk.jdk/Contents/Home; export PATH="$JAVA_HOME/bin:$PATH"; mvn --batch-mode --no-transfer-progress -pl integrations/embedding-models/bedrock -am -DskipTests compile
• git diff --check

API

This adds EmbeddingTokenUsage and protected token-usage recording helpers for
embedding model connections. Existing embed(...) return types are unchanged.

Documentation

• doc-needed
• doc-not-needed
• doc-included

[weiqingy]

Thanks for taking this on — A few things inline.

1. The two correctness issues I flagged inline (the Bedrock batch losing metrics, and a thrown embed() skipping the consume) share one root: the per-thread side channel — ThreadLocal in Java, ContextVar in Python — assumes the connection records usage on the same thread the setup later consumes it on, and that nothing throws between record and consume. Both assumptions hold for the single-text, single-thread happy path, but Bedrock's batch fans work out to embedPool workers (record lands on the wrong thread), and a record-then-throw connection leaves the slot populated for the next call. Chat sidesteps both by carrying usage on the response object and recording at the action layer (BaseChatModelSetup.java:118), never a thread-bound channel. Would returning usage alongside the embedding result — a small result-plus-usage value, or usage attached to the response, consumed by the setup from that returned value — be viable here? That would make batch/pooled connections and the exception path correct by construction rather than by convention. I realize the current design deliberately keeps the embed(...) return types unchanged for API compatibility, so this is a real tradeoff, not a free win — curious how you're weighing it.

[weiqingy]

This records into the current thread's ThreadLocal, which is right for the single-text path — but this same connection's embed(List) fans a batch of size > 1 out to embedPool workers (CompletableFuture.supplyAsync(() -> embed(text, parameters), embedPool)), so each per-text recordTokenUsage lands on a worker thread while BaseEmbeddingModelSetup.embed(List) consumes on the calling (mailbox) thread. For every Bedrock batch of size > 1 the consume then returns null and the tokens are silently dropped (exactly the RAG add(List<Document>) path #858 targets), while the workers' ThreadLocals are never consumed and — since recordTokenUsage accumulates onto any existing value — grow across batches. The size <= 1 path stays on the calling thread, so it's fine. This is the batch symptom of the side-channel root in the top-level comment; would carrying usage back with the result close it here?

[zxs1633079383]

Resolved in e023346 by removing the ThreadLocal side channel. BedrockEmbeddingModelConnection.embedWithUsage(List) now collects EmbeddingResult<float[]> from the worker futures and merges token usage on the caller path. BedrockEmbeddingModelTest.testBatchEmbeddingAggregatesTokenUsage covers the batch size > 1 path.

[weiqingy]

There's no try/finally around embed() and the consume, so if embed() throws after the connection already recorded usage, consumeTokenUsage() is skipped and the value persists on the thread. Bedrock hits this concretely: BedrockEmbeddingModelConnection records at line 139 before result.get("embedding") can return null and NPE at line 142, so a failed request leaves {tokens, tokens} in the ThreadLocal, and the next successful embed() on that thread folds those tokens in via the accumulate branch — inflating the metric with a plausible-looking number. The List overload at line 132 has the same shape, and embedding_model.py mirrors it. Would wrapping the consume in finally on both overloads (both languages) close the gap? The result-plus-usage change in the top-level comment would remove the need entirely, so worth deciding that direction first.

[zxs1633079383]

Resolved via the result-plus-usage path rather than try/finally. BaseEmbeddingModelSetup records token metrics only after embedWithUsage(...) returns a successful result, so a provider exception has no side-channel slot to leak into the next call. Added Java and Python no-leak regression tests.

[weiqingy]

Because consumeTokenUsage() clears the slot after every embed(), in the happy path current is always null and this branch never runs — it only fires when a prior consume was skipped (the leak conditions above), and when it does it adds to the stale value rather than resetting, turning a dropped metric into a wrong one that's harder to spot. If the side channel stays, is overwrite (set, not +=) the safer default here than accumulate? I looked for a real multiple-record-per-consume case that would motivate the += and couldn't find one — Bedrock's size <= 1 loop runs at most once — but if there is one I've missed, it'd be worth a comment here since the branch reads as intentional.

[zxs1633079383]

Removed with the side channel. There is no per-thread accumulate branch now; provider usage is returned with the embedding result, and metric accumulation only happens in the Flink counters after setup receives that result.

[weiqingy]

The mocks here are honest, so this isn't about the existing assertions — it's the failure modes that aren't reached. Despite its name, testEmbeddingTokenMetricsAccumulateAcrossRequests makes two separate embed() calls with a consume clearing between them, so it proves the Flink Counter is monotonic but never enters the if (current != null) branch in BaseEmbeddingModelConnection. The paths most likely to regress are untested: a record-then-throw followed by a successful embed() (asserting the second call's metric is not inflated), and the Bedrock batch path. Are those two cases worth adding? They're the ones that would catch the correctness issues above if the side channel stays.

[zxs1633079383]

Added the two regression cases called out here: prior provider failure followed by success does not inflate metrics in Java/Python setup tests, and Bedrock batch size > 1 aggregates usage from worker results. The existing cross-request test now only proves Flink counter monotonicity.

[zxs1633079383]

Updated in e023346.

Design change:

• Removed the Java ThreadLocal and Python ContextVar token-usage side channel.
• Added a result-plus-usage path: Java EmbeddingResult<T> / embedWithUsage(...), Python EmbeddingResult / embed_with_usage(...).
• Kept existing public embed(...) return types unchanged; default wrappers return embeddings with no usage.
• BaseEmbeddingModelSetup now records metrics only from a successful returned result.
• Bedrock batch embedding now merges token usage from worker results on the caller path.

Regression coverage added/updated:

• Java/Python usage recorded and no-usage no-op cases.
• Java/Python provider-failure-then-success cases, verifying no stale usage leaks into the next success.
• Bedrock batch size > 1 usage aggregation.
• OpenAI/Tongyi usage extraction remains covered.

Local verification:

• uv run --python 3.12 --extra test pytest flink_agents/api/embedding_models/tests/test_token_metrics.py flink_agents/integrations/embedding_models/tests/test_openai_embedding_model.py flink_agents/integrations/embedding_models/tests/test_tongyi_embedding_model.py -q
• uv run --python 3.12 --extra lint ruff format --check ... && uv run --python 3.12 --extra lint ruff check ...
• mvn --batch-mode --no-transfer-progress -pl api -Dtest=BaseEmbeddingModelSetupTokenMetricsTest test
• mvn --batch-mode --no-transfer-progress -pl integrations/embedding-models/bedrock -am -Dtest=BedrockEmbeddingModelTest -Dsurefire.failIfNoSpecifiedTests=false test
• git diff --check

[zxs1633079383]

Hi @weiqingy, just a gentle follow-up when you have a chance.

I updated this PR in e023346 to address the root correctness issue from your inline review:

• Removed the Java ThreadLocal and Python ContextVar token-usage side channel.
• Added a result-plus-usage path for embeddings while keeping the existing public embed(...) return types unchanged.
• Moved metric recording to setup after a successful returned result.
• Added regression coverage for provider-failure-then-success and Bedrock batch usage aggregation.

The PR is currently clean and all checks are green. If you still see any correctness/API concerns, I am happy to iterate; otherwise I would appreciate another review when convenient. Thanks again for the detailed review.

[joeyutong]

Thanks @weiqingy for the detailed review.

One additional concern: this PR records embedding token metrics inside BaseEmbeddingModelSetup.embed(...). Similar to chat model setup, the setup API itself is not inherently single-threaded, while the metric group assumes single-threaded access when resolving subgroups/counters.

If the same setup resource is invoked concurrently, setup-level metric recording can race during metric lookup/creation. Should embedding token recording stay at a caller/action-thread boundary where the single-threaded access guarantee is explicit?

[zxs1633079383]

Updated in cf989b4 to address the request-scoped metric-group concern.

What changed:

• BaseEmbeddingModelSetup.embed(...) no longer records metrics from the setup resource's mutable bound metric group. The plain embedding API now only returns embeddings, so direct/shared setup calls cannot race through Resource.metricGroup.
• Added explicit result/recording APIs: setup-level embedWithUsage(...) keeps provider usage with the returned embedding result, and recordTokenMetrics(requestMetricGroup, usage) records only against a caller-provided request/action metric group.
• Java and Python vector-store auto-embedding now capture the vector store's current request metric group at the caller boundary and pass that explicit group into embedding metric recording.
• Added regression tests proving metrics use the request-scoped group, skip when the group is absent, and do not read a later mutable bound metric group. Existing provider-failure/no-usage/batch usage coverage remains.

Local verification:

• ./tools/lint.sh --check
• mvn --batch-mode --no-transfer-progress -pl api -Dtest=BaseEmbeddingModelSetupTokenMetricsTest test
• mvn --batch-mode --no-transfer-progress -pl integrations/embedding-models/bedrock -am -Dtest=BedrockEmbeddingModelTest -Dsurefire.failIfNoSpecifiedTests=false test
• uv run --python 3.12 --extra test pytest flink_agents/api/embedding_models/tests/test_token_metrics.py flink_agents/integrations/embedding_models/tests/test_openai_embedding_model.py flink_agents/integrations/embedding_models/tests/test_tongyi_embedding_model.py -q
• uv run --python 3.12 --extra lint ruff format --check ... && uv run --python 3.12 --extra lint ruff check ...
• git diff --check

[joeyutong]

@zxs1633079383 Thanks for the update. Explicitly passing the request-scoped metric group fixes the mutable setup-bound metric group issue, but I think the thread-safety concern still remains.

The current path still records metrics from the caller thread. For example, vector-store auto-embedding records inside query(...) / auto-embedding, which are ordinary resource APIs and may be called from durableExecuteAsync or other concurrent paths. The metric group assumes single-threaded access: subgroup/counter lookup is backed by lazy maps, and Flink's default MetricGroup.counter(name) uses SimpleCounter, which is explicitly not thread-safe.

Without changing the broader resource abstraction, one possible direction is to pre-resolve the embedding metric handles at the existing action-thread binding point, e.g. from RunnerContextImpl.getResource(...) when it calls Resource.setMetricGroup(...) / Python set_metric_group(...). BaseEmbeddingModelSetup could resolve the model subgroup and the promptTokens / totalTokens counters there, and later recordTokenMetrics(...) / vector-store auto-embedding would use only this pre-resolved recorder instead of calling getSubGroup(...) / getCounter(...) again. The recorder should also avoid updating SimpleCounter concurrently, or use a thread-safe counter implementation such as ThreadSafeSimpleCounter.

Do you have any thoughts on where this should fit in the current abstraction?

[zxs1633079383]

Thanks @joeyutong, I agree with the concern.

cf989b4 fixes the previous mutable setup-bound metric-group problem, but it still lets ordinary caller threads perform metric lookup/recording. For vector-store auto-embedding, that means query(...) / auto-embedding can still call into lazy subgroup/counter resolution, and the counter update itself may happen from a concurrent resource API path. So this should not be treated as fully solved by just passing the request-scoped metric group explicitly.

Before I make another implementation change, I think the right next step is to settle the abstraction boundary. Two possible shapes I am considering:

1. Introduce an EmbeddingTokenMetricRecorder (Java) / equivalent lightweight recorder on the Python side. The recorder would be created or refreshed at the existing action-thread resource binding point, where the action/request metric group is known. Later embedding/vector-store paths would only call recorder.record(usage) and would not perform getSubGroup(...) / getCounter(...) lazy lookup from arbitrary caller threads.
2. Keep the public recordTokenMetrics(...) shape, but make it use pre-resolved counter handles captured at binding time rather than resolving counters on each recording call.

Either way, I think the recorder should avoid updating a non-thread-safe SimpleCounter concurrently. If the record path can still be reached from concurrent resource calls, then the handle should either wrap a thread-safe counter implementation such as ThreadSafeSimpleCounter, or the design should move the actual counter update back to a single-threaded action/mailbox boundary.

This also seems related to the broader direction discussed in #861: the cached Resource should stay a stable capability, while per-action metric state belongs to a request/use context or bound recorder. I will hold off on rebasing/churning the PR until this direction is clear, then update the patch and rerun the branch against the latest origin/main.

[joeyutong]

Thanks for laying out the options. After thinking more about this, I do not have a fully comprehensive solution within the current resource abstraction either.

I agree that a recorder-style direction could be the right long-term shape, but I do not think it should be embedding-specific. This seems more like a framework-level metric/effect abstraction: resource APIs may run from arbitrary caller threads, while actual MetricGroup access should happen only at a safe action/mailbox-thread boundary. Designing that properly is probably beyond the scope of this PR.

The second option is more focused, but if the record path can still be called concurrently, it would either still race on the default SimpleCounter, or require a special thread-safe counter path. That seems to weaken the current metric-group single-threaded access contract.

Given that, my suggestion is to narrow this PR to provider usage extraction and usage-returning APIs only, without promising automatic metric recording from ordinary resource APIs such as vector-store auto-embedding. We can track the framework-level metric recording abstraction separately and solve automatic recording there.

@weiqingy @zxs1633079383 , what do you think?

[weiqingy]

@joeyutong Agreed with narrowing.

One thing worth pulling on before the recorder shape settles: I'm not sure ThreadSafeSimpleCounter would actually close the race. Flink's AbstractMetricGroup.addMetric / addGroup are already synchronized (this) — the unsynchronized lazy maps look like ours, FlinkAgentsMetricGroupImpl.java:42-50 (plain HashMaps, check-then-put in getSubGroup / getCounter). If that reading is right, a thread-safe counter would fix inc() and leave the lookup racing, so pre-resolving the handles at binding time may be doing more of the work than the counter swap. It's also @Internal in flink-metrics-core. Curious whether you see it the same way.

@zxs1633079383 Also, cf989b4 is currently red — 6 checks (5× it-python, 1× cross-language). The "all checks green" note was about the earlier e023346. Both failures trace to two lines in the Python vector store; left inline.

[weiqingy]

This now calls self._get_embedding_model() unconditionally before the loop, only so request_metric_group = self.metric_group can be captured once. On origin/main the resolution was lazy, inside the branch that actually needs it:

for doc in documents:
    if doc.embedding is None:
        doc.embedding = self._get_embedding_model().embed(doc.content)

The effect is that a store which never auto-embeds — mem0 embeds server-side — now raises on add/get/delete. That's the 13 failures in test_mem0_vector_store.py: TypeError: No embedding model configured on this vector store. Could the resolution move back inside the if doc.embedding is None branch, and the metric group be read only on the path that embeds?

[weiqingy]

The call moved from embed(...) to embed_with_usage(...), and that quietly changes the extension point. BaseEmbeddingModelSetup.embed() delegates to embed_with_usage(), but embed_with_usage (embedding_model.py:157) goes straight to self._get_connection() — so a subclass that overrides embed() is no longer on the vector-store path. That's the 7 test_chroma_vector_store.py failures: TypeError: Expect BaseEmbeddingModelConnection, but is str, because MockEmbeddingModel.embed() is bypassed and the mock's connection is the string 'mock'. Beyond the test mock, this is a public-API behavior change — is embed() still meant to be an override point? If so, having the vector store go through it (with usage surfaced some other way) would keep that contract intact.

[joeyutong]

I agree that pre-resolving the metric handles at binding time addresses the lazy-map race, as I mentioned above. My remaining concern is that metric recording would still update the pre-resolved counters from arbitrary caller threads. Although this avoids concurrent MetricGroup lookup, it still requires the counters to support concurrent updates and moves metric mutation outside the single-threaded action/mailbox boundary.

I am not sure whether that is consistent with the intended single-threaded access contract of the current metric abstraction, so I would prefer not to pursue this direction in this PR for now. This is also why I suggested narrowing the PR to usage extraction and usage-returning APIs.

[zxs1633079383]

@joeyutong @weiqingy Thanks, I agree with narrowing this PR. I pushed 5af8560, which removes automatic embedding token metric recording and leaves this PR focused on provider usage extraction plus the EmbeddingResult / embedWithUsage(...) and Python equivalent APIs.

The vector-store paths are back to the existing embed(...) API. That restores the public override point, and model resolution is again lazy inside the missing-embedding branch, so Mem0 stores without a configured embedding model keep their prior behavior. This directly addresses the Python CI failures you pointed out.

I also agree that pre-resolving handles would only address the lazy lookup race, not the broader question of mutating counters from arbitrary resource threads. I will not pursue a ThreadSafeSimpleCounter or embedding-specific recorder in this PR; a framework-level safe metric/effect boundary should be designed separately.

Local verification for the affected paths is green: 49 Python tests (including Chroma and Mem0 vector stores), the API embedding-result test, and the Bedrock batch-usage test. CI should now rerun on the narrowed patch. Thanks again for steering this toward a cleaner boundary.