tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-02 06:22:48 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(guardrails): return HTTP 400 instead of 500 for Model Armor streaming blocks (#24693)

Browse Source 
When Model Armor blocks a streaming response, it correctly raises
HTTPException(status_code=400) but create_response() catches it with a
bare except Exception and hardcodes a 500 response, discarding the
original status code.

Fix create_response() to preserve status_code from HTTPException instead
of hardcoding 500. Also update Model Armor's streaming hook to yield an
SSE error event instead of raising (matching the Prisma Airs pattern),
and fix make_model_armor_request() to return 400 for upstream API
failures instead of passing through the upstream status code.
This commit is contained in:
michelligabriele
2026-04-02 21:28:52 -07:00
committed by GitHub
parent 52a596d2a4
commit a6dfd02610
4 changed files with 179 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/proxy/common_request_processing.py
+8 -3
View File
@@ -221,16 +221,21 @@ async def create_response(
f"Error consuming first chunk from generator: {e}"
)
# Fallback to a generic error stream
# Preserve status code from HTTPException (e.g., guardrail blocks)
error_status = getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR)
error_detail = getattr(e, "detail", "Error processing stream start")
if not isinstance(error_detail, str):
error_detail = str(error_detail)
async def error_gen_message() -> AsyncGenerator[str, None]:
yield f"data: {json.dumps({'error': {'message': 'Error processing stream start', 'code': status.HTTP_500_INTERNAL_SERVER_ERROR}})}\n\n"
yield f"data: {json.dumps({'error': {'message': error_detail, 'code': error_status}})}\n\n"
yield "data: [DONE]\n\n"
return StreamingResponse(
error_gen_message(),
media_type=media_type,
headers=headers,
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
status_code=error_status,
)
async def combined_generator() -> AsyncGenerator[str, None]:
litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py
+19 -4
View File
@@ -14,6 +14,8 @@ from fastapi import HTTPException
if TYPE_CHECKING:
from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel
import json
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.caching import DualCache
@@ -203,8 +205,8 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
response.text,
)
raise HTTPException(
status_code=response.status_code,
detail=f"Model Armor API error: {response.text}",
status_code=400,
detail=f"Model Armor API error (upstream {response.status_code}): {response.text}",
)
json_response = response.json()
@@ -746,8 +748,21 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
yield chunk
return
except HTTPException:
raise
except HTTPException as e:
# Yield error as SSE event so create_response() detects it and
# returns a proper JSON error response with the correct status code.
# (Raising from a generator hits create_response's generic except → 500.)
detail = (
e.detail if isinstance(e.detail, dict) else {"message": str(e.detail)}
)
error_value = detail.get("error", detail)
if isinstance(error_value, dict):
error_obj = dict(error_value)
else:
error_obj = {"message": str(error_value)}
error_obj["code"] = e.status_code
yield f"data: {json.dumps({'error': error_obj})}\n\n"
return
except Exception as e:
verbose_proxy_logger.error(
"Model Armor streaming error: %s", str(e), exc_info=True
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py
+124 -1
View File
@@ -380,8 +380,9 @@ async def test_model_armor_api_error_handling():
call_type="completion"
)
assert exc_info.value.status_code == 500
assert exc_info.value.status_code == 400
assert "Model Armor API error" in str(exc_info.value.detail)
assert "upstream 500" in str(exc_info.value.detail)
@pytest.mark.asyncio
@@ -485,6 +486,128 @@ async def test_model_armor_streaming_response():
assert len(result_chunks) > 0
mock_post.assert_called()
@pytest.mark.asyncio
async def test_model_armor_streaming_block_yields_sse_error():
"""Test that streaming content block yields SSE error event instead of raising HTTPException."""
mock_user_api_key_dict = UserAPIKeyAuth()
guardrail = ModelArmorGuardrail(
template_id="test-template",
project_id="test-project",
location="us-central1",
guardrail_name="model-armor-test",
)
# Mock Model Armor API response that triggers a block (SDP MATCH_FOUND)
mock_response = AsyncMock()
mock_response.status_code = 200
mock_response.json = AsyncMock(
return_value={
"sanitizationResult": {
"filterMatchState": "MATCH_FOUND",
"filterResults": {
"sdp": {
"sdpFilterResult": {
"inspectResult": {
"matchState": "MATCH_FOUND",
"findings": [
{
"infoType": "PASSWORD",
"likelihood": "VERY_LIKELY",
}
],
}
}
}
},
}
}
)
guardrail._ensure_access_token_async = AsyncMock(
return_value=("test-token", "test-project")
)
with patch.object(
guardrail.async_handler, "post", AsyncMock(return_value=mock_response)
):
async def mock_stream():
chunks = [
litellm.ModelResponseStream(
choices=[
litellm.types.utils.StreamingChoices(
delta=litellm.types.utils.Delta(
content="My password is "
)
)
]
),
litellm.ModelResponseStream(
choices=[
litellm.types.utils.StreamingChoices(
delta=litellm.types.utils.Delta(content="hunter2")
)
]
),
]
for chunk in chunks:
yield chunk
request_data = {
"model": "gpt-4",
"messages": [{"role": "user", "content": "What's your password?"}],
"metadata": {"guardrails": ["model-armor-test"]},
}
result_chunks = []
async for chunk in guardrail.async_post_call_streaming_iterator_hook(
user_api_key_dict=mock_user_api_key_dict,
response=mock_stream(),
request_data=request_data,
):
result_chunks.append(chunk)
# Should yield exactly one SSE error event (not raise HTTPException)
assert len(result_chunks) == 1
error_data = json.loads(result_chunks[0].removeprefix("data: "))
assert "error" in error_data
assert error_data["error"]["code"] == 400
@pytest.mark.asyncio
async def test_model_armor_api_failure_returns_400():
"""Test that Model Armor API failures raise HTTP 400, not the upstream status code."""
guardrail = ModelArmorGuardrail(
template_id="test-template",
project_id="test-project",
location="us-central1",
guardrail_name="model-armor-test",
)
# Mock a 500 response from the Model Armor GCP API
mock_response = AsyncMock()
mock_response.status_code = 500
mock_response.text = "Internal Server Error"
guardrail._ensure_access_token_async = AsyncMock(
return_value=("test-token", "test-project")
)
with patch.object(
guardrail.async_handler, "post", AsyncMock(return_value=mock_response)
):
with pytest.raises(HTTPException) as exc_info:
await guardrail.make_model_armor_request(
content="test content",
source="user_prompt",
)
# Should be 400, NOT the upstream 500
assert exc_info.value.status_code == 400
assert "upstream 500" in str(exc_info.value.detail)
def test_model_armor_ui_friendly_name():
"""Test the UI-friendly name of the Model Armor guardrail"""
from litellm.types.proxy.guardrails.guardrail_hooks.model_armor import (
tests/test_litellm/proxy/test_common_request_processing.py
+28 -1
View File
@@ -4,7 +4,7 @@ from typing import AsyncGenerator
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from fastapi import Request, status
from fastapi import HTTPException, Request, status
from fastapi.responses import JSONResponse, StreamingResponse
import litellm
@@ -899,6 +899,33 @@ class TestCommonRequestProcessingHelpers:
assert content[0] == f"data: {json.dumps(expected_error_data)}\n\n"
assert content[1] == "data: [DONE]\n\n"
async def test_create_streaming_response_generator_raises_http_exception(
self,
):
"""
Test that when a generator raises HTTPException, the response preserves
the original status code instead of hardcoding 500.
"""
mock_gen = AsyncMock()
mock_gen.__anext__.side_effect = HTTPException(
status_code=400, detail="Content blocked by guardrail"
)
response = await create_response(mock_gen, "text/event-stream", {})
assert response.status_code == 400
content = await self.consume_stream(response)
import json
expected_error_data = {
"error": {
"message": "Content blocked by guardrail",
"code": 400,
}
}
assert len(content) == 2
assert content[0] == f"data: {json.dumps(expected_error_data)}\n\n"
assert content[1] == "data: [DONE]\n\n"
async def test_create_streaming_response_first_chunk_error_string_code(self):
"""
Test that when the first chunk contains a string error code, a JSON error response is returned