mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-12 03:48:56 +00:00
Merge 63cd43143e into 780c504e00
This commit is contained in:
moyai-devin-berriai[bot]
GitHub
commit
e78fac2e9a
3 changed files with 2003 additions and 206 deletions
|
|
@ -1,13 +1,34 @@
|
|||
from collections.abc import Mapping, Sequence
|
||||
import copy
|
||||
import json
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from functools import partial, reduce
|
||||
from itertools import chain, groupby
|
||||
from typing import TYPE_CHECKING, Any, Final, Optional, Protocol, TypeAlias
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
|
||||
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,
|
||||
scoped_structured_message_indices,
|
||||
unappliable_request_rewrite,
|
||||
)
|
||||
from litellm.types.llms.openai import (
|
||||
AllMessageValues,
|
||||
ChatCompletionAssistantMessage,
|
||||
ChatCompletionAssistantToolCall,
|
||||
ChatCompletionSystemMessage,
|
||||
ChatCompletionTextObject,
|
||||
ChatCompletionToolCallFunctionChunk,
|
||||
ChatCompletionToolMessage,
|
||||
ChatCompletionToolParam,
|
||||
ChatCompletionToolParamFunctionChunk,
|
||||
ChatCompletionUserMessage,
|
||||
)
|
||||
from litellm.types.utils import GenericGuardrailAPIInputs
|
||||
|
||||
|
|
@ -83,6 +104,7 @@ def _extract_converse_texts(
|
|||
body: dict,
|
||||
skip_system: bool,
|
||||
skip_tool: bool,
|
||||
scan_only_tool_results: bool = False,
|
||||
) -> tuple[list[str], list[_StringHolder]]:
|
||||
"""
|
||||
Walk a Bedrock Converse request body and collect text content.
|
||||
|
|
@ -102,8 +124,8 @@ def _extract_converse_texts(
|
|||
inflating guardrail usage for a request whose only prompt is one user
|
||||
message. No other guardrail translation handler puts tool definitions in
|
||||
``texts``; the chat and messages handlers carry them in the structured
|
||||
``tools`` input instead, which this handler does not populate because a
|
||||
Bedrock ``toolSpec`` is not the OpenAI tool shape those consumers expect.
|
||||
``tools`` input instead, which this handler also populates by converting
|
||||
each Bedrock ``toolSpec`` to the OpenAI tool shape.
|
||||
|
||||
``additionalModelRequestFields`` is treated differently on purpose. Bedrock
|
||||
gives ``toolConfig.tools`` a fixed schema whose contents are tool metadata by
|
||||
|
|
@ -113,7 +135,7 @@ def _extract_converse_texts(
|
|||
"""
|
||||
holders: Final[list[_StringHolder]] = []
|
||||
|
||||
if not skip_system:
|
||||
if not skip_system and not scan_only_tool_results:
|
||||
for block in body.get("system") or []:
|
||||
if isinstance(block, dict):
|
||||
_collect_block_text(block, holders)
|
||||
|
|
@ -126,6 +148,8 @@ def _extract_converse_texts(
|
|||
continue
|
||||
if skip_tool and ("toolUse" in block or "toolResult" in block):
|
||||
continue
|
||||
if scan_only_tool_results and "toolResult" not in block:
|
||||
continue
|
||||
_collect_block_text(block, holders)
|
||||
tool_use = block.get("toolUse")
|
||||
if isinstance(tool_use, dict):
|
||||
|
|
@ -137,12 +161,507 @@ def _extract_converse_texts(
|
|||
_collect_block_text(inner, holders)
|
||||
_collect_strings(inner.get("json"), holders)
|
||||
|
||||
_collect_strings(body.get("additionalModelRequestFields"), holders)
|
||||
if not scan_only_tool_results:
|
||||
_collect_strings(body.get("additionalModelRequestFields"), holders)
|
||||
|
||||
texts: Final = [container[key] for container, key in holders]
|
||||
return texts, holders
|
||||
|
||||
|
||||
class _ConverseModel(BaseModel):
|
||||
model_config = ConfigDict(extra="ignore", frozen=True)
|
||||
|
||||
|
||||
class _ConverseGuardText(_ConverseModel):
|
||||
text: str = ""
|
||||
|
||||
|
||||
class _ConverseGuardContent(_ConverseModel):
|
||||
text: _ConverseGuardText | None = None
|
||||
|
||||
|
||||
class _ConverseToolUse(_ConverseModel):
|
||||
tool_use_id: str = Field(default="", alias="toolUseId")
|
||||
name: str = ""
|
||||
tool_input: object = Field(default=None, alias="input")
|
||||
|
||||
|
||||
class _ConverseToolResultContent(_ConverseModel):
|
||||
text: str = ""
|
||||
json_value: object = Field(default=None, alias="json")
|
||||
|
||||
|
||||
class _ConverseToolResult(_ConverseModel):
|
||||
tool_use_id: str = Field(default="", alias="toolUseId")
|
||||
content: tuple[_ConverseToolResultContent, ...] = ()
|
||||
|
||||
|
||||
class _ConverseContentBlock(_ConverseModel):
|
||||
text: str = ""
|
||||
guard_content: _ConverseGuardContent | None = Field(default=None, alias="guardContent")
|
||||
tool_use: _ConverseToolUse | None = Field(default=None, alias="toolUse")
|
||||
tool_result: _ConverseToolResult | None = Field(default=None, alias="toolResult")
|
||||
|
||||
|
||||
class _ConverseMessage(_ConverseModel):
|
||||
role: str = ""
|
||||
content: tuple[_ConverseContentBlock, ...] = ()
|
||||
|
||||
|
||||
class _ConverseToolInputSchema(_ConverseModel):
|
||||
json_schema: Mapping[str, object] | None = Field(default=None, alias="json")
|
||||
|
||||
|
||||
class _ConverseToolSpec(_ConverseModel):
|
||||
name: str = ""
|
||||
description: str = ""
|
||||
input_schema: _ConverseToolInputSchema | None = Field(default=None, alias="inputSchema")
|
||||
|
||||
|
||||
class _ConverseTool(_ConverseModel):
|
||||
tool_spec: _ConverseToolSpec | None = Field(default=None, alias="toolSpec")
|
||||
|
||||
|
||||
class _ConverseToolConfig(_ConverseModel):
|
||||
tools: tuple[_ConverseTool, ...] = ()
|
||||
|
||||
|
||||
class _ConverseRequest(_ConverseModel):
|
||||
system: tuple[_ConverseContentBlock, ...] = ()
|
||||
messages: tuple[_ConverseMessage, ...] = ()
|
||||
tool_config: _ConverseToolConfig | None = Field(default=None, alias="toolConfig")
|
||||
|
||||
|
||||
def _parse_converse_request(validate: Callable[[], _ConverseRequest]) -> _ConverseRequest | None:
|
||||
try:
|
||||
return validate()
|
||||
except ValidationError:
|
||||
verbose_proxy_logger.debug("BedrockPassthroughGuardrailHandler: Converse body not convertible to messages")
|
||||
return None
|
||||
|
||||
|
||||
def _converse_block_text(block: _ConverseContentBlock) -> str:
|
||||
if block.text:
|
||||
return block.text
|
||||
guard_text: Final = block.guard_content.text if block.guard_content else None
|
||||
return guard_text.text if guard_text else ""
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _BlockRef:
|
||||
message_index: int | None
|
||||
block_index: int
|
||||
path: tuple[str | int, ...] = ()
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _ConvertedMessage:
|
||||
message: AllMessageValues
|
||||
text_blocks: tuple[_BlockRef, ...]
|
||||
tool_result: tuple[_BlockRef, _ConverseToolResult] | None = None
|
||||
tool_uses: tuple[tuple[_BlockRef, _ConverseToolUse], ...] = ()
|
||||
|
||||
|
||||
_IndexedBlock: TypeAlias = tuple[_BlockRef, _ConverseContentBlock]
|
||||
_TEXT_PARTS: Final = TypeAdapter(tuple[ChatCompletionTextObject, ...])
|
||||
_OBJECT: Final = TypeAdapter(Mapping[str, object])
|
||||
_ITEMS: Final = TypeAdapter(Sequence[object])
|
||||
_PARSED_JSON: Final = TypeAdapter(object)
|
||||
_TextEdit: TypeAlias = tuple[_BlockRef, object]
|
||||
_JsonPair: TypeAlias = tuple[tuple[str | int, ...], object, object]
|
||||
|
||||
|
||||
def _indexed_blocks(message_index: int | None, blocks: Sequence[_ConverseContentBlock]) -> tuple[_IndexedBlock, ...]:
|
||||
return tuple((_BlockRef(message_index, block_index), block) for block_index, block in enumerate(blocks))
|
||||
|
||||
|
||||
def _text_blocks(blocks: Sequence[_IndexedBlock]) -> tuple[tuple[_BlockRef, str], ...]:
|
||||
return tuple((ref, text) for ref, block in blocks if (text := _converse_block_text(block)))
|
||||
|
||||
|
||||
def _text_content(text_blocks: Sequence[tuple[_BlockRef, str]]) -> tuple[ChatCompletionTextObject, ...]:
|
||||
return tuple(ChatCompletionTextObject(type="text", text=text) for _, text in text_blocks)
|
||||
|
||||
|
||||
def _refs(text_blocks: Sequence[tuple[_BlockRef, str]]) -> tuple[_BlockRef, ...]:
|
||||
return tuple(ref for ref, _ in text_blocks)
|
||||
|
||||
|
||||
def _scannable_tool_parts(tool_result: _ConverseToolResult) -> tuple[tuple[int, _ConverseToolResultContent], ...]:
|
||||
return tuple(
|
||||
(index, part) for index, part in enumerate(tool_result.content) if part.text or part.json_value is not None
|
||||
)
|
||||
|
||||
|
||||
def _converse_tool_message(tool_result: _ConverseToolResult) -> ChatCompletionToolMessage:
|
||||
return ChatCompletionToolMessage(
|
||||
role="tool",
|
||||
tool_call_id=tool_result.tool_use_id,
|
||||
content="\n".join(part.text or json.dumps(part.json_value) for _, part in _scannable_tool_parts(tool_result)),
|
||||
)
|
||||
|
||||
|
||||
def _converse_tool_call(tool_use: _ConverseToolUse) -> ChatCompletionAssistantToolCall:
|
||||
return ChatCompletionAssistantToolCall(
|
||||
id=tool_use.tool_use_id,
|
||||
type="function",
|
||||
function=ChatCompletionToolCallFunctionChunk(
|
||||
name=tool_use.name,
|
||||
arguments="{}" if tool_use.tool_input is None else json.dumps(tool_use.tool_input),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _converse_user_group(is_tool_result: bool, group: Sequence[_IndexedBlock]) -> tuple[_ConvertedMessage, ...]:
|
||||
if is_tool_result:
|
||||
return tuple(
|
||||
_ConvertedMessage(_converse_tool_message(block.tool_result), (), (ref, block.tool_result))
|
||||
for ref, block in group
|
||||
if block.tool_result
|
||||
)
|
||||
text_blocks: Final = _text_blocks(group)
|
||||
if not text_blocks:
|
||||
return ()
|
||||
message: Final = ChatCompletionUserMessage(role="user", content=list(_text_content(text_blocks)))
|
||||
return (_ConvertedMessage(message, _refs(text_blocks)),)
|
||||
|
||||
|
||||
def _converse_user_messages(blocks: Sequence[_IndexedBlock]) -> tuple[_ConvertedMessage, ...]:
|
||||
return tuple(
|
||||
chain.from_iterable(
|
||||
_converse_user_group(is_tool_result, tuple(group))
|
||||
for is_tool_result, group in groupby(blocks, key=lambda item: item[1].tool_result is not None)
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _converse_assistant_messages(
|
||||
blocks: Sequence[_IndexedBlock], include_tool_calls: bool
|
||||
) -> tuple[_ConvertedMessage, ...]:
|
||||
text_blocks: Final = _text_blocks(blocks)
|
||||
content: Final = list(_text_content(text_blocks))
|
||||
tool_calls: Final = [
|
||||
_converse_tool_call(block.tool_use) for _, block in blocks if block.tool_use and include_tool_calls
|
||||
]
|
||||
if tool_calls:
|
||||
message: Final = ChatCompletionAssistantMessage(
|
||||
role="assistant", content=content or None, tool_calls=tool_calls
|
||||
)
|
||||
return (
|
||||
_ConvertedMessage(
|
||||
message,
|
||||
_refs(text_blocks),
|
||||
tool_uses=tuple((ref, block.tool_use) for ref, block in blocks if block.tool_use),
|
||||
),
|
||||
)
|
||||
if not content:
|
||||
return ()
|
||||
return (_ConvertedMessage(ChatCompletionAssistantMessage(role="assistant", content=content), _refs(text_blocks)),)
|
||||
|
||||
|
||||
def _converse_message(
|
||||
message_index: int, message: _ConverseMessage, include_tool_calls: bool
|
||||
) -> tuple[_ConvertedMessage, ...]:
|
||||
blocks: Final = _indexed_blocks(message_index, message.content)
|
||||
if message.role == "assistant":
|
||||
return _converse_assistant_messages(blocks, include_tool_calls)
|
||||
return _converse_user_messages(blocks)
|
||||
|
||||
|
||||
def _converse_structured_messages(request: _ConverseRequest, include_tool_calls: bool) -> tuple[_ConvertedMessage, ...]:
|
||||
converted: Final = tuple(
|
||||
chain.from_iterable(
|
||||
_converse_message(index, message, include_tool_calls) for index, message in enumerate(request.messages)
|
||||
)
|
||||
)
|
||||
system_blocks: Final = _text_blocks(_indexed_blocks(None, request.system))
|
||||
if not system_blocks:
|
||||
return converted
|
||||
system: Final = ChatCompletionSystemMessage(role="system", content=list(_text_content(system_blocks)))
|
||||
return (_ConvertedMessage(system, _refs(system_blocks)), *converted)
|
||||
|
||||
|
||||
def _without(mapping: Mapping[str, object], key: str) -> tuple[tuple[str, object], ...]:
|
||||
return tuple(sorted(((name, value) for name, value in mapping.items() if name != key), key=lambda item: item[0]))
|
||||
|
||||
|
||||
def _without_content(message: Mapping[str, object]) -> tuple[tuple[str, object], ...]:
|
||||
return _without(message, "content")
|
||||
|
||||
|
||||
def _fields(message: AllMessageValues) -> Mapping[str, object]:
|
||||
return message
|
||||
|
||||
|
||||
def _as_object(value: object) -> Mapping[str, object] | None:
|
||||
try:
|
||||
return _OBJECT.validate_python(value)
|
||||
except ValidationError:
|
||||
return None
|
||||
|
||||
|
||||
def _returned_text_parts(message: Mapping[str, object]) -> tuple[str, ...] | None:
|
||||
content: Final = message.get("content")
|
||||
if isinstance(content, str):
|
||||
return (content,)
|
||||
try:
|
||||
parts: Final = _TEXT_PARTS.validate_python(content)
|
||||
except ValidationError:
|
||||
return None
|
||||
return tuple(part["text"] for part in parts)
|
||||
|
||||
|
||||
def _merged_edits(children: Sequence[tuple[_TextEdit, ...] | None]) -> tuple[_TextEdit, ...] | None:
|
||||
if any(child is None for child in children):
|
||||
return None
|
||||
return tuple(chain.from_iterable(child for child in children if child is not None))
|
||||
|
||||
|
||||
def _json_children(path: tuple[str | int, ...], original: object, returned: object) -> tuple[_JsonPair, ...] | None:
|
||||
original_object: Final = _as_object(original)
|
||||
returned_object: Final = _as_object(returned)
|
||||
if original_object is not None and returned_object is not None:
|
||||
if original_object.keys() != returned_object.keys():
|
||||
return None
|
||||
return tuple(((*path, key), value, returned_object[key]) for key, value in original_object.items())
|
||||
original_items: Final = _as_items(original)
|
||||
returned_items: Final = _as_items(returned)
|
||||
if original_items is None or returned_items is None or len(original_items) != len(returned_items):
|
||||
return None
|
||||
return tuple(((*path, index), value, returned_items[index]) for index, value in enumerate(original_items))
|
||||
|
||||
|
||||
def _json_edits(
|
||||
ref: _BlockRef, path: tuple[str | int, ...], original: object, returned: object
|
||||
) -> tuple[_TextEdit, ...] | None:
|
||||
"""Only write changed leaves, preserving independent flat-text redactions."""
|
||||
pending: Final[list[_JsonPair]] = [(path, original, returned)] # mutable-ok: work stack instead of recursion
|
||||
edits: Final[list[_TextEdit]] = [] # mutable-ok: collects the leaves the work stack visits
|
||||
while pending:
|
||||
leaf_path, before, after = pending.pop()
|
||||
if before == after:
|
||||
continue
|
||||
if isinstance(before, str) and isinstance(after, str):
|
||||
edits.append((_BlockRef(ref.message_index, ref.block_index, leaf_path), after))
|
||||
continue
|
||||
if (children := _json_children(leaf_path, before, after)) is None:
|
||||
return None
|
||||
pending.extend(children)
|
||||
return tuple(edits)
|
||||
|
||||
|
||||
def _as_items(value: object) -> Sequence[object] | None:
|
||||
if not isinstance(value, list):
|
||||
return None
|
||||
try:
|
||||
return _ITEMS.validate_python(value)
|
||||
except ValidationError:
|
||||
return None
|
||||
|
||||
|
||||
def _tool_part_edits(
|
||||
ref: _BlockRef, index: int, part: _ConverseToolResultContent, text: str
|
||||
) -> tuple[_TextEdit, ...] | None:
|
||||
if part.text:
|
||||
return ((_BlockRef(ref.message_index, ref.block_index, ("toolResult", "content", index, "text")), text),)
|
||||
try:
|
||||
value: Final = _PARSED_JSON.validate_json(text)
|
||||
except ValidationError:
|
||||
return None
|
||||
return _json_edits(ref, ("toolResult", "content", index, "json"), part.json_value, value)
|
||||
|
||||
|
||||
def _tool_result_edits(converted: _ConvertedMessage, returned: Mapping[str, object]) -> tuple[_TextEdit, ...] | None:
|
||||
if converted.tool_result is None or _without_content(returned) != _without_content(converted.message):
|
||||
return None
|
||||
ref, result = converted.tool_result
|
||||
texts: Final = _returned_text_parts(returned)
|
||||
parts: Final = _scannable_tool_parts(result)
|
||||
if texts is None or len(texts) != len(parts):
|
||||
return None
|
||||
return _merged_edits(tuple(_tool_part_edits(ref, index, part, text) for (index, part), text in zip(parts, texts)))
|
||||
|
||||
|
||||
def _call_function(call: object) -> Mapping[str, object] | None:
|
||||
call_object: Final = _as_object(call)
|
||||
return _as_object(call_object.get("function")) if call_object is not None else None
|
||||
|
||||
|
||||
def _tool_call_edits(
|
||||
tool_use: tuple[_BlockRef, _ConverseToolUse], before: object, after: object
|
||||
) -> tuple[_TextEdit, ...] | None:
|
||||
if before == after:
|
||||
return ()
|
||||
before_call: Final = _as_object(before)
|
||||
after_call: Final = _as_object(after)
|
||||
before_function: Final = _call_function(before)
|
||||
after_function: Final = _call_function(after)
|
||||
if before_call is None or after_call is None or before_function is None or after_function is None:
|
||||
return None
|
||||
arguments: Final = after_function.get("arguments")
|
||||
if (
|
||||
not isinstance(arguments, str)
|
||||
or _without(after_call, "function") != _without(before_call, "function")
|
||||
or _without(after_function, "arguments") != _without(before_function, "arguments")
|
||||
):
|
||||
return None
|
||||
try:
|
||||
value: Final = _PARSED_JSON.validate_json(arguments)
|
||||
except ValidationError:
|
||||
return None
|
||||
ref, tool = tool_use
|
||||
return _json_edits(ref, ("toolUse", "input"), tool.tool_input, value)
|
||||
|
||||
|
||||
def _tool_calls_edits(
|
||||
tool_uses: Sequence[tuple[_BlockRef, _ConverseToolUse]], before: object, after: object
|
||||
) -> tuple[_TextEdit, ...] | None:
|
||||
before_calls: Final = _as_items(before)
|
||||
after_calls: Final = _as_items(after)
|
||||
if before_calls is None or after_calls is None or len(before_calls) != len(after_calls):
|
||||
return None
|
||||
return _merged_edits(
|
||||
tuple(
|
||||
_tool_call_edits(tool_use, original, returned)
|
||||
for tool_use, original, returned in zip(tool_uses, before_calls, after_calls)
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _content_edits(converted: _ConvertedMessage, returned: Mapping[str, object]) -> tuple[_TextEdit, ...] | None:
|
||||
original: Final = _fields(converted.message)
|
||||
if returned.get("content") == original.get("content"):
|
||||
return ()
|
||||
texts: Final = _returned_text_parts(returned)
|
||||
before_texts: Final = _returned_text_parts(original)
|
||||
if texts is None or before_texts is None or len(texts) != len(converted.text_blocks):
|
||||
return None
|
||||
return tuple((ref, text) for ref, before, text in zip(converted.text_blocks, before_texts, texts) if text != before)
|
||||
|
||||
|
||||
def _message_text_edits(converted: _ConvertedMessage, returned: Mapping[str, object]) -> tuple[_TextEdit, ...] | None:
|
||||
if returned == converted.message:
|
||||
return ()
|
||||
if converted.tool_result is not None:
|
||||
return _tool_result_edits(converted, returned)
|
||||
original_calls: Final = _fields(converted.message).get("tool_calls")
|
||||
call_edits: Final = (
|
||||
_tool_calls_edits(converted.tool_uses, original_calls, returned.get("tool_calls"))
|
||||
if converted.tool_uses
|
||||
else ()
|
||||
)
|
||||
comparable: Final = {**returned, "tool_calls": original_calls} if converted.tool_uses else returned
|
||||
if call_edits is None or _without_content(comparable) != _without_content(converted.message):
|
||||
return None
|
||||
return _merged_edits((call_edits, _content_edits(converted, comparable)))
|
||||
|
||||
|
||||
def _leaf(body: Mapping[str, Any], ref: _BlockRef) -> _StringHolder:
|
||||
blocks: Final = body["system"] if ref.message_index is None else body["messages"][ref.message_index]["content"]
|
||||
block: Final = blocks[ref.block_index]
|
||||
if not ref.path:
|
||||
return (block, "text") if block.get("text") else (block["guardContent"]["text"], "text")
|
||||
return reduce(lambda node, key: node[key], ref.path[:-1], block), ref.path[-1]
|
||||
|
||||
|
||||
def _leaf_identity(holder: _StringHolder) -> tuple[int, str | int]:
|
||||
container, key = holder
|
||||
return id(container), key
|
||||
|
||||
|
||||
def _changed_leaves(
|
||||
holders: Sequence[_StringHolder], texts: Sequence[str], guardrailed_texts: Sequence[str]
|
||||
) -> frozenset[tuple[int, str | int]]:
|
||||
return frozenset(
|
||||
_leaf_identity(holder) for holder, before, after in zip(holders, texts, guardrailed_texts) if after != before
|
||||
)
|
||||
|
||||
|
||||
def _texts_applied_view(
|
||||
body: Mapping[str, Any], texts: Sequence[str], skip_system: bool, skip_tool: bool, scan_only_tool_results: bool
|
||||
) -> tuple[_ConvertedMessage, ...]:
|
||||
rewritten: Final = copy.deepcopy(dict(body))
|
||||
_, holders = _extract_converse_texts(rewritten, skip_system, skip_tool, scan_only_tool_results)
|
||||
_write_back_texts(texts, holders)
|
||||
return _converse_structured_messages(_ConverseRequest.model_validate(rewritten), not skip_tool)
|
||||
|
||||
|
||||
def _same_content(expected: Mapping[str, object], returned: Mapping[str, object]) -> bool:
|
||||
if expected.get("content") == returned.get("content"):
|
||||
return True
|
||||
expected_parts: Final = _returned_text_parts(expected)
|
||||
returned_parts: Final = _returned_text_parts(returned)
|
||||
return (
|
||||
expected_parts is not None
|
||||
and returned_parts is not None
|
||||
and "\n".join(expected_parts) == "\n".join(returned_parts)
|
||||
)
|
||||
|
||||
|
||||
def _reproduces(rewritten: _ConvertedMessage, returned: Mapping[str, object] | None) -> bool:
|
||||
"""Whether the flat-text rewrite already yields the row the guardrail returned."""
|
||||
return (
|
||||
returned is not None
|
||||
and _without_content(returned) == _without_content(rewritten.message)
|
||||
and _same_content(rewritten.message, returned)
|
||||
)
|
||||
|
||||
|
||||
def _structured_text_edits(
|
||||
leaf: Callable[[_BlockRef], _StringHolder],
|
||||
converted: Sequence[_ConvertedMessage],
|
||||
scoped_indices: Sequence[int],
|
||||
returned: Sequence[object],
|
||||
flat_changes: frozenset[tuple[int, str | int]],
|
||||
flat_view: Callable[[], tuple[_ConvertedMessage, ...]],
|
||||
) -> tuple[tuple[_StringHolder, object], ...] | None:
|
||||
"""
|
||||
Structured edits for the leaves the flat texts left unchanged.
|
||||
|
||||
A row whose rewrite cannot be mapped back onto its Converse blocks is accepted only when
|
||||
the flat-text rewrite already produces it.
|
||||
"""
|
||||
if len(returned) != len(scoped_indices):
|
||||
return None
|
||||
rows: Final = tuple(_as_object(message) for message in returned)
|
||||
proposed: Final = tuple(
|
||||
_message_text_edits(converted[index], row) if row is not None else None
|
||||
for index, row in zip(scoped_indices, rows)
|
||||
)
|
||||
unmapped: Final = tuple((index, row) for index, row, edits in zip(scoped_indices, rows, proposed) if edits is None)
|
||||
if unmapped:
|
||||
view: Final = flat_view()
|
||||
if len(view) != len(converted):
|
||||
return None
|
||||
if not all(_reproduces(view[index], row) for index, row in unmapped):
|
||||
return None
|
||||
resolved: Final = tuple(
|
||||
(leaf(ref), text) for ref, text in chain.from_iterable(edits for edits in proposed if edits is not None)
|
||||
)
|
||||
return tuple((holder, text) for holder, text in resolved if _leaf_identity(holder) not in flat_changes)
|
||||
|
||||
|
||||
def _write_back_structured_texts(edits: Sequence[tuple[_StringHolder, object]]) -> None:
|
||||
for (container, key), text in edits:
|
||||
container[key] = text
|
||||
|
||||
|
||||
def _converse_tools(request: _ConverseRequest) -> tuple[ChatCompletionToolParam, ...]:
|
||||
tools: Final = request.tool_config.tools if request.tool_config else ()
|
||||
return tuple(
|
||||
ChatCompletionToolParam(
|
||||
type="function",
|
||||
function=ChatCompletionToolParamFunctionChunk(
|
||||
name=spec.name,
|
||||
description=spec.description,
|
||||
parameters=dict(spec.input_schema.json_schema),
|
||||
),
|
||||
)
|
||||
for tool in tools
|
||||
if (spec := tool.tool_spec) and spec.name and spec.input_schema and spec.input_schema.json_schema is not None
|
||||
)
|
||||
|
||||
|
||||
def _extract_converse_output_texts(
|
||||
content_blocks: Sequence[object],
|
||||
) -> tuple[list[str], list[_StringHolder]]:
|
||||
|
|
@ -177,8 +696,8 @@ def _extract_converse_output_texts(
|
|||
|
||||
|
||||
def _write_back_texts(
|
||||
guardrailed_texts: list[str],
|
||||
holders: list[_StringHolder],
|
||||
guardrailed_texts: Sequence[str],
|
||||
holders: Sequence[_StringHolder],
|
||||
) -> None:
|
||||
if len(guardrailed_texts) < len(holders):
|
||||
verbose_proxy_logger.warning(
|
||||
|
|
@ -446,12 +965,29 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation):
|
|||
skip_system: Final = effective_skip_system_message_for_guardrail(guardrail_to_apply)
|
||||
skip_tool: Final = effective_skip_tool_message_for_guardrail(guardrail_to_apply)
|
||||
|
||||
texts, holders = _extract_converse_texts(body, skip_system, skip_tool)
|
||||
scan_only_tool_results: Final = effective_scan_only_tool_results_for_guardrail(guardrail_to_apply)
|
||||
texts, holders = _extract_converse_texts(body, skip_system, skip_tool, scan_only_tool_results)
|
||||
converse_request: Final = _parse_converse_request(lambda: _ConverseRequest.model_validate(body))
|
||||
converted: Final = _converse_structured_messages(converse_request, not skip_tool) if converse_request else ()
|
||||
structured_messages: Final = tuple(item.message for item in converted)
|
||||
scoped_message_indices: Final = scoped_structured_message_indices(
|
||||
structured_messages,
|
||||
scan_only_tool_results=scan_only_tool_results,
|
||||
skip_system=skip_system,
|
||||
skip_tool=skip_tool,
|
||||
)
|
||||
|
||||
if not texts:
|
||||
if not texts and not scoped_message_indices:
|
||||
return data
|
||||
|
||||
inputs: Final = GenericGuardrailAPIInputs(texts=texts)
|
||||
inputs: Final = GenericGuardrailAPIInputs(texts=list(texts))
|
||||
if scoped_message_indices:
|
||||
inputs["structured_messages"] = copy.deepcopy(
|
||||
[structured_messages[index] for index in scoped_message_indices]
|
||||
)
|
||||
tools: Final = _converse_tools(converse_request) if converse_request and not scan_only_tool_results else ()
|
||||
if tools:
|
||||
inputs["tools"] = list(tools)
|
||||
model: Final = data.get("model")
|
||||
if model:
|
||||
inputs["model"] = model
|
||||
|
|
@ -464,8 +1000,33 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation):
|
|||
)
|
||||
|
||||
guardrailed_texts: Final = guardrailed_inputs.get("texts", [])
|
||||
if guardrailed_texts:
|
||||
_write_back_texts(guardrailed_texts, holders)
|
||||
guardrailed_messages: Final = guardrailed_inputs.get("structured_messages")
|
||||
sent_messages: Final = [structured_messages[index] for index in scoped_message_indices]
|
||||
rewritten_rows: Final = (
|
||||
guardrailed_messages
|
||||
if sent_messages and isinstance(guardrailed_messages, list) and guardrailed_messages != sent_messages
|
||||
else None
|
||||
)
|
||||
applied_texts: Final = (
|
||||
texts if rewritten_rows is not None and len(guardrailed_texts) != len(texts) else guardrailed_texts
|
||||
)
|
||||
edits: Final = (
|
||||
()
|
||||
if rewritten_rows is None
|
||||
else _structured_text_edits(
|
||||
partial(_leaf, body),
|
||||
converted,
|
||||
scoped_message_indices,
|
||||
rewritten_rows,
|
||||
_changed_leaves(holders, texts, applied_texts),
|
||||
partial(_texts_applied_view, body, applied_texts, skip_system, skip_tool, scan_only_tool_results),
|
||||
)
|
||||
)
|
||||
if edits is None:
|
||||
raise unappliable_request_rewrite(guardrail_to_apply.guardrail_name)
|
||||
if applied_texts:
|
||||
_write_back_texts(applied_texts, holders)
|
||||
_write_back_structured_texts(edits)
|
||||
|
||||
return data
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,451 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import re
|
||||
from collections.abc import Callable, Iterator, Mapping
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import yaml
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
from tests.integration._support.client import JSON_OBJECT, Gateway, gateway_from_environment, list_value, object_value
|
||||
from tests.integration._support.process import owned_proxy
|
||||
from tests.integration._support.wire import Reply, Request, Wire, wire_server
|
||||
|
||||
_MODEL_ID: Final = "anthropic.claude-sonnet-5-v1:0"
|
||||
_MODEL: Final = "converse-guarded"
|
||||
_CARD: Final = "4111 1111 1111 1111"
|
||||
_CARD_PATTERN: Final = re.compile(r"\b(?:\d{4} ){3}\d{4}\b")
|
||||
_MASK: Final = "****"
|
||||
_FLAGGED: Final = "flagged-by-moderation"
|
||||
_GUARD_PATH: Final = "/beta/litellm_basic_guardrail_api"
|
||||
_PNG: Final = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNkYAAAAAYAAjCB0C8AAAAASUVORK5CYII="
|
||||
_JSON: Final = TypeAdapter(JsonValue)
|
||||
_UNAPPLIABLE: Final = "rewrote the request in a way this endpoint cannot apply"
|
||||
_GENERIC_MODES: Final = (
|
||||
"masksm",
|
||||
"masksm_flat",
|
||||
"smonly",
|
||||
"sm_partial",
|
||||
"echo_sm",
|
||||
"mask",
|
||||
"block",
|
||||
"rowcount",
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Rig:
|
||||
gateway: Gateway
|
||||
vendor: Wire
|
||||
bedrock: Wire
|
||||
keys: Mapping[str, str]
|
||||
|
||||
|
||||
def _masked(value: JsonValue) -> JsonValue:
|
||||
return _JSON.validate_json(_CARD_PATTERN.sub(_MASK, json.dumps(value)))
|
||||
|
||||
|
||||
def _flattened(row: Mapping[str, JsonValue]) -> dict[str, JsonValue]:
|
||||
content: Final = row.get("content")
|
||||
if not isinstance(content, list):
|
||||
return dict(row)
|
||||
parts: Final = tuple(object_value(part).get("text") for part in content)
|
||||
return {**row, "content": "\n".join(part for part in parts if isinstance(part, str))}
|
||||
|
||||
|
||||
def _normalised_first_part(row: Mapping[str, JsonValue]) -> dict[str, JsonValue]:
|
||||
content: Final = row.get("content")
|
||||
if not isinstance(content, list) or not content:
|
||||
return dict(row)
|
||||
first: Final = object_value(content[0])
|
||||
text: Final = first.get("text")
|
||||
if not isinstance(text, str):
|
||||
return dict(row)
|
||||
return {**row, "content": [{**first, "text": text.strip() + " "}, *content[1:]]}
|
||||
|
||||
|
||||
def _generic_verdict(mode: str, body: Mapping[str, JsonValue]) -> dict[str, JsonValue]:
|
||||
texts: Final = list_value(body.get("texts") or [])
|
||||
rows: Final = tuple(object_value(row) for row in list_value(body.get("structured_messages") or []))
|
||||
masked_texts: Final = _masked(list(texts))
|
||||
masked_rows: Final = list_value(_masked([dict(row) for row in rows]))
|
||||
if _CARD_PATTERN.search(json.dumps(body)) is None:
|
||||
return {"action": "NONE"}
|
||||
verdicts: Final[Mapping[str, Callable[[], dict[str, JsonValue]]]] = {
|
||||
"masksm": lambda: {"action": "GUARDRAIL_INTERVENED", "texts": masked_texts, "structured_messages": masked_rows},
|
||||
"masksm_flat": lambda: {
|
||||
"action": "GUARDRAIL_INTERVENED",
|
||||
"texts": masked_texts,
|
||||
"structured_messages": [_flattened(object_value(row)) for row in masked_rows],
|
||||
},
|
||||
"smonly": lambda: {"action": "GUARDRAIL_INTERVENED", "structured_messages": masked_rows},
|
||||
"sm_partial": lambda: {
|
||||
"action": "GUARDRAIL_INTERVENED",
|
||||
"texts": masked_texts,
|
||||
"structured_messages": [_normalised_first_part(row) for row in rows],
|
||||
},
|
||||
"echo_sm": lambda: {
|
||||
"action": "GUARDRAIL_INTERVENED",
|
||||
"texts": masked_texts,
|
||||
"structured_messages": [dict(row) for row in rows],
|
||||
},
|
||||
"mask": lambda: {"action": "GUARDRAIL_INTERVENED", "texts": masked_texts},
|
||||
"block": lambda: {"action": "BLOCKED", "blocked_reason": "card number"},
|
||||
"rowcount": lambda: {
|
||||
"action": "GUARDRAIL_INTERVENED",
|
||||
"texts": masked_texts,
|
||||
"structured_messages": masked_rows[:-1],
|
||||
},
|
||||
}
|
||||
return verdicts[mode]()
|
||||
|
||||
|
||||
def _vendor(request: Request) -> Reply:
|
||||
body: Final = JSON_OBJECT.validate_json(request.body)
|
||||
if request.target == "/moderation/moderations":
|
||||
flagged: Final = _FLAGGED in str(body.get("input"))
|
||||
result: Final = {
|
||||
"flagged": flagged,
|
||||
"categories": {"violence": flagged},
|
||||
"category_scores": {"violence": 0.99 if flagged else 0.0},
|
||||
"category_applied_input_types": {"violence": ["text"]},
|
||||
}
|
||||
return Reply(
|
||||
body=json.dumps({"id": "modr-scripted", "model": "omni-moderation-latest", "results": [result]}).encode()
|
||||
)
|
||||
if request.target == "/promptguard/api/v1/guard":
|
||||
messages: Final = body.get("messages")
|
||||
if _CARD_PATTERN.search(json.dumps(messages)) is None:
|
||||
return Reply(body=json.dumps({"decision": "allow"}).encode())
|
||||
return Reply(body=json.dumps({"decision": "redact", "redacted_messages": _masked(messages)}).encode())
|
||||
mode: Final = request.target.removeprefix("/gga/").removesuffix(_GUARD_PATH)
|
||||
return Reply(body=json.dumps(_generic_verdict(mode, body)).encode())
|
||||
|
||||
|
||||
def _bedrock(request: Request) -> Reply:
|
||||
assert request.target == f"/model/{_MODEL_ID}/converse", request.target
|
||||
return Reply(
|
||||
body=json.dumps(
|
||||
{
|
||||
"output": {"message": {"role": "assistant", "content": [{"text": "guarded reply"}]}},
|
||||
"stopReason": "end_turn",
|
||||
"usage": {"inputTokens": 7, "outputTokens": 2, "totalTokens": 9},
|
||||
"metrics": {"latencyMs": 1},
|
||||
}
|
||||
).encode()
|
||||
)
|
||||
|
||||
|
||||
def _guardrail(name: str, params: Mapping[str, JsonValue]) -> dict[str, JsonValue]:
|
||||
return {"guardrail_name": name, "litellm_params": {"mode": "pre_call", "default_on": False, **params}}
|
||||
|
||||
|
||||
def _config(directory: Path, vendor: Wire, bedrock: Wire) -> Path:
|
||||
configuration: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
|
||||
configuration["model_list"] = [
|
||||
*configuration["model_list"],
|
||||
{
|
||||
"model_name": _MODEL,
|
||||
"litellm_params": {
|
||||
"model": f"bedrock/{_MODEL_ID}",
|
||||
"api_base": bedrock.url,
|
||||
"aws_access_key_id": "AKIASCRIPTEDPROVIDER",
|
||||
"aws_secret_access_key": "scripted-secret",
|
||||
"aws_region_name": "us-east-1",
|
||||
},
|
||||
},
|
||||
]
|
||||
generic: Final = {"guardrail": "generic_guardrail_api", "api_key": "synthetic-guardrail-key"}
|
||||
configuration["guardrails"] = [
|
||||
*(_guardrail(f"gga-{mode}", {**generic, "api_base": f"{vendor.url}/gga/{mode}"}) for mode in _GENERIC_MODES),
|
||||
_guardrail(
|
||||
"gga-block-skip-tool",
|
||||
{**generic, "api_base": f"{vendor.url}/gga/block", "skip_tool_message_in_guardrail": True},
|
||||
),
|
||||
_guardrail(
|
||||
"gga-mask-tool-only",
|
||||
{**generic, "api_base": f"{vendor.url}/gga/mask", "scan_only_tool_results": True},
|
||||
),
|
||||
_guardrail(
|
||||
"promptguard",
|
||||
{"guardrail": "promptguard", "api_base": f"{vendor.url}/promptguard", "api_key": "synthetic-pg-key"},
|
||||
),
|
||||
_guardrail(
|
||||
"moderation",
|
||||
{"guardrail": "openai_moderation", "api_base": f"{vendor.url}/moderation", "api_key": "synthetic-key"},
|
||||
),
|
||||
]
|
||||
path: Final = directory / "converse-guardrails.yaml"
|
||||
path.write_text(yaml.safe_dump(configuration))
|
||||
return path
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]:
|
||||
directory: Final = tmp_path_factory.mktemp("converse-guardrail-rewrites")
|
||||
with (
|
||||
gateway_from_environment() as gateway,
|
||||
wire_server(_vendor) as vendor,
|
||||
wire_server(_bedrock) as bedrock,
|
||||
owned_proxy(gateway, directory, {}, config=_config(directory, vendor, bedrock), workers=2) as owned,
|
||||
owned.scenario() as scenario,
|
||||
):
|
||||
names: Final = (
|
||||
*(f"gga-{mode}" for mode in _GENERIC_MODES),
|
||||
"gga-block-skip-tool",
|
||||
"gga-mask-tool-only",
|
||||
"promptguard",
|
||||
"moderation",
|
||||
)
|
||||
yield _Rig(owned, vendor, bedrock, {name: scenario.key(guardrails=[name]) for name in names})
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _drained(rig: _Rig) -> Iterator[None]:
|
||||
rig.vendor.drain()
|
||||
rig.bedrock.drain()
|
||||
yield
|
||||
rig.vendor.drain()
|
||||
rig.bedrock.drain()
|
||||
|
||||
|
||||
def _converse(rig: _Rig, guardrail: str, body: Mapping[str, JsonValue]) -> httpx.Response:
|
||||
return rig.gateway.request("POST", f"/bedrock/model/{_MODEL}/converse", body, key=rig.keys[guardrail])
|
||||
|
||||
|
||||
def _forwarded(rig: _Rig, response: httpx.Response) -> dict[str, JsonValue]:
|
||||
assert response.status_code == 200, response.text
|
||||
sent: Final = rig.bedrock.drain()
|
||||
assert len(sent) == 1, response.text
|
||||
return JSON_OBJECT.validate_json(sent[0].body)
|
||||
|
||||
|
||||
def _refused(rig: _Rig, response: httpx.Response) -> str:
|
||||
assert response.status_code >= 400, response.text
|
||||
assert rig.bedrock.drain() == (), response.text
|
||||
return response.text
|
||||
|
||||
|
||||
def _vendor_bodies(rig: _Rig) -> tuple[dict[str, JsonValue], ...]:
|
||||
return tuple(JSON_OBJECT.validate_json(request.body) for request in rig.vendor.drain())
|
||||
|
||||
|
||||
def _user(*blocks: JsonValue) -> dict[str, JsonValue]:
|
||||
return {"role": "user", "content": list(blocks)}
|
||||
|
||||
|
||||
def _assistant(*blocks: JsonValue) -> dict[str, JsonValue]:
|
||||
return {"role": "assistant", "content": list(blocks)}
|
||||
|
||||
|
||||
def _tool_use(card: str) -> dict[str, JsonValue]:
|
||||
return {"toolUse": {"toolUseId": "tooluse_lookup", "name": "lookup_card", "input": {"card": card, "limit": 2}}}
|
||||
|
||||
|
||||
def _tool_result(*content: JsonValue) -> dict[str, JsonValue]:
|
||||
return {"toolResult": {"toolUseId": "tooluse_lookup", "content": list(content), "status": "success"}}
|
||||
|
||||
|
||||
def test_multi_block_user_message_with_flattened_vendor_reply_reaches_bedrock_masked(rig: _Rig) -> None:
|
||||
body: Final = {"messages": [_user({"text": "Order details below."}, {"text": f"Card {_CARD} please."})]}
|
||||
sent: Final = _forwarded(rig, _converse(rig, "gga-masksm_flat", body))
|
||||
assert sent["messages"] == [_user({"text": "Order details below."}, {"text": f"Card {_MASK} please."})]
|
||||
|
||||
|
||||
def test_multi_block_system_prompt_with_flattened_vendor_reply_reaches_bedrock_masked(rig: _Rig) -> None:
|
||||
body: Final = {
|
||||
"system": [{"text": "You are a billing assistant."}, {"text": f"Card on file {_CARD}."}],
|
||||
"messages": [_user({"text": "What card is on file?"})],
|
||||
}
|
||||
sent: Final = _forwarded(rig, _converse(rig, "gga-masksm_flat", body))
|
||||
assert sent["system"] == [{"text": "You are a billing assistant."}, {"text": f"Card on file {_MASK}."}]
|
||||
assert sent["messages"] == body["messages"]
|
||||
|
||||
|
||||
def test_tool_result_with_text_and_json_parts_reaches_bedrock_masked(rig: _Rig) -> None:
|
||||
body: Final = {
|
||||
"messages": [
|
||||
_user({"text": "Look up my card."}),
|
||||
_assistant(_tool_use("last four 1111")),
|
||||
_user(_tool_result({"text": f"found {_CARD}"}, {"json": {"card": _CARD, "status": "active"}})),
|
||||
]
|
||||
}
|
||||
sent: Final = _forwarded(rig, _converse(rig, "gga-masksm", body))
|
||||
assert sent["messages"] == [
|
||||
_user({"text": "Look up my card."}),
|
||||
_assistant(_tool_use("last four 1111")),
|
||||
_user(_tool_result({"text": f"found {_MASK}"}, {"json": {"card": _MASK, "status": "active"}})),
|
||||
]
|
||||
|
||||
|
||||
def test_tool_result_with_text_and_image_parts_reaches_bedrock_masked_with_image_intact(rig: _Rig) -> None:
|
||||
image: Final = {"image": {"format": "png", "source": {"bytes": _PNG}}}
|
||||
body: Final = {
|
||||
"messages": [
|
||||
_user({"text": "Read the screenshot."}),
|
||||
_assistant(_tool_use("screenshot")),
|
||||
_user(_tool_result({"text": f"screenshot shows {_CARD}"}, image)),
|
||||
]
|
||||
}
|
||||
sent: Final = _forwarded(rig, _converse(rig, "gga-masksm", body))
|
||||
assert sent["messages"] == [
|
||||
_user({"text": "Read the screenshot."}),
|
||||
_assistant(_tool_use("screenshot")),
|
||||
_user(_tool_result({"text": f"screenshot shows {_MASK}"}, image)),
|
||||
]
|
||||
|
||||
|
||||
def test_flat_text_mask_wins_over_an_unmasked_structured_leaf(rig: _Rig) -> None:
|
||||
body: Final = {"messages": [_user({"text": f"Card {_CARD} on file."}, {"text": "Thanks."})]}
|
||||
sent: Final = _forwarded(rig, _converse(rig, "gga-sm_partial", body))
|
||||
assert sent["messages"] == [_user({"text": f"Card {_MASK} on file."}, {"text": "Thanks."})]
|
||||
|
||||
|
||||
def test_promptguard_redaction_keeps_every_turn_in_place(rig: _Rig) -> None:
|
||||
body: Final = {
|
||||
"system": [{"text": "You are a support bot."}],
|
||||
"messages": [
|
||||
_user({"text": "My name is Ann."}),
|
||||
_assistant({"text": "Hi Ann, how can I help?"}),
|
||||
_user({"text": f"Repeat exactly: card {_CARD} ok"}),
|
||||
],
|
||||
}
|
||||
sent: Final = _forwarded(rig, _converse(rig, "promptguard", body))
|
||||
assert sent["system"] == [{"text": "You are a support bot."}]
|
||||
assert sent["messages"] == [
|
||||
_user({"text": "My name is Ann."}),
|
||||
_assistant({"text": "Hi Ann, how can I help?"}),
|
||||
_user({"text": f"Repeat exactly: card {_MASK} ok"}),
|
||||
]
|
||||
|
||||
|
||||
def test_moderation_blocks_a_flagged_latest_user_turn_before_bedrock(rig: _Rig) -> None:
|
||||
body: Final = {
|
||||
"system": [{"text": "You are a support bot."}],
|
||||
"messages": [
|
||||
_user({"text": "Say hi."}),
|
||||
_assistant({"text": "Hi."}),
|
||||
_user({"text": f"please {_FLAGGED} now"}),
|
||||
],
|
||||
}
|
||||
refused: Final = _refused(rig, _converse(rig, "moderation", body))
|
||||
assert "Violated OpenAI moderation policy" in refused
|
||||
|
||||
|
||||
def test_skip_tool_messages_keeps_tool_call_arguments_away_from_the_guardrail(rig: _Rig) -> None:
|
||||
body: Final = {
|
||||
"messages": [
|
||||
_user({"text": "Check my card."}),
|
||||
_assistant(_tool_use(_CARD)),
|
||||
_user(_tool_result({"text": "card is active"})),
|
||||
]
|
||||
}
|
||||
sent: Final = _forwarded(rig, _converse(rig, "gga-block-skip-tool", body))
|
||||
assert sent["messages"] == body["messages"]
|
||||
assert all(_CARD not in json.dumps(scanned) for scanned in _vendor_bodies(rig))
|
||||
|
||||
|
||||
def test_tool_only_guardrail_receives_no_tool_definitions(rig: _Rig) -> None:
|
||||
body: Final = {
|
||||
"messages": [
|
||||
_user({"text": f"My card is {_CARD}."}),
|
||||
_assistant(_tool_use("lookup")),
|
||||
_user(_tool_result({"text": f"found {_CARD}"})),
|
||||
],
|
||||
"toolConfig": {
|
||||
"tools": [
|
||||
{
|
||||
"toolSpec": {
|
||||
"name": "lookup_card",
|
||||
"description": "Look up a card",
|
||||
"inputSchema": {"json": {"type": "object", "properties": {"card": {"type": "string"}}}},
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
}
|
||||
sent: Final = _forwarded(rig, _converse(rig, "gga-mask-tool-only", body))
|
||||
scanned: Final = _vendor_bodies(rig)
|
||||
assert len(scanned) == 1
|
||||
assert scanned[0].get("tools") is None, scanned[0]
|
||||
assert scanned[0]["texts"] == [f"found {_CARD}"]
|
||||
assert sent["messages"] == [
|
||||
_user({"text": f"My card is {_CARD}."}),
|
||||
_assistant(_tool_use("lookup")),
|
||||
_user(_tool_result({"text": f"found {_MASK}"})),
|
||||
]
|
||||
|
||||
|
||||
def test_structured_only_vendor_mask_reaches_bedrock(rig: _Rig) -> None:
|
||||
body: Final = {"messages": [_user({"text": f"Card {_CARD} please."})]}
|
||||
sent: Final = _forwarded(rig, _converse(rig, "gga-smonly", body))
|
||||
assert sent["messages"] == [_user({"text": f"Card {_MASK} please."})]
|
||||
|
||||
|
||||
def test_guard_content_turn_reaches_bedrock_masked(rig: _Rig) -> None:
|
||||
body: Final = {"messages": [_user({"text": "Check this."}, {"guardContent": {"text": {"text": f"Card {_CARD}."}}})]}
|
||||
sent: Final = _forwarded(rig, _converse(rig, "gga-masksm", body))
|
||||
assert sent["messages"] == [_user({"text": "Check this."}, {"guardContent": {"text": {"text": f"Card {_MASK}."}}})]
|
||||
|
||||
|
||||
def test_tool_call_arguments_in_history_reach_bedrock_masked(rig: _Rig) -> None:
|
||||
body: Final = {
|
||||
"messages": [
|
||||
_user({"text": "Check my card."}),
|
||||
_assistant(_tool_use(_CARD)),
|
||||
_user(_tool_result({"text": "card is active"})),
|
||||
]
|
||||
}
|
||||
sent: Final = _forwarded(rig, _converse(rig, "gga-masksm", body))
|
||||
assert sent["messages"] == [
|
||||
_user({"text": "Check my card."}),
|
||||
_assistant(_tool_use(_MASK)),
|
||||
_user(_tool_result({"text": "card is active"})),
|
||||
]
|
||||
|
||||
|
||||
def test_echoed_structured_rows_keep_the_text_mask(rig: _Rig) -> None:
|
||||
body: Final = {"messages": [_user({"text": f"Card {_CARD} please."})]}
|
||||
sent: Final = _forwarded(rig, _converse(rig, "gga-echo_sm", body))
|
||||
assert sent["messages"] == [_user({"text": f"Card {_MASK} please."})]
|
||||
|
||||
|
||||
def test_single_block_string_reply_reaches_bedrock_masked(rig: _Rig) -> None:
|
||||
body: Final = {"messages": [_user({"text": f"Card {_CARD} please."})]}
|
||||
sent: Final = _forwarded(rig, _converse(rig, "gga-masksm_flat", body))
|
||||
assert sent["messages"] == [_user({"text": f"Card {_MASK} please."})]
|
||||
|
||||
|
||||
def test_missing_structured_row_is_refused_before_bedrock(rig: _Rig) -> None:
|
||||
body: Final = {"messages": [_user({"text": "Hello."}), _assistant({"text": "Hi."}), _user({"text": f"{_CARD}"})]}
|
||||
assert _UNAPPLIABLE in _refused(rig, _converse(rig, "gga-rowcount", body))
|
||||
|
||||
|
||||
def test_blocked_verdict_never_reaches_bedrock(rig: _Rig) -> None:
|
||||
body: Final = {"messages": [_user({"text": f"Card {_CARD} please."})]}
|
||||
_refused(rig, _converse(rig, "gga-block", body))
|
||||
|
||||
|
||||
def test_concurrent_requests_keep_their_own_rewrites(rig: _Rig) -> None:
|
||||
cards: Final = tuple(f"4111 1111 1111 {index:04d}" for index in range(24))
|
||||
bodies: Final = tuple(
|
||||
{"messages": [_user({"text": f"request {index}"}, {"text": f"card {card}"})]}
|
||||
for index, card in enumerate(cards)
|
||||
)
|
||||
with ThreadPoolExecutor(max_workers=12) as pool:
|
||||
responses: Final = tuple(pool.map(lambda body: _converse(rig, "gga-masksm_flat", body), bodies))
|
||||
assert [response.status_code for response in responses] == [200] * len(bodies)
|
||||
sent: Final = sorted(
|
||||
(JSON_OBJECT.validate_json(request.body)["messages"] for request in rig.bedrock.drain()), key=json.dumps
|
||||
)
|
||||
expected: Final = sorted(
|
||||
([_user({"text": f"request {index}"}, {"text": f"card {_MASK}"})] for index in range(len(cards))),
|
||||
key=json.dumps,
|
||||
)
|
||||
assert sent == expected
|
||||
File diff suppressed because it is too large
Load diff
Loading…
Add table
Reference in a new issue