scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-10 22:41:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 83

fix: preserve prompt_tokens_details when raw PromptTokensDetails reaches chunk aggregator

Browse source 
## TLDR

Signed-off-by: Ishaan <ishaangupta0408@gmail.com>
This commit is contained in:
Ishaan 2026-08-14 09:11:02 +00:00
parent 423b791ee0
commit 294bc7eba0
2 changed files with 24 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
litellm/litellm_core_utils/streaming_chunk_builder_utils.py
View file

@ -27,6 +27,7 @@ from litellm.types.utils import (
ServerToolUse,
Usage,
)
from openai.types.completion_usage import PromptTokensDetails
from litellm.utils import print_verbose, token_counter
if TYPE_CHECKING:
@ -650,6 +651,8 @@ class ChunkProcessor:
prompt_tokens_details = PromptTokensDetailsWrapper(**usage_chunk.prompt_tokens_details)
elif isinstance(usage_chunk.prompt_tokens_details, PromptTokensDetailsWrapper):
prompt_tokens_details = usage_chunk.prompt_tokens_details
elif isinstance(usage_chunk.prompt_tokens_details, PromptTokensDetails):
prompt_tokens_details = PromptTokensDetailsWrapper(**usage_chunk.prompt_tokens_details.model_dump())
return {
"prompt_tokens": prompt_tokens,

21
tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py
View file

@ -1076,6 +1076,27 @@ def test_prompt_tokens_details_survive_later_usage_chunk_without_details():
assert usage.prompt_tokens_details.cache_write_tokens == 10
def test_prompt_tokens_details_base_class_preserved_in_usage_chunk():
"""Regression for #36882: when the raw streaming usage chunk carries a
base-class PromptTokensDetails (not PromptTokensDetailsWrapper), the
helper must still extract cached_tokens rather than silently dropping the
whole prompt_tokens_details field."""
from openai.types.completion_usage import PromptTokensDetails as RawPromptTokensDetails
from litellm.litellm_core_utils.streaming_chunk_builder_utils import ChunkProcessor
from litellm.types.utils import Usage
raw_details = RawPromptTokensDetails(cached_tokens=0)
usage_chunk = Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150)
usage_chunk.prompt_tokens_details = raw_details # type: ignore[assignment]
processor = ChunkProcessor.__new__(ChunkProcessor)
result = processor._usage_chunk_calculation_helper(usage_chunk)
assert result["prompt_tokens_details"] is not None
assert result["prompt_tokens_details"].cached_tokens == 0
def test_get_combined_tool_content_custom_tool_call():
from litellm.litellm_core_utils.streaming_chunk_builder_utils import ChunkProcessor
from litellm.types.utils import ChatCompletionMessageCustomToolCall