From 542fe70edca852ef70ab51e5f67d506dc4ea57bf Mon Sep 17 00:00:00 2001 From: rahul dhanawade Date: Sun, 15 Feb 2026 20:11:11 +0530 Subject: [PATCH] Address PR feedback: Add error type to status code mapping and comprehensive unit tests --- litellm/llms/bedrock/chat/invoke_handler.py | 56 +++++++----- .../llms/bedrock/chat/test_invoke_handler.py | 87 +++++++++++++++++++ 2 files changed, 121 insertions(+), 22 deletions(-) diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py index 0b140ccb2b9..c62b0f35514 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -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) diff --git a/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py b/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py index a415d550215..cd31e7ae507 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py +++ b/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py @@ -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 \ No newline at end of file