scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-08 08:08:37 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 80

fix(bedrock): run realtime guardrails on Nova Sonic sessions (#49833)

Browse source 
* fix(bedrock): run realtime guardrails on Nova Sonic sessions

* fix(bedrock): quote the spoken violation message and skip it after a blocked Nova Sonic tool output

* fix(bedrock): screen every message role, end Nova Sonic sessions on blocked transcripts, drop unparseable frames

* fix(realtime): apply guardrail masks, keep block reasons out of spoken notices, fail closed on guardrail crashes

* fix(realtime): screen joined text parts, keep the documented notice, honor key opt-outs

Text parts of one message are joined before screening again, so a phrase split
across parts is still blocked, and a mask collapses them into one part. The
spoken notice is realtime_violation_message or the block reason, quoted only on
Nova Sonic. Client-written output_text and text parts are screened. Key and team
guardrail opt-outs reach regular and translation realtime sessions. A non-string
upstream transcript skips screening as before.

* fix(realtime): honor guardrail role scope, clear stale response.create drops, screen split keywords, report Nova guardrail outages

* fix(realtime): keep every per-part mask when a joined text masks differently

* fix(realtime): screen one joined text per item and pass the session's tenant context to guardrails

Multi-part items reach each guardrail once as the space-joined text, as on main, and only an item a
guardrail rewrites gets its parts screened again to keep per-part masks. Guardrails now receive the
session request data with its metadata, so key and team policies apply, and a ModifyResponseException
block is reported as a block. Nova scopes every message item as the user turn it sends to Bedrock.

* fix(realtime): catch phrases split mid-word across text parts without re-screening single-part items

* fix(realtime): screen each joined text in its own guardrail call and keep a Nova opening message for the retry while it is screened

* fix(realtime): screen text parts in a loop so large messages keep their masks

* fix(realtime): fail closed on a non-string transcript under a guardrail and give each turn its own guardrail metadata

* fix(realtime): fail a non-string transcript only when a guardrail would screen the user's turn

---------

Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
tavren27 2026-11-06 13:18:48 -08:00 • committed by GitHub
parent b5a72c2042
commit 2d2290067c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
12 changed files with 2294 additions and 247 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

316
litellm/litellm_core_utils/realtime_guardrails.py Normal file
View file

@ -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"}),
)

296
litellm/litellm_core_utils/realtime_streaming.py
View file

@ -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

266
litellm/llms/bedrock/realtime/handler.py
View file

@ -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

8
litellm/proxy/_lazy_openapi_snapshot.json
View file

@ -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": {

10
litellm/realtime_api/main.py
View file

@ -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 = (

3
litellm/types/guardrails.py
View file

@ -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

268
tests/integration/providers/test_bedrock_nova_sonic_realtime_guardrails.py Normal file
View file

@ -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

14
tests/integration/providers/test_realtime_transcription_guardrail_session_update.py
View file

@ -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))

481
tests/unit/litellm_core_utils/test_realtime_guardrails.py Normal file
View file

@ -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

217
tests/unit/litellm_core_utils/test_realtime_streaming.py
View file

@ -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):
"""

654
tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py
View file

@ -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"])

8
ui/litellm-dashboard/src/lib/http/schema.d.ts generated vendored
View file

@ -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;
/**