mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-05 02:41:56 +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
e3b4ef7cad
commit
49c6afe64f
3 changed files with 82 additions and 2 deletions
|
|
@ -2416,6 +2416,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 ##
|
||||
|
|
|
|||
|
|
@ -179,13 +179,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
|
||||
|
||||
|
|
|
|||
|
|
@ -1058,3 +1058,77 @@ def test_append_system_prompt_messages():
|
|||
kwargs=None, messages=messages
|
||||
)
|
||||
assert result == messages
|
||||
|
||||
|
||||
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