From 6b91f48c644228f23e68165f79e7163040c2ed1c Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 13 Feb 2024 22:11:03 -0800 Subject: [PATCH] fix(presidio_pii_masking.py): fix conditional check --- litellm/proxy/hooks/presidio_pii_masking.py | 4 ++-- litellm/tests/test_presidio_masking.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/hooks/presidio_pii_masking.py b/litellm/proxy/hooks/presidio_pii_masking.py index 25cd3c54c..85e626074 100644 --- a/litellm/proxy/hooks/presidio_pii_masking.py +++ b/litellm/proxy/hooks/presidio_pii_masking.py @@ -154,8 +154,8 @@ class _OPTIONAL_PresidioPIIMasking(CustomLogger): if litellm.output_parse_pii == False: return response - if isinstance(response, ModelResponse) and isinstance( - response.choices, StreamingChoices + if isinstance(response, ModelResponse) and not isinstance( + response.choices[0], StreamingChoices ): # /chat/completions requests if isinstance(response.choices[0].message.content, str): verbose_proxy_logger.debug( diff --git a/litellm/tests/test_presidio_masking.py b/litellm/tests/test_presidio_masking.py index 2275e78e9..f7e1ac12d 100644 --- a/litellm/tests/test_presidio_masking.py +++ b/litellm/tests/test_presidio_masking.py @@ -53,7 +53,7 @@ async def test_output_parsing(): mock_response="Hello ! How can I assist you today?", ) new_response = await pii_masking.async_post_call_success_hook( - user_api_key_dict=UserAPIKeyAuth(), response=response, call_type="completion" + user_api_key_dict=UserAPIKeyAuth(), response=response ) assert (