add tests for callback timing, base64 truncation, and detailed timing

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

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}"