mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-18 04:44:37 +00:00
fix(responses): preserve provider request IDs in streaming failure logs (#45453)
* fix(responses): preserve provider request IDs in streaming failure logs * fix(bedrock): log the provider request id for in-stream errors on bridged Bedrock streams Bedrock models that run through the chat-completions bridge (Claude, converse, invoke) raise in-stream failures from the event-stream decoder as a BedrockError without headers, so their failure logs kept error_provider_request_id null. The decoder now passes the stream's request-id headers to those errors. Both the native Responses iterator and the decoder pass only the PROVIDER_REQUEST_ID_HEADERS subset, so error responses and failure logs pick up the request id and none of the 200 response's other headers. MidStreamFallbackError gets its response at construction instead of mutating it afterwards. --------- Co-authored-by: moyai-devin-berriai[bot] <336287033+moyai-devin-berriai[bot]@users.noreply.github.com> Co-authored-by: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com>
This commit is contained in:
moyai-devin-berriai[bot]
GitHub
parent
b21e659c24
commit
a3251a3bfd
8 changed files with 977 additions and 20 deletions
|
|
@ -1,8 +1,11 @@
|
|||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
from litellm.constants import PROVIDER_REQUEST_ID_HEADERS
|
||||
|
||||
_HEADERS: Final = TypeAdapter(Mapping[str, str])
|
||||
_REQUEST_ID_HEADERS: Final = ("request-id", "x-request-id", "llm_provider-request-id", "llm_provider-x-request-id")
|
||||
|
||||
|
|
@ -16,6 +19,16 @@ def get_provider_request_id(headers: object) -> str | None:
|
|||
return next((normalized[name] for name in _REQUEST_ID_HEADERS if normalized.get(name)), None)
|
||||
|
||||
|
||||
def provider_request_id_headers(headers: Mapping[str, object] | None) -> httpx.Headers | None:
|
||||
"""The provider request id headers in `headers` whose values are ASCII strings, or None when there are none."""
|
||||
request_id_headers: Final = tuple(
|
||||
(name, value)
|
||||
for name, value in (headers or {}).items()
|
||||
if name.lower() in PROVIDER_REQUEST_ID_HEADERS and isinstance(value, str) and value.isascii()
|
||||
)
|
||||
return httpx.Headers(request_id_headers) if request_id_headers else None
|
||||
|
||||
|
||||
def get_response_headers(_response_headers: Mapping[str, str] | None = None) -> dict:
|
||||
"""
|
||||
|
||||
|
|
|
|||
|
|
@ -13,6 +13,7 @@ from litellm.caching.caching import InMemoryCache
|
|||
from litellm.constants import RESPONSE_FORMAT_TOOL_NAME
|
||||
from litellm.litellm_core_utils.core_helpers import map_finish_reason
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.litellm_core_utils.llm_response_utils.get_headers import provider_request_id_headers
|
||||
from litellm.litellm_core_utils.optional_imports import ensure_optional_import
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import native_redacted_thinking_data
|
||||
from litellm.llms.anthropic.chat.handler import (
|
||||
|
|
@ -448,6 +449,7 @@ class _EventStreamTally:
|
|||
f"x-amzn-requestid={_response_header(response_headers, 'x-amzn-requestid')!r}, "
|
||||
f"{self.bytes_received} bytes received"
|
||||
)
|
||||
request_id_headers: Final = provider_request_id_headers(response_headers)
|
||||
if not self.events:
|
||||
return BedrockError(
|
||||
status_code=502,
|
||||
|
|
@ -455,6 +457,7 @@ class _EventStreamTally:
|
|||
"Bedrock answered the stream with HTTP 200 but its body decoded to no events "
|
||||
f"({detail}, first bytes={self.head!r})"
|
||||
),
|
||||
headers=request_id_headers,
|
||||
)
|
||||
if not self.recognized_events:
|
||||
return BedrockError(
|
||||
|
|
@ -464,10 +467,12 @@ class _EventStreamTally:
|
|||
f"event type (event types={sorted(self.unrecognized_event_types)}, {detail}, "
|
||||
f"first payload={self.unrecognized_head!r})"
|
||||
),
|
||||
headers=request_id_headers,
|
||||
)
|
||||
return BedrockError(
|
||||
status_code=502,
|
||||
message=f"Bedrock stream ended with {undecoded} undecoded bytes after {self.events} events ({detail})",
|
||||
headers=request_id_headers,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -837,11 +842,12 @@ class AWSEventStreamDecoder:
|
|||
|
||||
event_stream_buffer: Final = EventStreamBuffer()
|
||||
tally: Final = _EventStreamTally(get_bedrock_stream_event_statuses())
|
||||
request_id_headers: Final = provider_request_id_headers(response_headers)
|
||||
for chunk in iterator:
|
||||
event_stream_buffer.add_data(chunk)
|
||||
tally.add_chunk(chunk)
|
||||
for event in event_stream_buffer:
|
||||
message = self._decode_event(event, tally)
|
||||
message = self._decode_event(event, tally, request_id_headers)
|
||||
if message:
|
||||
# sse_event = ServerSentEvent(data=message, event="completion")
|
||||
_data = json.loads(message)
|
||||
|
|
@ -858,11 +864,12 @@ class AWSEventStreamDecoder:
|
|||
|
||||
event_stream_buffer: Final = EventStreamBuffer()
|
||||
tally: Final = _EventStreamTally(get_bedrock_stream_event_statuses())
|
||||
request_id_headers: Final = provider_request_id_headers(response_headers)
|
||||
async for chunk in iterator:
|
||||
event_stream_buffer.add_data(chunk)
|
||||
tally.add_chunk(chunk)
|
||||
for event in event_stream_buffer:
|
||||
message = self._decode_event(event, tally)
|
||||
message = self._decode_event(event, tally, request_id_headers)
|
||||
if message:
|
||||
_data = json.loads(message)
|
||||
yield self.chunk_parser(chunk_data=_data)
|
||||
|
|
@ -882,27 +889,34 @@ class AWSEventStreamDecoder:
|
|||
)
|
||||
return response_stream_shape
|
||||
|
||||
def _decode_event(self, event: "EventStreamMessage", tally: _EventStreamTally) -> str | None:
|
||||
def _decode_event(
|
||||
self, event: "EventStreamMessage", tally: _EventStreamTally, request_id_headers: httpx.Headers | None
|
||||
) -> str | None:
|
||||
response_stream_shape: Final = self._response_stream_shape()
|
||||
response_dict: Final = bedrock_event_stream_response(event)
|
||||
tally.add_event(event, response_dict["headers"])
|
||||
return self._parse_message_from_response(response_dict, response_stream_shape)
|
||||
return self._parse_message_from_response(response_dict, response_stream_shape, request_id_headers)
|
||||
|
||||
def _parse_message_from_event(self, event: "EventStreamMessage") -> str | None:
|
||||
response_stream_shape: Final = self._response_stream_shape()
|
||||
return self._parse_message_from_response(bedrock_event_stream_response(event), response_stream_shape)
|
||||
|
||||
def _parse_message_from_response(
|
||||
self, response_dict: BedrockEventStreamResponseDict, response_stream_shape: "Shape"
|
||||
self,
|
||||
response_dict: BedrockEventStreamResponseDict,
|
||||
response_stream_shape: "Shape",
|
||||
request_id_headers: httpx.Headers | None = None,
|
||||
) -> str | None:
|
||||
parsed_response: Final = self.parser.parse(response_dict, response_stream_shape)
|
||||
|
||||
if response_dict["status_code"] != 200:
|
||||
raise build_bedrock_stream_error(response_dict, response_stream_shape)
|
||||
raise build_bedrock_stream_error(response_dict, response_stream_shape, headers=request_id_headers)
|
||||
event_type: Final = bedrock_event_stream_header(response_dict["headers"], ":event-type")
|
||||
event_error_status: Final = bedrock_stream_event_error_status(event_type)
|
||||
if event_type is not None and event_error_status is not None:
|
||||
raise build_bedrock_stream_event_error(event_type, event_error_status, response_dict["body"])
|
||||
raise build_bedrock_stream_event_error(
|
||||
event_type, event_error_status, response_dict["body"], headers=request_id_headers
|
||||
)
|
||||
if "chunk" in parsed_response:
|
||||
chunk = parsed_response.get("chunk")
|
||||
if not chunk:
|
||||
|
|
|
|||
|
|
@ -1913,13 +1913,18 @@ def bedrock_event_stream_header(headers: Mapping[str, object], name: str) -> str
|
|||
return value if isinstance(value, str) else None
|
||||
|
||||
|
||||
def build_bedrock_stream_event_error(event_type: str, status_code: int, body: bytes) -> BedrockError:
|
||||
return BedrockError(status_code=status_code, message=f"{event_type} {body.decode(errors='replace')}")
|
||||
def build_bedrock_stream_event_error(
|
||||
event_type: str, status_code: int, body: bytes, headers: httpx.Headers | None = None
|
||||
) -> BedrockError:
|
||||
return BedrockError(
|
||||
status_code=status_code, message=f"{event_type} {body.decode(errors='replace')}", headers=headers
|
||||
)
|
||||
|
||||
|
||||
def build_bedrock_stream_error(
|
||||
response_dict: BedrockEventStreamResponseDict,
|
||||
response_stream_shape: Shape | None,
|
||||
headers: httpx.Headers | None = None,
|
||||
) -> BedrockError:
|
||||
"""Build a BedrockError for a non-200 event-stream error event.
|
||||
|
||||
|
|
@ -1929,12 +1934,14 @@ def build_bedrock_stream_error(
|
|||
"""
|
||||
exception_type: Final = bedrock_event_stream_header(response_dict["headers"], ":exception-type")
|
||||
if exception_type is None:
|
||||
return BedrockError(status_code=response_dict["status_code"], message=response_dict["body"].decode())
|
||||
return BedrockError(
|
||||
status_code=response_dict["status_code"], message=response_dict["body"].decode(), headers=headers
|
||||
)
|
||||
|
||||
member: Final = _structure_members(response_stream_shape).get(exception_type)
|
||||
modeled_status: Final = None if member is None else _modeled_error_status(member)
|
||||
status_code: Final = response_dict["status_code"] if modeled_status is None else modeled_status
|
||||
return build_bedrock_stream_event_error(exception_type, status_code, response_dict["body"])
|
||||
return build_bedrock_stream_event_error(exception_type, status_code, response_dict["body"], headers=headers)
|
||||
|
||||
|
||||
class BedrockEventStreamDecoderBase:
|
||||
|
|
|
|||
|
|
@ -31,6 +31,7 @@ from litellm.litellm_core_utils.core_helpers import process_response_headers
|
|||
from litellm.litellm_core_utils.internal_call_metadata import without_budget_reservation
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.litellm_core_utils.llm_response_utils.get_api_base import get_api_base
|
||||
from litellm.litellm_core_utils.llm_response_utils.get_headers import provider_request_id_headers
|
||||
from litellm.litellm_core_utils.llm_response_utils.response_metadata import (
|
||||
update_response_metadata,
|
||||
)
|
||||
|
|
@ -262,7 +263,9 @@ def stream_error_status_and_message(error_obj: object) -> tuple[int, str]:
|
|||
return _status_code_for_error_fields(error_type, error_code), message
|
||||
|
||||
|
||||
def _map_stream_error_to_exception(error_obj: object, model: str, custom_llm_provider: str) -> Exception:
|
||||
def _map_stream_error_to_exception(
|
||||
error_obj: object, model: str, custom_llm_provider: str, headers: httpx.Headers | None = None
|
||||
) -> Exception:
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
|
||||
error_message, error_type, error_code = _error_event_fields(error_obj)
|
||||
|
|
@ -272,6 +275,7 @@ def _map_stream_error_to_exception(error_obj: object, model: str, custom_llm_pro
|
|||
status_code=status_code,
|
||||
message=f"Error code: {status_code} - {{'error': {error_body}}}",
|
||||
body=error_body,
|
||||
headers=headers,
|
||||
)
|
||||
try:
|
||||
return litellm.exception_type(
|
||||
|
|
@ -667,7 +671,12 @@ class BaseResponsesAPIStreamingIterator:
|
|||
)
|
||||
|
||||
def _map_error_event_exception(self, error_obj: object) -> Exception:
|
||||
return _map_stream_error_to_exception(error_obj, self.model or "", self.custom_llm_provider or "")
|
||||
return _map_stream_error_to_exception(
|
||||
error_obj,
|
||||
self.model or "",
|
||||
self.custom_llm_provider or "",
|
||||
headers=provider_request_id_headers(self._raw_response_headers),
|
||||
)
|
||||
|
||||
def _maybe_raise_for_error_event(self, result: object) -> None:
|
||||
chunk_type: Final = getattr(result, "type", None)
|
||||
|
|
@ -683,11 +692,17 @@ class BaseResponsesAPIStreamingIterator:
|
|||
mapped_exception: Final = self._map_error_event_exception(error_obj)
|
||||
if not _mid_stream_fallback_eligible(mapped_exception):
|
||||
raise mapped_exception
|
||||
mapped_status: Final = getattr(mapped_exception, "status_code", None)
|
||||
raise MidStreamFallbackError(
|
||||
message=str(mapped_exception),
|
||||
model=self.model or "",
|
||||
llm_provider=self.custom_llm_provider or "",
|
||||
original_exception=mapped_exception,
|
||||
response=httpx.Response(
|
||||
status_code=int(mapped_status) if mapped_status is not None else 503,
|
||||
headers=provider_request_id_headers(self._raw_response_headers),
|
||||
request=httpx.Request(method="POST", url=f"https://{self.custom_llm_provider or ''}.com/v1/"),
|
||||
),
|
||||
generated_content=self._generated_content,
|
||||
is_pre_first_chunk=not self._yielded_first_chunk,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1,13 +1,208 @@
|
|||
import asyncio
|
||||
import base64
|
||||
import json
|
||||
import uuid
|
||||
from typing import Final
|
||||
from collections.abc import Mapping
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from types import MappingProxyType
|
||||
from typing import Final, Literal, NamedTuple
|
||||
from urllib.parse import unquote
|
||||
|
||||
from integration._support.client import Gateway, eventually, object_value
|
||||
import anthropic
|
||||
import httpx
|
||||
import openai
|
||||
import pytest
|
||||
from integration._support.client import Gateway, Scenario, eventually, object_value
|
||||
from integration._support.database import read_rows
|
||||
from integration._support.responses_stream import (
|
||||
created,
|
||||
delta,
|
||||
error_event,
|
||||
failed,
|
||||
frame,
|
||||
healthy_stream,
|
||||
openai_error_event,
|
||||
spend_row,
|
||||
)
|
||||
from integration._support.upstream import _aws_event_frame, aws_event_stream_frame
|
||||
from integration._support.wire import Reply, Request, wire_server
|
||||
from pydantic import JsonValue
|
||||
|
||||
import litellm
|
||||
from litellm.exceptions import MidStreamFallbackError
|
||||
from litellm.litellm_core_utils.litellm_logging import _get_provider_request_id
|
||||
|
||||
_MODEL: Final = "bedrock/converse/anthropic.claude-sonnet-4-5-20250929-v1:0"
|
||||
_BEDROCK_RESPONSES_MODEL: Final = "bedrock/us.openai.gpt-5.6-sol"
|
||||
_TOKEN: Final = "synthetic-bedrock-bearer"
|
||||
_SHAPES: Final = ("nested-error", "top-level-error", "response.failed", "error-before-output")
|
||||
_ERROR_CODES: Final = ("server_error", "rate_limit_exceeded", "invalid_request_error", "content_policy_violation")
|
||||
_INVOKE_MODEL_ID: Final = "anthropic.claude-3-haiku-20240307-v1:0"
|
||||
_BEDROCK_AWS: Final[Mapping[str, JsonValue]] = MappingProxyType(
|
||||
{
|
||||
"aws_access_key_id": "AKIASCRIPTEDPROVIDER",
|
||||
"aws_secret_access_key": "scripted-secret",
|
||||
"aws_region_name": "us-east-1",
|
||||
}
|
||||
)
|
||||
_EVENT_STREAM: Final = "application/vnd.amazon.eventstream"
|
||||
_BEDROCK_STREAM_TARGETS: Final = frozenset(
|
||||
(
|
||||
"/model/anthropic.claude-sonnet-4-5-20250929-v1:0/converse-stream",
|
||||
f"/model/{_INVOKE_MODEL_ID}/invoke-with-response-stream",
|
||||
)
|
||||
)
|
||||
|
||||
BedrockFailure = Literal["exception-event", "exception-message", "no-known-event", "truncated-frame"]
|
||||
|
||||
|
||||
class _FailureShape(NamedTuple):
|
||||
bedrock_status: int
|
||||
converse_class: str
|
||||
after_output: tuple[bool, ...]
|
||||
|
||||
|
||||
BedrockEndpoint = Literal["chat", "responses", "messages", "messages-invoke"]
|
||||
_BEDROCK_ENDPOINTS: Final[tuple[BedrockEndpoint, ...]] = ("chat", "responses", "messages", "messages-invoke")
|
||||
_BEDROCK_FAILURES: Final[Mapping[BedrockFailure, _FailureShape]] = MappingProxyType(
|
||||
{
|
||||
"exception-event": _FailureShape(500, "ServiceUnavailableError", (False, True)),
|
||||
"exception-message": _FailureShape(429, "RateLimitError", (False, True)),
|
||||
"no-known-event": _FailureShape(502, "BadGatewayError", (False,)),
|
||||
"truncated-frame": _FailureShape(502, "BadGatewayError", (True,)),
|
||||
}
|
||||
)
|
||||
_BEDROCK_CELLS: Final = tuple(
|
||||
pytest.param(
|
||||
endpoint, failure, after_output, id=f"{endpoint}-{failure}-{'after' if after_output else 'before'}-output"
|
||||
)
|
||||
for endpoint in _BEDROCK_ENDPOINTS
|
||||
for failure, shape in _BEDROCK_FAILURES.items()
|
||||
for after_output in shape.after_output
|
||||
)
|
||||
|
||||
|
||||
def _stream(identity: str, shape: str, error_code: str, message: str) -> tuple[bytes, ...]:
|
||||
opening: Final = (frame(created(identity)), frame(delta(identity, "partial output")))
|
||||
nested: Final = frame(error_event({"type": error_code, "code": error_code, "message": message}))
|
||||
match shape:
|
||||
case "nested-error":
|
||||
return (*opening, nested)
|
||||
case "top-level-error":
|
||||
return (*opening, frame(openai_error_event(error_code, message)))
|
||||
case "response.failed":
|
||||
return (*opening, frame(failed(identity, error_code, message)))
|
||||
case "error-before-output":
|
||||
return (nested,)
|
||||
raise AssertionError(shape)
|
||||
|
||||
|
||||
def _bedrock_frame(event_type: str, payload: dict[str, JsonValue]) -> bytes:
|
||||
return _aws_event_frame(event_type, payload, "sc", "u")
|
||||
|
||||
|
||||
def _invoke_chunk(event: dict[str, JsonValue]) -> bytes:
|
||||
return _bedrock_frame("chunk", {"bytes": base64.b64encode(json.dumps(event).encode()).decode()})
|
||||
|
||||
|
||||
def _bedrock_output(invoke: bool, text: str = "partial output") -> tuple[bytes, ...]:
|
||||
if not invoke:
|
||||
return (
|
||||
_bedrock_frame("messageStart", {"role": "assistant"}),
|
||||
_bedrock_frame("contentBlockDelta", {"delta": {"text": text}, "contentBlockIndex": 0}),
|
||||
)
|
||||
started: Final[dict[str, JsonValue]] = {
|
||||
"id": f"msg_{uuid.uuid4().hex}",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"content": [],
|
||||
"model": _INVOKE_MODEL_ID,
|
||||
"stop_reason": None,
|
||||
"usage": {"input_tokens": 3, "output_tokens": 1},
|
||||
}
|
||||
return (
|
||||
_invoke_chunk({"type": "message_start", "message": started}),
|
||||
_invoke_chunk({"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}),
|
||||
_invoke_chunk({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": text}}),
|
||||
)
|
||||
|
||||
|
||||
def _bedrock_failure_frame(failure: BedrockFailure, message: str) -> bytes:
|
||||
match failure:
|
||||
case "exception-event":
|
||||
return _bedrock_frame("internalServerException", {"message": message})
|
||||
case "exception-message":
|
||||
headers: Final = {
|
||||
":exception-type": "throttlingException",
|
||||
":content-type": "application/json",
|
||||
":message-type": "exception",
|
||||
}
|
||||
return aws_event_stream_frame(headers, json.dumps({"message": message}).encode())
|
||||
case "no-known-event":
|
||||
return _bedrock_frame("somethingBedrockAddedLater", {"message": message})
|
||||
case "truncated-frame":
|
||||
return _bedrock_frame("contentBlockDelta", {"delta": {"text": message}, "contentBlockIndex": 0})[:-7]
|
||||
|
||||
|
||||
def _bedrock_failure_text(failure: BedrockFailure, message: str) -> str:
|
||||
return "undecoded bytes after" if failure == "truncated-frame" else message
|
||||
|
||||
|
||||
def _bedrock_stream_request(
|
||||
endpoint: BedrockEndpoint, model: str, prompt: str, *, stream: bool = True
|
||||
) -> tuple[str, dict[str, JsonValue]]:
|
||||
extra: Final[dict[str, JsonValue]] = {"stream": stream, "num_retries": 0, "cache": {"no-cache": True}}
|
||||
turn: Final[list[JsonValue]] = [{"role": "user", "content": prompt}]
|
||||
match endpoint:
|
||||
case "chat":
|
||||
return "/v1/chat/completions", {"model": model, "messages": turn, **extra}
|
||||
case "responses":
|
||||
return "/v1/responses", {"model": model, "input": prompt, **extra}
|
||||
case "messages" | "messages-invoke":
|
||||
return "/v1/messages", {"model": model, "messages": turn, "max_tokens": 16, **extra}
|
||||
|
||||
|
||||
def _failure_information(call_id: str) -> dict[str, JsonValue]:
|
||||
row: Final = spend_row(call_id)
|
||||
assert row["status"] == "failure", row
|
||||
metadata: Final = row["metadata"]
|
||||
parsed: Final = json.loads(metadata) if isinstance(metadata, str) else object_value(metadata)
|
||||
return object_value(parsed["error_information"])
|
||||
|
||||
|
||||
def _streamed_response(gateway: Gateway, model: str, prompt: str) -> httpx.Response:
|
||||
response: Final = gateway.request(
|
||||
"POST",
|
||||
"/v1/responses",
|
||||
{"model": model, "input": prompt, "stream": True, "num_retries": 0, "cache": {"no-cache": True}},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
return response
|
||||
|
||||
|
||||
async def _consume_sdk_stream(asynchronous: bool, api_base: str, identity: str) -> None:
|
||||
if asynchronous:
|
||||
async_stream: Final = await litellm.aresponses(
|
||||
model="openai/gpt-6",
|
||||
input=identity,
|
||||
stream=True,
|
||||
api_base=api_base,
|
||||
api_key="synthetic-provider-key",
|
||||
num_retries=0,
|
||||
)
|
||||
async for _ in async_stream:
|
||||
pass
|
||||
return
|
||||
tuple(
|
||||
litellm.responses(
|
||||
model="openai/gpt-6",
|
||||
input=identity,
|
||||
stream=True,
|
||||
api_base=api_base,
|
||||
api_key="synthetic-provider-key",
|
||||
num_retries=0,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def test_bedrock_500_keeps_amzn_request_id_on_error_headers_and_failure_log(gateway: Gateway) -> None:
|
||||
|
|
@ -54,3 +249,556 @@ def test_bedrock_500_keeps_amzn_request_id_on_error_headers_and_failure_log(gate
|
|||
parsed: Final = json.loads(metadata) if isinstance(metadata, str) else object_value(metadata)
|
||||
error_information: Final = object_value(parsed["error_information"])
|
||||
assert error_information["error_provider_request_id"] == amzn_request_id, error_information
|
||||
|
||||
|
||||
@pytest.mark.parametrize("shape", _SHAPES)
|
||||
@pytest.mark.parametrize("error_code", _ERROR_CODES)
|
||||
def test_responses_stream_failure_log_retains_provider_request_id(
|
||||
gateway: Gateway, shape: str, error_code: str
|
||||
) -> None:
|
||||
identity: Final = f"resp_{uuid.uuid4().hex}"
|
||||
provider_request_id: Final = str(uuid.uuid4())
|
||||
message: Final = f"stream failure {identity}"
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
assert request.method == "POST" and request.target == "/responses", request.target
|
||||
assert json.loads(request.body)["stream"] is True
|
||||
return Reply(
|
||||
headers={"x-amzn-RequestId": provider_request_id},
|
||||
content_type="text/event-stream",
|
||||
chunks=_stream(identity, shape, error_code, message),
|
||||
)
|
||||
|
||||
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model="openai/gpt-6", api_base=wire.url, num_retries=0)
|
||||
response: Final = _streamed_response(gateway, model, identity)
|
||||
assert message in response.text, response.text
|
||||
assert ("partial output" in response.text) is (shape != "error-before-output"), response.text
|
||||
assert response.headers.get("llm_provider-x-amzn-requestid") == provider_request_id, dict(response.headers)
|
||||
assert len(wire.drain()) == 1
|
||||
information: Final = _failure_information(response.headers["x-litellm-call-id"])
|
||||
assert information["error_provider_request_id"] == provider_request_id, information
|
||||
|
||||
|
||||
@pytest.mark.parametrize("shape", _SHAPES)
|
||||
def test_responses_stream_failure_without_provider_request_id_logs_none(gateway: Gateway, shape: str) -> None:
|
||||
identity: Final = f"resp_{uuid.uuid4().hex}"
|
||||
message: Final = f"stream failure {identity}"
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
assert request.method == "POST" and request.target == "/responses", request.target
|
||||
return Reply(content_type="text/event-stream", chunks=_stream(identity, shape, "server_error", message))
|
||||
|
||||
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model="openai/gpt-6", api_base=wire.url, num_retries=0)
|
||||
response: Final = _streamed_response(gateway, model, identity)
|
||||
assert message in response.text, response.text
|
||||
assert "llm_provider-x-amzn-requestid" not in response.headers, dict(response.headers)
|
||||
assert len(wire.drain()) == 1
|
||||
information: Final = _failure_information(response.headers["x-litellm-call-id"])
|
||||
assert information["error_class"] == "InternalServerError", information
|
||||
assert information["error_provider_request_id"] is None, information
|
||||
|
||||
|
||||
@pytest.mark.parametrize("shape", ("nested-error", "response.failed"))
|
||||
def test_bedrock_responses_stream_server_error_logs_amzn_request_id(gateway: Gateway, shape: str) -> None:
|
||||
identity: Final = f"resp_{uuid.uuid4().hex}"
|
||||
amzn_request_id: Final = str(uuid.uuid4())
|
||||
message: Final = f"bedrock stream failure {identity}"
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
assert request.method == "POST" and request.target == "/openai/v1/responses", request.target
|
||||
assert request.headers["authorization"] == f"Bearer {_TOKEN}", request.headers
|
||||
assert json.loads(request.body)["stream"] is True
|
||||
return Reply(
|
||||
headers={"x-amzn-RequestId": amzn_request_id},
|
||||
content_type="text/event-stream",
|
||||
chunks=_stream(identity, shape, "server_error", message),
|
||||
)
|
||||
|
||||
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(
|
||||
model=_BEDROCK_RESPONSES_MODEL,
|
||||
api_key=_TOKEN,
|
||||
api_base=wire.url,
|
||||
aws_region_name="us-east-1",
|
||||
num_retries=0,
|
||||
)
|
||||
response: Final = _streamed_response(gateway, model, identity)
|
||||
assert "partial output" in response.text and message in response.text, response.text
|
||||
assert response.headers.get("llm_provider-x-amzn-requestid") == amzn_request_id, dict(response.headers)
|
||||
assert len(wire.drain()) == 1
|
||||
information: Final = _failure_information(response.headers["x-litellm-call-id"])
|
||||
assert information["llm_provider"] == "bedrock", information
|
||||
assert information["error_provider_request_id"] == amzn_request_id, information
|
||||
|
||||
|
||||
def test_responses_stream_failure_log_reads_openai_x_request_id(gateway: Gateway) -> None:
|
||||
identity: Final = f"resp_{uuid.uuid4().hex}"
|
||||
provider_request_id: Final = f"req_{uuid.uuid4().hex}"
|
||||
message: Final = f"stream failure {identity}"
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
assert request.method == "POST" and request.target == "/responses", request.target
|
||||
return Reply(
|
||||
headers={"x-request-id": provider_request_id},
|
||||
content_type="text/event-stream",
|
||||
chunks=_stream(identity, "top-level-error", "server_error", message),
|
||||
)
|
||||
|
||||
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model="openai/gpt-6", api_base=wire.url, num_retries=0)
|
||||
response: Final = _streamed_response(gateway, model, identity)
|
||||
assert message in response.text, response.text
|
||||
assert len(wire.drain()) == 1
|
||||
information: Final = _failure_information(response.headers["x-litellm-call-id"])
|
||||
assert information["error_provider_request_id"] == provider_request_id, information
|
||||
|
||||
|
||||
@pytest.mark.parametrize("asynchronous", (False, True), ids=("sync", "async"))
|
||||
@pytest.mark.parametrize("event_type", ("error", "response.failed"))
|
||||
async def test_responses_sdk_stream_error_retains_provider_request_id(asynchronous: bool, event_type: str) -> None:
|
||||
identity: Final = f"resp_{uuid.uuid4().hex}"
|
||||
provider_request_id: Final = str(uuid.uuid4())
|
||||
shape: Final = "nested-error" if event_type == "error" else "response.failed"
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
assert request.method == "POST" and request.target == "/responses", request.target
|
||||
return Reply(
|
||||
headers={"x-amzn-RequestId": provider_request_id, "retry-after": "7"},
|
||||
content_type="text/event-stream",
|
||||
chunks=_stream(identity, shape, "server_error", identity),
|
||||
)
|
||||
|
||||
with wire_server(respond) as wire:
|
||||
with pytest.raises(MidStreamFallbackError) as raised:
|
||||
await _consume_sdk_stream(asynchronous, wire.url, identity)
|
||||
assert raised.value.status_code == 500
|
||||
assert raised.value.response.status_code == 500
|
||||
assert str(raised.value.response.request.url) == "https://openai.com/v1/"
|
||||
assert raised.value.generated_content == "partial output"
|
||||
assert dict(raised.value.response.headers) == {"x-amzn-requestid": provider_request_id}
|
||||
assert _get_provider_request_id(raised.value) == provider_request_id
|
||||
assert raised.value.original_exception is not None
|
||||
assert _get_provider_request_id(raised.value.original_exception) == provider_request_id
|
||||
assert len(wire.drain()) == 1
|
||||
|
||||
|
||||
def _concurrent_reply(request: Request, failing: tuple[str, ...], recovery: str) -> Reply:
|
||||
prompt: Final = json.loads(request.body)["input"]
|
||||
assert prompt in (*failing, recovery)
|
||||
chunks: Final = (
|
||||
healthy_stream(f"resp_{prompt}", "recovered")
|
||||
if prompt == recovery
|
||||
else _stream(f"resp_{prompt}", "nested-error", "server_error", prompt)
|
||||
)
|
||||
return Reply(content_type="text/event-stream", headers={"x-amzn-requestid": prompt}, chunks=chunks)
|
||||
|
||||
|
||||
def test_concurrent_responses_failures_keep_distinct_request_ids_and_recover(gateway: Gateway) -> None:
|
||||
identity: Final = uuid.uuid4().hex
|
||||
prompts: Final = tuple(f"{identity}-{index}" for index in range(4))
|
||||
recovery: Final = f"{identity}-recovery"
|
||||
|
||||
with (
|
||||
wire_server(lambda request: _concurrent_reply(request, prompts, recovery)) as wire,
|
||||
gateway.scenario() as scenario,
|
||||
):
|
||||
model: Final = scenario.model(model="openai/gpt-6", api_base=wire.url, num_retries=0)
|
||||
|
||||
def call(prompt: str) -> str:
|
||||
response: Final = _streamed_response(gateway, model, prompt)
|
||||
assert (prompt if prompt != recovery else "recovered") in response.text, response.text
|
||||
return response.headers["x-litellm-call-id"]
|
||||
|
||||
with ThreadPoolExecutor(max_workers=4) as workers:
|
||||
call_ids: Final = tuple(workers.map(call, prompts))
|
||||
recovered_call_id: Final = call(recovery)
|
||||
assert len(wire.drain()) == 5
|
||||
for prompt, call_id in zip(prompts, call_ids, strict=True):
|
||||
assert _failure_information(call_id)["error_provider_request_id"] == prompt
|
||||
assert spend_row(recovered_call_id)["status"] == "success"
|
||||
|
||||
|
||||
def _bedrock_logged_error(endpoint: BedrockEndpoint, failure: BedrockFailure) -> tuple[str, str]:
|
||||
shape: Final = _BEDROCK_FAILURES[failure]
|
||||
if endpoint == "messages-invoke":
|
||||
return "BedrockError", str(shape.bedrock_status)
|
||||
return shape.converse_class, "503" if failure == "exception-event" else str(shape.bedrock_status)
|
||||
|
||||
|
||||
def _bedrock_client_status(endpoint: BedrockEndpoint, failure: BedrockFailure, after_output: bool) -> int:
|
||||
logged_code: Final = int(_bedrock_logged_error(endpoint, failure)[1])
|
||||
if after_output or endpoint in ("responses", "messages"):
|
||||
return 200
|
||||
if endpoint == "messages-invoke" and failure == "no-known-event":
|
||||
return 200
|
||||
return logged_code
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("endpoint", "failure", "after_output"), _BEDROCK_CELLS)
|
||||
def test_bedrock_event_stream_failure_log_retains_amzn_request_id(
|
||||
gateway: Gateway, endpoint: BedrockEndpoint, failure: BedrockFailure, after_output: bool
|
||||
) -> None:
|
||||
amzn_request_id: Final = str(uuid.uuid4())
|
||||
message: Final = f"bedrock stream failure {amzn_request_id}"
|
||||
invoke: Final = endpoint == "messages-invoke"
|
||||
frames: Final = (*(_bedrock_output(invoke) if after_output else ()), _bedrock_failure_frame(failure, message))
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
assert request.method == "POST" and unquote(request.target) in _BEDROCK_STREAM_TARGETS, request.target
|
||||
return Reply(
|
||||
headers={"x-amzn-RequestId": amzn_request_id, "retry-after": "7"},
|
||||
content_type=_EVENT_STREAM,
|
||||
body=b"".join(frames),
|
||||
)
|
||||
|
||||
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(
|
||||
model=f"bedrock/invoke/{_INVOKE_MODEL_ID}" if invoke else _MODEL,
|
||||
api_base=wire.url,
|
||||
num_retries=0,
|
||||
**_BEDROCK_AWS,
|
||||
)
|
||||
path, body = _bedrock_stream_request(endpoint, model, message)
|
||||
with gateway.client.stream(
|
||||
"POST", path, json=body, headers={"Authorization": f"Bearer {gateway.key}"}
|
||||
) as response:
|
||||
text: Final = "\n".join(line for line in response.iter_lines() if line)
|
||||
status: Final = response.status_code
|
||||
response_headers: Final = dict(response.headers)
|
||||
assert status == _bedrock_client_status(endpoint, failure, after_output), text
|
||||
assert _bedrock_failure_text(failure, message) in text and ("partial output" in text) is after_output, text
|
||||
assert len(wire.drain()) == 1
|
||||
information: Final = _failure_information(response_headers["x-litellm-call-id"])
|
||||
assert (information["error_class"], information["error_code"]) == _bedrock_logged_error(endpoint, failure)
|
||||
assert information["error_provider_request_id"] == amzn_request_id, information
|
||||
if status != 200 and not invoke:
|
||||
assert response_headers.get("llm_provider-x-amzn-requestid") == amzn_request_id, response_headers
|
||||
assert "llm_provider-retry-after" not in response_headers, response_headers
|
||||
|
||||
|
||||
def _bedrock_model(scenario: Scenario, endpoint: BedrockEndpoint, api_base: str) -> str:
|
||||
return scenario.model(
|
||||
model=f"bedrock/invoke/{_INVOKE_MODEL_ID}" if endpoint == "messages-invoke" else _MODEL,
|
||||
api_base=api_base,
|
||||
num_retries=0,
|
||||
**_BEDROCK_AWS,
|
||||
)
|
||||
|
||||
|
||||
def _bedrock_stream(
|
||||
gateway: Gateway, endpoint: BedrockEndpoint, model: str, prompt: str
|
||||
) -> tuple[int, str, dict[str, str]]:
|
||||
path, body = _bedrock_stream_request(endpoint, model, prompt)
|
||||
with gateway.client.stream("POST", path, json=body, headers={"Authorization": f"Bearer {gateway.key}"}) as response:
|
||||
text: Final = "\n".join(line for line in response.iter_lines() if line)
|
||||
return response.status_code, text, dict(response.headers)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("after_output", (False, True), ids=("before-output", "after-output"))
|
||||
@pytest.mark.parametrize("endpoint", _BEDROCK_ENDPOINTS)
|
||||
def test_bedrock_event_stream_failure_without_amzn_request_id_logs_none(
|
||||
gateway: Gateway, endpoint: BedrockEndpoint, after_output: bool
|
||||
) -> None:
|
||||
message: Final = f"bedrock stream failure {uuid.uuid4()}"
|
||||
frames: Final = (
|
||||
*(_bedrock_output(endpoint == "messages-invoke") if after_output else ()),
|
||||
_bedrock_failure_frame("exception-event", message),
|
||||
)
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
assert unquote(request.target) in _BEDROCK_STREAM_TARGETS, request.target
|
||||
return Reply(headers={"retry-after": "7"}, content_type=_EVENT_STREAM, body=b"".join(frames))
|
||||
|
||||
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
||||
status, text, response_headers = _bedrock_stream(
|
||||
gateway, endpoint, _bedrock_model(scenario, endpoint, wire.url), message
|
||||
)
|
||||
assert status == _bedrock_client_status(endpoint, "exception-event", after_output), text
|
||||
assert message in text, text
|
||||
assert not any("amzn" in name for name in response_headers), response_headers
|
||||
assert len(wire.drain()) == 1
|
||||
information: Final = _failure_information(response_headers["x-litellm-call-id"])
|
||||
assert (information["error_class"], information["error_code"]) == _bedrock_logged_error(
|
||||
endpoint, "exception-event"
|
||||
)
|
||||
assert information["error_provider_request_id"] is None, information
|
||||
|
||||
|
||||
@pytest.mark.parametrize("endpoint", ("responses", "messages", "messages-invoke"))
|
||||
def test_bedrock_non_stream_500_logs_amzn_request_id_on_every_endpoint(
|
||||
gateway: Gateway, endpoint: BedrockEndpoint
|
||||
) -> None:
|
||||
amzn_request_id: Final = str(uuid.uuid4())
|
||||
targets: Final = frozenset(
|
||||
(
|
||||
"/model/anthropic.claude-sonnet-4-5-20250929-v1:0/converse",
|
||||
f"/model/{_INVOKE_MODEL_ID}/invoke",
|
||||
)
|
||||
)
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
assert request.method == "POST" and unquote(request.target) in targets, request.target
|
||||
return Reply(
|
||||
status=500,
|
||||
headers={"x-amzn-RequestId": amzn_request_id},
|
||||
body=b'{"message":"synthetic bedrock failure"}',
|
||||
)
|
||||
|
||||
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
||||
path, body = _bedrock_stream_request(
|
||||
endpoint, _bedrock_model(scenario, endpoint, wire.url), amzn_request_id, stream=False
|
||||
)
|
||||
response: Final = gateway.request("POST", path, body)
|
||||
assert response.status_code == 503, response.text
|
||||
if endpoint == "responses":
|
||||
assert response.headers.get("llm_provider-x-amzn-requestid") == amzn_request_id, dict(response.headers)
|
||||
assert len(wire.drain()) == 1
|
||||
information: Final = _failure_information(response.headers["x-litellm-call-id"])
|
||||
assert information["error_provider_request_id"] == amzn_request_id, information
|
||||
|
||||
|
||||
@pytest.mark.parametrize("amzn_request_id", ("", "a" * 5000, "dup-1, dup-1"), ids=("empty", "5kb", "joined-duplicate"))
|
||||
def test_bedrock_stream_failure_logs_odd_request_ids_like_the_non_stream_error(
|
||||
gateway: Gateway, amzn_request_id: str
|
||||
) -> None:
|
||||
stream_body: Final = b"".join((*_bedrock_output(False), _bedrock_failure_frame("exception-event", "odd id")))
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
if unquote(request.target) in _BEDROCK_STREAM_TARGETS:
|
||||
return Reply(headers={"x-amzn-RequestId": amzn_request_id}, content_type=_EVENT_STREAM, body=stream_body)
|
||||
return Reply(status=500, headers={"x-amzn-RequestId": amzn_request_id}, body=b'{"message":"odd id"}')
|
||||
|
||||
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _bedrock_model(scenario, "chat", wire.url)
|
||||
status, text, stream_headers = _bedrock_stream(gateway, "chat", model, "odd id")
|
||||
path, body = _bedrock_stream_request("chat", model, "odd id", stream=False)
|
||||
non_stream: Final = gateway.request("POST", path, body)
|
||||
assert status == 200 and "odd id" in text, text
|
||||
assert non_stream.status_code == 503, non_stream.text
|
||||
assert len(wire.drain()) == 2
|
||||
streamed: Final = _failure_information(stream_headers["x-litellm-call-id"])
|
||||
unstreamed: Final = _failure_information(non_stream.headers["x-litellm-call-id"])
|
||||
assert streamed["error_provider_request_id"] == unstreamed["error_provider_request_id"], (streamed, unstreamed)
|
||||
|
||||
|
||||
_BURST_CALLS: Final = tuple(
|
||||
(endpoint, failure, after_output)
|
||||
for endpoint in _BEDROCK_ENDPOINTS
|
||||
for failure, shape in _BEDROCK_FAILURES.items()
|
||||
for after_output in shape.after_output
|
||||
)
|
||||
|
||||
|
||||
def _healthy_bedrock_stream(invoke: bool) -> bytes:
|
||||
closing: Final = (
|
||||
(
|
||||
_invoke_chunk({"type": "content_block_stop", "index": 0}),
|
||||
_invoke_chunk(
|
||||
{"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 1}}
|
||||
),
|
||||
_invoke_chunk({"type": "message_stop"}),
|
||||
)
|
||||
if invoke
|
||||
else (
|
||||
_bedrock_frame("contentBlockStop", {"contentBlockIndex": 0}),
|
||||
_bedrock_frame("messageStop", {"stopReason": "end_turn"}),
|
||||
_bedrock_frame("metadata", {"usage": {"inputTokens": 3, "outputTokens": 1, "totalTokens": 4}}),
|
||||
)
|
||||
)
|
||||
return b"".join((*_bedrock_output(invoke, "recovered"), *closing))
|
||||
|
||||
|
||||
def test_concurrent_bedrock_stream_failures_on_every_endpoint_keep_distinct_request_ids(gateway: Gateway) -> None:
|
||||
identity: Final = uuid.uuid4().hex
|
||||
failing: Final = {f"{identity}-{index:02d}": call for index, call in enumerate(_BURST_CALLS)}
|
||||
recoveries: Final = {f"{identity}-r{index}": endpoint for index, endpoint in enumerate(_BEDROCK_ENDPOINTS)}
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
(token,) = (token for token in (*failing, *recoveries) if token in request.body.decode())
|
||||
if token in recoveries:
|
||||
body = _healthy_bedrock_stream(recoveries[token] == "messages-invoke")
|
||||
else:
|
||||
endpoint, failure, after_output = failing[token]
|
||||
opening = _bedrock_output(endpoint == "messages-invoke") if after_output else ()
|
||||
body = b"".join((*opening, _bedrock_failure_frame(failure, token)))
|
||||
return Reply(headers={"x-amzn-RequestId": token, "retry-after": "7"}, content_type=_EVENT_STREAM, body=body)
|
||||
|
||||
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
||||
models: Final = {endpoint: _bedrock_model(scenario, endpoint, wire.url) for endpoint in _BEDROCK_ENDPOINTS}
|
||||
|
||||
def call(token: str, endpoint: BedrockEndpoint) -> tuple[int, str, dict[str, str]]:
|
||||
return _bedrock_stream(gateway, endpoint, models[endpoint], token)
|
||||
|
||||
with ThreadPoolExecutor(max_workers=8) as workers:
|
||||
outcomes: Final = dict(
|
||||
zip(failing, workers.map(call, failing, (endpoint for endpoint, _, _ in failing.values())), strict=True)
|
||||
)
|
||||
recovered: Final = {token: call(token, endpoint) for token, endpoint in recoveries.items()}
|
||||
assert len(wire.drain()) == len(failing) + len(recoveries)
|
||||
for token, (endpoint, failure, after_output) in failing.items():
|
||||
status, text, response_headers = outcomes[token]
|
||||
assert status == _bedrock_client_status(endpoint, failure, after_output), (token, text)
|
||||
assert _bedrock_failure_text(failure, token) in text, (token, text)
|
||||
information = _failure_information(response_headers["x-litellm-call-id"])
|
||||
assert information["error_provider_request_id"] == token, (token, information)
|
||||
for token, (status, text, response_headers) in recovered.items():
|
||||
assert status == 200 and "recovered" in text, (token, text)
|
||||
assert spend_row(response_headers["x-litellm-call-id"])["status"] == "success"
|
||||
|
||||
|
||||
def _sdk_stream_text(
|
||||
client_kind: str, asynchronous: bool, gateway: Gateway, model: str, prompt: str
|
||||
) -> tuple[str, str]:
|
||||
extra_body: Final[dict[str, JsonValue]] = {"num_retries": 0, "cache": {"no-cache": True}}
|
||||
turn: Final = [{"role": "user", "content": prompt}]
|
||||
|
||||
def open_stream(client: openai.OpenAI | openai.AsyncOpenAI | anthropic.Anthropic | anthropic.AsyncAnthropic):
|
||||
match client_kind:
|
||||
case "openai-chat":
|
||||
return client.chat.completions.with_raw_response.create(
|
||||
model=model, messages=turn, stream=True, extra_body=extra_body
|
||||
)
|
||||
case "openai-responses":
|
||||
return client.responses.with_raw_response.create(
|
||||
model=model, input=prompt, stream=True, extra_body=extra_body
|
||||
)
|
||||
case "anthropic-messages":
|
||||
return client.messages.with_raw_response.create(
|
||||
model=model, messages=turn, max_tokens=16, stream=True, extra_body=extra_body
|
||||
)
|
||||
raise AssertionError(client_kind)
|
||||
|
||||
seen: Final[list[str]] = []
|
||||
base_url: Final = str(gateway.client.base_url).rstrip("/")
|
||||
if not asynchronous:
|
||||
with httpx.Client(timeout=15, trust_env=False) as transport:
|
||||
client = (
|
||||
anthropic.Anthropic(api_key=gateway.key, base_url=base_url, max_retries=0, http_client=transport)
|
||||
if client_kind == "anthropic-messages"
|
||||
else openai.OpenAI(api_key=gateway.key, base_url=f"{base_url}/v1", max_retries=0, http_client=transport)
|
||||
)
|
||||
raw = open_stream(client)
|
||||
try:
|
||||
seen.extend(repr(event) for event in raw.parse())
|
||||
except (openai.APIError, anthropic.APIError) as error:
|
||||
seen.append(str(error))
|
||||
return raw.headers["x-litellm-call-id"], "\n".join(seen)
|
||||
|
||||
async def run() -> str:
|
||||
async with httpx.AsyncClient(timeout=15, trust_env=False) as transport:
|
||||
client = (
|
||||
anthropic.AsyncAnthropic(api_key=gateway.key, base_url=base_url, max_retries=0, http_client=transport)
|
||||
if client_kind == "anthropic-messages"
|
||||
else openai.AsyncOpenAI(
|
||||
api_key=gateway.key, base_url=f"{base_url}/v1", max_retries=0, http_client=transport
|
||||
)
|
||||
)
|
||||
raw = await open_stream(client)
|
||||
try:
|
||||
async for event in raw.parse():
|
||||
seen.append(repr(event))
|
||||
except (openai.APIError, anthropic.APIError) as error:
|
||||
seen.append(str(error))
|
||||
return raw.headers["x-litellm-call-id"]
|
||||
|
||||
return asyncio.run(run()), "\n".join(seen)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("asynchronous", (False, True), ids=("sync", "async"))
|
||||
@pytest.mark.parametrize("client_kind", ("openai-chat", "openai-responses", "anthropic-messages"))
|
||||
def test_sdk_clients_streaming_bedrock_failure_through_proxy_log_amzn_request_id(
|
||||
gateway: Gateway, client_kind: str, asynchronous: bool
|
||||
) -> None:
|
||||
amzn_request_id: Final = str(uuid.uuid4())
|
||||
message: Final = f"bedrock stream failure {amzn_request_id}"
|
||||
frames: Final = (*_bedrock_output(False), _bedrock_failure_frame("exception-event", message))
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
assert unquote(request.target) in _BEDROCK_STREAM_TARGETS, request.target
|
||||
return Reply(headers={"x-amzn-RequestId": amzn_request_id}, content_type=_EVENT_STREAM, body=b"".join(frames))
|
||||
|
||||
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
||||
call_id, seen = _sdk_stream_text(
|
||||
client_kind, asynchronous, gateway, _bedrock_model(scenario, "chat", wire.url), message
|
||||
)
|
||||
assert "partial output" in seen and message in seen, seen
|
||||
assert len(wire.drain()) == 1
|
||||
information: Final = _failure_information(call_id)
|
||||
assert information["error_provider_request_id"] == amzn_request_id, information
|
||||
|
||||
|
||||
async def _consume_bedrock_sdk_stream(api: Literal["chat", "responses"], asynchronous: bool, api_base: str) -> str:
|
||||
common: Final[dict[str, JsonValue]] = {
|
||||
"model": _MODEL,
|
||||
"stream": True,
|
||||
"api_base": api_base,
|
||||
"num_retries": 0,
|
||||
**_BEDROCK_AWS,
|
||||
}
|
||||
turn: Final[list[JsonValue]] = [{"role": "user", "content": "probe"}]
|
||||
match api, asynchronous:
|
||||
case "chat", True:
|
||||
return "".join([repr(chunk) async for chunk in await litellm.acompletion(messages=turn, **common)])
|
||||
case "chat", False:
|
||||
return "".join(repr(chunk) for chunk in litellm.completion(messages=turn, **common))
|
||||
case "responses", True:
|
||||
return "".join([repr(event) async for event in await litellm.aresponses(input="probe", **common)])
|
||||
case "responses", False:
|
||||
return "".join(repr(event) for event in litellm.responses(input="probe", **common))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("asynchronous", (False, True), ids=("sync", "async"))
|
||||
@pytest.mark.parametrize("api", ("chat", "responses"))
|
||||
async def test_bedrock_sdk_stream_exception_carries_amzn_request_id(
|
||||
api: Literal["chat", "responses"], asynchronous: bool
|
||||
) -> None:
|
||||
amzn_request_id: Final = str(uuid.uuid4())
|
||||
message: Final = f"bedrock stream failure {uuid.uuid4()}"
|
||||
frames: Final = (*_bedrock_output(False), _bedrock_failure_frame("exception-event", message))
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
assert unquote(request.target) in _BEDROCK_STREAM_TARGETS, request.target
|
||||
return Reply(
|
||||
headers={"x-amzn-RequestId": amzn_request_id, "retry-after": "7"},
|
||||
content_type=_EVENT_STREAM,
|
||||
body=b"".join(frames),
|
||||
)
|
||||
|
||||
with wire_server(respond) as wire:
|
||||
with pytest.raises(MidStreamFallbackError) as raised:
|
||||
await _consume_bedrock_sdk_stream(api, asynchronous, wire.url)
|
||||
assert len(wire.drain()) == 1
|
||||
provider_error: Final = raised.value.original_exception
|
||||
assert isinstance(provider_error, litellm.ServiceUnavailableError), provider_error
|
||||
assert raised.value.generated_content == "partial output"
|
||||
assert message in str(provider_error)
|
||||
assert _get_provider_request_id(provider_error) == amzn_request_id
|
||||
assert "retry-after" not in provider_error.response.headers, provider_error.response.headers
|
||||
|
||||
|
||||
@pytest.mark.parametrize("api", ("chat", "responses"))
|
||||
@pytest.mark.parametrize("healthy", (True, False), ids=("healthy", "failing"))
|
||||
async def test_sync_bedrock_sdk_stream_with_non_ascii_request_id_header_keeps_streaming(
|
||||
api: Literal["chat", "responses"], healthy: bool
|
||||
) -> None:
|
||||
body: Final = (
|
||||
_healthy_bedrock_stream(False)
|
||||
if healthy
|
||||
else b"".join((*_bedrock_output(False), _bedrock_failure_frame("exception-event", "non-ascii id failure")))
|
||||
)
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
assert unquote(request.target) in _BEDROCK_STREAM_TARGETS, request.target
|
||||
return Reply(headers={"x-amzn-RequestId": "req-\u00fc-1"}, content_type=_EVENT_STREAM, body=body)
|
||||
|
||||
with wire_server(respond) as wire:
|
||||
if healthy:
|
||||
assert "recovered" in await _consume_bedrock_sdk_stream(api, False, wire.url)
|
||||
else:
|
||||
with pytest.raises(MidStreamFallbackError) as raised:
|
||||
await _consume_bedrock_sdk_stream(api, False, wire.url)
|
||||
assert "non-ascii id failure" in str(raised.value.original_exception)
|
||||
assert _get_provider_request_id(raised.value.original_exception) is None
|
||||
assert len(wire.drain()) == 1
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ from litellm.litellm_core_utils.llm_response_utils.get_headers import (
|
|||
_get_llm_provider_headers,
|
||||
get_response_headers,
|
||||
get_provider_request_id,
|
||||
provider_request_id_headers,
|
||||
)
|
||||
from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
|
||||
|
||||
|
|
@ -89,6 +90,31 @@ def test_invalid_provider_request_ids_remain_absent(headers: object) -> None:
|
|||
assert get_provider_request_id(headers) is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("headers", "expected"),
|
||||
(
|
||||
(
|
||||
{"X-Amzn-RequestId": "amzn-1", "retry-after": "7", "content-type": "text/event-stream"},
|
||||
{"x-amzn-requestid": "amzn-1"},
|
||||
),
|
||||
(
|
||||
{"x-request-id": "req_1", "cf-ray": "ray-1", "x-ratelimit-remaining-requests": "0"},
|
||||
{"x-request-id": "req_1", "cf-ray": "ray-1"},
|
||||
),
|
||||
({"retry-after": "7"}, None),
|
||||
({"x-request-id": 7, "request-id": None, "x-amzn-requestid": "amzn-2"}, {"x-amzn-requestid": "amzn-2"}),
|
||||
({"x-amzn-requestid": "amzn-\u00fc", "x-request-id": "req_2"}, {"x-request-id": "req_2"}),
|
||||
({}, None),
|
||||
(None, None),
|
||||
),
|
||||
)
|
||||
def test_provider_request_id_headers_keeps_only_request_id_headers(
|
||||
headers: dict[str, object] | None, expected: dict[str, str] | None
|
||||
) -> None:
|
||||
result = provider_request_id_headers(headers)
|
||||
assert (None if result is None else dict(result)) == expected
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _vcr_outcome_gate(request, vcr):
|
||||
install_live_call_probe(request, vcr)
|
||||
|
|
|
|||
|
|
@ -16,6 +16,7 @@ import litellm
|
|||
from litellm import AmazonInvokeConfig
|
||||
from litellm.exceptions import MidStreamFallbackError
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.litellm_core_utils.litellm_logging import _get_provider_request_id
|
||||
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
|
||||
from litellm.llms.bedrock.chat.invoke_handler import (
|
||||
AmazonOpenAICompatibleStreamDecoder,
|
||||
|
|
@ -864,14 +865,16 @@ async def test_async_invoke_streaming_non_200_forwards_bedrock_response_headers(
|
|||
|
||||
|
||||
def _event_stream_frame(event_type: str, payload: bytes) -> bytes:
|
||||
return _typed_event_stream_frame(
|
||||
{":event-type": event_type, ":content-type": "application/json", ":message-type": "event"}, payload
|
||||
)
|
||||
|
||||
|
||||
def _typed_event_stream_frame(frame_headers: Mapping[str, 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(":event-type", event_type)
|
||||
+ header(":content-type", "application/json")
|
||||
+ header(":message-type", "event")
|
||||
)
|
||||
headers: Final = b"".join(header(name, value) for name, value in frame_headers.items())
|
||||
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))
|
||||
|
|
@ -1367,6 +1370,87 @@ def test_iter_bytes_keeps_a_stream_whose_unknown_event_sits_beside_known_frames(
|
|||
assert "stop" in finish_reasons
|
||||
|
||||
|
||||
_STREAM_FAILURE_BODIES: Final = (
|
||||
pytest.param(
|
||||
_converse_event_frame("internalServerException", {"message": _UPSTREAM_REJECTION}), 500, id="exception-event"
|
||||
),
|
||||
pytest.param(
|
||||
_typed_event_stream_frame(
|
||||
{
|
||||
":exception-type": "throttlingException",
|
||||
":content-type": "application/json",
|
||||
":message-type": "exception",
|
||||
},
|
||||
json.dumps({"message": _UPSTREAM_REJECTION}).encode(),
|
||||
),
|
||||
429,
|
||||
id="exception-message",
|
||||
),
|
||||
pytest.param(_unknown_event_frame(), 502, id="no-known-event"),
|
||||
pytest.param(b"", 502, id="no-events"),
|
||||
pytest.param(_CONVERSE_CONTENT_FRAMES[0] + _truncated_frame(), 502, id="undecoded-tail"),
|
||||
)
|
||||
|
||||
|
||||
def _stream_headers(request_id: str | None) -> httpx.Headers:
|
||||
request_id_header: Final = {} if request_id is None else {"x-amzn-RequestId": request_id}
|
||||
return httpx.Headers(
|
||||
{"content-type": "application/vnd.amazon.eventstream", "retry-after": "7", **request_id_header}
|
||||
)
|
||||
|
||||
|
||||
def _assert_carries_only_the_request_id(error: BedrockError, status: int, request_id: str | None) -> None:
|
||||
assert error.status_code == status
|
||||
assert _get_provider_request_id(error) == request_id
|
||||
if request_id is None:
|
||||
assert error.headers is None
|
||||
return
|
||||
assert dict(error.headers or {}) == {"x-amzn-requestid": request_id}
|
||||
assert dict(error.response.headers) == {"x-amzn-requestid": request_id}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("request_id", ("req-stream-1", None))
|
||||
@pytest.mark.parametrize(("body", "status"), _STREAM_FAILURE_BODIES)
|
||||
def test_iter_bytes_stream_error_carries_only_the_provider_request_id_header(
|
||||
body: bytes, status: int, request_id: str | None
|
||||
) -> None:
|
||||
decoder: Final = AWSEventStreamDecoder(model="us.moonshotai.kimi-k3")
|
||||
|
||||
with pytest.raises(BedrockError) as exc_info:
|
||||
list(decoder.iter_bytes(iter([body]), response_headers=_stream_headers(request_id)))
|
||||
|
||||
_assert_carries_only_the_request_id(exc_info.value, status, request_id)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(("body", "status"), _STREAM_FAILURE_BODIES)
|
||||
async def test_aiter_bytes_stream_error_carries_only_the_provider_request_id_header(body: bytes, status: int) -> 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=_stream_headers("req-stream-2"))]
|
||||
|
||||
_assert_carries_only_the_request_id(exc_info.value, status, "req-stream-2")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("request_id", ("req-\u00fc-1", "req-\udcfc-1"), ids=("latin-1", "surrogate-escaped"))
|
||||
def test_iter_bytes_streams_and_fails_normally_with_a_non_ascii_provider_request_id(request_id: str) -> None:
|
||||
decoder: Final = AWSEventStreamDecoder(model="us.moonshotai.kimi-k3")
|
||||
headers: Final = {"content-type": "application/vnd.amazon.eventstream", "x-amzn-RequestId": request_id}
|
||||
|
||||
chunks: Final = list(decoder.iter_bytes(iter(_CONVERSE_CONTENT_FRAMES), response_headers=headers))
|
||||
with pytest.raises(BedrockError) as exc_info:
|
||||
list(decoder.iter_bytes(iter([_unknown_event_frame()]), response_headers=headers))
|
||||
|
||||
texts: Final = [chunk.choices[0].delta.content for chunk in chunks if isinstance(chunk, ModelResponseStream)]
|
||||
assert "".join(text or "" for text in texts) == "hi"
|
||||
assert exc_info.value.status_code == 502
|
||||
assert exc_info.value.headers is None
|
||||
|
||||
|
||||
def _assert_exception_event_surfaced_with_its_modeled_status(error: BaseException, event_type: str) -> None:
|
||||
assert not isinstance(error, litellm.BadGatewayError)
|
||||
assert getattr(error, "status_code", None) == _modeled_status(event_type)
|
||||
|
|
|
|||
|
|
@ -20,14 +20,18 @@ Pydantic ValidationError (previously typed as Optional[str]).
|
|||
"""
|
||||
|
||||
import json
|
||||
from datetime import datetime
|
||||
from importlib import import_module
|
||||
from typing import Final
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.exceptions import MidStreamFallbackError
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.litellm_core_utils.litellm_logging import _get_provider_request_id
|
||||
from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig
|
||||
from litellm.responses.streaming_iterator import (
|
||||
_ERROR_CODE_HTTP_STATUS,
|
||||
|
|
@ -64,6 +68,52 @@ def _make_error_chunk(error_type: str, code: str, message: str = "err") -> Error
|
|||
return ErrorEvent(type=ResponsesAPIStreamEvents.ERROR, sequence_number=0, error=error_obj)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("provider", ("openai", "azure", "bedrock"))
|
||||
@pytest.mark.parametrize("header", ("x-amzn-RequestId", "x-request-id", None))
|
||||
@pytest.mark.parametrize(
|
||||
"code,status", (("server_error", 500), ("rate_limit_exceeded", 429), ("invalid_request_error", 400))
|
||||
)
|
||||
def test_stream_error_preserves_provider_request_id(provider: str, header: str | None, code: str, status: int) -> None:
|
||||
stream_headers: Final = {"retry-after": "7", "content-type": "text/event-stream"}
|
||||
response: Final = httpx.Response(
|
||||
200, headers={**stream_headers, header: "provider-request-id"} if header else stream_headers
|
||||
)
|
||||
logging: Final = LiteLLMLoggingObj(
|
||||
model="test-model",
|
||||
messages=[],
|
||||
stream=True,
|
||||
call_type="responses",
|
||||
start_time=datetime(2026, 1, 1),
|
||||
litellm_call_id="test-call",
|
||||
function_id="test-function",
|
||||
)
|
||||
iterator: Final = BaseResponsesAPIStreamingIterator(
|
||||
response=response,
|
||||
model="test-model",
|
||||
responses_api_provider_config=None,
|
||||
logging_obj=logging,
|
||||
custom_llm_provider=provider,
|
||||
)
|
||||
expected_id: Final = "provider-request-id" if header else None
|
||||
expected_headers: Final = {header.lower(): "provider-request-id"} if header else {}
|
||||
mapped: Final = iterator._map_error_event_exception({"type": code, "code": code, "message": "failure"})
|
||||
assert _get_provider_request_id(mapped) == expected_id
|
||||
assert dict(getattr(mapped, "litellm_response_headers", None) or {}) == expected_headers
|
||||
expected_status: Final = 503 if provider == "bedrock" and status == 500 else status
|
||||
with pytest.raises((litellm.BadRequestError, MidStreamFallbackError)) as raised:
|
||||
iterator._maybe_raise_for_error_event(_make_error_chunk(code, code, "failure"))
|
||||
assert raised.value.status_code == expected_status
|
||||
assert _get_provider_request_id(raised.value) == expected_id
|
||||
assert response.status_code == 200
|
||||
if status in (429, 500):
|
||||
assert isinstance(raised.value, MidStreamFallbackError)
|
||||
assert dict(raised.value.response.headers) == expected_headers
|
||||
assert raised.value.response.status_code == expected_status
|
||||
assert str(raised.value.response.request.url) == f"https://{provider}.com/v1/"
|
||||
assert raised.value.generated_content == ""
|
||||
assert raised.value.is_pre_first_chunk is True
|
||||
|
||||
|
||||
def test_maybe_raise_for_error_event_wraps_unknown_error_in_mid_stream_fallback():
|
||||
iterator = _make_iterator()
|
||||
chunk = _make_error_chunk("server_error", "internal_error", "something went wrong")
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue