phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

test(test_custom_callback.py): add test for message redaction to standard logging object

Browse source
This commit is contained in:
Krrish Dholakia 2024-08-19 17:06:41 -07:00
parent 8ce8680a9a
commit 0ce476a7f6
1 changed files with 16 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

17
litellm/tests/test_custom_callback_input.py
View file

@ -1172,7 +1172,13 @@ def test_turn_off_message_logging():
@pytest.mark.parametrize("model", ["gpt-3.5-turbo", "azure/chatgpt-v-2"]) @pytest.mark.parametrize("model", ["gpt-3.5-turbo", "azure/chatgpt-v-2"])
def test_standard_logging_payload(model): @pytest.mark.parametrize(
"turn_off_message_logging",
[
True,
],
) # False
def test_standard_logging_payload(model, turn_off_message_logging):
""" """
Ensure valid standard_logging_payload is passed for logging calls to s3 Ensure valid standard_logging_payload is passed for logging calls to s3
@ -1184,6 +1190,8 @@ def test_standard_logging_payload(model):
customHandler = CompletionCustomHandler() customHandler = CompletionCustomHandler()
litellm.callbacks = [customHandler] litellm.callbacks = [customHandler]
litellm.turn_off_message_logging = turn_off_message_logging
with patch.object( with patch.object(
customHandler, "log_success_event", new=MagicMock() customHandler, "log_success_event", new=MagicMock()
) as mock_client: ) as mock_client:
@ -1237,3 +1245,10 @@ def test_standard_logging_payload(model):
]["model_map_value"] ]["model_map_value"]
is not None is not None
) )
## turn off message logging
slobject: StandardLoggingPayload = mock_client.call_args.kwargs["kwargs"][
"standard_logging_object"
]
if turn_off_message_logging:
assert "redacted-by-litellm" == slobject["messages"][0]["content"]