diff --git a/litellm/llms/anthropic/messages/native/handler.py b/litellm/llms/anthropic/messages/native/handler.py index 922de386fe4..bf1353044ca 100644 --- a/litellm/llms/anthropic/messages/native/handler.py +++ b/litellm/llms/anthropic/messages/native/handler.py @@ -597,18 +597,39 @@ async def anthropic_messages( ) ctx: Final = contextvars.copy_context() func_with_context: Final = partial(ctx.run, func) + + def map_exception(e: BaseLLMException) -> Exception: + try: + exception_type( + model=model, + custom_llm_provider=custom_llm_provider, + original_exception=e, + extra_kwargs=kwargs, + ) + except Exception as mapped: # noqa: BLE001 # exception_type hands the mapped error back by raising it + return mapped + return e + try: init_response: Final = await loop.run_in_executor(None, func_with_context) - if asyncio.iscoroutine(init_response): - return await init_response - return init_response + response: Final = await init_response if asyncio.iscoroutine(init_response) else init_response except BaseLLMException as e: - raise exception_type( - model=model, - custom_llm_provider=custom_llm_provider, - original_exception=e, - extra_kwargs=kwargs, - ) + raise map_exception(e) + return _with_stream_exceptions_mapped(response, map_exception) + + +def _with_stream_exceptions_mapped( + response: AnthropicMessagesResponse | Iterator[bytes] | AsyncIterator[object], + map_exception: Callable[[BaseLLMException], Exception], +) -> AnthropicMessagesResponse | Iterator[bytes] | AsyncIterator[object]: + from .streaming_iterator import AnthropicMessagesStreamingResponse, ExceptionMappedAnthropicSSEStream + + if not isinstance(response, AnthropicMessagesStreamingResponse): + return response + return AnthropicMessagesStreamingResponse( + completion_stream=ExceptionMappedAnthropicSSEStream(response.completion_stream, map_exception), + hidden_params=response.hidden_params, + ) def validate_anthropic_api_metadata(metadata: dict | None = None) -> dict | None: diff --git a/litellm/llms/anthropic/messages/native/streaming_iterator.py b/litellm/llms/anthropic/messages/native/streaming_iterator.py index 52b3f1042f8..7e96f97ed56 100644 --- a/litellm/llms/anthropic/messages/native/streaming_iterator.py +++ b/litellm/llms/anthropic/messages/native/streaming_iterator.py @@ -1,7 +1,7 @@ import asyncio import json import re -from collections.abc import AsyncGenerator, AsyncIterator, Mapping, Sequence +from collections.abc import AsyncGenerator, AsyncIterator, Callable, Mapping, Sequence from contextvars import ContextVar from datetime import datetime from typing import TYPE_CHECKING, Any, Final, Protocol, runtime_checkable @@ -25,6 +25,7 @@ from litellm.llms.anthropic.messages.native.utils import ( ) from litellm.llms.anthropic.messages.safeguards import StreamedSafeguardResults from litellm.llms.anthropic.messages.utils import anthropic_sse_frame +from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.proxy.pass_through_endpoints.success_handler import ( PassThroughEndpointLogging, ) @@ -493,6 +494,36 @@ class SafeguardedAnthropicSSEStream: await aclose_if_supported(self._stream) +class ExceptionMappedAnthropicSSEStream: + """ + Maps a provider exception raised while the stream is read (e.g. a Bedrock + ``exception`` event decoded into a ``BedrockError``) the same way + ``anthropic_messages`` maps one raised before the stream is returned, so + callers catch the litellm exception types instead of the internal + provider class. + """ + + def __init__(self, stream: AsyncIterator[bytes], map_exception: Callable[[BaseLLMException], Exception]) -> None: + self._stream: Final = stream + self._map_exception: Final = map_exception + + @property + def has_buffered_provider_output(self) -> bool: + return getattr(self._stream, "has_buffered_provider_output", False) is True + + def __aiter__(self) -> "ExceptionMappedAnthropicSSEStream": + return self + + async def __anext__(self) -> bytes: + try: + return await self._stream.__anext__() + except BaseLLMException as e: + raise self._map_exception(e) + + async def aclose(self) -> None: + await aclose_if_supported(self._stream) + + class BaseAnthropicMessagesStreamingIterator: """ Base class for Anthropic Messages streaming iterators that provides common logic diff --git a/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_handler.py b/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_handler.py index 8a924dd32d6..cbf5c8a561f 100644 --- a/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_handler.py +++ b/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_handler.py @@ -2,7 +2,9 @@ import asyncio import copy import json import os +import struct import uuid +import zlib from collections.abc import Mapping from dataclasses import dataclass from pathlib import Path @@ -2716,3 +2718,53 @@ async def test_anthropic_messages_strict_tools_rejection_is_sent_again_without_s assert first["tools"] == [tool] assert second == {**first, "tools": [{"name": "lookup_record", "input_schema": {"type": "object"}}]} assert _JSON_BODY.validate_python(response)["id"] == "msg_strict_retry" + + +def _bedrock_exception_frame(exception_type: str, payload: bytes) -> bytes: + def _header(name: str, value: str) -> bytes: + return bytes([len(name)]) + name.encode() + bytes([7]) + struct.pack(">H", len(value)) + value.encode() + + headers: Final = ( + _header(":message-type", "exception") + + _header(":exception-type", exception_type) + + _header(":content-type", "application/json") + ) + prelude: Final = struct.pack(">II", 16 + len(headers) + len(payload), len(headers)) + message: Final = prelude + struct.pack(">I", zlib.crc32(prelude)) + headers + payload + return message + struct.pack(">I", zlib.crc32(message)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("bedrock_exception_type", "mapped_exception", "status_code"), + [ + ("internalServerException", litellm.ServiceUnavailableError, 503), + ("throttlingException", litellm.RateLimitError, 429), + ], +) +async def test_bedrock_exception_event_mid_stream_raises_the_mapped_exception( + bedrock_exception_type: str, mapped_exception: type[Exception], status_code: int +) -> None: + frame: Final = _bedrock_exception_frame(bedrock_exception_type, b'{"message":"boom"}') + bedrock: Final = httpx.MockTransport( + lambda request: httpx.Response( + 200, headers={"content-type": "application/vnd.amazon.eventstream"}, content=frame, request=request + ) + ) + + stream: Final = await litellm.anthropic_messages( + max_tokens=16, + messages=[{"role": "user", "content": "hi"}], + model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", + stream=True, + aws_access_key_id="test-access-key", + aws_secret_access_key="test-secret-key", + aws_region_name="us-east-1", + client=AsyncHTTPHandler(transport=bedrock), + ) + + with pytest.raises(mapped_exception) as raised: + async for _ in stream: + pass + assert raised.value.status_code == status_code + assert bedrock_exception_type in str(raised.value)