Repository navigation
fix(models): restore usage_metadata in MindIE tool-mode simulated streaming - #5195
Merged
Merged
Conversation
…eaming In tool-enabled requests MindIEChatModel._astream falls back to awaiting the full _agenerate response and re-emitting it as simulated AIMessageChunks. The full response carries usage_metadata, but none of the simulated chunks copied it, so chunk aggregation (add_ai_message_chunks) produced a final message with usage_metadata=None. Token usage therefore vanished from token accounting, run stats, persistence and the UI for every tool-enabled streamed turn. Mirror OpenAI's terminal-usage-frame convention: attach msg.usage_metadata to exactly the last simulated chunk (the trailing tool-call chunk when present, else the last text chunk / the single tool-only chunk) so the aggregated message carries it exactly once. add_usage() is per-chunk additive, so attaching usage to every chunk would multiply the totals. Scope: MindIEChatModel only; other providers keep native streaming and ainvoke/non-tool astream were already correct. Tests: regression guard asserting exactly one carrier chunk equals the last one and that merged usage equals the original across all three simulated-stream branches, plus chain-level tests driving the public astream() wrapper and asserting the persisted model_dump() shape. Closes bytedance#5192
|
|
github-actions
Bot
added
area:backend
WillemJiang
approved these changes
Sep 5, 2026
1 of 9 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
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.
Fixes #5192.
What happens today
For tool-enabled requests MindIEChatModel._astream cannot use native
streaming: the MindIE engine drops
choiceswhenstream=Trueand tools arepresent, so the model awaits the full
_agenerateresponse and re-emits it assimulated
AIMessageChunks. The full response carriesusage_metadata, but thesimulated chunks never copy it. Chunk aggregation (
add_ai_message_chunks)only produces
usage_metadatawhen at least one chunk carries it, so the mergedfinal message ends up with
usage_metadata=None— every tool-enabled streamedturn silently loses token usage from token accounting (
RunJournal), runcompletion stats, thread-message persistence, and the frontend token-usage
display, even though the provider still bills for those tokens.
Fix
The native
ChatOpenAIstream already solves this exact problem: it emits aterminal usage frame as its last chunk, and aggregation carries it forward. This
PR mirrors that convention in the simulated stream — attach
msg.usage_metadatato exactly the last chunk it yields: the trailing tool-call chunk when tool
calls are present, otherwise the final text chunk, or the single tool-only
chunk. Deliberately attached to one chunk only:
add_usage()sums per chunk, soattaching it to every chunk would multiply the totals.
Tests
New regression test mirrors the issue's deterministic shape across the three
simulated-stream branches (text-only, text + tool call, tool-call only): the
underlying response carries usage, and the test asserts exactly one carrier
chunk — the last one — and that the merged message reproduces the original
usage_metadata. A chain-level test drives the publicastream()wrapper (thepath LangGraph state accumulation,
RunJournal.on_llm_end, and the frontendactually consume) and asserts the persisted
model_dump()shape keeps usagedespite the wrapper's synthetic trailing
chunk_position="last"chunk. Fails onmain before the fix (no chunk carries usage and the aggregated message is
None), passes with it. Also run:tests/test_mindie_provider.py, 47 passed.ruff format/check clean.
AI disclosure
AI tool used: yes (Trae IDE)
Extent: implementation and tests prepared by the agent; root cause verified by
reading langchain-core chunk-aggregation semantics (
add_ai_message_chunks) andthe native
ChatOpenAIterminal-usage frame before writing the fixHuman review: changes verified by running
tests/test_mindie_provider.pylocally (47 passed) and
ruff check/ruff format --check