diff --git a/tests/benchmarks/test_redact_messages_benchmarks.py b/tests/benchmarks/test_redact_messages_benchmarks.py new file mode 100644 index 00000000000..d360e5cb9ef --- /dev/null +++ b/tests/benchmarks/test_redact_messages_benchmarks.py @@ -0,0 +1,22 @@ +from typing import Final + +import pytest + +from litellm.litellm_core_utils.redact_messages import request_texts, scrub_request_text + +REQUEST_TEXTS: Final = request_texts("hi there") + + +@pytest.mark.benchmark +@pytest.mark.parametrize( + ("error", "scrubbed"), + [ + ("." * 100_000, "." * 100_000), + (" ...'x" * 20_000, " ...'x" * 20_000), + ("input_value=" * 10_000, "input_value=redacted-by-litellm"), + ("input_value=x, input_type=y\n" * 4_000, "input_value=redacted-by-litellm, input_type=y\n" * 4_000), + ], + ids=["dots", "quoted_cuts", "input_value_labels_on_one_line", "input_values_on_many_lines"], +) +def test_scrubbing_a_large_error(error: str, scrubbed: str) -> None: + assert scrub_request_text(error, REQUEST_TEXTS) == scrubbed diff --git a/tests/unit/litellm_core_utils/test_redact_messages.py b/tests/unit/litellm_core_utils/test_redact_messages.py index d9fe5a56f82..b2daddf2d09 100644 --- a/tests/unit/litellm_core_utils/test_redact_messages.py +++ b/tests/unit/litellm_core_utils/test_redact_messages.py @@ -11,7 +11,6 @@ import importlib import json import os import threading -import time import tracemalloc from collections.abc import AsyncIterator, Callable, Mapping from datetime import datetime @@ -2435,19 +2434,6 @@ def test_scrub_request_text_keeps_an_input_value_when_the_request_has_no_text() assert scrub_request_text(error, request_texts(None)) == error -@pytest.mark.parametrize( - "error", - ["." * 1_000_000, "input_value=" * 100_000, "input_value=x, input_type=y\n" * 40_000, " ...'x" * 200_000], - ids=["dots", "input_value_labels_on_one_line", "input_values_on_many_lines", "quoted_cuts"], -) -def test_scrub_request_text_stays_fast_on_a_large_error(error: str) -> None: - started: Final = time.perf_counter() - - scrub_request_text(error, request_texts("hi there")) - - assert time.perf_counter() - started < 0.5 - - def test_scrub_request_text_leaves_an_error_that_does_not_quote_the_prompt() -> None: assert scrub_request_text("Incorrect API key provided", request_texts(PROMPT, "hi")) == "Incorrect API key provided"