[api][plan][integrations] Record built-in chat token metrics outside the async call boundary#712
Merged
Merged
Conversation
…the async call boundary The Java built-in chat action recorded token metrics from inside the durableExecuteAsync callable: each chat-model connection called BaseChatModelConnection.recordTokenMetrics(...) within chat(), which runs on a durable-execution pool thread. The runtime metric group is meant to be used from the operator/mailbox (action) thread. Mirror the Python path: connections now stash model_name/promptTokens/ completionTokens into the response ChatMessage.extraArgs, and ChatModelAction records after the durable call returns via a new BaseChatModelSetup.recordTokenMetrics(...). The setup's bound metric group is the action metric group, so the emitted metric path and counter names are unchanged. This also captures token metrics for Python-backed models invoked from the Java action, which previously recorded none. Recording matches Python's guard: model name non-empty and both token counts greater than zero; values are read as Number and converted with longValue() to tolerate Integer/Long across the Pemja bridge and durable recovery.
Clarify that the metric groups returned by getAgentMetricGroup() and getActionMetricGroup() must only be accessed from the operator/mailbox (action) thread, not from inside a durableExecute / durableExecuteAsync callable, which runs on a separate thread pool.
github-actions
Bot
added
doc-not-needed
2 tasks
Contributor
|
@xintongsong Should this PR be backported to the release-0.2 branch? |
Contributor
|
Exactly. That's what I was wondering by asking for the |
3 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Linked issue: #706
Purpose of change
The Java built-in chat action recorded token metrics from inside the durable async callable: each chat-model connection called
BaseChatModelConnection.recordTokenMetrics(...)within itschat()method, which runs on a durable-execution pool thread. The runtime metric group is meant to be used from the operator/mailbox (action) thread, so touching it from the callable crosses that boundary.This brings Java to the same execution boundary as Python. Mirroring
chat_model_action.py:model_name/promptTokens/completionTokensinto the responseChatMessage.extraArgsinstead of recording insidechat().ChatModelActionrecords after the durable call returns (before structured-output reassignment, which would drop the keys) via a newBaseChatModelSetup.recordTokenMetrics(...)— the Python-parity record site (chat_model._record_token_metrics). The setup's bound metric group is the action metric group, so the emitted metric path and counter names are unchanged.BaseChatModelConnection.recordTokenMetrics(...)and its now-deadconnection.setMetricGroup(...)forwarding are removed.RunnerContextmetric-group getter Javadoc now documents that the returned group must only be accessed from the operator/mailbox thread, not inside a durable callable.Recording matches Python's guard exactly: the model name must be non-empty and both token counts greater than zero; values are read as
Numberand converted withlongValue()to tolerateInteger/Longacross the Pemja bridge and durable recovery.This also fixes a latent gap: Python-backed chat models invoked from the Java action previously recorded no token metrics (the path bypasses the Java connection's recording); they are now captured once.
Tests
BaseChatModelSetupTokenMetricsTest), mirroring Python'stest_token_metrics.py: records to the per-model sub-group counters, no-ops without a metric group, separate sub-groups per model,getResourceType() == CHAT_MODEL.ChatModelActionTestwith cases for the newrecordChatTokenMetricshelper: records once when all keys are present and positive;Integer-typed token values still recorded viaNumber.longValue(); skips when a key is missing or non-numeric; skips when a token is0or the model name is empty (Python parity)../tools/build.sh -j,./tools/ut.sh -j, and./tools/lint.sh -call pass.API
Adds
public BaseChatModelSetup.recordTokenMetrics(String, long, long)(the Python-parity record site) and removes the previouslyprotectedBaseChatModelConnection.recordTokenMetrics(...). No change to public configuration, event, or resource APIs. Emitted metric names and paths are unchanged.Documentation
doc-neededdoc-not-neededdoc-included