mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
fix(guardrails): apply guardrail masks to Bedrock InvokeModel Claude requests and replies (#49947)
This commit is contained in:
shrey
GitHub
parent
74783d38ea
commit
85fdea0fd1
4 changed files with 415 additions and 22 deletions
|
|
@ -781,12 +781,17 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
for event in AnthropicMessagesHandler._parse_sse_data_line(line)
|
||||
)
|
||||
|
||||
def _translate_to_openai(self, data: dict) -> ChatCompletionRequest:
|
||||
def _translate_to_openai(
|
||||
self,
|
||||
data: dict, # mutable-ok: the request body the adapter reads
|
||||
*,
|
||||
keep_cache_control: bool = False,
|
||||
) -> ChatCompletionRequest:
|
||||
"""Translate Anthropic request to OpenAI chat completion format."""
|
||||
(
|
||||
chat_completion_compatible_request,
|
||||
_tool_name_mapping,
|
||||
) = LiteLLMAnthropicMessagesAdapter().translate_anthropic_to_openai(
|
||||
) = LiteLLMAnthropicMessagesAdapter(keep_cache_control=keep_cache_control).translate_anthropic_to_openai(
|
||||
anthropic_message_request=cast(AnthropicMessagesRequest, scannable_tool_result_request(data)),
|
||||
preserve_midturn_system=True,
|
||||
)
|
||||
|
|
@ -813,9 +818,13 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
data: dict,
|
||||
guardrail_to_apply: "CustomGuardrail",
|
||||
litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
|
||||
*,
|
||||
request_data: dict[str, object] | None = None, # mutable-ok: the request dict guardrails receive
|
||||
keep_cache_control: bool = False,
|
||||
) -> Mapping[str, object]:
|
||||
"""
|
||||
Process input messages by applying guardrails to text content.
|
||||
Process input messages by applying guardrails to text content. ``request_data`` is what the guardrail sees as
|
||||
the request when ``data`` is only its body, e.g. a provider pass-through route that nests the body
|
||||
"""
|
||||
messages: Final = data.get("messages")
|
||||
if messages is None:
|
||||
|
|
@ -827,8 +836,17 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
|
||||
# The top-level prompt is translated on its own below so it can be hoisted in front of
|
||||
# any mid-turn system entries and scanned first, aligned with that structured position.
|
||||
translation_source: Final = {key: value for key, value in data.items() if key != "system"}
|
||||
chat_completion_compatible_request: Final = self._translate_to_openai(translation_source)
|
||||
body_model: Final = data.get("model")
|
||||
request_model: Final = (request_data or {}).get("model")
|
||||
model: Final[str] = (
|
||||
body_model
|
||||
if isinstance(body_model, str) and body_model
|
||||
else (request_model if isinstance(request_model, str) else "")
|
||||
)
|
||||
translation_source: Final = {**{key: value for key, value in data.items() if key != "system"}, "model": model}
|
||||
chat_completion_compatible_request: Final = self._translate_to_openai(
|
||||
translation_source, keep_cache_control=keep_cache_control
|
||||
)
|
||||
|
||||
full_structured_messages: Final = cast(
|
||||
list[AllMessageValues],
|
||||
|
|
@ -837,7 +855,9 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
has_midturn_system_message: Final = any(
|
||||
str(message.get("role") or "").lower() == "system" for message in full_structured_messages
|
||||
)
|
||||
hoisted_system_message: Final = None if skip_system else self._hoisted_top_level_system_message(data)
|
||||
hoisted_system_message: Final = (
|
||||
None if skip_system else self._hoisted_top_level_system_message(data, model, keep_cache_control)
|
||||
)
|
||||
if hoisted_system_message is not None:
|
||||
full_structured_messages.insert(0, hoisted_system_message)
|
||||
# skip_system already excluded the trusted top-level prompt (it is simply not hoisted);
|
||||
|
|
@ -899,7 +919,6 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
if structured_messages:
|
||||
inputs["structured_messages"] = structured_messages
|
||||
# Include model information if available
|
||||
model: Final = data.get("model")
|
||||
if model:
|
||||
inputs["model"] = model
|
||||
original_tool_names: Final = _original_tool_names(
|
||||
|
|
@ -907,7 +926,7 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
)
|
||||
guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail(
|
||||
inputs=inputs,
|
||||
request_data=data,
|
||||
request_data=data if request_data is None else request_data,
|
||||
input_type="request",
|
||||
logging_obj=litellm_logging_obj,
|
||||
)
|
||||
|
|
@ -955,6 +974,7 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
hoisted_system_message=hoisted_system_message,
|
||||
preserve_system_messages=has_midturn_system_message,
|
||||
original_tool_names=original_tool_names.used,
|
||||
model=model,
|
||||
)
|
||||
else:
|
||||
if guardrailed_texts and len(guardrailed_texts) != len(scanned):
|
||||
|
|
@ -977,17 +997,20 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
|
||||
return data
|
||||
|
||||
def _hoisted_top_level_system_message(self, data: Mapping[str, object]) -> AllMessageValues | None:
|
||||
def _hoisted_top_level_system_message(
|
||||
self, data: Mapping[str, object], model: str, keep_cache_control: bool
|
||||
) -> AllMessageValues | None:
|
||||
"""Return the system message produced by translating the top-level prompt."""
|
||||
system: Final = data.get("system")
|
||||
if not system:
|
||||
return None
|
||||
probe: Final = self._translate_to_openai(
|
||||
{
|
||||
"model": data.get("model") or "",
|
||||
"model": model,
|
||||
"messages": [],
|
||||
"system": system,
|
||||
}
|
||||
},
|
||||
keep_cache_control=keep_cache_control,
|
||||
)
|
||||
hoisted: Final = probe.get("messages") or []
|
||||
return hoisted[0] if hoisted else None
|
||||
|
|
@ -1121,6 +1144,7 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
preserve_system_messages: bool = False,
|
||||
*,
|
||||
original_tool_names: Mapping[str, str],
|
||||
model: str,
|
||||
) -> None:
|
||||
"""Write a guardrail's structured-message rewrite back without losing corrections."""
|
||||
from litellm.litellm_core_utils.prompt_templates.factory import (
|
||||
|
|
@ -1129,7 +1153,6 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
)
|
||||
|
||||
_is_system: Final = AnthropicMessagesHandler._is_system
|
||||
model: Final = str(data.get("model") or "")
|
||||
converted: Final[list] = [] # mutable-ok: API message payload
|
||||
|
||||
def _convert_run(run: list) -> None: # mutable-ok: API message payload
|
||||
|
|
|
|||
|
|
@ -380,8 +380,8 @@ _CLAUDE_OVER_REASONING_EFFORT_PROVIDERS: Final = frozenset({LlmProviders.KYMA.va
|
|||
|
||||
|
||||
class LiteLLMAnthropicMessagesAdapter:
|
||||
def __init__(self):
|
||||
pass
|
||||
def __init__(self, *, keep_cache_control: bool = False) -> None:
|
||||
self.keep_cache_control: Final = keep_cache_control
|
||||
|
||||
### FOR [BETA] `/v1/messages` endpoint support
|
||||
|
||||
|
|
@ -437,7 +437,7 @@ class LiteLLMAnthropicMessagesAdapter:
|
|||
cache_control: Final = (
|
||||
source.get("cache_control") if isinstance(source, dict) else getattr(source, "cache_control", None)
|
||||
)
|
||||
if cache_control and model and self.target_consumes_cache_control(model):
|
||||
if cache_control and (self.keep_cache_control or (model and self.target_consumes_cache_control(model))):
|
||||
# TypedDict objects support dict operations at runtime
|
||||
# Use type ignore consistent with codebase pattern (see anthropic/chat/transformation.py:432)
|
||||
if isinstance(target, dict):
|
||||
|
|
|
|||
|
|
@ -14,6 +14,7 @@ from pydantic import InstanceOf, TypeAdapter, ValidationError
|
|||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.litellm_core_utils.core_helpers import as_str_mapping
|
||||
from litellm.llms.anthropic.chat.guardrail_translation.handler import AnthropicMessagesHandler
|
||||
from litellm.llms.base_llm.guardrail_translation.base_translation import (
|
||||
BaseTranslation,
|
||||
|
|
@ -60,6 +61,8 @@ if TYPE_CHECKING:
|
|||
|
||||
_CONVERSE_ACTIONS: Final = frozenset({"converse", "converse-stream"})
|
||||
_INVOKE_STREAM_ACTION: Final = "invoke-with-response-stream"
|
||||
_INVOKE_ACTIONS: Final = frozenset({"invoke", _INVOKE_STREAM_ACTION})
|
||||
_BLOCK_TYPES_THE_ANTHROPIC_HANDLER_DOES_NOT_SCAN: Final = frozenset({"document", "search_result"})
|
||||
_EVENT_STREAM_CONTENT_TYPE: Final = "vnd.amazon.eventstream"
|
||||
_EVENT_STREAM_MEDIA_TYPE: Final = "application/vnd.amazon.eventstream"
|
||||
|
||||
|
|
@ -73,6 +76,83 @@ def _is_invoke_stream_endpoint(endpoint: str) -> bool:
|
|||
return endpoint.rstrip("/").split("/")[-1] == _INVOKE_STREAM_ACTION
|
||||
|
||||
|
||||
def _anthropic_invoke_body(
|
||||
data: Mapping[str, object],
|
||||
) -> dict[str, object] | None: # mutable-ok: the body guardrails write back into
|
||||
"""The body of an InvokeModel call in Anthropic Messages format, which Bedrock requires to carry anthropic_version"""
|
||||
endpoint: Final = data.get("endpoint")
|
||||
body: Final = data.get("data")
|
||||
if not isinstance(endpoint, str) or endpoint.rstrip("/").split("/")[-1] not in _INVOKE_ACTIONS:
|
||||
return None
|
||||
if not isinstance(body, dict):
|
||||
return None
|
||||
invoke_body: Final = _BODY_BY_REFERENCE.validate_python(body)
|
||||
if not isinstance(invoke_body.get("anthropic_version"), str) or not isinstance(invoke_body.get("messages"), list):
|
||||
return None
|
||||
messages: Final = _ARRAY_BY_REFERENCE.validate_python(invoke_body["messages"])
|
||||
if not all(_is_well_formed_message(message) for message in messages):
|
||||
return None
|
||||
blocks: Final = _anthropic_message_blocks(messages)
|
||||
if any(block.get("type") in _BLOCK_TYPES_THE_ANTHROPIC_HANDLER_DOES_NOT_SCAN for block in blocks):
|
||||
return None
|
||||
return invoke_body
|
||||
|
||||
|
||||
def _body_level_request_data(
|
||||
passthrough_request: Mapping[str, object],
|
||||
) -> dict[str, object]: # mutable-ok: the request dict guardrails receive
|
||||
"""What `/v1/messages` hands a guardrail as `request_data`: the Anthropic body at the top level, with the proxy's own keys (model, metadata, headers) on top"""
|
||||
body: Final = as_str_mapping(passthrough_request.get("data")) or {}
|
||||
proxy_keys: Final = {key: value for key, value in passthrough_request.items() if key != "data"}
|
||||
return {**body, **proxy_keys}
|
||||
|
||||
|
||||
def _is_well_formed_message(message: object) -> bool:
|
||||
"""Whether the Messages translation can read `message`, so a malformed body still reaches Bedrock to be rejected there"""
|
||||
fields: Final = as_str_mapping(message)
|
||||
if fields is None or not isinstance(fields.get("role"), str):
|
||||
return False
|
||||
content: Final = fields.get("content")
|
||||
if isinstance(content, str):
|
||||
return True
|
||||
return isinstance(content, list) and all(
|
||||
as_str_mapping(block) is not None for block in _ARRAY_BY_REFERENCE.validate_python(content)
|
||||
)
|
||||
|
||||
|
||||
def _anthropic_message_blocks(messages: Sequence[object]) -> tuple[Mapping[str, object], ...]:
|
||||
fields: Final = (as_str_mapping(message) for message in messages)
|
||||
return tuple(chain.from_iterable(_anthropic_content_blocks(field.get("content")) for field in fields if field))
|
||||
|
||||
|
||||
def _anthropic_content_blocks(content: object) -> tuple[Mapping[str, object], ...]:
|
||||
blocks: Final[list[Mapping[str, object]]] = [] # mutable-ok: collected while walking the content
|
||||
pending: Final[list[object]] = [content] # mutable-ok: content lists still to walk, nested tool results included
|
||||
while pending:
|
||||
current = pending.pop()
|
||||
if not isinstance(current, list):
|
||||
continue
|
||||
for item in _ARRAY_BY_REFERENCE.validate_python(current):
|
||||
block = as_str_mapping(item)
|
||||
if block is None:
|
||||
continue
|
||||
blocks.append(block)
|
||||
if block.get("type") == "tool_result":
|
||||
pending.append(block.get("content"))
|
||||
return tuple(blocks)
|
||||
|
||||
|
||||
def _anthropic_invoke_reply(
|
||||
response: object,
|
||||
request_data: object,
|
||||
) -> dict[str, object] | None: # mutable-ok: the reply guardrails write back into
|
||||
if not isinstance(response, dict) or not isinstance(request_data, dict):
|
||||
return None
|
||||
if _anthropic_invoke_body(_OBJECT_BY_REFERENCE.validate_python(request_data)) is None:
|
||||
return None
|
||||
return _BODY_BY_REFERENCE.validate_python(response)
|
||||
|
||||
|
||||
def _generic_passthrough_handler() -> "PassThroughEndpointHandler":
|
||||
"""
|
||||
Fallback for non-Converse Bedrock routes (e.g. invoke). The generic
|
||||
|
|
@ -1285,6 +1365,17 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation):
|
|||
body: Final = data.get("data")
|
||||
|
||||
if not _is_converse_endpoint(endpoint):
|
||||
passthrough_request: Final = _BODY_BY_REFERENCE.validate_python(data)
|
||||
anthropic_body: Final = _anthropic_invoke_body(passthrough_request)
|
||||
if anthropic_body is not None:
|
||||
await AnthropicMessagesHandler().process_input_messages( # pyright: ignore[reportUnknownMemberType] # takes the untyped request dict
|
||||
data=anthropic_body,
|
||||
guardrail_to_apply=guardrail_to_apply,
|
||||
litellm_logging_obj=litellm_logging_obj,
|
||||
request_data=_body_level_request_data(passthrough_request),
|
||||
keep_cache_control=True,
|
||||
)
|
||||
return passthrough_request
|
||||
return await _generic_passthrough_handler().process_input_messages(
|
||||
data=data,
|
||||
guardrail_to_apply=guardrail_to_apply,
|
||||
|
|
@ -1410,6 +1501,15 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation):
|
|||
and not _is_converse_endpoint(endpoint)
|
||||
and not _is_invoke_stream_endpoint(endpoint)
|
||||
):
|
||||
anthropic_reply: Final = _anthropic_invoke_reply(response, request_data) # pyright: ignore[reportUnknownArgumentType] # the base signature types request_data as a bare dict
|
||||
if anthropic_reply is not None:
|
||||
return await AnthropicMessagesHandler().process_output_response( # pyright: ignore[reportUnknownMemberType] # takes the untyped request dict
|
||||
response=anthropic_reply, # pyright: ignore[reportArgumentType] # the Anthropic reply dict, written back in place
|
||||
guardrail_to_apply=guardrail_to_apply,
|
||||
litellm_logging_obj=litellm_logging_obj,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
request_data=_body_level_request_data(_OBJECT_BY_REFERENCE.validate_python(request_data or {})),
|
||||
)
|
||||
return await _generic_passthrough_handler().process_output_response(
|
||||
response=response,
|
||||
guardrail_to_apply=guardrail_to_apply,
|
||||
|
|
|
|||
|
|
@ -1394,6 +1394,7 @@ class _ScriptedGuardrail(CustomGuardrail):
|
|||
self.skip_system_message_in_guardrail = skip_system
|
||||
self.scan_only_tool_results = scan_only_tool_results
|
||||
self.received: tuple[GenericGuardrailAPIInputs, ...] = ()
|
||||
self.request_data: tuple[Mapping[str, object], ...] = ()
|
||||
|
||||
async def apply_guardrail(
|
||||
self,
|
||||
|
|
@ -1403,6 +1404,7 @@ class _ScriptedGuardrail(CustomGuardrail):
|
|||
logging_obj: LiteLLMLoggingObj | None = None,
|
||||
) -> Mapping[str, object]:
|
||||
self.received = (*self.received, copy.deepcopy(inputs))
|
||||
self.request_data = (*self.request_data, request_data)
|
||||
return self.reply(inputs)
|
||||
|
||||
|
||||
|
|
@ -1429,16 +1431,30 @@ def _converse_request(body: dict) -> dict:
|
|||
|
||||
|
||||
async def _sent_to_bedrock(data: dict) -> dict:
|
||||
"""The Converse body Bedrock receives when the passthrough route sends ``data``."""
|
||||
"""The body Bedrock receives when the passthrough route sends ``data``."""
|
||||
received: list[httpx.Request] = []
|
||||
|
||||
action: Final = data["endpoint"].rsplit("/", 1)[-1]
|
||||
|
||||
def respond(request: httpx.Request) -> httpx.Response:
|
||||
received.append(request)
|
||||
reply = {
|
||||
"output": {"message": {"role": "assistant", "content": [{"text": "ok"}]}},
|
||||
"stopReason": "end_turn",
|
||||
"usage": {"inputTokens": 1, "outputTokens": 1, "totalTokens": 2},
|
||||
}
|
||||
reply: Final = (
|
||||
{
|
||||
"id": "msg_bdrk_1",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "claude-haiku-4-5",
|
||||
"content": [{"type": "text", "text": "ok"}],
|
||||
"stop_reason": "end_turn",
|
||||
"usage": {"input_tokens": 1, "output_tokens": 1},
|
||||
}
|
||||
if action == "invoke"
|
||||
else {
|
||||
"output": {"message": {"role": "assistant", "content": [{"text": "ok"}]}},
|
||||
"stopReason": "end_turn",
|
||||
"usage": {"inputTokens": 1, "outputTokens": 1, "totalTokens": 2},
|
||||
}
|
||||
)
|
||||
return httpx.Response(200, json=reply, request=request)
|
||||
|
||||
router = litellm.Router(
|
||||
|
|
@ -1457,7 +1473,7 @@ async def _sent_to_bedrock(data: dict) -> dict:
|
|||
await router.allm_passthrough_route(
|
||||
model="bedrock-model",
|
||||
method="POST",
|
||||
endpoint="/model/bedrock-model/converse",
|
||||
endpoint=f"/model/bedrock-model/{action}",
|
||||
data=data["data"],
|
||||
client=AsyncHTTPHandler(transport=httpx.MockTransport(respond)),
|
||||
)
|
||||
|
|
@ -3843,6 +3859,260 @@ class TestInvokeStreamGuardrailRelay:
|
|||
assert rewritten_key == received_key
|
||||
|
||||
|
||||
def _anthropic_invoke_data(
|
||||
action: str = "invoke",
|
||||
) -> dict[str, object]: # mutable-ok: tests write into the request the handler gets
|
||||
return {
|
||||
"endpoint": f"model/anthropic.claude-sonnet/{action}",
|
||||
"custom_llm_provider": "bedrock",
|
||||
"model": "anthropic.claude-sonnet",
|
||||
"litellm_call_id": "call-1",
|
||||
"data": {
|
||||
"anthropic_version": "bedrock-2023-05-31",
|
||||
"max_tokens": 512,
|
||||
"system": "You are a travel assistant. Card on file 4111 1111 1111 1111",
|
||||
"tools": [
|
||||
{
|
||||
"name": "find_trips",
|
||||
"description": "Search trips",
|
||||
"input_schema": {"type": "object", "properties": {"port": {"type": "string"}}},
|
||||
}
|
||||
],
|
||||
"messages": [
|
||||
{"role": "user", "content": "Pay with 4111 1111 1111 1111"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "tool_use", "id": "toolu_1", "name": "find_trips", "input": {"port": "Miami"}}
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "tool_result", "tool_use_id": "toolu_1", "content": "Card 4111 1111 1111 1111 charged"}
|
||||
],
|
||||
},
|
||||
],
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _masked_anthropic_invoke_body() -> Mapping[str, object]:
|
||||
body: Final = _anthropic_invoke_data()["data"]
|
||||
return {
|
||||
**body,
|
||||
"system": "You are a travel assistant. Card on file [CARD] 1111 1111 1111",
|
||||
"messages": [
|
||||
{"role": "user", "content": "Pay with [CARD] 1111 1111 1111"},
|
||||
body["messages"][1],
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "tool_result", "tool_use_id": "toolu_1", "content": "Card [CARD] 1111 1111 1111 charged"}
|
||||
],
|
||||
},
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
def _rows_masked(inputs: GenericGuardrailAPIInputs) -> Mapping[str, object]:
|
||||
rows = [
|
||||
{**row, "content": row["content"].replace("4111", "[CARD]")} if isinstance(row.get("content"), str) else row
|
||||
for row in inputs["structured_messages"]
|
||||
]
|
||||
return {"structured_messages": rows}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("action", ["invoke", "invoke-with-response-stream"])
|
||||
async def test_anthropic_invoke_masks_are_written_back_into_the_body(action: str) -> None:
|
||||
guardrail = _ScriptedGuardrail(lambda inputs: {"texts": _masked_texts(inputs)})
|
||||
|
||||
result = await BedrockPassthroughGuardrailHandler().process_input_messages(
|
||||
_anthropic_invoke_data(action), guardrail
|
||||
)
|
||||
|
||||
assert result["data"] == _masked_anthropic_invoke_body()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_anthropic_invoke_masks_reach_bedrock() -> None:
|
||||
data: Final = _anthropic_invoke_data()
|
||||
|
||||
await BedrockPassthroughGuardrailHandler().process_input_messages(
|
||||
data, _ScriptedGuardrail(lambda inputs: {"texts": _masked_texts(inputs)})
|
||||
)
|
||||
|
||||
assert await _sent_to_bedrock(data) == _masked_anthropic_invoke_body()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_anthropic_invoke_guardrail_gets_roles_tools_and_the_body_at_the_top_level_of_request_data() -> None:
|
||||
data: Final = _anthropic_invoke_data()
|
||||
guardrail = _replying({})
|
||||
|
||||
await BedrockPassthroughGuardrailHandler().process_input_messages(data, guardrail)
|
||||
|
||||
inputs = guardrail.received[-1]
|
||||
assert [row["role"] for row in inputs["structured_messages"]] == ["system", "user", "assistant", "tool"]
|
||||
assert [tool["function"]["name"] for tool in inputs["tools"]] == ["find_trips"]
|
||||
assert inputs["model"] == data["model"]
|
||||
request_data: Final = guardrail.request_data[-1]
|
||||
assert request_data["messages"] is data["data"]["messages"]
|
||||
assert request_data["endpoint"] == data["endpoint"]
|
||||
assert request_data["litellm_call_id"] == data["litellm_call_id"]
|
||||
assert "data" not in request_data
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_anthropic_invoke_structured_message_rewrites_are_written_back() -> None:
|
||||
result = await BedrockPassthroughGuardrailHandler().process_input_messages(
|
||||
_anthropic_invoke_data(), _ScriptedGuardrail(_rows_masked)
|
||||
)
|
||||
|
||||
assert result["data"]["messages"][0] == {
|
||||
"role": "user",
|
||||
"content": [{"type": "text", "text": "Pay with [CARD] 1111 1111 1111"}],
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model", ["anthropic.claude-sonnet", "prod-model", ""])
|
||||
async def test_anthropic_invoke_structured_rewrite_keeps_cache_control_whatever_the_model_is_called(
|
||||
model: str,
|
||||
) -> None:
|
||||
cached: Final = {"type": "ephemeral"}
|
||||
data: Final = _anthropic_invoke_data()
|
||||
data["model"] = model
|
||||
data["data"]["system"] = [{"type": "text", "text": "Card on file 4111 1111 1111 1111", "cache_control": cached}]
|
||||
data["data"]["tools"] = [
|
||||
{
|
||||
"name": "find_trips",
|
||||
"description": "Search trips",
|
||||
"input_schema": {"type": "object", "properties": {}},
|
||||
"cache_control": cached,
|
||||
}
|
||||
]
|
||||
data["data"]["messages"] = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Pay with 4111 1111 1111 1111", "cache_control": {**cached, "ttl": "1h"}}
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
result = await BedrockPassthroughGuardrailHandler().process_input_messages(
|
||||
data,
|
||||
_ScriptedGuardrail(
|
||||
lambda inputs: {
|
||||
"structured_messages": json.loads(json.dumps(inputs["structured_messages"]).replace("4111", "[CARD]")),
|
||||
"tools": json.loads(json.dumps(inputs["tools"])),
|
||||
}
|
||||
),
|
||||
)
|
||||
|
||||
assert result["data"]["system"] == [
|
||||
{"type": "text", "text": "Card on file [CARD] 1111 1111 1111", "cache_control": cached}
|
||||
]
|
||||
assert result["data"]["messages"][0]["content"] == [
|
||||
{"type": "text", "text": "Pay with [CARD] 1111 1111 1111", "cache_control": {**cached, "ttl": "1h"}}
|
||||
]
|
||||
assert result["data"]["tools"][0]["cache_control"] == cached
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_invoke_body_without_anthropic_version_is_scanned_whole() -> None:
|
||||
data: Final = {**_anthropic_invoke_data(), "data": {"messages": [{"role": "user", "content": "hello"}]}}
|
||||
guardrail = _replying({})
|
||||
|
||||
await BedrockPassthroughGuardrailHandler().process_input_messages(data, guardrail)
|
||||
|
||||
assert json.loads(guardrail.received[-1]["texts"][0])["data"] == data["data"]
|
||||
assert "structured_messages" not in guardrail.received[-1]
|
||||
|
||||
|
||||
_PLAIN_TEXT_DOCUMENT: Final = {
|
||||
"type": "document",
|
||||
"source": {"type": "text", "media_type": "text/plain", "data": "Card 4111 1111 1111 1111"},
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"content",
|
||||
[
|
||||
pytest.param([{"type": "text", "text": "Summarize this"}, _PLAIN_TEXT_DOCUMENT], id="document"),
|
||||
pytest.param(
|
||||
[
|
||||
{
|
||||
"type": "search_result",
|
||||
"source": "https://example.com/cards",
|
||||
"title": "Cards",
|
||||
"content": [{"type": "text", "text": "Card 4111 1111 1111 1111"}],
|
||||
}
|
||||
],
|
||||
id="search_result",
|
||||
),
|
||||
pytest.param(
|
||||
[{"type": "tool_result", "tool_use_id": "toolu_1", "content": [_PLAIN_TEXT_DOCUMENT]}],
|
||||
id="document_in_tool_result",
|
||||
),
|
||||
],
|
||||
)
|
||||
async def test_anthropic_invoke_body_with_blocks_the_messages_handler_does_not_scan_is_scanned_whole(content: Sequence[object]) -> None:
|
||||
invoke: Final = _anthropic_invoke_data()
|
||||
data: Final = {**invoke, "data": {**invoke["data"], "messages": [{"role": "user", "content": content}]}}
|
||||
guardrail = _replying({})
|
||||
|
||||
await BedrockPassthroughGuardrailHandler().process_input_messages(data, guardrail)
|
||||
|
||||
assert json.loads(guardrail.received[-1]["texts"][0])["data"] == data["data"]
|
||||
assert "structured_messages" not in guardrail.received[-1]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"messages",
|
||||
[
|
||||
pytest.param([{"content": "hi"}], id="message_without_role"),
|
||||
pytest.param([{"role": "user", "content": ["hi"]}], id="bare_string_in_content"),
|
||||
pytest.param([{"role": "user", "content": None}], id="null_content"),
|
||||
pytest.param(["hi"], id="message_is_a_string"),
|
||||
],
|
||||
)
|
||||
async def test_anthropic_invoke_malformed_messages_are_scanned_whole_and_left_for_bedrock_to_reject(
|
||||
messages: Sequence[object],
|
||||
) -> None:
|
||||
invoke: Final = _anthropic_invoke_data()
|
||||
data: Final = {**invoke, "data": {**invoke["data"], "messages": messages}}
|
||||
guardrail = _replying({})
|
||||
|
||||
result: Final = await BedrockPassthroughGuardrailHandler().process_input_messages(data, guardrail)
|
||||
|
||||
assert json.loads(guardrail.received[-1]["texts"][0])["data"] == data["data"]
|
||||
assert result["data"]["messages"] == messages
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_anthropic_invoke_response_masks_are_written_back() -> None:
|
||||
response = {
|
||||
"id": "msg_bdrk_1",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"content": [{"type": "text", "text": "Charged 4111 1111 1111 1111"}],
|
||||
"stop_reason": "end_turn",
|
||||
}
|
||||
|
||||
result: Final = await BedrockPassthroughGuardrailHandler().process_output_response(
|
||||
response=response,
|
||||
guardrail_to_apply=_ScriptedGuardrail(lambda inputs: {"texts": _masked_texts(inputs)}),
|
||||
request_data=_anthropic_invoke_data(),
|
||||
)
|
||||
|
||||
assert result == {**response, "content": [{"type": "text", "text": "Charged [CARD] 1111 1111 1111"}]}
|
||||
|
||||
|
||||
def test_sse_events_keeps_a_data_line_whose_json_holds_unicode_line_separators() -> None:
|
||||
event: Final = {
|
||||
"type": "content_block_delta",
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue