From d5cae91fdc2b48e37fef981a5370df93134183c0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 09:06:45 -0700 Subject: [PATCH] perf(logging): check data URI markers before the scan and hop only for data URIs Check for ';', 'data:', and ';base64,' before running the data URI regex, and offload to a worker thread only when a request holds a data URI. Tests now watch every regex call and confirm plain text is copied without yielding to the loop. --- litellm/litellm_core_utils/logging_utils.py | 17 ++-- .../litellm_core_utils/test_logging_utils.py | 86 ++++++++++++++----- 2 files changed, 72 insertions(+), 31 deletions(-) diff --git a/litellm/litellm_core_utils/logging_utils.py b/litellm/litellm_core_utils/logging_utils.py index 8e9e23c5cb7..9a139b2f67d 100644 --- a/litellm/litellm_core_utils/logging_utils.py +++ b/litellm/litellm_core_utils/logging_utils.py @@ -44,6 +44,7 @@ Helper utils used for logging callbacks # Regex matching data-URI base64 content: "data:;base64," # Captures: group(1)=mime_type, group(2)=base64_payload _DATA_URI_RE: Final = re.compile(r"data:([^;,\s]{1,255});base64,([A-Za-z0-9+/=]+)") +_DATA_URI_PREFIX: Final = "data:" _DATA_URI_BASE64_MARKER: Final = ";base64," # Maximum nesting depth for _truncate_base64_in_value to guard against @@ -62,7 +63,7 @@ def _base64_data_uri_replacer(match: re.Match[str]) -> str: def _may_contain_base64_data_uri(value: str) -> bool: - return ";" in value and _DATA_URI_BASE64_MARKER in value + return ";" in value and _DATA_URI_PREFIX in value and _DATA_URI_BASE64_MARKER in value def _truncate_base64_in_string(value: str) -> str: @@ -160,16 +161,16 @@ async def truncate_base64_in_messages_async( messages: str | list | dict | None, # mutable-ok: same contract as truncate_base64_in_messages ) -> str | list | dict | None: # mutable-ok: same contract as truncate_base64_in_messages """ - Same result as truncate_base64_in_messages, but payloads whose data-URI-bearing - strings reach BASE64_TRUNCATION_OFFLOAD_THRESHOLD_CHARS are scanned in a worker - thread so the regex pass over multi-MB base64 images does not block the event loop. + Same result as truncate_base64_in_messages, but payloads with a data URI whose + string content reaches BASE64_TRUNCATION_OFFLOAD_THRESHOLD_CHARS are scanned in a + worker thread so the regex pass over multi-MB base64 images does not block the event loop. """ if messages is None or MAX_BASE64_LENGTH_FOR_LOGGING <= 0: return messages - data_uri_chars: Final = sum( - len(leaf) for leaf in _iter_string_leaves(messages) if _may_contain_base64_data_uri(leaf) - ) - if data_uri_chars < BASE64_TRUNCATION_OFFLOAD_THRESHOLD_CHARS: + leaves: Final = tuple(_iter_string_leaves(messages)) + if sum(len(leaf) for leaf in leaves) < BASE64_TRUNCATION_OFFLOAD_THRESHOLD_CHARS or not any( + _may_contain_base64_data_uri(leaf) for leaf in leaves + ): return truncate_base64_in_messages(messages) 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 5cfd470cb9a..69bec2cba63 100644 --- a/tests/unit/litellm_core_utils/test_logging_utils.py +++ b/tests/unit/litellm_core_utils/test_logging_utils.py @@ -3,6 +3,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 unittest.mock import MagicMock @@ -17,7 +18,7 @@ from litellm.litellm_core_utils.logging_utils import ( truncate_base64_in_messages, truncate_base64_in_messages_async, ) -from collections.abc import AsyncIterator +from collections.abc import AsyncIterator, Callable, Coroutine from datetime import datetime as datetime_assemble_streaming from litellm import( Choices, @@ -99,8 +100,8 @@ class TestTruncateBase64InString: @pytest.mark.parametrize( "text", [ - 'data: {"choices": [{"delta": {"content": "hi"}}]}\n\n' * 50_000, - "data:" * 200_000, + 'data: {"choices": [{"delta": {"content": "hi"}}]}\n\n' * 50_000 + ";base64,", + "data:" * 200_000 + ";base64,", ], ids=["sse_lines", "whitespace_free_prefixes"], ) @@ -241,19 +242,30 @@ def scan_threads(monkeypatch): return threads -@pytest.fixture -def regex_scans(monkeypatch): - """Record every string the data URI regex runs over.""" - scanned: list[str] = [] - original = logging_utils._DATA_URI_RE +def _run_without_yielding(coroutine: Coroutine[object, object, object]) -> object: + try: + coroutine.send(None) + except StopIteration as finished: + return finished.value + coroutine.close() + raise AssertionError("the coroutine yielded to the event loop") - class RecordingPattern: - def sub(self, repl, string: str) -> str: - scanned.append(string) - return original.sub(repl, string) - monkeypatch.setattr(logging_utils, "_DATA_URI_RE", RecordingPattern()) - return scanned +def _regex_calls_while(run: Callable[[], object]) -> tuple[object, list[str]]: + calls: list[str] = [] + previous = sys.getprofile() + + def record(frame, event: str, arg: object) -> None: + name = getattr(arg, "__qualname__", "") + if event == "c_call" and name.startswith("Pattern."): + calls.append(name) + + sys.setprofile(record) + try: + result = run() + finally: + sys.setprofile(previous) + return result, calls def _tool_result_messages(tool_output: str) -> list: @@ -300,28 +312,56 @@ class TestTruncateBase64InMessagesAsync: assert scan_threads assert set(scan_threads) == {threading.get_ident()} - @pytest.mark.asyncio @pytest.mark.parametrize( "row", [ '{"id": 7, "name": "widget", "status": "ok"}\n', "data: {\"choices\": [{\"delta\": {\"content\": \"hi\"}}]}\n\n", "const total = rows.map((r) => r.price).reduce((a, b) => a + b, 0);\n", + 'data: {"choices": [{"delta": {"content": "a; b"}}]}\n\n', + "Content-Type: text/plain;base64,aGVsbG8=\n", ], - ids=["json_rows", "sse_lines", "code_with_semicolons"], + ids=["json_rows", "sse_lines", "code_with_semicolons", "sse_lines_with_semicolons", "base64_marker_without_data"], ) - async def test_tool_message_without_data_uri_skips_the_regex_and_the_thread_hop( - self, scan_threads, regex_scans, row - ): + @pytest.mark.parametrize( + "truncate", + [lambda messages: _run_without_yielding(truncate_base64_in_messages_async(messages)), truncate_base64_in_messages], + ids=["async", "sync"], + ) + def test_tool_message_without_data_uri_is_copied_inline_without_a_regex_scan(self, row, truncate): tool_output = (row * (_PLAIN_TOOL_OUTPUT_CHARS // len(row) + 1))[:_PLAIN_TOOL_OUTPUT_CHARS] messages = _tool_result_messages(tool_output) - result = await truncate_base64_in_messages_async(messages) + result, regex_calls = _regex_calls_while(lambda: truncate(messages)) assert result == messages - assert regex_scans == [] - assert scan_threads - assert set(scan_threads) == {threading.get_ident()} + assert result is not messages + assert regex_calls == [] + + def test_small_data_uri_is_truncated_inline_with_a_regex_scan(self): + payload = "L" * 2_000 + messages = _image_messages(payload) + + result, regex_calls = _regex_calls_while( + lambda: _run_without_yielding(truncate_base64_in_messages_async(messages)) + ) + + assert regex_calls + assert payload not in result[0]["content"][1]["image_url"]["url"] + assert "base64_data truncated" in result[0]["content"][1]["image_url"]["url"] + + @pytest.mark.asyncio + async def test_small_data_uri_beside_large_text_is_scanned_off_the_event_loop(self, monkeypatch, scan_threads): + monkeypatch.setattr(logging_utils, "BASE64_TRUNCATION_OFFLOAD_THRESHOLD_CHARS", 1_000) + payload = "M" * 200 + messages = _image_messages(payload) + _tool_result_messages("plain tool output\n" * 100) + + result = await truncate_base64_in_messages_async(messages) + offload_threads = tuple(scan_threads) + + assert payload not in result[0]["content"][1]["image_url"]["url"] + assert offload_threads + assert threading.get_ident() not in offload_threads @pytest.mark.asyncio async def test_none_and_disabled_truncation_short_circuit(self, monkeypatch, scan_threads):