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: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
