Repository navigation
feat: report Pi model usage per invocation - #419
Open
mnajafian-nv wants to merge 84 commits into
Open
mnajafian-nv wants to merge 84 commits into
mnajafian-nv wants to merge 84 commits into
Conversation
The Hermes adapter returned no usage, so every Hermes run read unknown tokens. Hermes keeps token counters per session, and the adapter keeps one session per runtime, so each turn's result holds session totals. Each invocation now reports the change since the previous invocation's totals: input_tokens with cache reads and cache writes included, cached_input_tokens, output_tokens and total_tokens. The baseline is 0 for a new Hermes agent, is cleared while an invocation runs, and is set again only from a returned result. After an invocation that raised, was cancelled or timed out, or returned no counters, the next invocation reports no counters instead of a carried-over or whole-session total. A counter that went down is unknown for that invocation. This follows the rule the Codex adapter already uses for thread totals. Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
The Pi adapter returned no usage although Pi reports token counts for every call it bills. Pi reports usage per call, so the adapter now sums what Pi bills during the invocation, the way Pi's own usage totals count it: assistant messages, tool results that report usage, finished compaction summaries and usage entries. The result holds input_tokens with cache reads and cache writes included, cached_input_tokens, output_tokens and total_tokens. A call without a complete usage report, or a compaction that was aborted or finished without usage, leaves the invocation without usage, so a partial sum is never reported as the invocation's total. Pi's cost figures come from its model catalog's rates, so they are not reported as cost_usd. Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
Pi keeps no usage for a provider call it stopped, so an aborted assistant message reports zero tokens although its input may already have been billed. Adding that zero made a cancelled invocation report a complete-looking total that left out the aborted call. An aborted assistant message now leaves the invocation without usage, the same rule as an aborted compaction. The comments and README now say "records" instead of "bills", because the adapter sums what Pi recorded, and the README lists every case in which the result has no usage. Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
A Hermes counter above the 64-bit range would make AgentUsage raise and turn a finished invocation into an error, so such a counter is now ignored like any other invalid value, matching the Codex adapter. Hermes records auxiliary calls such as context-compression summaries outside its session counters, so the README now says the usage covers the main model's calls. A test now checks that the invocation after a cancelled one reports no usage, which is the rule that keeps a cancelled turn's tokens out of the next result. Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
Hermes counts a provider response without usage as an API call but adds no tokens to its session counters, so the adapter reported such an invocation as a measured 0, and a mixed invocation as a complete subtotal. Count those responses through the context compressor's note_usage_less_response hook, which both Hermes runtimes call, and report no usage for an invocation that had one. A response that reported zero tokens still reports 0. Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
Pi fills a model call's usage with zeros and overwrites it only when the provider reports usage, so a response without usage was summed as a measured 0. Treat a model call whose counts are all zero as unreported. Pi's compaction summary retry also returns only the final attempt's usage, and failed attempts emit no usage, so a retried compaction undercounted. Leave the invocation's usage unknown in both cases. Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
When a compaction cuts inside a turn, Pi summarizes the earlier history and the turn prefix in two model calls and reports only their combined usage on compaction_end. A summary call without a usage report is zero-filled, so the combined usage can look complete while one call is missing. Pi exposes the compaction plan, including isSplitTurn, only to extensions through session_before_compact. Register a hidden adapter extension that forwards that event to the running prompt, and leave the prompt's usage unknown unless the plan says the compaction is not split. State in the README and code comments that an all-zero provider report is discarded together with an unreported call. Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughHermes and Pi adapters now report token usage per invocation. Hermes uses validated session-counter deltas and provider observations. Pi aggregates provider, tool, and compaction usage before mapping it to invocation results. ChangesHermes token usage
Pi token usage
Priority: ➖ Normal Estimated code review effort: 4 (Complex) | ~60 minutes Change: Feature Sequence Diagram(s)sequenceDiagram
participant HermesRuntime
participant HermesAgent
participant AgentRunResult
HermesRuntime->>HermesAgent: Invoke a turn
HermesAgent-->>HermesRuntime: Return response and session counters
HermesRuntime->>HermesRuntime: Validate counters and provider observations
HermesRuntime-->>AgentRunResult: Include available usage
sequenceDiagram
participant PiSession
participant PiSdkSessionHandle
participant PiPromptOutcome
participant PiRuntime
PiSession->>PiSdkSessionHandle: Emit provider and compaction events
PiSdkSessionHandle->>PiSdkSessionHandle: Aggregate reported counts
PiSdkSessionHandle-->>PiPromptOutcome: Return prompt usage
PiRuntime->>PiRuntime: Normalize usage to AgentUsage
PiRuntime-->>PiRuntime: Attach usage to invocation outcome
Merge Risk: 🟡 Moderate · up to
|
| Check name | Status | Explanation | Resolution |
|---|---|---|---|
| Docstring Coverage |  |
Docstring coverage is 37.04% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 108 functions across 9 files. (2 skipped:… | Write docstrings for the functions missing them to satisfy the coverage threshold. |
-
Check name Status Explanation Linked Issues check Check skipped because no linked issues were found for this pull request. Out of Scope Changes check Check skipped because no linked issues were found for this pull request. Title check The title follows Conventional Commits format, uses an allowed lowercase type, stays under 72 characters, and accurately describes the Pi usage-reporting change. It does not mention the Hermes changes… Description check The pull request description includes an overview, reviewer starting point, detailed implementation context, validation details, a Related Issues section using the Relates to keyword, and both require… Full details: Title check
-
Explanation
The title follows Conventional Commits format, uses an allowed lowercase type, stays under 72 characters, and accurately describes the Pi usage-reporting change. It does not mention the Hermes changes, but it remains related to a main part of the changeset.
Full details: Description check
-
Explanation
The pull request description includes an overview, reviewer starting point, detailed implementation context, validation details, a Related Issues section using the Relates to keyword, and both required confirmation checkboxes.
-
Full details: Docstring Coverage
-
Explanation
Docstring coverage is 37.04% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 108 functions across 9 files. (2 skipped: 2 unsupported.)
Resolution
Write docstrings for the functions missing them to satisfy the coverage threshold.
✨ Finishing Touches
-
🧪 Generate unit tests (beta)
-
- Create a new PR
-
- Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Comment @coderabbitai help to get the list of available commands.
|
Fern docs preview: https://nvidia-preview-pull-request-419.docs.buildwithfern.com/nemo/fabric |
Collaborator
|
Partial provider usage is still reported as measured zeros in both adapters. I reproduced this on f734edf with the pinned SDKs: a response with Hermes's Can we preserve whether each count was reported before SDK normalization, leave missing dimensions and derived totals unknown, and add partial-report cases alongside the explicit-zero tests? The current checks handle an entirely absent report, but can't distinguish an omitted dimension from a reported zero once the SDK has normalized it. |
Bring in the Claude adapter's per-invocation cost and cache token usage, plus the other main changes since this branch was cut. No conflicts. Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
A provider response can report usage without one of its counts, such as prompt_tokens with no completion_tokens. Hermes' normalize_usage and Pi's provider parsers fill the missing count with 0, so both adapters reported it as a measured 0 and derived total_tokens from it. The usage-less checks only caught a response with no usage at all. Both adapters now read the provider's raw usage before the SDK normalizes it and keep only the counts it reported. Hermes wraps turn_response_check.record_response_usage for the agent it owns and counts responses whose usage lacks an input or output count. Pi wraps the agent's stream function, reads each call's raw stream events through onProviderStreamEvent for every supported wire protocol, and keys the result by the call's final usage object, which covers assistant messages and compaction summaries. An omitted input count leaves out input_tokens, cached_input_tokens and total_tokens; an omitted output count leaves out output_tokens and total_tokens. A reported zero stays 0, which also lets Pi keep a provider report of all zeros that it used to discard. Tests cover a prompt-only report, a completion-only report, reported zeros and a partial response inside a multi-call invocation, through Hermes' real usage accounting and the real Pi SDK against a local streaming endpoint. Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
Anthropic Messages reports input_tokens without cache reads and writes, and both Hermes and Pi add a missing cache field back as 0. A response with input_tokens but no cache fields therefore still produced a measured input count. Input is now reported only when input_tokens and both cache fields are present, the rule the Claude adapter uses. When a provider sends no cache details, both SDKs count 0 cached tokens. cached_input_tokens is now left out in that case. For the protocols whose input count includes the cache, input_tokens and total_tokens stay known. The Hermes adapter replaced turn_response_check.record_response_usage on every start and never restored it, and nothing showed that Hermes' turn loop still calls that name. The wrapper is now installed once per process and dispatches to observers registered per agent, and stop() removes the agent's observer. Start fails with a lifecycle error unless record_response_usage is still Hermes' turn_usage function and check_api_response still calls it by that name. The mock endpoint takes a usage object per scenario, and the Hermes e2e test sends a partial report through Hermes' real turn loop. For Pi, only a message_delta reports the final Anthropic output count, because message_start carries a placeholder. Google output counts as reported when thoughtsTokenCount is present without candidatesTokenCount. The protocol comes from the model that served the stream event, and a protocol outside Pi's built-in ones reports no counts. PiTokenUsage now holds input with the cache included, so the cache reads can be unknown while input stays known. The READMEs now say which Hermes paths the check covers. Hermes' native Gemini and Bedrock transports, and its Anthropic transport when a stream omits usage, fill missing counts before the adapter sees them, and its Codex app-server runtime records usage elsewhere. They also say that a Pi cache warm-up during an invocation leaves the result without usage: the adapter receives the warm-up's raw events but cannot tie them to the warm-up's usage entry. Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
Hermes' check_api_response returns on a reply that ends with finish_reason length or content_filter before it calls record_response_usage, so the provider bills that call but Hermes' session counters never include it. A turn that hit max_tokens and continued therefore reported measured input, output and total tokens that left out the truncated call. The adapter's process-wide wrapper now also covers the two handlers check_api_response calls in those cases, recover_from_truncation and handle_content_policy_refusal, with the same once-per-process install and start-time check. When either handler returns before usage is recorded, the response counts like a response without usage, and the invocation reports no usage. A truncation that falls through to record_response_usage is still counted. The mock endpoint takes a finish reason per scenario, and the Hermes e2e test sends a length reply and a content_filter reply through Hermes' real turn loop. A terminal content_filter result already had no token counters, so that case reported no usage before this change; the fix matters when a configured fallback answers after the refusal. Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
Bedrock Converse reports input without the cache, like Anthropic Messages, but its API documents cacheReadInputTokens and cacheWriteInputTokens as optional. The adapter therefore takes inputTokens alone as the input count, and the README now says why. Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
The Deep Agents adapter summed usage_metadata from the AI messages that reached graph state through the updates stream. Three cases came out wrong: - Deep Agents adds SummarizationMiddleware to every agent and subagent. Its summary call returns only text, so the summary model's message and usage never enter graph state, and an invocation that summarized under-reported input, output and total tokens. - When the invocation raised, including GraphRecursionError, the adapter returned no messages, so usage was None even though model calls had finished. - A count was summed over the calls that reported it and returned as known. ChatOpenAI also turns an omitted Chat Completions count into 0 and derives a missing total, so a partial provider report came out as measured zeros. The adapter now records usage from LangChain's chat-model callbacks, which fire for every chat-model call that runs in the adapter process: the main agent, declarative subagents and summarization. Agent Protocol subagents run remotely and are not included. A new recorder is created for each invocation, so earlier turns are never counted again. It outlives an exception, so a failed invocation still reports the fields that every counted call reported. Each token count and the cost are known only when every counted call in the invocation reported them, and an explicit zero stays a known zero. A request the provider rejects with a 4xx status (other than 408) contributes no usage. A 5xx response, timeout, cancellation or unfinished call leaves the invocation's counts and cost unknown, because the provider may have done work without reporting usage. OpenAI SDK retries happen inside one LangChain call, so a 5xx the SDK retries to success is not visible to the recorder, and only the final response's reported usage counts. When a response carries OpenAI-shaped raw usage, the raw counts decide what was reported: an omitted count stays unknown, the provider's total is used when present, and a derived total counts only when input and output were both reported. The graph tests run the real Deep Agents graph and ChatOpenAI against a scripted Chat Completions endpoint for summarization, parallel task subagents, 400, 408, 503 and timeout failures, the recursion limit, partial reports and explicit zeros. Separate ChatOpenAI.with_retry tests cover 429 then success and 503 then success. The existing usage tests now produce their messages through a LangChain chat-model call so the callbacks fire. Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
Hermes' check_api_response hands a response that fails validation to retry_invalid_response and returns before record_response_usage. The provider can still bill that call, and a successful retry then produced a measured invocation total that left it out. The adapter's process-wide wrapper now covers retry_invalid_response next to the truncation and refusal handlers, with the same start-time check. Unless the handler falls through to usage recording, the response counts like a response without usage and the invocation reports no usage. The start check now also rejects a handler name that the installed Hermes does not define. The Hermes e2e test answers with the router timeout sentinel that Hermes' chat-completions transport rejects in validate_response, with usage that bills the prompt and no completion tokens, followed by a successful retry through Hermes' real turn loop. It fails without the wrapper. The content_filter e2e case is replaced by a unit test of the refusal handler's dispatch with a fallback reply: a terminal refusal result carries no counters, so the e2e case passed without the wrapper, and Fabric's Hermes config enables no fallback. Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
The Hermes README still said the adapter wraps only record_response_usage and listed only truncated and refused replies. The adapter also observes retry_invalid_response, and a streamed call that Hermes retries inside its API call, such as an empty stream, never reaches these handlers, so its billed usage is not observed. The README now names all three handlers and states that limit. Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
Contributor
Author
|
Thanks, I reproduced both cases. For every raw report that reaches the adapter's observer, an omitted input or output count now stays unknown in both adapters, and so does the total; a reported zero stays 0. Anthropic-shaped input needs both cache fields, as in the Claude adapter. The Hermes adapter wraps |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
The Hermes adapter's response observer skipped a falsy usage object, so an empty usage dict never reached the omitted-count check. Hermes' record_response_usage also treats an empty usage object as missing and calls note_usage_less_response for it, so the invocation's usage was already unknown through the compressor hook. The observer now checks for None explicitly and counts an empty usage object as a report with input, cached input and output left out, so the result no longer depends on that hook alone. The new test runs an empty usage object through Hermes' real record_response_usage and expects unknown usage, followed by a counted invocation. The Pi README said a tool result without a complete usage report leaves the result without usage. A tool result with no usage field adds nothing; only a usage object with a missing or invalid count makes the invocation's usage unknown. The README now says both. Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
Use the system Bash path so macOS tests exercise Bash 3.2 even when PATH contains a newer Bash. This keeps the empty-array regression meaningful in CI. Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
This branch was successfully deployed
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
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.
Overview
Pi's final token counters can hide fields that the provider omitted and attempts that failed before a retry succeeded. Treating those counters as complete usage can make a run look cheaper than the available evidence supports.
Collect usage from supported provider responses before Pi replaces missing fields with zero. Keep available counts when a prompt fails, and identify work whose usage cannot be fully accounted for.
Dependencies: #442 supplies the shared accounting contract and normalizer. The diff against
mainincludes unmerged prerequisites.Details
The collector reads raw usage, records fetch attempts, and observes compaction and tool-provided usage. Shared accounting in #442 normalizes the counts. Cache reads and writes remain independently available. An HTTP error alone never establishes zero usage.
Each prompt gets its own accounting snapshot. Returned snapshots do not change, and later reports remain associated with the prompt that started the work. The fetch wrapper keeps a stable identity and preserves a caller-supplied fetch for cache-warming requests.
Coverage depends on the protocol. Pi 1.0.3 hides Bedrock's internal attempt count, so its reports remain observations. Google Generative AI and Vertex can retain complete reports with their default zero-retry setting; enabling provider retries leaves coverage unverified. A recovered Codex WebSocket error also leaves a gap. Failed terminal responses retain usable reports when no unaccounted retry occurred.
Validation
Regression tests cover raw omissions versus reported zero, independent cache counts, failed responses, hidden retries, compaction, tool reports and late warm-up reports. Cache-warming tests check that the shared fetch wrapper preserves the caller's fetch and its stable identity.
Where should the reviewer start?
Start with
adapters/typescript/pi/src/usage.tsfor raw-field parsing, fetch ownership and snapshot boundaries. Then readpi-sdk.tsandtest/usage.test.mjs, especially retry, compaction and cache-warming cases.Related Issues: (use one of the action keywords Closes / Fixes / Resolves / Relates to)
Relates to: none
I confirm this contribution is my own work, or I have the right to submit it under this project's license.
I searched existing issues and open pull requests, and this does not duplicate existing work.