mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-05 20:23:18 +00:00
fix(guardrails): return HTTP 400 instead of 500 for Model Armor streaming blocks (#24693)
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:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user