From 1908245c668b5b74891641d9e02ad0b794f8b947 Mon Sep 17 00:00:00 2001 From: Erick Aleman Date: Tue, 21 Apr 2026 21:35:44 -0400 Subject: [PATCH] fix(cost_calculator): robust handling of server_tool_use in streaming responses --- .../llm_cost_calc/tool_call_cost_tracking.py | 6 ++- .../streaming_chunk_builder_utils.py | 12 +++-- litellm/llms/anthropic/cost_calculation.py | 6 ++- .../test_cost_calculator_attributeerror.py | 48 +++++++++++++++++++ litellm/types/utils.py | 5 +- 5 files changed, 67 insertions(+), 10 deletions(-) create mode 100644 litellm/tests/test_cost_calculator_attributeerror.py diff --git a/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py b/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py index 8da66d4600d..6aa3862fece 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py +++ b/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py @@ -340,7 +340,8 @@ class StandardBuiltInToolCostTracking: if ( hasattr(usage, "server_tool_use") and usage.server_tool_use is not None - and usage.server_tool_use.web_search_requests is not None + and getattr(usage.server_tool_use, "web_search_requests", None) + is not None ): return True return False @@ -353,7 +354,8 @@ class StandardBuiltInToolCostTracking: if ( hasattr(usage, "server_tool_use") and usage.server_tool_use is not None - and usage.server_tool_use.web_search_requests is not None + and getattr(usage.server_tool_use, "web_search_requests", None) + is not None ): return True elif ( 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/llms/anthropic/cost_calculation.py b/litellm/llms/anthropic/cost_calculation.py index 3882d8f978c..e47f3311ae7 100644 --- a/litellm/llms/anthropic/cost_calculation.py +++ b/litellm/llms/anthropic/cost_calculation.py @@ -113,7 +113,7 @@ def get_cost_for_anthropic_web_search( if ( usage is None or usage.server_tool_use is None - or usage.server_tool_use.web_search_requests is None + or getattr(usage.server_tool_use, "web_search_requests", None) is None ): return 0.0 @@ -128,5 +128,7 @@ def get_cost_for_anthropic_web_search( return 0.0 ## Calculate the total cost - total_cost = cost_per_web_search_request * usage.server_tool_use.web_search_requests + total_cost = cost_per_web_search_request * getattr( + usage.server_tool_use, "web_search_requests", 0 + ) return total_cost diff --git a/litellm/tests/test_cost_calculator_attributeerror.py b/litellm/tests/test_cost_calculator_attributeerror.py new file mode 100644 index 00000000000..ecde75e76a5 --- /dev/null +++ b/litellm/tests/test_cost_calculator_attributeerror.py @@ -0,0 +1,48 @@ +import sys +import os +from unittest.mock import MagicMock +import pytest + +# 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 + print(f"Cost calculated: {cost}") + +if __name__ == "__main__": + test_usage_coercion_from_dict() + test_completion_cost_with_dict_usage() + print("Tests passed!") diff --git a/litellm/types/utils.py b/litellm/types/utils.py index d0bc9b78941..7e57ff5cab4 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -1652,7 +1652,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