diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index c808cfca457..4103b5227d9 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -583,11 +583,13 @@ class ChunkProcessor: completion_tokens_details = usage_chunk_dict[ "completion_tokens_details" ] - if ( - hasattr(usage_chunk, "server_tool_use") - and usage_chunk.server_tool_use is not None - ): - server_tool_use = usage_chunk.server_tool_use + _server_tool_use = ( + usage_chunk.get("server_tool_use") + if isinstance(usage_chunk, dict) + else getattr(usage_chunk, "server_tool_use", None) + ) + if _server_tool_use is not None: + server_tool_use = _server_tool_use if ( usage_chunk_dict["prompt_tokens_details"] is not None and getattr( diff --git a/litellm/tests/test_cost_calculator_attributeerror.py b/litellm/tests/test_cost_calculator_attributeerror.py new file mode 100644 index 00000000000..4b2bb7ef23d --- /dev/null +++ b/litellm/tests/test_cost_calculator_attributeerror.py @@ -0,0 +1,49 @@ +import sys +import os + +# Add the project root to the path +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../.."))) + +from litellm.cost_calculator import completion_cost +from litellm.types.utils import Usage, ServerToolUse, ModelResponse + + +def test_usage_coercion_from_dict(): + """ + Verify that Usage correctly coerces server_tool_use from a dict to a ServerToolUse object. + """ + usage = Usage( + prompt_tokens=10, + completion_tokens=20, + server_tool_use={"web_search_requests": 5}, + ) + + assert isinstance(usage.server_tool_use, ServerToolUse) + assert usage.server_tool_use.web_search_requests == 5 + + +def test_completion_cost_with_dict_usage(): + """ + Verify that completion_cost handles a response where server_tool_use is a dict. + This simulates the bug reported in #26153. + """ + # Create a usage object and manually set server_tool_use to a dict to simulate the state after stream assembly + usage = Usage(prompt_tokens=10, completion_tokens=20) + usage.server_tool_use = { + "web_search_requests": 5 + } # Manually bypass the __init__ coercion for testing defensive checks + + response = ModelResponse( + id="test-id", + choices=[{"message": {"role": "assistant", "content": "hello"}}], + usage=usage, + ) + + # This should not raise AttributeError + cost = completion_cost(completion_response=response, model="gpt-3.5-turbo") + assert cost is not None + + +if __name__ == "__main__": + test_usage_coercion_from_dict() + test_completion_cost_with_dict_usage() diff --git a/litellm/types/utils.py b/litellm/types/utils.py index e3058d106a6..a52b9cbef1c 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -1300,6 +1300,12 @@ class Delta(SafeAttributeModel, OpenAIObject): self.images: Optional[List[ImageURLListItem]] = None self.annotations: Optional[List[ChatCompletionAnnotation]] = None + self._handle_reasoning(reasoning_content, thinking_blocks, reasoning_items) + self._handle_extra_fields(annotations, images) + self._handle_tool_calls(function_call, tool_calls) + self.audio = audio + + def _handle_reasoning(self, reasoning_content, thinking_blocks, reasoning_items): if reasoning_content is not None: self.reasoning_content = reasoning_content else: @@ -1319,6 +1325,7 @@ class Delta(SafeAttributeModel, OpenAIObject): if hasattr(self, "reasoning_items"): del self.reasoning_items + def _handle_extra_fields(self, annotations, images): # Add annotations to the delta, ensure they are only on Delta if they exist (Match OpenAI spec) if annotations is not None: self.annotations = annotations @@ -1330,6 +1337,7 @@ class Delta(SafeAttributeModel, OpenAIObject): else: del self.images + def _handle_tool_calls(self, function_call, tool_calls): if function_call is not None and isinstance(function_call, dict): self.function_call = FunctionCall(**function_call) else: @@ -1350,8 +1358,6 @@ class Delta(SafeAttributeModel, OpenAIObject): else: self.tool_calls = tool_calls - self.audio = audio - def __contains__(self, key): # Define custom behavior for the 'in' operator return hasattr(self, key) @@ -1538,7 +1544,7 @@ class Usage(SafeAttributeModel, CompletionUsage): prompt_tokens_details: Optional[PromptTokensDetailsWrapper] = None """Breakdown of tokens used in the prompt.""" - def __init__( # noqa: PLR0915 + def __init__( self, prompt_tokens: Optional[int] = None, completion_tokens: Optional[int] = None, @@ -1554,7 +1560,45 @@ class Usage(SafeAttributeModel, CompletionUsage): cost: Optional[float] = None, **params, ): - # handle reasoning_tokens + _completion_tokens_details = self._handle_completion_tokens_details( + completion_tokens_details, reasoning_tokens, completion_tokens + ) + _prompt_tokens_details = self._handle_prompt_tokens_details( + prompt_tokens_details, params + ) + + super().__init__( + prompt_tokens=prompt_tokens or 0, + completion_tokens=completion_tokens or 0, + total_tokens=total_tokens or 0, + completion_tokens_details=_completion_tokens_details or None, + prompt_tokens_details=_prompt_tokens_details or None, + ) + + if server_tool_use is not None: + if isinstance(server_tool_use, dict): + self.server_tool_use = ServerToolUse(**server_tool_use) + else: + self.server_tool_use = server_tool_use + else: # maintain openai compatibility in usage object if possible + del self.server_tool_use + + if cost is not None: + self.cost = cost + else: + del self.cost + + self._handle_extra_mappings(params) + + for k, v in params.items(): + setattr(self, k, v) + + def _handle_completion_tokens_details( + self, + completion_tokens_details, + reasoning_tokens: Optional[int], + completion_tokens: Optional[int], + ) -> Optional[CompletionTokensDetailsWrapper]: _completion_tokens_details: Optional[CompletionTokensDetailsWrapper] = None # First, handle existing completion_tokens_details @@ -1592,7 +1636,11 @@ class Usage(SafeAttributeModel, CompletionUsage): # Prevent negative token counts from inconsistent data _completion_tokens_details.text_tokens = max(0, calculated_text_tokens) + return _completion_tokens_details + def _handle_prompt_tokens_details( + self, prompt_tokens_details, params: dict + ) -> Optional[PromptTokensDetailsWrapper]: # handle prompt_tokens_details _prompt_tokens_details: Optional[PromptTokensDetailsWrapper] = None @@ -1642,25 +1690,9 @@ class Usage(SafeAttributeModel, CompletionUsage): _prompt_tokens_details.cache_creation_tokens = params[ "cache_creation_input_tokens" ] + return _prompt_tokens_details - super().__init__( - prompt_tokens=prompt_tokens or 0, - completion_tokens=completion_tokens or 0, - total_tokens=total_tokens or 0, - completion_tokens_details=_completion_tokens_details or None, - prompt_tokens_details=_prompt_tokens_details or None, - ) - - if server_tool_use is not None: - self.server_tool_use = server_tool_use - else: # maintain openai compatibility in usage object if possible - del self.server_tool_use - - if cost is not None: - self.cost = cost - else: - del self.cost - + def _handle_extra_mappings(self, params: dict): ## ANTHROPIC MAPPING ## if "cache_creation_input_tokens" in params and isinstance( params["cache_creation_input_tokens"], int diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py index e40a0817fd9..f24464d77bf 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py @@ -1,8 +1,6 @@ -import json import os import sys -import pytest sys.path.insert( 0, os.path.abspath("../../..") @@ -520,7 +518,7 @@ def test_stream_chunk_builder_anthropic_web_search(): assert usage.prompt_tokens == 50 assert usage.completion_tokens == 27 assert usage.total_tokens == 77 - assert usage.server_tool_use["web_search_requests"] == 2 + assert usage.server_tool_use.web_search_requests == 2 def test_sort_chunks_handles_dict_hidden_params_created_at():