mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-18 04:44:37 +00:00
fix(bedrock)!: scan Claude invoke-with-response-stream passthrough output with post_call guardrails (#46524)
* fix(bedrock): scan Claude invoke-with-response-stream passthrough frames with post_call guardrails (LIT-9618) Invoke streams of Claude models now go through the same frame-by-frame relay as converse-stream. Each chunk frame's Anthropic event is scanned by the /v1/messages stream translation, a block ends the stream with an AWS exception frame, and a held stream gets its masks written back into re-encoded chunk frames. Claude is told by its model id or, behind an application inference profile, by the anthropic_version its request body carries. Streams of other model families stay as they were * test(bedrock): cover invoke stream vendor outage, upstream exception and release on scan (LIT-9618) * fix(bedrock): hold Claude invoke streams for guardrails the relay cannot run and mark each scan as streaming (LIT-9618) A guardrail released on scan or run in parallel left a Claude invoke-with-response-stream unscanned, while converse-stream holds the whole stream and scans it once. Claude invoke streams now take the same buffered path, with the rewrites written back into the chunk frames that carried each text. Each invoke scan also sends stream true and the events relayed so far as responses, as /v1/messages does, so a vendor can tie the scans of one streamed reply together. * refactor(bedrock): mark the request dicts the buffered stream paths hand to post_call hooks (LIT-9618) * fix(bedrock): scan a Claude invoke stream the upstream breaks off as a finished stream (LIT-9618) An invoke stream that ended on an upstream exception, without a message_delta stop_reason, was scanned as still in flight, so its tool input never reached the guardrail and a held stream released it unscanned. The scan and its scan key now close such a stream the way converse-stream treats the end of its upstream. * test(bedrock): pass the converse-stream endpoint where the tool call without input test frames its stream (LIT-9618) * test(bedrock): expect the ValidationException main now returns when a guardrail blocks a Claude invoke stream before its first frame (LIT-9618) * test(bedrock): type the Claude invoke event helpers in the passthrough guardrail tests * test(bedrock): type the remaining parameters of the Claude invoke stream tests * refactor(bedrock): bind the buffered invoke rewrite in its condition --------- Co-authored-by: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com>
This commit is contained in:
yucheng-berri
GitHub
parent
ef1cb19351
commit
9f115be88d
7 changed files with 2214 additions and 100 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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/<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/<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
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue