From 11c11f372405fc6b5ab9959d950b72389081c709 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 30 Nov 2024 00:18:03 -0800 Subject: [PATCH] fix(http_handler.py): fix error message masking --- litellm/llms/custom_httpx/http_handler.py | 11 +++++++---- tests/otel_tests/test_guardrails.py | 2 +- 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index bbbfce0bf..0186907a6 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -33,7 +33,10 @@ import re def mask_sensitive_info(error_message): # Find the start of the key parameter - key_index = error_message.find("key=") + if isinstance(error_message, str): + key_index = error_message.find("key=") + else: + return error_message # If key is found if key_index != -1: @@ -212,8 +215,8 @@ class AsyncHTTPHandler: setattr(e, "message", await e.response.aread()) setattr(e, "text", await e.response.aread()) else: - setattr(e, "message", e.response.text) - setattr(e, "text", e.response.text) + setattr(e, "message", mask_sensitive_info(e.response.text)) + setattr(e, "text", mask_sensitive_info(e.response.text)) e = MaskedHTTPStatusError( e, message=getattr(e, "message", None), text=getattr(e, "text", None) ) @@ -455,12 +458,12 @@ class HTTPHandler: llm_provider="litellm-httpx-handler", ) except httpx.HTTPStatusError as e: - error_text = mask_sensitive_info(e.response.text) if stream is True: setattr(e, "message", mask_sensitive_info(e.response.read())) setattr(e, "text", mask_sensitive_info(e.response.read())) else: + error_text = mask_sensitive_info(e.response.text) setattr(e, "message", error_text) setattr(e, "text", error_text) diff --git a/tests/otel_tests/test_guardrails.py b/tests/otel_tests/test_guardrails.py index 342ce33b9..12d9d1c38 100644 --- a/tests/otel_tests/test_guardrails.py +++ b/tests/otel_tests/test_guardrails.py @@ -212,7 +212,7 @@ async def test_bedrock_guardrail_triggered(): session, "sk-1234", model="fake-openai-endpoint", - messages=[{"role": "user", "content": f"Hello do you like coffee?"}], + messages=[{"role": "user", "content": "Hello do you like coffee?"}], guardrails=["bedrock-pre-guard"], ) pytest.fail("Should have thrown an exception")