From 4853cfd25d3d266c9a56c92155c12df4c9830714 Mon Sep 17 00:00:00 2001 From: oss-agent-shin <279349115+oss-agent-shin@users.noreply.github.com> Date: Wed, 6 May 2026 17:47:07 +0000 Subject: [PATCH] Fix Anthropic streaming web search cost usage Co-authored-by: ishaan-berri --- .../streaming_chunk_builder_utils.py | 19 ++-- litellm/types/utils.py | 7 +- .../test_streaming_chunk_builder_utils.py | 92 ++++++++++++++++++- 3 files changed, 105 insertions(+), 13 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index fe7c62c3842..c5d4ab668be 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -464,12 +464,15 @@ class ChunkProcessor: id=id, ) - def _usage_chunk_calculation_helper(self, usage_chunk: Usage) -> dict: + def _usage_chunk_calculation_helper( + self, usage_chunk: Union[Usage, Dict[str, Any]] + ) -> dict: prompt_tokens = 0 completion_tokens = 0 ## anthropic prompt caching information ## cache_creation_input_tokens: Optional[int] = None cache_read_input_tokens: Optional[int] = None + server_tool_use: Optional[ServerToolUse] = None completion_tokens_details: Optional[CompletionTokensDetails] = None prompt_tokens_details: Optional[PromptTokensDetailsWrapper] = None @@ -499,12 +502,19 @@ class ChunkProcessor: usage_chunk.prompt_tokens_details, PromptTokensDetailsWrapper ): prompt_tokens_details = usage_chunk.prompt_tokens_details + if "server_tool_use" in usage_chunk: + raw_server_tool_use = usage_chunk.get("server_tool_use") + if isinstance(raw_server_tool_use, dict): + server_tool_use = ServerToolUse(**raw_server_tool_use) + elif isinstance(raw_server_tool_use, ServerToolUse): + server_tool_use = raw_server_tool_use return { "prompt_tokens": prompt_tokens, "completion_tokens": completion_tokens, "cache_creation_input_tokens": cache_creation_input_tokens, "cache_read_input_tokens": cache_read_input_tokens, + "server_tool_use": server_tool_use, "completion_tokens_details": completion_tokens_details, "prompt_tokens_details": prompt_tokens_details, } @@ -584,11 +594,8 @@ 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 + if usage_chunk_dict["server_tool_use"] is not None: + server_tool_use = usage_chunk_dict["server_tool_use"] if ( usage_chunk_dict["prompt_tokens_details"] is not None and getattr( diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 00a7748309b..e8de84a410d 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -1553,7 +1553,7 @@ class Usage(SafeAttributeModel, CompletionUsage): completion_tokens_details: Optional[ Union[CompletionTokensDetailsWrapper, dict] ] = None, - server_tool_use: Optional[ServerToolUse] = None, + server_tool_use: Optional[Union[ServerToolUse, dict]] = None, cost: Optional[float] = None, **params, ): @@ -1655,7 +1655,10 @@ class Usage(SafeAttributeModel, CompletionUsage): ) if server_tool_use is not None: - self.server_tool_use = server_tool_use + 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 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..e940e0d0f41 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,14 +1,11 @@ -import json import os import sys -import pytest - sys.path.insert( 0, os.path.abspath("../../..") ) # Adds the parent directory to the system path -from litellm import stream_chunk_builder +from litellm import completion_cost, stream_chunk_builder from litellm.litellm_core_utils.streaming_chunk_builder_utils import ChunkProcessor from litellm.types.utils import ( ChatCompletionDeltaToolCall, @@ -520,7 +517,92 @@ 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 is not None + assert usage.server_tool_use.web_search_requests == 2 + + +def test_stream_chunk_builder_anthropic_web_search_completion_cost(): + chunks = [ + ModelResponseStream( + id="chatcmpl-mocked-usage-1", + created=1745513206, + model="claude-sonnet-4-6", + object="chat.completion.chunk", + system_fingerprint=None, + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + provider_specific_fields=None, + content="hello", + role="assistant", + function_call=None, + tool_calls=None, + audio=None, + ), + logprobs=None, + ) + ], + provider_specific_fields=None, + stream_options={"include_usage": True}, + usage=Usage( + completion_tokens=0, + prompt_tokens=50, + total_tokens=50, + completion_tokens_details=None, + server_tool_use=ServerToolUse(web_search_requests=2), + prompt_tokens_details=None, + ), + ), + ModelResponseStream( + id="chatcmpl-mocked-usage-1", + created=1745513207, + model="claude-sonnet-4-6", + object="chat.completion.chunk", + system_fingerprint=None, + choices=[ + StreamingChoices( + finish_reason="stop", + index=0, + delta=Delta( + provider_specific_fields=None, + content=None, + role=None, + function_call=None, + tool_calls=None, + audio=None, + ), + logprobs=None, + ) + ], + provider_specific_fields=None, + stream_options={"include_usage": True}, + usage=Usage( + completion_tokens=27, + prompt_tokens=0, + total_tokens=27, + completion_tokens_details=None, + prompt_tokens_details=None, + ), + ), + ] + + response = stream_chunk_builder(chunks) + + assert response is not None + assert response.usage is not None + assert response.usage.server_tool_use is not None + assert response.usage.server_tool_use.web_search_requests == 2 + assert completion_cost(completion_response=response) > 0 + + dict_response = stream_chunk_builder([chunk.model_dump() for chunk in chunks]) + + assert dict_response is not None + assert dict_response.usage is not None + assert dict_response.usage.server_tool_use is not None + assert dict_response.usage.server_tool_use.web_search_requests == 2 + assert completion_cost(completion_response=dict_response) > 0 def test_sort_chunks_handles_dict_hidden_params_created_at():