This commit is contained in:
Ishaan 2026-09-13 00:00:00 -07:00 • committed by GitHub
commit 7a3c5bf0f0
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 24 additions and 0 deletions

View file

@ -34,6 +34,7 @@ from litellm.types.utils import (
StreamingChoices,
Usage,
)
from openai.types.completion_usage import PromptTokensDetails
from litellm.utils import print_verbose, token_counter
if TYPE_CHECKING:
@ -779,6 +780,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,

View file

@ -1212,6 +1212,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