mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-28 01:32:17 +00:00
b6fcd03848
* fix(bedrock): surface a converse-stream 200 that decodes to no events as a 502 instead of an empty turn * fix(bedrock): quote the body head only when a stream decoded no events The leftover-bytes error keeps the byte and event counts, the content type and the request id but no longer quotes the first bytes of a stream that already decoded events, since that head is the start of a healthy stream and can hold model output. The anthropic_messages empty-stream warning no longer prints the request's model name. --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
927 lines
35 KiB
Python
927 lines
35 KiB
Python
import base64
|
|
import binascii
|
|
import itertools
|
|
import datetime
|
|
import json
|
|
import struct
|
|
from collections.abc import AsyncIterator, Mapping, Sequence
|
|
from typing import Final
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
import httpx
|
|
import pytest
|
|
|
|
import litellm
|
|
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
|
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
|
|
from litellm.llms.bedrock.chat.invoke_handler import (
|
|
AmazonOpenAICompatibleStreamDecoder,
|
|
AWSEventStreamDecoder,
|
|
make_call,
|
|
make_sync_call,
|
|
)
|
|
from litellm.exceptions import MidStreamFallbackError
|
|
from litellm.llms.bedrock.common_utils import BedrockError
|
|
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
|
|
from litellm.types.utils import ModelResponseStream
|
|
|
|
|
|
def test_transform_thinking_blocks_with_redacted_content():
|
|
thinking_block = {"redactedContent": "This is a redacted content"}
|
|
decoder = AWSEventStreamDecoder(model="test")
|
|
transformed_thinking_blocks = decoder.translate_thinking_blocks(thinking_block)
|
|
assert len(transformed_thinking_blocks) == 1
|
|
assert transformed_thinking_blocks[0]["type"] == "redacted_thinking"
|
|
assert transformed_thinking_blocks[0]["data"] == "This is a redacted content"
|
|
|
|
|
|
def test_transform_tool_calls_index():
|
|
chunks = [
|
|
{
|
|
"delta": {"text": "Certainly! I can help you with the"},
|
|
"contentBlockIndex": 0,
|
|
},
|
|
{
|
|
"delta": {"text": " current weather and time in Tokyo."},
|
|
"contentBlockIndex": 0,
|
|
},
|
|
{"delta": {"text": " To get this information, I'll"}, "contentBlockIndex": 0},
|
|
{"delta": {"text": " need to use two"}, "contentBlockIndex": 0},
|
|
{"delta": {"text": " different tools: one"}, "contentBlockIndex": 0},
|
|
{"delta": {"text": " for the weather and one for"}, "contentBlockIndex": 0},
|
|
{"delta": {"text": " the time. Let me fetch"}, "contentBlockIndex": 0},
|
|
{"delta": {"text": " that data for you."}, "contentBlockIndex": 0},
|
|
{
|
|
"start": {
|
|
"toolUse": {
|
|
"toolUseId": "tooluse_JX1wqyUvRjyTcVSg_6-JwA",
|
|
"name": "Weather_Tool",
|
|
}
|
|
},
|
|
"contentBlockIndex": 1,
|
|
},
|
|
{"delta": {"toolUse": {"input": ""}}, "contentBlockIndex": 1},
|
|
{"delta": {"toolUse": {"input": '{"locatio'}}, "contentBlockIndex": 1},
|
|
{"delta": {"toolUse": {"input": 'n": "Toky'}}, "contentBlockIndex": 1},
|
|
{"delta": {"toolUse": {"input": 'o"}'}}, "contentBlockIndex": 1},
|
|
{
|
|
"start": {
|
|
"toolUse": {
|
|
"toolUseId": "tooluse_rxDBNjDMQ-mqA-YOp9_3cQ",
|
|
"name": "Query_Time_Tool",
|
|
}
|
|
},
|
|
"contentBlockIndex": 2,
|
|
},
|
|
{"delta": {"toolUse": {"input": ""}}, "contentBlockIndex": 2},
|
|
{"delta": {"toolUse": {"input": '{"locati'}}, "contentBlockIndex": 2},
|
|
{"delta": {"toolUse": {"input": 'on"'}}, "contentBlockIndex": 2},
|
|
{"delta": {"toolUse": {"input": ': "Tokyo"}'}}, "contentBlockIndex": 2},
|
|
{"stopReason": "tool_use"},
|
|
]
|
|
decoder = AWSEventStreamDecoder(model="test")
|
|
parsed_chunks = []
|
|
for chunk in chunks:
|
|
parsed_chunk = decoder._chunk_parser(chunk)
|
|
parsed_chunks.append(parsed_chunk)
|
|
tool_call_chunks1 = parsed_chunks[8:12]
|
|
tool_call_chunks2 = parsed_chunks[13:17]
|
|
for tool_call_hunk in tool_call_chunks1:
|
|
tool_call_hunk_dict = tool_call_hunk.model_dump()
|
|
for tool_call in tool_call_hunk_dict["choices"][0]["delta"]["tool_calls"]:
|
|
assert tool_call["index"] == 0
|
|
for tool_call_hunk in tool_call_chunks2:
|
|
tool_call_hunk_dict = tool_call_hunk.model_dump()
|
|
for tool_call in tool_call_hunk_dict["choices"][0]["delta"]["tool_calls"]:
|
|
assert tool_call["index"] == 1
|
|
|
|
|
|
def test_transform_tool_calls_index_with_optional_arg_func():
|
|
chunks = [
|
|
{
|
|
"contentBlockIndex": 0,
|
|
"delta": {"text": "To"},
|
|
"p": "abcdefghijklmnopqrstuv",
|
|
},
|
|
{
|
|
"contentBlockIndex": 0,
|
|
"delta": {"text": " get the current time, I"},
|
|
"p": "abcdefghijklmnopqrstuvwxyzABCD",
|
|
},
|
|
{
|
|
"contentBlockIndex": 0,
|
|
"delta": {"text": ' can use the "get_time"'},
|
|
"p": "abcdefghijkl",
|
|
},
|
|
{
|
|
"contentBlockIndex": 0,
|
|
"delta": {"text": " function. Since the user"},
|
|
"p": "abcdefghijkl",
|
|
},
|
|
{
|
|
"contentBlockIndex": 0,
|
|
"delta": {"text": " didn't specify whether"},
|
|
"p": "abcdefghijklmnopqrstuvw",
|
|
},
|
|
{
|
|
"contentBlockIndex": 0,
|
|
"delta": {"text": " they want UTC time or local time,"},
|
|
"p": "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUV",
|
|
},
|
|
{
|
|
"contentBlockIndex": 0,
|
|
"delta": {"text": " I'll assume they"},
|
|
"p": "abcdefghijkl",
|
|
},
|
|
{
|
|
"contentBlockIndex": 0,
|
|
"delta": {"text": " want the local time. Here's"},
|
|
"p": "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMN",
|
|
},
|
|
{
|
|
"contentBlockIndex": 0,
|
|
"delta": {"text": " how I"},
|
|
"p": "abcdefghijklmnopqrstuvw",
|
|
},
|
|
{
|
|
"contentBlockIndex": 0,
|
|
"delta": {"text": "'ll make the function call:"},
|
|
"p": "abcdefghijklmnopqrstuvwxyzAB",
|
|
},
|
|
{
|
|
"contentBlockIndex": 0,
|
|
"p": "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ",
|
|
},
|
|
{
|
|
"contentBlockIndex": 1,
|
|
"p": "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNO",
|
|
"start": {
|
|
"toolUse": {
|
|
"name": "get_time",
|
|
"toolUseId": "tooluse_htgmgeJATsKTl4s_LW77sQ",
|
|
}
|
|
},
|
|
},
|
|
{
|
|
"contentBlockIndex": 1,
|
|
"delta": {"toolUse": {"input": ""}},
|
|
"p": "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUV",
|
|
},
|
|
{"contentBlockIndex": 1, "p": "abcdefghijklmnopqrstuvw"},
|
|
{"p": "abcdefghijklmnopqrstuvwxyzABCDEFGHIJK", "stopReason": "tool_use"},
|
|
]
|
|
decoder = AWSEventStreamDecoder(model="test")
|
|
parsed_chunks = []
|
|
for chunk in chunks:
|
|
parsed_chunk = decoder._chunk_parser(chunk)
|
|
parsed_chunks.append(parsed_chunk)
|
|
tool_call_chunks = parsed_chunks[11:14]
|
|
for tool_call_hunk in tool_call_chunks:
|
|
tool_call_hunk_dict = tool_call_hunk.model_dump()
|
|
for tool_call in tool_call_hunk_dict["choices"][0]["delta"]["tool_calls"]:
|
|
assert tool_call["index"] == 0
|
|
|
|
|
|
def test_bedrock_converse_streaming_consistent_id():
|
|
"""
|
|
Tests that all chunks in a Bedrock Converse stream response share the same ID,
|
|
capturing the ID from the initial 'messageStart' event.
|
|
"""
|
|
# Simulate a realistic Bedrock Converse stream
|
|
native_conversation_id = "a1b2c3d4-e5f6-7890-1234-56789abcdef0"
|
|
mock_stream_chunks = [
|
|
{
|
|
"messageStart": {
|
|
"conversationId": native_conversation_id,
|
|
"role": "assistant",
|
|
}
|
|
},
|
|
{"delta": {"text": "Hello"}, "contentBlockIndex": 0},
|
|
{"delta": {"text": " world!"}, "contentBlockIndex": 0},
|
|
{"stopReason": "stop"},
|
|
]
|
|
|
|
decoder = AWSEventStreamDecoder(model="bedrock/anthropic.claude-3-sonnet-v1:0")
|
|
|
|
# Process each chunk and collect the parsed responses
|
|
parsed_responses = []
|
|
for chunk in mock_stream_chunks:
|
|
parsed_responses.append(decoder.converse_chunk_parser(chunk))
|
|
|
|
# Verify that all parsed responses have the same, non-null ID derived from the native ID
|
|
assert len(parsed_responses) > 1, "Should have processed multiple chunks"
|
|
|
|
expected_id = f"chatcmpl-{native_conversation_id}"
|
|
|
|
for response in parsed_responses:
|
|
assert (
|
|
response.id == expected_id
|
|
), "All chunk IDs must match the one captured from the messageStart event"
|
|
|
|
|
|
def test_converse_streaming_usage_uses_provider_thinking_tokens():
|
|
"""Regression LIT-5714: the messageStop event carries provider thinking tokens
|
|
under ``additionalModelResponseFields``; the usage chunk must report them instead
|
|
of a token_counter estimate."""
|
|
chunks = [
|
|
{
|
|
"contentBlockIndex": 0,
|
|
"delta": {"reasoningContent": {"text": "thinking about it"}},
|
|
},
|
|
{
|
|
"stopReason": "end_turn",
|
|
"additionalModelResponseFields": {"usage": {"output_tokens_details": {"thinking_tokens": 1033}}},
|
|
},
|
|
{"usage": {"inputTokens": 40, "outputTokens": 3002, "totalTokens": 3042}},
|
|
]
|
|
|
|
decoder = AWSEventStreamDecoder(model="bedrock/anthropic.claude-opus-4-7")
|
|
parsed = [decoder.converse_chunk_parser(chunk) for chunk in chunks]
|
|
|
|
usage = parsed[-1].usage
|
|
assert usage.completion_tokens_details.reasoning_tokens == 1033
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_make_call_does_not_rechunk_stream_by_default():
|
|
"""Re-chunking the event stream into fixed 1024-byte blocks holds small
|
|
early events (messageStart, contentBlockStart) in httpx's ByteChunker until
|
|
1024 bytes accumulate, delaying time-to-first-chunk by the whole generation
|
|
when Bedrock trickles bytes (e.g. buffered tool-use streams)."""
|
|
response = MagicMock()
|
|
response.status_code = 200
|
|
client = MagicMock()
|
|
client.post = AsyncMock(return_value=response)
|
|
|
|
await make_call(
|
|
client=client,
|
|
api_base="https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-sonnet-4-6/converse-stream",
|
|
headers={},
|
|
data="{}",
|
|
model="anthropic.claude-sonnet-4-6",
|
|
messages=[],
|
|
logging_obj=MagicMock(),
|
|
)
|
|
|
|
response.aiter_bytes.assert_called_once_with(chunk_size=None)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_make_call_honors_explicit_stream_chunk_size():
|
|
response = MagicMock()
|
|
response.status_code = 200
|
|
client = MagicMock()
|
|
client.post = AsyncMock(return_value=response)
|
|
|
|
await make_call(
|
|
client=client,
|
|
api_base="https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-sonnet-4-6/converse-stream",
|
|
headers={},
|
|
data="{}",
|
|
model="anthropic.claude-sonnet-4-6",
|
|
messages=[],
|
|
logging_obj=MagicMock(),
|
|
stream_chunk_size=2048,
|
|
)
|
|
|
|
response.aiter_bytes.assert_called_once_with(chunk_size=2048)
|
|
|
|
|
|
def test_make_sync_call_does_not_rechunk_stream_by_default():
|
|
response = MagicMock()
|
|
response.status_code = 200
|
|
client = MagicMock()
|
|
client.post = MagicMock(return_value=response)
|
|
|
|
make_sync_call(
|
|
client=client,
|
|
api_base="https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-sonnet-4-6/converse-stream",
|
|
headers={},
|
|
data="{}",
|
|
signed_json_body=None,
|
|
model="anthropic.claude-sonnet-4-6",
|
|
messages=[],
|
|
logging_obj=MagicMock(),
|
|
)
|
|
|
|
response.iter_bytes.assert_called_once_with(chunk_size=None)
|
|
|
|
|
|
def test_make_sync_call_honors_explicit_stream_chunk_size():
|
|
response = MagicMock()
|
|
response.status_code = 200
|
|
client = MagicMock()
|
|
client.post = MagicMock(return_value=response)
|
|
|
|
make_sync_call(
|
|
client=client,
|
|
api_base="https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-sonnet-4-6/converse-stream",
|
|
headers={},
|
|
data="{}",
|
|
signed_json_body=None,
|
|
model="anthropic.claude-sonnet-4-6",
|
|
messages=[],
|
|
logging_obj=MagicMock(),
|
|
stream_chunk_size=2048,
|
|
)
|
|
|
|
response.iter_bytes.assert_called_once_with(chunk_size=2048)
|
|
|
|
|
|
CONVERSE_MODEL = "anthropic.claude-sonnet-4-6"
|
|
CONVERSE_METADATA_EVENT = {
|
|
"usage": {"inputTokens": 10, "outputTokens": 5, "totalTokens": 15},
|
|
"metrics": {"latencyMs": 100},
|
|
}
|
|
|
|
|
|
def _converse_stream_wrapper(events, model=CONVERSE_MODEL):
|
|
async def bedrock_stream():
|
|
decoder = AWSEventStreamDecoder(model=model)
|
|
for event in events:
|
|
yield decoder._chunk_parser(chunk_data=event)
|
|
|
|
return CustomStreamWrapper(
|
|
completion_stream=bedrock_stream(),
|
|
model=model,
|
|
custom_llm_provider="bedrock",
|
|
logging_obj=LiteLLMLoggingObj(
|
|
model=model,
|
|
messages=[{"role": "user", "content": "hi"}],
|
|
stream=True,
|
|
call_type="completion",
|
|
start_time=datetime.datetime.now(),
|
|
litellm_call_id="1234",
|
|
function_id="1234",
|
|
),
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"events, expected_finish_reason",
|
|
[
|
|
pytest.param(
|
|
(
|
|
{"role": "assistant"},
|
|
{"contentBlockIndex": 0, "delta": {"text": "Hello"}},
|
|
{"contentBlockIndex": 0, "delta": {"text": " world"}},
|
|
{"contentBlockIndex": 0},
|
|
{"stopReason": "end_turn"},
|
|
CONVERSE_METADATA_EVENT,
|
|
),
|
|
"stop",
|
|
id="text",
|
|
),
|
|
pytest.param(
|
|
(
|
|
{"role": "assistant"},
|
|
{"contentBlockIndex": 0, "start": {"toolUse": {"toolUseId": "t1", "name": "get_weather"}}},
|
|
{"contentBlockIndex": 0, "delta": {"toolUse": {"input": '{"city": "SF"}'}}},
|
|
{"contentBlockIndex": 0},
|
|
{"contentBlockIndex": 1, "start": {"toolUse": {"toolUseId": "t2", "name": "get_time"}}},
|
|
{"contentBlockIndex": 1, "delta": {"toolUse": {"input": '{"tz": "PT"}'}}},
|
|
{"contentBlockIndex": 1},
|
|
{"stopReason": "tool_use"},
|
|
CONVERSE_METADATA_EVENT,
|
|
),
|
|
"tool_calls",
|
|
id="multiple_tool_calls",
|
|
),
|
|
pytest.param(
|
|
(
|
|
{"role": "assistant"},
|
|
{"contentBlockIndex": 0, "start": {}},
|
|
{"contentBlockIndex": 0, "delta": {"text": "Let me check."}},
|
|
{"contentBlockIndex": 0},
|
|
{"contentBlockIndex": 1, "start": {"toolUse": {"toolUseId": "t1", "name": "get_weather"}}},
|
|
{"contentBlockIndex": 1, "delta": {"toolUse": {"input": '{"city": "SF"}'}}},
|
|
{"contentBlockIndex": 1},
|
|
{"stopReason": "tool_use"},
|
|
CONVERSE_METADATA_EVENT,
|
|
),
|
|
"tool_calls",
|
|
id="text_then_tool_call",
|
|
),
|
|
pytest.param(
|
|
(
|
|
{"role": "assistant"},
|
|
{"contentBlockIndex": 0, "start": {}},
|
|
{"contentBlockIndex": 0, "delta": {"reasoningContent": {"text": "thinking hard"}}},
|
|
{"contentBlockIndex": 0, "delta": {"reasoningContent": {"signature": "sig123"}}},
|
|
{"contentBlockIndex": 0},
|
|
{"contentBlockIndex": 1, "start": {}},
|
|
{"contentBlockIndex": 1, "delta": {"text": "Answer"}},
|
|
{"contentBlockIndex": 1},
|
|
{"stopReason": "end_turn"},
|
|
CONVERSE_METADATA_EVENT,
|
|
),
|
|
"stop",
|
|
id="reasoning_then_text",
|
|
),
|
|
],
|
|
)
|
|
@pytest.mark.asyncio
|
|
async def test_converse_stream_ends_on_finish_reason_chunk(events, expected_finish_reason):
|
|
"""The usage-only metadata event Bedrock sends after messageStop must not reach the caller as an extra
|
|
assistant delta following the finish_reason chunk."""
|
|
wrapper = _converse_stream_wrapper(events)
|
|
|
|
chunks = [chunk async for chunk in wrapper]
|
|
|
|
finish_reasons = [choice.finish_reason for chunk in chunks for choice in chunk.choices if choice.finish_reason]
|
|
assert finish_reasons == [expected_finish_reason]
|
|
assert chunks[-1].choices[0].finish_reason == expected_finish_reason, (
|
|
f"stream must end on the finish_reason chunk, got trailing {chunks[-1].model_dump(exclude_none=True)}"
|
|
)
|
|
roles = [choice.delta.role for chunk in chunks for choice in chunk.choices if choice.delta.role]
|
|
assert roles == ["assistant"]
|
|
assert any(getattr(chunk, "usage", None) is not None for chunk in wrapper.chunks)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_nova_invoke_stream_reports_bedrock_usage_and_finish_reason():
|
|
"""InvokeModel Nova wraps every Converse event under its event-type key and reports usage
|
|
without ``totalTokens``; the stream must end on Bedrock's finish reason and surface the
|
|
cached tokens instead of a token-count estimate."""
|
|
events = (
|
|
{"messageStart": {"role": "assistant"}},
|
|
{"contentBlockDelta": {"delta": {"text": "OK"}, "contentBlockIndex": 0}},
|
|
{"contentBlockDelta": {"delta": {"text": "."}, "contentBlockIndex": 0}},
|
|
{"contentBlockStop": {"contentBlockIndex": 0}},
|
|
{"messageStop": {"stopReason": "end_turn"}},
|
|
{
|
|
"metadata": {
|
|
"usage": {
|
|
"inputTokens": 5,
|
|
"outputTokens": 3,
|
|
"cacheReadInputTokenCount": 12262,
|
|
"cacheWriteInputTokenCount": 0,
|
|
},
|
|
"metrics": {},
|
|
"trace": {},
|
|
}
|
|
},
|
|
)
|
|
wrapper = _converse_stream_wrapper(events, model="bedrock/invoke/us.amazon.nova-pro-v1:0")
|
|
|
|
chunks = [chunk async for chunk in wrapper]
|
|
|
|
assert "".join(choice.delta.content or "" for chunk in chunks for choice in chunk.choices) == "OK."
|
|
finish_reasons = [choice.finish_reason for chunk in chunks for choice in chunk.choices if choice.finish_reason]
|
|
assert finish_reasons == ["stop"]
|
|
assert chunks[-1].choices[0].finish_reason == "stop"
|
|
usages = [chunk.usage for chunk in wrapper.chunks if getattr(chunk, "usage", None) is not None]
|
|
assert len(usages) == 1
|
|
assert usages[0].prompt_tokens == 12267
|
|
assert usages[0].prompt_tokens_details.cached_tokens == 12262
|
|
assert usages[0].completion_tokens == 3
|
|
assert usages[0].total_tokens == 12270
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_converse_stream_still_emits_guardrail_trace_after_finish_reason():
|
|
"""Guardrail metadata events carry a trace payload alongside usage; that chunk must still reach the caller
|
|
after the finish_reason chunk, as it did before the regression."""
|
|
trace = {"guardrail": {"inputAssessment": {"g1": {}}}}
|
|
events = (
|
|
{"role": "assistant"},
|
|
{"contentBlockIndex": 0, "delta": {"text": "Hello"}},
|
|
{"contentBlockIndex": 0},
|
|
{"stopReason": "end_turn"},
|
|
{**CONVERSE_METADATA_EVENT, "trace": trace},
|
|
)
|
|
wrapper = _converse_stream_wrapper(events)
|
|
|
|
chunks = [chunk async for chunk in wrapper]
|
|
|
|
finish_reasons = [choice.finish_reason for chunk in chunks for choice in chunk.choices if choice.finish_reason]
|
|
assert finish_reasons == ["stop"]
|
|
assert chunks[-1].provider_specific_fields == {"trace": trace}
|
|
assert chunks[-1].choices[0].delta.content == ""
|
|
assert chunks[-1].choices[0].delta.role == "assistant"
|
|
|
|
|
|
def test_invoke_streaming_forwards_bedrock_response_headers():
|
|
response = MagicMock()
|
|
response.status_code = 200
|
|
response.iter_bytes = MagicMock(return_value=iter([]))
|
|
response.headers = httpx.Headers({"x-amzn-requestid": "req-789"})
|
|
client = HTTPHandler()
|
|
client.post = MagicMock(return_value=response)
|
|
|
|
stream = litellm.completion(
|
|
model="bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0",
|
|
messages=[{"role": "user", "content": "hi"}],
|
|
stream=True,
|
|
client=client,
|
|
aws_access_key_id="fake",
|
|
aws_secret_access_key="fake",
|
|
aws_region_name="us-east-1",
|
|
)
|
|
|
|
assert stream._hidden_params["additional_headers"]["llm_provider-x-amzn-requestid"] == "req-789"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_invoke_streaming_forwards_bedrock_response_headers():
|
|
async def _no_bytes(chunk_size=None):
|
|
return
|
|
yield b""
|
|
|
|
response = MagicMock()
|
|
response.status_code = 200
|
|
response.aiter_bytes = _no_bytes
|
|
response.headers = httpx.Headers({"x-amzn-requestid": "req-987"})
|
|
client = AsyncHTTPHandler()
|
|
client.post = AsyncMock(return_value=response)
|
|
|
|
stream = await litellm.acompletion(
|
|
model="bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0",
|
|
messages=[{"role": "user", "content": "hi"}],
|
|
stream=True,
|
|
client=client,
|
|
aws_access_key_id="fake",
|
|
aws_secret_access_key="fake",
|
|
aws_region_name="us-east-1",
|
|
)
|
|
|
|
assert stream._hidden_params["additional_headers"]["llm_provider-x-amzn-requestid"] == "req-987"
|
|
|
|
|
|
def _bedrock_stream_error_response(status_code: int, request_id: str) -> httpx.Response:
|
|
return httpx.Response(
|
|
status_code=status_code,
|
|
headers={
|
|
"x-amzn-RequestId": request_id,
|
|
"x-amzn-ErrorType": "InternalServerException",
|
|
},
|
|
text='{"message":"Amazon Bedrock is unable to process your request."}',
|
|
request=httpx.Request("POST", "https://bedrock-runtime.us-east-1.amazonaws.com/"),
|
|
)
|
|
|
|
|
|
def test_invoke_streaming_error_forwards_bedrock_response_headers():
|
|
error_response = _bedrock_stream_error_response(500, "req-stream-err-1")
|
|
client = HTTPHandler()
|
|
client.post = MagicMock(
|
|
side_effect=httpx.HTTPStatusError(
|
|
"server error",
|
|
request=error_response.request,
|
|
response=error_response,
|
|
)
|
|
)
|
|
|
|
with pytest.raises(litellm.ServiceUnavailableError) as exc_info:
|
|
litellm.completion(
|
|
model="bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0",
|
|
messages=[{"role": "user", "content": "hi"}],
|
|
stream=True,
|
|
client=client,
|
|
aws_access_key_id="fake",
|
|
aws_secret_access_key="fake",
|
|
aws_region_name="us-east-1",
|
|
)
|
|
|
|
assert exc_info.value.response.headers["x-amzn-requestid"] == "req-stream-err-1"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_invoke_streaming_error_forwards_bedrock_response_headers():
|
|
error_response = _bedrock_stream_error_response(500, "req-stream-err-2")
|
|
client = AsyncHTTPHandler()
|
|
client.post = AsyncMock(
|
|
side_effect=httpx.HTTPStatusError(
|
|
"server error",
|
|
request=error_response.request,
|
|
response=error_response,
|
|
)
|
|
)
|
|
|
|
with pytest.raises(litellm.ServiceUnavailableError) as exc_info:
|
|
await litellm.acompletion(
|
|
model="bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0",
|
|
messages=[{"role": "user", "content": "hi"}],
|
|
stream=True,
|
|
client=client,
|
|
aws_access_key_id="fake",
|
|
aws_secret_access_key="fake",
|
|
aws_region_name="us-east-1",
|
|
)
|
|
|
|
assert exc_info.value.response.headers["x-amzn-requestid"] == "req-stream-err-2"
|
|
|
|
|
|
def _unread_bedrock_stream_error_response(status_code: int, request_id: str) -> httpx.Response:
|
|
return httpx.Response(
|
|
status_code=status_code,
|
|
headers={
|
|
"x-amzn-RequestId": request_id,
|
|
"x-amzn-ErrorType": "InternalServerException",
|
|
},
|
|
stream=httpx.ByteStream(b'{"message":"Amazon Bedrock is unable to process your request."}'),
|
|
request=httpx.Request("POST", "https://bedrock-runtime.us-east-1.amazonaws.com/"),
|
|
)
|
|
|
|
|
|
def test_invoke_streaming_error_forwards_headers_when_body_was_never_read():
|
|
"""A retried streamed request raises HTTPStatusError over a body nobody read, so
|
|
reading it for the error message throws and loses the request id (LIT-5428)."""
|
|
error_response = _unread_bedrock_stream_error_response(500, "req-unread-sync")
|
|
client = HTTPHandler()
|
|
client.post = MagicMock(
|
|
side_effect=httpx.HTTPStatusError(
|
|
"server error",
|
|
request=error_response.request,
|
|
response=error_response,
|
|
)
|
|
)
|
|
|
|
with pytest.raises(litellm.ServiceUnavailableError) as exc_info:
|
|
litellm.completion(
|
|
model="bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0",
|
|
messages=[{"role": "user", "content": "hi"}],
|
|
stream=True,
|
|
client=client,
|
|
aws_access_key_id="fake",
|
|
aws_secret_access_key="fake",
|
|
aws_region_name="us-east-1",
|
|
)
|
|
|
|
assert exc_info.value.response.headers["x-amzn-requestid"] == "req-unread-sync"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_invoke_streaming_error_forwards_headers_when_body_was_never_read():
|
|
error_response = _unread_bedrock_stream_error_response(500, "req-unread-async")
|
|
client = AsyncHTTPHandler()
|
|
client.post = AsyncMock(
|
|
side_effect=httpx.HTTPStatusError(
|
|
"server error",
|
|
request=error_response.request,
|
|
response=error_response,
|
|
)
|
|
)
|
|
|
|
with pytest.raises(litellm.ServiceUnavailableError) as exc_info:
|
|
await litellm.acompletion(
|
|
model="bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0",
|
|
messages=[{"role": "user", "content": "hi"}],
|
|
stream=True,
|
|
client=client,
|
|
aws_access_key_id="fake",
|
|
aws_secret_access_key="fake",
|
|
aws_region_name="us-east-1",
|
|
)
|
|
|
|
assert exc_info.value.response.headers["x-amzn-requestid"] == "req-unread-async"
|
|
|
|
|
|
def test_invoke_streaming_non_200_forwards_bedrock_response_headers():
|
|
"""A caller-supplied client that returns a failure instead of raising still reaches the
|
|
provider's headers, and reading the streamed body for the message must not throw (LIT-5428)."""
|
|
error_response = _unread_bedrock_stream_error_response(500, "req-non200-sync")
|
|
client = HTTPHandler()
|
|
client.post = MagicMock(return_value=error_response)
|
|
|
|
with pytest.raises(litellm.ServiceUnavailableError) as exc_info:
|
|
litellm.completion(
|
|
model="bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0",
|
|
messages=[{"role": "user", "content": "hi"}],
|
|
stream=True,
|
|
client=client,
|
|
aws_access_key_id="fake",
|
|
aws_secret_access_key="fake",
|
|
aws_region_name="us-east-1",
|
|
)
|
|
|
|
assert exc_info.value.response.headers["x-amzn-requestid"] == "req-non200-sync"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_invoke_streaming_non_200_forwards_bedrock_response_headers():
|
|
error_response = _unread_bedrock_stream_error_response(500, "req-non200-async")
|
|
client = AsyncHTTPHandler()
|
|
client.post = AsyncMock(return_value=error_response)
|
|
|
|
with pytest.raises(litellm.ServiceUnavailableError) as exc_info:
|
|
await litellm.acompletion(
|
|
model="bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0",
|
|
messages=[{"role": "user", "content": "hi"}],
|
|
stream=True,
|
|
client=client,
|
|
aws_access_key_id="fake",
|
|
aws_secret_access_key="fake",
|
|
aws_region_name="us-east-1",
|
|
)
|
|
|
|
assert exc_info.value.response.headers["x-amzn-requestid"] == "req-non200-async"
|
|
|
|
|
|
def _bedrock_event_stream_frame(chunk: Mapping[str, object]) -> 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(":event-type", "chunk") + header(":content-type", "application/json") + header(
|
|
":message-type", "event"
|
|
)
|
|
payload: Final = json.dumps({"bytes": base64.b64encode(json.dumps(chunk).encode()).decode()}).encode()
|
|
prelude: Final = struct.pack(">II", 12 + len(headers) + len(payload) + 4, len(headers))
|
|
body: Final = prelude + struct.pack(">I", binascii.crc32(prelude)) + headers + payload
|
|
return body + struct.pack(">I", binascii.crc32(body))
|
|
|
|
|
|
def _openai_stream_chunk(delta: Mapping[str, str], finish_reason: str | None = None) -> Mapping[str, object]:
|
|
return {
|
|
"id": "chatcmpl-1",
|
|
"object": "chat.completion.chunk",
|
|
"created": 1,
|
|
"model": "moonshot.kimi-k2-thinking",
|
|
"choices": [{"index": 0, "delta": delta, "finish_reason": finish_reason}],
|
|
}
|
|
|
|
|
|
_MOONSHOT_RAW_STREAM: Final = b"".join(
|
|
_bedrock_event_stream_frame(chunk)
|
|
for chunk in (
|
|
_openai_stream_chunk({"role": "assistant", "reasoning_content": "thinking"}),
|
|
_openai_stream_chunk({"content": '{"city": '}),
|
|
_openai_stream_chunk({"content": '"San Francisco"}'}),
|
|
_openai_stream_chunk({}, "stop"),
|
|
)
|
|
)
|
|
|
|
|
|
def _assert_moonshot_stream_content(chunks: Sequence[ModelResponseStream]) -> None:
|
|
assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == '{"city": "San Francisco"}'
|
|
assert "".join(getattr(chunk.choices[0].delta, "reasoning_content", None) or "" for chunk in chunks) == "thinking"
|
|
assert [chunk.choices[0].finish_reason for chunk in chunks if chunk.choices[0].finish_reason] == ["stop"]
|
|
|
|
|
|
@pytest.fixture
|
|
def _aws_test_credentials(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setenv("AWS_ACCESS_KEY_ID", "AKIATEST")
|
|
monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "test-secret")
|
|
monkeypatch.setenv("AWS_REGION_NAME", "us-east-1")
|
|
|
|
|
|
@pytest.mark.parametrize("response_format", [None, {"type": "json_object"}])
|
|
def test_moonshot_invoke_stream_yields_openai_shaped_chunks(
|
|
_aws_test_credentials: None, response_format: Mapping[str, str] | None
|
|
) -> None:
|
|
raw_stream: Final = _MOONSHOT_RAW_STREAM
|
|
response: Final = MagicMock(status_code=200, headers={})
|
|
response.iter_bytes = lambda chunk_size=None: iter([raw_stream])
|
|
client: Final = HTTPHandler()
|
|
client.post = MagicMock(return_value=response)
|
|
|
|
stream: Final = litellm.completion(
|
|
model="bedrock/invoke/moonshot.kimi-k2-thinking",
|
|
messages=[{"role": "user", "content": "weather as json"}],
|
|
stream=True,
|
|
client=client,
|
|
**({"response_format": response_format} if response_format else {}),
|
|
)
|
|
_assert_moonshot_stream_content(list(stream))
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_moonshot_invoke_async_stream_yields_openai_shaped_chunks(_aws_test_credentials: None) -> None:
|
|
async def _aiter_bytes(chunk_size: int | None = None) -> AsyncIterator[bytes]:
|
|
yield _MOONSHOT_RAW_STREAM
|
|
|
|
response: Final = MagicMock(status_code=200, headers={})
|
|
response.aiter_bytes = _aiter_bytes
|
|
client: Final = AsyncHTTPHandler()
|
|
client.post = AsyncMock(return_value=response)
|
|
|
|
stream: Final = await litellm.acompletion(
|
|
model="bedrock/invoke/moonshot.kimi-k2-thinking",
|
|
messages=[{"role": "user", "content": "weather as json"}],
|
|
stream=True,
|
|
response_format={"type": "json_object"},
|
|
client=client,
|
|
)
|
|
|
|
_assert_moonshot_stream_content([chunk async for chunk in stream])
|
|
|
|
|
|
def _truncated_frame() -> bytes:
|
|
return _bedrock_event_stream_frame(_openai_stream_chunk({"role": "assistant"}))[:-8]
|
|
|
|
|
|
def _event_stream_headers() -> httpx.Headers:
|
|
return httpx.Headers({"content-type": "application/vnd.amazon.eventstream", "x-amzn-RequestId": "req-empty-1"})
|
|
|
|
|
|
_UNDECODABLE_STREAM_BODIES: Final = (
|
|
pytest.param(b"", id="empty"),
|
|
pytest.param(b"\x00\x00\x00\x05", id="shorter-than-a-prelude"),
|
|
pytest.param(_truncated_frame(), id="truncated-first-message"),
|
|
)
|
|
|
|
|
|
def _assert_no_events_error(error: BedrockError, body: bytes) -> None:
|
|
assert error.status_code == 502
|
|
assert "HTTP 200" in error.message
|
|
assert "decoded to no events" in error.message
|
|
assert f"{len(body)} bytes received" in error.message
|
|
assert "application/vnd.amazon.eventstream" in error.message
|
|
assert "req-empty-1" in error.message
|
|
assert f"first bytes={body[:200]!r}" in error.message
|
|
|
|
|
|
@pytest.mark.parametrize("body", _UNDECODABLE_STREAM_BODIES)
|
|
def test_iter_bytes_raises_when_a_200_body_decodes_to_no_events(body: bytes) -> None:
|
|
decoder: Final = AWSEventStreamDecoder(model="us.moonshotai.kimi-k3")
|
|
|
|
with pytest.raises(BedrockError) as exc_info:
|
|
list(decoder.iter_bytes(iter([body]), response_headers=_event_stream_headers()))
|
|
|
|
_assert_no_events_error(exc_info.value, body)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("body", _UNDECODABLE_STREAM_BODIES)
|
|
async def test_aiter_bytes_raises_when_a_200_body_decodes_to_no_events(body: bytes) -> None:
|
|
async def _chunks() -> AsyncIterator[bytes]:
|
|
yield body
|
|
|
|
decoder: Final = AWSEventStreamDecoder(model="us.moonshotai.kimi-k3")
|
|
|
|
with pytest.raises(BedrockError) as exc_info:
|
|
_ = [chunk async for chunk in decoder.aiter_bytes(_chunks(), response_headers=_event_stream_headers())]
|
|
|
|
_assert_no_events_error(exc_info.value, body)
|
|
|
|
|
|
def test_iter_bytes_raises_when_the_stream_ends_mid_message() -> None:
|
|
decoder: Final = AmazonOpenAICompatibleStreamDecoder(model="moonshot.kimi-k2-thinking", sync_stream=True)
|
|
stream: Final = decoder.iter_bytes(iter([_MOONSHOT_RAW_STREAM, _truncated_frame()]))
|
|
|
|
chunks: Final = list(itertools.islice(stream, 4))
|
|
with pytest.raises(BedrockError) as exc_info:
|
|
next(stream)
|
|
|
|
_assert_moonshot_stream_content(chunks)
|
|
assert exc_info.value.status_code == 502
|
|
assert f"{len(_truncated_frame())} undecoded bytes after 4 events" in exc_info.value.message
|
|
assert "first bytes=" not in exc_info.value.message
|
|
|
|
|
|
def test_iter_bytes_yields_a_complete_stream_without_raising() -> None:
|
|
decoder: Final = AmazonOpenAICompatibleStreamDecoder(model="moonshot.kimi-k2-thinking", sync_stream=True)
|
|
|
|
chunks: Final = list(decoder.iter_bytes(iter([_MOONSHOT_RAW_STREAM[:100], _MOONSHOT_RAW_STREAM[100:]])))
|
|
|
|
_assert_moonshot_stream_content(chunks)
|
|
|
|
|
|
def _assert_empty_stream_surfaced_as_bad_gateway(error: MidStreamFallbackError) -> None:
|
|
assert error.status_code == 502
|
|
assert error.is_pre_first_chunk is True
|
|
assert isinstance(error.original_exception, litellm.BadGatewayError)
|
|
assert "decoded to no events" in str(error)
|
|
assert "req-empty-1" in str(error)
|
|
|
|
|
|
def test_converse_stream_with_an_empty_200_body_raises_instead_of_an_empty_turn(_aws_test_credentials: None) -> None:
|
|
response: Final = MagicMock(status_code=200, headers=_event_stream_headers())
|
|
response.iter_bytes = lambda chunk_size=None: iter([b""])
|
|
client: Final = HTTPHandler()
|
|
client.post = MagicMock(return_value=response)
|
|
|
|
with pytest.raises(MidStreamFallbackError) as exc_info:
|
|
list(
|
|
litellm.completion(
|
|
model="bedrock/us.moonshotai.kimi-k3",
|
|
messages=[{"role": "user", "content": "hi"}],
|
|
stream=True,
|
|
client=client,
|
|
)
|
|
)
|
|
|
|
_assert_empty_stream_surfaced_as_bad_gateway(exc_info.value)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_converse_stream_with_an_empty_200_body_raises_instead_of_an_empty_turn(
|
|
_aws_test_credentials: None,
|
|
) -> None:
|
|
async def _aiter_bytes(chunk_size: int | None = None) -> AsyncIterator[bytes]:
|
|
yield b""
|
|
|
|
response: Final = MagicMock(status_code=200, headers=_event_stream_headers())
|
|
response.aiter_bytes = _aiter_bytes
|
|
client: Final = AsyncHTTPHandler()
|
|
client.post = AsyncMock(return_value=response)
|
|
|
|
stream: Final = await litellm.acompletion(
|
|
model="bedrock/us.moonshotai.kimi-k3",
|
|
messages=[{"role": "user", "content": "hi"}],
|
|
stream=True,
|
|
client=client,
|
|
)
|
|
with pytest.raises(MidStreamFallbackError) as exc_info:
|
|
_ = [chunk async for chunk in stream]
|
|
|
|
_assert_empty_stream_surfaced_as_bad_gateway(exc_info.value)
|