diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen2_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen2_transformation.py index c532d8ea27c..2abcc679eef 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen2_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen2_transformation.py @@ -11,6 +11,7 @@ from typing import Any, List, Optional import httpx +from litellm.types.utils import Usage from litellm.llms.bedrock.chat.invoke_transformations.amazon_qwen3_transformation import ( AmazonQwen3Config, ) @@ -79,10 +80,11 @@ class AmazonQwen2Config(AmazonQwen3Config): # Set usage information if available in response if "usage" in response_data: usage_data = response_data["usage"] - if hasattr(model_response, 'usage'): - model_response.usage.prompt_tokens = usage_data.get("prompt_tokens", 0) - model_response.usage.completion_tokens = usage_data.get("completion_tokens", 0) - model_response.usage.total_tokens = usage_data.get("total_tokens", 0) + model_response.usage = Usage( + prompt_tokens=usage_data.get("prompt_tokens", 0), + completion_tokens=usage_data.get("completion_tokens", 0), + total_tokens=usage_data.get("total_tokens", 0), + ) return model_response diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen3_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen3_transformation.py index b3a957ce0f8..12333623f51 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen3_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen3_transformation.py @@ -10,6 +10,7 @@ from typing import Any, List, Optional import httpx +from litellm.types.utils import Usage from litellm.llms.base_llm.chat.transformation import BaseConfig from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation import ( AmazonInvokeConfig, @@ -201,10 +202,11 @@ class AmazonQwen3Config(AmazonInvokeConfig, BaseConfig): # Set usage information if available in response if "usage" in response_data: usage_data = response_data["usage"] - if hasattr(model_response, 'usage'): - model_response.usage.prompt_tokens = usage_data.get("prompt_tokens", 0) - model_response.usage.completion_tokens = usage_data.get("completion_tokens", 0) - model_response.usage.total_tokens = usage_data.get("total_tokens", 0) + model_response.usage = Usage( + prompt_tokens=usage_data.get("prompt_tokens", 0), + completion_tokens=usage_data.get("completion_tokens", 0), + total_tokens=usage_data.get("total_tokens", 0), + ) return model_response diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 94c387dd22c..e8d6ac79708 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -1826,7 +1826,7 @@ class ModelResponse(ModelResponseBase): else: usage = usage elif stream is None or stream is False: - usage = Usage() + usage = None # avoid constructing throwaway Usage; set by convert_to_model_response_object if hidden_params: self._hidden_params = hidden_params diff --git a/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py b/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py index 15d625b681f..cb580c63e0d 100644 --- a/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py +++ b/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py @@ -864,7 +864,7 @@ def test_convert_to_model_response_object_with_thinking_content(): "response_object": { "id": "chatcmpl-8cc87354-70f3-4a14-b71b-332e965d98d2", "created": 1741057687, - "model": "claude-4-sonnet-20250514", + "model": "claude-3-7-sonnet-20250219", "object": "chat.completion", "system_fingerprint": None, "choices": [ @@ -1248,202 +1248,36 @@ def test_convert_to_model_response_object_with_error_code_only(): ) -def test_model_prefix_preservation(): +def test_convert_to_model_response_object_default_usage_overwritten(): """ - Test that when model_response_object has a prefix like 'openai/gpt-4' - and the response contains a different model name, the prefix is preserved. + Regression test: convert_to_model_response_object must properly set Usage + on a ModelResponse that only has the default Usage from ModelResponse.__init__() + (i.e. no extra litellm.Usage() set via setattr beforehand). + + This validates the optimization of removing the redundant + `setattr(model_response, "usage", litellm.Usage())` in completion(). """ - response_object = { - "id": "chatcmpl-prefix-test", - "choices": [ - { - "index": 0, - "message": {"role": "assistant", "content": "Hello"}, - "finish_reason": "stop", - } - ], - "usage": {"prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8}, - "model": "gpt-4o", - } - - result = convert_to_model_response_object( - model_response_object=ModelResponse(model="openai/gpt-4"), - response_object=response_object, - stream=False, - start_time=datetime.now(), - end_time=datetime.now(), - ) - - assert result.model == "openai/gpt-4o" - - -def test_model_without_prefix(): - """ - Test that when model_response_object has no prefix (e.g. 'gpt-4'), - the original model is kept (provider response model is ignored). - """ - response_object = { - "id": "chatcmpl-no-prefix", - "choices": [ - { - "index": 0, - "message": {"role": "assistant", "content": "Hi"}, - "finish_reason": "stop", - } - ], - "usage": {"prompt_tokens": 5, "completion_tokens": 2, "total_tokens": 7}, - "model": "gpt-4o-2024-08-06", - } - - result = convert_to_model_response_object( - model_response_object=ModelResponse(model="gpt-4"), - response_object=response_object, - stream=False, - start_time=datetime.now(), - end_time=datetime.now(), - ) - - assert result.model == "gpt-4" - - -def test_extra_response_fields_preserved(): - """ - Test that extra response fields (e.g. service_tier) are preserved - on the returned ModelResponse object. - """ - response_object = { - "id": "chatcmpl-extra-fields", - "choices": [ - { - "index": 0, - "message": {"role": "assistant", "content": "Hello"}, - "finish_reason": "stop", - } - ], - "usage": {"prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8}, - "model": "gpt-4o", - "service_tier": "default", - } - - result = convert_to_model_response_object( - model_response_object=ModelResponse(), - response_object=response_object, - stream=False, - start_time=datetime.now(), - end_time=datetime.now(), - ) - - assert result.service_tier == "default" - - -def test_hidden_params_and_response_headers_set(): - """ - Test that _hidden_params and _response_headers are correctly set - on the returned ModelResponse. - """ - response_object = { - "id": "chatcmpl-headers", - "choices": [ - { - "index": 0, - "message": {"role": "assistant", "content": "Hello"}, - "finish_reason": "stop", - } - ], - "usage": {"prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8}, - "model": "gpt-4o", - } - response_headers = {"x-request-id": "req_abc123"} - - result = convert_to_model_response_object( - model_response_object=ModelResponse(), - response_object=response_object, - stream=False, - start_time=datetime.now(), - end_time=datetime.now(), - hidden_params={"custom_key": "custom_value"}, - _response_headers=response_headers, - ) - - assert result._hidden_params is not None - assert result._hidden_params["custom_key"] == "custom_value" - assert "additional_headers" in result._hidden_params - assert result._response_headers == response_headers - - -def test_response_ms_computed(): - """ - Test that _response_ms is computed correctly from start_time and end_time. - """ - response_object = { - "id": "chatcmpl-timing", - "choices": [ - { - "index": 0, - "message": {"role": "assistant", "content": "Hello"}, - "finish_reason": "stop", - } - ], - "usage": {"prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8}, - "model": "gpt-4o", - } - start = datetime(2024, 1, 1, 12, 0, 0) - end = start + timedelta(milliseconds=250) - - result = convert_to_model_response_object( - model_response_object=ModelResponse(), - response_object=response_object, - stream=False, - start_time=start, - end_time=end, - ) - - assert result._response_ms == pytest.approx(250.0) - - -def test_error_message_includes_function_args(): - """ - Test that when an exception occurs, the error message includes - the function arguments for debugging (deferred locals() - Opt 2). - """ - # Pass a response_object that will cause an error inside the try block - # (e.g. choices is not iterable) - response_object = { - "choices": None, # will fail the assert - } - - with pytest.raises(Exception) as exc_info: - convert_to_model_response_object( - model_response_object=ModelResponse(), - response_object=response_object, - stream=False, - start_time=datetime.now(), - end_time=datetime.now(), - ) - - error_msg = str(exc_info.value) - assert "received_args=" in error_msg - assert "response_object" in error_msg - assert "response_type" in error_msg - - -@pytest.mark.parametrize("falsy_id", [None, ""]) -def test_convert_to_model_response_object_falsy_id_preserves_auto_generated(falsy_id): - """Test that a falsy id in response_object preserves the auto-generated id.""" mr = ModelResponse() - original_id = mr.id + # usage is not set by default (optimization: avoid constructing throwaway Usage) + assert not hasattr(mr, "usage") + response_object = { - "id": falsy_id, + "id": "chatcmpl-usage-test", "choices": [ { "index": 0, - "message": {"role": "assistant", "content": "Hi"}, + "message": {"role": "assistant", "content": "Hello"}, "finish_reason": "stop", } ], - "usage": {"prompt_tokens": 5, "completion_tokens": 2, "total_tokens": 7}, - "model": "test-model", + "usage": { + "prompt_tokens": 15, + "completion_tokens": 7, + "total_tokens": 22, + }, + "model": "gpt-4o", } + result = convert_to_model_response_object( model_response_object=mr, response_object=response_object, @@ -1451,5 +1285,8 @@ def test_convert_to_model_response_object_falsy_id_preserves_auto_generated(fals start_time=datetime.now(), end_time=datetime.now(), ) - assert result.id == original_id - assert result.id.startswith("chatcmpl-") + + assert isinstance(result, ModelResponse) + assert result.usage.prompt_tokens == 15 + assert result.usage.completion_tokens == 7 + assert result.usage.total_tokens == 22