scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-05 08:07:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 137

Address PR feedback: Add error type to status code mapping and comprehensive unit tests

Browse source
This commit is contained in:
rahul dhanawade 2026-02-15 20:11:11 +05:30
parent 2e8a4d052a
commit 542fe70edc
2 changed files with 121 additions and 22 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

56
litellm/llms/bedrock/chat/invoke_handler.py
View file

@ -15,12 +15,21 @@ from typing import (
cast,
get_args,
)
BEDROCK_ERROR_TYPE_TO_STATUS = {
"ValidationException": 400,
"AccessDeniedException": 403,
"ThrottlingException": 429,
"ServiceQuotaExceededException": 429,
"ResourceNotFoundException": 404,
"ModelNotReadyException": 503,
}
import httpx # type: ignore
import litellm
from litellm import verbose_logger
from litellm._uuid import uuid
import json
from litellm.caching.caching import InMemoryCache
from litellm.litellm_core_utils.core_helpers import map_finish_reason
from litellm.litellm_core_utils.litellm_logging import Logging
@ -67,7 +76,6 @@ from litellm.utils import CustomStreamWrapper, get_secret
from ..base_aws_llm import BaseAWSLLM
from ..common_utils import BedrockError, ModelResponseIterator, get_bedrock_tool_name
_response_stream_shape_cache = None
bedrock_tool_name_mappings: InMemoryCache = InMemoryCache(
max_size_in_memory=50, default_ttl=600
@ -1656,7 +1664,6 @@ class AWSEventStreamDecoder:
) -> Iterator[Union[GChunk, ModelResponseStream, dict]]:
"""Given an iterator that yields lines, iterate over it & yield every event encountered"""
from botocore.eventstream import EventStreamBuffer
import json
event_stream_buffer = EventStreamBuffer()
is_first_chunk = True
@ -1666,15 +1673,19 @@ class AWSEventStreamDecoder:
is_first_chunk = False
if chunk.strip().startswith(b"{"):
try:
from litellm.exceptions import BadRequestError
error_data = json.loads(chunk)
error_msg = error_data.get("message") or error_data.get("Message") or str(error_data)
raise BadRequestError(
message=f"Bedrock Error: {error_msg}",
model=getattr(self, "model", "bedrock-unknown"),
llm_provider="bedrock"
)
except (json.JSONDecodeError, UnicodeDecodeError):
if "message" in error_data:
# Extract error type
error_type = error_data.get("type") or error_data.get("__type", "")
# Get status code, default to 400 if type not found
status_code = BEDROCK_ERROR_TYPE_TO_STATUS.get(error_type, 400)
# Raise BedrockError with correct status code
raise BedrockError(
status_code=status_code,
message=error_data.get("message", "Unknown error"),
)
except (json.JSONDecodeError, KeyError):
pass
event_stream_buffer.add_data(chunk)
@ -1690,7 +1701,6 @@ class AWSEventStreamDecoder:
) -> AsyncIterator[Union[GChunk, ModelResponseStream, dict]]:
"""Given an async iterator that yields lines, iterate over it & yield every event encountered"""
from botocore.eventstream import EventStreamBuffer
import json
event_stream_buffer = EventStreamBuffer()
is_first_chunk = True
@ -1702,18 +1712,20 @@ class AWSEventStreamDecoder:
# If it starts with '{', it's a JSON error, not a binary stream
if chunk.strip().startswith(b"{"):
try:
from litellm.exceptions import BadRequestError
error_data = json.loads(chunk)
# Extract message and raise a clear exception
error_msg = error_data.get("message") or error_data.get("Message") or str(error_data)
raise BadRequestError(
message=f"Bedrock Error: {error_msg}",
model=getattr(self, "model", "bedrock-unknown"),
llm_provider="bedrock"
)
except (json.JSONDecodeError, UnicodeDecodeError):
# If parsing fails, fall back to normal stream processing
if "message" in error_data:
# Extract error type
error_type = error_data.get("type") or error_data.get("__type", "")
# Get status code, default to 400 if type not found
status_code = BEDROCK_ERROR_TYPE_TO_STATUS.get(error_type, 400)
# Raise BedrockError with correct status code
raise BedrockError(
status_code=status_code,
message=error_data.get("message", "Unknown error"),
)
except (json.JSONDecodeError, KeyError):
pass
event_stream_buffer.add_data(chunk)

87
tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py
View file

@ -1,5 +1,6 @@
import os
import sys
import pytest
sys.path.insert(
@ -7,6 +8,7 @@ sys.path.insert(
) # Adds the parent directory to the system path
from litellm.llms.bedrock.chat.invoke_handler import AWSEventStreamDecoder
from litellm.llms.bedrock.common_utils import BedrockError
def test_transform_thinking_blocks_with_redacted_content():
@ -200,3 +202,88 @@ def test_bedrock_converse_streaming_consistent_id():
assert (
response.id == expected_id
), "All chunk IDs must match the one captured from the messageStart event"
import pytest
from litellm.llms.bedrock.common_utils import BedrockError
def test_json_error_detection_validation_exception():
"""Test that ValidationException errors are properly detected with 400 status"""
error_chunk = b'{"message": "Invalid request", "type": "ValidationException"}'
decoder = AWSEventStreamDecoder(model="anthropic.claude-v2")
# Create a mock iterator that yields the error chunk
def mock_iterator():
yield error_chunk
with pytest.raises(BedrockError) as exc_info:
list(decoder.iter_bytes(mock_iterator()))
assert exc_info.value.status_code == 400
assert "Invalid request" in str(exc_info.value.message)
def test_json_error_detection_throttling_exception():
"""Test that ThrottlingException gets 429 status code"""
error_chunk = b'{"message": "Rate exceeded", "type": "ThrottlingException"}'
decoder = AWSEventStreamDecoder(model="anthropic.claude-v2")
def mock_iterator():
yield error_chunk
with pytest.raises(BedrockError) as exc_info:
list(decoder.iter_bytes(mock_iterator()))
assert exc_info.value.status_code == 429
assert "Rate exceeded" in str(exc_info.value.message)
def test_json_error_detection_access_denied():
"""Test that AccessDeniedException gets 403 status code"""
error_chunk = b'{"message": "Access denied", "type": "AccessDeniedException"}'
decoder = AWSEventStreamDecoder(model="anthropic.claude-v2")
def mock_iterator():
yield error_chunk
with pytest.raises(BedrockError) as exc_info:
list(decoder.iter_bytes(mock_iterator()))
assert exc_info.value.status_code == 403
assert "Access denied" in str(exc_info.value.message)
def test_json_error_detection_unknown_type_defaults_to_400():
"""Test that unknown error types default to 400 status code"""
error_chunk = b'{"message": "Unknown error type", "type": "UnknownException"}'
decoder = AWSEventStreamDecoder(model="anthropic.claude-v2")
def mock_iterator():
yield error_chunk
with pytest.raises(BedrockError) as exc_info:
list(decoder.iter_bytes(mock_iterator()))
assert exc_info.value.status_code == 400 # Should default to 400
assert "Unknown error type" in str(exc_info.value.message)
async def test_async_json_error_detection():
"""Test that async iterator also properly detects JSON errors"""
error_chunk = b'{"message": "Async error", "type": "ThrottlingException"}'
decoder = AWSEventStreamDecoder(model="anthropic.claude-v2")
async def mock_async_iterator():
yield error_chunk
with pytest.raises(BedrockError) as exc_info:
async for _ in decoder.aiter_bytes(mock_async_iterator()):
pass
assert exc_info.value.status_code == 429