diff --git a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py index 1acad8db8e8..48925c06408 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py @@ -42,7 +42,7 @@ from litellm.types.proxy.guardrails.guardrail_hooks.lakera_ai_v2 import ( LakeraAIRequest, LakeraAIResponse, ) -from litellm.types.utils import CallTypesLiteral, GuardrailStatus, ModelResponse +from litellm.types.utils import CallTypes, CallTypesLiteral, GuardrailStatus, ModelResponse _DECODED_JSON: Final = TypeAdapter(object) _RESPONSE_FIELDS: Final = TypeAdapter(Mapping[str, object]) @@ -203,6 +203,7 @@ def _unmaskable_reason( guardrail: "LakeraAIGuardrail", data: dict[str, object], lakera_response: LakeraAIResponse | None, + call_type: str, ) -> str | None: """Why a PII-only violation on ``data`` can't be masked in place, or None when it can.""" if has_non_string_content(data): @@ -211,7 +212,7 @@ def _unmaskable_reason( return "messages and input are both present, so the write-back is positionally ambiguous" if "messages" in data and not isinstance(data.get("messages"), list): return "a messages key that isn't a list, so there's nothing to merge the redacted content into" - if _bedrock_passthrough_body(data) is not None: + if _bedrock_passthrough_body(data, call_type) is not None: return "a /bedrock pass-through body, which has no write-back path for the redacted content" if not _has_responses_instructions(guardrail, data): return "no write-back path for the redacted content" @@ -220,9 +221,14 @@ def _unmaskable_reason( return None -def _bedrock_passthrough_body(data: Mapping[str, object]) -> Mapping[str, object] | None: - """The client's request body on a /bedrock pass-through call, which the route nests under ``data``.""" - if data.get("custom_llm_provider") != "bedrock" or not isinstance(data.get("endpoint"), str): +def _bedrock_passthrough_body(data: Mapping[str, object], call_type: str | None) -> Mapping[str, object] | None: + """The client's request body on a /bedrock pass-through call (call type ``allm_passthrough_route``), which the + route nests under ``data``.""" + if ( + call_type != CallTypes.allm_passthrough_route.value + or data.get("custom_llm_provider") != "bedrock" + or not isinstance(data.get("endpoint"), str) + ): return None return as_str_mapping(data.get("data")) @@ -268,11 +274,11 @@ def _bedrock_turns(message: object) -> Sequence[Mapping[str, object]]: ] -def _bedrock_inspection_messages(data: Mapping[str, object]) -> Sequence[Mapping[str, str]]: +def _bedrock_inspection_messages(data: Mapping[str, object], call_type: str | None) -> Sequence[Mapping[str, str]]: """Chat-style messages for a /bedrock invoke or Converse body: ``system`` first as a system message, then each ``messages`` turn with its text parts joined and its tool results as tool messages. Turns with no text are dropped. """ - body: Final = _bedrock_passthrough_body(data) + body: Final = _bedrock_passthrough_body(data, call_type) if body is None: return [] system: Final = body.get("system") @@ -283,7 +289,9 @@ def _bedrock_inspection_messages(data: Mapping[str, object]) -> Sequence[Mapping ) -def _build_lakera_inspection_messages(data: Mapping[str, object]) -> Sequence[Mapping[str, str]]: +def _build_lakera_inspection_messages( + data: Mapping[str, object], *, call_type: str | None = None +) -> Sequence[Mapping[str, str]]: """Like build_inspection_messages, but also covers the Responses-API ``instructions`` field, placed first since litellm later converts it into the model's leading system message and a prompt-injection detector @@ -304,7 +312,7 @@ def _build_lakera_inspection_messages(data: Mapping[str, object]) -> Sequence[Ma ) return [ *leading, - *_bedrock_inspection_messages(data), + *_bedrock_inspection_messages(data, call_type), *build_inspection_messages(dict(data)), ] @@ -580,6 +588,7 @@ class LakeraAIGuardrail(CustomGuardrail): inspected_messages: Sequence[AllMessageValues], lakera_response: LakeraAIResponse | None, masked_entity_count: dict[str, int], + call_type: str, ) -> bool: """Mask a body whose only obstacle to mask-in-place is the Responses-API ``instructions`` field, writing the redacted instructions straight into @@ -587,7 +596,7 @@ class LakeraAIGuardrail(CustomGuardrail): that field and would fold the instructions text into ``data["input"]``. Returns False without masking anything when _unmaskable_reason names an obstacle this can't get around.""" - if _unmaskable_reason(self, data, lakera_response) is not None: + if _unmaskable_reason(self, data, lakera_response, call_type) is not None: return False redacted: Final = self._mask_pii_in_messages( messages=inspected_messages, @@ -619,7 +628,7 @@ class LakeraAIGuardrail(CustomGuardrail): return data # Covers multimodal list content + Responses-API input/instructions. - inspection_messages: Final = _build_lakera_inspection_messages(data) + inspection_messages: Final = _build_lakera_inspection_messages(data, call_type=call_type) if not inspection_messages: verbose_proxy_logger.warning("Lakera AI: not running guardrail. No inspectable text in data") return data @@ -646,7 +655,7 @@ class LakeraAIGuardrail(CustomGuardrail): # handled safely by _apply_redacted_messages_back_preserving_fields's # scope-index merge, which never touches a message outside the scope # it actually redacted instead of reconstructing the list from scratch. - is_bedrock_body: Final = _bedrock_passthrough_body(data) is not None # pyright: ignore[reportUnknownArgumentType] # hook data is an untyped dict + is_bedrock_body: Final = _bedrock_passthrough_body(data, call_type) is not None # pyright: ignore[reportUnknownArgumentType] # hook data is an untyped dict is_multimodal_input: Final = ( has_non_string_content(data) or _has_combined_messages_and_input(data) @@ -720,6 +729,7 @@ class LakeraAIGuardrail(CustomGuardrail): inspected_messages=new_messages, lakera_response=lakera_guardrail_response, masked_entity_count=masked_entity_count, + call_type=call_type, ) if masked_in_place: verbose_proxy_logger.warning( @@ -729,7 +739,7 @@ class LakeraAIGuardrail(CustomGuardrail): verbose_proxy_logger.error( "Lakera Guardrail: Monitoring mode - PII detected but NOT masked, forwarding unredacted " "content to the model (reason: %s)", - _unmaskable_reason(self, data, lakera_guardrail_response), + _unmaskable_reason(self, data, lakera_guardrail_response, call_type), ) else: verbose_proxy_logger.warning( @@ -760,7 +770,7 @@ class LakeraAIGuardrail(CustomGuardrail): return # Covers multimodal list content + Responses-API input/instructions. - inspection_messages: Final = _build_lakera_inspection_messages(data) + inspection_messages: Final = _build_lakera_inspection_messages(data, call_type=call_type) if not inspection_messages: verbose_proxy_logger.warning("Lakera AI: not running guardrail. No inspectable text in data") return diff --git a/tests/integration/authorization/test_lakera_v2_bedrock_passthrough.py b/tests/integration/authorization/test_lakera_v2_bedrock_passthrough.py index 2eeb717b5fd..832bc7054f1 100644 --- a/tests/integration/authorization/test_lakera_v2_bedrock_passthrough.py +++ b/tests/integration/authorization/test_lakera_v2_bedrock_passthrough.py @@ -83,6 +83,7 @@ def _bedrock_reply(request: Request) -> Reply: class BedrockLakeraRig: owned: OwnedProxy lakera: Wire + lakera_down: Wire upstream: Wire @@ -91,6 +92,7 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I directory: Final = tmp_path_factory.mktemp("lakera-bedrock-passthrough") with ExitStack() as stack: lakera: Final = stack.enter_context(wire_server(_lakera_verdict)) + lakera_down: Final = stack.enter_context(wire_server(lambda _: Reply(status=500, body=b'{"error":"down"}'))) upstream: Final = stack.enter_context(wire_server(_bedrock_reply)) config: Final = _JSON.validate_python(yaml.safe_load(_CONFIG.read_text())) aws: Final = { @@ -127,12 +129,36 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I }, } for on_flagged in ON_FLAGGED + ] + + [ + { + "guardrail_name": "lakera-during", + "litellm_params": { + "guardrail": "lakera_v2", + "mode": "during_call", + "default_on": False, + "on_flagged": "block", + "api_base": lakera.url, + "api_key": "synthetic-lakera-key", + }, + }, + { + "guardrail_name": "lakera-down", + "litellm_params": { + "guardrail": "lakera_v2", + "mode": "pre_call", + "default_on": False, + "on_flagged": "block", + "api_base": lakera_down.url, + "api_key": "synthetic-lakera-key", + }, + }, ], } ) ) owned: Final = stack.enter_context(owned_proxy_process(session_gateway, directory, {}, config=path)) - yield BedrockLakeraRig(owned, lakera, upstream) + yield BedrockLakeraRig(owned, lakera, lakera_down, upstream) def _invoke_body(prompt: str, system: str = SYSTEM) -> dict[str, JsonValue]: @@ -164,11 +190,12 @@ def _body(action: str, prompt: str, system: str = SYSTEM) -> dict[str, JsonValue return _converse_body(prompt, system) if action.startswith("converse") else _invoke_body(prompt, system) -def _send(rig: BedrockLakeraRig, on_flagged: str, action: str, body: dict[str, JsonValue]) -> httpx.Response: +def _send(rig: BedrockLakeraRig, guardrail: str, action: str, body: dict[str, JsonValue]) -> httpx.Response: gateway: Final = rig.owned.gateway with gateway.scenario() as scenario: - key: Final = scenario.key(metadata={"guardrails": [f"lakera-{on_flagged}"]}) + key: Final = scenario.key(metadata={"guardrails": [f"lakera-{guardrail}"]}) rig.lakera.drain() + rig.lakera_down.drain() rig.upstream.drain() return gateway.request("POST", f"/bedrock/model/{MODEL}/{action}", body, key=key) @@ -225,7 +252,7 @@ def test_card_in_the_system_prompt_is_sent_to_lakera(rig: BedrockLakeraRig) -> N ids=["guard-content", "tool-result-json"], ) def test_card_in_a_nested_converse_block_is_blocked(rig: BedrockLakeraRig, block: dict[str, JsonValue]) -> None: - body: Final = { + body: Final[dict[str, JsonValue]] = { **_converse_body("Hi again"), "messages": [{"role": "user", "content": [{"text": "Hi again"}, block]}], } @@ -251,3 +278,68 @@ def test_monitor_mode_forwards_the_body_unchanged(rig: BedrockLakeraRig) -> None assert response.status_code == 200, response.text assert len(_lakera_messages(rig)) == 1 assert [_JSON.validate_json(request.body) for request in rig.upstream.drain()] == [body] + + +@pytest.mark.parametrize( + "content", + [PII_PROMPT, [{"type": "text", "text": PII_PROMPT}]], + ids=["string", "text-blocks"], +) +def test_card_in_an_anthropic_tool_result_is_blocked(rig: BedrockLakeraRig, content: JsonValue) -> None: + body: Final[dict[str, JsonValue]] = { + **_invoke_body("Hi"), + "messages": [ + {"role": "user", "content": "Look up my billing details."}, + {"role": "assistant", "content": [{"type": "tool_use", "id": "t1", "name": "billing", "input": {}}]}, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "t1", "content": content}]}, + ], + } + + response: Final = _send(rig, "block", "invoke", body) + + assert response.status_code == 400, response.text + assert [role for role, text in _lakera_messages(rig)[0] if CARD in text] == ["tool"] + assert rig.upstream.drain() == () + + +@pytest.mark.parametrize("action", ["invoke", "converse"]) +def test_during_call_guardrail_reads_the_bedrock_body(rig: BedrockLakeraRig, action: str) -> None: + response: Final = _send(rig, "during", action, _body(action, PII_PROMPT)) + + assert response.status_code == 400, response.text + assert ("user", PII_PROMPT) in _lakera_messages(rig)[0] + + +@pytest.mark.parametrize("action", ["invoke", "converse"]) +def test_lakera_outage_fails_the_bedrock_request_closed(rig: BedrockLakeraRig, action: str) -> None: + response: Final = _send(rig, "down", action, _body(action, "Hi again")) + + assert response.status_code == 500, response.text + assert len(rig.lakera_down.drain()) >= 1 + assert rig.upstream.drain() == () + + +def test_a_chat_call_carrying_bedrock_keys_is_masked_not_blocked(rig: BedrockLakeraRig) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = scenario.key(metadata={"guardrails": ["lakera-mask"]}) + rig.lakera.drain() + rig.upstream.drain() + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": MODEL, + "messages": [{"role": "user", "content": PII_PROMPT}], + "custom_llm_provider": "bedrock", + "endpoint": f"/model/{MODEL}/converse", + "data": _converse_body("Hi"), + }, + key=key, + ) + + assert response.status_code == 200, response.text + assert _lakera_messages(rig) == ((("user", PII_PROMPT),),) + sent: Final = b"".join(request.body for request in rig.upstream.drain()).decode() + assert "Please update the billing" in sent + assert CARD not in sent diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py index 6e267af183d..6f1fe3256f4 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py @@ -27,7 +27,7 @@ from litellm.proxy.guardrails.guardrail_hooks.lakera_ai_v2 import ( humanize_lakera_block_reasons, ) from litellm.types.guardrails import LitellmParams, Mode -from litellm.types.utils import ModelResponse +from litellm.types.utils import CallTypesLiteral, ModelResponse from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome @@ -2483,7 +2483,10 @@ def _converse_body() -> dict[str, object]: "role": "user", "content": [{"text": "hello world"}, {"image": {"format": "png", "source": {"bytes": "AA"}}}], }, - {"role": "assistant", "content": [{"text": "hi"}, {"toolUse": {"toolUseId": "t1", "name": "f", "input": {}}}]}, + { + "role": "assistant", + "content": [{"text": "hi"}, {"toolUse": {"toolUseId": "t1", "name": "f", "input": {}}}], + }, {"role": "user", "content": [{"toolResult": {"toolUseId": "t1", "content": [{"text": "card on file"}]}}]}, ], } @@ -2513,6 +2516,11 @@ BEDROCK_PII_LAKERA_RESPONSE: Final = { "breakdown": [{"detector_type": "pii/credit_card", "detected": True, "message_id": 1}], "payload": [{"detector_type": "pii/credit_card", "start": 0, "end": 5, "message_id": 1}], } +EMAIL_AT_START_LAKERA_RESPONSE: Final = { + "flagged": True, + "breakdown": [{"detector_type": "pii/email", "detected": True, "message_id": 0}], + "payload": [{"detector_type": "pii/email", "start": 0, "end": 7, "message_id": 0}], +} BEDROCK_PROMPT_ATTACK_LAKERA_RESPONSE: Final = { "flagged": True, "breakdown": [{"detector_type": "prompt_attack", "detected": True, "message_id": 1}], @@ -2520,7 +2528,10 @@ BEDROCK_PROMPT_ATTACK_LAKERA_RESPONSE: Final = { async def _run_bedrock_pre_call( - guardrail: LakeraAIGuardrail, lakera_response: Mapping[str, object], data: dict[str, object] + guardrail: LakeraAIGuardrail, + lakera_response: Mapping[str, object], + data: dict[str, object], + call_type: CallTypesLiteral = "allm_passthrough_route", ) -> Sequence[object]: """Run the pre-call hook against a Lakera double and return the messages each guard request carried.""" sent: Final = list[object]() # mutable-ok: the Lakera double records each guard request it receives @@ -2534,14 +2545,16 @@ async def _run_bedrock_pre_call( user_api_key_dict=UserAPIKeyAuth(api_key="test_key"), cache=MagicMock(), data=data, - call_type="allm_passthrough_route", + call_type=call_type, ) return sent class TestBedrockPassthroughBody: def test_reads_converse_system_text_and_tool_results(self): - assert _build_lakera_inspection_messages(_bedrock_data(_converse_body())) == [ + assert _build_lakera_inspection_messages( + _bedrock_data(_converse_body()), call_type="allm_passthrough_route" + ) == [ {"role": "system", "content": "be nice"}, {"role": "user", "content": "hello world"}, {"role": "assistant", "content": "hi"}, @@ -2549,7 +2562,7 @@ class TestBedrockPassthroughBody: ] def test_reads_anthropic_invoke_system_text_and_tool_results(self): - assert _build_lakera_inspection_messages(_bedrock_data(_invoke_body())) == [ + assert _build_lakera_inspection_messages(_bedrock_data(_invoke_body()), call_type="allm_passthrough_route") == [ {"role": "system", "content": "be nice"}, {"role": "user", "content": "hello world"}, {"role": "assistant", "content": "hi"}, @@ -2568,7 +2581,7 @@ class TestBedrockPassthroughBody: }, ], } - assert _build_lakera_inspection_messages(_bedrock_data(body)) == [ + assert _build_lakera_inspection_messages(_bedrock_data(body), call_type="allm_passthrough_route") == [ {"role": "system", "content": "system card"}, {"role": "user", "content": "hello\nguarded card"}, {"role": "tool", "content": '{"card": "on file"}'}, @@ -2576,13 +2589,44 @@ class TestBedrockPassthroughBody: def test_a_nested_body_for_another_provider_is_not_read(self): data: Final = {**_bedrock_data(_invoke_body()), "custom_llm_provider": "vertex_ai"} - assert _build_lakera_inspection_messages(data) == [] + assert _build_lakera_inspection_messages(data, call_type="allm_passthrough_route") == [] + + @pytest.mark.parametrize("call_type", ["acompletion", "aresponses", "anthropic_messages", None]) + def test_bedrock_keys_on_another_call_type_are_not_read(self, call_type: str | None): + data: Final = {**_bedrock_data(_invoke_body()), "messages": [USER_MSG]} + assert _build_lakera_inspection_messages(data, call_type=call_type) == [USER_MSG] + + @pytest.mark.asyncio + async def test_a_chat_call_carrying_bedrock_keys_is_still_masked(self): + data: Final = {**_bedrock_data(_invoke_body()), "messages": [{"role": "user", "content": "a@b.com hi"}]} + sent: Final = await _run_bedrock_pre_call( + LakeraAIGuardrail(api_key="test_key", on_flagged="mask"), + EMAIL_AT_START_LAKERA_RESPONSE, + data, + "acompletion", + ) + assert sent == [[{"role": "user", "content": "a@b.com hi"}]] + assert data["messages"] == [{"role": "user", "content": "[MASKED EMAIL] hi"}] + assert data["data"] == _invoke_body() + + @pytest.mark.asyncio + async def test_monitor_still_masks_responses_instructions_next_to_bedrock_keys(self): + data: Final = {**_bedrock_data(_invoke_body()), "instructions": "a@b.com is the contact", "input": "hi"} + await _run_bedrock_pre_call( + LakeraAIGuardrail(api_key="test_key", on_flagged="monitor"), + EMAIL_AT_START_LAKERA_RESPONSE, + data, + "responses", + ) + assert data["instructions"] == "[MASKED EMAIL] is the contact" @pytest.mark.asyncio async def test_clean_body_reaches_lakera_and_is_forwarded_unchanged(self): data: Final = _bedrock_data(_converse_body()) sent: Final = await _run_bedrock_pre_call(LakeraAIGuardrail(api_key="test_key"), {"flagged": False}, data) - assert sent == [_build_lakera_inspection_messages(_bedrock_data(_converse_body()))] + assert sent == [ + _build_lakera_inspection_messages(_bedrock_data(_converse_body()), call_type="allm_passthrough_route") + ] assert data["data"] == _converse_body() @pytest.mark.asyncio