mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
fix(proxy): carry the guardrail headers on a guardrail's own ProxyException
This commit is contained in:
parent
4a6cba89f5
commit
5172a03a4d
2 changed files with 35 additions and 1 deletions
|
|
@ -2263,7 +2263,8 @@ async def pass_through_request(
|
|||
cache_key=None,
|
||||
api_base=str(url._uri_reference) if url else None,
|
||||
)
|
||||
custom_headers.update(_guardrail_response_headers(_parsed_body, kwargs, post_call_guardrail_data)) # pyright: ignore[reportUnknownArgumentType, reportUnknownMemberType] # untyped request and header dicts
|
||||
guardrail_headers: Final = _guardrail_response_headers(_parsed_body, kwargs, post_call_guardrail_data) # pyright: ignore[reportUnknownArgumentType] # untyped request dicts
|
||||
custom_headers.update(guardrail_headers) # pyright: ignore[reportUnknownMemberType] # untyped header dict
|
||||
if CustomGuardrail._is_guardrail_intervention(e):
|
||||
verbose_proxy_logger.warning(
|
||||
"pass_through_endpoint: request blocked by guardrail - %s",
|
||||
|
|
@ -2300,6 +2301,7 @@ async def pass_through_request(
|
|||
#########################################################
|
||||
|
||||
if isinstance(e, ProxyException):
|
||||
e.headers = {**e.headers, **guardrail_headers}
|
||||
raise
|
||||
if isinstance(e, HTTPException):
|
||||
raise ProxyException(
|
||||
|
|
|
|||
|
|
@ -486,6 +486,38 @@ async def test_pass_through_request_preserves_proxy_exception_status():
|
|||
assert exc.value.param == "limit"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_guardrail_proxy_exception_carries_the_guardrails_that_ran():
|
||||
blocked: Final = ProxyException(message="blocked by guardrail", type="None", param="None", code=400)
|
||||
|
||||
async def record_then_block(
|
||||
user_api_key_dict: UserAPIKeyAuth, data: dict[str, object], call_type: str, endpoint_type: object
|
||||
) -> None:
|
||||
data["metadata"] = {"applied_guardrails": ["passed-guard"]} # mutable-ok: a guardrail records on the live body
|
||||
raise blocked
|
||||
|
||||
with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging:
|
||||
mock_proxy_logging.post_call_failure_hook = AsyncMock()
|
||||
mock_proxy_logging.pre_call_hook = AsyncMock(side_effect=record_then_block)
|
||||
|
||||
mock_request = MagicMock(spec=Request)
|
||||
mock_request.method = "POST"
|
||||
mock_request.body = AsyncMock(return_value=b'{"model": "gpt-5.6", "messages": []}')
|
||||
mock_request.headers = Headers({"content-type": "application/json"})
|
||||
mock_request.query_params = QueryParams({})
|
||||
|
||||
with pytest.raises(ProxyException) as exc:
|
||||
await pass_through_request(
|
||||
request=mock_request,
|
||||
target="http://test.com/v1/chat/completions",
|
||||
custom_headers={},
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"),
|
||||
)
|
||||
|
||||
assert exc.value is blocked
|
||||
assert exc.value.headers["x-litellm-applied-guardrails"] == "passed-guard"
|
||||
|
||||
|
||||
def test_is_langfuse_route():
|
||||
"""
|
||||
Test that the is_langfuse_route method correctly identifies Langfuse routes
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue