diff --git a/litellm/tests/test_cost_calculator_attributeerror.py b/litellm/tests/test_cost_calculator_attributeerror.py index 4b2bb7ef23d..a8b7bbaf943 100644 --- a/litellm/tests/test_cost_calculator_attributeerror.py +++ b/litellm/tests/test_cost_calculator_attributeerror.py @@ -4,8 +4,8 @@ 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 +from litellm.litellm_core_utils.streaming_chunk_builder_utils import ChunkProcessor +from litellm.types.utils import Usage, ServerToolUse def test_usage_coercion_from_dict(): @@ -22,28 +22,45 @@ def test_usage_coercion_from_dict(): assert usage.server_tool_use.web_search_requests == 5 -def test_completion_cost_with_dict_usage(): +def test_chunk_processor_coerces_dict_server_tool_use_from_stream_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 + Verify streaming usage assembly handles server_tool_use from a dict. - response = ModelResponse( - id="test-id", - choices=[{"message": {"role": "assistant", "content": "hello"}}], - usage=usage, + Anthropic streaming chunks can arrive with usage as plain dicts. This + exercises the exact dict path that previously raised AttributeError when + server_tool_use was accessed with attribute syntax. + """ + chunks = [ + { + "usage": { + "prompt_tokens": 10, + "completion_tokens": 0, + "total_tokens": 10, + "server_tool_use": {"web_search_requests": 5}, + } + }, + { + "usage": { + "prompt_tokens": 0, + "completion_tokens": 20, + "total_tokens": 20, + } + }, + ] + + usage = ChunkProcessor(chunks=chunks).calculate_usage( + chunks=chunks, + model="claude-sonnet-4-5-20250929", + completion_output="hello", ) - # This should not raise AttributeError - cost = completion_cost(completion_response=response, model="gpt-3.5-turbo") - assert cost is not None + assert usage.prompt_tokens == 10 + assert usage.completion_tokens == 20 + assert usage.total_tokens == 30 + assert isinstance(usage.server_tool_use, ServerToolUse) + assert usage.server_tool_use.web_search_requests == 5 if __name__ == "__main__": test_usage_coercion_from_dict() - test_completion_cost_with_dict_usage() + test_chunk_processor_coerces_dict_server_tool_use_from_stream_usage()