diff --git a/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py b/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py index 317b95f18f0..ded3ce2415f 100644 --- a/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py +++ b/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py @@ -1,15 +1,18 @@ +import base64 import copy import json import struct from binascii import crc32 -from collections.abc import AsyncGenerator, AsyncIterable, Mapping, Sequence +from collections.abc import AsyncGenerator, AsyncIterable, Callable, Mapping, Sequence from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, Final, Optional, Protocol, TypeAlias +from types import MappingProxyType +from typing import TYPE_CHECKING, Any, Final, Optional, Protocol, TypeAlias, TypeVar from pydantic import TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger +from litellm.llms.anthropic.chat.guardrail_translation.handler import AnthropicMessagesHandler from litellm.llms.base_llm.guardrail_translation.base_translation import ( BaseTranslation, StreamingScanKey, @@ -19,6 +22,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import ( effective_skip_system_message_for_guardrail, effective_skip_tool_message_for_guardrail, ) +from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: @@ -31,6 +35,7 @@ if TYPE_CHECKING: from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth _CONVERSE_ACTIONS: Final = frozenset({"converse", "converse-stream"}) +_INVOKE_STREAM_ACTION: Final = "invoke-with-response-stream" _EVENT_STREAM_CONTENT_TYPE: Final = "vnd.amazon.eventstream" _EVENT_STREAM_MEDIA_TYPE: Final = "application/vnd.amazon.eventstream" @@ -40,6 +45,10 @@ def _is_converse_endpoint(endpoint: str) -> bool: return bool(parts) and parts[-1] in _CONVERSE_ACTIONS +def _is_invoke_stream_endpoint(endpoint: str) -> bool: + return endpoint.rstrip("/").split("/")[-1] == _INVOKE_STREAM_ACTION + + def _generic_passthrough_handler() -> "PassThroughEndpointHandler": """ Fallback for non-Converse Bedrock routes (e.g. invoke). The generic @@ -360,6 +369,17 @@ def _converse_exception(message: str, status_code: int) -> bytes: ) +def _base64_json(event: Mapping[str, object]) -> str: + return base64.b64encode(json.dumps(event, separators=(",", ":")).encode()).decode() + + +def _invoke_chunk(event: Mapping[str, object]) -> bytes: + return _event_stream_message( + _string_headers({":event-type": "chunk", ":content-type": "application/json", ":message-type": "event"}), + json.dumps({"bytes": _base64_json(event)}, separators=(",", ":")).encode(), + ) + + @dataclass(slots=True) class ConverseStreamFrame: """One frame of a Converse event stream as the proxy relays it. @@ -400,10 +420,52 @@ class ConverseStreamFrame: ) +def _anthropic_sse(event: Mapping[str, object] | None) -> bytes: + if event is None: + return b"" + return f"event: {event['type']}\ndata: {json.dumps(event, separators=(',', ':'))}\n\n".encode() + + +@dataclass(slots=True) +class InvokeStreamFrame: + """One frame of an invoke-with-response-stream as the proxy relays it. + + ``event`` is the Anthropic Messages stream event a ``chunk`` frame carries base64 encoded in the ``bytes`` of its + ``payload``. Exception frames and bytes that do not decode keep it None and are relayed as received. + ``ends_stream`` marks the empty item the framer adds once the upstream has ended, and ``guardrailed_event`` is set + when a guardrail rewrote the event.""" + + raw: bytes + payload: Mapping[str, object] | None = None + event: Mapping[str, object] | None = None + guardrailed_event: Mapping[str, object] | None = None + ends_stream: bool = False + + def sse(self) -> bytes: + """The event as a ``/v1/messages`` stream carries it, the shape its guardrail translation scans.""" + return _anthropic_sse(self.guardrailed_event if self.guardrailed_event is not None else self.event) + + def received_sse(self) -> bytes: + """The event as received from the upstream, before any guardrail rewrote it.""" + return _anthropic_sse(self.event) + + def encode(self) -> bytes: + if self.guardrailed_event is None or self.payload is None: + return self.raw + headers_length: Final = _unpack_uint32(self.raw[4:8]) + return _event_stream_message( + self.raw[12 : 12 + headers_length], + json.dumps({**self.payload, "bytes": _base64_json(self.guardrailed_event)}, separators=(",", ":")).encode(), + ) + + _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object]) +_Frame: Final = TypeVar("_Frame", ConverseStreamFrame, InvokeStreamFrame) -def _decode_converse_frame(raw: bytes) -> ConverseStreamFrame: +def _event_message(raw: bytes) -> tuple[str, bytes] | None: + """The event type and payload of one event-stream frame, or None for an exception frame or bytes that do not + decode.""" from botocore.eventstream import EventStreamBuffer, ParserError buffer: Final = EventStreamBuffer() @@ -414,20 +476,47 @@ def _decode_converse_frame(raw: bytes) -> ConverseStreamFrame: verbose_proxy_logger.debug( "BedrockPassthroughGuardrailHandler: could not decode event-stream frame, relaying it unmodified: %s", e ) - return ConverseStreamFrame(raw=raw) + return None # botocore-stubs types EventStreamMessage.headers as int; at runtime it is the decoded header dict. headers: Final = _JSON_OBJECT.validate_python(message.headers) event_type: Final = headers.get(":event-type") if headers.get(":message-type") != "event" or not isinstance(event_type, str): + return None + return event_type, message.payload + + +def _decode_converse_frame(raw: bytes) -> ConverseStreamFrame: + message: Final = _event_message(raw) + if message is None: return ConverseStreamFrame(raw=raw) + event_type, payload_bytes = message try: - payload: Final = _JSON_OBJECT.validate_python(json.loads(message.payload)) + payload: Final = _JSON_OBJECT.validate_python(json.loads(payload_bytes)) except (json.JSONDecodeError, UnicodeDecodeError, ValidationError): return ConverseStreamFrame(raw=raw, event_type=event_type) return ConverseStreamFrame(raw=raw, event_type=event_type, payload=payload) -async def _converse_stream_frames(chunks: AsyncIterable[bytes]) -> AsyncGenerator[ConverseStreamFrame, None]: +def _decode_invoke_frame(raw: bytes) -> InvokeStreamFrame: + message: Final = _event_message(raw) + if message is None or message[0] != "chunk": + return InvokeStreamFrame(raw=raw) + try: + payload: Final = _JSON_OBJECT.validate_python(json.loads(message[1])) + encoded: Final = payload.get("bytes") + event: Final = ( + _JSON_OBJECT.validate_python(json.loads(base64.b64decode(encoded))) if isinstance(encoded, str) else None + ) + except (ValueError, ValidationError): + return InvokeStreamFrame(raw=raw) + if event is None or not isinstance(event.get("type"), str): + return InvokeStreamFrame(raw=raw) + return InvokeStreamFrame(raw=raw, payload=payload, event=event) + + +async def _event_stream_frames( + chunks: AsyncIterable[bytes], frame_type: type[_Frame], decode: Callable[[bytes], _Frame] +) -> AsyncGenerator[_Frame, None]: """Cut the upstream bytes into frames as soon as each frame has fully arrived, then add an empty item that ends the stream once the upstream has ended. @@ -437,7 +526,7 @@ async def _converse_stream_frames(chunks: AsyncIterable[bytes]) -> AsyncGenerato framing = True # rebind-ok: cleared once a frame length is unusable async for chunk in chunks: if not framing: - yield ConverseStreamFrame(raw=chunk) + yield frame_type(raw=chunk) continue pending += chunk while len(pending) >= 4: @@ -447,14 +536,14 @@ async def _converse_stream_frames(chunks: AsyncIterable[bytes]) -> AsyncGenerato break if len(pending) < total_length: break - yield _decode_converse_frame(pending[:total_length]) + yield decode(pending[:total_length]) pending = pending[total_length:] if not framing and pending: - yield ConverseStreamFrame(raw=pending) + yield frame_type(raw=pending) pending = b"" if pending: - yield ConverseStreamFrame(raw=pending) - yield ConverseStreamFrame(raw=b"", ends_stream=True) + yield frame_type(raw=pending) + yield frame_type(raw=b"", ends_stream=True) def _converse_frames(items: Sequence[object]) -> tuple[ConverseStreamFrame, ...]: @@ -543,6 +632,90 @@ def _stream_marked_request_data( return {**request_data, "stream": True, "response": response} +def _invoke_frames(items: Sequence[object]) -> tuple[InvokeStreamFrame, ...]: + return tuple(item for item in items if isinstance(item, InvokeStreamFrame)) + + +def _sse_events(sse: bytes) -> tuple[Mapping[str, object], ...]: + return tuple( + _JSON_OBJECT.validate_python(json.loads(line[len("data:") :])) + for line in sse.decode().splitlines() + if line.startswith("data:") + ) + + +def _invoke_message_stopped(frames: Sequence[InvokeStreamFrame]) -> bool: + """Whether a ``message_delta`` with a ``stop_reason`` came, as a ``/v1/messages`` stream ends.""" + return any( + frame.event is not None + and frame.event.get("type") == "message_delta" + and isinstance(delta := frame.event.get("delta"), Mapping) + and _JSON_OBJECT.validate_python(delta).get("stop_reason") is not None + for frame in frames + ) + + +def _invoke_stream_has_ended(frames: Sequence[InvokeStreamFrame]) -> bool: + """Ended on a ``message_delta`` with a ``stop_reason`` or once the upstream has ended.""" + return any(frame.ends_stream for frame in frames) or _invoke_message_stopped(frames) + + +_UPSTREAM_ENDED_SSE: Final = _anthropic_sse( + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 0}, + } +) + + +def _invoke_scan_sse(frames: Sequence[InvokeStreamFrame], events: Sequence[bytes]) -> tuple[bytes, ...]: + """The ``/v1/messages`` events to scan for these frames, closed with a ``message_delta`` when the upstream ended + without one so the translation scans them as a whole stream, tool calls included.""" + if any(frame.ends_stream for frame in frames) and not _invoke_message_stopped(frames): + return (*events, _UPSTREAM_ENDED_SSE) + return tuple(events) + + +async def _scan_invoke_stream( + frames: Sequence[InvokeStreamFrame], + guardrail_to_apply: "CustomGuardrail", + litellm_logging_obj: Optional["LiteLLMLoggingObj"], + user_api_key_dict: Optional["UserAPIKeyAuth"], + request_data: dict[str, object] | None, # mutable-ok: the request dict guardrails receive +) -> None: + """Scan the Anthropic events relayed so far as ``/v1/messages`` scans its stream, with the request marked as + streaming and those events as its ``responses``. Once the stream has ended, the rewrites are written back into the + frames, which reach the caller only when the stream is held until moderated.""" + received: Final = tuple(frame.sse() for frame in frames) + scan_events: Final = _invoke_scan_sse(frames, received) + scanned: Final[list[object]] = list(scan_events) # mutable-ok: the /v1/messages translation rewrites it in place + if request_data is not None: + request_data["stream"] = True # rebind-ok: marks the request guardrails receive as a stream + request_data["responses"] = scanned # rebind-ok: each scan sees every event relayed so far + await AnthropicMessagesHandler().process_output_streaming_response( # pyright: ignore[reportUnknownMemberType] # takes the untyped request dict + responses_so_far=scanned, + guardrail_to_apply=guardrail_to_apply, + litellm_logging_obj=litellm_logging_obj, + user_api_key_dict=user_api_key_dict, + request_data=request_data, + deliver_ended_stream_rewrites=_invoke_stream_has_ended(frames), + ) + for frame, before, after in zip(frames, received, scanned): + if isinstance(after, bytes) and after != before and (rewritten := _sse_events(after)): + frame.guardrailed_event = rewritten[0] + + +def _invoke_block_frames( + exc: "ModifyResponseException", stream_started: bool, frames: Sequence[InvokeStreamFrame] +) -> tuple[bytes, ...]: + """The ``/v1/messages`` block events for the stream relayed so far, each in its own ``chunk`` frame.""" + sse_chunks: Final = AnthropicMessagesHandler().build_block_sse_chunks( + exc, stream_started=stream_started, responses_so_far=[frame.sse() for frame in frames] + ) + return tuple(_invoke_chunk(event) for event in _sse_events(b"".join(sse_chunks))) + + def _write_back_stream_texts( frames: Sequence[ConverseStreamFrame], groups: _StreamTextGroups, rewritten_texts: Sequence[str] ) -> None: @@ -562,6 +735,113 @@ def _write_back_stream_texts( frame.guardrailed_payload = rewritten +_INVOKE_DELTA_TEXT_FIELDS: Final = MappingProxyType( + { + "text_delta": ("text", "text"), + "input_json_delta": ("tool", "partial_json"), + "thinking_delta": ("reasoning", "thinking"), + } +) + + +@dataclass(frozen=True, slots=True) +class _InvokeDeltaText: + """The guardrailable text one Anthropic event's delta carries, keyed by its content block and kind like the + Converse holders: text, tool input and thinking. Signatures and citations are never rewritten.""" + + group: _StreamGroupKey + delta: Mapping[str, object] + field: str + text: str + + +def _invoke_delta_text(event: Mapping[str, object] | None) -> _InvokeDeltaText | None: + if event is None: + return None + try: + delta: Final = _JSON_OBJECT.validate_python(event.get("delta")) + except ValidationError: + return None + kind: Final = delta.get("type") + field: Final = _INVOKE_DELTA_TEXT_FIELDS.get(kind) if isinstance(kind, str) else None + text: Final = delta.get(field[1]) if field is not None else None + if field is None or not isinstance(text, str): + return None + index: Final = event.get("index") + return _InvokeDeltaText((index if isinstance(index, int) else None, field[0]), delta, field[1], text) + + +async def _post_call_rewritten_texts( + groups: _StreamTextGroups, + proxy_logging_obj: "ProxyLogging", + user_api_key_dict: "UserAPIKeyAuth", + data: dict[str, object], # mutable-ok: the request dict post_call hooks receive +) -> Sequence[str] | None: + """Run the post_call guardrails on a buffered stream's texts as one Converse response with a content block per + group, and return each group's text as the guardrails left it, or None when the response came back in another + shape.""" + synthetic_response: Final[dict] = { + "output": { + "message": { + "role": "assistant", + "content": [{"text": text} for text in groups.texts()], + } + }, + "stopReason": "end_turn", + } + + processed: Final = await proxy_logging_obj.post_call_success_hook( + data=data, + user_api_key_dict=user_api_key_dict, + response=synthetic_response, + ) + + if not isinstance(processed, dict): + verbose_proxy_logger.debug( + "BedrockPassthroughGuardrailHandler: post_call_success_hook returned %s, leaving event stream unmodified", + type(processed).__name__, + ) + return None + + try: + processed_blocks: Final = processed["output"]["message"]["content"] + return [processed_blocks[i]["text"] for i in range(len(groups.groups))] + except (KeyError, IndexError, TypeError): + return None + + +async def _guardrail_buffered_invoke_stream( + body_bytes: bytes, + proxy_logging_obj: "ProxyLogging", + user_api_key_dict: "UserAPIKeyAuth", + data: dict[str, object], # mutable-ok: the request dict post_call hooks receive +) -> bytes: + """Scan a whole Claude invoke stream the way the buffered Converse path scans its stream, and write the guardrail + rewrites back into the ``chunk`` frames that carried each text.""" + + async def received() -> AsyncGenerator[bytes, None]: + yield body_bytes + + frames: Final = [frame async for frame in _event_stream_frames(received(), InvokeStreamFrame, _decode_invoke_frame)] + texts: Final = [_invoke_delta_text(frame.event) for frame in frames] + groups: Final = _group_stream_texts([[] if text is None else [(text.group, text.text)] for text in texts]) + if not groups.groups: + return body_bytes + rewritten_texts: Final = await _post_call_rewritten_texts(groups, proxy_logging_obj, user_api_key_dict, data) + if rewritten_texts is None: + return body_bytes + split: Final = _split_group_texts(groups, rewritten_texts) + for frame_idx, (frame, text) in enumerate(zip(frames, texts)): + if ( + text is not None + and frame.event is not None + and (new_text := split.get((frame_idx, 0))) is not None + and new_text != text.text + ): + frame.guardrailed_event = {**frame.event, "delta": {**text.delta, text.field: new_text}} + return b"".join(frame.encode() for frame in frames) + + class BedrockPassthroughGuardrailHandler(BaseTranslation): @staticmethod def is_event_stream_content_type(content_type: str) -> bool: @@ -572,16 +852,33 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation): return _EVENT_STREAM_MEDIA_TYPE @staticmethod - def event_stream_endpoint_is_de_anonymizable(endpoint: str) -> bool: - return _is_converse_endpoint(endpoint) + def event_stream_endpoint_is_de_anonymizable(endpoint: str, model: str, body: object) -> bool: + """The streams the buffered path can scan whole and rewrite, the same ones the relay scans.""" + return BedrockPassthroughGuardrailHandler.event_stream_endpoint_is_relayable(endpoint, model, body) + + @staticmethod + def event_stream_endpoint_is_relayable(endpoint: str, model: str, body: object) -> bool: + """Converse streams of every model, and invoke streams of Claude, whose chunks carry Anthropic Messages stream + events. Claude is told by its model id or, for ids that do not name it such as application inference profile + ARNs, by the ``anthropic_version`` Bedrock requires in every Claude Messages request body.""" + return _is_converse_endpoint(endpoint) or ( + _is_invoke_stream_endpoint(endpoint) + and ( + BaseAWSLLM.get_bedrock_invoke_provider(model) == "anthropic" + or (isinstance(body, Mapping) and "anthropic_version" in body) + ) + ) @staticmethod async def de_anonymize_event_stream( body_bytes: bytes, proxy_logging_obj: "ProxyLogging", user_api_key_dict: "UserAPIKeyAuth", - data: dict, + data: dict[str, object], # mutable-ok: the request dict post_call hooks receive ) -> bytes: + if _is_invoke_stream_endpoint(str(data.get("endpoint") or "")): + return await _guardrail_buffered_invoke_stream(body_bytes, proxy_logging_obj, user_api_key_dict, data) + from botocore.eventstream import EventStreamBuffer frames: Final[list[_StreamFrame]] = [] @@ -629,37 +926,13 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation): if not groups.groups: return body_bytes - synthetic_response: Final[dict] = { - "output": { - "message": { - "role": "assistant", - "content": [{"text": text} for text in groups.texts()], - } - }, - "stopReason": "end_turn", - } - - processed: Final = await proxy_logging_obj.post_call_success_hook( - data=data, - user_api_key_dict=user_api_key_dict, - response=synthetic_response, + de_anonymized_texts: Final = await _post_call_rewritten_texts( + groups, proxy_logging_obj, user_api_key_dict, data ) - - if not isinstance(processed, dict): - verbose_proxy_logger.debug( - "BedrockPassthroughGuardrailHandler: post_call_success_hook returned %s, " - "leaving event stream unmodified", - type(processed).__name__, - ) + if de_anonymized_texts is None: return body_bytes - try: - processed_blocks: Final = processed["output"]["message"]["content"] - de_anonymized_texts: Final = [processed_blocks[i]["text"] for i in range(len(groups.groups))] - except (KeyError, IndexError, TypeError): - return body_bytes - - new_text_map: Final = _split_group_texts(groups, de_anonymized_texts) # pyright: ignore[reportUnknownArgumentType] # the hook returns an untyped response + new_text_map: Final = _split_group_texts(groups, de_anonymized_texts) result_parts: Final[list[bytes]] = [] @@ -744,8 +1017,13 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation): user_api_key_dict: Optional["UserAPIKeyAuth"] = None, request_data: dict | None = None, ) -> object: - endpoint: Final = (request_data or {}).get("endpoint", "") - if endpoint and not _is_converse_endpoint(endpoint): + endpoint: Final[object] = (request_data or {}).get("endpoint", "") + if ( + isinstance(endpoint, str) + and endpoint + and not _is_converse_endpoint(endpoint) + and not _is_invoke_stream_endpoint(endpoint) + ): return await _generic_passthrough_handler().process_output_response( response=response, guardrail_to_apply=guardrail_to_apply, @@ -825,8 +1103,14 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation): the caller only when the stream is held until moderated. A guardrail that sets ``reads_stream_marks`` gets a copy of the request marked as chat completions marks a streamed answer; the others get it unchanged. Such a guardrail also gets the ended scan of a reply with no text, such as a call to a tool that takes no - input, with no texts to scan. + input, with no texts to scan. Invoke frames are scanned as Anthropic Messages events. """ + invoke_frames: Final = _invoke_frames(responses_so_far) + if invoke_frames: + await _scan_invoke_stream( + invoke_frames, guardrail_to_apply, litellm_logging_obj, user_api_key_dict, request_data + ) + return responses_so_far frames: Final = _converse_frames(responses_so_far) groups: Final = _group_stream_texts([frame.scanned_texts() for frame in frames]) ended: Final = self._check_streaming_has_ended(frames) @@ -862,6 +1146,11 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation): return _reads_stream_marks(guardrail) def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None: + invoke_frames: Final = _invoke_frames(responses_so_far) + if invoke_frames: + return AnthropicMessagesHandler().get_streaming_scan_key( + list(_invoke_scan_sse(invoke_frames, [frame.received_sse() for frame in invoke_frames])) + ) frames: Final = _converse_frames(responses_so_far) return StreamingScanKey( texts=_group_stream_texts([frame.received_texts() for frame in frames]).texts(), @@ -869,7 +1158,7 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation): ) def _check_streaming_has_ended(self, responses_so_far: Sequence[object]) -> bool: - return any( + return _invoke_stream_has_ended(_invoke_frames(responses_so_far)) or any( frame.event_type == "messageStop" or frame.ends_stream for frame in _converse_frames(responses_so_far) ) @@ -881,8 +1170,12 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation): ) -> Sequence[bytes]: """ End the message with the guardrail's message as a new text block and ``stopReason`` ``guardrail_intervened``, - closing any content block left open and starting the message when nothing was relayed yet. + closing any content block left open and starting the message when nothing was relayed yet. An invoke stream + gets the ``/v1/messages`` block events instead. """ + invoke_frames: Final = _invoke_frames(responses_so_far or ()) + if invoke_frames or _is_invoke_stream_endpoint(str(exc.request_data.get("endpoint") or "")): + return _invoke_block_frames(exc, stream_started, invoke_frames) frames: Final = _converse_frames(responses_so_far or ()) started: Final = frozenset( index @@ -911,19 +1204,23 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation): @staticmethod def owns_event_stream_item(item: object) -> bool: - return isinstance(item, ConverseStreamFrame) + return isinstance(item, (ConverseStreamFrame, InvokeStreamFrame)) @staticmethod - def event_stream_items(chunks: AsyncIterable[bytes]) -> AsyncGenerator[ConverseStreamFrame, None]: - return _converse_stream_frames(chunks) + def event_stream_items( + chunks: AsyncIterable[bytes], endpoint: str + ) -> AsyncGenerator[ConverseStreamFrame | InvokeStreamFrame, None]: + if _is_invoke_stream_endpoint(endpoint): + return _event_stream_frames(chunks, InvokeStreamFrame, _decode_invoke_frame) + return _event_stream_frames(chunks, ConverseStreamFrame, _decode_converse_frame) @staticmethod def encode_event_stream_item(item: object) -> bytes: - if isinstance(item, ConverseStreamFrame): + if isinstance(item, (ConverseStreamFrame, InvokeStreamFrame)): return item.encode() if isinstance(item, bytes): return item - raise TypeError(f"cannot relay {type(item).__name__} on a Converse event stream") + raise TypeError(f"cannot relay {type(item).__name__} on a Bedrock event stream") @staticmethod def event_stream_error_item(message: str, status_code: int) -> bytes: diff --git a/litellm/llms/pass_through/guardrail_translation/handler.py b/litellm/llms/pass_through/guardrail_translation/handler.py index d57175fa76e..885fc98ffea 100644 --- a/litellm/llms/pass_through/guardrail_translation/handler.py +++ b/litellm/llms/pass_through/guardrail_translation/handler.py @@ -212,7 +212,7 @@ class _EventStreamMediaTypeProvider(Protocol): @runtime_checkable class _EventStreamEndpointChecker(Protocol): - def event_stream_endpoint_is_de_anonymizable(self, endpoint: str) -> bool: ... + def event_stream_endpoint_is_de_anonymizable(self, endpoint: str, model: str, body: object) -> bool: ... _PROVIDER_HANDLERS: dict[str, type[BaseTranslation]] = {} @@ -357,13 +357,15 @@ class LlmPassthroughRouteHandler(BaseTranslation): return translation.build_block_sse_chunks(exc, stream_started=stream_started, responses_so_far=responses_so_far) @staticmethod - def event_stream_items(provider: object, chunks: AsyncIterable[bytes]) -> AsyncIterator[object] | None: - """The upstream event stream cut into the items the provider's guardrail translation scans, or None when the - provider has no streaming guardrail translation.""" - framer: Final[Callable[[AsyncIterable[bytes]], AsyncIterator[object]] | None] = getattr( + def event_stream_items( + provider: object, endpoint: object, chunks: AsyncIterable[bytes] + ) -> AsyncIterator[object] | None: + """The upstream event stream of ``endpoint`` cut into the items the provider's guardrail translation scans, or + None when the provider has no streaming guardrail translation.""" + framer: Final[Callable[[AsyncIterable[bytes], str], AsyncIterator[object]] | None] = getattr( _provider_handler(provider), "event_stream_items", None ) - return framer(chunks) if framer is not None else None + return framer(chunks, endpoint if isinstance(endpoint, str) else "") if framer is not None else None @staticmethod def encode_event_stream_item(provider: object, item: object) -> bytes: @@ -402,11 +404,26 @@ class LlmPassthroughRouteHandler(BaseTranslation): return getattr(handler_cls, "de_anonymize_event_stream", None) @staticmethod - def supports_event_stream_de_anonymization(provider: str | None, endpoint: str | None) -> bool: - handler_cls: Final = _get_provider_handlers().get(provider or "") + def supports_event_stream_relay(provider: object, endpoint: object, model: object, body: object) -> bool: + """Whether the provider's guardrail translation can scan this endpoint's event stream for ``model`` and the + request ``body`` frame by frame as it is relayed.""" + relayable: Final[Callable[[str, str, object], bool] | None] = getattr( + _provider_handler(provider), "event_stream_endpoint_is_relayable", None + ) + return relayable is not None and relayable( + endpoint if isinstance(endpoint, str) else "", model if isinstance(model, str) else "", body + ) + + @staticmethod + def supports_event_stream_de_anonymization(provider: object, endpoint: object, model: object, body: object) -> bool: + """Whether the provider's guardrail translation can scan this endpoint's whole event stream for ``model`` and + the request ``body`` and rewrite its frames.""" + handler_cls: Final = _provider_handler(provider) if not isinstance(handler_cls, _EventStreamEndpointChecker): return False - return handler_cls.event_stream_endpoint_is_de_anonymizable(endpoint or "") + return handler_cls.event_stream_endpoint_is_de_anonymizable( + endpoint if isinstance(endpoint, str) else "", model if isinstance(model, str) else "", body + ) @staticmethod async def de_anonymize_event_stream( diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 9020b89b766..74ee87c772b 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -2982,23 +2982,25 @@ class ProxyBaseLLMRequestProcessing: else: generator = response - if ( - self._has_post_call_guardrails_for_passthrough() - and self._passthrough_endpoint_has_stream_guardrail_handler() - ): - relayed_guardrails: Final = self._passthrough_stream_guardrails( + post_call_guarded: Final = self._has_post_call_guardrails_for_passthrough() + relayed_guardrails: Final = ( + self._passthrough_stream_guardrails( self.data # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # self.data is the untyped request dict ) - if relayed_guardrails is not None: - return await self._relay_guarded_passthrough_event_stream( - request=request, - request_data=self.data, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # self.data is the untyped request dict - generator=generator, # pyright: ignore[reportUnknownArgumentType] # raw upstream byte stream - guardrails=relayed_guardrails, - proxy_logging_obj=proxy_logging_obj, - user_api_key_dict=user_api_key_dict, - headers=streaming_headers, # pyright: ignore[reportUnknownArgumentType] # merged upstream headers - ) + if post_call_guarded and self._passthrough_stream_is_relayable() + else None + ) + if relayed_guardrails is not None: + return await self._relay_guarded_passthrough_event_stream( + request=request, + request_data=self.data, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # self.data is the untyped request dict + generator=generator, # pyright: ignore[reportUnknownArgumentType] # raw upstream byte stream + guardrails=relayed_guardrails, + proxy_logging_obj=proxy_logging_obj, + user_api_key_dict=user_api_key_dict, + headers=streaming_headers, # pyright: ignore[reportUnknownArgumentType] # merged upstream headers + ) + if post_call_guarded and self._passthrough_endpoint_has_stream_guardrail_handler(): body_bytes: Final = b"".join([chunk async for chunk in generator]) modified_bytes: Final = await self._handle_event_stream_allm_passthrough_route( body_bytes=body_bytes, @@ -3506,16 +3508,30 @@ class ProxyBaseLLMRequestProcessing: endpoints may have their stream buffered for post-call guardrails; every other endpoint must keep streaming so the response is not silently turned into a non-streaming body when no content modification would occur (e.g. - Bedrock invoke-with-response-stream, whose frames the Converse handler - leaves untouched). + Bedrock invoke-with-response-stream of a model other than Claude). """ from litellm.llms.pass_through.guardrail_translation.handler import ( LlmPassthroughRouteHandler, ) return LlmPassthroughRouteHandler.supports_event_stream_de_anonymization( - self.data.get("custom_llm_provider"), - self.data.get("endpoint"), + self.data.get("custom_llm_provider"), # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # self.data is the untyped request dict + self.data.get("endpoint"), # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # self.data is the untyped request dict + getattr(self.data.get("litellm_logging_obj"), "model", None), # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # self.data is the untyped request dict + self.data.get("data"), # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # self.data is the untyped request dict + ) + + def _passthrough_stream_is_relayable(self) -> bool: + """True when the passthrough provider's guardrail translation scans this endpoint's stream as it relays.""" + from litellm.llms.pass_through.guardrail_translation.handler import ( + LlmPassthroughRouteHandler, + ) + + return LlmPassthroughRouteHandler.supports_event_stream_relay( + self.data.get("custom_llm_provider"), # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # self.data is the untyped request dict + self.data.get("endpoint"), # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # self.data is the untyped request dict + getattr(self.data.get("litellm_logging_obj"), "model", None), # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # self.data is the untyped request dict + self.data.get("data"), # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # self.data is the untyped request dict ) @staticmethod @@ -3601,7 +3617,7 @@ class ProxyBaseLLMRequestProcessing: ) provider: Final = request_data.get("custom_llm_provider") - items: Final = LlmPassthroughRouteHandler.event_stream_items(provider, generator) + items: Final = LlmPassthroughRouteHandler.event_stream_items(provider, request_data.get("endpoint"), generator) if items is None: raise TypeError(f"provider {provider!r} has no event-stream framer") guarded: Final = proxy_logging_obj.async_post_call_event_stream_iterator_hook( diff --git a/tests/integration/observability/test_bedrock_invoke_stream_post_call_guardrail.py b/tests/integration/observability/test_bedrock_invoke_stream_post_call_guardrail.py new file mode 100644 index 00000000000..3ef73f613fe --- /dev/null +++ b/tests/integration/observability/test_bedrock_invoke_stream_post_call_guardrail.py @@ -0,0 +1,870 @@ +"""A post_call guardrail scans a Claude model's Bedrock invoke-with-response-stream passthrough as it relays (LIT-9618). + +The caller is boto3 ``InvokeModelWithResponseStream`` (or raw HTTP) against +``/bedrock/model//invoke-with-response-stream``. Each ``chunk`` frame carries one base64 Anthropic Messages +stream event, and a guardrail that scans responses through ``apply_guardrail`` runs on those events the way it runs +on ``/bedrock/model//converse-stream`` and ``/v1/messages``: frames reach the caller as the upstream sends +them, every fifth frame and the end of the stream are scanned, a block ends the stream with an AWS exception frame, +and a held stream gets its masks written back into re-encoded ``chunk`` frames. A guardrail the relay cannot run on +frames, one released on scan or run in parallel, gets the whole stream held and scanned once as converse-stream does. +Streams of other model families keep going out as received""" + +from __future__ import annotations + +import base64 +import json +import os +import threading +import uuid +from collections.abc import Iterator, Mapping, Sequence +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from functools import partial +from pathlib import Path +from typing import Final +from urllib.parse import unquote + +import boto3 +import httpx +import pytest +import yaml +from botocore.config import Config +from botocore.eventstream import EventStreamBuffer +from botocore.exceptions import ClientError, EventStreamError +from integration._support.budget import case_timeout_seconds, import_seconds +from integration._support.client import Gateway, JsonValue, Scenario, eventually +from integration._support.database import read_rows +from integration._support.process import owned_proxy_process +from integration._support.upstream import aws_event_stream_frame +from integration._support.wire import Reply, Request, Wire, wire_server + +CLAUDE_ID: Final = "anthropic.claude-sonnet-5-v1:0" +LLAMA_ID: Final = "meta.llama4-maverick-17b-instruct-v1:0" +PROFILE_ID: Final = "arn:aws:bedrock:us-east-1:000000000000:application-inference-profile/relayprofile" +CLAUDE: Final = "invoke-relay" +LLAMA: Final = "invoke-relay-llama" +PROFILE: Final = "invoke-relay-profile" +EVENT_STREAM: Final = "application/vnd.amazon.eventstream" +SAMPLED: Final = "invoke-sampled" +HELD: Final = "invoke-held" +END_ONLY: Final = "invoke-end-only" +CODE: Final = "invoke-code" +BEDROCK: Final = "invoke-bedrock" +BEDROCK_MESSAGE: Final = "invoke-bedrock-message" +UNGUARDED: Final = "invoke-unguarded" +OUTAGE: Final = "invoke-outage" +RELEASE_ON_SCAN: Final = "invoke-release-on-scan" +PARALLEL: Final = "invoke-parallel" +NOMA: Final = "invoke-noma" +NOMA_HELD: Final = "invoke-noma-held" +CODE_BLOCK_REASON: Final = "code says no" +BLOCK_REASON: Final = "Violated guardrail policy" +BEDROCK_BLOCK_MESSAGE: Final = "Sorry, the model cannot answer this question." +MASKED_WORD: Final = "Paris" +MASK: Final = "[CITY]" +BLOCKED_WORD: Final = "FORBIDDEN" + + +def chunk_frame(event: Mapping[str, JsonValue]) -> bytes: + encoded: Final = base64.b64encode(json.dumps(event, separators=(",", ":")).encode()).decode() + return aws_event_stream_frame( + {":event-type": "chunk", ":content-type": "application/json", ":message-type": "event"}, + json.dumps({"bytes": encoded, "p": "abcdefghij"}, separators=(",", ":")).encode(), + ) + + +def claude_frames(*deltas: str) -> tuple[bytes, ...]: + """A Claude invoke stream with one text content block, one ``text_delta`` event per string.""" + return ( + chunk_frame( + { + "type": "message_start", + "message": { + "id": "msg_bdrk_relay", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5", + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 9, "output_tokens": 1}, + }, + } + ), + chunk_frame({"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}), + *( + chunk_frame({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": delta}}) + for delta in deltas + ), + chunk_frame({"type": "content_block_stop", "index": 0}), + chunk_frame( + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 4}, + } + ), + chunk_frame( + { + "type": "message_stop", + "amazon-bedrock-invocationMetrics": { + "inputTokenCount": 9, + "outputTokenCount": 4, + "invocationLatency": 120, + "firstByteLatency": 40, + }, + } + ), + ) + + +def llama_frames(*generations: str) -> tuple[bytes, ...]: + return tuple( + chunk_frame({"generation": text, "prompt_token_count": None, "generation_token_count": 1, "stop_reason": None}) + for text in generations + ) + + +@dataclass(frozen=True, slots=True) +class Frame: + message_type: str + kind: str + payload: Mapping[str, JsonValue] + + +def _frame(headers: Mapping[str, str], payload: bytes) -> Frame: + body: Final = json.loads(payload) + if headers.get(":event-type") != "chunk": + return Frame(headers[":message-type"], headers.get(":event-type") or headers[":exception-type"], body) + event: Final = json.loads(base64.b64decode(body["bytes"])) + return Frame("event", str(event.get("type")), event) + + +def decode(raw: bytes) -> tuple[Frame, ...]: + buffer: Final = EventStreamBuffer() + buffer.add_data(raw) + return tuple(_frame(message.headers, message.payload) for message in buffer) + + +def streamed_text(frames: Sequence[Frame]) -> str: + return "".join( + str(delta["text"]) + for frame in frames + if frame.kind == "content_block_delta" and isinstance(delta := frame.payload["delta"], dict) + ) + + +def verdict(request: Request) -> Reply: + body: Final = json.loads(request.body) + texts: Final = [str(text) for text in body["texts"]] + if any(BLOCKED_WORD in text for text in texts) or BLOCKED_WORD in json.dumps(body.get("tool_calls")): + return Reply(body=json.dumps({"action": "BLOCKED", "blocked_reason": BLOCK_REASON}).encode()) + if any(MASKED_WORD in text for text in texts): + masked: Final = [text.replace(MASKED_WORD, MASK) for text in texts] + return Reply(body=json.dumps({"action": "GUARDRAIL_INTERVENED", "texts": masked}).encode()) + return Reply(body=json.dumps({"action": "NONE"}).encode()) + + +def bedrock_texts(request: Request) -> tuple[str, ...]: + return tuple(str(item["text"]["text"]) for item in json.loads(request.body)["content"]) + + +def bedrock_verdict(request: Request) -> Reply: + """ApplyGuardrail that anonymizes ``MASKED_WORD`` and blocks ``BLOCKED_WORD``, as Bedrock policies do.""" + texts: Final = bedrock_texts(request) + if any(BLOCKED_WORD in text for text in texts): + word: Final = {"match": BLOCKED_WORD, "action": "BLOCKED"} + return Reply( + body=json.dumps( + { + "action": "GUARDRAIL_INTERVENED", + "outputs": [{"text": BEDROCK_BLOCK_MESSAGE}], + "assessments": [{"wordPolicy": {"customWords": [word]}}], + } + ).encode() + ) + if not any(MASKED_WORD in text for text in texts): + return Reply(body=json.dumps({"action": "NONE", "outputs": [], "assessments": []}).encode()) + pii: Final = {"match": MASKED_WORD, "type": "ADDRESS", "action": "ANONYMIZED"} + return Reply( + body=json.dumps( + { + "action": "GUARDRAIL_INTERVENED", + "outputs": [{"text": text.replace(MASKED_WORD, MASK)} for text in texts], + "assessments": [{"sensitiveInformationPolicy": {"piiEntities": [pii]}}], + } + ).encode() + ) + + +@dataclass(frozen=True, slots=True) +class StreamOutcome: + events: tuple[str, ...] + text: str + error_code: str | None + error_message: str | None + rejected_before_stream: str | None + + +@dataclass(frozen=True, slots=True) +class Rig: + gateway: Gateway + scenario: Scenario + provider: Wire + sink: Wire + scripts: dict[str, Reply] + keys: Mapping[str, str] + outage: threading.Event + + def scans(self, guardrail: str, marker: str) -> tuple[tuple[str, ...], ...]: + return tuple( + tuple(json.loads(request.body)["texts"]) + for request in self.sink.drain() + if request.target.startswith(f"/{guardrail}/") and marker.encode() in request.body + ) + + def bedrock_scans(self, guardrail: str, marker: str) -> tuple[tuple[str, ...], ...]: + return tuple( + bedrock_texts(request) + for request in self.sink.drain() + if request.target.startswith(f"/{guardrail}/") and marker.encode() in request.body + ) + + def guardrail_request_data(self, guardrail: str, marker: str) -> tuple[Mapping[str, JsonValue], ...]: + """The request data each Noma scan of the marker's stream carried, in the order the scans ran.""" + return tuple( + json.loads(request.body)["request_data"] + for request in self.sink.drain() + if request.target.startswith(f"/{guardrail}/") and marker.encode() in request.body + ) + + def invoke_stream(self, guardrail: str, marker: str, model: str = CLAUDE) -> StreamOutcome: + client: Final = boto3.client( + "bedrock-runtime", + region_name="us-east-1", + endpoint_url=str(self.gateway.client.base_url).rstrip("/") + "/bedrock", + aws_access_key_id="AKIASYNTHETICRELAY", + aws_secret_access_key="synthetic-relay-secret", + config=Config(read_timeout=case_timeout_seconds(), retries={"max_attempts": 1}), + ) + key: Final = self.keys[guardrail] + + def authenticate(request: object, **_: object) -> None: + getattr(request, "headers")["x-litellm-api-key"] = key + + client.meta.events.register("before-send.bedrock-runtime.*", authenticate) + try: + response: Final = client.invoke_model_with_response_stream( + modelId=model, body=json.dumps(_invoke_body(marker)) + ) + except ClientError as error: + return StreamOutcome((), "", None, None, str(error)) + events: Final[list[Mapping[str, JsonValue]]] = [] + error: EventStreamError | None = None + try: + for event in response["body"]: + events.append(json.loads(event["chunk"]["bytes"])) + except EventStreamError as raised: + error = raised + frames: Final = tuple(Frame("event", str(event["type"]), event) for event in events) + return StreamOutcome( + tuple(frame.kind for frame in frames), + streamed_text(frames), + None if error is None else error.response["Error"]["Code"], + None if error is None else error.response["Error"]["Message"], + None, + ) + + +def _guardrail(name: str, sink: Wire, **settings: JsonValue) -> dict[str, JsonValue]: + return { + "guardrail_name": name, + "litellm_params": { + "guardrail": "generic_guardrail_api", + "mode": "post_call", + "default_on": False, + "api_base": f"{sink.url}/{name}", + "api_key": "synthetic-guardrail-key", + **settings, + }, + } + + +def _bedrock_guardrail(name: str, sink: Wire, **settings: JsonValue) -> dict[str, JsonValue]: + return { + "guardrail_name": name, + "litellm_params": { + "guardrail": "bedrock", + "mode": "post_call", + "default_on": False, + "guardrailIdentifier": name, + "guardrailVersion": "1", + "aws_access_key_id": "AKIASYNTHETICGUARD", + "aws_secret_access_key": "synthetic-guardrail-secret", + "aws_region_name": "us-east-1", + "aws_bedrock_runtime_endpoint": f"{sink.url}/{name}", + **settings, + }, + } + + +def _noma_guardrail(name: str, sink: Wire, **settings: JsonValue) -> dict[str, JsonValue]: + return { + "guardrail_name": name, + "litellm_params": { + "guardrail": "noma_v2", + "mode": "post_call", + "default_on": False, + "api_base": f"{sink.url}/{name}", + "api_key": "synthetic-guardrail-key", + **settings, + }, + } + + +def _deployment(name: str, model_id: str, upstream: Wire) -> dict[str, JsonValue]: + return { + "model_name": name, + "litellm_params": { + "model": f"bedrock/{model_id}", + "api_base": upstream.url, + "aws_access_key_id": "AKIASYNTHETICRELAY", + "aws_secret_access_key": "synthetic-relay-secret", + "aws_region_name": "us-east-1", + }, + } + + +def _invoke_body(marker: str) -> dict[str, JsonValue]: + return { + "anthropic_version": "bedrock-2023-05-31", + "max_tokens": 64, + "messages": [{"role": "user", "content": [{"type": "text", "text": marker}]}], + } + + +def _marker(request: Request) -> str: + body: Final = json.loads(request.body) + if "prompt" in body: + return str(body["prompt"]) + return str(body["messages"][0]["content"][0]["text"]) + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + scripts: Final[dict[str, Reply]] = {} + outage: Final = threading.Event() + + def provider(request: Request) -> Reply: + assert unquote(request.target) in ( + f"/model/{CLAUDE_ID}/invoke-with-response-stream", + f"/model/{CLAUDE_ID}/invoke", + f"/model/{LLAMA_ID}/invoke-with-response-stream", + f"/model/{PROFILE_ID}/invoke-with-response-stream", + ), request.target + return scripts[_marker(request)] + + def vendor(request: Request) -> Reply: + if outage.is_set() and request.target.startswith(f"/{OUTAGE}/"): + return Reply(status=503, body=b'{"error":"guardrail vendor unavailable"}') + if request.target.startswith((f"/{BEDROCK}/", f"/{BEDROCK_MESSAGE}/", f"/{RELEASE_ON_SCAN}/")): + return bedrock_verdict(request) + if request.target.startswith((f"/{NOMA}/", f"/{NOMA_HELD}/")): + return Reply(body=json.dumps({"action": "NONE"}).encode()) + return verdict(request) + + with httpx.Client(base_url=os.environ["INTEGRATION_PROXY_URL"], timeout=30, trust_env=False) as root_client: + root: Final = Gateway( + root_client, + os.environ.get("INTEGRATION_MASTER_KEY", "sk-integration-master"), + os.environ["INTEGRATION_UPSTREAM_URL"], + ) + directory: Final = tmp_path_factory.mktemp("invoke-stream-relay") + with wire_server(provider) as upstream, wire_server(vendor) as sink: + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["guardrails"] = [ + _guardrail(SAMPLED, sink), + _guardrail(HELD, sink, streaming_buffer_until_moderated=True), + _guardrail(END_ONLY, sink, streaming_end_of_stream_only=True), + _guardrail(OUTAGE, sink), + _guardrail(PARALLEL, sink, run_in_parallel=True), + _noma_guardrail(NOMA, sink, streaming_buffer_until_moderated=False), + _noma_guardrail(NOMA_HELD, sink, streaming_buffer_until_moderated=True), + _bedrock_guardrail(BEDROCK, sink), + _bedrock_guardrail(BEDROCK_MESSAGE, sink, disable_exception_on_block=True), + _bedrock_guardrail(RELEASE_ON_SCAN, sink, streaming_buffer_release_on_scan=True), + { + "guardrail_name": CODE, + "litellm_params": { + "guardrail": "custom_code", + "mode": "post_call", + "default_on": False, + "custom_code": ( + "def apply_guardrail(inputs, request_data, input_type):\n" + " for text in inputs['texts']:\n" + f" if '{BLOCKED_WORD}' in text:\n" + f" return block('{CODE_BLOCK_REASON}')\n" + " return allow()\n" + ), + }, + }, + ] + config["model_list"] = [ + _deployment(CLAUDE, CLAUDE_ID, upstream), + _deployment(LLAMA, LLAMA_ID, upstream), + _deployment(PROFILE, PROFILE_ID, upstream), + ] + config_path: Final = directory / "invoke-stream-relay.yaml" + config_path.write_text(yaml.safe_dump(config)) + with owned_proxy_process(root, directory, {}, config=config_path, workers=2) as owned: + with owned.gateway.scenario() as scenario: + guarded: Final = ( + SAMPLED, + HELD, + END_ONLY, + CODE, + BEDROCK, + BEDROCK_MESSAGE, + OUTAGE, + RELEASE_ON_SCAN, + PARALLEL, + NOMA, + NOMA_HELD, + ) + yield Rig( + owned.gateway, + scenario, + upstream, + sink, + scripts, + { + **{name: scenario.key(guardrails=[name]) for name in guarded}, + UNGUARDED: scenario.key(), + }, + outage, + ) + + +def _stream_with_call_id(rig: Rig, guardrail: str, marker: str, model: str = CLAUDE) -> tuple[int, str, bytes, str]: + with rig.gateway.client.stream( + "POST", + f"/bedrock/model/{model}/invoke-with-response-stream", + json={"prompt": marker} if model == LLAMA else _invoke_body(marker), + headers={"Authorization": f"Bearer {rig.keys[guardrail]}"}, + timeout=case_timeout_seconds(), + ) as response: + return ( + response.status_code, + response.headers.get("content-type", ""), + response.read(), + response.headers.get("x-litellm-call-id", ""), + ) + + +def _stream(rig: Rig, guardrail: str, marker: str, model: str = CLAUDE) -> tuple[int, str, bytes]: + status, content_type, body, _ = _stream_with_call_id(rig, guardrail, marker, model) + return status, content_type, body + + +def _exception(frames: Sequence[Frame]) -> tuple[str, str] | None: + last: Final = frames[-1] if frames else None + if last is None or last.message_type != "exception": + return None + return last.kind, str(last.payload.get("message")) + + +def _spend_statuses(call_id: str, at_least: int) -> list[JsonValue]: + rows: Final = eventually( + lambda: read_rows( + 'SELECT status FROM "LiteLLM_SpendLogs" WHERE litellm_call_id=%s OR request_id=%s ' + "OR metadata->>'litellm_call_id'=%s ORDER BY status", + (call_id, call_id, call_id), + ), + lambda rows: len(rows) >= at_least, + seconds=70, + ) + return [row["status"] for row in rows] + + +def test_first_frames_reach_the_caller_while_the_upstream_is_still_streaming_and_get_scanned(rig: Rig) -> None: + gate: Final = threading.Event() + marker: Final = f"invoke-{uuid.uuid4().hex}" + frames: Final = claude_frames(f"Hello {marker} ", "world") + first_transport_chunk: Final = frames[0] + frames[1] + rig.scripts[marker] = Reply( + content_type=EVENT_STREAM, chunks=(first_transport_chunk, *frames[2:]), gate_after_first=gate + ) + release_held_upstream: Final = threading.Timer(import_seconds(), gate.set) + release_held_upstream.start() + try: + with rig.gateway.client.stream( + "POST", + f"/bedrock/model/{CLAUDE}/invoke-with-response-stream", + json=_invoke_body(marker), + headers={"Authorization": f"Bearer {rig.keys[SAMPLED]}"}, + timeout=case_timeout_seconds(), + ) as response: + chunks: Final = response.iter_raw() + received: Final = bytearray() + for chunk in chunks: + received.extend(chunk) + if len(received) >= len(first_transport_chunk): + break + upstream_still_held: Final = not gate.is_set() + gate.set() + received.extend(b"".join(chunks)) + status: Final = response.status_code + finally: + release_held_upstream.cancel() + gate.set() + assert status == 200, received + assert upstream_still_held, "the first frames reached the caller only after the upstream finished its stream" + assert bytes(received) == b"".join(frames), decode(bytes(received)) + assert rig.scans(SAMPLED, marker) == ((f"Hello {marker} world",),) + + +def test_block_mid_stream_reaches_boto3_as_an_event_stream_exception(rig: Rig) -> None: + marker: Final = f"invoke-{uuid.uuid4().hex}" + rig.scripts[marker] = Reply( + content_type=EVENT_STREAM, chunks=claude_frames(f"The secret {marker} is ", BLOCKED_WORD, ".") + ) + outcome: Final = rig.invoke_stream(SAMPLED, marker) + assert outcome.rejected_before_stream is None, outcome + assert outcome.events[:2] == ("message_start", "content_block_start"), outcome + assert "message_stop" not in outcome.events, outcome + assert (outcome.error_code, outcome.error_message) == ("validationException", BLOCK_REASON), outcome + + +def test_claude_behind_an_application_inference_profile_is_scanned_by_its_request_body(rig: Rig) -> None: + marker: Final = f"invoke-{uuid.uuid4().hex}" + rig.scripts[marker] = Reply( + content_type=EVENT_STREAM, chunks=claude_frames(f"The secret {marker} is ", BLOCKED_WORD, ".") + ) + outcome: Final = rig.invoke_stream(SAMPLED, marker, model=PROFILE) + assert outcome.rejected_before_stream is None, outcome + assert "message_stop" not in outcome.events, outcome + assert (outcome.error_code, outcome.error_message) == ("validationException", BLOCK_REASON), outcome + + +def test_clean_stream_reaches_boto3_unchanged_after_its_scan(rig: Rig) -> None: + marker: Final = f"invoke-{uuid.uuid4().hex}" + rig.scripts[marker] = Reply(content_type=EVENT_STREAM, chunks=claude_frames(f"Fine {marker} ", "answer")) + outcome: Final = rig.invoke_stream(SAMPLED, marker) + assert outcome == StreamOutcome( + ( + "message_start", + "content_block_start", + "content_block_delta", + "content_block_delta", + "content_block_stop", + "message_delta", + "message_stop", + ), + f"Fine {marker} answer", + None, + None, + None, + ), outcome + assert rig.scans(SAMPLED, marker) == ((f"Fine {marker} answer",),) + + +def test_held_stream_writes_the_guardrail_masks_back_into_the_chunk_frames(rig: Rig) -> None: + marker: Final = f"invoke-{uuid.uuid4().hex}" + rig.scripts[marker] = Reply( + content_type=EVENT_STREAM, chunks=claude_frames(f"{marker} {MASKED_WORD[:3]}", f"{MASKED_WORD[3:]} is big") + ) + status, content_type, body = _stream(rig, HELD, marker) + assert (status, content_type) == (200, EVENT_STREAM), body + frames: Final = decode(body) + assert [frame.kind for frame in frames] == [ + "message_start", + "content_block_start", + "content_block_delta", + "content_block_delta", + "content_block_stop", + "message_delta", + "message_stop", + ], frames + assert streamed_text(frames) == f"{marker} {MASK} is big", frames + assert rig.scans(HELD, marker) == ((f"{marker} {MASKED_WORD} is big",),) + + +def test_held_stream_block_returns_an_error_before_any_frame(rig: Rig) -> None: + marker: Final = f"invoke-{uuid.uuid4().hex}" + rig.scripts[marker] = Reply( + content_type=EVENT_STREAM, chunks=claude_frames(f"The secret {marker} is ", BLOCKED_WORD) + ) + outcome: Final = rig.invoke_stream(HELD, marker) + assert outcome.events == () and outcome.rejected_before_stream is not None, outcome + assert outcome.rejected_before_stream.endswith( + f"(ValidationException) when calling the InvokeModelWithResponseStream operation: {BLOCK_REASON}" + ), outcome + + +def test_held_stream_the_upstream_breaks_off_still_scans_its_tool_input(rig: Rig) -> None: + marker: Final = f"invoke-{uuid.uuid4().hex}" + text_frames: Final = claude_frames(f"Looking up {marker}") + tool_use: Final = {"type": "tool_use", "id": "toolu_relay", "name": "lookup", "input": {}} + upstream_error: Final = aws_event_stream_frame( + { + ":exception-type": "modelStreamErrorException", + ":content-type": "application/json", + ":message-type": "exception", + }, + json.dumps({"message": "upstream stream broke"}).encode(), + ) + rig.scripts[marker] = Reply( + content_type=EVENT_STREAM, + chunks=( + *text_frames[:4], + chunk_frame({"type": "content_block_start", "index": 1, "content_block": tool_use}), + chunk_frame( + { + "type": "content_block_delta", + "index": 1, + "delta": {"type": "input_json_delta", "partial_json": json.dumps({"query": BLOCKED_WORD})}, + } + ), + chunk_frame({"type": "content_block_stop", "index": 1}), + upstream_error, + ), + ) + outcome: Final = rig.invoke_stream(HELD, marker) + assert outcome.events == () and outcome.rejected_before_stream is not None, outcome + assert outcome.rejected_before_stream.endswith( + f"(ValidationException) when calling the InvokeModelWithResponseStream operation: {BLOCK_REASON}" + ), outcome + + +def test_bedrock_guardrail_holds_the_stream_and_anonymizes_it(rig: Rig) -> None: + marker: Final = f"invoke-{uuid.uuid4().hex}" + rig.scripts[marker] = Reply( + content_type=EVENT_STREAM, chunks=claude_frames(f"{marker} ", MASKED_WORD, " is", " the", " capital") + ) + outcome: Final = rig.invoke_stream(BEDROCK, marker) + assert (outcome.error_code, outcome.rejected_before_stream) == (None, None), outcome + assert outcome.text == f"{marker} {MASK} is the capital", outcome + assert rig.bedrock_scans(BEDROCK, marker) == ((f"{marker} {MASKED_WORD} is the capital",),) + + +def test_bedrock_guardrail_block_message_replaces_a_held_stream(rig: Rig) -> None: + marker: Final = f"invoke-{uuid.uuid4().hex}" + rig.scripts[marker] = Reply( + content_type=EVENT_STREAM, chunks=claude_frames(f"The secret {marker} is ", BLOCKED_WORD) + ) + outcome: Final = rig.invoke_stream(BEDROCK_MESSAGE, marker) + assert (outcome.error_code, outcome.rejected_before_stream) == (None, None), outcome + assert outcome.events[0] == "message_start" and outcome.events[-1] == "message_stop", outcome + assert outcome.text == BEDROCK_BLOCK_MESSAGE, outcome + + +def test_end_of_stream_only_relays_every_content_frame_then_blocks(rig: Rig) -> None: + marker: Final = f"invoke-{uuid.uuid4().hex}" + frames: Final = claude_frames(f"Late {marker} ", "one", " two", " three", " four", f" {BLOCKED_WORD}") + rig.scripts[marker] = Reply(content_type=EVENT_STREAM, chunks=frames) + status, content_type, body = _stream(rig, END_ONLY, marker) + assert (status, content_type) == (200, EVENT_STREAM), body + content_frames: Final = b"".join(frames[:-2]) + assert body.startswith(content_frames), decode(body) + assert decode(body[len(content_frames) :]) == ( + Frame("exception", "validationException", {"message": BLOCK_REASON}), + ), decode(body) + assert rig.scans(END_ONLY, marker) == ((f"Late {marker} one two three four {BLOCKED_WORD}",),) + + +def test_block_raised_mid_stream_logs_the_relayed_spend_and_the_failure(rig: Rig) -> None: + marker: Final = f"invoke-{uuid.uuid4().hex}" + rig.scripts[marker] = Reply( + content_type=EVENT_STREAM, + chunks=claude_frames(f"The secret {marker} is ", BLOCKED_WORD, *(f" word{index}" for index in range(30))), + pause_between_chunks=0.05, + ) + status, content_type, body, call_id = _stream_with_call_id(rig, CODE, marker) + assert (status, content_type) == (200, EVENT_STREAM), body + frames: Final = decode(body) + assert frames[0].kind == "message_start", frames + assert _exception(frames) == ("validationException", CODE_BLOCK_REASON), frames + assert call_id, "the relayed stream carries no x-litellm-call-id" + assert _spend_statuses(call_id, 2) == ["failure", "success"] + + +def test_non_claude_invoke_stream_is_relayed_as_received_without_a_scan(rig: Rig) -> None: + marker: Final = f"invoke-{uuid.uuid4().hex}" + frames: Final = llama_frames(f"Llama {marker} ", BLOCKED_WORD, " end") + rig.scripts[marker] = Reply(content_type=EVENT_STREAM, chunks=frames) + status, content_type, body = _stream(rig, SAMPLED, marker, model=LLAMA) + assert (status, content_type) == (200, EVENT_STREAM), body + assert body == b"".join(frames), decode(body) + assert rig.scans(SAMPLED, marker) == () + + +def test_invoke_stream_without_a_guardrail_is_relayed_as_received(rig: Rig) -> None: + marker: Final = f"invoke-{uuid.uuid4().hex}" + frames: Final = claude_frames(f"Open {marker} ", BLOCKED_WORD, " end") + rig.scripts[marker] = Reply(content_type=EVENT_STREAM, chunks=frames) + status, content_type, body = _stream(rig, UNGUARDED, marker) + assert (status, content_type) == (200, EVENT_STREAM), body + assert body == b"".join(frames), decode(body) + + +def test_held_stream_released_on_scan_stays_buffered_and_applies_masks(rig: Rig) -> None: + marker: Final = f"invoke-{uuid.uuid4().hex}" + rig.scripts[marker] = Reply( + content_type=EVENT_STREAM, chunks=claude_frames(f"{marker} ", MASKED_WORD, " is", " the", " capital") + ) + outcome: Final = rig.invoke_stream(RELEASE_ON_SCAN, marker) + assert (outcome.error_code, outcome.rejected_before_stream) == (None, None), outcome + assert outcome.events[0] == "message_start" and outcome.events[-1] == "message_stop", outcome + assert outcome.text == f"{marker} {MASK} is the capital", outcome + assert rig.bedrock_scans(RELEASE_ON_SCAN, marker) == ((f"{marker} {MASKED_WORD} is the capital",),) + + +def test_held_stream_released_on_scan_with_a_block_returns_an_error_before_any_event(rig: Rig) -> None: + marker: Final = f"invoke-{uuid.uuid4().hex}" + rig.scripts[marker] = Reply( + content_type=EVENT_STREAM, chunks=claude_frames(f"Release {marker} ", BLOCKED_WORD, " end") + ) + outcome: Final = rig.invoke_stream(RELEASE_ON_SCAN, marker) + assert outcome.events == () and outcome.rejected_before_stream is not None, outcome + assert outcome.rejected_before_stream.endswith( + f"(ValidationException) when calling the InvokeModelWithResponseStream operation: {BLOCK_REASON}" + ), outcome + assert rig.bedrock_scans(RELEASE_ON_SCAN, marker) == ((f"Release {marker} {BLOCKED_WORD} end",),) + + +def test_parallel_guardrail_scans_the_whole_invoke_stream_before_any_event(rig: Rig) -> None: + marker: Final = f"invoke-{uuid.uuid4().hex}" + rig.scripts[marker] = Reply( + content_type=EVENT_STREAM, chunks=claude_frames(f"Parallel {marker} ", BLOCKED_WORD, " end") + ) + outcome: Final = rig.invoke_stream(PARALLEL, marker) + assert outcome.events == () and outcome.rejected_before_stream is not None, outcome + assert outcome.rejected_before_stream.endswith( + f"(ValidationException) when calling the InvokeModelWithResponseStream operation: {BLOCK_REASON}" + ), outcome + assert rig.scans(PARALLEL, marker) == ((f"Parallel {marker} {BLOCKED_WORD} end",),) + + +def test_each_sampled_scan_sends_a_streaming_request_with_every_event_relayed_so_far(rig: Rig) -> None: + marker: Final = f"invoke-{uuid.uuid4().hex}" + frames: Final = claude_frames(f"Noma {marker} ", *(f" word{index}" for index in range(9))) + rig.scripts[marker] = Reply(content_type=EVENT_STREAM, chunks=frames) + status, content_type, body = _stream(rig, NOMA, marker) + assert (status, content_type) == (200, EVENT_STREAM), body + assert body == b"".join(frames), decode(body) + scans: Final = rig.guardrail_request_data(NOMA, marker) + assert len(scans) > 1, scans + assert [scan.get("stream") for scan in scans] == [True] * len(scans), scans + relayed: Final = [len(scan["responses"]) for scan in scans if isinstance(scan.get("responses"), list)] + assert len(relayed) == len(scans) and relayed == sorted(set(relayed)), relayed + + +def test_the_end_of_stream_scan_sends_the_whole_stream_and_its_finish_reason(rig: Rig) -> None: + marker: Final = f"invoke-{uuid.uuid4().hex}" + frames: Final = claude_frames(f"Noma {marker} ", "held", " reply", " here") + rig.scripts[marker] = Reply(content_type=EVENT_STREAM, chunks=frames) + outcome: Final = rig.invoke_stream(NOMA_HELD, marker) + assert (outcome.error_code, outcome.rejected_before_stream) == (None, None), outcome + assert outcome.text == f"Noma {marker} held reply here", outcome + scans: Final = rig.guardrail_request_data(NOMA_HELD, marker) + assert [scan.get("stream") for scan in scans] == [True] * len(scans) and scans, scans + scan: Final = scans[-1] + responses: Final = scan.get("responses") + assert isinstance(responses, list) and len(responses) >= len(frames), responses + response: Final = scan.get("response") + assert isinstance(response, dict) and isinstance(choices := response.get("choices"), list), response + assert [choice.get("finish_reason") for choice in choices if isinstance(choice, dict)] == ["stop"], response + + +def test_upstream_exception_mid_stream_reaches_the_caller_after_the_scanned_frames(rig: Rig) -> None: + marker: Final = f"invoke-{uuid.uuid4().hex}" + upstream_error: Final = aws_event_stream_frame( + { + ":exception-type": "modelStreamErrorException", + ":content-type": "application/json", + ":message-type": "exception", + }, + json.dumps({"message": "upstream stream broke"}).encode(), + ) + frames: Final = claude_frames(f"Partial {marker} ", "answer") + rig.scripts[marker] = Reply(content_type=EVENT_STREAM, chunks=(*frames[:3], upstream_error)) + status, content_type, body = _stream(rig, SAMPLED, marker) + assert (status, content_type) == (200, EVENT_STREAM), body + assert body == b"".join((*frames[:3], upstream_error)), decode(body) + assert _exception(decode(body)) == ("modelStreamErrorException", "upstream stream broke") + assert rig.scans(SAMPLED, marker) == ((f"Partial {marker} ",),) + + +def test_guardrail_vendor_outage_during_a_burst_fails_closed_then_recovers(rig: Rig) -> None: + def burst(count: int) -> list[tuple[str, tuple[int, str, bytes, str]]]: + markers = [f"invoke-{uuid.uuid4().hex}" for _ in range(count)] + for marker in markers: + rig.scripts[marker] = Reply( + content_type=EVENT_STREAM, chunks=claude_frames(f"Outage {marker} ", "one", " two", " three") + ) + with ThreadPoolExecutor(max_workers=count) as pool: + return list(zip(markers, pool.map(partial(_stream_with_call_id, rig, OUTAGE), markers))) + + rig.outage.set() + try: + during: Final = burst(8) + finally: + rig.outage.clear() + after: Final = burst(8) + for marker, (status, _, body, _) in during: + exception = _exception(decode(body)) + assert status == 200 and exception is not None and exception[0] == "internalServerException", (marker, body) + for marker, (status, _, body, _) in after: + assert status == 200 and _exception(decode(body)) is None, (marker, body) + assert body == b"".join(rig.scripts[marker].chunks or ()), (marker, decode(body)) + for marker, (_, _, _, call_id) in (*during, *after): + assert _spend_statuses(call_id, 1).count("success") == 1, marker + + +def test_client_disconnect_stops_the_proxy_reading_the_upstream(rig: Rig) -> None: + marker: Final = f"invoke-{uuid.uuid4().hex}" + frames: Final = claude_frames(f"Long {marker} ", *(f" word{index}" for index in range(60))) + rig.scripts[marker] = Reply(content_type=EVENT_STREAM, chunks=frames, pause_between_chunks=0.1) + while not rig.provider.disconnected.empty(): + rig.provider.disconnected.get_nowait() + with rig.gateway.client.stream( + "POST", + f"/bedrock/model/{CLAUDE}/invoke-with-response-stream", + json=_invoke_body(marker), + headers={"Authorization": f"Bearer {rig.keys[SAMPLED]}"}, + timeout=case_timeout_seconds(), + ) as response: + assert response.status_code == 200 + received: Final = bytearray() + for chunk in response.iter_raw(): + received.extend(chunk) + if len(received) >= len(b"".join(frames[:3])): + break + closed: Final = eventually( + lambda: rig.provider.disconnected.qsize(), lambda count: count > 0, seconds=4, return_last_on_timeout=True + ) + assert closed, "the proxy kept reading the upstream after the caller went away" + assert rig.provider.disconnected.get_nowait() == f"/model/{CLAUDE_ID}/invoke-with-response-stream" + + +def test_non_streaming_invoke_block_is_unchanged(rig: Rig) -> None: + marker: Final = f"invoke-{uuid.uuid4().hex}" + reply: Final = { + "id": "msg_bdrk_relay", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5", + "content": [{"type": "text", "text": f"The secret {marker} is {BLOCKED_WORD}"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 9, "output_tokens": 4}, + } + rig.scripts[marker] = Reply(body=json.dumps(reply).encode()) + response: Final = rig.gateway.client.post( + f"/bedrock/model/{CLAUDE}/invoke", + json=_invoke_body(marker), + headers={"Authorization": f"Bearer {rig.keys[SAMPLED]}"}, + timeout=case_timeout_seconds(), + ) + assert response.status_code == 400, response.text + assert BLOCK_REASON in response.text, response.text diff --git a/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py b/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py index 1bf85b89232..85de558a8b2 100644 --- a/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py +++ b/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py @@ -11,12 +11,15 @@ Validates that: import copy import json +from collections.abc import Mapping, Sequence from datetime import datetime -from typing import Literal +from typing import Literal, TypeAlias from unittest.mock import AsyncMock, MagicMock import pytest import respx +from pydantic import JsonValue +from typing_extensions import ReadOnly, TypedDict import litellm from litellm.caching.caching import DualCache @@ -858,7 +861,10 @@ class _CardMaskingGuardrail(CustomGuardrail): async def _run_post_call_guardrail_on_stream( - stream_bytes: bytes, guardrail: CustomGuardrail, monkeypatch: pytest.MonkeyPatch + stream_bytes: bytes, + guardrail: CustomGuardrail, + monkeypatch: pytest.MonkeyPatch, + endpoint: str = "model/claude/converse-stream", ) -> bytes: monkeypatch.setattr(litellm, "callbacks", [guardrail]) logging_obj = LiteLLMLoggingObj( @@ -877,7 +883,7 @@ async def _run_post_call_guardrail_on_stream( data={ "model": "claude", "custom_llm_provider": "bedrock", - "endpoint": "model/claude/converse-stream", + "endpoint": endpoint, "litellm_logging_obj": logging_obj, }, ) @@ -1365,6 +1371,11 @@ def _decoded_frames(raw: bytes) -> list[tuple[str, str, dict]]: ] +_CONVERSE_STREAM = "model/claude/converse-stream" +_INVOKE_STREAM = "model/anthropic.claude-sonnet-5-v1:0/invoke-with-response-stream" +_APPLICATION_PROFILE = "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abc" + + async def _chunks(*chunks: bytes): for chunk in chunks: yield chunk @@ -1377,6 +1388,7 @@ async def _relay_through_post_call_guardrail( holds_until_moderated: bool = False, received: list[bytes] | None = None, outer: tuple[CustomGuardrail, bool] | None = None, + endpoint: str = _CONVERSE_STREAM, ) -> bytes: layers = [(guardrail, holds_until_moderated), *([outer] if outer is not None else [])] monkeypatch.setattr(litellm, "callbacks", [layer for layer, _ in layers]) @@ -1390,12 +1402,12 @@ async def _relay_through_post_call_guardrail( function_id="test-function", ) relayed = ProxyLogging(user_api_key_cache=DualCache()).async_post_call_event_stream_iterator_hook( - response=BedrockPassthroughGuardrailHandler.event_stream_items(_chunks(*chunks)), + response=BedrockPassthroughGuardrailHandler.event_stream_items(_chunks(*chunks), endpoint), user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), request_data={ "model": "claude", "custom_llm_provider": "bedrock", - "endpoint": "model/claude/converse-stream", + "endpoint": endpoint, "litellm_logging_obj": logging_obj, }, guardrails=layers, @@ -1505,7 +1517,12 @@ class TestConverseStreamFrames: stream = b"".join(_card_stream()) transport = tuple(stream[i : i + 7] for i in range(0, len(stream), 7)) - frames = [frame async for frame in BedrockPassthroughGuardrailHandler.event_stream_items(_chunks(*transport))] + frames = [ + frame + async for frame in BedrockPassthroughGuardrailHandler.event_stream_items( + _chunks(*transport), _CONVERSE_STREAM + ) + ] assert [frame.raw for frame in frames] == [*_card_stream(), b""] assert [frame.event_type for frame in frames] == [ @@ -1527,7 +1544,9 @@ class TestConverseStreamFrames: frames = [ frame - async for frame in BedrockPassthroughGuardrailHandler.event_stream_items(_chunks(frame + garbage, b"tail")) + async for frame in BedrockPassthroughGuardrailHandler.event_stream_items( + _chunks(frame + garbage, b"tail"), _CONVERSE_STREAM + ) ] assert [item.raw for item in frames] == [frame, garbage, b"tail", b""] @@ -1539,7 +1558,9 @@ class TestConverseStreamFrames: frames = [ frame - async for frame in BedrockPassthroughGuardrailHandler.event_stream_items(_chunks(complete + partial[:-3])) + async for frame in BedrockPassthroughGuardrailHandler.event_stream_items( + _chunks(complete + partial[:-3]), _CONVERSE_STREAM + ) ] assert [frame.raw for frame in frames] == [complete, partial[:-3], b""] @@ -1551,7 +1572,7 @@ class TestConverseStreamFrames: frames = [ frame async for frame in BedrockPassthroughGuardrailHandler.event_stream_items( - _chunks(_build_event_stream_frame("contentBlockDelta", delta)) + _chunks(_build_event_stream_frame("contentBlockDelta", delta)), _CONVERSE_STREAM ) ] @@ -1564,7 +1585,10 @@ class TestConverseStreamFrames: b'{"contentBlockIndex":0,"delta":{"text":"\xff"}}', ) - frames = [frame async for frame in BedrockPassthroughGuardrailHandler.event_stream_items(_chunks(frame))] + frames = [ + frame + async for frame in BedrockPassthroughGuardrailHandler.event_stream_items(_chunks(frame), _CONVERSE_STREAM) + ] assert [(item.raw, item.event_type, item.payload) for item in frames] == [ (frame, "contentBlockDelta", None), @@ -1901,7 +1925,7 @@ class TestConverseStreamGuardrailRelay: frames = [ frame async for frame in BedrockPassthroughGuardrailHandler.event_stream_items( - _chunks(*_tool_without_input_stream()) + _chunks(*_tool_without_input_stream()), _CONVERSE_STREAM ) ] @@ -2001,3 +2025,662 @@ class TestConverseStreamGuardrailRelay: item = BedrockPassthroughGuardrailHandler.event_stream_error_item("scan failed", status_code) assert _decoded_frames(item) == [("exception", exception_type, {"message": "scan failed"})] + + +class _InvokeTextBlock(TypedDict): + type: ReadOnly[Literal["text"]] + text: ReadOnly[str] + + +class _InvokeToolUseBlock(TypedDict): + type: ReadOnly[Literal["tool_use"]] + id: ReadOnly[str] + name: ReadOnly[str] + input: ReadOnly[Mapping[str, JsonValue]] + + +class _InvokeUsage(TypedDict, total=False): + input_tokens: ReadOnly[int] + output_tokens: ReadOnly[int] + + +class _InvokeMessage(TypedDict, total=False): + id: ReadOnly[str] + type: ReadOnly[Literal["message"]] + role: ReadOnly[Literal["assistant"]] + content: ReadOnly[Sequence[_InvokeTextBlock | _InvokeToolUseBlock]] + model: ReadOnly[str] + stop_reason: ReadOnly[str | None] + usage: ReadOnly[_InvokeUsage] + + +class _InvokeTextDelta(TypedDict): + type: ReadOnly[Literal["text_delta"]] + text: ReadOnly[str] + + +class _InvokeJsonDelta(TypedDict): + type: ReadOnly[Literal["input_json_delta"]] + partial_json: ReadOnly[str] + + +class _InvokeThinkingDelta(TypedDict): + type: ReadOnly[Literal["thinking_delta"]] + thinking: ReadOnly[str] + + +class _InvokeSignatureDelta(TypedDict): + type: ReadOnly[Literal["signature_delta"]] + signature: ReadOnly[str] + + +_InvokeDelta: TypeAlias = _InvokeTextDelta | _InvokeJsonDelta | _InvokeThinkingDelta | _InvokeSignatureDelta + + +class _InvokeMessageStart(TypedDict): + type: ReadOnly[Literal["message_start"]] + message: ReadOnly[_InvokeMessage] + + +class _InvokeBlockStart(TypedDict): + type: ReadOnly[Literal["content_block_start"]] + index: ReadOnly[int] + content_block: ReadOnly[_InvokeTextBlock | _InvokeToolUseBlock] + + +class _InvokeBlockDelta(TypedDict): + type: ReadOnly[Literal["content_block_delta"]] + index: ReadOnly[int] + delta: ReadOnly[_InvokeDelta] + + +class _InvokeBlockStop(TypedDict): + type: ReadOnly[Literal["content_block_stop"]] + index: ReadOnly[int] + + +class _InvokeStopReason(TypedDict): + stop_reason: ReadOnly[str] + + +class _InvokeMessageDelta(TypedDict): + type: ReadOnly[Literal["message_delta"]] + delta: ReadOnly[_InvokeStopReason] + usage: ReadOnly[_InvokeUsage] + + +class _InvokeMessageStop(TypedDict): + type: ReadOnly[Literal["message_stop"]] + + +_InvokeEvent: TypeAlias = ( + _InvokeMessageStart + | _InvokeBlockStart + | _InvokeBlockDelta + | _InvokeBlockStop + | _InvokeMessageDelta + | _InvokeMessageStop +) + + +def _invoke_raw_chunk_frame(event_json: bytes) -> bytes: + import base64 + + return _build_frame( + {":event-type": "chunk", ":content-type": "application/json", ":message-type": "event"}, + json.dumps({"bytes": base64.b64encode(event_json).decode(), "p": "abcdef"}).encode(), + ) + + +def _invoke_chunk_frame(event: _InvokeEvent) -> bytes: + return _invoke_raw_chunk_frame(json.dumps(event).encode()) + + +def _invoke_stream(*deltas: str) -> tuple[bytes, ...]: + return ( + _invoke_chunk_frame( + { + "type": "message_start", + "message": {"id": "msg_1", "type": "message", "role": "assistant", "content": [], "model": "claude"}, + } + ), + _invoke_chunk_frame({"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}), + *( + _invoke_chunk_frame({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": t}}) + for t in deltas + ), + _invoke_chunk_frame({"type": "content_block_stop", "index": 0}), + _invoke_chunk_frame( + {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 4}} + ), + _invoke_chunk_frame({"type": "message_stop"}), + ) + + +def _invoke_card_stream() -> tuple[bytes, ...]: + return _invoke_stream("Card: 4111 1111 ", "1111 1111.", " Thanks!") + + +def _decoded_invoke_frames(raw: bytes) -> list[tuple[str, dict[str, JsonValue]]]: + import base64 + import json + + return [ + (message_type, json.loads(base64.b64decode(payload["bytes"])) if kind == "chunk" else payload) + for message_type, kind, payload in _decoded_frames(raw) + ] + + +def _delta_text(event: Mapping[str, JsonValue]) -> str: + delta = event["delta"] + assert isinstance(delta, dict) and delta.get("type") == "text_delta", event + text = delta["text"] + assert isinstance(text, str), event + return text + + +def _streamed_text(raw: bytes) -> str: + return "".join( + _delta_text(event) + for _, event in _decoded_invoke_frames(raw) + if event.get("type") == "content_block_delta" + and isinstance(delta := event.get("delta"), dict) + and delta.get("type") == "text_delta" + ) + + +class _RequestDataRecordingGuardrail(CustomGuardrail): + def __init__(self) -> None: + super().__init__(guardrail_name="request-recorder", event_hook=GuardrailEventHooks.post_call, default_on=True) + self.scans: list[tuple[object, int, object]] = [] + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict[str, object], + input_type: Literal["request", "response"], + logging_obj: LiteLLMLoggingObj | None = None, + ) -> GenericGuardrailAPIInputs: + response = request_data.get("response") + self.scans.append( + ( + request_data.get("stream"), + len(request_data.get("responses") or ()), + response.choices[0].finish_reason if response is not None else None, + ) + ) + return inputs + + +class _ToolInputRecordingGuardrail(CustomGuardrail): + def __init__(self) -> None: + super().__init__(guardrail_name="tool-recorder", event_hook=GuardrailEventHooks.post_call, default_on=True) + self.tool_arguments: list[list[str]] = [] + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict[str, object], + input_type: Literal["request", "response"], + logging_obj: LiteLLMLoggingObj | None = None, + ) -> GenericGuardrailAPIInputs: + self.tool_arguments.append([call.function.arguments for call in inputs.get("tool_calls") or []]) + return inputs + + +def _invoke_delta_frame(index: int, delta: _InvokeDelta) -> bytes: + return _invoke_chunk_frame({"type": "content_block_delta", "index": index, "delta": delta}) + + +class TestBufferedInvokeStream: + @pytest.mark.asyncio + async def test_masks_are_written_back_into_the_chunk_frames_that_carried_the_text( + self, monkeypatch: pytest.MonkeyPatch + ): + guardrail = _CardMaskingGuardrail() + stream = b"".join(_invoke_card_stream()) + + result = await _run_post_call_guardrail_on_stream(stream, guardrail, monkeypatch, endpoint=_INVOKE_STREAM) + + assert guardrail.received_texts == [["Card: 4111 1111 1111 1111. Thanks!"]] + assert _streamed_text(result) == "Card: ****. Thanks!" + assert [event for _, event in _decoded_invoke_frames(result) if event["type"] != "content_block_delta"] == [ + event for _, event in _decoded_invoke_frames(stream) if event["type"] != "content_block_delta" + ] + + @pytest.mark.asyncio + async def test_each_content_block_is_scanned_and_masked_on_its_own(self, monkeypatch: pytest.MonkeyPatch): + guardrail = _CardMaskingGuardrail() + stream = b"".join( + ( + _invoke_delta_frame(0, {"type": "thinking_delta", "thinking": "The card is 4111 1111 1111 1111"}), + _invoke_delta_frame(0, {"type": "signature_delta", "signature": "sig-4111 1111 1111 1111"}), + _invoke_delta_frame(1, {"type": "text_delta", "text": "Paying with 4111 1111 "}), + _invoke_delta_frame(1, {"type": "text_delta", "text": "1111 1111 now"}), + _invoke_delta_frame(2, {"type": "input_json_delta", "partial_json": '{"card": "4111 1111 1111 1111"}'}), + _invoke_chunk_frame({"type": "message_stop"}), + ) + ) + + result = await _run_post_call_guardrail_on_stream(stream, guardrail, monkeypatch, endpoint=_INVOKE_STREAM) + + events = [event for _, event in _decoded_invoke_frames(result) if event["type"] == "content_block_delta"] + assert events[0]["delta"] == {"type": "thinking_delta", "thinking": "The card is ****"} + assert events[1]["delta"] == {"type": "signature_delta", "signature": "sig-4111 1111 1111 1111"} + assert _delta_text(events[2]) + _delta_text(events[3]) == "Paying with **** now" + assert events[4]["delta"] == {"type": "input_json_delta", "partial_json": '{"card": "****"}'} + + @pytest.mark.asyncio + async def test_two_text_blocks_reach_the_guardrail_as_two_texts(self, monkeypatch: pytest.MonkeyPatch): + guardrail = _CardMaskingGuardrail() + stream = b"".join( + ( + _invoke_delta_frame(0, {"type": "text_delta", "text": "First 4111 1111 1111 1111"}), + _invoke_delta_frame(1, {"type": "text_delta", "text": "Second 4111 1111 1111 1111"}), + _invoke_chunk_frame({"type": "message_stop"}), + ) + ) + + result = await _run_post_call_guardrail_on_stream(stream, guardrail, monkeypatch, endpoint=_INVOKE_STREAM) + + assert guardrail.received_texts == [["First 4111 1111 1111 1111", "Second 4111 1111 1111 1111"]] + events = [event for _, event in _decoded_invoke_frames(result)] + assert [_delta_text(event) for event in events[:2]] == ["First ****", "Second ****"] + + @pytest.mark.asyncio + async def test_a_delta_whose_text_is_not_a_string_is_left_as_received(self, monkeypatch: pytest.MonkeyPatch): + guardrail = _CardMaskingGuardrail() + odd = _invoke_raw_chunk_frame( + b'{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":4111}}' + ) + stream = b"".join( + ( + _invoke_delta_frame(0, {"type": "text_delta", "text": "Card 4111 1111 1111 1111"}), + odd, + _invoke_chunk_frame({"type": "message_stop"}), + ) + ) + + result = await _run_post_call_guardrail_on_stream(stream, guardrail, monkeypatch, endpoint=_INVOKE_STREAM) + + assert guardrail.received_texts == [["Card 4111 1111 1111 1111"]] + frames = _decoded_invoke_frames(result) + assert _delta_text(frames[0][1]) == "Card ****" + assert odd in result + + @pytest.mark.asyncio + async def test_a_stream_without_text_skips_the_post_call_hooks(self): + hook = AsyncMock() + stream = b"".join( + ( + _invoke_chunk_frame({"type": "message_start", "message": {"id": "msg_1", "content": []}}), + _invoke_chunk_frame({"type": "message_stop"}), + ) + ) + + result = await BedrockPassthroughGuardrailHandler.de_anonymize_event_stream( + body_bytes=stream, + proxy_logging_obj=MagicMock(post_call_success_hook=hook), + user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), + data={"custom_llm_provider": "bedrock", "endpoint": _INVOKE_STREAM}, + ) + + hook.assert_not_awaited() + assert result == stream + + @pytest.mark.asyncio + async def test_a_stream_the_guardrail_leaves_alone_is_returned_byte_for_byte(self, monkeypatch: pytest.MonkeyPatch): + guardrail = _RecordingGuardrail() + stream = b"".join(_invoke_stream("Hello ", "there")) + b"trailing" + + result = await _run_post_call_guardrail_on_stream(stream, guardrail, monkeypatch, endpoint=_INVOKE_STREAM) + + assert guardrail.received_texts == [["Hello there"]] + assert result == stream + + @pytest.mark.asyncio + async def test_a_block_raises_before_any_frame_is_returned(self, monkeypatch: pytest.MonkeyPatch): + guardrail = _RaisingGuardrail(GuardrailBlocked("Blocked"), "Card") + + with pytest.raises(GuardrailBlocked): + await _run_post_call_guardrail_on_stream( + b"".join(_invoke_card_stream()), guardrail, monkeypatch, endpoint=_INVOKE_STREAM + ) + + +class TestInvokeStreamGuardrailRelay: + @pytest.mark.asyncio + async def test_frames_are_cut_at_their_boundaries_whatever_the_transport_chunks(self): + stream = b"".join(_invoke_card_stream()) + transport = tuple(stream[i : i + 7] for i in range(0, len(stream), 7)) + + frames = [ + frame + async for frame in BedrockPassthroughGuardrailHandler.event_stream_items( + _chunks(*transport), _INVOKE_STREAM + ) + ] + + assert [frame.raw for frame in frames] == [*_invoke_card_stream(), b""] + assert [frame.event["type"] for frame in frames[:-1]] == [ + event["type"] for _, event in _decoded_invoke_frames(stream) + ] + assert [BedrockPassthroughGuardrailHandler.encode_event_stream_item(frame) for frame in frames] == [ + *_invoke_card_stream(), + b"", + ] + + @pytest.mark.parametrize( + "payload", + [ + pytest.param(b"not json", id="payload-not-json"), + pytest.param(b'{"bytes": "%%%"}', id="bytes-not-base64"), + pytest.param(b'{"bytes": "bm90IGpzb24="}', id="decoded-not-json"), + pytest.param(b'{"bytes": "eyJpbmRleCI6IDB9"}', id="event-without-type"), + pytest.param(b'{"p": "abc"}', id="no-bytes"), + pytest.param(b'{"bytes": 5}', id="bytes-not-a-string"), + ], + ) + @pytest.mark.asyncio + async def test_a_chunk_that_does_not_decode_is_relayed_as_received(self, payload: bytes): + frame = _build_frame( + {":event-type": "chunk", ":content-type": "application/json", ":message-type": "event"}, payload + ) + + frames = [ + item async for item in BedrockPassthroughGuardrailHandler.event_stream_items(_chunks(frame), _INVOKE_STREAM) + ] + + assert [(item.raw, item.event) for item in frames] == [(frame, None), (b"", None)] + assert BedrockPassthroughGuardrailHandler.encode_event_stream_item(frames[0]) == frame + + @pytest.mark.asyncio + async def test_sampled_stream_relays_frames_as_received_and_scans_the_text(self, monkeypatch: pytest.MonkeyPatch): + guardrail = _CardMaskingGuardrail() + + relayed = await _relay_through_post_call_guardrail( + _invoke_card_stream(), guardrail, monkeypatch, endpoint=_INVOKE_STREAM + ) + + assert relayed == b"".join(_invoke_card_stream()) + assert guardrail.received_texts == [[f"Card: {_TEST_CARD}. Thanks!"]] + + @pytest.mark.asyncio + async def test_stream_held_until_moderated_writes_the_masks_back_into_the_chunk_frames( + self, monkeypatch: pytest.MonkeyPatch + ): + stream = _invoke_card_stream() + + relayed = await _relay_through_post_call_guardrail( + stream, _CardMaskingGuardrail(), monkeypatch, holds_until_moderated=True, endpoint=_INVOKE_STREAM + ) + + assert [event["type"] for _, event in _decoded_invoke_frames(relayed)] == [ + event["type"] for _, event in _decoded_invoke_frames(b"".join(stream)) + ] + assert _streamed_text(relayed) == "Card: ****. Thanks!" + assert {payload.get("p") for _, _, payload in _decoded_frames(relayed)} == {"abcdef"} + assert relayed.startswith(b"".join(stream[:2])) and relayed.endswith(b"".join(stream[5:])) + + @pytest.mark.parametrize( + "ending", + [ + pytest.param( + BedrockPassthroughGuardrailHandler.event_stream_error_item("Model stream error", 500), + id="exception-frame", + ), + pytest.param(b"", id="upstream-closed-without-a-stop-reason"), + ], + ) + @pytest.mark.asyncio + async def test_held_stream_that_ends_without_a_stop_reason_still_applies_masks( + self, monkeypatch: pytest.MonkeyPatch, ending: bytes + ): + stream = (*_invoke_card_stream()[:4], ending) + + relayed = await _relay_through_post_call_guardrail( + stream, _CardMaskingGuardrail(), monkeypatch, holds_until_moderated=True, endpoint=_INVOKE_STREAM + ) + + assert _streamed_text(relayed) == "Card: ****." + assert relayed.endswith(ending) + + @pytest.mark.asyncio + async def test_a_later_guardrail_scans_the_masks_a_held_guardrail_wrote_back(self, monkeypatch: pytest.MonkeyPatch): + held = _RecordingGuardrail() + + relayed = await _relay_through_post_call_guardrail( + _invoke_card_stream(), + _CardMaskingGuardrail(), + monkeypatch, + holds_until_moderated=True, + outer=(held, True), + endpoint=_INVOKE_STREAM, + ) + + assert _streamed_text(relayed) == "Card: ****. Thanks!" + assert held.received_texts == [["Card: ****. Thanks!"]] + + @pytest.mark.asyncio + async def test_end_of_stream_only_block_withholds_the_frames_after_the_stop_reason( + self, monkeypatch: pytest.MonkeyPatch + ): + from fastapi import HTTPException + + guardrail = _RaisingGuardrail(HTTPException(status_code=400, detail="Violated guardrail policy"), "Thanks") + guardrail.streaming_end_of_stream_only = True + relayed: list[bytes] = [] + + with pytest.raises(HTTPException): + await _relay_through_post_call_guardrail( + _invoke_card_stream(), guardrail, monkeypatch, received=relayed, endpoint=_INVOKE_STREAM + ) + + assert relayed == list(_invoke_card_stream()[:-2]) + + @pytest.mark.asyncio + async def test_block_after_frames_were_relayed_raises_for_the_relay_to_report( + self, monkeypatch: pytest.MonkeyPatch + ): + from fastapi import HTTPException + + block = HTTPException(status_code=400, detail="Violated guardrail policy") + relayed: list[bytes] = [] + + with pytest.raises(HTTPException) as raised: + await _relay_through_post_call_guardrail( + _invoke_card_stream(), + _RaisingGuardrail(block, "Thanks"), + monkeypatch, + received=relayed, + endpoint=_INVOKE_STREAM, + ) + + assert raised.value is block + assert relayed == list(_invoke_card_stream()[:4]) + + @pytest.mark.asyncio + async def test_guardrail_message_ends_the_message_after_closing_the_open_block( + self, monkeypatch: pytest.MonkeyPatch + ): + from litellm.exceptions import ModifyResponseException + + guardrail = _RaisingGuardrail( + ModifyResponseException(message="Sorry, I can't share that.", model="claude", request_data={}), + "1111 1111.", + ) + stream = _invoke_stream("Card: ", "4111 1111 ", "1111 1111.", " Done") + + relayed = await _relay_through_post_call_guardrail(stream, guardrail, monkeypatch, endpoint=_INVOKE_STREAM) + + assert relayed.startswith(b"".join(stream[:4])) + assert [event for _, event in _decoded_invoke_frames(relayed[len(b"".join(stream[:4])) :])] == [ + {"type": "content_block_stop", "index": 0}, + {"type": "content_block_start", "index": 1, "content_block": {"type": "text", "text": ""}}, + { + "type": "content_block_delta", + "index": 1, + "delta": {"type": "text_delta", "text": "Sorry, I can't share that."}, + }, + {"type": "content_block_stop", "index": 1}, + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 0}, + }, + {"type": "message_stop"}, + ] + + @pytest.mark.asyncio + async def test_guardrail_message_on_a_held_stream_is_a_whole_message(self, monkeypatch: pytest.MonkeyPatch): + from litellm.exceptions import ModifyResponseException + + guardrail = _RaisingGuardrail( + ModifyResponseException( + message="Sorry, I can't share that.", + model="claude", + request_data={"custom_llm_provider": "bedrock", "endpoint": _INVOKE_STREAM}, + ), + "Card", + ) + + relayed = await _relay_through_post_call_guardrail( + _invoke_card_stream(), guardrail, monkeypatch, holds_until_moderated=True, endpoint=_INVOKE_STREAM + ) + + assert [event["type"] for _, event in _decoded_invoke_frames(relayed)] == [ + "message_start", + "content_block_start", + "content_block_delta", + "content_block_stop", + "message_delta", + "message_stop", + ] + assert _streamed_text(relayed) == "Sorry, I can't share that." + + @pytest.mark.parametrize( + ("endpoint", "model", "body", "relayable"), + [ + ("model/x/converse-stream", "meta.llama4-maverick-17b-instruct-v1:0", {}, True), + (_INVOKE_STREAM, "anthropic.claude-sonnet-5-v1:0", {}, True), + (_INVOKE_STREAM, "us.anthropic.claude-sonnet-5-v1:0", {}, True), + (_INVOKE_STREAM, "invoke/anthropic.claude-sonnet-5-v1:0", {}, True), + (_INVOKE_STREAM, "meta.llama4-maverick-17b-instruct-v1:0", {"prompt": "hi"}, False), + (_INVOKE_STREAM, _APPLICATION_PROFILE, {"anthropic_version": "bedrock-2023-05-31"}, True), + (_INVOKE_STREAM, _APPLICATION_PROFILE, {"prompt": "hi"}, False), + (_INVOKE_STREAM, _APPLICATION_PROFILE, None, False), + (_INVOKE_STREAM, _APPLICATION_PROFILE, ["anthropic_version"], False), + ("model/anthropic.claude-sonnet-5-v1:0/invoke", "anthropic.claude-sonnet-5-v1:0", {}, False), + ("model/x/invoke", _APPLICATION_PROFILE, {"anthropic_version": "bedrock-2023-05-31"}, False), + ], + ) + def test_relayed_and_buffered_streams_are_converse_and_claude_invoke_streams( + self, endpoint: str, model: str, body: object, relayable: bool + ): + assert BedrockPassthroughGuardrailHandler.event_stream_endpoint_is_relayable(endpoint, model, body) is relayable + assert ( + BedrockPassthroughGuardrailHandler.event_stream_endpoint_is_de_anonymizable(endpoint, model, body) + is relayable + ) + + @pytest.mark.asyncio + async def test_each_sampled_scan_sees_a_streaming_request_and_every_event_relayed_so_far( + self, monkeypatch: pytest.MonkeyPatch + ): + guardrail = _RequestDataRecordingGuardrail() + guardrail.streaming_sampling_rate = 2 + + await _relay_through_post_call_guardrail( + _invoke_stream("Card: ", "4111 1111 ", "1111 1111", "."), guardrail, monkeypatch, endpoint=_INVOKE_STREAM + ) + + assert len(guardrail.scans) > 1 + assert all(stream is True for stream, _, _ in guardrail.scans) + responses_seen = [responses for _, responses, _ in guardrail.scans] + assert responses_seen == sorted(set(responses_seen)), responses_seen + assert [finish_reason for _, _, finish_reason in guardrail.scans[:-1]] == [None] * (len(guardrail.scans) - 1) + + @pytest.mark.asyncio + async def test_the_held_scan_sees_the_whole_stream_and_its_finish_reason(self, monkeypatch: pytest.MonkeyPatch): + guardrail = _RequestDataRecordingGuardrail() + stream = ( + _invoke_chunk_frame( + { + "type": "message_start", + "message": { + "id": "msg_1", + "type": "message", + "role": "assistant", + "content": [], + "model": "claude", + "stop_reason": None, + "usage": {"input_tokens": 5, "output_tokens": 0}, + }, + } + ), + *_invoke_card_stream()[1:], + ) + + await _relay_through_post_call_guardrail( + stream, guardrail, monkeypatch, holds_until_moderated=True, endpoint=_INVOKE_STREAM + ) + + assert guardrail.scans == [(True, len(stream) + 1, "stop")] + + @pytest.mark.parametrize("holds_until_moderated", [True, False], ids=["held", "sampled"]) + @pytest.mark.asyncio + async def test_a_stream_the_upstream_breaks_off_still_gets_its_tool_input_scanned( + self, monkeypatch: pytest.MonkeyPatch, holds_until_moderated: bool + ): + guardrail = _ToolInputRecordingGuardrail() + stream = ( + _invoke_chunk_frame( + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}} + ), + _invoke_delta_frame(0, {"type": "text_delta", "text": "Calling the tool"}), + _invoke_chunk_frame({"type": "content_block_stop", "index": 0}), + _invoke_chunk_frame( + { + "type": "content_block_start", + "index": 1, + "content_block": {"type": "tool_use", "id": "toolu_1", "name": "lookup", "input": {}}, + } + ), + _invoke_delta_frame(1, {"type": "input_json_delta", "partial_json": '{"card": "4111 1111 1111 1111"}'}), + _invoke_chunk_frame({"type": "content_block_stop", "index": 1}), + _build_frame( + { + ":exception-type": "modelStreamErrorException", + ":content-type": "application/json", + ":message-type": "exception", + }, + b'{"message":"Model stream error"}', + ), + ) + + await _relay_through_post_call_guardrail( + stream, guardrail, monkeypatch, holds_until_moderated=holds_until_moderated, endpoint=_INVOKE_STREAM + ) + + assert guardrail.tool_arguments[-1] == ['{"card": "4111 1111 1111 1111"}'] + + @pytest.mark.asyncio + async def test_sampled_scan_key_reads_the_received_text_once_a_guardrail_rewrote_it(self): + delta: _InvokeBlockDelta = { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": "In Paris"}, + } + frames = [ + frame + async for frame in BedrockPassthroughGuardrailHandler.event_stream_items( + _chunks(_invoke_chunk_frame(delta)), _INVOKE_STREAM + ) + ] + received_key = BedrockPassthroughGuardrailHandler().get_streaming_scan_key(frames) + frames[0].guardrailed_event = {**delta, "delta": {"type": "text_delta", "text": "In [CITY]"}} + + rewritten_key = BedrockPassthroughGuardrailHandler().get_streaming_scan_key(frames) + + assert received_key is not None and received_key.texts == ("In Paris",) + assert rewritten_key == received_key diff --git a/tests/unit/llms/pass_through/guardrail_translation/test_handler.py b/tests/unit/llms/pass_through/guardrail_translation/test_handler.py index 9b5bfe5534f..0dba0d7b90f 100644 --- a/tests/unit/llms/pass_through/guardrail_translation/test_handler.py +++ b/tests/unit/llms/pass_through/guardrail_translation/test_handler.py @@ -1,7 +1,8 @@ import json -from typing import Final +from typing import Final, Literal import pytest +from typing_extensions import ReadOnly, TypedDict from litellm.constants import SERVER_STREAMING_CLASSIFICATION_KEY, SERVER_STREAMING_CLASSIFICATION_MARKER from litellm.llms.pass_through.guardrail_translation.handler import PassThroughEndpointHandler @@ -77,7 +78,7 @@ async def _frames(*raw: bytes) -> list[object]: for chunk in raw: yield chunk - return [frame async for frame in BedrockPassthroughGuardrailHandler.event_stream_items(chunks())] + return [frame async for frame in BedrockPassthroughGuardrailHandler.event_stream_items(chunks(), "converse-stream")] @pytest.mark.asyncio @@ -116,3 +117,96 @@ def test_passthrough_route_block_before_any_frame_uses_the_request_provider(prov ) assert (chunks is not None) is ends_message + + +_INVOKE_STREAM: Final = "model/anthropic.claude-sonnet-5-v1:0/invoke-with-response-stream" + + +class _TextDelta(TypedDict): + type: ReadOnly[Literal["text_delta"]] + text: ReadOnly[str] + + +class _TextDeltaEvent(TypedDict): + type: ReadOnly[Literal["content_block_delta"]] + index: ReadOnly[int] + delta: ReadOnly[_TextDelta] + + +def _invoke_frame(event: _TextDeltaEvent) -> bytes: + import base64 + + return _converse_frame("chunk", {"bytes": base64.b64encode(json.dumps(event).encode()).decode()}) + + +@pytest.mark.parametrize( + ("provider", "endpoint", "model", "body", "relayable"), + [ + ("bedrock", "model/x/converse-stream", "meta.llama4-maverick-17b-instruct-v1:0", None, True), + ("bedrock", _INVOKE_STREAM, "anthropic.claude-sonnet-5-v1:0", None, True), + ("bedrock", _INVOKE_STREAM, "meta.llama4-maverick-17b-instruct-v1:0", {"prompt": "hi"}, False), + ("bedrock", _INVOKE_STREAM, None, None, False), + ("bedrock", _INVOKE_STREAM, None, {"anthropic_version": "bedrock-2023-05-31"}, True), + ("bedrock", None, "anthropic.claude-sonnet-5-v1:0", None, False), + ("vertex_ai", "model/x/converse-stream", "anthropic.claude-sonnet-5-v1:0", None, False), + ], +) +def test_passthrough_route_relays_and_buffers_the_streams_its_provider_translation_scans( + provider: str, endpoint: str | None, model: str | None, body: object, relayable: bool +): + from litellm.llms.pass_through.guardrail_translation.handler import LlmPassthroughRouteHandler + + assert LlmPassthroughRouteHandler.supports_event_stream_relay(provider, endpoint, model, body) is relayable + assert ( + LlmPassthroughRouteHandler.supports_event_stream_de_anonymization(provider, endpoint, model, body) is relayable + ) + + +@pytest.mark.asyncio +async def test_passthrough_route_frames_an_invoke_stream_for_the_bedrock_translation(): + from litellm.llms.pass_through.guardrail_translation.handler import LlmPassthroughRouteHandler + + async def chunks(): + yield _invoke_frame({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "Hi "}}) + yield _invoke_frame({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "you"}}) + + items: Final = LlmPassthroughRouteHandler.event_stream_items("bedrock", _INVOKE_STREAM, chunks()) + assert items is not None + frames: Final = [frame async for frame in items] + + scan_key: Final = LlmPassthroughRouteHandler().get_streaming_scan_key(frames) + + assert scan_key is not None and scan_key.texts == ("Hi you",) + assert LlmPassthroughRouteHandler()._check_streaming_has_ended(frames) is True + assert LlmPassthroughRouteHandler()._check_streaming_has_ended(frames[:2]) is False + + +def test_passthrough_route_block_before_any_invoke_frame_is_a_whole_message_in_chunk_frames(): + import base64 + + from botocore.eventstream import EventStreamBuffer + + from litellm.exceptions import ModifyResponseException + from litellm.llms.pass_through.guardrail_translation.handler import LlmPassthroughRouteHandler + + chunks: Final = LlmPassthroughRouteHandler().build_block_sse_chunks( + ModifyResponseException( + message="blocked", model="m", request_data={"custom_llm_provider": "bedrock", "endpoint": _INVOKE_STREAM} + ), + stream_started=False, + responses_so_far=[], + ) + + assert chunks is not None + buffer: Final = EventStreamBuffer() + buffer.add_data(b"".join(chunks)) + events: Final = [json.loads(base64.b64decode(json.loads(message.payload)["bytes"])) for message in buffer] + assert [event["type"] for event in events] == [ + "message_start", + "content_block_start", + "content_block_delta", + "content_block_stop", + "message_delta", + "message_stop", + ] + assert events[2]["delta"]["text"] == "blocked" diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 37cc67af143..6cc4e22a26a 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -2,7 +2,7 @@ import asyncio import copy import datetime import json -from collections.abc import Mapping +from collections.abc import AsyncIterable, Mapping from types import MappingProxyType, SimpleNamespace from typing import AsyncGenerator, Callable, Final, Iterator, Literal, Optional, Sequence from urllib.parse import unquote_plus @@ -6268,8 +6268,14 @@ class TestAllmPassthroughStreamingProviderGate: failure_hook=None, read_body=True, request=None, + endpoint="model/us.amazon.nova-lite-v1:0/converse-stream", + model="us.amazon.nova-lite-v1:0", + post_call_guarded=True, + request_body=None, ): - processing_obj = self._build_processing_obj("bedrock", "model/us.amazon.nova-lite-v1:0/converse-stream") + processing_obj = self._build_processing_obj("bedrock", endpoint) + processing_obj.data["litellm_logging_obj"].model = model + processing_obj.data["data"] = request_body monkeypatch.setattr(litellm, "callbacks", [guardrail]) recorded: dict = {} @@ -6288,7 +6294,7 @@ class TestAllmPassthroughStreamingProviderGate: ), patch.object( ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", - return_value=True, + return_value=post_call_guarded, ), patch( "litellm.llms.bedrock.passthrough.guardrail_translation.handler." "BedrockPassthroughGuardrailHandler.de_anonymize_event_stream", @@ -6323,8 +6329,30 @@ class TestAllmPassthroughStreamingProviderGate: "_parallel_guardrail", ], ) - async def test_bedrock_converse_stream_is_buffered_for_a_guardrail_that_needs_the_whole_response( - self, monkeypatch, guardrail_factory + @pytest.mark.parametrize( + ("endpoint", "model", "request_body"), + [ + ("model/us.amazon.nova-lite-v1:0/converse-stream", "us.amazon.nova-lite-v1:0", None), + ( + "model/us.anthropic.claude-sonnet-5-v1:0/invoke-with-response-stream", + "us.anthropic.claude-sonnet-5-v1:0", + None, + ), + ( + "model/abc/invoke-with-response-stream", + "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abc", + {"anthropic_version": "bedrock-2023-05-31", "max_tokens": 8, "messages": []}, + ), + ], + ids=["converse-stream", "claude-invoke-stream", "claude-invoke-stream-behind-an-application-inference-profile"], + ) + async def test_bedrock_stream_is_buffered_for_a_guardrail_that_needs_the_whole_response( + self, + monkeypatch: pytest.MonkeyPatch, + guardrail_factory: str, + endpoint: str, + model: str, + request_body: Mapping[str, object] | None, ): guardrail = ( self._apply_guardrail(mask_response_content=True) @@ -6333,7 +6361,12 @@ class TestAllmPassthroughStreamingProviderGate: ) result, body, mock_handler, recorded = await self._run_converse_stream( - monkeypatch, guardrail, [b"raw-1", b"raw-2"] + monkeypatch, + guardrail, + [b"raw-1", b"raw-2"], + endpoint=endpoint, + model=model, + request_body=request_body, ) assert isinstance(result, Response) @@ -6635,6 +6668,110 @@ class TestAllmPassthroughStreamingProviderGate: assert "upstream closed" in events self._proxy_logging_obj.post_call_failure_hook.assert_not_awaited() + def _invoke_stream(self) -> list[bytes]: + import base64 + + return [ + self._converse_frame("chunk", {"bytes": base64.b64encode(json.dumps(event).encode()).decode()}) + for event in ( + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "Hello"}}, + {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 1}}, + ) + ] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("model", "request_body"), + [ + ("us.anthropic.claude-sonnet-5-v1:0", None), + ( + "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abc", + {"anthropic_version": "bedrock-2023-05-31", "max_tokens": 8, "messages": []}, + ), + ], + ids=["claude-model-id", "application-inference-profile-with-a-claude-body"], + ) + async def test_bedrock_invoke_stream_of_a_claude_model_is_relayed_and_scanned_for_its_text( + self, monkeypatch: pytest.MonkeyPatch, model: str, request_body: Mapping[str, object] | None + ): + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.llms.pass_through.guardrail_translation.handler import LlmPassthroughRouteHandler + + guardrail = self._apply_guardrail() + scanned: dict[str, object] = {} + + async def scan_relayed( + response: AsyncIterable[object], + user_api_key_dict: ProxyUserAPIKeyAuth, + request_data: dict[str, object], + guardrails: Sequence[tuple[CustomGuardrail, bool]], + ) -> AsyncGenerator[object, None]: + items = [item async for item in response] + scanned["guardrails"] = list(guardrails) + scanned["texts"] = LlmPassthroughRouteHandler().get_streaming_scan_key(items).texts + for item in items: + yield item + + stream = b"".join(self._invoke_stream()) + result, body, mock_handler, _ = await self._run_converse_stream( + monkeypatch, + guardrail, + [stream[:20], stream[20:]], + event_stream_hook=scan_relayed, + endpoint=f"model/{model}/invoke-with-response-stream", + model=model, + request_body=request_body, + ) + + assert isinstance(result, _UpstreamClosingStreamingResponse) + assert result.headers["content-type"] == "application/vnd.amazon.eventstream" + assert body == stream + assert scanned == {"guardrails": [(guardrail, False)], "texts": ("Hello",)} + mock_handler.assert_not_awaited() + + @pytest.mark.asyncio + async def test_bedrock_invoke_stream_of_a_claude_model_without_a_post_call_guardrail_streams_as_received( + self, monkeypatch: pytest.MonkeyPatch + ): + result, body, mock_handler, recorded = await self._run_converse_stream( + monkeypatch, + self._apply_guardrail(), + [b"raw-1", b"raw-2"], + endpoint="model/us.anthropic.claude-sonnet-5-v1:0/invoke-with-response-stream", + model="us.anthropic.claude-sonnet-5-v1:0", + post_call_guarded=False, + ) + + assert isinstance(result, _UpstreamClosingStreamingResponse) + assert body == b"raw-1raw-2" + assert recorded == {} + mock_handler.assert_not_awaited() + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("model", "guardrail_factory"), + [ + ("meta.llama4-maverick-17b-instruct-v1:0", "_apply_guardrail"), + ("meta.llama4-maverick-17b-instruct-v1:0", "_success_hook_guardrail"), + ], + ids=["guardrail-the-relay-scans", "guardrail-that-needs-the-whole-response"], + ) + async def test_bedrock_invoke_stream_of_a_model_without_a_stream_translation_streams_as_received( + self, monkeypatch: pytest.MonkeyPatch, model: str, guardrail_factory: str + ): + result, body, mock_handler, recorded = await self._run_converse_stream( + monkeypatch, + getattr(self, guardrail_factory)(), + [b"raw-1", b"raw-2"], + endpoint=f"model/{model}/invoke-with-response-stream", + model=model, + ) + + assert isinstance(result, _UpstreamClosingStreamingResponse) + assert body == b"raw-1raw-2" + assert recorded == {} + mock_handler.assert_not_awaited() + @pytest.mark.asyncio async def test_bedrock_invoke_stream_is_not_buffered(self, monkeypatch): processing_obj = self._build_processing_obj(