mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-06 08:16:43 +00:00
perf: skip throwaway Usage() construction in ModelResponse.__init__
Avoid constructing a default Usage() object that gets immediately overwritten by convert_to_model_response_object. Set usage=None instead; the real Usage is assigned via setattr later. Also fix Bedrock Qwen2/Qwen3 transform_response to assign a new Usage object instead of mutating a potentially missing one.
This commit is contained in:
Ryan Crabbe
Sameer Kankute
parent
bacfa9653a
commit
b07078ae75
4 changed files with 38 additions and 197 deletions
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue