diff --git a/litellm/litellm_core_utils/logging_utils.py b/litellm/litellm_core_utils/logging_utils.py index 52a32a71d53..c74b9bc7b73 100644 --- a/litellm/litellm_core_utils/logging_utils.py +++ b/litellm/litellm_core_utils/logging_utils.py @@ -3,7 +3,7 @@ import functools import inspect import re import time -from collections.abc import Iterator, Mapping, Sequence +from collections.abc import Callable, Iterator, Mapping, Sequence from datetime import datetime from typing import TYPE_CHECKING, Any, Final, TypeAlias, cast @@ -76,6 +76,7 @@ def _truncate_base64_in_string(value: str) -> str: def _truncate_base64_in_value( value: str | dict[str, object] | list[object] | None, + truncate_string: Callable[[str], str], ) -> str | dict[str, object] | list[object] | None: """Iteratively truncate base64 data URIs in a JSON-like value (str/list/dict). @@ -85,13 +86,13 @@ def _truncate_base64_in_value( # Stack entries: (source_value, depth, parent_container, key_or_index) # We mutate *copies* of dicts/lists in-place via parent references. if isinstance(value, str): - return _truncate_base64_in_string(value) + return truncate_string(value) if not isinstance(value, (dict, list)): return value # Shallow-copy the root so we don't mutate the caller's data. root: Final = {k: v for k, v in value.items()} if isinstance(value, dict) else list(value) - stack: Final[list] = [(root, 0)] + stack: Final[list[tuple[dict[str, object] | list[object], int]]] = [(root, 0)] # mutable-ok: explicit stack while stack: container, depth = stack.pop() @@ -100,19 +101,19 @@ def _truncate_base64_in_value( if isinstance(container, dict): for k, v in container.items(): if isinstance(v, str): - container[k] = _truncate_base64_in_string(v) + container[k] = truncate_string(v) elif isinstance(v, dict): - copy: dict | list = {ck: cv for ck, cv in v.items()} + copy: dict[str, object] | list[object] = {ck: cv for ck, cv in v.items()} container[k] = copy stack.append((copy, depth + 1)) elif isinstance(v, list): copy = list(v) container[k] = copy stack.append((copy, depth + 1)) - elif isinstance(container, list): + else: for i, v in enumerate(container): if isinstance(v, str): - container[i] = _truncate_base64_in_string(v) + container[i] = truncate_string(v) elif isinstance(v, dict): copy = {ck: cv for ck, cv in v.items()} container[i] = copy @@ -134,8 +135,15 @@ def truncate_base64_in_messages( """ if messages is None or MAX_BASE64_LENGTH_FOR_LOGGING <= 0: return messages + return _copy_messages(messages, _truncate_base64_in_string) + + +def _copy_messages( + messages: str | dict[str, object] | list[object], # mutable-ok: same contract as truncate_base64_in_messages + truncate_string: Callable[[str], str], +) -> str | dict[str, object] | list[object] | None: # mutable-ok: same contract as truncate_base64_in_messages try: - return _truncate_base64_in_value(messages) + return _truncate_base64_in_value(messages, truncate_string) except Exception as e: verbose_logger.debug("Failed to truncate base64 in messages: %s", e) return messages @@ -177,8 +185,10 @@ async def truncate_base64_in_messages_async( if messages is None or MAX_BASE64_LENGTH_FOR_LOGGING <= 0: return messages leaves: Final = tuple(_iter_string_leaves(messages)) - if sum(len(leaf) for leaf in leaves) < BASE64_TRUNCATION_OFFLOAD_THRESHOLD_CHARS or _is_few_plain_strings(leaves): + if sum(len(leaf) for leaf in leaves) < BASE64_TRUNCATION_OFFLOAD_THRESHOLD_CHARS: return truncate_base64_in_messages(messages) + if _is_few_plain_strings(leaves): + return _copy_messages(messages, lambda text: text) return await asyncio.to_thread(truncate_base64_in_messages, messages) diff --git a/tests/unit/litellm_core_utils/test_logging_utils.py b/tests/unit/litellm_core_utils/test_logging_utils.py index 70b0be401ac..83c1f6e2f6c 100644 --- a/tests/unit/litellm_core_utils/test_logging_utils.py +++ b/tests/unit/litellm_core_utils/test_logging_utils.py @@ -5,6 +5,7 @@ Tests for litellm.litellm_core_utils.logging_utils — base64 truncation helpers import asyncio, datetime, importlib, litellm, os, pytest_asyncio import sys import threading +from itertools import chain from types import FrameType from unittest.mock import MagicMock @@ -19,7 +20,7 @@ from litellm.litellm_core_utils.logging_utils import ( truncate_base64_in_messages, truncate_base64_in_messages_async, ) -from collections.abc import AsyncIterator, Callable, Coroutine +from collections.abc import AsyncIterator, Callable, Coroutine, Sequence from datetime import datetime as datetime_assemble_streaming from litellm import( Choices, @@ -325,7 +326,7 @@ class TestTruncateBase64InMessagesAsync: ], ids=["json_rows", "sse_lines", "code_with_semicolons", "sse_lines_with_semicolons", "base64_marker_without_data"], ) - def test_tool_message_without_data_uri_is_copied_without_a_regex_scan(self, row): + def test_tool_message_without_data_uri_is_copied_without_a_regex_scan(self, row: str): tool_output: Final = (row * (_PLAIN_TOOL_OUTPUT_CHARS // len(row) + 1))[:_PLAIN_TOOL_OUTPUT_CHARS] messages: Final = _tool_result_messages(tool_output) @@ -347,7 +348,9 @@ class TestTruncateBase64InMessagesAsync: assert payload not in result[0]["content"][1]["image_url"]["url"] assert "base64_data truncated" in result[0]["content"][1]["image_url"]["url"] - def test_large_plain_text_is_copied_inline_without_a_thread_hop(self, monkeypatch): + def test_large_plain_text_is_copied_inline_without_rescanning_or_a_thread_hop( + self, monkeypatch: pytest.MonkeyPatch, scan_threads: Sequence[int] + ): monkeypatch.setattr(logging_utils, "BASE64_TRUNCATION_OFFLOAD_THRESHOLD_CHARS", 1_000) messages: Final = _tool_result_messages("plain tool output\n" * 100) @@ -355,12 +358,15 @@ class TestTruncateBase64InMessagesAsync: assert result == messages assert result is not messages + assert scan_threads == [] @pytest.mark.asyncio - async def test_long_plain_text_history_is_copied_off_the_event_loop(self, monkeypatch, scan_threads): + async def test_long_plain_text_history_is_copied_off_the_event_loop( + self, monkeypatch: pytest.MonkeyPatch, scan_threads: Sequence[int] + ): monkeypatch.setattr(logging_utils, "BASE64_TRUNCATION_OFFLOAD_THRESHOLD_CHARS", 1_000) monkeypatch.setattr(logging_utils, "BASE64_TRUNCATION_OFFLOAD_THRESHOLD_STRINGS", 10) - messages: Final = [message for turn in range(5) for message in _tool_result_messages(f"rows {turn}\n" * 100)] + messages: Final = list(chain.from_iterable(_tool_result_messages(f"rows {turn}\n" * 100) for turn in range(5))) result: Final = await truncate_base64_in_messages_async(messages) offload_threads: Final = tuple(scan_threads)