mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-15 23:31:29 +00:00
perf: skip redundant redaction + avoid double check
- Add global_redaction_applied flag to skip per-callback redaction - Add should_redact param to avoid calling should_redact_message_logging twice - Add tests for both flag=True (skip) and flag=False (proceed) cases
This commit is contained in:
parent
aaef0ce3ea
commit
d527a042ae
3 changed files with 82 additions and 2 deletions
|
|
@ -2489,6 +2489,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
result = redact_message_input_output_from_logging(
|
||||
model_call_details=_model_call_details,
|
||||
result=result,
|
||||
should_redact=global_redaction_applied,
|
||||
)
|
||||
|
||||
## LOGGING HOOK ##
|
||||
|
|
|
|||
|
|
@ -189,13 +189,18 @@ def should_redact_message_logging(model_call_details: dict) -> bool:
|
|||
|
||||
|
||||
def redact_message_input_output_from_logging(
|
||||
model_call_details: dict, result, input: Optional[Any] = None
|
||||
model_call_details: dict,
|
||||
result,
|
||||
input: Optional[Any] = None,
|
||||
should_redact: Optional[bool] = None,
|
||||
) -> Any:
|
||||
"""
|
||||
Removes messages, prompts, input, response from logging. This modifies the data in-place
|
||||
only redacts when litellm.turn_off_message_logging == True
|
||||
"""
|
||||
if should_redact_message_logging(model_call_details):
|
||||
if should_redact is None:
|
||||
should_redact = should_redact_message_logging(model_call_details)
|
||||
if should_redact:
|
||||
return perform_redaction(model_call_details, result)
|
||||
return result
|
||||
|
||||
|
|
|
|||
|
|
@ -1355,3 +1355,77 @@ def test_get_error_information_error_code_priority():
|
|||
result = StandardLoggingPayloadSetup.get_error_information(no_code_exception)
|
||||
assert result["error_code"] == ""
|
||||
assert result["error_class"] == "NoCodeException"
|
||||
|
||||
|
||||
def test_global_redaction_skips_per_callback_redaction():
|
||||
"""
|
||||
When global_redaction_applied=True, per-callback redaction functions should
|
||||
return early without processing to avoid redundant work.
|
||||
"""
|
||||
from litellm.litellm_core_utils.redact_messages import (
|
||||
redact_message_input_output_from_custom_logger,
|
||||
)
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
|
||||
# Test redact_message_input_output_from_custom_logger skips when global redaction applied
|
||||
custom_logger = CustomLogger()
|
||||
custom_logger.message_logging = False # Would normally trigger redaction
|
||||
|
||||
mock_logging_obj = MagicMock()
|
||||
mock_logging_obj.model_call_details = {"messages": [{"content": "secret"}]}
|
||||
|
||||
original_result = {"response": "already-redacted"}
|
||||
|
||||
result = redact_message_input_output_from_custom_logger(
|
||||
litellm_logging_obj=mock_logging_obj,
|
||||
result=original_result,
|
||||
custom_logger=custom_logger,
|
||||
global_redaction_applied=True,
|
||||
)
|
||||
|
||||
assert result is original_result # Should return unchanged (early return)
|
||||
|
||||
# Test CustomLogger.redact_standard_logging_payload_from_model_call_details skips
|
||||
custom_logger.turn_off_message_logging = True
|
||||
|
||||
model_call_details = {
|
||||
"messages": [{"content": "secret"}],
|
||||
"standard_logging_object": {"response": "sensitive"},
|
||||
}
|
||||
|
||||
result = custom_logger.redact_standard_logging_payload_from_model_call_details(
|
||||
model_call_details=model_call_details,
|
||||
global_redaction_applied=True,
|
||||
)
|
||||
|
||||
assert result is model_call_details # Should return unchanged (early return)
|
||||
|
||||
|
||||
def test_per_callback_redaction_proceeds_when_global_redaction_not_applied():
|
||||
"""
|
||||
When global_redaction_applied=False, per-callback redaction should still
|
||||
proceed normally if the callback has redaction enabled.
|
||||
"""
|
||||
from litellm.litellm_core_utils.redact_messages import (
|
||||
redact_message_input_output_from_custom_logger,
|
||||
)
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
|
||||
# Test redact_message_input_output_from_custom_logger proceeds when global redaction NOT applied
|
||||
custom_logger = CustomLogger()
|
||||
custom_logger.message_logging = False # Triggers redaction
|
||||
|
||||
mock_logging_obj = MagicMock()
|
||||
mock_logging_obj.model_call_details = {"messages": [{"content": "secret"}]}
|
||||
|
||||
original_result = {"response": "sensitive-data"}
|
||||
|
||||
result = redact_message_input_output_from_custom_logger(
|
||||
litellm_logging_obj=mock_logging_obj,
|
||||
result=original_result,
|
||||
custom_logger=custom_logger,
|
||||
global_redaction_applied=False,
|
||||
)
|
||||
|
||||
# Result should be different (redacted), not the same object
|
||||
assert result is not original_result
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue