diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 12f7fe2c3c..40fae4e4a5 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -1,6 +1,7 @@ import asyncio import json import logging +import time import traceback from datetime import datetime from typing import ( @@ -441,10 +442,18 @@ class ProxyBaseLLMRequestProcessing: ), **( { - "x-litellm-timing-pre-processing-ms": str(hidden_params.get("timing_pre_processing_ms", None)), - "x-litellm-timing-llm-api-ms": str(hidden_params.get("timing_llm_api_ms", None)), - "x-litellm-timing-post-processing-ms": str(hidden_params.get("timing_post_processing_ms", None)), - "x-litellm-timing-message-copy-ms": str(hidden_params.get("timing_message_copy_ms", None)), + "x-litellm-timing-pre-processing-ms": str( + hidden_params.get("timing_pre_processing_ms", None) + ), + "x-litellm-timing-llm-api-ms": str( + hidden_params.get("timing_llm_api_ms", None) + ), + "x-litellm-timing-post-processing-ms": str( + hidden_params.get("timing_post_processing_ms", None) + ), + "x-litellm-timing-message-copy-ms": str( + hidden_params.get("timing_message_copy_ms", None) + ), } if LITELLM_DETAILED_TIMING else {} @@ -564,16 +573,6 @@ class ProxyBaseLLMRequestProcessing: ) -> Tuple[dict, LiteLLMLoggingObj]: start_time = datetime.now() # start before calling guardrail hooks - # Calculate request queue time if arrival_time is available - # Use start_time.timestamp() to avoid extra time.time() call for better performance - proxy_server_request = self.data.get("proxy_server_request", {}) - arrival_time = proxy_server_request.get("arrival_time") - queue_time_seconds = None - if arrival_time is not None: - # Convert start_time (datetime) to timestamp for calculation - processing_start_time = start_time.timestamp() - queue_time_seconds = processing_start_time - arrival_time - self.data = await add_litellm_data_to_request( data=self.data, request=request, @@ -583,6 +582,15 @@ class ProxyBaseLLMRequestProcessing: proxy_config=proxy_config, ) + # Calculate request queue time after add_litellm_data_to_request + # which sets arrival_time in proxy_server_request + proxy_server_request = self.data.get("proxy_server_request", {}) + arrival_time = proxy_server_request.get("arrival_time") + queue_time_seconds = None + if arrival_time is not None: + processing_start_time = time.time() + queue_time_seconds = processing_start_time - arrival_time + # Store queue time in metadata after add_litellm_data_to_request to ensure it's preserved if queue_time_seconds is not None: from litellm.proxy.litellm_pre_call_utils import _get_metadata_variable_name @@ -634,7 +642,7 @@ class ProxyBaseLLMRequestProcessing: self.data["litellm_call_id"] = request.headers.get( "x-litellm-call-id", str(uuid.uuid4()) ) - + ### AUTO STREAM USAGE TRACKING ### # If always_include_stream_usage is enabled and this is a streaming request # automatically add stream_options={'include_usage': True} if not already set @@ -650,7 +658,7 @@ class ProxyBaseLLMRequestProcessing: and "include_usage" not in self.data["stream_options"] ): self.data["stream_options"]["include_usage"] = True - + ### CALL HOOKS ### - modify/reject incoming data before calling the model ## LOGGING OBJECT ## - initialize logging object for logging success/failure events for call @@ -710,7 +718,9 @@ class ProxyBaseLLMRequestProcessing: "Request received by LiteLLM: payload too large to log (%d bytes, limit %d). Keys: %s", len(_payload_str), MAX_PAYLOAD_SIZE_FOR_DEBUG_LOG, - list(self.data.keys()) if isinstance(self.data, dict) else type(self.data).__name__, + list(self.data.keys()) + if isinstance(self.data, dict) + else type(self.data).__name__, ) else: verbose_proxy_logger.debug( @@ -913,9 +923,9 @@ class ProxyBaseLLMRequestProcessing: # aliasing/routing, but the OpenAI-compatible response `model` field should reflect # what the client sent. if requested_model_from_client: - self.data["_litellm_client_requested_model"] = ( - requested_model_from_client - ) + self.data[ + "_litellm_client_requested_model" + ] = requested_model_from_client if route_type == "allm_passthrough_route": # Check if response is an async generator if self._is_streaming_response(response): @@ -1510,9 +1520,9 @@ class ProxyBaseLLMRequestProcessing: # Add cache-related fields to **params (handled by Usage.__init__) if cache_creation_input_tokens is not None: - usage_kwargs["cache_creation_input_tokens"] = ( - cache_creation_input_tokens - ) + usage_kwargs[ + "cache_creation_input_tokens" + ] = cache_creation_input_tokens if cache_read_input_tokens is not None: usage_kwargs["cache_read_input_tokens"] = cache_read_input_tokens diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 7bebe00d61..977304f732 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -84,7 +84,7 @@ class TestProxyBaseLLMRequestProcessing: """ Test that hierarchical router settings are stored as router_settings_override instead of creating a full user_config with model_list. - + This approach avoids expensive per-request Router instantiation by passing settings as kwargs overrides to the main router. """ @@ -114,7 +114,7 @@ class TestProxyBaseLLMRequestProcessing: mock_general_settings = {} mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth) mock_proxy_config = MagicMock(spec=ProxyConfig) - + mock_router_settings = { "routing_strategy": "least-busy", "timeout": 30.0, @@ -134,7 +134,10 @@ class TestProxyBaseLLMRequestProcessing: route_type = "acompletion" - returned_data, logging_obj = await processing_obj.common_processing_pre_call_logic( + ( + returned_data, + logging_obj, + ) = await processing_obj.common_processing_pre_call_logic( request=mock_request, general_settings=mock_general_settings, user_api_key_dict=mock_user_api_key_dict, @@ -156,7 +159,7 @@ class TestProxyBaseLLMRequestProcessing: # This allows passing them as kwargs to the main router instead of creating a new one assert "router_settings_override" in returned_data assert "user_config" not in returned_data - + router_settings_override = returned_data["router_settings_override"] assert router_settings_override["routing_strategy"] == "least-busy" assert router_settings_override["timeout"] == 30.0 @@ -173,34 +176,39 @@ class TestProxyBaseLLMRequestProcessing: # Test with stream timeout header headers_with_timeout = {"x-litellm-stream-timeout": "30.5"} - result = LiteLLMProxyRequestSetup._get_stream_timeout_from_request(headers_with_timeout) + result = LiteLLMProxyRequestSetup._get_stream_timeout_from_request( + headers_with_timeout + ) assert result == 30.5 - + # Test without stream timeout header headers_without_timeout = {} - result = LiteLLMProxyRequestSetup._get_stream_timeout_from_request(headers_without_timeout) + result = LiteLLMProxyRequestSetup._get_stream_timeout_from_request( + headers_without_timeout + ) assert result is None - + # Test with invalid header value (should raise ValueError when converting to float) headers_with_invalid = {"x-litellm-stream-timeout": "invalid"} with pytest.raises(ValueError): - LiteLLMProxyRequestSetup._get_stream_timeout_from_request(headers_with_invalid) + LiteLLMProxyRequestSetup._get_stream_timeout_from_request( + headers_with_invalid + ) @pytest.mark.asyncio async def test_add_litellm_data_to_request_with_stream_timeout_header(self): """ - Test that x-litellm-stream-timeout header gets processed and added to request data + Test that x-litellm-stream-timeout header gets processed and added to request data when calling add_litellm_data_to_request. """ - from litellm.integrations.opentelemetry import UserAPIKeyAuth from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request # Create test data with a basic completion request test_data = { "model": "gpt-3.5-turbo", - "messages": [{"role": "user", "content": "Hello"}] + "messages": [{"role": "user", "content": "Hello"}], } - + # Mock request with stream timeout header mock_request = MagicMock(spec=Request) mock_request.headers = {"x-litellm-stream-timeout": "45.0"} @@ -208,7 +216,7 @@ class TestProxyBaseLLMRequestProcessing: mock_request.method = "POST" mock_request.query_params = {} mock_request.client = None - + # Create a minimal mock with just the required attributes mock_user_api_key_dict = MagicMock() mock_user_api_key_dict.api_key = "test_api_key_hash" @@ -232,10 +240,10 @@ class TestProxyBaseLLMRequestProcessing: mock_user_api_key_dict.model_max_budget = None mock_user_api_key_dict.parent_otel_span = None mock_user_api_key_dict.team_model_aliases = None - + general_settings = {} mock_proxy_config = MagicMock() - + # Call the actual function that processes headers and adds data result_data = await add_litellm_data_to_request( data=test_data, @@ -245,11 +253,11 @@ class TestProxyBaseLLMRequestProcessing: version=None, proxy_config=mock_proxy_config, ) - + # Verify that stream_timeout was extracted from header and added to request data assert "stream_timeout" in result_data assert result_data["stream_timeout"] == 45.0 - + # Verify that the original test data is preserved assert result_data["model"] == "gpt-3.5-turbo" assert result_data["messages"] == [{"role": "user", "content": "Hello"}] @@ -269,7 +277,7 @@ class TestProxyBaseLLMRequestProcessing: mock_user_api_key_dict.rpm_limit = None mock_user_api_key_dict.max_budget = None mock_user_api_key_dict.spend = 0 - + # Create logging object with cost breakdown including discount logging_obj = LiteLLMLoggingObj( model="vertex_ai/gemini-pro", @@ -280,7 +288,7 @@ class TestProxyBaseLLMRequestProcessing: litellm_call_id="test-call-id", function_id="test-function-id", ) - + # Set cost breakdown with discount information logging_obj.set_cost_breakdown( input_cost=0.00005, @@ -291,7 +299,7 @@ class TestProxyBaseLLMRequestProcessing: discount_percent=0.05, discount_amount=0.000005, ) - + # Call get_custom_headers with discount info headers = ProxyBaseLLMRequestProcessing.get_custom_headers( user_api_key_dict=mock_user_api_key_dict, @@ -299,14 +307,14 @@ class TestProxyBaseLLMRequestProcessing: response_cost=0.000095, litellm_logging_obj=logging_obj, ) - + # Verify discount headers are present assert "x-litellm-response-cost" in headers assert float(headers["x-litellm-response-cost"]) == 0.000095 - + assert "x-litellm-response-cost-original" in headers assert float(headers["x-litellm-response-cost-original"]) == 0.0001 - + assert "x-litellm-response-cost-discount-amount" in headers assert float(headers["x-litellm-response-cost-discount-amount"]) == 0.000005 @@ -324,7 +332,7 @@ class TestProxyBaseLLMRequestProcessing: mock_user_api_key_dict.rpm_limit = None mock_user_api_key_dict.max_budget = None mock_user_api_key_dict.spend = 0 - + # Create logging object without discount logging_obj = LiteLLMLoggingObj( model="gpt-3.5-turbo", @@ -335,7 +343,7 @@ class TestProxyBaseLLMRequestProcessing: litellm_call_id="test-call-id", function_id="test-function-id", ) - + # Set cost breakdown without discount information logging_obj.set_cost_breakdown( input_cost=0.00005, @@ -343,7 +351,7 @@ class TestProxyBaseLLMRequestProcessing: total_cost=0.0001, cost_for_built_in_tools_cost_usd_dollar=0.0, ) - + # Call get_custom_headers headers = ProxyBaseLLMRequestProcessing.get_custom_headers( user_api_key_dict=mock_user_api_key_dict, @@ -351,11 +359,11 @@ class TestProxyBaseLLMRequestProcessing: response_cost=0.0001, litellm_logging_obj=logging_obj, ) - + # Verify discount headers are NOT present assert "x-litellm-response-cost" in headers assert float(headers["x-litellm-response-cost"]) == 0.0001 - + # Discount headers should not be in the final dict assert "x-litellm-response-cost-original" not in headers assert "x-litellm-response-cost-discount-amount" not in headers @@ -374,7 +382,7 @@ class TestProxyBaseLLMRequestProcessing: mock_user_api_key_dict.rpm_limit = None mock_user_api_key_dict.max_budget = None mock_user_api_key_dict.spend = 0 - + # Create logging object with margin logging_obj = LiteLLMLoggingObj( model="gpt-4", @@ -394,20 +402,20 @@ class TestProxyBaseLLMRequestProcessing: margin_percent=0.10, margin_total_amount=0.00001, ) - + headers = ProxyBaseLLMRequestProcessing.get_custom_headers( user_api_key_dict=mock_user_api_key_dict, response_cost=0.00011, litellm_logging_obj=logging_obj, ) - + # Verify margin headers are present assert "x-litellm-response-cost" in headers assert float(headers["x-litellm-response-cost"]) == 0.00011 - + assert "x-litellm-response-cost-margin-amount" in headers assert float(headers["x-litellm-response-cost-margin-amount"]) == 0.00001 - + assert "x-litellm-response-cost-margin-percent" in headers assert float(headers["x-litellm-response-cost-margin-percent"]) == 0.10 @@ -425,7 +433,7 @@ class TestProxyBaseLLMRequestProcessing: mock_user_api_key_dict.rpm_limit = None mock_user_api_key_dict.max_budget = None mock_user_api_key_dict.spend = 0 - + # Create logging object without margin logging_obj = LiteLLMLoggingObj( model="gpt-4", @@ -442,13 +450,13 @@ class TestProxyBaseLLMRequestProcessing: total_cost=0.0001, cost_for_built_in_tools_cost_usd_dollar=0.0, ) - + headers = ProxyBaseLLMRequestProcessing.get_custom_headers( user_api_key_dict=mock_user_api_key_dict, response_cost=0.0001, litellm_logging_obj=logging_obj, ) - + # Verify margin headers are not present assert "x-litellm-response-cost-margin-amount" not in headers assert "x-litellm-response-cost-margin-percent" not in headers @@ -480,13 +488,18 @@ class TestProxyBaseLLMRequestProcessing: discount_percent=0.05, discount_amount=0.000005, ) - - original_cost, discount_amount, margin_total_amount, margin_percent = _get_cost_breakdown_from_logging_obj(logging_obj) + + ( + original_cost, + discount_amount, + margin_total_amount, + margin_percent, + ) = _get_cost_breakdown_from_logging_obj(logging_obj) assert original_cost == 0.0001 assert discount_amount == 0.000005 assert margin_total_amount is None assert margin_percent is None - + # Test with margin info logging_obj_with_margin = LiteLLMLoggingObj( model="gpt-4", @@ -506,13 +519,18 @@ class TestProxyBaseLLMRequestProcessing: margin_percent=0.10, margin_total_amount=0.00001, ) - - original_cost, discount_amount, margin_total_amount, margin_percent = _get_cost_breakdown_from_logging_obj(logging_obj_with_margin) + + ( + original_cost, + discount_amount, + margin_total_amount, + margin_percent, + ) = _get_cost_breakdown_from_logging_obj(logging_obj_with_margin) assert original_cost == 0.0001 assert discount_amount is None assert margin_total_amount == 0.00001 assert margin_percent == 0.10 - + # Test with no discount or margin info logging_obj_no_discount = LiteLLMLoggingObj( model="gpt-3.5-turbo", @@ -529,15 +547,25 @@ class TestProxyBaseLLMRequestProcessing: total_cost=0.0001, cost_for_built_in_tools_cost_usd_dollar=0.0, ) - - original_cost, discount_amount, margin_total_amount, margin_percent = _get_cost_breakdown_from_logging_obj(logging_obj_no_discount) + + ( + original_cost, + discount_amount, + margin_total_amount, + margin_percent, + ) = _get_cost_breakdown_from_logging_obj(logging_obj_no_discount) assert original_cost is None assert discount_amount is None assert margin_total_amount is None assert margin_percent is None - + # Test with None logging object - original_cost, discount_amount, margin_total_amount, margin_percent = _get_cost_breakdown_from_logging_obj(None) + ( + original_cost, + discount_amount, + margin_total_amount, + margin_percent, + ) = _get_cost_breakdown_from_logging_obj(None) assert original_cost is None assert discount_amount is None assert margin_total_amount is None @@ -546,7 +574,7 @@ class TestProxyBaseLLMRequestProcessing: def test_get_custom_headers_key_spend_includes_response_cost(self): """ Test that x-litellm-key-spend header includes the current request's response_cost. - + This ensures that the spend header reflects the updated spend including the current request, even though spend tracking updates happen asynchronously after the response. """ @@ -564,10 +592,12 @@ class TestProxyBaseLLMRequestProcessing: call_id="test-call-id-1", response_cost=response_cost_1, ) - + assert "x-litellm-key-spend" in headers_1 expected_spend_1 = 0.001 + 0.0005 # Initial spend + current request cost - assert float(headers_1["x-litellm-key-spend"]) == pytest.approx(expected_spend_1, abs=1e-10) + assert float(headers_1["x-litellm-key-spend"]) == pytest.approx( + expected_spend_1, abs=1e-10 + ) assert float(headers_1["x-litellm-response-cost"]) == response_cost_1 # Test case 2: response_cost is provided as string @@ -577,10 +607,12 @@ class TestProxyBaseLLMRequestProcessing: call_id="test-call-id-2", response_cost=response_cost_2, ) - + assert "x-litellm-key-spend" in headers_2 expected_spend_2 = 0.001 + 0.0003 # Initial spend + current request cost - assert float(headers_2["x-litellm-key-spend"]) == pytest.approx(expected_spend_2, abs=1e-10) + assert float(headers_2["x-litellm-key-spend"]) == pytest.approx( + expected_spend_2, abs=1e-10 + ) # Test case 3: response_cost is None (should use original spend) headers_3 = ProxyBaseLLMRequestProcessing.get_custom_headers( @@ -588,9 +620,11 @@ class TestProxyBaseLLMRequestProcessing: call_id="test-call-id-3", response_cost=None, ) - + assert "x-litellm-key-spend" in headers_3 - assert float(headers_3["x-litellm-key-spend"]) == 0.001 # Should use original spend + assert ( + float(headers_3["x-litellm-key-spend"]) == 0.001 + ) # Should use original spend # Test case 4: response_cost is 0 (should not change spend) headers_4 = ProxyBaseLLMRequestProcessing.get_custom_headers( @@ -598,9 +632,11 @@ class TestProxyBaseLLMRequestProcessing: call_id="test-call-id-4", response_cost=0.0, ) - + assert "x-litellm-key-spend" in headers_4 - assert float(headers_4["x-litellm-key-spend"]) == 0.001 # Should remain unchanged for 0 cost + assert ( + float(headers_4["x-litellm-key-spend"]) == 0.001 + ) # Should remain unchanged for 0 cost # Test case 5: user_api_key_dict.spend is None (should default to 0.0) mock_user_api_key_dict.spend = None @@ -609,7 +645,7 @@ class TestProxyBaseLLMRequestProcessing: call_id="test-call-id-5", response_cost=0.0002, ) - + assert "x-litellm-key-spend" in headers_5 assert float(headers_5["x-litellm-key-spend"]) == 0.0002 # 0.0 + 0.0002 @@ -620,9 +656,11 @@ class TestProxyBaseLLMRequestProcessing: call_id="test-call-id-6", response_cost=-0.0001, # Negative cost (should not be added) ) - + assert "x-litellm-key-spend" in headers_6 - assert float(headers_6["x-litellm-key-spend"]) == 0.001 # Should use original spend + assert ( + float(headers_6["x-litellm-key-spend"]) == 0.001 + ) # Should use original spend # Test case 7: response_cost is invalid string (should fallback to original spend) headers_7 = ProxyBaseLLMRequestProcessing.get_custom_headers( @@ -630,9 +668,77 @@ class TestProxyBaseLLMRequestProcessing: call_id="test-call-id-7", response_cost="invalid", # Invalid string ) - + assert "x-litellm-key-spend" in headers_7 - assert float(headers_7["x-litellm-key-spend"]) == 0.001 # Should use original spend on error + assert ( + float(headers_7["x-litellm-key-spend"]) == 0.001 + ) # Should use original spend on error + + @pytest.mark.asyncio + async def test_queue_time_seconds_is_set_in_metadata(self, monkeypatch): + """ + Test that queue_time_seconds is correctly calculated and stored in metadata + after add_litellm_data_to_request populates arrival_time. + + This verifies the fix for the bug where queue_time_seconds was always None + because arrival_time was read BEFORE add_litellm_data_to_request set it. + """ + processing_obj = ProxyBaseLLMRequestProcessing(data={}) + mock_request = MagicMock(spec=Request) + mock_request.headers = {} + mock_request.url = MagicMock() + mock_request.url.path = "/v1/chat/completions" + + async def mock_add_litellm_data_to_request(*args, **kwargs): + data = kwargs.get("data", args[0] if args else {}) + # Simulate what add_litellm_data_to_request does: set arrival_time + import time + + data["proxy_server_request"] = { + "url": "/v1/chat/completions", + "method": "POST", + "headers": {}, + "body": {}, + "arrival_time": time.time() - 0.5, # Simulate request arrived 0.5s ago + } + data["metadata"] = data.get("metadata", {}) + return data + + async def mock_pre_call_hook(user_api_key_dict, data, call_type): + return copy.deepcopy(data) + + mock_proxy_logging_obj = MagicMock(spec=ProxyLogging) + mock_proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=mock_pre_call_hook) + monkeypatch.setattr( + litellm.proxy.common_request_processing, + "add_litellm_data_to_request", + mock_add_litellm_data_to_request, + ) + mock_general_settings = {} + mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth) + mock_proxy_config = MagicMock(spec=ProxyConfig) + route_type = "acompletion" + + ( + returned_data, + logging_obj, + ) = await processing_obj.common_processing_pre_call_logic( + request=mock_request, + general_settings=mock_general_settings, + user_api_key_dict=mock_user_api_key_dict, + proxy_logging_obj=mock_proxy_logging_obj, + proxy_config=mock_proxy_config, + route_type=route_type, + ) + + # Verify queue_time_seconds is set and non-negative + metadata = returned_data.get("metadata", {}) + assert ( + "queue_time_seconds" in metadata + ), "queue_time_seconds should be set in metadata" + assert ( + metadata["queue_time_seconds"] >= 0.5 + ), f"queue_time_seconds should be at least 0.5, got {metadata['queue_time_seconds']}" @pytest.mark.asyncio @@ -695,19 +801,19 @@ class TestCommonRequestProcessingHelpers: Test that when the first chunk is an error, a JSON error response is returned instead of an SSE streaming response """ + async def mock_generator(): yield 'data: {"error": {"code": 403, "message": "forbidden"}}\n\n' yield 'data: {"content": "more data"}\n\n' yield "data: [DONE]\n\n" - response = await create_response( - mock_generator(), "text/event-stream", {} - ) + response = await create_response(mock_generator(), "text/event-stream", {}) # Should return JSONResponse instead of StreamingResponse assert isinstance(response, JSONResponse) assert response.status_code == status.HTTP_403_FORBIDDEN # Verify the response is in standard JSON error format import json + body = json.loads(response.body.decode()) assert "error" in body assert body["error"]["code"] == 403 @@ -719,9 +825,7 @@ class TestCommonRequestProcessingHelpers: yield 'data: {"content": "second part"}\n\n' yield "data: [DONE]\n\n" - response = await create_response( - mock_generator(), "text/event-stream", {} - ) + response = await create_response(mock_generator(), "text/event-stream", {}) assert response.status_code == status.HTTP_200_OK content = await self.consume_stream(response) assert content == [ @@ -736,9 +840,7 @@ class TestCommonRequestProcessingHelpers: yield # Implicitly raises StopAsyncIteration - response = await create_response( - mock_generator(), "text/event-stream", {} - ) + response = await create_response(mock_generator(), "text/event-stream", {}) assert response.status_code == status.HTTP_200_OK content = await self.consume_stream(response) assert content == [] @@ -780,17 +882,17 @@ class TestCommonRequestProcessingHelpers: """ Test that when the first chunk contains a string error code, a JSON error response is returned """ + async def mock_generator(): yield 'data: {"error": {"code": "429", "message": "too many requests"}}\n\n' yield "data: [DONE]\n\n" - response = await create_response( - mock_generator(), "text/event-stream", {} - ) + response = await create_response(mock_generator(), "text/event-stream", {}) assert isinstance(response, JSONResponse) assert response.status_code == status.HTTP_429_TOO_MANY_REQUESTS # Verify the response is in standard JSON error format import json + body = json.loads(response.body.decode()) assert "error" in body assert body["error"]["code"] == "429" @@ -829,9 +931,7 @@ class TestCommonRequestProcessingHelpers: async def mock_generator(): yield "data: [DONE]\n\n" - response = await create_response( - mock_generator(), "text/event-stream", {} - ) + response = await create_response(mock_generator(), "text/event-stream", {}) assert response.status_code == status.HTTP_200_OK # Default status content = await self.consume_stream(response) assert content == ["data: [DONE]\n\n"] @@ -842,9 +942,7 @@ class TestCommonRequestProcessingHelpers: yield 'data: {"content": "actual data"}\n\n' yield "data: [DONE]\n\n" - response = await create_response( - mock_generator(), "text/event-stream", {} - ) + response = await create_response(mock_generator(), "text/event-stream", {}) assert response.status_code == status.HTTP_200_OK # Default status content = await self.consume_stream(response) assert content == [ @@ -855,7 +953,6 @@ class TestCommonRequestProcessingHelpers: async def test_create_streaming_response_all_chunks_have_dd_trace(self): """Test that all stream chunks are wrapped with dd trace at the streaming generator level""" - import json from unittest.mock import patch # Create a mock tracer @@ -873,9 +970,7 @@ class TestCommonRequestProcessingHelpers: # Patch the tracer in the common_request_processing module with patch("litellm.proxy.common_request_processing.tracer", mock_tracer): - response = await create_response( - mock_generator(), "text/event-stream", {} - ) + response = await create_response(mock_generator(), "text/event-stream", {}) assert response.status_code == 200 @@ -930,9 +1025,7 @@ class TestCommonRequestProcessingHelpers: # Patch the tracer in the common_request_processing module with patch("litellm.proxy.common_request_processing.tracer", mock_tracer): - response = await create_response( - mock_generator(), "text/event-stream", {} - ) + response = await create_response(mock_generator(), "text/event-stream", {}) # Should return JSONResponse instead of StreamingResponse assert isinstance(response, JSONResponse) @@ -940,6 +1033,7 @@ class TestCommonRequestProcessingHelpers: # Verify the response is in standard JSON error format import json + body = json.loads(response.body.decode()) assert "error" in body assert body["error"]["code"] == 400 @@ -1000,7 +1094,7 @@ class TestExtractErrorFromSSEChunk: def test_extract_error_from_sse_chunk_with_invalid_json(self): """Test invalid JSON should return default error""" - chunk = 'data: {invalid json}\n\n' + chunk = "data: {invalid json}\n\n" error = _extract_error_from_sse_chunk(chunk) assert error["message"] == "Unknown error" @@ -1037,35 +1131,35 @@ class TestExtractErrorFromSSEChunk: class TestOverrideOpenAIResponseModel: """Tests for _override_openai_response_model function""" - def test_override_model_preserves_fallback_model_when_fallback_occurred_object(self): + def test_override_model_preserves_fallback_model_when_fallback_occurred_object( + self, + ): """ Test that when a fallback occurred (x-litellm-attempted-fallbacks > 0), the actual model used (fallback model) is preserved instead of being overridden with the requested model. - + This is the regression test to ensure the model being called is properly displayed when a fallback happens. """ requested_model = "gpt-4" fallback_model = "gpt-3.5-turbo" - + # Create a mock object response with fallback model # _hidden_params is an attribute (not a dict key) accessed via getattr response_obj = MagicMock() response_obj.model = fallback_model response_obj._hidden_params = { - "additional_headers": { - "x-litellm-attempted-fallbacks": 1 - } + "additional_headers": {"x-litellm-attempted-fallbacks": 1} } - + # Call the function - should preserve fallback model _override_openai_response_model( response_obj=response_obj, requested_model=requested_model, log_context="test_context", ) - + # Verify the model was NOT overridden - should still be the fallback model assert response_obj.model == fallback_model assert response_obj.model != requested_model @@ -1077,7 +1171,7 @@ class TestOverrideOpenAIResponseModel: """ requested_model = "gpt-4" fallback_model = "claude-haiku-4-5-20251001" - + # Create a mock object response with fallback model response_obj = MagicMock() response_obj.model = fallback_model @@ -1086,14 +1180,14 @@ class TestOverrideOpenAIResponseModel: "x-litellm-attempted-fallbacks": 2 # Multiple fallbacks } } - + # Call the function - should preserve fallback model _override_openai_response_model( response_obj=response_obj, requested_model=requested_model, log_context="test_context", ) - + # Verify the model was NOT overridden - should still be the fallback model assert response_obj.model == fallback_model assert response_obj.model != requested_model @@ -1105,19 +1199,19 @@ class TestOverrideOpenAIResponseModel: """ requested_model = "gpt-4" downstream_model = "gpt-3.5-turbo" - + # Create a dict response without fallback # For dict responses, _hidden_params won't be found via getattr, # so the fallback check won't trigger and model will be overridden response_obj = {"model": downstream_model} - + # Call the function - should override to requested model _override_openai_response_model( response_obj=response_obj, requested_model=requested_model, log_context="test_context", ) - + # Verify the model WAS overridden to requested model assert response_obj["model"] == requested_model @@ -1128,21 +1222,21 @@ class TestOverrideOpenAIResponseModel: """ requested_model = "gpt-4" downstream_model = "gpt-3.5-turbo" - + # Create a mock object response without fallback response_obj = MagicMock() response_obj.model = downstream_model response_obj._hidden_params = { "additional_headers": {} # No attempted_fallbacks header } - + # Call the function - should override to requested model _override_openai_response_model( response_obj=response_obj, requested_model=requested_model, log_context="test_context", ) - + # Verify the model WAS overridden to requested model assert response_obj.model == requested_model @@ -1153,7 +1247,7 @@ class TestOverrideOpenAIResponseModel: """ requested_model = "gpt-4" downstream_model = "gpt-3.5-turbo" - + # Create a mock object response response_obj = MagicMock() response_obj.model = downstream_model @@ -1162,14 +1256,14 @@ class TestOverrideOpenAIResponseModel: "x-litellm-attempted-fallbacks": 0 # Zero means no fallback occurred } } - + # Call the function - should override to requested model _override_openai_response_model( response_obj=response_obj, requested_model=requested_model, log_context="test_context", ) - + # Verify the model WAS overridden to requested model assert response_obj.model == requested_model @@ -1180,23 +1274,21 @@ class TestOverrideOpenAIResponseModel: """ requested_model = "gpt-4" downstream_model = "gpt-3.5-turbo" - + # Create a mock object response response_obj = MagicMock() response_obj.model = downstream_model response_obj._hidden_params = { - "additional_headers": { - "x-litellm-attempted-fallbacks": None - } + "additional_headers": {"x-litellm-attempted-fallbacks": None} } - + # Call the function - should override to requested model _override_openai_response_model( response_obj=response_obj, requested_model=requested_model, log_context="test_context", ) - + # Verify the model WAS overridden to requested model assert response_obj.model == requested_model @@ -1207,19 +1299,19 @@ class TestOverrideOpenAIResponseModel: """ requested_model = "gpt-4" downstream_model = "gpt-3.5-turbo" - + # Create a mock object response without _hidden_params response_obj = MagicMock() response_obj.model = downstream_model # Don't set _hidden_params - getattr will return {} - + # Call the function - should override to requested model _override_openai_response_model( response_obj=response_obj, requested_model=requested_model, log_context="test_context", ) - + # Verify the model WAS overridden to requested model assert response_obj.model == requested_model @@ -1229,34 +1321,30 @@ class TestOverrideOpenAIResponseModel: without modifying the response. """ fallback_model = "gpt-3.5-turbo" - + # Create a mock object response response_obj = MagicMock() response_obj.model = fallback_model response_obj._hidden_params = { - "additional_headers": { - "x-litellm-attempted-fallbacks": 1 - } + "additional_headers": {"x-litellm-attempted-fallbacks": 1} } - + # Call the function with None requested_model _override_openai_response_model( response_obj=response_obj, requested_model=None, log_context="test_context", ) - + # Verify the model was not changed assert response_obj.model == fallback_model - + # Call with empty string _override_openai_response_model( response_obj=response_obj, requested_model="", log_context="test_context", ) - + # Verify the model was not changed assert response_obj.model == fallback_model - -