diff --git a/tests/test_litellm/litellm_core_utils/llm_response_utils/test_response_metadata.py b/tests/test_litellm/litellm_core_utils/llm_response_utils/test_response_metadata.py new file mode 100644 index 00000000000..e6147d16d19 --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/llm_response_utils/test_response_metadata.py @@ -0,0 +1,291 @@ +""" +Tests for litellm.litellm_core_utils.llm_response_utils.response_metadata + +Covers the callback_duration_ms timing metric that flows from the Logging object +through _hidden_params to the x-litellm-callback-duration-ms response header. +""" + +import datetime +from unittest.mock import MagicMock + +from litellm.litellm_core_utils.llm_response_utils.response_metadata import ( + ResponseMetadata, + update_response_metadata, +) +from litellm.types.utils import ModelResponse + + +class TestCallbackDurationMs: + """Tests for the callback_duration_ms metric in ResponseMetadata.""" + + def _make_logging_obj(self, callback_duration_ms=None, llm_api_duration_ms=None): + """Build a minimal mock logging object.""" + logging_obj = MagicMock() + logging_obj.model_call_details = {} + if llm_api_duration_ms is not None: + logging_obj.model_call_details["llm_api_duration_ms"] = llm_api_duration_ms + logging_obj.caching_details = None + if callback_duration_ms is not None: + logging_obj.callback_duration_ms = callback_duration_ms + else: + # Simulate a Logging object that has no callback_duration_ms + del logging_obj.callback_duration_ms + return logging_obj + + def test_callback_duration_ms_set_in_hidden_params(self): + """When logging_obj has callback_duration_ms, it should appear in _hidden_params.""" + result = ModelResponse() + logging_obj = self._make_logging_obj(callback_duration_ms=12.3456) + + metadata = ResponseMetadata(result) + start = datetime.datetime(2025, 1, 1, 0, 0, 0) + end = datetime.datetime(2025, 1, 1, 0, 0, 1) + metadata.set_timing_metrics(start, end, logging_obj) + metadata.apply() + + hidden = result._hidden_params + assert hidden.get("callback_duration_ms") == 12.3456 + + def test_callback_duration_ms_absent_when_not_on_logging_obj(self): + """When logging_obj lacks callback_duration_ms, hidden_params should not have it.""" + result = ModelResponse() + logging_obj = self._make_logging_obj(callback_duration_ms=None) + + metadata = ResponseMetadata(result) + start = datetime.datetime(2025, 1, 1, 0, 0, 0) + end = datetime.datetime(2025, 1, 1, 0, 0, 1) + metadata.set_timing_metrics(start, end, logging_obj) + metadata.apply() + + hidden = result._hidden_params + assert hidden.get("callback_duration_ms") is None + + def test_update_response_metadata_includes_callback_duration(self): + """End-to-end: update_response_metadata should propagate callback_duration_ms.""" + result = ModelResponse() + logging_obj = self._make_logging_obj( + callback_duration_ms=5.5, llm_api_duration_ms=800.0 + ) + logging_obj._response_cost_calculator = MagicMock(return_value=0.001) + logging_obj.litellm_call_id = "test-call-id" + + start = datetime.datetime(2025, 1, 1, 0, 0, 0) + end = datetime.datetime(2025, 1, 1, 0, 0, 1) + + update_response_metadata( + result=result, + logging_obj=logging_obj, + model="gpt-4", + kwargs={}, + start_time=start, + end_time=end, + ) + + hidden = result._hidden_params + assert hidden.get("callback_duration_ms") == 5.5 + # overhead should also be set + assert hidden.get("litellm_overhead_time_ms") is not None + + +class TestCallbackDurationInCustomHeaders: + """Test that callback_duration_ms flows into get_custom_headers.""" + + def test_header_present_when_callback_duration_in_hidden_params(self): + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.common_request_processing import ( + ProxyBaseLLMRequestProcessing, + ) + + user_api_key_dict = UserAPIKeyAuth(api_key="sk-test") + hidden_params = { + "_response_ms": 1000.0, + "litellm_overhead_time_ms": 50.0, + "callback_duration_ms": 7.25, + } + + headers = ProxyBaseLLMRequestProcessing.get_custom_headers( + user_api_key_dict=user_api_key_dict, + hidden_params=hidden_params, + ) + + assert "x-litellm-callback-duration-ms" in headers + assert headers["x-litellm-callback-duration-ms"] == "7.25" + + def test_header_absent_when_no_callback_duration(self): + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.common_request_processing import ( + ProxyBaseLLMRequestProcessing, + ) + + user_api_key_dict = UserAPIKeyAuth(api_key="sk-test") + hidden_params = { + "_response_ms": 1000.0, + } + + headers = ProxyBaseLLMRequestProcessing.get_custom_headers( + user_api_key_dict=user_api_key_dict, + hidden_params=hidden_params, + ) + + # Should be excluded because value is "None" which is in exclude_values + assert "x-litellm-callback-duration-ms" not in headers + + +class TestDetailedTiming: + """Tests for detailed per-phase timing headers behind LITELLM_DETAILED_TIMING.""" + + def _make_logging_obj( + self, + llm_api_duration_ms=500.0, + message_copy_duration_ms=2.5, + api_call_start_time=None, + ): + logging_obj = MagicMock() + logging_obj.model_call_details = { + "llm_api_duration_ms": llm_api_duration_ms, + } + if api_call_start_time is not None: + logging_obj.model_call_details["api_call_start_time"] = api_call_start_time + logging_obj.caching_details = None + logging_obj.callback_duration_ms = 1.0 + logging_obj.message_copy_duration_ms = message_copy_duration_ms + return logging_obj + + def test_detailed_timing_headers_present_when_enabled(self, monkeypatch): + """When LITELLM_DETAILED_TIMING is true, detailed timing keys appear in hidden_params.""" + import litellm.litellm_core_utils.llm_response_utils.response_metadata as mod + + monkeypatch.setattr(mod, "LITELLM_DETAILED_TIMING", True) + + result = ModelResponse() + start = datetime.datetime(2025, 1, 1, 0, 0, 0) + api_call_start = datetime.datetime(2025, 1, 1, 0, 0, 0, 20000) # +20ms + end = datetime.datetime(2025, 1, 1, 0, 0, 0, 530000) # +530ms total + + logging_obj = self._make_logging_obj( + llm_api_duration_ms=500.0, + message_copy_duration_ms=2.5, + api_call_start_time=api_call_start, + ) + + metadata = ResponseMetadata(result) + metadata.set_timing_metrics(start, end, logging_obj) + metadata.apply() + + hidden = result._hidden_params + assert hidden.get("timing_llm_api_ms") == 500.0 + assert hidden.get("timing_message_copy_ms") == 2.5 + assert hidden.get("timing_pre_processing_ms") == 20.0 + assert hidden.get("timing_post_processing_ms") == 10.0 # 530 - 20 - 500 + + def test_detailed_timing_absent_when_disabled(self, monkeypatch): + """When LITELLM_DETAILED_TIMING is false, no detailed timing keys.""" + import litellm.litellm_core_utils.llm_response_utils.response_metadata as mod + + monkeypatch.setattr(mod, "LITELLM_DETAILED_TIMING", False) + + result = ModelResponse() + start = datetime.datetime(2025, 1, 1, 0, 0, 0) + end = datetime.datetime(2025, 1, 1, 0, 0, 1) + logging_obj = self._make_logging_obj() + + metadata = ResponseMetadata(result) + metadata.set_timing_metrics(start, end, logging_obj) + metadata.apply() + + hidden = result._hidden_params + assert hidden.get("timing_llm_api_ms") is None + assert hidden.get("timing_pre_processing_ms") is None + + def test_detailed_timing_headers_in_custom_headers(self, monkeypatch): + """When LITELLM_DETAILED_TIMING is true, headers flow to get_custom_headers.""" + import litellm.proxy.common_request_processing as crp_mod + + monkeypatch.setattr(crp_mod, "LITELLM_DETAILED_TIMING", True) + + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.common_request_processing import ( + ProxyBaseLLMRequestProcessing, + ) + + user_api_key_dict = UserAPIKeyAuth(api_key="sk-test") + hidden_params = { + "_response_ms": 530.0, + "timing_llm_api_ms": 500.0, + "timing_pre_processing_ms": 20.0, + "timing_post_processing_ms": 10.0, + "timing_message_copy_ms": 2.5, + } + + headers = ProxyBaseLLMRequestProcessing.get_custom_headers( + user_api_key_dict=user_api_key_dict, + hidden_params=hidden_params, + ) + + assert headers["x-litellm-timing-llm-api-ms"] == "500.0" + assert headers["x-litellm-timing-pre-processing-ms"] == "20.0" + assert headers["x-litellm-timing-post-processing-ms"] == "10.0" + assert headers["x-litellm-timing-message-copy-ms"] == "2.5" + + def test_detailed_timing_headers_absent_when_disabled(self, monkeypatch): + """When LITELLM_DETAILED_TIMING is false, no timing headers emitted.""" + import litellm.proxy.common_request_processing as crp_mod + + monkeypatch.setattr(crp_mod, "LITELLM_DETAILED_TIMING", False) + + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.common_request_processing import ( + ProxyBaseLLMRequestProcessing, + ) + + user_api_key_dict = UserAPIKeyAuth(api_key="sk-test") + hidden_params = { + "_response_ms": 530.0, + "timing_llm_api_ms": 500.0, + } + + headers = ProxyBaseLLMRequestProcessing.get_custom_headers( + user_api_key_dict=user_api_key_dict, + hidden_params=hidden_params, + ) + + assert "x-litellm-timing-llm-api-ms" not in headers + assert "x-litellm-timing-pre-processing-ms" not in headers + + +class TestLoggingInitCallbackDuration: + """Test that Logging.__init__ tracks deep copy time in callback_duration_ms.""" + + def test_logging_init_sets_callback_duration_ms(self): + from litellm.litellm_core_utils.litellm_logging import Logging + + obj = Logging( + model="gpt-4", + messages=[{"role": "user", "content": "hello " * 100}], + stream=False, + call_type="acompletion", + start_time=datetime.datetime.now(), + litellm_call_id="test-123", + function_id="func-123", + ) + + # callback_duration_ms should be set and non-negative + assert hasattr(obj, "callback_duration_ms") + assert obj.callback_duration_ms >= 0 + + def test_logging_init_callback_duration_zero_for_none_messages(self): + from litellm.litellm_core_utils.litellm_logging import Logging + + obj = Logging( + model="gpt-4", + messages=None, + stream=False, + call_type="acompletion", + start_time=datetime.datetime.now(), + litellm_call_id="test-456", + function_id="func-456", + ) + + # Should still be set (deep copy of None is essentially a no-op) + assert hasattr(obj, "callback_duration_ms") + assert obj.callback_duration_ms >= 0 diff --git a/tests/test_litellm/litellm_core_utils/test_logging_utils.py b/tests/test_litellm/litellm_core_utils/test_logging_utils.py new file mode 100644 index 00000000000..59061b6c68f --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/test_logging_utils.py @@ -0,0 +1,156 @@ +""" +Tests for litellm.litellm_core_utils.logging_utils — base64 truncation helpers. +""" + +import pytest + +from litellm.litellm_core_utils.logging_utils import ( + _format_base64_size, + _truncate_base64_in_string, + truncate_base64_in_messages, +) + +# --------------------------------------------------------------------------- +# _format_base64_size +# --------------------------------------------------------------------------- + + +class TestFormatBase64Size: + def test_bytes_range(self): + assert _format_base64_size(4) == "3B" + + def test_kb_range(self): + # 2000 base64 chars ~ 1500 bytes ~ 1.5KB + assert "KB" in _format_base64_size(2000) + + def test_mb_range(self): + # 2_000_000 base64 chars ~ 1.5MB + result = _format_base64_size(2_000_000) + assert "MB" in result + + +# --------------------------------------------------------------------------- +# _truncate_base64_in_string +# --------------------------------------------------------------------------- + + +class TestTruncateBase64InString: + def test_short_data_uri_not_truncated(self): + uri = "data:image/png;base64,AAAA" + assert _truncate_base64_in_string(uri) == uri + + def test_long_data_uri_truncated(self): + payload = "A" * 200 + uri = f"data:application/pdf;base64,{payload}" + result = _truncate_base64_in_string(uri) + assert "base64_data truncated" in result + assert "application/pdf" in result + assert payload not in result + + def test_multiple_data_uris(self): + payload = "B" * 200 + text = f"first: data:image/png;base64,{payload} second: data:image/jpeg;base64,{payload}" + result = _truncate_base64_in_string(text) + assert result.count("base64_data truncated") == 2 + + def test_no_data_uri(self): + text = "hello world, no base64 here" + assert _truncate_base64_in_string(text) == text + + +# --------------------------------------------------------------------------- +# truncate_base64_in_messages +# --------------------------------------------------------------------------- + + +class TestTruncateBase64InMessages: + def test_none_input(self): + assert truncate_base64_in_messages(None) is None + + def test_string_messages(self): + payload = "C" * 200 + msg = f"Look at data:image/png;base64,{payload}" + result = truncate_base64_in_messages(msg) + assert isinstance(result, str) + assert "base64_data truncated" in result + + def test_openai_vision_format(self): + """Typical OpenAI multimodal message with image_url containing base64.""" + payload = "D" * 500 + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is in this image?"}, + { + "type": "image_url", + "image_url": { + "url": f"data:image/png;base64,{payload}", + "detail": "auto", + }, + }, + ], + } + ] + result = truncate_base64_in_messages(messages) + # Original must not be mutated + assert payload in messages[0]["content"][1]["image_url"]["url"] + # Result should be truncated + url = result[0]["content"][1]["image_url"]["url"] + assert "base64_data truncated" in url + assert payload not in url + # Non-base64 parts preserved + assert result[0]["content"][0]["text"] == "What is in this image?" + + def test_multiple_images(self): + """Two base64 images in one message.""" + payload1 = "E" * 300 + payload2 = "F" * 400 + messages = [ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": {"url": f"data:image/png;base64,{payload1}"}, + }, + { + "type": "image_url", + "image_url": {"url": f"data:application/pdf;base64,{payload2}"}, + }, + ], + } + ] + result = truncate_base64_in_messages(messages) + for part in result[0]["content"]: + assert "base64_data truncated" in part["image_url"]["url"] + + def test_does_not_mutate_original(self): + payload = "G" * 200 + messages = [{"role": "user", "content": f"data:image/png;base64,{payload}"}] + truncate_base64_in_messages(messages) + # Original unchanged + assert payload in messages[0]["content"] + + def test_dict_messages(self): + payload = "H" * 200 + messages = {"prompt": f"data:image/png;base64,{payload}"} + result = truncate_base64_in_messages(messages) + assert "base64_data truncated" in result["prompt"] + + def test_preserves_short_base64(self): + """Short base64 under threshold should not be truncated.""" + short = "AAAA" + messages = [ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": {"url": f"data:image/png;base64,{short}"}, + } + ], + } + ] + result = truncate_base64_in_messages(messages) + assert result[0]["content"][0]["image_url"]["url"] == f"data:image/png;base64,{short}"