Repository navigation
fix: preserve Qwen token usage per invocation - #448
mnajafian-nv wants to merge 48 commits into
Conversation
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>
Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
Qwen Code's SDK result carries cumulative usage for the live session, including cache_read_input_tokens. The adapter dropped the cache reads, so a run that mostly hit the prompt cache looked the same as one that did not. It also clamped a counter that went down to 0, so the invocation in which Qwen Code reset its counters, such as a /clear prompt, reported 0 tokens. And because Qwen Code records each count a provider omitted as 0, a provider that sent no usage produced 0 input and output tokens. Report cache reads as cached_input_tokens when the cumulative cache counter grows, and mark input as including cache, because Qwen Code counts prompt tokens with their cache reads included. Treat a counter that did not grow as unknown, one count at a time, and leave a total derived from an unknown count, or above the safe integer range, unknown. Qwen cannot distinguish a reported zero from an omitted count, and its result cannot show that an invocation made no model calls: slash commands report zero turns and zero API duration, and /compress can still send the conversation to the model. So zero growth is never reported as a measured 0. Leave usage unknown when a counter goes down or the session ID changes. Qwen Code resets its counters on /clear and replays a resumed session's history into them, so the difference then does not describe the invocation. Qwen's cumulative results do not identify individual responses with missing usage. If some calls report tokens and others do not, positive deltas can undercount the invocation. The README and the Qwen docs page state this, and a test pins it. Let the shared mock provider serve a configured chat-completion usage object through the scenario's usage field, keeping its zero default. The Qwen end-to-end tests now check per-invocation counts on two prompts with positive usage, and expect no usage when the provider reports zeros. Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
This comment was marked as outdated.
This comment was marked as outdated.
|
Fern docs preview: https://nvidia-preview-pull-request-448.docs.buildwithfern.com/nemo/fabric |
LangChain carries cache reads per model call in usage_metadata.input_token_details.cache_read, and OpenAI raw usage carries them in prompt_tokens_details.cached_tokens, but the adapter never reported cached_input_tokens. Sum cached input per call with the same rule as the other counts: one call without a cache count leaves the sum unknown, and a reported 0 stays 0. Raw OpenAI usage is read first because ChatOpenAI renames the cache key for priority and flex service tiers. Report input_tokens_include_cache=True whenever input tokens are known. LangChain's UsageMetadata.input_tokens counts every input token type; langchain-anthropic adds cache reads and writes back into it, and OpenAI's prompt_tokens already includes cached tokens. Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
The adapter drops a model call that the provider rejects with a 4xx status other than 408, because the provider spent no tokens on it. When every call was rejected, or no call started, the call list was empty and the usage was reported as unknown, although the invocation provably spent no tokens. The adapter now reports 0 for every token count and the cost in that case. Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
Deep Agents routes a subagent spec with graph_id to its async subagent middleware, which starts and updates tasks on a remote Agent Protocol server. Those model calls never reach the adapter's callbacks, so the reported counts left them out while appearing complete. The adapter blocked such subagents only when a tools policy was configured. When the configured subagents include one with graph_id, the adapter now reports the invocation's token counts and cost as unknown. The local sums stay in output.usage. Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
The usage notes said that after a 5xx response, timeout, or cancelled call the provider may have done work. The reason the counts stay unknown is narrower: the provider may have billed the call without reporting usage. The README, the recorder docstring, and a test comment now say that. Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
Apply ruff format to the check added in the previous commit. Signed-off-by: mnajafian-nv <mnajafian@nvidia.com>
…gent The adapter reports 0 when its callbacks observed no billed model call. That 0 holds for the invocation only when every model call runs through those callbacks. Deep Agents runs declarative subagents and summarization inside the graph, where the run config's callbacks reach them, but an Agent Protocol subagent (graph_id) runs on a remote server that the callbacks never see. The adapter already reports unknown usage for a runtime with such a subagent; this test pins that for an invocation whose only observed call was rejected, and the docstring and README state the condition. 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>
Separate complete invocation totals from available observations and preserve accounting gaps. Add the first collector in Deep Agents. 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>
Keep SDK fixtures in their unique temporary workspace and use the neighboring tests' cleanup pattern. Avoid writing to the user home directory or renaming across filesystems. 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>
Include the current shared accounting validation while preserving the adapter collectors and canonical catalog declarations. 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>
AjayThorve
left a comment
There was a problem hiding this comment.
Reviewed Qwen cumulative deltas, reset/session-change handling, interrupted stream receipts and failed-query retirement. Local validation: 36 adapter tests passed, including the pinned Qwen emitter; descriptors/catalog passed. Positive observations remain scope_unverified rather than guessed complete totals. Product CI/DCO are green; NVSkills evidence is missing. Merge after #442. The retained B200 bridge/provider mismatch (nvidia versus the adapter schema requiring openai) remains separate. Reviewed commit: 0c7ec8d.
|
This PR preserves the Qwen per-invocation usage approach. I’m closing the current branch so future work can follow the shared accounting design. |
Overview
Qwen Code reports cumulative session counters and fills omitted usage values with zero. Public assistant messages also do not establish coverage of every model call. Those sources cannot establish a complete invocation total.
Retain the positive counts available for each NVIDIA NeMo Fabric invocation as observations in
RunResult.model_usage. Keep the coverage gap visible instead of treating a reset, omission or hidden call as measured zero.Dependencies: #442 supplies the shared accounting contract and normalizer. The diff against
mainincludes unmerged prerequisites.Details
Take differences between successive session counters, including cache reads. A session change or decreasing counter invalidates the affected comparison rather than producing a clamped zero. Pass the observations to shared accounting in #442.
If the stream ends without a terminal result, retain useful final assistant-message counts with an unfinished-work gap. Failed terminal results also keep their usage. Retire failed queries so later activity cannot enter another invocation.
The cumulative counters cannot show whether every provider response reported usage. Complete totals therefore remain unknown even when a positive subtotal is available. No provider cost is inferred from those counters.
Validation
The Qwen Code adapter build and tests passed, along with focused Python adapter, shared-accounting and catalog checks. With the adapter contract and common packages built, reproduce the adapter checks from the repository root:
npm run build --prefix adapters/typescript/qwen node --test adapters/typescript/qwen/test/*.test.mjsRegression tests cover session changes, counter resets, independent cache growth, zero growth and interrupted streams. SDK cases cover normal repeated prompts, compression, omitted usage and failed-query retirement.
Where should the reviewer start?
Start with
QwenUsageTrackerandQwenStreamUsageinadapters/typescript/qwen/src/qwen-sdk.ts. Then review failed-query retirement inruntime.tsand reset, session-change and interrupted-stream cases intest/usage.test.mjs.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.