diff --git a/litellm/llms/azure_ai/common_utils.py b/litellm/llms/azure_ai/common_utils.py index daf83120710..e8431f29f07 100644 --- a/litellm/llms/azure_ai/common_utils.py +++ b/litellm/llms/azure_ai/common_utils.py @@ -50,6 +50,7 @@ def get_azure_ai_auth_headers( "identity with `litellm.enable_azure_ad_token_refresh = True`)" ) + AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY: Final = "azure_model_router_selected_model" diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 25387cc38df..cc738d68016 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -273,7 +273,9 @@ def test_response_cost_calculator_uses_router_model_id_from_litellm_metadata(): assert cost is not None, "Cost should not be None" expected_cost = (10 * custom_input_cost) + (5 * custom_output_cost) - assert cost == pytest.approx(expected_cost), f"Expected {expected_cost}, got {cost}" + assert cost == pytest.approx( + expected_cost + ), f"Expected {expected_cost}, got {cost}" finally: litellm.model_cost.pop(custom_model_id, None) @@ -870,8 +872,13 @@ async def test_datadog_logger_not_shadowed_by_llm_obs(monkeypatch): # Regression check: we expect a distinct DataDogLogger, not the LLM Obs logger assert type(datadog_logger) is DataDogLogger - assert any(isinstance(cb, DataDogLLMObsLogger) for cb in logging_module._in_memory_loggers) - assert any(type(cb) is DataDogLogger for cb in logging_module._in_memory_loggers) + assert any( + isinstance(cb, DataDogLLMObsLogger) + for cb in logging_module._in_memory_loggers + ) + assert any( + type(cb) is DataDogLogger for cb in logging_module._in_memory_loggers + ) finally: logging_module._in_memory_loggers.clear() @@ -882,7 +889,9 @@ async def test_logfire_logger_accepts_env_vars_for_base_url(monkeypatch): # Required env vars for Logfire integration monkeypatch.setenv("LOGFIRE_TOKEN", "test-token") - monkeypatch.setenv("LOGFIRE_BASE_URL", "https://logfire-api-custom.pydantic.dev") # no trailing slash on purpose + monkeypatch.setenv( + "LOGFIRE_BASE_URL", "https://logfire-api-custom.pydantic.dev" + ) # no trailing slash on purpose # Import after env vars are set (important if module-level caching exists) from litellm.integrations.opentelemetry import OpenTelemetry # logger class @@ -901,7 +910,9 @@ async def test_logfire_logger_accepts_env_vars_for_base_url(monkeypatch): # Sanity: we got the right logger type and it is cached assert type(logger) is OpenTelemetry - assert any(type(cb) is OpenTelemetry for cb in logging_module._in_memory_loggers) + assert any( + type(cb) is OpenTelemetry for cb in logging_module._in_memory_loggers + ) # Core regression check: base URL env var should influence the exporter endpoint. # @@ -912,7 +923,9 @@ async def test_logfire_logger_accepts_env_vars_for_base_url(monkeypatch): or getattr(logger, "config", None) or getattr(logger, "_otel_config", None) ) - assert cfg is not None, "Expected OpenTelemetry logger to keep an otel config on the instance" + assert ( + cfg is not None + ), "Expected OpenTelemetry logger to keep an otel config on the instance" endpoint = getattr(cfg, "endpoint", None) or getattr(cfg, "otlp_endpoint", None) assert endpoint is not None, "Expected otel config to expose the OTLP endpoint" @@ -1070,7 +1083,9 @@ async def test_logging_non_streaming_request(): # Use the filtered call for assertions call_args = calls_with_expected_input[0] - standard_logging_object = call_args.kwargs["kwargs"]["standard_logging_object"] + standard_logging_object = call_args.kwargs["kwargs"][ + "standard_logging_object" + ] assert standard_logging_object["stream"] is not True finally: # Restore original callbacks to ensure test isolation @@ -1088,14 +1103,18 @@ async def test_logging_non_streaming_request(): "agenerate_content_stream", ], ) -def test_success_handler_skips_sync_callbacks_for_async_requests(logging_obj, async_flag): +def test_success_handler_skips_sync_callbacks_for_async_requests( + logging_obj, async_flag +): """Ensure sync success callbacks are skipped when async call type flags are set.""" from litellm.integrations.custom_logger import CustomLogger class DummyLogger(CustomLogger): pass - logging_obj.stream = False # simulate non-streaming request where sync callbacks would normally run + logging_obj.stream = ( + False # simulate non-streaming request where sync callbacks would normally run + ) logging_obj.model_call_details["litellm_params"] = {async_flag: True} logging_obj.litellm_params = logging_obj.model_call_details["litellm_params"] @@ -1171,11 +1190,21 @@ def test_success_handler_runs_sync_callbacks_for_sync_requests(logging_obj, call def test_is_sync_litellm_request(): assert LitellmLogging._is_sync_litellm_request({}) is True assert LitellmLogging._is_sync_litellm_request({"acompletion": True}) is False - assert LitellmLogging._is_sync_litellm_request({"allm_passthrough_route": True}) is False - assert LitellmLogging._is_sync_litellm_request({"aanthropic_messages": True}) is False + assert ( + LitellmLogging._is_sync_litellm_request({"allm_passthrough_route": True}) + is False + ) + assert ( + LitellmLogging._is_sync_litellm_request({"aanthropic_messages": True}) is False + ) assert LitellmLogging._is_sync_litellm_request({"agenerate_content": True}) is False - assert LitellmLogging._is_sync_litellm_request({"agenerate_content_stream": True}) is False - assert LitellmLogging._is_sync_litellm_request({"aanthropic_messages": False}) is True + assert ( + LitellmLogging._is_sync_litellm_request({"agenerate_content_stream": True}) + is False + ) + assert ( + LitellmLogging._is_sync_litellm_request({"aanthropic_messages": False}) is True + ) def test_get_litellm_params_propagates_allm_passthrough_route(): @@ -1222,7 +1251,9 @@ async def test_dispatch_success_handlers_invokes_callbacks_once_for_final_stream logging_obj.model_call_details["litellm_params"] = {"acompletion": True} with ( - patch.object(mock_callback, "async_log_success_event", new_callable=AsyncMock) as mock_async_log, + patch.object( + mock_callback, "async_log_success_event", new_callable=AsyncMock + ) as mock_async_log, patch.object(mock_callback, "log_success_event") as mock_sync_log, patch.object( logging_obj, @@ -1283,7 +1314,9 @@ async def test_dispatch_success_handlers_sync_path_invokes_callback_once_for_fin with ( patch.object(mock_callback, "log_success_event") as mock_sync_log, - patch.object(mock_callback, "async_log_success_event", new_callable=AsyncMock) as mock_async_log, + patch.object( + mock_callback, "async_log_success_event", new_callable=AsyncMock + ) as mock_async_log, patch.object( logging_obj, "_success_handler_helper_fn", @@ -1325,14 +1358,20 @@ async def test_dispatch_prefer_async_handlers_runs_legacy_callbacks( logging_obj.model_call_details["litellm_params"] = {} with ( - patch.object(logging_obj, "async_success_handler", new_callable=AsyncMock) as mock_async, - patch.object(logging_obj, "success_handler", new_callable=MagicMock) as mock_sync, + patch.object( + logging_obj, "async_success_handler", new_callable=AsyncMock + ) as mock_async, + patch.object( + logging_obj, "success_handler", new_callable=MagicMock + ) as mock_sync, patch.object( logging_obj, "_should_run_sync_callbacks_for_async_calls", return_value=True, ), - patch("litellm.litellm_core_utils.litellm_logging.executor.submit") as mock_submit, + patch( + "litellm.litellm_core_utils.litellm_logging.executor.submit" + ) as mock_submit, ): await logging_obj.dispatch_success_handlers( result=result, @@ -1366,7 +1405,9 @@ async def test_dispatch_success_handlers_invokes_async_callback_for_pass_through try: with ( - patch.object(mock_callback, "async_log_success_event", new_callable=AsyncMock) as mock_async_log, + patch.object( + mock_callback, "async_log_success_event", new_callable=AsyncMock + ) as mock_async_log, patch.object(mock_callback, "log_success_event") as mock_sync_log, ): await logging_obj.dispatch_success_handlers(result={"id": "pt-1"}) @@ -1393,14 +1434,20 @@ async def test_dispatch_failure_handlers_prefer_async_does_not_submit_sync_handl logging_obj.model_call_details["litellm_params"] = {} with ( - patch.object(logging_obj, "async_failure_handler", new_callable=AsyncMock) as mock_async, - patch.object(logging_obj, "failure_handler", new_callable=MagicMock) as mock_sync, + patch.object( + logging_obj, "async_failure_handler", new_callable=AsyncMock + ) as mock_async, + patch.object( + logging_obj, "failure_handler", new_callable=MagicMock + ) as mock_sync, patch.object( logging_obj, "_should_run_sync_failure_callbacks_for_async_calls", return_value=False, ), - patch("litellm.litellm_core_utils.litellm_logging.executor.submit") as mock_submit, + patch( + "litellm.litellm_core_utils.litellm_logging.executor.submit" + ) as mock_submit, ): await logging_obj.dispatch_failure_handlers( exception, @@ -1483,8 +1530,12 @@ async def test_dispatch_failure_handlers_submits_sync_handler_for_failure_only_c patch.object(litellm, "success_callback", []), patch.object(litellm, "failure_callback", [_sync_failure_callback]), patch.object(logging_obj, "async_failure_handler", new_callable=AsyncMock), - patch.object(logging_obj, "failure_handler", new_callable=MagicMock) as mock_sync, - patch("litellm.litellm_core_utils.litellm_logging.executor.submit") as mock_submit, + patch.object( + logging_obj, "failure_handler", new_callable=MagicMock + ) as mock_sync, + patch( + "litellm.litellm_core_utils.litellm_logging.executor.submit" + ) as mock_submit, ): await logging_obj.dispatch_failure_handlers( exception, @@ -1511,9 +1562,15 @@ async def test_dispatch_failure_handlers_sync_sdk_shortcut_runs_sync_handler_inl logging_obj.model_call_details["litellm_params"] = {} with ( - patch.object(logging_obj, "async_failure_handler", new_callable=AsyncMock) as mock_async, - patch.object(logging_obj, "failure_handler", new_callable=MagicMock) as mock_sync, - patch("litellm.litellm_core_utils.litellm_logging.executor.submit") as mock_submit, + patch.object( + logging_obj, "async_failure_handler", new_callable=AsyncMock + ) as mock_async, + patch.object( + logging_obj, "failure_handler", new_callable=MagicMock + ) as mock_sync, + patch( + "litellm.litellm_core_utils.litellm_logging.executor.submit" + ) as mock_submit, ): await logging_obj.dispatch_failure_handlers( exception, @@ -1560,10 +1617,14 @@ def test_success_handler_skips_guardrail_logging_hook_when_disabled(logging_obj) event_hook=GuardrailEventHooks.logging_only, ) guardrail.should_run_guardrail = MagicMock(return_value=False) - guardrail.logging_hook = MagicMock(return_value=(logging_obj.model_call_details, model_response)) + guardrail.logging_hook = MagicMock( + return_value=(logging_obj.model_call_details, model_response) + ) dummy_logger = DummyLogger() - dummy_logger.logging_hook = MagicMock(return_value=(logging_obj.model_call_details, model_response)) + dummy_logger.logging_hook = MagicMock( + return_value=(logging_obj.model_call_details, model_response) + ) with patch.object( logging_obj, @@ -1697,7 +1758,11 @@ def test_get_request_tags_from_metadata_and_litellm_metadata(): # Test case 2: Tags in litellm_metadata only tags = StandardLoggingPayloadSetup._get_request_tags( - litellm_params={"litellm_metadata": {"tags": ["litellm-metadata-tag-1", "litellm-metadata-tag-2"]}}, + litellm_params={ + "litellm_metadata": { + "tags": ["litellm-metadata-tag-1", "litellm-metadata-tag-2"] + } + }, proxy_server_request={}, ) assert "litellm-metadata-tag-1" in tags @@ -1802,9 +1867,15 @@ def test_get_request_tags_does_not_mutate_original_tags(): user_agent_count_2 = len([t for t in tags2 if t.startswith("User-Agent:")]) user_agent_count_3 = len([t for t in tags3 if t.startswith("User-Agent:")]) - assert user_agent_count_1 == 2, f"Expected 2 User-Agent tags, got {user_agent_count_1}" - assert user_agent_count_2 == 2, f"Expected 2 User-Agent tags, got {user_agent_count_2}" - assert user_agent_count_3 == 2, f"Expected 2 User-Agent tags, got {user_agent_count_3}" + assert ( + user_agent_count_1 == 2 + ), f"Expected 2 User-Agent tags, got {user_agent_count_1}" + assert ( + user_agent_count_2 == 2 + ), f"Expected 2 User-Agent tags, got {user_agent_count_2}" + assert ( + user_agent_count_3 == 2 + ), f"Expected 2 User-Agent tags, got {user_agent_count_3}" # Verify all returned lists are independent (different objects) assert tags1 is not tags2 @@ -1837,7 +1908,9 @@ def test_get_extra_header_tags(): # Test case 3: Extra headers configured but request has no headers dict litellm.extra_spend_tag_headers = ["x-custom", "x-tenant"] - result = StandardLoggingPayloadSetup._get_extra_header_tags(proxy_server_request={"headers": "not-a-dict"}) + result = StandardLoggingPayloadSetup._get_extra_header_tags( + proxy_server_request={"headers": "not-a-dict"} + ) assert result is None # Test case 4: Extra headers configured but none match request headers @@ -2138,7 +2211,9 @@ def test_get_masked_values(): "presidio_anonymizer_api_base": None, "vertex_credentials": "{sensitive_api_key}", } - masked_values = _get_masked_values(sensitive_object, unmasked_length=4, number_of_asterisks=4) + masked_values = _get_masked_values( + sensitive_object, unmasked_length=4, number_of_asterisks=4 + ) assert masked_values["presidio_anonymizer_api_base"] is None assert masked_values["vertex_credentials"] == "{s****y}" @@ -2163,7 +2238,9 @@ async def test_e2e_generate_cold_storage_object_key_successful(): patch("litellm.integrations.s3.get_s3_object_key") as mock_get_s3_key, ): # Mock the S3 object key generation to return a predictable result - mock_get_s3_key.return_value = "2025-01-15/time-10-30-45-123456_chatcmpl-test-12345.json" + mock_get_s3_key.return_value = ( + "2025-01-15/time-10-30-45-123456_chatcmpl-test-12345.json" + ) # Call the function result = StandardLoggingPayloadSetup._generate_cold_storage_object_key( @@ -2204,12 +2281,16 @@ async def test_e2e_generate_cold_storage_object_key_with_custom_logger_s3_path() with ( patch("litellm.cold_storage_custom_logger", "s3_v2"), - patch("litellm.logging_callback_manager.get_active_custom_logger_for_callback_name") as mock_get_logger, + patch( + "litellm.logging_callback_manager.get_active_custom_logger_for_callback_name" + ) as mock_get_logger, patch("litellm.integrations.s3.get_s3_object_key") as mock_get_s3_key, ): # Setup mocks mock_get_logger.return_value = mock_custom_logger - mock_get_s3_key.return_value = "storage/2025-01-15/time-10-30-45-123456_chatcmpl-test-12345.json" + mock_get_s3_key.return_value = ( + "storage/2025-01-15/time-10-30-45-123456_chatcmpl-test-12345.json" + ) # Call the function result = StandardLoggingPayloadSetup._generate_cold_storage_object_key( @@ -2228,7 +2309,9 @@ async def test_e2e_generate_cold_storage_object_key_with_custom_logger_s3_path() ) # Verify the result - assert result == "storage/2025-01-15/time-10-30-45-123456_chatcmpl-test-12345.json" + assert ( + result == "storage/2025-01-15/time-10-30-45-123456_chatcmpl-test-12345.json" + ) @pytest.mark.asyncio @@ -2251,12 +2334,16 @@ async def test_e2e_generate_cold_storage_object_key_with_logger_no_s3_path(): with ( patch("litellm.cold_storage_custom_logger", "s3_v2"), - patch("litellm.logging_callback_manager.get_active_custom_logger_for_callback_name") as mock_get_logger, + patch( + "litellm.logging_callback_manager.get_active_custom_logger_for_callback_name" + ) as mock_get_logger, patch("litellm.integrations.s3.get_s3_object_key") as mock_get_s3_key, ): # Setup mocks mock_get_logger.return_value = mock_custom_logger - mock_get_s3_key.return_value = "2025-01-15/time-10-30-45-123456_chatcmpl-test-12345.json" + mock_get_s3_key.return_value = ( + "2025-01-15/time-10-30-45-123456_chatcmpl-test-12345.json" + ) # Call the function result = StandardLoggingPayloadSetup._generate_cold_storage_object_key( @@ -2372,7 +2459,9 @@ def test_get_usage_as_dict(): assert result == {"prompt_tokens": 20, "completion_tokens": 30} # Test case 5: response_obj with no usage key returns empty - result = StandardLoggingPayloadSetup.get_usage_as_dict(response_obj={"id": "resp-1", "choices": []}) + result = StandardLoggingPayloadSetup.get_usage_as_dict( + response_obj={"id": "resp-1", "choices": []} + ) assert result == {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0} @@ -2385,20 +2474,26 @@ def test_append_system_prompt_messages(): # Test case 1: system in kwargs with existing messages kwargs = {"system": "You are a helpful assistant"} messages = [{"role": "user", "content": "Hello"}] - result = StandardLoggingPayloadSetup.append_system_prompt_messages(kwargs=kwargs, messages=messages) + result = StandardLoggingPayloadSetup.append_system_prompt_messages( + kwargs=kwargs, messages=messages + ) assert len(result) == 2 assert result[0] == {"role": "system", "content": "You are a helpful assistant"} assert result[1] == {"role": "user", "content": "Hello"} # Test case 2: system in kwargs with None messages kwargs = {"system": "You are a helpful assistant"} - result = StandardLoggingPayloadSetup.append_system_prompt_messages(kwargs=kwargs, messages=None) + result = StandardLoggingPayloadSetup.append_system_prompt_messages( + kwargs=kwargs, messages=None + ) assert len(result) == 1 assert result[0] == {"role": "system", "content": "You are a helpful assistant"} # Test case 3: system in kwargs with empty messages list kwargs = {"system": "You are a helpful assistant"} - result = StandardLoggingPayloadSetup.append_system_prompt_messages(kwargs=kwargs, messages=[]) + result = StandardLoggingPayloadSetup.append_system_prompt_messages( + kwargs=kwargs, messages=[] + ) assert len(result) == 1 assert result[0] == {"role": "system", "content": "You are a helpful assistant"} @@ -2408,18 +2503,24 @@ def test_append_system_prompt_messages(): {"role": "system", "content": "You are a helpful assistant"}, {"role": "user", "content": "Hello"}, ] - result = StandardLoggingPayloadSetup.append_system_prompt_messages(kwargs=kwargs, messages=messages) + result = StandardLoggingPayloadSetup.append_system_prompt_messages( + kwargs=kwargs, messages=messages + ) assert len(result) == 2 assert result[0] == {"role": "system", "content": "You are a helpful assistant"} # Test case 5: no system in kwargs returns messages unchanged kwargs = {} messages = [{"role": "user", "content": "Hello"}] - result = StandardLoggingPayloadSetup.append_system_prompt_messages(kwargs=kwargs, messages=messages) + result = StandardLoggingPayloadSetup.append_system_prompt_messages( + kwargs=kwargs, messages=messages + ) assert result == messages # Test case 6: None kwargs returns messages unchanged - result = StandardLoggingPayloadSetup.append_system_prompt_messages(kwargs=None, messages=messages) + result = StandardLoggingPayloadSetup.append_system_prompt_messages( + kwargs=None, messages=messages + ) assert result == messages @@ -2480,11 +2581,12 @@ async def test_async_success_handler_sets_standard_logging_object_for_pass_throu # Verify that standard_logging_object was set assert "standard_logging_object" in logging_obj.model_call_details, ( - "standard_logging_object should be set for pass-through endpoints even when complete_streaming_response is None" - ) - assert logging_obj.model_call_details["standard_logging_object"] is not None, ( - "standard_logging_object should not be None for pass-through endpoints" + "standard_logging_object should be set for pass-through endpoints " + "even when complete_streaming_response is None" ) + assert ( + logging_obj.model_call_details["standard_logging_object"] is not None + ), "standard_logging_object should not be None for pass-through endpoints" # Verify that async_complete_streaming_response was set to prevent re-processing # This is consistent with the existing code pattern for regular streaming @@ -2492,13 +2594,15 @@ async def test_async_success_handler_sets_standard_logging_object_for_pass_throu "async_complete_streaming_response should be set to prevent re-processing, " "consistent with the existing code pattern" ) - assert logging_obj.model_call_details["async_complete_streaming_response"] is result, ( - "async_complete_streaming_response should be set to the result" - ) + assert ( + logging_obj.model_call_details["async_complete_streaming_response"] is result + ), "async_complete_streaming_response should be set to the result" # Verify that response_cost is set to None (cost calculation not possible for pass-through) # This is consistent with the error handling in the non-pass-through code path - assert "response_cost" in logging_obj.model_call_details, "response_cost should be set for pass-through endpoints" + assert ( + "response_cost" in logging_obj.model_call_details + ), "response_cost should be set for pass-through endpoints" assert logging_obj.model_call_details["response_cost"] is None, ( "response_cost should be None for pass-through endpoints since " "StandardPassThroughResponseObject doesn't have standard usage info" @@ -2557,10 +2661,14 @@ async def test_async_success_handler_prevents_reprocessing_for_pass_through_endp # Verify first call set the values assert "standard_logging_object" in logging_obj.model_call_details assert "async_complete_streaming_response" in logging_obj.model_call_details - first_standard_logging_object = logging_obj.model_call_details["standard_logging_object"] + first_standard_logging_object = logging_obj.model_call_details[ + "standard_logging_object" + ] # Second call - should return early due to async_complete_streaming_response guard - with patch.object(logging_obj, "get_combined_callback_list", return_value=[]) as mock_callbacks: + with patch.object( + logging_obj, "get_combined_callback_list", return_value=[] + ) as mock_callbacks: await logging_obj.async_success_handler( result=result, start_time=start_time, @@ -2571,9 +2679,10 @@ async def test_async_success_handler_prevents_reprocessing_for_pass_through_endp mock_callbacks.assert_not_called() # Verify standard_logging_object wasn't modified by second call - assert logging_obj.model_call_details["standard_logging_object"] is first_standard_logging_object, ( - "standard_logging_object should not be modified on re-processing" - ) + assert ( + logging_obj.model_call_details["standard_logging_object"] + is first_standard_logging_object + ), "standard_logging_object should not be modified on re-processing" @pytest.mark.asyncio @@ -2612,7 +2721,9 @@ async def test_async_success_handler_sets_standard_logging_object_for_streaming_ } # Create a pass-through response object (simulating unparseable streaming response) - result = StandardPassThroughResponseObject(response='data: {"chunk": 1}\ndata: {"chunk": 2}\ndata: [DONE]') + result = StandardPassThroughResponseObject( + response='data: {"chunk": 1}\ndata: {"chunk": 2}\ndata: [DONE]' + ) start_time = datetime.now() end_time = datetime.now() @@ -2632,9 +2743,9 @@ async def test_async_success_handler_sets_standard_logging_object_for_streaming_ "standard_logging_object should be set for streaming pass-through endpoints " "even when the response cannot be parsed into a ModelResponse" ) - assert logging_obj.model_call_details["standard_logging_object"] is not None, ( - "standard_logging_object should not be None for streaming pass-through endpoints" - ) + assert ( + logging_obj.model_call_details["standard_logging_object"] is not None + ), "standard_logging_object should not be None for streaming pass-through endpoints" def test_get_error_information_error_code_priority(): @@ -2676,22 +2787,30 @@ def test_get_error_information_error_code_priority(): self.message = message super().__init__(message) - both_exception = BothAttributesException(code="400", status_code=500, message="Bad Request") + both_exception = BothAttributesException( + code="400", status_code=500, message="Bad Request" + ) result = StandardLoggingPayloadSetup.get_error_information(both_exception) assert result["error_code"] == "400" # Should prefer 'code' over 'status_code' # Test case 4: Exception with 'code' as empty string - should fall back to 'status_code' - empty_code_exception = BothAttributesException(code="", status_code=404, message="Not Found") + empty_code_exception = BothAttributesException( + code="", status_code=404, message="Not Found" + ) result = StandardLoggingPayloadSetup.get_error_information(empty_code_exception) assert result["error_code"] == "404" # Should fall back to status_code # Test case 5: Exception with 'code' as "None" string - should fall back to 'status_code' - none_string_exception = BothAttributesException(code="None", status_code=503, message="Service Unavailable") + none_string_exception = BothAttributesException( + code="None", status_code=503, message="Service Unavailable" + ) result = StandardLoggingPayloadSetup.get_error_information(none_string_exception) assert result["error_code"] == "503" # Should fall back to status_code # Test case 6: Exception with 'code' as None - should fall back to 'status_code' - none_code_exception = BothAttributesException(code=None, status_code=401, message="Unauthorized") + none_code_exception = BothAttributesException( + code=None, status_code=401, message="Unauthorized" + ) result = StandardLoggingPayloadSetup.get_error_information(none_code_exception) assert result["error_code"] == "401" # Should fall back to status_code @@ -2740,7 +2859,9 @@ def test_get_error_information_prefers_message_attribute_over_str(): ) result = StandardLoggingPayloadSetup.get_error_information(exc) - assert result["error_message"] == msg, f"expected message from .message attribute, got {result['error_message']!r}" + assert ( + result["error_message"] == msg + ), f"expected message from .message attribute, got {result['error_message']!r}" assert result["error_code"] == "401" assert result["error_class"] == "ProxyExceptionLike" @@ -2815,7 +2936,8 @@ def test_get_error_information_preserves_explicit_empty_message(): exc = ProxyExceptionLike(message="", code=500) result = StandardLoggingPayloadSetup.get_error_information(exc) assert result["error_message"] == "", ( - f"explicit empty .message must survive verbatim; got {result['error_message']!r}" + "explicit empty .message must survive verbatim; got " + f"{result['error_message']!r}" ) @@ -3078,7 +3200,9 @@ def test_process_hidden_params_recalculates_cost_after_failure_handler_zero(): choices=[{"message": {"role": "assistant", "content": "ok"}}], usage=Usage(prompt_tokens=9698, completion_tokens=30, total_tokens=9728), ) - logging_obj._process_hidden_params_and_response_cost(result, datetime.now(), datetime.now()) + logging_obj._process_hidden_params_and_response_cost( + result, datetime.now(), datetime.now() + ) cost = logging_obj.model_call_details.get("response_cost") assert cost is not None and cost > 0 @@ -3102,7 +3226,9 @@ def test_process_hidden_params_preserves_zero_cost_in_hidden_params(): litellm_call_id="test-hidden-zero-cost", function_id="test-hidden-zero-cost", ) - logging_obj.model_call_details["litellm_params"] = {"model": "gemini-2.5-flash-lite"} + logging_obj.model_call_details["litellm_params"] = { + "model": "gemini-2.5-flash-lite" + } logging_obj.optional_params = {} result = ModelResponse( @@ -3112,7 +3238,9 @@ def test_process_hidden_params_preserves_zero_cost_in_hidden_params(): ) result._hidden_params = {"response_cost": 0.0} - logging_obj._process_hidden_params_and_response_cost(result, datetime.now(), datetime.now()) + logging_obj._process_hidden_params_and_response_cost( + result, datetime.now(), datetime.now() + ) assert logging_obj.model_call_details.get("response_cost") == 0.0 slo = logging_obj.model_call_details.get("standard_logging_object") or {} @@ -3161,7 +3289,9 @@ def test_process_hidden_params_uses_hidden_params_cost_after_failure_handler_zer ) result._hidden_params = {"response_cost": passthrough_cost} - logging_obj._process_hidden_params_and_response_cost(result, datetime.now(), datetime.now()) + logging_obj._process_hidden_params_and_response_cost( + result, datetime.now(), datetime.now() + ) assert logging_obj.model_call_details.get("response_cost") == passthrough_cost slo = logging_obj.model_call_details.get("standard_logging_object") or {} @@ -3218,7 +3348,9 @@ def test_function_setup_litellm_metadata_populates_metadata(): assert litellm_metadata.get("user_api_key_hash") == test_api_key_hash # metadata should be a COPY, not an alias — mutating one must not affect the other - assert metadata is not litellm_metadata, "litellm_params['metadata'] should be a copy, not the same object" + assert ( + metadata is not litellm_metadata + ), "litellm_params['metadata'] should be a copy, not the same object" def test_function_setup_litellm_metadata_guardrail_writes_visible_after_setup(): @@ -3263,9 +3395,9 @@ def test_function_setup_litellm_metadata_guardrail_writes_visible_after_setup(): litellm_params = logging_obj.model_call_details.get("litellm_params", {}) litellm_metadata = litellm_params.get("litellm_metadata") assert litellm_metadata is not None - assert litellm_metadata.get("standard_logging_guardrail_information") == [guardrail_entry], ( - "guardrail writes after function_setup must be visible to the logging object" - ) + assert litellm_metadata.get("standard_logging_guardrail_information") == [ + guardrail_entry + ], "guardrail writes after function_setup must be visible to the logging object" assert litellm_metadata.get("applied_guardrails") == ["pam-ethical-request"] merged = StandardLoggingPayloadSetup.merge_litellm_metadata(litellm_params) @@ -3434,7 +3566,9 @@ def test_failure_handler_skips_sync_callbacks_for_pass_through_requests(logging_ @pytest.mark.parametrize("call_type", ["completion", "acompletion"]) -def test_failure_handler_runs_sync_callbacks_for_non_pass_through_requests(logging_obj, call_type): +def test_failure_handler_runs_sync_callbacks_for_non_pass_through_requests( + logging_obj, call_type +): """Ensure sync failure callbacks still fire for normal (non-pass-through) requests.""" from litellm.integrations.custom_logger import CustomLogger @@ -3595,7 +3729,9 @@ def test_standard_logging_hidden_params_backfills_response_cost_without_mutating ) response._hidden_params = {"response_cost": None, "model_id": "mid-test"} - payload = logging_obj._build_standard_logging_payload(response, datetime.now(), datetime.now()) + payload = logging_obj._build_standard_logging_payload( + response, datetime.now(), datetime.now() + ) assert payload is not None assert payload["hidden_params"]["response_cost"] == 0.002 @@ -3649,7 +3785,10 @@ def test_merge_hidden_params_from_response_into_metadata_no_op_when_empty(): _hidden_params = {} logging_obj._merge_hidden_params_from_response_into_metadata(_NoHp()) - assert "hidden_params" not in logging_obj.model_call_details["litellm_params"]["metadata"] + assert ( + "hidden_params" + not in logging_obj.model_call_details["litellm_params"]["metadata"] + ) # ── StandardLoggingPayloadSetup.get_additional_headers ─────────────────────── @@ -3738,7 +3877,9 @@ def _model_router_response(selected_model: str, stamp: bool): from litellm.types.utils import ModelResponse response = ModelResponse(model=selected_model) - response._hidden_params = {AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY: selected_model} if stamp else {} + response._hidden_params = ( + {AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY: selected_model} if stamp else {} + ) return response @@ -3762,7 +3903,9 @@ def test_standard_logging_payload_uses_stamped_model_router_model(logging_obj): "messages": [], "litellm_params": {"metadata": {}}, }, - init_response_obj=_model_router_response("azure_ai/grok-4-1-fast-reasoning", stamp=True), + init_response_obj=_model_router_response( + "azure_ai/grok-4-1-fast-reasoning", stamp=True + ), start_time=now, end_time=now, logging_obj=logging_obj, @@ -3773,7 +3916,9 @@ def test_standard_logging_payload_uses_stamped_model_router_model(logging_obj): assert payload["model"] == "azure_ai/grok-4-1-fast-reasoning" -def test_standard_logging_payload_keeps_requested_model_without_router_stamp(logging_obj): +def test_standard_logging_payload_keeps_requested_model_without_router_stamp( + logging_obj, +): """ Control for the test above: an ordinary azure_ai deployment is unaffected, so the stamp is what redirects attribution rather than the response model winning unconditionally. @@ -3792,7 +3937,9 @@ def test_standard_logging_payload_keeps_requested_model_without_router_stamp(log "messages": [], "litellm_params": {"metadata": {}}, }, - init_response_obj=_model_router_response("azure_ai/grok-4-1-fast-reasoning", stamp=False), + init_response_obj=_model_router_response( + "azure_ai/grok-4-1-fast-reasoning", stamp=False + ), start_time=now, end_time=now, logging_obj=logging_obj, @@ -3838,7 +3985,9 @@ def test_success_handler_computes_cost_for_dict_response(): "_build_standard_logging_payload", return_value={"response_cost": expected_cost}, ), - patch("litellm.litellm_core_utils.litellm_logging.emit_standard_logging_payload"), + patch( + "litellm.litellm_core_utils.litellm_logging.emit_standard_logging_payload" + ), patch.object( logging_obj, "_is_recognized_call_type_for_logging", @@ -3875,7 +4024,9 @@ def test_success_handler_preserves_precomputed_cost_for_dict_response(): "_build_standard_logging_payload", return_value={"response_cost": precomputed_cost}, ), - patch("litellm.litellm_core_utils.litellm_logging.emit_standard_logging_payload"), + patch( + "litellm.litellm_core_utils.litellm_logging.emit_standard_logging_payload" + ), patch.object( logging_obj, "_is_recognized_call_type_for_logging", @@ -3914,7 +4065,9 @@ def test_success_handler_unified_helper_runs_for_typed_results(): "_build_standard_logging_payload", return_value={"response_cost": expected_cost}, ), - patch("litellm.litellm_core_utils.litellm_logging.emit_standard_logging_payload"), + patch( + "litellm.litellm_core_utils.litellm_logging.emit_standard_logging_payload" + ), patch.object( logging_obj, "_is_recognized_call_type_for_logging", @@ -3969,7 +4122,9 @@ class TestFirstApiCallStartTimeSetOnce: assert first == obj.model_call_details["api_call_start_time"] # Set on the logging object only — user metadata untouched. assert user_meta == {} - assert "first_api_call_start_time" not in obj.model_call_details["litellm_params"] + assert ( + "first_api_call_start_time" not in obj.model_call_details["litellm_params"] + ) time.sleep(0.002) # ensure a distinct retry timestamp obj.pre_call(input="hi", api_key="sk-test") @@ -3986,16 +4141,18 @@ def test_get_error_information_for_logging_payload_ignores_spoofed_disconnect_wi baseline = StandardLoggingPayloadSetup.get_error_information( original_exception=ValueError("provider failure"), ) - error_information, error_str = StandardLoggingPayloadSetup.get_error_information_for_logging_payload( - metadata={ - "error_information": { - "error_code": "499", - "error_message": "Client disconnected the request", - "error_class": "ClientDisconnected", - } - }, - original_exception=ValueError("provider failure"), - error_str="provider failure", + error_information, error_str = ( + StandardLoggingPayloadSetup.get_error_information_for_logging_payload( + metadata={ + "error_information": { + "error_code": "499", + "error_message": "Client disconnected the request", + "error_class": "ClientDisconnected", + } + }, + original_exception=ValueError("provider failure"), + error_str="provider failure", + ) ) assert error_information == baseline assert error_str == "provider failure" @@ -4009,18 +4166,22 @@ def test_get_error_information_for_logging_payload_client_disconnect(): "error_message": "Client disconnected the request", "error_class": "ClientDisconnected", } - error_information, error_str = StandardLoggingPayloadSetup.get_error_information_for_logging_payload( - metadata={"client_disconnected": True, "error_information": custom_error}, - original_exception=None, - error_str=None, + error_information, error_str = ( + StandardLoggingPayloadSetup.get_error_information_for_logging_payload( + metadata={"client_disconnected": True, "error_information": custom_error}, + original_exception=None, + error_str=None, + ) ) assert error_information == custom_error assert error_str == "Client disconnected the request" - error_information, error_str = StandardLoggingPayloadSetup.get_error_information_for_logging_payload( - metadata={"client_disconnected": True}, - original_exception=None, - error_str="existing error", + error_information, error_str = ( + StandardLoggingPayloadSetup.get_error_information_for_logging_payload( + metadata={"client_disconnected": True}, + original_exception=None, + error_str="existing error", + ) ) assert error_information["error_code"] == "499" assert error_str == "existing error" @@ -4028,10 +4189,12 @@ def test_get_error_information_for_logging_payload_client_disconnect(): baseline = StandardLoggingPayloadSetup.get_error_information( original_exception=None, ) - error_information, error_str = StandardLoggingPayloadSetup.get_error_information_for_logging_payload( - metadata={}, - original_exception=None, - error_str=None, + error_information, error_str = ( + StandardLoggingPayloadSetup.get_error_information_for_logging_payload( + metadata={}, + original_exception=None, + error_str=None, + ) ) assert error_information == baseline assert error_str is None @@ -4066,7 +4229,9 @@ def test_get_error_information_prefers_message_attribute_over_empty_str(): def __str__(self): return "" - info = StandardLoggingPayloadSetup.get_error_information(original_exception=_SilentExc()) + info = StandardLoggingPayloadSetup.get_error_information( + original_exception=_SilentExc() + ) assert info["error_message"] == "real failure detail" assert info["error_code"] == "401" @@ -4097,7 +4262,9 @@ def _responses_api_response_with_text(text="hello world"): type="message", role="assistant", status="completed", - content=[ResponseOutputText(annotations=[], text=text, type="output_text")], + content=[ + ResponseOutputText(annotations=[], text=text, type="output_text") + ], ) ], usage=ResponseAPIUsage(input_tokens=11, output_tokens=7, total_tokens=18), @@ -4112,7 +4279,9 @@ def _responses_api_response_with_text(text="hello world"): ("ResponseFailedEvent", "response.failed"), ], ) -def test_handle_anthropic_messages_response_logging_translates_terminal_responses_api_event(event_cls, event_type): +def test_handle_anthropic_messages_response_logging_translates_terminal_responses_api_event( + event_cls, event_type +): """Regression for #28595 / #28943. When anthropic_messages routes to the OpenAI Responses backend and stream=True, success_handler receives a terminal Responses API event. The handler must translate it to a ModelResponse whose choices carry @@ -4151,7 +4320,10 @@ def test_handle_anthropic_messages_response_logging_passes_model_response_throug """Anthropic-native path already yields a ModelResponse; it must be returned unchanged.""" logging_obj = _anthropic_messages_logging_obj() model_response = ModelResponse() - assert logging_obj._handle_anthropic_messages_response_logging(result=model_response) is model_response + assert ( + logging_obj._handle_anthropic_messages_response_logging(result=model_response) + is model_response + ) def test_handle_anthropic_messages_response_logging_degrades_on_unparseable_responses_payload(): @@ -4447,7 +4619,9 @@ def test_non_image_response_has_no_output_image_count(logging_obj): def test_zero_token_video_usage_preserves_duration_seconds(logging_obj): """Video usage bills by duration; the payload must keep duration_seconds even with zero tokens.""" - payload = _build_payload_for_media_response(logging_obj, {"id": "video-1", "usage": {"duration_seconds": 4.0}}) + payload = _build_payload_for_media_response( + logging_obj, {"id": "video-1", "usage": {"duration_seconds": 4.0}} + ) assert payload is not None assert payload["metadata"]["usage_object"]["duration_seconds"] == 4.0 diff --git a/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py b/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py index f6b131311f2..11a727c9635 100644 --- a/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py +++ b/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py @@ -115,7 +115,9 @@ def test_azure_ai_grok_stop_parameter_handling(): # Test supported parameters for Grok models for model in ("grok-4-fast", "grok-4.3"): grok_params = config.get_supported_openai_params(model) - assert "stop" not in grok_params, "Grok models should not support stop parameter" + assert ( + "stop" not in grok_params + ), "Grok models should not support stop parameter" # Test supported parameters for non-Grok models gpt_params = config.get_supported_openai_params("gpt-4") @@ -194,7 +196,8 @@ def test_azure_model_router_response_shows_actual_model(): # Verify that the response contains the actual model used, not the router model assert result.model == "azure_ai/gpt-5-nano-2025-08-07", ( - f"Expected model to be 'azure_ai/gpt-5-nano-2025-08-07' (actual model used), but got '{result.model}'" + f"Expected model to be 'azure_ai/gpt-5-nano-2025-08-07' (actual model used), " + f"but got '{result.model}'" ) @@ -252,11 +255,19 @@ def test_azure_model_router_stamps_selected_model_on_hidden_params(): ) assert result._hidden_params[AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY] == result.model - assert result._hidden_params[AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY] == "azure_ai/grok-4-1-fast-reasoning" - assert AzureFoundryModelInfo.get_model_router_selected_model(result._hidden_params) == ( - "azure_ai/grok-4-1-fast-reasoning" + assert ( + result._hidden_params[AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY] + == "azure_ai/grok-4-1-fast-reasoning" + ) + assert AzureFoundryModelInfo.get_model_router_selected_model( + result._hidden_params + ) == ("azure_ai/grok-4-1-fast-reasoning") + assert ( + AzureFoundryModelInfo.is_model_router_call( + model="smart-pick", hidden_params=result._hidden_params + ) + is True ) - assert AzureFoundryModelInfo.is_model_router_call(model="smart-pick", hidden_params=result._hidden_params) is True def test_azure_model_router_stamp_does_not_leak_across_responses(): @@ -264,7 +275,9 @@ def test_azure_model_router_stamp_does_not_leak_across_responses(): ModelResponse declares _hidden_params as a class-level dict, so the stamp has to be written as a fresh dict. Mutating in place would bleed the selected model into unrelated responses. """ - from litellm.llms.azure_ai.common_utils import AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY + from litellm.llms.azure_ai.common_utils import ( + AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY, + ) from litellm.types.utils import ModelResponse untouched = ModelResponse() @@ -292,10 +305,14 @@ def test_drop_tool_level_extra_fields_strips_copilot_mcp_server_name(): mock_response.text = error_text mock_response.json.return_value = json.loads(error_text) mock_response.status_code = 400 - e = httpx.HTTPStatusError(message="400", request=MagicMock(), response=mock_response) + e = httpx.HTTPStatusError( + message="400", request=MagicMock(), response=mock_response + ) assert config._error_has_tool_level_extra_fields(error_text) is True - assert config.should_retry_llm_api_inside_llm_translation_on_http_error(e, {}) is True + assert ( + config.should_retry_llm_api_inside_llm_translation_on_http_error(e, {}) is True + ) request_data = { "model": "FW-Kimi-K2.6", @@ -418,7 +435,9 @@ def test_azure_ai_stripping_does_not_mutate_caller_messages(): { "role": "assistant", "content": "I can help.", - "thinking_blocks": [{"type": "thinking", "thinking": "Reading the file.", "signature": "sig"}], + "thinking_blocks": [ + {"type": "thinking", "thinking": "Reading the file.", "signature": "sig"} + ], "provider_specific_fields": {"thought_signature": "sig-top"}, "tool_calls": [ { diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 51928ef67e5..c9b3e5faf8b 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -129,12 +129,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", @@ -174,10 +178,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", @@ -219,11 +227,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", @@ -242,7 +254,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, @@ -2280,7 +2294,9 @@ class TestOverrideOpenAIResponseModel: assert response_obj.model == actual_model_used assert response_obj.model != requested_model - def test_override_model_preserves_model_router_model_for_alias_without_router_in_name(self): + def test_override_model_preserves_model_router_model_for_alias_without_router_in_name( + self, + ): """ The client sends a model group alias, which carries no model_router/ prefix, so the name check alone only fires when the operator happened to put "model-router" in the @@ -2898,7 +2914,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() @@ -2929,7 +2947,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() @@ -3000,7 +3020,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() @@ -3196,7 +3218,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, ) @@ -3212,7 +3236,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() @@ -3228,7 +3254,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( @@ -3280,7 +3308,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" @@ -3476,7 +3506,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 @@ -3505,7 +3537,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( @@ -3523,7 +3557,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 @@ -3548,7 +3584,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", @@ -3609,7 +3647,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( @@ -3660,7 +3700,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 @@ -3689,7 +3731,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) @@ -3726,13 +3770,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" ) @@ -3746,7 +3796,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"] @@ -3771,12 +3823,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): @@ -3792,11 +3854,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): @@ -3807,7 +3873,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, ) @@ -3839,7 +3907,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, ) @@ -3869,7 +3939,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, ) @@ -3884,7 +3956,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) @@ -3895,7 +3969,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( @@ -3914,8 +3990,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: @@ -3942,17 +4016,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() @@ -3971,7 +4051,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() @@ -3986,7 +4068,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() @@ -3995,7 +4079,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) @@ -4003,7 +4089,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() @@ -4082,7 +4170,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, ) @@ -4148,9 +4238,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, @@ -4200,9 +4288,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, @@ -4240,9 +4326,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, @@ -4283,9 +4367,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, @@ -4397,9 +4479,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, @@ -4430,7 +4510,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 @@ -4481,17 +4563,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) @@ -4500,27 +4579,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) @@ -4536,23 +4615,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) @@ -4574,17 +4649,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) @@ -4603,17 +4675,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) @@ -4954,6 +5023,7 @@ class TestResponseCostHeaderForTypedDictResponses: 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 @@ -5105,7 +5175,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, @@ -5136,7 +5208,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( @@ -5262,7 +5336,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}}, @@ -5270,7 +5346,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}, @@ -5301,7 +5380,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", @@ -5331,7 +5412,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, @@ -5692,12 +5775,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 @@ -6467,7 +6554,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 @@ -6496,7 +6585,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 @@ -6791,7 +6882,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.""" @@ -6937,7 +7030,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"] @@ -6954,7 +7049,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()) @@ -7008,7 +7105,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.""" @@ -7054,7 +7153,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()) @@ -7093,7 +7194,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()) @@ -7110,7 +7213,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())