diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index a33b491fce2..812fd8ed47d 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -127,12 +127,16 @@ class TestProxyBaseLLMRequestProcessing: assert json.loads(result.body) == guardrailed_body @pytest.mark.asyncio - async def test_handle_non_streaming_allm_passthrough_route_forwards_upstream_headers(self, monkeypatch): + async def test_handle_non_streaming_allm_passthrough_route_forwards_upstream_headers( + self, monkeypatch + ): """The guardrail JSON path must forward upstream response headers (e.g. x-amzn-requestid) alongside the x-litellm-* headers, matching the non-guardrail passthrough path, while dropping length headers that no longer match the rewritten body.""" - processing_obj = ProxyBaseLLMRequestProcessing(data={"custom_llm_provider": "bedrock"}) + processing_obj = ProxyBaseLLMRequestProcessing( + data={"custom_llm_provider": "bedrock"} + ) monkeypatch.setattr( processing_obj, "_has_post_call_guardrails_for_passthrough", @@ -172,10 +176,14 @@ class TestProxyBaseLLMRequestProcessing: assert result.headers["content-length"] == str(len(result.body)) @pytest.mark.asyncio - async def test_handle_event_stream_allm_passthrough_route_forwards_upstream_headers(self, monkeypatch): + async def test_handle_event_stream_allm_passthrough_route_forwards_upstream_headers( + self, monkeypatch + ): """The guardrail event-stream branch must also forward upstream response headers alongside the x-litellm-* headers.""" - processing_obj = ProxyBaseLLMRequestProcessing(data={"custom_llm_provider": "bedrock"}) + processing_obj = ProxyBaseLLMRequestProcessing( + data={"custom_llm_provider": "bedrock"} + ) monkeypatch.setattr( processing_obj, "_has_post_call_guardrails_for_passthrough", @@ -217,11 +225,15 @@ class TestProxyBaseLLMRequestProcessing: assert result.headers["x-litellm-call-id"] == "test-call-id" @pytest.mark.asyncio - async def test_handle_non_streaming_allm_passthrough_route_applies_response_headers_hook(self, monkeypatch): + async def test_handle_non_streaming_allm_passthrough_route_applies_response_headers_hook( + self, monkeypatch + ): """Guardrailed non-streaming passthrough responses must include headers injected by post_call_response_headers_hook, matching the headers a non-guardrailed passthrough response would carry.""" - processing_obj = ProxyBaseLLMRequestProcessing(data={"custom_llm_provider": "bedrock"}) + processing_obj = ProxyBaseLLMRequestProcessing( + data={"custom_llm_provider": "bedrock"} + ) monkeypatch.setattr( processing_obj, "_has_post_call_guardrails_for_passthrough", @@ -240,7 +252,9 @@ class TestProxyBaseLLMRequestProcessing: return kwargs["response"] proxy_logging_obj.post_call_success_hook = fake_post_call_success_hook - proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value={"x-litellm-custom": "from-hook"}) + proxy_logging_obj.post_call_response_headers_hook = AsyncMock( + return_value={"x-litellm-custom": "from-hook"} + ) result = await processing_obj._handle_non_streaming_allm_passthrough_route( response=upstream, @@ -364,7 +378,9 @@ class TestProxyBaseLLMRequestProcessing: json.dumps(persisted_body) @pytest.mark.asyncio - async def test_common_processing_pre_call_logic_arms_auto_router_compression_before_guardrails(self, monkeypatch): + async def test_common_processing_pre_call_logic_arms_auto_router_compression_before_guardrails( + self, monkeypatch + ): """arm_pre_call must run before pre_call_hook: an auto router's own compression policy has to be in `data["metadata"]` (naming the model-side guardrail so it runs even if it isn't default_on) by the time guardrails see the request.""" @@ -2175,10 +2191,16 @@ class TestCommonRequestProcessingHelpers: def _stringified_none_paths(node: object, path: str = "error") -> tuple[str, ...]: if isinstance(node, dict): - return tuple(found for key, value in node.items() for found in _stringified_none_paths(value, f"{path}.{key}")) + return tuple( + found + for key, value in node.items() + for found in _stringified_none_paths(value, f"{path}.{key}") + ) if isinstance(node, (list, tuple)): return tuple( - found for index, value in enumerate(node) for found in _stringified_none_paths(value, f"{path}[{index}]") + found + for index, value in enumerate(node) + for found in _stringified_none_paths(value, f"{path}[{index}]") ) return (path,) if node == "None" else () @@ -2925,7 +2947,9 @@ class TestStreamingOverheadHeader: user_api_key_dict=mock_user_api_key_dict, call_id="test-call-id", hidden_params={}, - litellm_logging_obj=self._timing_logging_obj({"_response_ms": 500.0, "litellm_overhead_time_ms": 42.5}), + litellm_logging_obj=self._timing_logging_obj( + {"_response_ms": 500.0, "litellm_overhead_time_ms": 42.5} + ), ) assert headers["x-litellm-response-duration-ms"] == "500.0" @@ -2944,7 +2968,9 @@ class TestStreamingOverheadHeader: user_api_key_dict=mock_user_api_key_dict, call_id="test-call-id", hidden_params={}, - litellm_logging_obj=self._timing_logging_obj({"_response_ms": 500.0, "litellm_overhead_time_ms": 42.5}), + litellm_logging_obj=self._timing_logging_obj( + {"_response_ms": 500.0, "litellm_overhead_time_ms": 42.5} + ), read_timing_from_logging_obj=False, ) @@ -2965,7 +2991,9 @@ class TestStreamingOverheadHeader: user_api_key_dict=mock_user_api_key_dict, call_id="test-call-id", hidden_params={"_response_ms": 300.0}, - litellm_logging_obj=self._timing_logging_obj({"_response_ms": 500.0, "litellm_overhead_time_ms": 42.5}), + litellm_logging_obj=self._timing_logging_obj( + {"_response_ms": 500.0, "litellm_overhead_time_ms": 42.5} + ), ) assert headers["x-litellm-response-duration-ms"] == "300.0" @@ -3006,7 +3034,9 @@ class TestStreamingOverheadHeader: user_api_key_dict=mock_user_api_key_dict, call_id="test-call-id", hidden_params={"_response_ms": 300.0, "litellm_overhead_time_ms": 7.5}, - litellm_logging_obj=self._timing_logging_obj({"_response_ms": 500.0, "litellm_overhead_time_ms": 42.5}), + litellm_logging_obj=self._timing_logging_obj( + {"_response_ms": 500.0, "litellm_overhead_time_ms": 42.5} + ), ) assert headers["x-litellm-response-duration-ms"] == "300.0" @@ -3458,7 +3488,9 @@ class TestStreamCloseOnDisconnect: finally: closed.set() - response = _UpstreamClosingStreamingResponse(body(), media_type="text/event-stream") + response = _UpstreamClosingStreamingResponse( + body(), media_type="text/event-stream" + ) async def receive(): await asyncio.Event().wait() @@ -3489,7 +3521,9 @@ class TestStreamCloseOnDisconnect: finally: closed.set() - response = _UpstreamClosingStreamingResponse(body(), media_type="text/event-stream") + response = _UpstreamClosingStreamingResponse( + body(), media_type="text/event-stream" + ) async def receive(): await disconnected.wait() @@ -3560,7 +3594,9 @@ class TestStreamCloseOnDisconnect: finally: inner_closed.set() - response = await create_response(generator=wrapped(), media_type="text/event-stream", headers={}) + response = await create_response( + generator=wrapped(), media_type="text/event-stream", headers={} + ) async def receive(): await asyncio.Event().wait() @@ -3792,7 +3828,9 @@ class TestStreamCloseOnDisconnect: with pytest.raises(_ClientDisconnectedBeforeFirstChunk): await asyncio.wait_for( - _buffer_first_chunk_honoring_disconnect(AcloseRaises(), request=self._request_that_disconnects()), + _buffer_first_chunk_honoring_disconnect( + AcloseRaises(), request=self._request_that_disconnects() + ), timeout=5, ) @@ -3808,7 +3846,9 @@ class TestStreamCloseOnDisconnect: with pytest.raises(_ClientDisconnectedBeforeFirstChunk): await asyncio.wait_for( - _buffer_first_chunk_honoring_disconnect(blocking_gen(), request=self._request_that_disconnects()), + _buffer_first_chunk_honoring_disconnect( + blocking_gen(), request=self._request_that_disconnects() + ), timeout=5, ) assert closed.is_set() @@ -3824,7 +3864,9 @@ class TestHandleLLMApiExceptionRetryAfter: user_api_key_dict = UserAPIKeyAuth(api_key="sk-test") proxy_logging_obj = MagicMock() proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) - proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value=callback_headers or {}) + proxy_logging_obj.post_call_response_headers_hook = AsyncMock( + return_value=callback_headers or {} + ) try: await processor._handle_llm_api_exception( @@ -3876,7 +3918,9 @@ class TestHandleLLMApiExceptionRetryAfter: enable_pre_call_checks=False, cooldown_list=[], ) - proxy_exc = await self._invoke(exc, callback_headers={"retry-after": "", "x-custom": "1"}) + proxy_exc = await self._invoke( + exc, callback_headers={"retry-after": "", "x-custom": "1"} + ) assert proxy_exc.headers["retry-after"] == "43" assert proxy_exc.headers["x-custom"] == "1" @@ -4105,7 +4149,9 @@ class TestDisconnectGatherCleanup: return Request(scope={"type": "http", "headers": []}, receive=receive) @pytest.mark.asyncio - async def test_base_process_llm_request_raises_499_on_client_disconnect(self, monkeypatch): + async def test_base_process_llm_request_raises_499_on_client_disconnect( + self, monkeypatch + ): """With cancel_on_disconnect enabled, base_process_llm_request returns 499.""" import asyncio @@ -4134,7 +4180,9 @@ class TestDisconnectGatherCleanup: "common_processing_pre_call_logic", AsyncMock(return_value=({"model": "gemini-2.0-flash"}, mock_logging_obj)), ) - monkeypatch.setattr(processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False)) + monkeypatch.setattr( + processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False) + ) with pytest.raises(HTTPException) as exc_info: await processing_obj.base_process_llm_request( @@ -4152,7 +4200,9 @@ class TestDisconnectGatherCleanup: assert "disconnected" in exc_info.value.detail.lower() @pytest.mark.asyncio - async def test_base_process_llm_request_reraises_cancelled_error_without_client_disconnect(self, monkeypatch): + async def test_base_process_llm_request_reraises_cancelled_error_without_client_disconnect( + self, monkeypatch + ): import asyncio import litellm.proxy.common_request_processing as cpr @@ -4177,7 +4227,9 @@ class TestDisconnectGatherCleanup: "common_processing_pre_call_logic", AsyncMock(return_value=({"model": "gemini-2.0-flash"}, mock_logging_obj)), ) - monkeypatch.setattr(processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False)) + monkeypatch.setattr( + processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False) + ) monkeypatch.setattr( cpr, "route_request", @@ -4238,7 +4290,9 @@ class TestDisconnectGatherCleanup: "common_processing_pre_call_logic", AsyncMock(return_value=({"model": "gemini-2.0-flash"}, mock_logging_obj)), ) - monkeypatch.setattr(processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False)) + monkeypatch.setattr( + processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False) + ) with pytest.raises(HTTPException): await processing_obj.base_process_llm_request( @@ -4289,7 +4343,9 @@ class TestDisconnectGatherCleanup: assert task.done() @pytest.mark.asyncio - async def test_base_process_llm_request_preserves_llm_error_after_gather(self, monkeypatch): + async def test_base_process_llm_request_preserves_llm_error_after_gather( + self, monkeypatch + ): import litellm.proxy.common_request_processing as cpr from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing @@ -4318,7 +4374,9 @@ class TestDisconnectGatherCleanup: "common_processing_pre_call_logic", AsyncMock(return_value=({"model": "gemini-2.0-flash"}, mock_logging_obj)), ) - monkeypatch.setattr(processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False)) + monkeypatch.setattr( + processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False) + ) mock_request = MagicMock(spec=Request) mock_request.is_disconnected = AsyncMock(return_value=False) @@ -4355,13 +4413,19 @@ class TestStreamingClientDisconnectLogging: "litellm_params": {"metadata": {}}, } - recorded = await _record_streaming_client_disconnect_if_needed(mock_request, request_data) + recorded = await _record_streaming_client_disconnect_if_needed( + mock_request, request_data + ) assert recorded is True assert request_data["metadata"]["client_disconnected"] is True - assert request_data["metadata"]["error_information"]["error_code"] == "499" assert ( - mock_logging_obj.model_call_details["litellm_params"]["metadata"]["error_information"]["error_code"] + request_data["metadata"]["error_information"]["error_code"] == "499" + ) + assert ( + mock_logging_obj.model_call_details["litellm_params"]["metadata"][ + "error_information" + ]["error_code"] == "499" ) @@ -4375,7 +4439,9 @@ class TestStreamingClientDisconnectLogging: mock_request.is_disconnected = AsyncMock(return_value=False) request_data = {"metadata": {}} - recorded = await _record_streaming_client_disconnect_if_needed(mock_request, request_data) + recorded = await _record_streaming_client_disconnect_if_needed( + mock_request, request_data + ) assert recorded is False assert "client_disconnected" not in request_data["metadata"] @@ -4400,12 +4466,22 @@ class TestStreamingClientDisconnectLogging: "litellm_params": {"metadata": {}}, } - recorded = await _record_streaming_client_disconnect_if_needed(mock_request, request_data) + recorded = await _record_streaming_client_disconnect_if_needed( + mock_request, request_data + ) assert recorded is True assert request_data["metadata"]["client_disconnected"] is True - assert mock_logging_obj.model_call_details["litellm_params"]["metadata"]["client_disconnected"] is True - assert mock_logging_obj.model_call_details["metadata"]["client_disconnected"] is True + assert ( + mock_logging_obj.model_call_details["litellm_params"]["metadata"][ + "client_disconnected" + ] + is True + ) + assert ( + mock_logging_obj.model_call_details["metadata"]["client_disconnected"] + is True + ) @pytest.mark.asyncio async def test_record_streaming_client_disconnect_handles_none_request_data_metadata(self): @@ -4421,11 +4497,15 @@ class TestStreamingClientDisconnectLogging: "litellm_params": {"metadata": None}, } - recorded = await _record_streaming_client_disconnect_if_needed(mock_request, request_data) + recorded = await _record_streaming_client_disconnect_if_needed( + mock_request, request_data + ) assert recorded is True assert request_data["metadata"]["client_disconnected"] is True - assert request_data["litellm_params"]["metadata"]["client_disconnected"] is True + assert ( + request_data["litellm_params"]["metadata"]["client_disconnected"] is True + ) @pytest.mark.asyncio async def test_apply_client_disconnect_metadata_none_returns_early(self): @@ -4436,7 +4516,9 @@ class TestStreamingClientDisconnectLogging: _apply_client_disconnect_metadata(None) @pytest.mark.asyncio - async def test_finalize_streaming_generator_cleanup_fires_deferred_logging(self, monkeypatch): + async def test_finalize_streaming_generator_cleanup_fires_deferred_logging( + self, monkeypatch + ): from litellm.proxy.common_request_processing import ( ProxyBaseLLMRequestProcessing, ) @@ -4468,7 +4550,9 @@ class TestStreamingClientDisconnectLogging: assert request_data["metadata"]["error_information"]["error_code"] == "499" @pytest.mark.asyncio - async def test_finalize_streaming_generator_cleanup_skips_disconnect_after_completion(self, monkeypatch): + async def test_finalize_streaming_generator_cleanup_skips_disconnect_after_completion( + self, monkeypatch + ): from litellm.proxy.common_request_processing import ( ProxyBaseLLMRequestProcessing, ) @@ -4498,7 +4582,9 @@ class TestStreamingClientDisconnectLogging: assert "client_disconnected" not in request_data["metadata"] @pytest.mark.asyncio - async def test_async_streaming_data_generator_records_499_on_early_aclose(self, monkeypatch): + async def test_async_streaming_data_generator_records_499_on_early_aclose( + self, monkeypatch + ): from litellm.proxy.common_request_processing import ( ProxyBaseLLMRequestProcessing, ) @@ -4513,7 +4599,9 @@ class TestStreamingClientDisconnectLogging: yield {"choices": [{"delta": {"content": " there"}}]} mock_proxy_logging = MagicMock(spec=ProxyLogging) - mock_proxy_logging.async_post_call_streaming_iterator_hook = mock_streaming_iterator + mock_proxy_logging.async_post_call_streaming_iterator_hook = ( + mock_streaming_iterator + ) ProxyLogging._callback_capabilities_cache.clear() mock_request = MagicMock(spec=Request) @@ -4524,7 +4612,9 @@ class TestStreamingClientDisconnectLogging: "model": "gemini-2.0-flash", "metadata": {}, "litellm_params": {"metadata": {}}, - "litellm_logging_obj": MagicMock(model_call_details={"metadata": {}, "litellm_params": {}}), + "litellm_logging_obj": MagicMock( + model_call_details={"metadata": {}, "litellm_params": {}} + ), } gen = ProxyBaseLLMRequestProcessing.async_streaming_data_generator( @@ -4543,8 +4633,6 @@ class TestStreamingClientDisconnectLogging: assert request_data["metadata"]["error_information"]["error_code"] == "499" ProxyLogging._callback_capabilities_cache.clear() - - class TestCancelOnDisconnect: """ Coverage for the opt-in `general_settings.cancel_on_disconnect` flag: @@ -4571,17 +4659,23 @@ class TestCancelOnDisconnect: llm_call = asyncio.get_running_loop().create_future() disconnect_event = asyncio.Event() - await _cancel_llm_call_on_client_disconnect(request, llm_call, disconnect_event) + await _cancel_llm_call_on_client_disconnect( + request, llm_call, disconnect_event + ) assert llm_call.cancelled() assert disconnect_event.is_set() async def test_monitor_is_noop_while_client_stays_connected(self): - request = self._request([{"type": "http.request", "body": b"", "more_body": False}]) + request = self._request( + [{"type": "http.request", "body": b"", "more_body": False}] + ) llm_call = asyncio.get_running_loop().create_future() disconnect_event = asyncio.Event() - monitor = asyncio.create_task(_cancel_llm_call_on_client_disconnect(request, llm_call, disconnect_event)) + monitor = asyncio.create_task( + _cancel_llm_call_on_client_disconnect(request, llm_call, disconnect_event) + ) await asyncio.sleep(0.01) assert not monitor.done() @@ -4600,7 +4694,9 @@ class TestCancelOnDisconnect: llm_call = asyncio.get_running_loop().create_future() disconnect_event = asyncio.Event() - await _cancel_llm_call_on_client_disconnect(request, llm_call, disconnect_event) + await _cancel_llm_call_on_client_disconnect( + request, llm_call, disconnect_event + ) assert not llm_call.cancelled() assert not disconnect_event.is_set() @@ -4615,7 +4711,9 @@ class TestCancelOnDisconnect: with pytest.raises(asyncio.CancelledError): await _await_llm_call_cancelling_on_disconnect(request, llm_call) - async def _drive_base_process_llm_request(self, monkeypatch, general_settings: dict, llm_call, request: Request): + async def _drive_base_process_llm_request( + self, monkeypatch, general_settings: dict, llm_call, request: Request + ): from litellm.proxy._types import UserAPIKeyAuth logging_obj = MagicMock() @@ -4624,7 +4722,9 @@ class TestCancelOnDisconnect: logging_obj._on_deferred_stream_complete = None logging_obj.cost_breakdown = None - processor = ProxyBaseLLMRequestProcessing(data={"model": "fake-model", "litellm_logging_obj": logging_obj}) + processor = ProxyBaseLLMRequestProcessing( + data={"model": "fake-model", "litellm_logging_obj": logging_obj} + ) proxy_logging_obj = MagicMock(spec=ProxyLogging) proxy_logging_obj.during_call_hook = AsyncMock(return_value=None) @@ -4632,7 +4732,9 @@ class TestCancelOnDisconnect: proxy_logging_obj.post_call_success_hook = AsyncMock( side_effect=lambda data, user_api_key_dict, response: response ) - proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value=None) + proxy_logging_obj.post_call_response_headers_hook = AsyncMock( + return_value=None + ) async def fake_route_request(**kwargs): return llm_call() @@ -4711,7 +4813,9 @@ class TestCancelOnDisconnect: with pytest.raises(ProxyException) as exc_info: await processor._handle_llm_api_exception( - e=HTTPException(status_code=499, detail="Client disconnected the request"), + e=HTTPException( + status_code=499, detail="Client disconnected the request" + ), user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), proxy_logging_obj=proxy_logging_obj, ) @@ -4777,9 +4881,7 @@ class TestAllmPassthroughRoutePostCallGuardrails: proxy_logging_obj = ProxyLogging(user_api_key_cache=MagicMock()) monkeypatch.setattr(proxy_logging_obj, "post_call_success_hook", capture_hook) - with patch.object( - ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True - ): + with patch.object(ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True): processing_obj = ProxyBaseLLMRequestProcessing(data={}) result = await processing_obj._handle_non_streaming_allm_passthrough_route( response=httpx_response, @@ -4829,9 +4931,7 @@ class TestAllmPassthroughRoutePostCallGuardrails: proxy_logging_obj = ProxyLogging(user_api_key_cache=MagicMock()) monkeypatch.setattr(proxy_logging_obj, "post_call_success_hook", non_dict_hook) - with patch.object( - ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True - ): + with patch.object(ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True): processing_obj = ProxyBaseLLMRequestProcessing(data={}) result = await processing_obj._handle_non_streaming_allm_passthrough_route( response=httpx_response, @@ -4869,9 +4969,7 @@ class TestAllmPassthroughRoutePostCallGuardrails: hook_spy = AsyncMock() monkeypatch.setattr(proxy_logging_obj, "post_call_success_hook", hook_spy) - with patch.object( - ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True - ): + with patch.object(ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True): processing_obj = ProxyBaseLLMRequestProcessing(data={}) result = await processing_obj._handle_non_streaming_allm_passthrough_route( response=httpx_response, @@ -4912,9 +5010,7 @@ class TestAllmPassthroughRoutePostCallGuardrails: hook_spy = AsyncMock() monkeypatch.setattr(proxy_logging_obj, "post_call_success_hook", hook_spy) - with patch.object( - ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=False - ): + with patch.object(ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=False): processing_obj = ProxyBaseLLMRequestProcessing(data={}) result = await processing_obj._handle_non_streaming_allm_passthrough_route( response=httpx_response, @@ -5026,9 +5122,7 @@ class TestEventStreamAllmPassthroughRoute: "content-length": "99", } - with patch.object( - ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True - ): + with patch.object(ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True): processing_obj = ProxyBaseLLMRequestProcessing(data={}) result = await processing_obj._handle_non_streaming_allm_passthrough_route( response=mock_response, @@ -5059,7 +5153,9 @@ class TestAllmPassthroughStreamingProviderGate: de-anonymized. """ - def _build_processing_obj(self, custom_llm_provider: str, endpoint: str = "") -> ProxyBaseLLMRequestProcessing: + def _build_processing_obj( + self, custom_llm_provider: str, endpoint: str = "" + ) -> ProxyBaseLLMRequestProcessing: logging_obj = MagicMock() logging_obj.litellm_call_id = "call-123" logging_obj.cost_breakdown = None @@ -5146,17 +5242,14 @@ class TestAllmPassthroughStreamingProviderGate: processing_obj = self._build_processing_obj("anthropic") chunks = [b"chunk-1", b"chunk-2"] - with ( - patch.object( - ProxyBaseLLMRequestProcessing, - "_has_post_call_guardrails", - return_value=False, - ), - patch.object( - ProxyBaseLLMRequestProcessing, - "_has_post_call_guardrails_for_passthrough", - return_value=True, - ), + with patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails", + return_value=False, + ), patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails_for_passthrough", + return_value=True, ): result = await self._run(processing_obj, monkeypatch, chunks) @@ -5165,27 +5258,27 @@ class TestAllmPassthroughStreamingProviderGate: assert streamed == chunks @pytest.mark.asyncio - async def test_bedrock_converse_stream_is_buffered_through_handler(self, monkeypatch): - processing_obj = self._build_processing_obj("bedrock", "model/us.amazon.nova-lite-v1:0/converse-stream") + async def test_bedrock_converse_stream_is_buffered_through_handler( + self, monkeypatch + ): + processing_obj = self._build_processing_obj( + "bedrock", "model/us.amazon.nova-lite-v1:0/converse-stream" + ) chunks = [b"raw-1", b"raw-2"] - with ( - patch.object( - ProxyBaseLLMRequestProcessing, - "_has_post_call_guardrails", - return_value=False, - ), - patch.object( - ProxyBaseLLMRequestProcessing, - "_has_post_call_guardrails_for_passthrough", - return_value=True, - ), - patch( - "litellm.llms.bedrock.passthrough.guardrail_translation.handler." - "BedrockPassthroughGuardrailHandler.de_anonymize_event_stream", - new=AsyncMock(return_value=b"modified-body"), - ) as mock_handler, - ): + with patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails", + return_value=False, + ), patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails_for_passthrough", + return_value=True, + ), patch( + "litellm.llms.bedrock.passthrough.guardrail_translation.handler." + "BedrockPassthroughGuardrailHandler.de_anonymize_event_stream", + new=AsyncMock(return_value=b"modified-body"), + ) as mock_handler: result = await self._run(processing_obj, monkeypatch, chunks) assert isinstance(result, Response) @@ -5201,23 +5294,19 @@ class TestAllmPassthroughStreamingProviderGate: ) chunks = [b"raw-1", b"raw-2"] - with ( - patch.object( - ProxyBaseLLMRequestProcessing, - "_has_post_call_guardrails", - return_value=False, - ), - patch.object( - ProxyBaseLLMRequestProcessing, - "_has_post_call_guardrails_for_passthrough", - return_value=True, - ), - patch( - "litellm.llms.bedrock.passthrough.guardrail_translation.handler." - "BedrockPassthroughGuardrailHandler.de_anonymize_event_stream", - new=AsyncMock(return_value=b"modified-body"), - ) as mock_handler, - ): + with patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails", + return_value=False, + ), patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails_for_passthrough", + return_value=True, + ), patch( + "litellm.llms.bedrock.passthrough.guardrail_translation.handler." + "BedrockPassthroughGuardrailHandler.de_anonymize_event_stream", + new=AsyncMock(return_value=b"modified-body"), + ) as mock_handler: result = await self._run(processing_obj, monkeypatch, chunks) assert isinstance(result, StreamingResponse) @@ -5239,17 +5328,14 @@ class TestAllmPassthroughStreamingProviderGate: ) chunks = [b"raw-1", b"raw-2"] - with ( - patch.object( - ProxyBaseLLMRequestProcessing, - "_has_post_call_guardrails", - return_value=False, - ), - patch.object( - ProxyBaseLLMRequestProcessing, - "_has_post_call_guardrails_for_passthrough", - return_value=False, - ), + with patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails", + return_value=False, + ), patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails_for_passthrough", + return_value=False, ): result = await self._run(processing_obj, monkeypatch, chunks) @@ -5268,17 +5354,14 @@ class TestAllmPassthroughStreamingProviderGate: processing_obj = self._build_processing_obj("anthropic") chunks = [b"chunk-1", b"chunk-2"] - with ( - patch.object( - ProxyBaseLLMRequestProcessing, - "_has_post_call_guardrails", - return_value=False, - ), - patch.object( - ProxyBaseLLMRequestProcessing, - "_has_post_call_guardrails_for_passthrough", - return_value=False, - ), + with patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails", + return_value=False, + ), patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails_for_passthrough", + return_value=False, ): result = await self._run(processing_obj, monkeypatch, chunks) @@ -5726,7 +5809,9 @@ class TestCostHeadersForCallsPricedAtZero: fastapi_response = Response() processing_obj = ProxyBaseLLMRequestProcessing(data={"litellm_logging_obj": logging_obj}) - with patch.object(ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails", return_value=False): + with patch.object( + ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails", return_value=False + ): await processing_obj.base_process_llm_request( request=MagicMock(spec=Request, headers={}), fastapi_response=fastapi_response, @@ -5797,7 +5882,9 @@ class TestCostHeadersForCallsPricedAtZero: assert breakdown.tool_usage_cost == 0.0 def test_cost_breakdown_stays_empty_for_an_inference_call(self): - breakdown = _get_cost_breakdown_from_logging_obj(litellm_logging_obj=self._logging_obj(call_type="acompletion")) + breakdown = _get_cost_breakdown_from_logging_obj( + litellm_logging_obj=self._logging_obj(call_type="acompletion") + ) assert breakdown == CostBreakdownHeaderValues() @@ -5824,6 +5911,7 @@ class TestCostHeadersForCallsPricedAtZero: class TestPreCallWithFallbacksOnLocalRateLimit: + @pytest.mark.asyncio async def test_fallback_triggered_on_local_rate_limit(self): from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError @@ -5975,7 +6063,9 @@ class TestPreCallWithFallbacksOnLocalRateLimit: mock_router.fallbacks = [{"gpt-4": ["gpt-3.5-turbo"]}] user_api_key_dict = MagicMock() - user_api_key_dict.router_settings = {"fallbacks": [{"gpt-4": ["claude-3-haiku"]}]} + user_api_key_dict.router_settings = { + "fallbacks": [{"gpt-4": ["claude-3-haiku"]}] + } with patch.object( processor, @@ -6006,7 +6096,9 @@ class TestPreCallWithFallbacksOnLocalRateLimit: from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing - processor = ProxyBaseLLMRequestProcessing(data={"model": "gpt-4", "disable_fallbacks": True}) + processor = ProxyBaseLLMRequestProcessing( + data={"model": "gpt-4", "disable_fallbacks": True} + ) async def mock_pre_call_logic(**kwargs): raise ProxyRateLimitError( @@ -6132,7 +6224,9 @@ class TestPreCallWithFallbacksOnLocalRateLimit: # Real per-key per-model TPM limiter + a key carrying the customer's # `model_tpm_limit` metadata (only the primary is capped). - limiter = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) + limiter = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache()) + ) user_api_key_dict = UserAPIKeyAuth( api_key="sk-lit3890", metadata={"model_tpm_limit": {primary_model: 100}}, @@ -6140,7 +6234,10 @@ class TestPreCallWithFallbacksOnLocalRateLimit: # Pre-seed the primary's per-model token counter at the cap so the very # next request trips it. The counter key uses the *hashed* api_key. - counter_key = f"{user_api_key_dict.api_key}::{primary_model}::{precise_minute}::request_count" + counter_key = ( + f"{user_api_key_dict.api_key}::{primary_model}" + f"::{precise_minute}::request_count" + ) await limiter.internal_usage_cache.async_set_cache( key=counter_key, value={"current_requests": 0, "current_tpm": 100, "current_rpm": 0}, @@ -6171,7 +6268,9 @@ class TestPreCallWithFallbacksOnLocalRateLimit: mock_router = MagicMock() mock_router.fallbacks = [{primary_model: [fallback_model]}] - with patch("litellm.proxy.hooks.parallel_request_limiter.datetime", _FrozenClock): + with patch( + "litellm.proxy.hooks.parallel_request_limiter.datetime", _FrozenClock + ): with patch.object( processor, "common_processing_pre_call_logic", @@ -6201,7 +6300,9 @@ class TestPreCallWithFallbacksOnLocalRateLimit: # Sanity-check the premise: the limiter genuinely raises a # ProxyRateLimitError for the capped primary under the frozen clock. - with patch("litellm.proxy.hooks.parallel_request_limiter.datetime", _FrozenClock): + with patch( + "litellm.proxy.hooks.parallel_request_limiter.datetime", _FrozenClock + ): with pytest.raises(ProxyRateLimitError): await limiter.async_pre_call_hook( user_api_key_dict=user_api_key_dict, @@ -6562,12 +6663,16 @@ class TestStreamingClientDisconnectBilling: prompt_tokens=1000, completion_tokens=10, total_tokens=1010, - prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=500), + prompt_tokens_details=PromptTokensDetailsWrapper( + cached_tokens=500 + ), ), ) ) - event = await self._bill_and_collect_success_event(append_openai_style_cached_usage_chunk) + event = await self._bill_and_collect_success_event( + append_openai_style_cached_usage_chunk + ) usage = event["response_obj"].usage assert getattr(usage, "cache_read_input_tokens", None) == 500 @@ -7337,7 +7442,9 @@ class TestInjectCostIntoUsageDict: logging_obj.model_call_details["custom_llm_provider"] = "anthropic" assert logging_obj.cost_breakdown is None - model_response = ModelResponse(usage=Usage(prompt_tokens=3216, completion_tokens=8, total_tokens=3224)) + model_response = ModelResponse( + usage=Usage(prompt_tokens=3216, completion_tokens=8, total_tokens=3224) + ) cost = ProxyBaseLLMRequestProcessing._logging_obj_cost_or_none(model_response, logging_obj) assert cost is not None and cost > 0 @@ -7366,7 +7473,9 @@ class TestInjectCostIntoUsageDict: ) existing = logging_obj.cost_breakdown - model_response = ModelResponse(usage=Usage(prompt_tokens=3216, completion_tokens=8, total_tokens=3224)) + model_response = ModelResponse( + usage=Usage(prompt_tokens=3216, completion_tokens=8, total_tokens=3224) + ) ProxyBaseLLMRequestProcessing._logging_obj_cost_or_none(model_response, logging_obj) assert logging_obj.cost_breakdown is existing @@ -7661,7 +7770,9 @@ def test_ttft_keepalive_interval_only_arms_for_a_streaming_request(request_data, @pytest.mark.asyncio @pytest.mark.parametrize("stream_requested, expect_ping", [(True, True), (False, False)]) -async def test_base_process_llm_request_pings_while_the_upstream_call_is_still_running(stream_requested, expect_ping): +async def test_base_process_llm_request_pings_while_the_upstream_call_is_still_running( + stream_requested, expect_ping +): """The wiring, not the helper: every route funnels through this method, and the whole time-to-first-token is spent inside the call it wraps.""" @@ -7807,7 +7918,9 @@ async def test_a_late_failure_is_reported_to_the_failure_hook(): async def record(exc): audited.append(exc) - response = await open_sse_before_first_byte(slow_failure(), ping_interval_seconds=0.05, on_late_failure=record) + response = await open_sse_before_first_byte( + slow_failure(), ping_interval_seconds=0.05, on_late_failure=record + ) collected = await _drain(response) assert [type(exc).__name__ for exc in audited] == ["HTTPException"] @@ -7824,7 +7937,9 @@ async def test_a_failing_audit_hook_never_costs_the_client_its_error_frame(): async def broken_hook(exc): raise RuntimeError("the audit backend is down") - response = await open_sse_before_first_byte(slow_failure(), ping_interval_seconds=0.05, on_late_failure=broken_hook) + response = await open_sse_before_first_byte( + slow_failure(), ping_interval_seconds=0.05, on_late_failure=broken_hook + ) collected = await _drain(response) error_frame = json.loads(collected[-2].decode().removeprefix("data: ").strip()) @@ -7878,7 +7993,9 @@ async def test_base_process_llm_request_audits_a_failure_that_lands_after_its_ke [(0, False), (None, True)], ids=["operator-hard-disabled-this-deployment", "deployment-says-nothing"], ) -async def test_base_process_llm_request_honours_a_deployment_hard_disable(deployment_keepalive, expect_ping): +async def test_base_process_llm_request_honours_a_deployment_hard_disable( + deployment_keepalive, expect_ping +): """`keepalive_seconds: 0` is documented as a disable a request cannot lift. The funnel has to hand its router to the gate for that to hold before the upstream has answered, since no deployment has served the request yet.""" @@ -7924,7 +8041,9 @@ async def test_a_hook_returning_a_replacement_decides_what_the_client_sees(): async def sanitize(exc): return HTTPException(status_code=502, detail="upstream unavailable") - response = await open_sse_before_first_byte(slow_failure(), ping_interval_seconds=0.05, on_late_failure=sanitize) + response = await open_sse_before_first_byte( + slow_failure(), ping_interval_seconds=0.05, on_late_failure=sanitize + ) collected = await _drain(response) error_frame = json.loads(collected[-2].decode().removeprefix("data: ").strip()) @@ -7963,7 +8082,9 @@ async def test_a_hook_that_returns_nothing_leaves_the_real_error_intact(): async def audit_only(exc): return None - response = await open_sse_before_first_byte(slow_failure(), ping_interval_seconds=0.05, on_late_failure=audit_only) + response = await open_sse_before_first_byte( + slow_failure(), ping_interval_seconds=0.05, on_late_failure=audit_only + ) collected = await _drain(response) error_frame = json.loads(collected[-2].decode().removeprefix("data: ").strip()) @@ -7980,7 +8101,9 @@ async def test_a_broken_hook_does_not_replace_the_real_error_with_its_own_bug(): async def broken_hook(exc): raise RuntimeError("the audit backend is down") - response = await open_sse_before_first_byte(slow_failure(), ping_interval_seconds=0.05, on_late_failure=broken_hook) + response = await open_sse_before_first_byte( + slow_failure(), ping_interval_seconds=0.05, on_late_failure=broken_hook + ) collected = await _drain(response) error_frame = json.loads(collected[-2].decode().removeprefix("data: ").strip()) @@ -8193,7 +8316,9 @@ class TestStreamingResponseHeadersFollowFallback: proxy_logging_obj.post_call_success_hook = AsyncMock( side_effect=lambda data, user_api_key_dict, response: response ) - proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value={"x-callback-header": "kept"}) + proxy_logging_obj.post_call_response_headers_hook = AsyncMock( + return_value={"x-callback-header": "kept"} + ) async def fake_route_request(**kwargs): async def call(): @@ -8201,7 +8326,9 @@ class TestStreamingResponseHeadersFollowFallback: return call() - monkeypatch.setattr(litellm.proxy.common_request_processing, "route_request", fake_route_request) + monkeypatch.setattr( + litellm.proxy.common_request_processing, "route_request", fake_route_request + ) result = await processor.base_process_llm_request( request=Request(scope={"type": "http", "headers": []}),