diff --git a/litellm/litellm_core_utils/realtime_guardrails.py b/litellm/litellm_core_utils/realtime_guardrails.py new file mode 100644 index 00000000000..382eecc5a8c --- /dev/null +++ b/litellm/litellm_core_utils/realtime_guardrails.py @@ -0,0 +1,316 @@ +import json +from collections.abc import AsyncIterator, Iterable, Mapping, Sequence +from dataclasses import dataclass +from typing import TYPE_CHECKING, Final, TypeAlias + +from pydantic import TypeAdapter, ValidationError + +from litellm._logging import verbose_logger +from litellm.exceptions import ModifyResponseException + +if TYPE_CHECKING: + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.types.guardrails import GuardrailEventHooks + +_DEFAULT_VIOLATION_MESSAGE: Final = "I'm sorry, that request was blocked by the content filter." +BLOCKED_TOOL_OUTPUT: Final = json.dumps({"error": "Tool output blocked by content policy"}) +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object]) +_JSON_ARRAY: Final = TypeAdapter(tuple[object, ...]) +_TEXTS: Final = TypeAdapter(tuple[str, ...]) +_TYPED_TEXT_PART_TYPES: Final = frozenset({"input_text", "output_text", "text"}) + + +@dataclass(frozen=True, slots=True) +class RealtimeGuardrailBlock: + message: str + ends_session_on_violation: bool + end_session_after_n_fails: int | None + + def ends_session(self, violation_count: int) -> bool: + return self.ends_session_on_violation or ( + self.end_session_after_n_fails is not None and violation_count >= self.end_session_after_n_fails + ) + + +@dataclass(frozen=True, slots=True) +class RealtimeGuardrailRewrite: + texts: tuple[str, ...] + + +RealtimeGuardrailVerdict: TypeAlias = RealtimeGuardrailBlock | RealtimeGuardrailRewrite | None + + +@dataclass(frozen=True, slots=True) +class RealtimeMessageText: + role: str + texts: tuple[str, ...] + + +@dataclass(frozen=True, slots=True) +class RealtimeToolOutput: + texts: tuple[str, ...] + sanitized_event: str + + +RealtimeGuardedItem: TypeAlias = RealtimeMessageText | RealtimeToolOutput + + +def _json_object(value: object) -> Mapping[str, object] | None: + try: + return _JSON_OBJECT.validate_python(value) + except ValidationError: + return None + + +def _turn_copy(value: object) -> object: + mapping: Final = _json_object(value) + return value if mapping is None else dict(mapping) + + +def _typed_text(part: object) -> str | None: + part_object: Final = _json_object(part) + if part_object is None or part_object.get("type") not in _TYPED_TEXT_PART_TYPES: + return None + text: Final[object] = part_object.get("text", "") + return text if isinstance(text, str) else None + + +def _content_parts(content: object) -> Sequence[object]: + return _JSON_ARRAY.validate_python(content) if isinstance(content, list) else () + + +def _typed_texts(content: object) -> tuple[str, ...]: + return tuple(text for text in map(_typed_text, _content_parts(content)) if text is not None) + + +def _with_text(part: object, text: str | None) -> object: + part_object: Final = _json_object(part) + return part if text is None or part_object is None else {**part_object, "text": text} + + +def _tool_output(client_event: Mapping[str, object], item: Mapping[str, object]) -> RealtimeToolOutput | None: + output: Final[object] = item.get("output", "") + text: Final = output if isinstance(output, str) else json.dumps(output) + if not text: + return None + sanitized_event: Final = json.dumps({**client_event, "item": {**item, "output": BLOCKED_TOOL_OUTPUT}}) + return RealtimeToolOutput(texts=(text,), sanitized_event=sanitized_event) + + +def realtime_guarded_item(client_event: Mapping[str, object]) -> RealtimeGuardedItem | None: + if client_event.get("type") != "conversation.item.create": + return None + item: Final = _json_object(client_event.get("item")) + if item is None: + return None + if item.get("type") == "function_call_output": + return _tool_output(client_event, item) + texts: Final = _typed_texts(item.get("content")) + role: Final[object] = item.get("role") + return RealtimeMessageText(role=role if isinstance(role, str) else "user", texts=texts) if any(texts) else None + + +def rewritten_item(client_event: Mapping[str, object], texts: tuple[str, ...]) -> Mapping[str, object]: + item: Final = _json_object(client_event.get("item")) or {} + if item.get("type") == "function_call_output": + return {**item, "output": " ".join(texts)} + parts: Final = _content_parts(item.get("content")) + text_positions: Final = tuple(index for index, part in enumerate(parts) if _typed_text(part) is not None) + replacements: Final = dict(zip(text_positions, texts)) + merged_away: Final = frozenset(text_positions[len(texts) :]) + return { + **item, + "content": [ + _with_text(part, replacements.get(index)) for index, part in enumerate(parts) if index not in merged_away + ], + } + + +def _violation_text(error: Exception) -> str: + detail: Final[object] = getattr(error, "detail", None) + detail_object: Final = _json_object(detail) + if detail_object is not None: + reason: Final[object] = detail_object.get("error") + return reason if isinstance(reason, str) and reason else str(error) + if detail is not None: + return str(detail) + return str(error) or _DEFAULT_VIOLATION_MESSAGE + + +def _is_guardrail_block(error: Exception) -> bool: + return hasattr(error, "status_code") or isinstance(error, (ValueError, ModifyResponseException)) + + +def _screened_texts(screened: object, texts: tuple[str, ...]) -> tuple[str, ...]: + screened_object: Final = _json_object(screened) + try: + returned: Final = _TEXTS.validate_python(screened_object.get("texts") if screened_object else None) + except ValidationError: + return texts + return returned if len(returned) == len(texts) else texts + + +async def _guardrail_outcome( + guardrail: "CustomGuardrail", texts: tuple[str, ...], request_data: Mapping[str, object] +) -> RealtimeGuardrailBlock | tuple[str, ...]: + try: + screened: Final = await guardrail.apply_guardrail( + inputs={"texts": list(texts), "images": []}, + request_data=dict(request_data), + input_type="request", + ) + except Exception as error: + if not _is_guardrail_block(error): + verbose_logger.exception("[realtime guardrail] unexpected error in apply_guardrail: %s", error) + raise + return RealtimeGuardrailBlock( + message=guardrail.realtime_violation_message or _violation_text(error), + ends_session_on_violation=guardrail.on_violation == "end_session", + end_session_after_n_fails=guardrail.end_session_after_n_fails, + ) + return _screened_texts(screened, texts) + + +async def _screen( + guardrails: Sequence["CustomGuardrail"], texts: tuple[str, ...], request_data: Mapping[str, object] +) -> RealtimeGuardrailBlock | tuple[str, ...]: + if not guardrails: + return texts + outcome: Final = await _guardrail_outcome(guardrails[0], texts, request_data) + if isinstance(outcome, RealtimeGuardrailBlock): + return outcome + return await _screen(guardrails[1:], outcome, request_data) + + +async def _screened_until_block( + guardrails: Sequence["CustomGuardrail"], texts: tuple[str, ...], request_data: Mapping[str, object] +) -> AsyncIterator[RealtimeGuardrailBlock | str]: + for text in texts: + if isinstance(outcome := await _screen(guardrails, (text,), request_data), RealtimeGuardrailBlock): + yield outcome + return + for screened in outcome: + yield screened + + +async def _screen_each( + guardrails: Sequence["CustomGuardrail"], texts: tuple[str, ...], request_data: Mapping[str, object] +) -> RealtimeGuardrailBlock | tuple[str, ...]: + outcomes: Final = tuple([outcome async for outcome in _screened_until_block(guardrails, texts, request_data)]) + block: Final = next((outcome for outcome in outcomes if isinstance(outcome, RealtimeGuardrailBlock)), None) + return block if block is not None else tuple(outcome for outcome in outcomes if isinstance(outcome, str)) + + +def _joins(texts: tuple[str, ...]) -> tuple[str, ...]: + unspaced: Final = "".join(texts) + boundaries: Final = zip(texts, texts[1:]) + word_break_unknown: Final = any(not left[-1:].isspace() and not right[:1].isspace() for left, right in boundaries) + return (unspaced, " ".join(texts)) if word_break_unknown else (unspaced,) + + +async def _rewritten_parts( + guardrails: Sequence["CustomGuardrail"], texts: tuple[str, ...], request_data: Mapping[str, object] +) -> RealtimeGuardrailBlock | tuple[str, ...]: + parts: Final = await _screen_each(guardrails, texts, request_data) + if isinstance(parts, RealtimeGuardrailBlock): + return parts + rejoined: Final = _joins(parts) + screened_rejoined: Final = await _screen_each(guardrails, rejoined, request_data) + if isinstance(screened_rejoined, RealtimeGuardrailBlock): + return screened_rejoined + masked_join: Final = next((text for text, join in zip(screened_rejoined, rejoined) if text != join), None) + return parts if masked_join is None else (masked_join,) + + +def _in_role_scope(guardrail: "CustomGuardrail", role: str) -> bool: + from litellm.llms.base_llm.guardrail_translation.utils import ( + effective_scan_only_tool_results_for_guardrail, + effective_skip_system_message_for_guardrail, + effective_skip_tool_message_for_guardrail, + role_out_of_guardrail_scope, + ) + + return not role_out_of_guardrail_scope( + role, + skip_system_message=effective_skip_system_message_for_guardrail(guardrail), + skip_tool_message=effective_skip_tool_message_for_guardrail(guardrail), + scan_only_tool_results=effective_scan_only_tool_results_for_guardrail(guardrail), + ) + + +def realtime_runner_guardrail( + callback: object, + check_data: Mapping[str, object], + event_hooks: Sequence["GuardrailEventHooks"], +) -> "CustomGuardrail | None": + from litellm.integrations.custom_guardrail import CustomGuardrail + + if not isinstance(callback, CustomGuardrail) or callback.use_native_lifecycle_hooks: + return None + runs: Final = any( + callback.should_run_guardrail(data=check_data, event_type=event_hook) for event_hook in event_hooks + ) + return callback if runs else None + + +def realtime_screening_guardrails( + role: str, + event_hooks: Sequence["GuardrailEventHooks"], + check_data: Mapping[str, object], + callbacks: Iterable[object], +) -> tuple["CustomGuardrail", ...]: + runner_guardrails: Final = (realtime_runner_guardrail(callback, check_data, event_hooks) for callback in callbacks) + guardrails: Final = { + id(guardrail): guardrail + for guardrail in runner_guardrails + if guardrail is not None and _in_role_scope(guardrail, role) + } + return tuple(guardrails.values()) + + +async def realtime_guardrail_verdict( + texts: tuple[str, ...], + role: str, + event_hooks: Sequence["GuardrailEventHooks"], + request_data: Mapping[str, object], + user_api_key_dict: object | None, + callbacks: Iterable[object], +) -> RealtimeGuardrailVerdict: + check_data: Final = {**request_data, "transcript": " ".join(texts), "stream": True} + in_scope: Final = realtime_screening_guardrails(role, event_hooks, check_data, callbacks) + guardrail_request_data: Final = { + **{key: _turn_copy(value) for key, value in request_data.items()}, + "user_api_key_dict": user_api_key_dict, + } + joined: Final = _joins(texts) + screened: Final = await _screen_each(in_scope, joined, guardrail_request_data) + if isinstance(screened, RealtimeGuardrailBlock): + return screened + if screened == joined: + return None + rewritten: Final = screened if len(texts) < 2 else await _rewritten_parts(in_scope, texts, guardrail_request_data) + if isinstance(rewritten, RealtimeGuardrailBlock): + return rewritten + return None if rewritten == texts else RealtimeGuardrailRewrite(texts=rewritten) + + +def guardrail_violation_error_event(message: str) -> str: + return json.dumps( + { + "type": "error", + "error": {"type": "guardrail_violation", "message": message, "code": "content_policy_violation"}, + } + ) + + +def guardrail_voice_prompt_events(message: str) -> tuple[str, str]: + prompt: Final = f"Say exactly the following message to the user, word for word, do not add anything else: {message}" + return ( + json.dumps( + { + "type": "conversation.item.create", + "item": {"type": "message", "role": "user", "content": [{"type": "input_text", "text": prompt}]}, + } + ), + json.dumps({"type": "response.create"}), + ) diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index b6af181980e..01aa7296b0f 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -47,13 +47,27 @@ from litellm.types.realtime import ALL_DELTA_TYPES from .litellm_logging import Logging as LiteLLMLogging from .realtime_errors import client_close_code, realtime_error_event, websocket_close_reason +from .realtime_guardrails import ( + RealtimeGuardedItem, + RealtimeGuardrailBlock, + RealtimeGuardrailRewrite, + RealtimeGuardrailVerdict, + RealtimeMessageText, + RealtimeToolOutput, + guardrail_violation_error_event, + guardrail_voice_prompt_events, + realtime_guarded_item, + realtime_guardrail_verdict, + realtime_runner_guardrail, + realtime_screening_guardrails, + rewritten_item, +) if TYPE_CHECKING: from openai.types.realtime.realtime_audio_formats import RealtimeAudioFormats from websockets.asyncio.client import ClientConnection from websockets.exceptions import ConnectionClosed - from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.types.guardrails import GuardrailEventHooks CLIENT_CONNECTION_CLASS = ClientConnection @@ -437,21 +451,6 @@ def _opted_out_of_always_on_guardrail(callback: object, request_data: Mapping[st ) -def _realtime_runner_guardrail( - callback: object, - check_data: Mapping[str, object], - event_hooks: Sequence["GuardrailEventHooks"], -) -> "CustomGuardrail | None": - from litellm.integrations.custom_guardrail import CustomGuardrail - - if not isinstance(callback, CustomGuardrail) or callback.use_native_lifecycle_hooks: - return None - runs: Final = any( - callback.should_run_guardrail(data=check_data, event_type=event_hook) for event_hook in event_hooks - ) - return callback if runs else None - - def translation_guardrails_apply(request_data: Mapping[str, object]) -> bool: from litellm.types.guardrails import GuardrailEventHooks @@ -479,7 +478,7 @@ def realtime_guardrails_apply(request_data: Mapping[str, object], callbacks: Seq check_data: Final = {**request_data, "stream": True} runner_event_hooks: Final = (GuardrailEventHooks.realtime_input_transcription, GuardrailEventHooks.pre_call) return any( - _realtime_runner_guardrail(callback, check_data, runner_event_hooks) is not None for callback in callbacks + realtime_runner_guardrail(callback, check_data, runner_event_hooks) is not None for callback in callbacks ) @@ -1663,16 +1662,60 @@ class RealTimeStreaming: async def run_realtime_guardrails( self, - transcript: str, + transcript: object, item_id: str | None = None, pre_block_backend_message: str | None = None, event_hooks: Sequence["GuardrailEventHooks"] | None = None, ) -> bool: """ - Run registered guardrails on realtime text (transcript, user message, tool output). - Returns True if blocked (synthetic warning already sent to client). Returns False if clean (caller should send response.create to the backend). + """ + from litellm.types.guardrails import GuardrailEventHooks + + if isinstance(transcript, str): + verdict: Final = await self.apply_realtime_guardrails( + (transcript,), "user", item_id, pre_block_backend_message, event_hooks + ) + return isinstance(verdict, RealtimeGuardrailBlock) + if realtime_screening_guardrails( + "user", + event_hooks or (GuardrailEventHooks.realtime_input_transcription,), + {**self.request_data, "stream": True}, + litellm.callbacks, + ): + raise TypeError(f"Realtime guardrails cannot screen a {type(transcript).__name__} transcript") + return False + + async def _screen_client_item(self, guarded_item: RealtimeGuardedItem | None) -> RealtimeGuardrailVerdict: + from litellm.types.guardrails import GuardrailEventHooks + + match guarded_item: + case RealtimeToolOutput(texts=texts, sanitized_event=sanitized_event): + return await self.apply_realtime_guardrails( + texts, + "tool", + pre_block_backend_message=sanitized_event, + event_hooks=[GuardrailEventHooks.pre_call], + ) + case RealtimeMessageText(role=role, texts=texts): + return await self.apply_realtime_guardrails(texts, role, event_hooks=[GuardrailEventHooks.pre_call]) + case None: + return None + + async def apply_realtime_guardrails( + self, + texts: tuple[str, ...], + role: str, + item_id: str | None = None, + pre_block_backend_message: str | None = None, + event_hooks: Sequence["GuardrailEventHooks"] | None = None, + ) -> RealtimeGuardrailVerdict: + """ + Run registered guardrails on realtime text (transcript, user message, tool output). + + A block has already been reported to the client and the backend when this returns. + A rewrite carries the masked texts the caller forwards in place of the originals. ``pre_block_backend_message`` (if provided) is sent to the backend BEFORE any of the guardrail's own backend messages when a block is @@ -1687,102 +1730,33 @@ class RealTimeStreaming: """ from litellm.types.guardrails import GuardrailEventHooks - if event_hooks is None: - event_hooks = [GuardrailEventHooks.realtime_input_transcription] - _realtime_event_types: Final = event_hooks - _check_data: Final = {**self.request_data, "transcript": transcript, "stream": True} - _already_run: Final[set] = set() - - runner_guardrails: Final = ( - _realtime_runner_guardrail(callback, _check_data, _realtime_event_types) for callback in litellm.callbacks + block: Final = await realtime_guardrail_verdict( + texts, + role, + event_hooks if event_hooks is not None else (GuardrailEventHooks.realtime_input_transcription,), + self.request_data, + self.user_api_key_dict, + litellm.callbacks, ) - for callback in runner_guardrails: - if callback is None or id(callback) in _already_run: - continue - _already_run.add(id(callback)) - try: - await callback.apply_guardrail( - inputs={"texts": [transcript], "images": []}, - request_data={"user_api_key_dict": self.user_api_key_dict}, - input_type="request", - ) - except Exception as e: - # Re-raise unexpected errors (no status_code/detail = programming bug, not a block). - # HTTPException and guardrail-raised exceptions have a status_code or detail attr. - is_guardrail_block = hasattr(e, "status_code") or isinstance(e, ValueError) - if not is_guardrail_block: - verbose_logger.exception( - "[realtime guardrail] unexpected error in apply_guardrail: %s", - e, - ) - raise - # Extract the human-readable error from the detail dict (HTTPException) - # or fall back to str(e) for plain ValueError. - detail: object | None = getattr(e, "detail", None) - if isinstance(detail, dict): - safe_msg = detail.get("error") or str(e) - elif detail is not None: - safe_msg = str(detail) - else: - safe_msg = str(e) or "I'm sorry, that request was blocked by the content filter." + if not isinstance(block, RealtimeGuardrailBlock): + return block + await self._send_blocked_turn_to_backend(pre_block_backend_message, item_id) + await self.websocket.send_text(guardrail_violation_error_event(block.message)) + if not self._is_transcription_session: + for voice_prompt_event in guardrail_voice_prompt_events(block.message): + await self._send_to_backend(voice_prompt_event) - # Use realtime_violation_message if configured; fall back to guardrail error text. - error_msg = getattr(callback, "realtime_violation_message", None) or safe_msg + self._violation_count += 1 + if block.ends_session(self._violation_count): + verbose_logger.warning("[realtime guardrail] ending session after violation %d", self._violation_count) + await self.backend_ws.close() - await self._send_blocked_turn_to_backend(pre_block_backend_message, item_id) - # Send the policy violation hint (shows as small gray status text in UI). - await self.websocket.send_text( - json.dumps( - { - "type": "error", - "error": { - "type": "guardrail_violation", - "message": error_msg, - "code": "content_policy_violation", - }, - } - ) - ) - if not self._is_transcription_session: - # Ask the LLM to voice the exact guardrail message so the - # user hears it as audio in voice sessions (not just text). - guardrail_prompt = ( - f"Say exactly the following message to the user, word for word, " - f"do not add anything else: {error_msg}" - ) - await self._send_to_backend( - json.dumps( - { - "type": "conversation.item.create", - "item": { - "type": "message", - "role": "user", - "content": [{"type": "input_text", "text": guardrail_prompt}], - }, - } - ) - ) - await self._send_to_backend(json.dumps({"type": "response.create"})) - - self._violation_count += 1 - end_session_after: int | None = getattr(callback, "end_session_after_n_fails", None) - should_end = getattr(callback, "on_violation", None) == "end_session" or ( - end_session_after is not None and self._violation_count >= end_session_after - ) - if should_end: - verbose_logger.warning( - "[realtime guardrail] ending session after violation %d", - self._violation_count, - ) - await self.backend_ws.close() - - verbose_logger.warning( - "[realtime guardrail] BLOCKED transcript (violation %d): %r", - self._violation_count, - transcript[:80], - ) - return True - return False + verbose_logger.warning( + "[realtime guardrail] BLOCKED transcript (violation %d): %r", + self._violation_count, + " ".join(texts)[:80], + ) + return block async def _handle_provider_config_message(self, raw_response: str) -> None: """Process a backend message when a provider_config is set (transformed path).""" @@ -1866,7 +1840,7 @@ class RealTimeStreaming: self.store_message(event_str) self._capture_transcription_usage(event) await self._send_event_to_client(event, event_str) - blocked = await self.run_realtime_guardrails(cast(str, transcript)) + blocked = await self.run_realtime_guardrails(transcript) if not blocked and self._should_disable_vad_auto_response(): await self._send_to_backend(json.dumps({"type": "response.create"})) continue @@ -2267,92 +2241,22 @@ class RealTimeStreaming: guardrail_turn_detection_injected = False msg_type: str | None = None try: - from litellm.types.guardrails import GuardrailEventHooks - msg_obj = json.loads(message) client_event = msg_obj msg_type = client_event.get("type") - if msg_type == "conversation.item.create": - # Check user text messages for prompt injection - item = client_event.get("item", {}) - # Check function_call_output first so a client cannot - # bypass the tool-result guardrail by also setting - # role="user" on a function_call_output item. - if item.get("type") == "function_call_output": - # Tool results are client-controlled and fed to the - # model; check them with the same guardrail used for - # user text so an attacker cannot smuggle blocked - # content into a function_call_output. - output: object = item.get("output", "") - output_text = output if isinstance(output, str) else json.dumps(output) - if output_text: - # Build the sanitized function_call_output up - # front so we can hand it to the guardrail - # runner as the pre-block message. Providers - # that pair every toolCall with a toolResponse - # (e.g. Gemini/Vertex Live) require the - # toolResponse to arrive BEFORE any other - # client message — otherwise the guardrail's - # own clientContent would violate the - # pending-tool-call protocol contract and the - # backend could close the connection before - # the sanitized response ever lands. Dropping - # the blocked item outright would similarly - # leave such providers waiting indefinitely. - # The sanitized payload carries no blocked - # content — only a generic policy marker. - sanitized_msg = json.dumps( - { - **msg_obj, - "item": { - **item, - "output": json.dumps( - { - "error": "Tool output blocked by content policy", - } - ), - }, - } - ) - blocked = await self.run_realtime_guardrails( - output_text, - pre_block_backend_message=sanitized_msg, - event_hooks=[GuardrailEventHooks.pre_call], - ) - if blocked: - # ``_pending_guardrail_message`` is - # intentionally NOT set here. That flag - # exists to swallow the reflexive - # ``response.create`` an OpenAI client - # sends immediately after a user text - # message. In a tool-calling flow the - # client may not send a ``response.create`` - # at all (e.g. Gemini SDKs auto-respond), - # so leaving the flag set would - # incorrectly drop an unrelated - # ``response.create`` from a later - # interaction turn. - continue - elif item.get("role") == "user": - content_list: Sequence[object] = item.get("content", []) - texts = [ - c.get("text", "") - for c in content_list - if isinstance(c, dict) and c.get("type") == "input_text" - ] - combined_text = " ".join(texts) - if combined_text: - blocked = await self.run_realtime_guardrails( - combined_text, - event_hooks=[GuardrailEventHooks.pre_call], - ) - if blocked: - # Store the guardrail reason so the next response.create - # (sent automatically by the client) is rewritten to - # include it as response instructions. - self._pending_guardrail_message = combined_text - continue # don't forward the original blocked message + self._pending_guardrail_message = None + + match await self._screen_client_item(guarded_item := realtime_guarded_item(client_event)): + case RealtimeGuardrailBlock(): + if isinstance(guarded_item, RealtimeMessageText): + self._pending_guardrail_message = " ".join(guarded_item.texts) + continue + case RealtimeGuardrailRewrite(texts=masked_texts): + msg_obj["item"] = rewritten_item(client_event, masked_texts) + message = json.dumps(msg_obj) + case None: + pass if msg_type == "response.create" and self._pending_guardrail_message: # The guardrail already sent the synthetic AI bubble — drop this diff --git a/litellm/llms/bedrock/realtime/handler.py b/litellm/llms/bedrock/realtime/handler.py index 49939399ec8..9f4fb8cfecc 100644 --- a/litellm/llms/bedrock/realtime/handler.py +++ b/litellm/llms/bedrock/realtime/handler.py @@ -6,12 +6,14 @@ This uses aws_sdk_bedrock_runtime for bidirectional streaming with Nova Sonic. import asyncio import contextlib +import functools import importlib.metadata import json -from collections.abc import AsyncIterator, Awaitable, Callable, Mapping, MutableMapping, Sequence +from collections.abc import AsyncIterator, Awaitable, Callable, Iterable, Mapping, MutableMapping, Sequence from dataclasses import dataclass, field +from enum import Enum, auto from types import MappingProxyType -from typing import Final, NoReturn, Protocol, runtime_checkable +from typing import TYPE_CHECKING, Final, NoReturn, Protocol, runtime_checkable from pydantic import JsonValue, TypeAdapter @@ -24,11 +26,25 @@ from litellm.constants import ( BEDROCK_REALTIME_SDK_DISTRIBUTION, BEDROCK_REALTIME_SDK_SUPPORTED_RANGE, BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY, + REALTIME_SESSION_ENDED_ABNORMALLY_KEY, REALTIME_SESSION_SUCCESS_LOGGED_KEY, ) from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +from litellm.litellm_core_utils.realtime_errors import realtime_error_event +from litellm.litellm_core_utils.realtime_guardrails import ( + RealtimeGuardedItem, + RealtimeGuardrailBlock, + RealtimeGuardrailRewrite, + RealtimeGuardrailVerdict, + RealtimeToolOutput, + guardrail_violation_error_event, + guardrail_voice_prompt_events, + realtime_guarded_item, + realtime_guardrail_verdict, + rewritten_item, +) from litellm.litellm_core_utils.realtime_streaming import DefaultLoggedRealTimeEventTypes from litellm.llms.base_llm.realtime.transformation import RealtimeSessionUpdateRejection from litellm.types.llms.bedrock import AwsAuthParams @@ -40,9 +56,13 @@ from ..base_aws_llm import BaseAWSLLM, run_aws_signing from ..common_utils import BedrockError from .transformation import BedrockRealtimeConfig, requested_modalities +if TYPE_CHECKING: + from litellm.types.guardrails import GuardrailEventHooks + _CLIENT_MESSAGE_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) _REQUEST_MAPPING_ADAPTER: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter(Mapping[str, object] | None) _SPEND_LOG_METADATA_KEYS: Final = ("litellm_metadata", "metadata") +GUARDRAIL_UNAVAILABLE_MESSAGE: Final = "A guardrail could not check this request, so the session was ended" _EMPTY_JSON_OBJECT: Final[Mapping[str, JsonValue]] = MappingProxyType({}) @@ -259,11 +279,11 @@ def _raise_provider_failure(scope: MutableMapping[str, object], failure: BaseExc raise error from failure -def _parse_client_message(message: str) -> Mapping[str, JsonValue]: +def _parse_client_message(message: str) -> Mapping[str, JsonValue] | None: try: return _json_dict(_CLIENT_MESSAGE_ADAPTER.validate_json(message)) except ValueError: - return _EMPTY_JSON_OBJECT + return None async def _send_opening_message( @@ -282,6 +302,27 @@ async def _send_opening_message( client_ws.scope[BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY] = True # rebind-ok: scope outlives the attempt +async def _forward_client_message( + client_ws: RealtimeClientWebSocket, + bedrock_stream: BedrockBidirectionalStream, + transformation_config: BedrockRealtimeConfig, + model: str, + session_state: RealtimeResponseTransformInput, + send_to_bedrock: Callable[[str], Awaitable[None]], + client_message: str, +) -> None: + transformed_messages: Final = transformation_config.transform_realtime_request( + message=client_message, + model=model, + session_configuration_request=session_state.get("session_configuration_request"), + ) + if transformed_messages and not _is_session_committed(client_ws.scope): + await _send_opening_message(client_ws, bedrock_stream, client_message, transformed_messages, send_to_bedrock) + return + for bedrock_message in transformed_messages: + await send_to_bedrock(bedrock_message) + + async def _ack_session_update( client_ws: RealtimeClientWebSocket, transformation_config: BedrockRealtimeConfig, @@ -317,6 +358,161 @@ async def _reject_session_update( return True +async def _send_guardrail_violation(client_ws: RealtimeClientWebSocket, block: RealtimeGuardrailBlock) -> bool: + try: + await client_ws.send_text(guardrail_violation_error_event(block.message)) + except Exception as e: # noqa: BLE001 # any send failure means the client is gone + verbose_proxy_logger.debug("Bedrock Realtime: guardrail error not delivered: %s", e, exc_info=True) + return False + return True + + +async def _end_guarded_session(client_ws: RealtimeClientWebSocket) -> None: + verbose_proxy_logger.warning("Bedrock Realtime: ending session after a guardrail violation") + with contextlib.suppress(Exception): + await client_ws.close() + + +async def _end_unscreened_session(client_ws: RealtimeClientWebSocket) -> None: + with contextlib.suppress(Exception): + await client_ws.send_text(realtime_error_event(GUARDRAIL_UNAVAILABLE_MESSAGE, error_type="guardrail_error")) + with contextlib.suppress(Exception): + await client_ws.close() + + +@dataclass(frozen=True, slots=True) +class _GuardrailUnavailable: + pass + + +def _registered_callbacks() -> Iterable[object]: + return litellm.callbacks + + +class _ClientEventRoute(Enum): + FORWARD = auto() + DROP = auto() + END_SESSION = auto() + + +class BedrockRealtimeGuardrails: + """One Bedrock realtime session's guardrail checks, shared by both bridge directions. + + Nova Sonic has no way to cancel a response or remove a turn, so a blocked user transcript ends + the session before that turn's answer reaches the client. + """ + + def __init__( + self, + request_data: Mapping[str, object] = _EMPTY_JSON_OBJECT, + user_api_key_dict: object | None = None, + callbacks: Callable[[], Iterable[object]] = _registered_callbacks, + ) -> None: + self._request_data: Final = request_data + self._user_api_key_dict: Final = user_api_key_dict + self._callbacks: Final = callbacks + self._violation_count = 0 + self._drop_next_response_create = False + self._failed = False + + @property + def failed(self) -> bool: + return self._failed + + async def _verdict( + self, texts: tuple[str, ...], role: str, hook: "GuardrailEventHooks" + ) -> RealtimeGuardrailVerdict | _GuardrailUnavailable: + try: + return await realtime_guardrail_verdict( + texts, role, (hook,), self._request_data, self._user_api_key_dict, self._callbacks() + ) + except Exception as e: # noqa: BLE001 # a guardrail that cannot decide fails closed, never as a Bedrock failure + verbose_proxy_logger.warning("Bedrock Realtime: a guardrail failed, ending the session: %s", e) + self._failed = True + return _GuardrailUnavailable() + + def _ends_session(self, block: RealtimeGuardrailBlock) -> bool: + self._violation_count += 1 + return block.ends_session(self._violation_count) + + async def screen_client_event( + self, + client_event: Mapping[str, JsonValue], + client_ws: RealtimeClientWebSocket, + forward_to_bedrock: Callable[[str], Awaitable[None]], + ) -> _ClientEventRoute: + from litellm.types.guardrails import GuardrailEventHooks + + event_type: Final = client_event.get("type") + if event_type == "response.create" and self._drop_next_response_create: + self._drop_next_response_create = False + return _ClientEventRoute.DROP + if event_type == "conversation.item.create": + self._drop_next_response_create = False + guarded_item: Final = realtime_guarded_item(client_event) + if guarded_item is None: + return _ClientEventRoute.FORWARD + held: Final = json.dumps(client_event) + client_ws.scope[BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY] = held # rebind-ok: replayed on fallback + role: Final = "tool" if isinstance(guarded_item, RealtimeToolOutput) else "user" + match await self._verdict(guarded_item.texts, role, GuardrailEventHooks.pre_call): + case None: + return _ClientEventRoute.FORWARD + case RealtimeGuardrailRewrite(texts=texts): + await forward_to_bedrock(json.dumps({**client_event, "item": rewritten_item(client_event, texts)})) + return _ClientEventRoute.DROP + case RealtimeGuardrailBlock() as block: + return await self._route_blocked_item(block, guarded_item, client_ws, forward_to_bedrock) + case _GuardrailUnavailable(): + await _end_unscreened_session(client_ws) + return _ClientEventRoute.END_SESSION + + async def _route_blocked_item( + self, + block: RealtimeGuardrailBlock, + guarded_item: RealtimeGuardedItem, + client_ws: RealtimeClientWebSocket, + forward_to_bedrock: Callable[[str], Awaitable[None]], + ) -> _ClientEventRoute: + if isinstance(guarded_item, RealtimeToolOutput): + await forward_to_bedrock(guarded_item.sanitized_event) + if not await _send_guardrail_violation(client_ws, block): + return _ClientEventRoute.END_SESSION + if self._ends_session(block): + await _end_guarded_session(client_ws) + return _ClientEventRoute.END_SESSION + if isinstance(guarded_item, RealtimeToolOutput): + return _ClientEventRoute.DROP + for voice_prompt_event in guardrail_voice_prompt_events(f'"{block.message}"'): + await forward_to_bedrock(voice_prompt_event) + self._drop_next_response_create = True + return _ClientEventRoute.DROP + + async def continues_after( + self, provider_event: Mapping[str, object], client_ws: RealtimeClientWebSocket, client_audio_streamed: bool + ) -> bool: + from litellm.types.guardrails import GuardrailEventHooks + + transcript: Final = provider_event.get("transcript") + if ( + not client_audio_streamed + or provider_event.get("type") != "conversation.item.input_audio_transcription.completed" + or not isinstance(transcript, str) + or not transcript + ): + return True + match await self._verdict((transcript,), "user", GuardrailEventHooks.realtime_input_transcription): + case RealtimeGuardrailBlock() as block: + if await _send_guardrail_violation(client_ws, block): + await _end_guarded_session(client_ws) + return False + case _GuardrailUnavailable(): + await _end_unscreened_session(client_ws) + return False + case None | RealtimeGuardrailRewrite(): + return True + + async def _close_bedrock_input( bedrock_stream: BedrockBidirectionalStream, close_messages: Sequence[str], @@ -381,6 +577,8 @@ class BedrockRealtime(BaseAWSLLM): aws_external_id: str | None = None, aws_session_tags: object = None, drop_params: bool = False, + user_api_key_dict: object | None = None, + litellm_metadata: Mapping[str, object] | None = None, **kwargs: object, ): """ @@ -461,7 +659,15 @@ class BedrockRealtime(BaseAWSLLM): try: await self._run_session( - websocket, open_bidirectional_stream, model, logging_obj, pending_opening_message, drop_params + websocket, + open_bidirectional_stream, + model, + logging_obj, + pending_opening_message, + drop_params, + BedrockRealtimeGuardrails( + request_data={"litellm_metadata": litellm_metadata or {}}, user_api_key_dict=user_api_key_dict + ), ) finally: await _close_bedrock_client(bedrock_client) @@ -474,6 +680,7 @@ class BedrockRealtime(BaseAWSLLM): logging_obj: LiteLLMLogging, pending_opening_message: str | None, drop_params: bool = False, + guardrails: BedrockRealtimeGuardrails | None = None, ) -> None: transformation_config: Final = BedrockRealtimeConfig(drop_params=drop_params) @@ -496,6 +703,7 @@ class BedrockRealtime(BaseAWSLLM): "session_configuration_request": None, } + session_guardrails: Final = guardrails if guardrails is not None else BedrockRealtimeGuardrails() outcome: Final = await self._bridge( websocket, bedrock_stream, @@ -504,7 +712,10 @@ class BedrockRealtime(BaseAWSLLM): session_state, logging_obj, initial_message=pending_opening_message, + guardrails=session_guardrails, ) + if session_guardrails.failed: + logging_obj.model_call_details[REALTIME_SESSION_ENDED_ABNORMALLY_KEY] = True logged_events: Final = ( *outcome.logged_events, @@ -544,10 +755,12 @@ class BedrockRealtime(BaseAWSLLM): session_state: RealtimeResponseTransformInput, logging_obj: LiteLLMLogging, initial_message: str | None, + guardrails: BedrockRealtimeGuardrails | None = None, ) -> _BridgeOutcome: """Run both forwarding directions until the client leaves or either side fails.""" logged: Final[list[OpenAIRealtimeEvents]] = [] # mutable-ok: events forwarded before a failure are still spend signals: Final = SessionEndSignals() + session_guardrails: Final = guardrails if guardrails is not None else BedrockRealtimeGuardrails() async def collect_logged_events() -> None: try: @@ -559,6 +772,7 @@ class BedrockRealtime(BaseAWSLLM): logging_obj, session_state, signals=signals, + guardrails=session_guardrails, ): logged.append(event) except Exception: @@ -575,6 +789,7 @@ class BedrockRealtime(BaseAWSLLM): logging_obj, initial_message, signals=signals, + guardrails=session_guardrails, ) ) bedrock_task: Final = asyncio.create_task(collect_logged_events()) @@ -610,6 +825,7 @@ class BedrockRealtime(BaseAWSLLM): initial_message: str | None = None, *, signals: SessionEndSignals, + guardrails: BedrockRealtimeGuardrails | None = None, ) -> None: """Forward messages from client WebSocket to Bedrock stream. @@ -631,10 +847,23 @@ class BedrockRealtime(BaseAWSLLM): raise _stream_ended_before_session_closed() from e verbose_proxy_logger.debug("Bedrock Realtime: Sent to Bedrock: %s", bedrock_message[:200]) + forward_to_bedrock: Final = functools.partial( + _forward_client_message, + client_ws, + bedrock_stream, + transformation_config, + model, + session_state, + send_to_bedrock, + ) + session_guardrails: Final = guardrails if guardrails is not None else BedrockRealtimeGuardrails() try: async for message in _client_messages(client_ws, initial_message): verbose_proxy_logger.debug("Bedrock Realtime: Received from client: %s", message[:200]) parsed_client_message = _parse_client_message(message) + if parsed_client_message is None: + verbose_proxy_logger.warning("Bedrock Realtime: dropped a client message that is not valid JSON") + continue is_session_update = _json_str(parsed_client_message.get("type")) == "session.update" if is_session_update: rejection = transformation_config.session_update_rejection(parsed_client_message) @@ -643,18 +872,15 @@ class BedrockRealtime(BaseAWSLLM): break continue - transformed_messages = transformation_config.transform_realtime_request( - message=message, - model=model, - session_configuration_request=session_state.get("session_configuration_request"), - ) - if transformed_messages and not _is_session_committed(client_ws.scope): - await _send_opening_message( - client_ws, bedrock_stream, message, transformed_messages, send_to_bedrock - ) - else: - for bedrock_message in transformed_messages: - await send_to_bedrock(bedrock_message) + match await session_guardrails.screen_client_event( + parsed_client_message, client_ws, forward_to_bedrock + ): + case _ClientEventRoute.END_SESSION: + break + case _ClientEventRoute.DROP: + continue + case _ClientEventRoute.FORWARD: + await forward_to_bedrock(message) if is_session_update and not await _ack_session_update( client_ws, transformation_config, model, logging_obj, parsed_client_message @@ -679,6 +905,7 @@ class BedrockRealtime(BaseAWSLLM): session_state: RealtimeResponseTransformInput, *, signals: SessionEndSignals, + guardrails: BedrockRealtimeGuardrails | None = None, ) -> AsyncIterator[OpenAIRealtimeEvents]: """Forward messages from Bedrock to the client, yielding the ones to record for spend logging. @@ -694,6 +921,7 @@ class BedrockRealtime(BaseAWSLLM): verbose_proxy_logger.debug("Bedrock Realtime: Sent to client: %s", message_json[:200]) return True + session_guardrails: Final = guardrails if guardrails is not None else BedrockRealtimeGuardrails() output: Final = await bedrock_stream.await_output() while True: result = await output[1].receive() @@ -750,3 +978,7 @@ class BedrockRealtime(BaseAWSLLM): return if _should_log_event(openai_message): yield openai_message + if not await session_guardrails.continues_after( + openai_message, client_ws, transformation_config.client_audio_streamed + ): + return diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index ba472007175..0182c27fac4 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -16613,7 +16613,7 @@ "type": "null" } ], - "description": "For /v1/realtime sessions: 'warn' speaks the violation message and continues; 'end_session' speaks the message and closes the connection. For guardrail='mcp_security': 'block' rejects the request; 'alert' only logs a warning.", + "description": "For /v1/realtime sessions: 'warn' speaks the violation message and continues; 'end_session' speaks the message and closes the connection. Bedrock Nova Sonic cannot cancel a spoken turn's answer, so a blocked spoken turn always ends its session. For guardrail='mcp_security': 'block' rejects the request; 'alert' only logs a warning.", "title": "On Violation" }, "only_scan_new_messages": { @@ -16689,7 +16689,7 @@ "type": "null" } ], - "description": "The message the bot speaks aloud when a /v1/realtime guardrail fires. Falls back to violation_message_template if not set.", + "description": "The message the bot speaks aloud when a /v1/realtime guardrail fires. Falls back to the guardrail's block reason if not set.", "title": "Realtime Violation Message" }, "run_in_parallel": { @@ -19303,7 +19303,7 @@ "type": "null" } ], - "description": "For /v1/realtime sessions: 'warn' speaks the violation message and continues; 'end_session' speaks the message and closes the connection. For guardrail='mcp_security': 'block' rejects the request; 'alert' only logs a warning.", + "description": "For /v1/realtime sessions: 'warn' speaks the violation message and continues; 'end_session' speaks the message and closes the connection. Bedrock Nova Sonic cannot cancel a spoken turn's answer, so a blocked spoken turn always ends its session. For guardrail='mcp_security': 'block' rejects the request; 'alert' only logs a warning.", "title": "On Violation" }, "only_scan_new_messages": { @@ -19724,7 +19724,7 @@ "type": "null" } ], - "description": "The message the bot speaks aloud when a /v1/realtime guardrail fires. Falls back to violation_message_template if not set.", + "description": "The message the bot speaks aloud when a /v1/realtime guardrail fires. Falls back to the guardrail's block reason if not set.", "title": "Realtime Violation Message" }, "rules": { diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index be1df4aa0f6..85d87554023 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -149,7 +149,7 @@ def _build_litellm_metadata(kwargs: Mapping[str, object]) -> dict[str, object]: _GUARDRAIL_OPT_OUT_METADATA_KEYS: Final = ("user_api_key_metadata", "user_api_key_team_metadata") -def _translation_litellm_metadata(kwargs: Mapping[str, object]) -> Mapping[str, object]: +def _session_litellm_metadata(kwargs: Mapping[str, object]) -> Mapping[str, object]: request_metadata: Final = _OBJECT_MAPPING.validate_python(kwargs.get("metadata") or {}) guardrail_opt_outs: Final = { key: request_metadata[key] for key in _GUARDRAIL_OPT_OUT_METADATA_KEYS if key in request_metadata @@ -218,7 +218,7 @@ def _translation_rejection_error( def _raise_if_http_translation_rejected(custom_llm_provider: str, model: str, kwargs: Mapping[str, object]) -> None: rejection: Final = _translation_rejection( - custom_llm_provider, _HTTP_TRANSLATION_PROVIDERS, _translation_litellm_metadata(kwargs) + custom_llm_provider, _HTTP_TRANSLATION_PROVIDERS, _session_litellm_metadata(kwargs) ) if rejection is None: return @@ -858,9 +858,7 @@ async def _arealtime( ) session_litellm_metadata: Final = ( - dict(_translation_litellm_metadata(kwargs)) - if realtime_mode == "translation" - else _build_litellm_metadata(kwargs) + _build_litellm_metadata(kwargs) if live_session_start is not None else dict(_session_litellm_metadata(kwargs)) ) translation_rejection: Final = ( _translation_rejection(_custom_llm_provider, _WEBSOCKET_TRANSLATION_PROVIDERS, session_litellm_metadata) @@ -1014,6 +1012,8 @@ async def _arealtime( aws_external_id=aws_external_id, aws_session_tags=aws_session_tags, drop_params=drop_params, + user_api_key_dict=kwargs.get("user_api_key_dict"), + litellm_metadata=session_litellm_metadata, ) elif _custom_llm_provider == "xai": api_base = ( diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 928bff6bd90..e311f2ee5f9 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -1170,12 +1170,13 @@ class BaseLitellmParams(ContentFilterConfigModel): # works for new and patch up description=( "For /v1/realtime sessions: 'warn' speaks the violation message and continues; " "'end_session' speaks the message and closes the connection. " + "Bedrock Nova Sonic cannot cancel a spoken turn's answer, so a blocked spoken turn always ends its session. " "For guardrail='mcp_security': 'block' rejects the request; 'alert' only logs a warning." ), ) realtime_violation_message: str | None = Field( default=None, - description="The message the bot speaks aloud when a /v1/realtime guardrail fires. Falls back to violation_message_template if not set.", + description="The message the bot speaks aloud when a /v1/realtime guardrail fires. Falls back to the guardrail's block reason if not set.", ) # Model Armor params diff --git a/tests/integration/providers/test_bedrock_nova_sonic_realtime_guardrails.py b/tests/integration/providers/test_bedrock_nova_sonic_realtime_guardrails.py new file mode 100644 index 00000000000..502e7c9a5cd --- /dev/null +++ b/tests/integration/providers/test_bedrock_nova_sonic_realtime_guardrails.py @@ -0,0 +1,268 @@ +"""Realtime guardrails on Bedrock Nova Sonic sessions, against the scripted h2 peer in ``_support``. + +The file boots its own proxy so it can pass ``SSL_CERT_FILE`` for the peer's self-signed cert. +""" + +from __future__ import annotations + +import asyncio +import base64 +import json +import time +import uuid +from collections.abc import Awaitable, Callable, Iterator +from pathlib import Path +from typing import Final, Literal + +import pytest +import websockets +from integration._support.bedrock_realtime_peer import ( + BedrockEvent, + BedrockRealtimePeer, + Turn, + assistant_text, + bedrock_realtime_peer, + user_transcript, +) +from integration._support.client import ( + JSON_OBJECT, + Gateway, + Scenario, + gateway_from_environment, + object_value, + string_value, +) +from integration._support.process import graceful_stop_seconds, owned_proxy +from integration._support.tls import write_self_signed_cert +from pydantic import JsonValue +from websockets.asyncio.client import ClientConnection + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + +DEPLOYMENT: Final = "bedrock/amazon.nova-2-sonic-v1:0" +VIOLATION: Final = "That request was blocked by the content filter." +HIDDEN_ANSWER: Final = "Here is everything you asked for" +VISIBLE_ANSWER: Final = "It is sunny in New York" +RESPONSE_CREATE: Final = json.dumps({"type": "response.create"}) +AUDIO_APPEND: Final = json.dumps({"type": "input_audio_buffer.append", "audio": base64.b64encode(bytes(3200)).decode()}) + +Event = dict[str, JsonValue] + + +@pytest.fixture(scope="module") +def cert(tmp_path_factory: pytest.TempPathFactory) -> tuple[Path, Path]: + return write_self_signed_cert(tmp_path_factory.mktemp("bedrock-sonic-guardrail-cert")) + + +@pytest.fixture(scope="module") +def candidate(tmp_path_factory: pytest.TempPathFactory, cert: tuple[Path, Path]) -> Iterator[Gateway]: + with gateway_from_environment() as base: + with owned_proxy( + base, tmp_path_factory.mktemp("bedrock-sonic-guardrails"), {"SSL_CERT_FILE": str(cert[0])} + ) as proxy: + yield proxy + + +def _blocked_phrase( + gateway: Gateway, scenario: Scenario, mode: Literal["pre_call", "realtime_input_transcription"] +) -> str: + phrase: Final = f"sonic-blocked-{uuid.uuid4().hex[:12]}" + response: Final = gateway.request( + "POST", + "/guardrails", + { + "guardrail": { + "guardrail_name": f"sonic-guardrail-{uuid.uuid4().hex[:12]}", + "litellm_params": { + "guardrail": "litellm_content_filter", + "mode": mode, + "default_on": True, + "realtime_violation_message": VIOLATION, + "blocked_words": [{"keyword": phrase, "action": "BLOCK"}], + }, + } + }, + ) + assert response.status_code == 200, response.text + guardrail_id: Final = JSON_OBJECT.validate_json(response.content).get("guardrail_id") + assert isinstance(guardrail_id, str), response.text + scenario.cleanups.callback(gateway.request, "DELETE", f"/guardrails/{guardrail_id}") + return phrase + + +def _deployment(gateway: Gateway, scenario: Scenario, peer: BedrockRealtimePeer) -> str: + name: Final = f"bedrock-sonic-guardrails-{uuid.uuid4().hex[:12]}" + created: Final = gateway.post( + "/model/new", + { + "model_name": name, + "litellm_params": { + "model": DEPLOYMENT, + "aws_region_name": "us-east-1", + "aws_access_key_id": "integration-dummy-key", + "aws_secret_access_key": "integration-dummy-secret", + "api_base": peer.api_base, + }, + "model_info": {"mode": "realtime"}, + }, + ) + scenario.cleanups.callback(scenario.delete_model, string_value(object_value(created["model_info"])["id"])) + return name + + +def _user_text(text: str) -> str: + return json.dumps( + { + "type": "conversation.item.create", + "item": {"type": "message", "role": "user", "content": [{"type": "input_text", "text": text}]}, + } + ) + + +async def _events_until(socket: ClientConnection, last: Callable[[Event], bool]) -> tuple[Event, ...]: + collected: Final[list[Event]] = [] + deadline: Final = time.monotonic() + 60 + while not collected or not last(collected[-1]): + try: + raw: Final = await asyncio.wait_for(socket.recv(), deadline - time.monotonic()) + except (TimeoutError, asyncio.TimeoutError): + raise AssertionError(f"timed out waiting; types={[event.get('type') for event in collected]}") + collected.append(JSON_OBJECT.validate_json(raw)) + return tuple(collected) + + +async def _events_until_closed(socket: ClientConnection) -> tuple[Event, ...]: + return tuple([JSON_OBJECT.validate_json(raw) async for raw in socket]) + + +def _is_done(event: Event) -> bool: + return event.get("type") == "response.done" + + +def _is_guardrail_error(event: Event) -> bool: + error: Final = event.get("error") + return event.get("type") == "error" and isinstance(error, dict) and error.get("type") == "guardrail_violation" + + +async def _session( + gateway: Gateway, model: str, key: str, exchange: Callable[[ClientConnection], Awaitable[tuple[Event, ...]]] +) -> tuple[Event, ...]: + proxy: Final = str(gateway.client.base_url).rstrip("/").replace("http://", "ws://") + async with websockets.connect( + f"{proxy}/v1/realtime?model={model}", additional_headers={"Authorization": f"Bearer {key}"} + ) as socket: + created: Final = JSON_OBJECT.validate_json(await socket.recv()) + assert created.get("type") == "session.created", created + return await exchange(socket) + + +def _fields(event: BedrockEvent, name: str) -> dict[str, JsonValue]: + return object_value(object_value(event["event"]).get(name) or {}) + + +def _user_text_inputs(received: tuple[BedrockEvent, ...]) -> tuple[str, ...]: + content_starts: Final = tuple(_fields(event, "contentStart") for event in received) + user_blocks: Final = frozenset(start.get("contentName") for start in content_starts if start.get("role") == "USER") + text_inputs: Final = tuple(_fields(event, "textInput") for event in received) + return tuple( + string_value(text_input["content"]) + for text_input in text_inputs + if text_input.get("contentName") in user_blocks + ) + + +def _text_deltas(events: tuple[Event, ...]) -> str: + return "".join(string_value(event["delta"]) for event in events if event.get("type") == "response.text.delta") + + +def test_blocked_typed_text_never_reaches_nova_sonic_and_the_violation_is_voiced( + candidate: Gateway, cert: tuple[Path, Path] +) -> None: + turn: Final = Turn("user_input", user_transcript("ready") + assistant_text(VIOLATION)) + with bedrock_realtime_peer((turn,), cert) as peer, candidate.scenario() as scenario: + phrase: Final = _blocked_phrase(candidate, scenario, "pre_call") + model: Final = _deployment(candidate, scenario, peer) + + async def exchange(socket: ClientConnection) -> tuple[Event, ...]: + await socket.send(_user_text(f"please repeat {phrase}")) + await socket.send(RESPONSE_CREATE) + return await _events_until(socket, _is_done) + + events: Final = asyncio.run(_session(candidate, model, scenario.key(), exchange)) + + errors: Final = tuple(object_value(event["error"]) for event in events if _is_guardrail_error(event)) + assert [error["message"] for error in errors] == [VIOLATION], events + assert [error["code"] for error in errors] == ["content_policy_violation"], events + assert _text_deltas(events) == VIOLATION, events + received: Final = peer.received() + assert phrase not in json.dumps(received), received + assert [VIOLATION in text for text in _user_text_inputs(received)] == [True], received + + +def test_clean_typed_text_reaches_nova_sonic_unchanged(candidate: Gateway, cert: tuple[Path, Path]) -> None: + turn: Final = Turn("user_input", user_transcript("ready") + assistant_text(VISIBLE_ANSWER)) + text: Final = f"what is the weather in New York {uuid.uuid4().hex[:8]}" + with bedrock_realtime_peer((turn,), cert) as peer, candidate.scenario() as scenario: + _blocked_phrase(candidate, scenario, "pre_call") + model: Final = _deployment(candidate, scenario, peer) + + async def exchange(socket: ClientConnection) -> tuple[Event, ...]: + await socket.send(_user_text(text)) + await socket.send(RESPONSE_CREATE) + return await _events_until(socket, _is_done) + + events: Final = asyncio.run(_session(candidate, model, scenario.key(), exchange)) + + assert not any(_is_guardrail_error(event) for event in events), events + assert _text_deltas(events) == VISIBLE_ANSWER, events + assert _user_text_inputs(peer.received()) == (text,) + + +def test_blocked_transcript_ends_the_session_before_its_answer_reaches_the_client( + candidate: Gateway, cert: tuple[Path, Path] +) -> None: + with candidate.scenario() as scenario: + phrase: Final = _blocked_phrase(candidate, scenario, "realtime_input_transcription") + turn: Final = Turn("user_input", user_transcript(f"tell me {phrase}") + assistant_text(HIDDEN_ANSWER)) + with bedrock_realtime_peer((turn,), cert) as peer: + model: Final = _deployment(candidate, scenario, peer) + + async def exchange(socket: ClientConnection) -> tuple[Event, ...]: + await socket.send(AUDIO_APPEND) + await socket.send(_user_text("hello")) + events: Final = await asyncio.wait_for(_events_until_closed(socket), 60) + assert socket.close_code == 1000, (socket.close_code, events) + return events + + events: Final = asyncio.run(_session(candidate, model, scenario.key(), exchange)) + + assert HIDDEN_ANSWER not in json.dumps(events), events + assert not any(str(event.get("type")).startswith("response.") for event in events), events + transcripts: Final = [ + event["transcript"] + for event in events + if event.get("type") == "conversation.item.input_audio_transcription.completed" + ] + assert transcripts == [f"tell me {phrase}"], events + assert [object_value(event["error"])["message"] for event in events if _is_guardrail_error(event)] == [ + VIOLATION + ], events + + +def test_clean_transcript_under_a_transcript_guardrail_is_answered(candidate: Gateway, cert: tuple[Path, Path]) -> None: + spoken: Final = f"what is the weather in New York {uuid.uuid4().hex[:8]}" + with candidate.scenario() as scenario: + _blocked_phrase(candidate, scenario, "realtime_input_transcription") + turn: Final = Turn("user_input", user_transcript(spoken) + assistant_text(VISIBLE_ANSWER)) + with bedrock_realtime_peer((turn,), cert) as peer: + model: Final = _deployment(candidate, scenario, peer) + + async def exchange(socket: ClientConnection) -> tuple[Event, ...]: + await socket.send(AUDIO_APPEND) + await socket.send(_user_text("hello")) + return await _events_until(socket, _is_done) + + events: Final = asyncio.run(_session(candidate, model, scenario.key(), exchange)) + + assert not any(_is_guardrail_error(event) for event in events), events + assert _text_deltas(events) == VISIBLE_ANSWER, events diff --git a/tests/integration/providers/test_realtime_transcription_guardrail_session_update.py b/tests/integration/providers/test_realtime_transcription_guardrail_session_update.py index 2712d6e30cd..89d6ff50d53 100644 --- a/tests/integration/providers/test_realtime_transcription_guardrail_session_update.py +++ b/tests/integration/providers/test_realtime_transcription_guardrail_session_update.py @@ -1591,13 +1591,6 @@ def _assert_voice_session_ungated(session: Session, observed: tuple[dict[str, Js assert _session_updates(observed) == (VOICE_UPDATE,), _session_updates(observed) -@pytest.mark.xfail( - strict=True, - reason=( - "the realtime route hands the transcript guardrail only the request's guardrails list, so the key and " - "team metadata that opts out of a default-on guardrail never reaches it" - ), -) def test_key_opted_out_of_the_transcript_guardrail_is_not_gated_by_the_prompt_guardrail( guardrail_proxy: OwnedProxy, ) -> None: @@ -1609,13 +1602,6 @@ def test_key_opted_out_of_the_transcript_guardrail_is_not_gated_by_the_prompt_gu _assert_voice_session_ungated(session, _observed(guardrail_proxy.gateway.upstream_url, handle.scenario_id)) -@pytest.mark.xfail( - strict=True, - reason=( - "the realtime route hands the transcript guardrail only the request's guardrails list, so the key and " - "team metadata that opts out of a default-on guardrail never reaches it" - ), -) def test_team_opted_out_of_the_transcript_guardrail_is_not_gated(guardrail_proxy: OwnedProxy) -> None: with guardrail_proxy.gateway.scenario() as scenario: handle: Final = _scripted(scenario, _voice_scenario(BLOCKED_TRANSCRIPT)) diff --git a/tests/unit/litellm_core_utils/test_realtime_guardrails.py b/tests/unit/litellm_core_utils/test_realtime_guardrails.py new file mode 100644 index 00000000000..26ab2887883 --- /dev/null +++ b/tests/unit/litellm_core_utils/test_realtime_guardrails.py @@ -0,0 +1,481 @@ +import json +import sys +from typing import Final, Literal + +import pytest +from fastapi import HTTPException + +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging +from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ContentFilterGuardrail +from litellm.proxy.guardrails.guardrail_hooks.custom_code.custom_code_guardrail import CustomCodeGuardrail +from litellm.litellm_core_utils.realtime_guardrails import ( + RealtimeGuardrailBlock, + RealtimeGuardrailRewrite, + RealtimeMessageText, + realtime_guarded_item, + realtime_guardrail_verdict, + rewritten_item, +) +from litellm.types.guardrails import BlockedWord, ContentFilterAction, ContentFilterPattern, GuardrailEventHooks +from litellm.types.utils import GenericGuardrailAPIInputs + + +class RaisingGuardrail(CustomGuardrail): + def __init__(self, error: Exception, realtime_violation_message: str | None = None) -> None: + super().__init__( + guardrail_name="raising", + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + realtime_violation_message=realtime_violation_message, + ) + self.error = error + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: LiteLLMLogging | None = None, + ) -> GenericGuardrailAPIInputs: + raise self.error + + +class MaskingGuardrail(CustomGuardrail): + def __init__(self, word: str) -> None: + super().__init__(guardrail_name=f"mask-{word}", event_hook=GuardrailEventHooks.pre_call, default_on=True) + self.word = word + self.screened: list[tuple[str, ...]] = [] + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: LiteLLMLogging | None = None, + ) -> GenericGuardrailAPIInputs: + texts: Final = inputs.get("texts", []) + self.screened.append(tuple(texts)) + return {**inputs, "texts": [text.replace(self.word, "[MASKED]") for text in texts]} + + +class LastTextRedactingGuardrail(MaskingGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: LiteLLMLogging | None = None, + ) -> GenericGuardrailAPIInputs: + last: Final = inputs.get("texts", [""])[-1] + if self.word not in last: + return inputs + return {**inputs, "texts": [last.replace(self.word, "[MASKED]")]} + + +def _item_create(item: object) -> dict: + return {"type": "conversation.item.create", "item": item} + + +@pytest.mark.parametrize("role", ["user", "assistant", "system"]) +def test_only_the_typed_text_of_a_message_is_screened_whatever_its_role(role: str) -> None: + client_event: Final = _item_create( + { + "type": "message", + "role": role, + "content": [{"type": "input_audio", "audio": "AAAA"}, {"type": "input_text", "text": "hello"}], + } + ) + + assert realtime_guarded_item(client_event) == RealtimeMessageText(role=role, texts=("hello",)) + + +@pytest.mark.parametrize("part_type", ["output_text", "text"]) +def test_assistant_text_written_by_the_client_is_screened(part_type: str) -> None: + client_event: Final = _item_create( + {"type": "message", "role": "assistant", "content": [{"type": part_type, "text": "hello"}]} + ) + + assert realtime_guarded_item(client_event) == RealtimeMessageText(role="assistant", texts=("hello",)) + + +@pytest.mark.parametrize( + "client_event", + [ + pytest.param( + _item_create({"type": "message", "role": "user", "content": [{"type": "input_audio", "audio": "AAAA"}]}), + id="audio-only user message", + ), + pytest.param( + _item_create({"type": "function_call_output", "call_id": "call_1", "output": ""}), id="empty tool output" + ), + pytest.param(_item_create("hello"), id="item that is not an object"), + ], +) +def test_items_without_client_text_are_not_screened(client_event: dict) -> None: + assert realtime_guarded_item(client_event) is None + + +@pytest.mark.asyncio +async def test_a_plain_string_block_reason_becomes_the_violation_message() -> None: + guardrail: Final = RaisingGuardrail(HTTPException(status_code=400, detail="hello is not allowed")) + + block: Final = await realtime_guardrail_verdict( + ("hello",), "user", (GuardrailEventHooks.pre_call,), {}, None, (guardrail,) + ) + + assert isinstance(block, RealtimeGuardrailBlock) + assert block.message == "hello is not allowed" + + +@pytest.mark.asyncio +async def test_a_configured_violation_message_replaces_the_block_reason() -> None: + guardrail: Final = RaisingGuardrail(HTTPException(status_code=400, detail={"error": "keyword 'hello' detected"})) + configured: Final = RaisingGuardrail( + HTTPException(status_code=400, detail="hello is not allowed"), realtime_violation_message="Please rephrase" + ) + + block: Final = await realtime_guardrail_verdict( + ("hello",), "user", (GuardrailEventHooks.pre_call,), {}, None, (guardrail,) + ) + configured_block: Final = await realtime_guardrail_verdict( + ("hello",), "user", (GuardrailEventHooks.pre_call,), {}, None, (configured,) + ) + + assert isinstance(block, RealtimeGuardrailBlock) + assert block.message == "keyword 'hello' detected" + assert isinstance(configured_block, RealtimeGuardrailBlock) + assert configured_block.message == "Please rephrase" + + +@pytest.mark.asyncio +async def test_masked_texts_come_back_as_a_rewrite_and_reach_later_guardrails_masked() -> None: + first: Final = MaskingGuardrail("hunter2") + second: Final = MaskingGuardrail("swordfish") + + verdict: Final = await realtime_guardrail_verdict( + ("my hunter2", "and swordfish", "hi"), "user", (GuardrailEventHooks.pre_call,), {}, None, (first, second) + ) + + assert verdict == RealtimeGuardrailRewrite(texts=("my [MASKED]", "and [MASKED]", "hi")) + assert sorted(second.screened[:2]) == [("my [MASKED] and swordfish hi",), ("my [MASKED]and swordfishhi",)] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("parts", "screened"), + [ + pytest.param( + ("Compare this photo ", "with that one"), + ("Compare this photo with that one",), + id="a space at the boundary", + ), + pytest.param( + ("Compare this photo", "with that one"), + ("Compare this photowith that one", "Compare this photo with that one"), + id="no space at the boundary", + ), + ], +) +async def test_a_clean_multi_part_item_sends_each_guardrail_only_the_joins_it_needs_one_text_per_call( + parts: tuple[str, ...], screened: tuple[str, ...] +) -> None: + guardrail: Final = MaskingGuardrail("hunter2") + + verdict: Final = await realtime_guardrail_verdict( + parts, "user", (GuardrailEventHooks.pre_call,), {}, None, (guardrail,) + ) + + assert verdict is None + assert sorted(guardrail.screened) == sorted((text,) for text in screened) + + +@pytest.mark.asyncio +async def test_texts_a_guardrail_leaves_unchanged_need_no_rewrite() -> None: + verdict: Final = await realtime_guardrail_verdict( + ("hello",), "user", (GuardrailEventHooks.pre_call,), {}, None, (MaskingGuardrail("hunter2"),) + ) + + assert verdict is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("parts", "phrase"), + [ + pytest.param(("tell me the secret", "plan"), "secret plan", id="split at a word break"), + pytest.param(("tell me about secret-", "project"), "secret-project", id="split mid-word"), + ], +) +async def test_a_phrase_split_across_typed_text_parts_is_still_blocked(parts: tuple[str, str], phrase: str) -> None: + client_event: Final = _item_create( + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": part} for part in parts]} + ) + guarded_item: Final = realtime_guarded_item(client_event) + assert isinstance(guarded_item, RealtimeMessageText) + + phrase_filter: Final = ContentFilterGuardrail( + guardrail_name="phrase-filter", + blocked_words=[BlockedWord(keyword=phrase, action=ContentFilterAction.BLOCK)], + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + ) + + verdict: Final = await realtime_guardrail_verdict( + guarded_item.texts, guarded_item.role, (GuardrailEventHooks.pre_call,), {}, None, (phrase_filter,) + ) + + assert isinstance(verdict, RealtimeGuardrailBlock) + + +@pytest.mark.asyncio +async def test_a_mask_inside_one_part_keeps_every_part_in_place() -> None: + first_photo: Final = {"type": "input_image", "image_url": "data:image/png;base64,AAA"} + second_photo: Final = {"type": "input_image", "image_url": "data:image/png;base64,BBB"} + client_event: Final = _item_create( + { + "type": "message", + "role": "user", + "content": [ + {"type": "input_text", "text": "Compare this photo"}, + first_photo, + {"type": "input_text", "text": "with hunter2's photo"}, + second_photo, + ], + } + ) + guarded_item: Final = realtime_guarded_item(client_event) + assert isinstance(guarded_item, RealtimeMessageText) + + verdict: Final = await realtime_guardrail_verdict( + guarded_item.texts, "user", (GuardrailEventHooks.pre_call,), {}, None, (MaskingGuardrail("hunter2"),) + ) + + assert isinstance(verdict, RealtimeGuardrailRewrite) + assert rewritten_item(client_event, verdict.texts)["content"] == [ + {"type": "input_text", "text": "Compare this photo"}, + first_photo, + {"type": "input_text", "text": "with [MASKED]'s photo"}, + second_photo, + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("parts", "word", "masked"), + [ + pytest.param(("my secret", "plan is here"), "secret plan", "my [MASKED] is here", id="split at a word break"), + pytest.param(("my code is hun", "ter2 ok"), "hunter2", "my code is [MASKED] ok", id="split mid-word"), + ], +) +async def test_a_mask_across_parts_replaces_them_with_the_masked_text_the_model_would_have_read( + parts: tuple[str, ...], word: str, masked: str +) -> None: + verdict: Final = await realtime_guardrail_verdict( + parts, "user", (GuardrailEventHooks.pre_call,), {}, None, (MaskingGuardrail(word),) + ) + + assert verdict == RealtimeGuardrailRewrite(texts=(masked,)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "parts", + [ + pytest.param(("abc", "123-45-6789 hunter2 rocks"), id="masks inside one part"), + pytest.param(("abc", "123-45-6789 hunter2", "rocks"), id="a mask inside a part and one across parts"), + ], +) +async def test_no_mask_one_text_part_gets_is_lost_when_its_joined_text_masks_differently( + parts: tuple[str, ...], +) -> None: + content_filter: Final = ContentFilterGuardrail( + guardrail_name="ssn-and-keyword-mask", + patterns=[ + ContentFilterPattern(pattern_type="prebuilt", pattern_name="us_ssn", action=ContentFilterAction.MASK) + ], + blocked_words=[BlockedWord(keyword="hunter2 rocks", action=ContentFilterAction.MASK)], + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + ) + + verdict: Final = await realtime_guardrail_verdict( + parts, "user", (GuardrailEventHooks.pre_call,), {}, None, (content_filter,) + ) + + assert isinstance(verdict, RealtimeGuardrailRewrite) + forwarded: Final = " ".join(verdict.texts) + assert "123-45-6789" not in forwarded + assert "hunter2 rocks" not in forwarded + assert forwarded.startswith("abc") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("parts", "rewritten"), + [ + pytest.param(("my code is", "hunter2"), ("my code is", "[MASKED]"), id="mask inside one part"), + pytest.param(("my code is hun", "ter2"), ("my code is [MASKED]",), id="mask across parts"), + ], +) +async def test_a_guardrail_that_answers_with_one_text_still_masks_parts_joined_without_a_space( + parts: tuple[str, ...], rewritten: tuple[str, ...] +) -> None: + verdict: Final = await realtime_guardrail_verdict( + parts, "user", (GuardrailEventHooks.pre_call,), {}, None, (LastTextRedactingGuardrail("hunter2"),) + ) + + assert verdict == RealtimeGuardrailRewrite(texts=rewritten) + + +@pytest.mark.asyncio +async def test_a_mask_in_a_message_with_more_text_parts_than_the_recursion_limit_still_comes_back() -> None: + parts: Final = (*(f"part {index}." for index in range(sys.getrecursionlimit())), "my code is hunter2.") + + verdict: Final = await realtime_guardrail_verdict( + parts, "user", (GuardrailEventHooks.pre_call,), {}, None, (MaskingGuardrail("hunter2"),) + ) + + assert verdict == RealtimeGuardrailRewrite(texts=(*parts[:-1], "my code is [MASKED].")) + + +class GuardrailLogRecordingGuardrail(CustomGuardrail): + def __init__(self) -> None: + super().__init__(guardrail_name="log-recorder", event_hook=GuardrailEventHooks.pre_call, default_on=True) + self.guardrail_log_lengths: list[int] = [] + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: LiteLLMLogging | None = None, + ) -> GenericGuardrailAPIInputs: + self.guardrail_log_lengths.append( + len(request_data["litellm_metadata"].get("standard_logging_guardrail_information", [])) + ) + return inputs + + +@pytest.mark.asyncio +async def test_a_guardrail_sees_only_the_current_turns_guardrail_log_however_long_the_session_runs() -> None: + session_request_data: Final = {"litellm_metadata": {"user_api_key_team_id": "team-a"}} + content_filter: Final = ContentFilterGuardrail( + guardrail_name="phrase-filter", + blocked_words=[BlockedWord(keyword="hunter2", action=ContentFilterAction.BLOCK)], + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + ) + recorder: Final = GuardrailLogRecordingGuardrail() + + for turn in ("hello", "how are you", "goodbye"): + await realtime_guardrail_verdict( + (turn,), "user", (GuardrailEventHooks.pre_call,), session_request_data, None, (content_filter, recorder) + ) + + assert recorder.guardrail_log_lengths == [1, 1, 1] + assert session_request_data == {"litellm_metadata": {"user_api_key_team_id": "team-a"}} + + +class ScopedGuardrail(RaisingGuardrail): + def __init__( + self, *, skip_system_message: bool | None, skip_tool_message: bool | None, scan_only_tool_results: bool | None + ) -> None: + super().__init__(HTTPException(status_code=400, detail="blocked")) + self.skip_system_message_in_guardrail = skip_system_message + self.skip_tool_message_in_guardrail = skip_tool_message + self.scan_only_tool_results = scan_only_tool_results + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("scope", "screened_roles"), + [ + pytest.param({}, ("user", "assistant", "system", "tool"), id="no scope"), + pytest.param({"skip_system_message": True}, ("user", "assistant", "tool"), id="skip system"), + pytest.param({"skip_tool_message": True}, ("user", "assistant", "system"), id="skip tool"), + pytest.param({"scan_only_tool_results": True}, ("tool",), id="only tool results"), + ], +) +async def test_a_guardrail_screens_only_the_roles_its_scope_allows( + scope: dict[str, bool], screened_roles: tuple[str, ...] +) -> None: + guardrail: Final = ScopedGuardrail( + skip_system_message=scope.get("skip_system_message"), + skip_tool_message=scope.get("skip_tool_message"), + scan_only_tool_results=scope.get("scan_only_tool_results"), + ) + + verdicts: Final = { + role: await realtime_guardrail_verdict( + ("hello",), role, (GuardrailEventHooks.pre_call,), {}, None, (guardrail,) + ) + for role in ("user", "assistant", "system", "tool") + } + + assert tuple(role for role, verdict in verdicts.items() if verdict is not None) == screened_roles + + +def test_a_rewrite_puts_the_masked_text_in_place_of_the_typed_text_parts() -> None: + audio: Final = {"type": "input_audio", "audio": "AAAA"} + client_event: Final = _item_create( + { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "a"}, audio, {"type": "input_text", "text": "b"}], + } + ) + + assert rewritten_item(client_event, ("[MASKED] b",)) == { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "[MASKED] b"}, audio], + } + + +def test_a_rewrite_replaces_the_tool_output() -> None: + client_event: Final = _item_create({"type": "function_call_output", "call_id": "call_1", "output": "hunter2"}) + + assert rewritten_item(client_event, ("[MASKED]",)) == { + "type": "function_call_output", + "call_id": "call_1", + "output": "[MASKED]", + } + + +@pytest.mark.asyncio +async def test_a_guardrail_that_crashes_is_not_treated_as_a_block() -> None: + guardrail: Final = RaisingGuardrail(RuntimeError("guardrail backend down")) + + with pytest.raises(RuntimeError, match="guardrail backend down"): + await realtime_guardrail_verdict(("hello",), "user", (GuardrailEventHooks.pre_call,), {}, None, (guardrail,)) + + +_TEAM_POLICY_CODE: Final = ( + "def apply_guardrail(inputs, request_data, input_type):\n" + ' if request_data["team_id"] == "team-strict":\n' + ' return block("blocked for this team")\n' + " return allow()\n" +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("team_id", "blocked"), + [pytest.param("team-strict", True, id="policy team"), pytest.param("team-other", False, id="other team")], +) +async def test_a_team_scoped_guardrail_sees_the_session_team(team_id: str, blocked: bool) -> None: + team_policy = CustomCodeGuardrail( + custom_code=_TEAM_POLICY_CODE, guardrail_name="team-policy", event_hook="pre_call", default_on=True + ) + + verdict = await realtime_guardrail_verdict( + ("hello",), + "user", + (GuardrailEventHooks.pre_call,), + {"litellm_metadata": {"user_api_key_team_id": team_id}}, + None, + (team_policy,), + ) + + assert isinstance(verdict, RealtimeGuardrailBlock) is blocked diff --git a/tests/unit/litellm_core_utils/test_realtime_streaming.py b/tests/unit/litellm_core_utils/test_realtime_streaming.py index eec1dfcf459..7f205b0bd00 100644 --- a/tests/unit/litellm_core_utils/test_realtime_streaming.py +++ b/tests/unit/litellm_core_utils/test_realtime_streaming.py @@ -1056,6 +1056,117 @@ async def test_transcription_captured_in_backend_to_client(): assert logging_obj.model_call_details["messages"] == streaming.input_messages +@pytest.mark.asyncio +@pytest.mark.parametrize("transcript", [None, 123, ["a"], {"text": "a"}]) +async def test_a_non_string_transcript_does_not_end_the_session(transcript: object) -> None: + client_ws: Final = MagicMock() + client_ws.send_text = AsyncMock() + transcript_event: Final = json.dumps( + {"type": "conversation.item.input_audio_transcription.completed", "transcript": transcript, "item_id": "i"} + ) + later_event: Final = json.dumps({"type": "response.done", "response": {"id": "resp_1"}}) + backend_ws: Final = MagicMock() + backend_ws.recv = AsyncMock( + side_effect=[transcript_event.encode(), later_event.encode(), ConnectionClosed(None, None)] + ) + backend_ws.send = AsyncMock() + logging_obj: Final = _mock_logging() + logging_obj.model_call_details = {} + logging_obj.dispatch_success_handlers = AsyncMock() + + await RealTimeStreaming(client_ws, backend_ws, logging_obj).backend_to_client_send_messages() + + sent_types: Final = [json.loads(call.args[0])["type"] for call in client_ws.send_text.call_args_list] + assert sent_types == ["conversation.item.input_audio_transcription.completed", "response.done"] + backend_ws.send.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("transcript", [None, 123, ["a"]]) +async def test_a_non_string_transcript_fails_the_session_when_a_transcript_guardrail_would_screen_it( + transcript: object, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "callbacks", [_transcription_guardrail()]) + client_ws: Final = MagicMock() + client_ws.send_text = AsyncMock() + transcript_event: Final = json.dumps( + {"type": "conversation.item.input_audio_transcription.completed", "transcript": transcript, "item_id": "i"} + ) + later_event: Final = json.dumps({"type": "response.done", "response": {"id": "resp_1"}}) + backend_ws: Final = MagicMock() + backend_ws.recv = AsyncMock( + side_effect=[transcript_event.encode(), later_event.encode(), ConnectionClosed(None, None)] + ) + backend_ws.send = AsyncMock() + logging_obj: Final = _mock_logging() + logging_obj.model_call_details = {} + logging_obj.dispatch_success_handlers = AsyncMock() + + close: Final = await RealTimeStreaming(client_ws, backend_ws, logging_obj).backend_to_client_send_messages() + + assert close.code == 1011 + sent_types: Final = [json.loads(call.args[0])["type"] for call in client_ws.send_text.call_args_list] + assert sent_types == ["conversation.item.input_audio_transcription.completed"] + backend_ws.send.assert_not_called() + + +def _native_lifecycle_transcription_guardrail() -> CustomGuardrail: + class _NativeLifecycleGuardrail(CustomGuardrail): + use_native_lifecycle_hooks = True + + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + return inputs + + return _NativeLifecycleGuardrail( + guardrail_name="test_native_guard", + event_hook=GuardrailEventHooks.realtime_input_transcription, + default_on=True, + ) + + +def _tool_results_only_transcription_guardrail() -> CustomGuardrail: + class _ToolResultsOnlyGuardrail(CustomGuardrail): + scan_only_tool_results = True + + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + return inputs + + return _ToolResultsOnlyGuardrail( + guardrail_name="test_tool_results_guard", + event_hook=GuardrailEventHooks.realtime_input_transcription, + default_on=True, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "guardrail_factory", [_native_lifecycle_transcription_guardrail, _tool_results_only_transcription_guardrail] +) +async def test_a_non_string_transcript_keeps_the_session_when_no_guardrail_would_screen_the_users_turn( + guardrail_factory: Callable[[], CustomGuardrail], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "callbacks", [guardrail_factory()]) + client_ws: Final = MagicMock() + client_ws.send_text = AsyncMock() + transcript_event: Final = json.dumps( + {"type": "conversation.item.input_audio_transcription.completed", "transcript": None, "item_id": "i"} + ) + later_event: Final = json.dumps({"type": "response.done", "response": {"id": "resp_1"}}) + backend_ws: Final = MagicMock() + backend_ws.recv = AsyncMock( + side_effect=[transcript_event.encode(), later_event.encode(), ConnectionClosed(None, None)] + ) + backend_ws.send = AsyncMock() + logging_obj: Final = _mock_logging() + logging_obj.model_call_details = {} + logging_obj.dispatch_success_handlers = AsyncMock() + + await RealTimeStreaming(client_ws, backend_ws, logging_obj).backend_to_client_send_messages() + + sent_types: Final = [json.loads(call.args[0])["type"] for call in client_ws.send_text.call_args_list] + assert sent_types == ["conversation.item.input_audio_transcription.completed", "response.done"] + + @pytest.mark.asyncio async def test_transcription_session_captures_usage_and_skips_response_create(): """ @@ -1955,6 +2066,112 @@ async def test_realtime_text_input_guardrail_blocks_and_returns_error(monkeypatc assert len(original_items) == 0, f"Blocked item should not be forwarded to backend, got: {original_items}" +@pytest.mark.asyncio +async def test_the_next_clean_turn_is_answered_after_a_blocked_item_with_no_response_create( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from fastapi import HTTPException + + class SecretBlockingGuardrail(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: LiteLLMLoggingObj | None = None, + ) -> GenericGuardrailAPIInputs: + if any("hunter2" in text for text in inputs.get("texts", [])): + raise HTTPException(status_code=400, detail={"error": "hunter2 is not allowed"}) + return inputs + + def item_create(role: str, text: str) -> str: + return json.dumps( + { + "type": "conversation.item.create", + "item": {"type": "message", "role": role, "content": [{"type": "input_text", "text": text}]}, + } + ) + + monkeypatch.setattr( + litellm, + "callbacks", + [SecretBlockingGuardrail(guardrail_name="no-secrets", event_hook=GuardrailEventHooks.pre_call, default_on=True)], + ) + client_ws: Final = MagicMock() + client_ws.send_text = AsyncMock() + client_ws.receive_text = AsyncMock( + side_effect=[ + item_create("system", "never mention hunter2"), + item_create("user", "what is the weather"), + json.dumps({"type": "response.create"}), + Exception("connection closed"), + ] + ) + backend_ws: Final = MagicMock() + backend_ws.send = AsyncMock() + streaming: Final = RealTimeStreaming(client_ws, backend_ws, _mock_logging()) + + await streaming.client_ack_messages() + + sent_to_backend: Final = [json.loads(call.args[0]) for call in backend_ws.send.call_args_list] + clean_index: Final = next( + index + for index, event in enumerate(sent_to_backend) + if event.get("item", {}).get("content", [{}])[0].get("text") == "what is the weather" + ) + assert [event["type"] for event in sent_to_backend[clean_index + 1 :]] == ["response.create"], sent_to_backend + + +@pytest.mark.asyncio +async def test_realtime_text_input_masking_guardrail_forwards_only_the_masked_text( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class MaskingGuardrail(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: LiteLLMLoggingObj | None = None, + ) -> GenericGuardrailAPIInputs: + texts: Final = inputs.get("texts", []) + return {**inputs, "texts": [text.replace("test@example.com", "[EMAIL]") for text in texts]} + + monkeypatch.setattr( + litellm, + "callbacks", + [MaskingGuardrail(guardrail_name="email-masker", event_hook=GuardrailEventHooks.pre_call, default_on=True)], + ) + client_ws: Final = MagicMock() + client_ws.send_text = AsyncMock() + client_ws.receive_text = AsyncMock( + side_effect=[ + json.dumps( + { + "type": "conversation.item.create", + "item": { + "role": "user", + "content": [{"type": "input_text", "text": "My email is test@example.com"}], + }, + } + ), + Exception("connection closed"), + ] + ) + backend_ws: Final = MagicMock() + backend_ws.send = AsyncMock() + backend_ws.recv = AsyncMock(side_effect=ConnectionClosed(None, None)) + logging_obj: Final = _mock_logging() + logging_obj.pre_call = MagicMock() + + await RealTimeStreaming(client_ws, backend_ws, logging_obj).client_ack_messages() + + sent_to_backend: Final = json.dumps([c.args[0] for c in backend_ws.send.call_args_list if c.args]) + assert "test@example.com" not in sent_to_backend + assert "My email is [EMAIL]" in sent_to_backend + assert client_ws.send_text.call_args_list == [] + + @pytest.mark.asyncio async def test_realtime_function_call_output_guardrail_blocks_and_returns_error(monkeypatch: pytest.MonkeyPatch): """ diff --git a/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py b/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py index cf2f25faaa9..1221667cb51 100644 --- a/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py +++ b/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py @@ -10,23 +10,38 @@ from typing import Final, Literal from unittest.mock import MagicMock import pytest +from fastapi import HTTPException from pydantic import TypeAdapter +from typing_extensions import ReadOnly, TypedDict, Unpack import litellm from litellm.constants import ( BEDROCK_REALTIME_SDK_SUPPORTED_RANGE, + REALTIME_SESSION_ENDED_ABNORMALLY_KEY, REALTIME_SESSION_SUCCESS_LOGGED_KEY, WEBSOCKET_CLOSE_REASON_MAX_BYTES, ) +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.litellm_core_utils.realtime_errors import websocket_close_reason +from litellm.litellm_core_utils.realtime_guardrails import BLOCKED_TOOL_OUTPUT from litellm.llms.bedrock.common_utils import BedrockError -from litellm.llms.bedrock.realtime.handler import BedrockRealtime, SessionEndSignals -from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload +from litellm.llms.bedrock.realtime.handler import ( + GUARDRAIL_UNAVAILABLE_MESSAGE, + BedrockRealtime, + BedrockRealtimeGuardrails, + SessionEndSignals, +) from litellm.llms.bedrock.realtime.transformation import BedrockRealtimeConfig +from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload from litellm.realtime_api import main as realtime_main -from litellm.integrations.custom_logger import CustomLogger -from litellm.types.utils import StandardLoggingPayload, StandardLoggingPayloadErrorInformation +from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.utils import ( + GenericGuardrailAPIInputs, + StandardLoggingPayload, + StandardLoggingPayloadErrorInformation, +) from tests.unit.llms.bedrock.event_loop_probe import host_in_loop_probing_vault @@ -223,10 +238,19 @@ class DisconnectingClientWS: return self._messages.pop(0) raise RuntimeError("client disconnected") - async def send_text(self, message): + async def send_text(self, message: str) -> None: self.sent_to_client.append(message) +class GuardrailErrorUndeliveredClientWS(DisconnectingClientWS): + """Client that disconnects just as the proxy tries to tell it about a guardrail violation""" + + async def send_text(self, message: str) -> None: + if "guardrail_violation" in message: + raise RuntimeError("client disconnected") + await super().send_text(message) + + class ClosableClientWS: def __init__(self): self.closed = False @@ -1052,7 +1076,16 @@ class TestBedrockRealtimeProviderFailurePropagation: assert other_provider.calls == 0 @pytest.mark.asyncio - async def test_readiness_failure_on_an_input_first_start_replays_the_input(self, stub_aws_sdk_client): + @pytest.mark.parametrize( + "guardrail_awaits_its_backend", [False, True], ids=["no guardrail", "a guardrail that awaits its backend"] + ) + async def test_readiness_failure_on_an_input_first_start_replays_the_input( + self, stub_aws_sdk_client, monkeypatch, guardrail_awaits_its_backend: bool + ): + guardrail: Final = AwaitingGuardrail( + "hunter2", guardrail_name="remote", event_hook=GuardrailEventHooks.pre_call, default_on=True + ) + monkeypatch.setattr(litellm, "callbacks", [guardrail] if guardrail_awaits_its_backend else []) handler = BedrockRealtime() websocket = RealtimeClientWS([self.USER_TEXT]) healthy_stream = SessionBoundBedrockStream() @@ -1687,5 +1720,614 @@ class TestBedrockRealtimeSdkImportErrors: assert BEDROCK_REALTIME_SDK_SUPPORTED_RANGE in close_reason +class SessionEndOptions(TypedDict, total=False): + on_violation: ReadOnly[str] + end_session_after_n_fails: ReadOnly[int] + + +class BlockedWordGuardrailOptions(SessionEndOptions, total=False): + guardrail_name: ReadOnly[str] + event_hook: ReadOnly[GuardrailEventHooks] + default_on: ReadOnly[bool] + realtime_violation_message: ReadOnly[str] + + +class BlockedWordGuardrail(CustomGuardrail): + def __init__(self, blocked_word: str, **kwargs: Unpack[BlockedWordGuardrailOptions]) -> None: + super().__init__(**kwargs) + self.blocked_word = blocked_word + self.checked: list[tuple[tuple[str, ...], object]] = [] + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: LiteLLMLogging | None = None, + ) -> GenericGuardrailAPIInputs: + self.checked.append((tuple(inputs["texts"]), request_data.get("user_api_key_dict"))) + if any(self.blocked_word in text for text in inputs["texts"]): + raise HTTPException(status_code=400, detail={"error": f"{self.blocked_word} is not allowed"}) + return inputs + + +class MaskingGuardrail(BlockedWordGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: LiteLLMLogging | None = None, + ) -> GenericGuardrailAPIInputs: + return {**inputs, "texts": [text.replace(self.blocked_word, "[MASKED]") for text in inputs["texts"]]} + + +class AwaitingGuardrail(BlockedWordGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: LiteLLMLogging | None = None, + ) -> GenericGuardrailAPIInputs: + await asyncio.sleep(0.05) + return await super().apply_guardrail(inputs, request_data, input_type, logging_obj) + + +class SystemSkippingGuardrail(BlockedWordGuardrail): + skip_system_message_in_guardrail: Final = True + + +class CrashingGuardrail(BlockedWordGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: LiteLLMLogging | None = None, + ) -> GenericGuardrailAPIInputs: + raise RuntimeError("guardrail backend down") + + +class TestBedrockRealtimeGuardrails: + """Realtime guardrails must screen Nova Sonic sessions the way they screen OpenAI realtime sessions""" + + MODEL: Final = "amazon.nova-sonic-v1:0" + VIOLATION: Final = "I can't help with that." + BLOCKED_TEXT: Final = "my secret is hunter2" + CLEAN_TEXT: Final = "what is the weather like" + RESPONSE_CREATE: Final = json.dumps({"type": "response.create"}) + + @staticmethod + def _guardrail( + event_hook: GuardrailEventHooks = GuardrailEventHooks.pre_call, + default_on: bool = True, + **kwargs: Unpack[SessionEndOptions], + ) -> BlockedWordGuardrail: + return BlockedWordGuardrail( + "hunter2", + guardrail_name="no_secrets", + event_hook=event_hook, + default_on=default_on, + realtime_violation_message=TestBedrockRealtimeGuardrails.VIOLATION, + **kwargs, + ) + + @staticmethod + def _user_text(text: str, role: str = "user") -> str: + return json.dumps( + { + "type": "conversation.item.create", + "item": {"type": "message", "role": role, "content": [{"type": "input_text", "text": text}]}, + } + ) + + @staticmethod + def _bedrock_events(stream: FakeBedrockStream) -> list[dict]: + return [json.loads(chunk.value.bytes_.decode("utf-8"))["event"] for chunk in stream.input_stream.sent] + + @classmethod + def _text_inputs(cls, stream: FakeBedrockStream) -> list[str]: + events: Final = cls._bedrock_events(stream) + user_contents: Final = { + event["contentStart"]["contentName"] + for event in events + if event.get("contentStart", {}).get("role") == "USER" + } + return [ + event["textInput"]["content"] + for event in events + if event.get("textInput", {}).get("contentName") in user_contents + ] + + @staticmethod + def _client_events(client_ws: DisconnectingClientWS | RealtimeClientWS) -> list[dict]: + return [json.loads(message) for message in client_ws.sent_to_client] + + @classmethod + def _guardrail_errors(cls, client_ws: DisconnectingClientWS | RealtimeClientWS) -> list[dict]: + return [ + event["error"] + for event in cls._client_events(client_ws) + if event["type"] == "error" and event["error"]["type"] == "guardrail_violation" + ] + + async def _forward_client( + self, + messages: list[str], + guardrail: BlockedWordGuardrail, + client_ws_type: type[DisconnectingClientWS] | type[ConnectedClientWS] = DisconnectingClientWS, + ) -> tuple[FakeBedrockStream, DisconnectingClientWS | ConnectedClientWS]: + stream: Final = FakeBedrockStream() + client_ws: Final = client_ws_type(messages) + await asyncio.wait_for( + BedrockRealtime()._forward_client_to_bedrock( + client_ws, + stream, + BedrockRealtimeConfig(), + self.MODEL, + {}, + signals=SessionEndSignals(), + guardrails=BedrockRealtimeGuardrails(callbacks=lambda: (guardrail,)), + ), + timeout=1, + ) + return stream, client_ws + + @staticmethod + def _user_transcript(text: str) -> tuple[str, ...]: + return ( + json.dumps( + { + "event": { + "contentStart": { + "role": "USER", + "type": "TEXT", + "additionalModelFields": json.dumps({"generationStage": "FINAL"}), + } + } + } + ), + json.dumps({"event": {"textOutput": {"content": text}}}), + json.dumps({"event": {"contentEnd": {"stopReason": "PARTIAL_TURN"}}}), + ) + + @staticmethod + def _assistant_text(text: str) -> tuple[str, ...]: + return ( + json.dumps({"event": {"contentStart": {"role": "ASSISTANT", "type": "TEXT"}}}), + json.dumps({"event": {"textOutput": {"content": text}}}), + json.dumps({"event": {"contentEnd": {"stopReason": "END_TURN"}}}), + ) + + async def _forward_bedrock( + self, + payloads: tuple[str, ...], + guardrail: BlockedWordGuardrail, + client_audio_streamed: bool = True, + client_ws_type: type[ConnectedClientWS] | type[GuardrailErrorUndeliveredClientWS] = ConnectedClientWS, + ) -> tuple[ConnectedClientWS | GuardrailErrorUndeliveredClientWS, list[object]]: + config: Final = BedrockRealtimeConfig() + config.client_audio_streamed = client_audio_streamed + client_ws: Final = client_ws_type([]) + logged: Final = [ + event + async for event in BedrockRealtime()._forward_bedrock_to_client( + ScriptedBedrockStream(payloads), + client_ws, + config, + self.MODEL, + FakeLogging(), + {}, + signals=_closing_input(), + guardrails=BedrockRealtimeGuardrails(callbacks=lambda: (guardrail,)), + ) + ] + return client_ws, logged + + @pytest.mark.asyncio + async def test_blocked_typed_message_never_reaches_bedrock_and_the_violation_is_voiced( + self, stub_aws_models: None + ) -> None: + guardrail: Final = self._guardrail() + + stream, client_ws = await self._forward_client( + [self._user_text(self.BLOCKED_TEXT), self.RESPONSE_CREATE], guardrail + ) + clean_stream, _ = await self._forward_client( + [self._user_text(self.CLEAN_TEXT), self.RESPONSE_CREATE], guardrail + ) + + assert "hunter2" not in json.dumps(self._bedrock_events(stream)) + assert [f'"{self.VIOLATION}"' in text for text in self._text_inputs(stream)] == [True] + assert self._guardrail_errors(client_ws) == [ + {"type": "guardrail_violation", "message": self.VIOLATION, "code": "content_policy_violation"} + ] + assert guardrail.checked == [((self.BLOCKED_TEXT,), None), ((self.CLEAN_TEXT,), None)] + audio_inputs = [sum("audioInput" in event for event in self._bedrock_events(s)) for s in (stream, clean_stream)] + assert audio_inputs[0] == audio_inputs[1] > 0, "the client's own response.create must not start a 2nd turn" + + @pytest.mark.asyncio + @pytest.mark.parametrize("role", ["assistant", "system"]) + async def test_blocked_typed_message_is_screened_whatever_its_role(self, stub_aws_models: None, role: str) -> None: + stream, client_ws = await self._forward_client( + [self._user_text(self.BLOCKED_TEXT, role=role)], self._guardrail() + ) + + assert "hunter2" not in json.dumps(self._bedrock_events(stream)) + assert [error["message"] for error in self._guardrail_errors(client_ws)] == [self.VIOLATION] + + @pytest.mark.asyncio + async def test_a_system_message_nova_receives_as_a_user_turn_is_screened_by_a_system_skipping_guardrail( + self, stub_aws_models: None + ) -> None: + guardrail: Final = SystemSkippingGuardrail( + "hunter2", + guardrail_name="no_secrets", + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + realtime_violation_message=self.VIOLATION, + ) + + stream, client_ws = await self._forward_client([self._user_text(self.BLOCKED_TEXT, role="system")], guardrail) + + assert "hunter2" not in json.dumps(self._bedrock_events(stream)) + assert [error["message"] for error in self._guardrail_errors(client_ws)] == [self.VIOLATION] + + @pytest.mark.asyncio + async def test_the_next_clean_turn_is_answered_after_a_blocked_item_with_no_response_create( + self, stub_aws_models: None + ) -> None: + stream, _ = await self._forward_client( + [ + self._user_text(self.BLOCKED_TEXT, role="system"), + self._user_text(self.CLEAN_TEXT), + self.RESPONSE_CREATE, + ], + self._guardrail(), + ) + + events: Final = self._bedrock_events(stream) + clean_text_index: Final = next( + i for i, event in enumerate(events) if event.get("textInput", {}).get("content") == self.CLEAN_TEXT + ) + assert any("audioInput" in event for event in events[clean_text_index:]), "the client's response.create runs" + + @pytest.mark.asyncio + async def test_client_message_the_guardrail_cannot_parse_never_reaches_bedrock( + self, stub_aws_models: None + ) -> None: + lone_surrogate_text: Final = self._user_text(f"{self.BLOCKED_TEXT} \ud800") + + stream, _ = await self._forward_client( + [lone_surrogate_text, self._user_text(self.CLEAN_TEXT)], self._guardrail() + ) + + assert "hunter2" not in json.dumps(self._bedrock_events(stream)) + assert self._text_inputs(stream) == [self.CLEAN_TEXT] + + @pytest.mark.asyncio + async def test_clean_typed_message_reaches_bedrock_unchanged(self, stub_aws_models: None) -> None: + guardrail: Final = self._guardrail() + + stream, client_ws = await self._forward_client([self._user_text(self.CLEAN_TEXT)], guardrail) + + assert self._text_inputs(stream) == [self.CLEAN_TEXT] + assert client_ws.sent_to_client == [] + assert guardrail.checked == [((self.CLEAN_TEXT,), None)] + + @pytest.mark.asyncio + async def test_masked_typed_message_reaches_bedrock_masked_and_is_answered(self, stub_aws_models: None) -> None: + guardrail: Final = MaskingGuardrail( + "hunter2", guardrail_name="mask_secrets", event_hook=GuardrailEventHooks.pre_call, default_on=True + ) + + stream, client_ws = await self._forward_client( + [self._user_text(self.BLOCKED_TEXT), self.RESPONSE_CREATE], guardrail + ) + + events: Final = self._bedrock_events(stream) + assert "hunter2" not in json.dumps(events) + assert self._text_inputs(stream) == ["my secret is [MASKED]"] + assert client_ws.sent_to_client == [] + assert any("audioInput" in event for event in events), "the client's response.create runs" + + @pytest.mark.asyncio + async def test_masked_tool_output_reaches_bedrock_masked(self, stub_aws_models: None) -> None: + guardrail: Final = MaskingGuardrail( + "hunter2", guardrail_name="mask_secrets", event_hook=GuardrailEventHooks.pre_call, default_on=True + ) + tool_output: Final = json.dumps( + { + "type": "conversation.item.create", + "item": {"type": "function_call_output", "call_id": "tool_1", "output": self.BLOCKED_TEXT}, + } + ) + + stream, client_ws = await self._forward_client([tool_output], guardrail) + + events: Final = self._bedrock_events(stream) + assert [event["toolResult"]["content"] for event in events if "toolResult" in event] == [ + "my secret is [MASKED]" + ] + assert client_ws.sent_to_client == [] + + @pytest.mark.asyncio + async def test_the_block_reason_is_spoken_when_no_message_is_configured(self, stub_aws_models: None) -> None: + guardrail: Final = BlockedWordGuardrail( + "hunter2", guardrail_name="no_secrets", event_hook=GuardrailEventHooks.pre_call, default_on=True + ) + + stream, client_ws = await self._forward_client( + [self._user_text(self.BLOCKED_TEXT), self.RESPONSE_CREATE], guardrail + ) + + assert self.BLOCKED_TEXT not in json.dumps(self._bedrock_events(stream)) + assert ['"hunter2 is not allowed"' in text for text in self._text_inputs(stream)] == [True] + assert [error["message"] for error in self._guardrail_errors(client_ws)] == ["hunter2 is not allowed"] + + @pytest.mark.asyncio + async def test_guardrail_that_crashes_ends_the_session_as_an_outage_without_forwarding_the_text( + self, stub_aws_models: None + ) -> None: + guardrail: Final = CrashingGuardrail( + "hunter2", guardrail_name="flaky", event_hook=GuardrailEventHooks.pre_call, default_on=True + ) + stream: Final = SessionBoundBedrockStream() + client_ws: Final = ConnectedClientWS([self._user_text(self.CLEAN_TEXT), self.RESPONSE_CREATE]) + logging_obj: Final = FakeLogging() + + async def open_stream() -> SessionBoundBedrockStream: + return stream + + await asyncio.wait_for( + BedrockRealtime()._run_session( + client_ws, + open_stream, + self.MODEL, + logging_obj, + None, + guardrails=BedrockRealtimeGuardrails(callbacks=lambda: (guardrail,)), + ), + timeout=2, + ) + + assert self.CLEAN_TEXT not in json.dumps(self._bedrock_events(stream)) + assert [event["error"] for event in self._client_events(client_ws) if event["type"] == "error"] == [ + {"type": "guardrail_error", "message": GUARDRAIL_UNAVAILABLE_MESSAGE} + ] + assert client_ws.closed + assert logging_obj.model_call_details.get(REALTIME_SESSION_ENDED_ABNORMALLY_KEY) is True + + @pytest.mark.asyncio + async def test_blocked_tool_output_is_replaced_and_answered_by_the_model_itself( + self, stub_aws_models: None + ) -> None: + tool_output: Final = json.dumps( + { + "type": "conversation.item.create", + "item": {"type": "function_call_output", "call_id": "tool_1", "output": self.BLOCKED_TEXT}, + } + ) + + stream, client_ws = await self._forward_client([tool_output, self.RESPONSE_CREATE], self._guardrail()) + + events: Final = self._bedrock_events(stream) + assert "hunter2" not in json.dumps(events) + tool_result_index: Final = next(i for i, event in enumerate(events) if "toolResult" in event) + assert events[tool_result_index]["toolResult"]["content"] == BLOCKED_TOOL_OUTPUT + tool_starts: Final = [ + event["contentStart"] for event in events if event.get("contentStart", {}).get("type") == "TOOL" + ] + assert [start["toolResultInputConfiguration"]["toolUseId"] for start in tool_starts] == ["tool_1"] + assert self._text_inputs(stream) == [] + assert [error["message"] for error in self._guardrail_errors(client_ws)] == [self.VIOLATION] + assert any("audioInput" in event for event in events[tool_result_index:]), "the client's response.create runs" + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("event_hook", "default_on"), + [ + (GuardrailEventHooks.pre_call, False), + (GuardrailEventHooks.post_call, True), + (GuardrailEventHooks.realtime_input_transcription, True), + ], + ids=["not-default-on", "post-call-hook", "transcription-hook"], + ) + async def test_guardrail_that_does_not_apply_to_typed_text_changes_nothing( + self, stub_aws_models: None, event_hook: GuardrailEventHooks, default_on: bool + ) -> None: + guardrail: Final = self._guardrail(event_hook=event_hook, default_on=default_on) + + stream, client_ws = await self._forward_client([self._user_text(self.BLOCKED_TEXT)], guardrail) + + assert self._text_inputs(stream) == [self.BLOCKED_TEXT] + assert client_ws.sent_to_client == [] + assert guardrail.checked == [] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("guardrail_kwargs", "blocked_messages"), + [({"on_violation": "end_session"}, 1), ({"end_session_after_n_fails": 2}, 2)], + ids=["on-violation-end-session", "after-two-fails"], + ) + async def test_end_session_guardrail_closes_the_session( + self, stub_aws_models: None, guardrail_kwargs: SessionEndOptions, blocked_messages: int + ) -> None: + messages: Final = [self._user_text(self.BLOCKED_TEXT)] * blocked_messages + [self._user_text(self.CLEAN_TEXT)] + + stream, client_ws = await self._forward_client( + messages, self._guardrail(**guardrail_kwargs), client_ws_type=ConnectedClientWS + ) + + assert client_ws.closed + assert len(self._guardrail_errors(client_ws)) == blocked_messages + assert self.CLEAN_TEXT not in self._text_inputs(stream) + assert sum(self.VIOLATION in text for text in self._text_inputs(stream)) == blocked_messages - 1 + assert stream.input_stream.closed + + @pytest.mark.asyncio + @pytest.mark.parametrize("guardrail_kwargs", [{}, {"on_violation": "end_session"}], ids=["default", "end-session"]) + async def test_blocked_transcript_ends_the_session_before_its_answer_reaches_the_client( + self, guardrail_kwargs: SessionEndOptions + ) -> None: + payloads: Final = ( + *self._user_transcript(self.BLOCKED_TEXT), + *self._assistant_text("Your secret is safe with me"), + *self._user_transcript(self.CLEAN_TEXT), + *self._assistant_text("It is sunny"), + ) + guardrail: Final = self._guardrail( + event_hook=GuardrailEventHooks.realtime_input_transcription, **guardrail_kwargs + ) + + client_ws, _ = await self._forward_bedrock(payloads, guardrail) + + client_events: Final = self._client_events(client_ws) + assert client_ws.closed + assert not any(event["type"].startswith("response.") for event in client_events) + assert [error["message"] for error in self._guardrail_errors(client_ws)] == [self.VIOLATION] + assert guardrail.checked == [((self.BLOCKED_TEXT,), None)] + + @pytest.mark.asyncio + async def test_blocked_transcript_answer_never_reaches_the_client_when_the_previous_answer_closes_late( + self, + ) -> None: + earlier_answer: Final = self._assistant_text("Earlier answer") + payloads: Final = ( + *earlier_answer[:2], + *self._user_transcript(self.BLOCKED_TEXT), + earlier_answer[2], + *self._assistant_text("Your secret is safe with me"), + ) + + client_ws, _ = await self._forward_bedrock( + payloads, self._guardrail(event_hook=GuardrailEventHooks.realtime_input_transcription) + ) + + client_events: Final = self._client_events(client_ws) + assert "Your secret is safe with me" not in json.dumps(client_events) + assert [event["delta"] for event in client_events if event["type"] == "response.text.delta"] == [ + "Earlier answer" + ] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("event_hook", "default_on", "client_audio_streamed"), + [ + (GuardrailEventHooks.realtime_input_transcription, True, False), + (GuardrailEventHooks.pre_call, True, True), + (GuardrailEventHooks.realtime_input_transcription, False, True), + ], + ids=["typed-text-session", "pre-call-hook", "not-default-on"], + ) + async def test_transcript_guardrail_that_does_not_apply_changes_nothing( + self, event_hook: GuardrailEventHooks, default_on: bool, client_audio_streamed: bool + ) -> None: + guardrail: Final = self._guardrail(event_hook=event_hook, default_on=default_on) + + client_ws, _ = await self._forward_bedrock( + (*self._user_transcript(self.BLOCKED_TEXT), *self._assistant_text("Your secret is safe with me")), + guardrail, + client_audio_streamed=client_audio_streamed, + ) + + client_events: Final = self._client_events(client_ws) + assert [event["delta"] for event in client_events if event["type"] == "response.text.delta"] == [ + "Your secret is safe with me" + ] + assert self._guardrail_errors(client_ws) == [] + assert guardrail.checked == [] + + @pytest.mark.asyncio + async def test_session_ends_when_the_client_leaves_before_hearing_about_a_blocked_message( + self, stub_aws_models: None + ) -> None: + stream, client_ws = await self._forward_client( + [self._user_text(self.BLOCKED_TEXT), self._user_text(self.CLEAN_TEXT)], + self._guardrail(), + client_ws_type=GuardrailErrorUndeliveredClientWS, + ) + + assert self._text_inputs(stream) == [] + assert self._client_events(client_ws) == [] + + @pytest.mark.asyncio + async def test_answers_stop_when_the_client_leaves_before_hearing_about_a_blocked_transcript(self) -> None: + guardrail: Final = self._guardrail(event_hook=GuardrailEventHooks.realtime_input_transcription) + + _, logged = await self._forward_bedrock( + (*self._user_transcript(self.BLOCKED_TEXT), *self._assistant_text("Your secret is safe with me")), + guardrail, + client_ws_type=GuardrailErrorUndeliveredClientWS, + ) + + assert not any(event["type"].startswith("response.") for event in logged) + + @pytest.mark.asyncio + async def test_proxy_request_guardrails_and_key_reach_the_bedrock_session( + self, stub_aws_sdk_client: dict[str, object], monkeypatch: pytest.MonkeyPatch + ) -> None: + class ProxyLogging(FakeLogging): + def update_from_kwargs(self, **kwargs: object) -> None: + pass + + guardrail: Final = self._guardrail(default_on=False) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + stream: Final = ScriptedBedrockStream([], receiver_type=DrainedThenOpenBedrockReceiver) + stub_aws_sdk_client["streams"] = [stream] + user_api_key_dict: Final = object() + client_ws: Final = DisconnectingClientWS([self._user_text(self.BLOCKED_TEXT)]) + + await asyncio.wait_for( + realtime_main._arealtime.__wrapped__( + model=f"bedrock/{self.MODEL}", + websocket=client_ws, + litellm_logging_obj=ProxyLogging(), + aws_region_name="us-east-1", + aws_access_key_id="k", + aws_secret_access_key="s", + user_api_key_dict=user_api_key_dict, + metadata={"guardrails": ["no_secrets"]}, + ), + timeout=1, + ) + + assert guardrail.checked == [((self.BLOCKED_TEXT,), user_api_key_dict)] + assert "hunter2" not in json.dumps(self._bedrock_events(stream)) + assert [error["message"] for error in self._guardrail_errors(client_ws)] == [self.VIOLATION] + + @pytest.mark.asyncio + async def test_a_key_that_opted_out_of_global_guardrails_is_not_screened( + self, stub_aws_sdk_client: dict[str, object], monkeypatch: pytest.MonkeyPatch + ) -> None: + class ProxyLogging(FakeLogging): + def update_from_kwargs(self, **kwargs: object) -> None: + pass + + guardrail: Final = self._guardrail() + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + stream: Final = ScriptedBedrockStream([], receiver_type=DrainedThenOpenBedrockReceiver) + stub_aws_sdk_client["streams"] = [stream] + client_ws: Final = DisconnectingClientWS([self._user_text(self.BLOCKED_TEXT)]) + + await asyncio.wait_for( + realtime_main._arealtime.__wrapped__( + model=f"bedrock/{self.MODEL}", + websocket=client_ws, + litellm_logging_obj=ProxyLogging(), + aws_region_name="us-east-1", + aws_access_key_id="k", + aws_secret_access_key="s", + metadata={"user_api_key_metadata": {"disable_global_guardrails": True}}, + ), + timeout=1, + ) + + assert guardrail.checked == [] + assert self._guardrail_errors(client_ws) == [] + assert self.BLOCKED_TEXT in json.dumps(self._bedrock_events(stream)) + + if __name__ == "__main__": pytest.main([__file__, "-v"]) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 422e355a03c..6e0a7898b6e 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -28081,7 +28081,7 @@ export interface components { on_sensitive_data?: ("block" | "route") | null; /** * On Violation - * @description For /v1/realtime sessions: 'warn' speaks the violation message and continues; 'end_session' speaks the message and closes the connection. For guardrail='mcp_security': 'block' rejects the request; 'alert' only logs a warning. + * @description For /v1/realtime sessions: 'warn' speaks the violation message and continues; 'end_session' speaks the message and closes the connection. Bedrock Nova Sonic cannot cancel a spoken turn's answer, so a blocked spoken turn always ends its session. For guardrail='mcp_security': 'block' rejects the request; 'alert' only logs a warning. */ on_violation?: ("warn" | "end_session" | "block" | "alert") | null; /** @@ -28112,7 +28112,7 @@ export interface components { patterns?: components["schemas"]["ContentFilterPattern"][] | null; /** * Realtime Violation Message - * @description The message the bot speaks aloud when a /v1/realtime guardrail fires. Falls back to violation_message_template if not set. + * @description The message the bot speaks aloud when a /v1/realtime guardrail fires. Falls back to the guardrail's block reason if not set. */ realtime_violation_message?: string | null; /** @@ -38228,7 +38228,7 @@ export interface components { on_sensitive_data?: ("block" | "route") | null; /** * On Violation - * @description For /v1/realtime sessions: 'warn' speaks the violation message and continues; 'end_session' speaks the message and closes the connection. For guardrail='mcp_security': 'block' rejects the request; 'alert' only logs a warning. + * @description For /v1/realtime sessions: 'warn' speaks the violation message and continues; 'end_session' speaks the message and closes the connection. Bedrock Nova Sonic cannot cancel a spoken turn's answer, so a blocked spoken turn always ends its session. For guardrail='mcp_security': 'block' rejects the request; 'alert' only logs a warning. */ on_violation?: ("warn" | "end_session" | "block" | "alert") | null; /** @@ -38386,7 +38386,7 @@ export interface components { prompt_injections?: boolean | null; /** * Realtime Violation Message - * @description The message the bot speaks aloud when a /v1/realtime guardrail fires. Falls back to violation_message_template if not set. + * @description The message the bot speaks aloud when a /v1/realtime guardrail fires. Falls back to the guardrail's block reason if not set. */ realtime_violation_message?: string | null; /**