litellm/litellm/litellm_core_utils/redact_messages.py
2024-06-13 13:49:23 -07:00

65 lines
2.5 KiB
Python

# +-----------------------------------------------+
# | |
# | Give Feedback / Get Help |
# | https://github.com/BerriAI/litellm/issues/new |
# | |
# +-----------------------------------------------+
#
# Thank you users! We ❤️ you! - Krrish & Ishaan
import copy
from typing import TYPE_CHECKING, Any
import litellm
if TYPE_CHECKING:
from litellm.utils import Logging as _LiteLLMLoggingObject
LiteLLMLoggingObject = _LiteLLMLoggingObject
else:
LiteLLMLoggingObject = Any
def redact_message_input_output_from_logging(
litellm_logging_obj: LiteLLMLoggingObject, result
):
"""
Removes messages, prompts, input, response from logging. This modifies the data in-place
only redacts when litellm.turn_off_message_logging == True
"""
# check if user opted out of logging message/response to callbacks
if litellm.turn_off_message_logging is not True:
return result
_result = copy.deepcopy(result)
# remove messages, prompts, input, response from logging
litellm_logging_obj.model_call_details["messages"] = [
{"role": "user", "content": "redacted-by-litellm"}
]
litellm_logging_obj.model_call_details["prompt"] = ""
litellm_logging_obj.model_call_details["input"] = ""
# response cleaning
# ChatCompletion Responses
if (
litellm_logging_obj.stream is True
and "complete_streaming_response" in litellm_logging_obj.model_call_details
):
_streaming_response = litellm_logging_obj.model_call_details[
"complete_streaming_response"
]
for choice in _streaming_response.choices:
if isinstance(choice, litellm.Choices):
choice.message.content = "redacted-by-litellm"
elif isinstance(choice, litellm.utils.StreamingChoices):
choice.delta.content = "redacted-by-litellm"
else:
if _result is not None:
if isinstance(_result, litellm.ModelResponse):
if hasattr(_result, "choices") and _result.choices is not None:
for choice in _result.choices:
if isinstance(choice, litellm.Choices):
choice.message.content = "redacted-by-litellm"
elif isinstance(choice, litellm.utils.StreamingChoices):
choice.delta.content = "redacted-by-litellm"
return _result