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

fix(http_handler.py): fix error message masking

Browse source
This commit is contained in:
Krrish Dholakia 2024-11-30 00:18:03 -08:00
parent c6124984aa
commit 11c11f3724
2 changed files with 8 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

11
litellm/llms/custom_httpx/http_handler.py
View file

@ -33,7 +33,10 @@ import re
def mask_sensitive_info(error_message): def mask_sensitive_info(error_message):
# Find the start of the key parameter # 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 is found
if key_index != -1: if key_index != -1:
@ -212,8 +215,8 @@ class AsyncHTTPHandler:
setattr(e, "message", await e.response.aread()) setattr(e, "message", await e.response.aread())
setattr(e, "text", await e.response.aread()) setattr(e, "text", await e.response.aread())
else: else:
setattr(e, "message", e.response.text) setattr(e, "message", mask_sensitive_info(e.response.text))
setattr(e, "text", e.response.text) setattr(e, "text", mask_sensitive_info(e.response.text))
e = MaskedHTTPStatusError( e = MaskedHTTPStatusError(
e, message=getattr(e, "message", None), text=getattr(e, "text", None) e, message=getattr(e, "message", None), text=getattr(e, "text", None)
) )
@ -455,12 +458,12 @@ class HTTPHandler:
llm_provider="litellm-httpx-handler", llm_provider="litellm-httpx-handler",
) )
except httpx.HTTPStatusError as e: except httpx.HTTPStatusError as e:
error_text = mask_sensitive_info(e.response.text)
if stream is True: if stream is True:
setattr(e, "message", mask_sensitive_info(e.response.read())) setattr(e, "message", mask_sensitive_info(e.response.read()))
setattr(e, "text", mask_sensitive_info(e.response.read())) setattr(e, "text", mask_sensitive_info(e.response.read()))
else: else:
error_text = mask_sensitive_info(e.response.text)
setattr(e, "message", error_text) setattr(e, "message", error_text)
setattr(e, "text", error_text) setattr(e, "text", error_text)

2
tests/otel_tests/test_guardrails.py
View file

@ -212,7 +212,7 @@ async def test_bedrock_guardrail_triggered():
session, session,
"sk-1234", "sk-1234",
model="fake-openai-endpoint", 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"], guardrails=["bedrock-pre-guard"],
) )
pytest.fail("Should have thrown an exception") pytest.fail("Should have thrown an exception")