fix(logging): honor turn_off_message_logging for per-callback streaming copies

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yucheng 2026-09-14 19:08:12 +00:00
parent 1b04f2f2be
commit f637ea189f
2 changed files with 36 additions and 3 deletions

View file

@ -48,10 +48,15 @@ def redact_message_input_output_from_custom_logger(
def redact_streaming_responses_for_custom_logger(model_call_details: dict, custom_logger: CustomLogger) -> dict:
"""
Returns a copy of model_call_details whose streaming response entries are redacted deepcopies
when the custom logger has opted out of message logging. The shared model_call_details is left
untouched so other callbacks still receive the unredacted response.
when the custom logger has opted out of message logging via `message_logging=False` or
`turn_off_message_logging=True`. The shared model_call_details is left untouched so other
callbacks still receive the unredacted response.
"""
if not (hasattr(custom_logger, "message_logging") and custom_logger.message_logging is not True):
opted_out: Final = (
getattr(custom_logger, "message_logging", True) is not True
or getattr(custom_logger, "turn_off_message_logging", False) is True
)
if not opted_out:
return model_call_details
redacted_entries: Final = {
streaming_key: _redacted_streaming_response_copy(model_call_details[streaming_key])

View file

@ -850,6 +850,34 @@ class TestRedactStreamingResponsesForCustomLogger:
assert response_obj.choices[0].message.content == "secret content"
assert model_call_details["async_complete_streaming_response"] is response_obj
def test_turn_off_message_logging_gets_redacted_copy(self):
response_obj = litellm.ModelResponse(
choices=[
litellm.Choices(
message=litellm.Message(content="stream-secret", role="assistant")
)
]
)
model_call_details = {
"stream": True,
"complete_streaming_response": response_obj,
}
opted_out_logger = CustomLogger(turn_off_message_logging=True)
default_logger = CustomLogger()
redacted_details = redact_streaming_responses_for_custom_logger(
model_call_details=model_call_details, custom_logger=opted_out_logger
)
shared_details = redact_streaming_responses_for_custom_logger(
model_call_details=model_call_details, custom_logger=default_logger
)
assert redacted_details["complete_streaming_response"].choices[0].message.content == "redacted-by-litellm"
assert shared_details is model_call_details
assert shared_details["complete_streaming_response"].choices[0].message.content == "stream-secret"
assert model_call_details["complete_streaming_response"] is response_obj
assert response_obj.choices[0].message.content == "stream-secret"
def test_compliant_logger_gets_shared_response(self):
model_call_details, response_obj = self._model_call_details()
compliant_logger = CustomLogger()