From 4470973d00675e383de7fd55bcc3bcfc0f43ee27 Mon Sep 17 00:00:00 2001 From: yucheng Date: Sat, 31 Oct 2026 11:57:41 -0700 Subject: [PATCH] fix(guardrails): read Gemini contents only on pass-through and Google generateContent calls --- .../guardrail_hooks/lakera_ai_v2.py | 53 ++++++++++------ .../test_lakera_v2_gemini_native_route.py | 18 ++++++ .../guardrail_hooks/test_lakera_ai_v2.py | 61 ++++++++++++++++--- 3 files changed, 107 insertions(+), 25 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py index 234cac7847a..1abbca78692 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py @@ -40,13 +40,22 @@ 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]) _OBJECT_LIST: Final = TypeAdapter(list[object]) _GEMINI_ROLES: Final[Mapping[str, str]] = MappingProxyType({"user": "user", "model": "assistant", "function": "tool"}) +_GEMINI_BODY_CALL_TYPES: Final = frozenset( + { + CallTypes.pass_through.value, + CallTypes.generate_content.value, + CallTypes.agenerate_content.value, + CallTypes.generate_content_stream.value, + CallTypes.agenerate_content_stream.value, + } +) _DETECTOR_CATEGORY_PHRASES: Final[Mapping[str, str]] = MappingProxyType( { @@ -203,6 +212,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 +221,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 _has_gemini_contents(data): + if _has_gemini_contents(data, call_type): return "Gemini contents, which have no write-back path for the redacted content" if not _has_responses_instructions(guardrail, data): return "no write-back path for the redacted content" @@ -235,19 +245,22 @@ def _gemini_text(content: object) -> str: return "\n".join(text for text in texts if isinstance(text, str) and text) -def _has_gemini_contents(data: object) -> bool: - """True if ``data`` is a Gemini ``generateContent`` body, which carries its prompt in - ``contents`` and ``systemInstruction`` instead of ``messages``.""" +def _has_gemini_contents(data: object, call_type: str | None) -> bool: + """True if ``data`` is a Gemini ``generateContent`` body on a pass-through or Google ``generateContent`` call, + which carries its prompt in ``contents`` and ``systemInstruction`` instead of ``messages``.""" fields: Final = as_str_mapping(data) or {} - return isinstance(fields.get("contents"), list | dict) or isinstance( - fields.get("systemInstruction") or fields.get("system_instruction"), dict + return call_type in _GEMINI_BODY_CALL_TYPES and ( + isinstance(fields.get("contents"), list | dict) + or isinstance(fields.get("systemInstruction") or fields.get("system_instruction"), dict) ) -def _gemini_inspection_messages(data: Mapping[str, object]) -> Sequence[Mapping[str, str]]: - """Chat-style messages for a Gemini ``generateContent`` body: ``systemInstruction`` first as a - system message, then each ``contents`` turn with its text parts joined and ``model`` read as - ``assistant``. Turns with no text are dropped.""" +def _gemini_inspection_messages(data: Mapping[str, object], call_type: str | None) -> Sequence[Mapping[str, str]]: + """Chat-style messages for a Gemini ``generateContent`` body on a pass-through or Google ``generateContent`` + call: ``systemInstruction`` first as a system message, then each ``contents`` turn with its text parts joined and + ``model`` read as ``assistant``. Turns with no text are dropped.""" + if call_type not in _GEMINI_BODY_CALL_TYPES: + return [] system_text: Final = _gemini_text(data.get("systemInstruction") or data.get("system_instruction")) turns: Final = ( {"role": _GEMINI_ROLES.get(str(content.get("role") or "user"), "user"), "content": _gemini_text(content)} @@ -260,7 +273,9 @@ def _gemini_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 @@ -281,7 +296,7 @@ def _build_lakera_inspection_messages(data: Mapping[str, object]) -> Sequence[Ma ) return [ *leading, - *_gemini_inspection_messages(data), + *_gemini_inspection_messages(data, call_type), *build_inspection_messages(dict(data)), ] @@ -557,6 +572,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 @@ -564,7 +580,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, @@ -596,7 +612,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 @@ -623,7 +639,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_gemini_body: Final = _has_gemini_contents(data) # pyright: ignore[reportUnknownArgumentType] # hook data is an untyped dict + is_gemini_body: Final = _has_gemini_contents(data, call_type) # 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) @@ -697,6 +713,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( @@ -706,7 +723,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( @@ -737,7 +754,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_gemini_native_route.py b/tests/integration/authorization/test_lakera_v2_gemini_native_route.py index 47909cc81b9..33d08f1076e 100644 --- a/tests/integration/authorization/test_lakera_v2_gemini_native_route.py +++ b/tests/integration/authorization/test_lakera_v2_gemini_native_route.py @@ -200,3 +200,21 @@ def test_monitor_on_native_route_forwards_the_card_unchanged(rig: NativeGeminiRi assert response.status_code == 200, response.text assert len(_lakera_messages(rig)) == 1 assert [_JSON.validate_json(request.body)["contents"] for request in rig.upstream.drain()] == [body["contents"]] + + +def test_a_chat_call_carrying_contents_is_masked_not_blocked(rig: NativeGeminiRig) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = scenario.key(metadata={"guardrails": ["lakera-native-mask"]}) + rig.lakera.drain() + rig.upstream.drain() + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": MODEL, "messages": [{"role": "user", "content": PII_PROMPT}], "contents": []}, + key=key, + ) + + assert response.status_code == 200, response.text + assert _lakera_messages(rig) == ((("user", PII_PROMPT),),) + assert all(CARD.encode() not in request.body for request in rig.upstream.drain()) 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 0d98e4ace64..778e91e0c6b 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 @@ -1736,6 +1736,11 @@ GEMINI_INSPECTION_MESSAGES: Final = [ {"role": "user", "content": "hello world"}, {"role": "assistant", "content": "hi"}, ] +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}], +} PROMPT_ATTACK_ONLY_LAKERA_RESPONSE: Final = { "flagged": True, "breakdown": [{"detector_type": "prompt_attack", "detected": True, "message_id": 1}], @@ -1743,7 +1748,10 @@ PROMPT_ATTACK_ONLY_LAKERA_RESPONSE: Final = { async def _run_gemini_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 = "pass_through_endpoint", ) -> 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 @@ -1757,7 +1765,7 @@ async def _run_gemini_pre_call( user_api_key_dict=UserAPIKeyAuth(api_key="test_key"), cache=MagicMock(), data=data, - call_type="pass_through_endpoint", + call_type=call_type, ) return sent @@ -1769,18 +1777,23 @@ def _without_metadata(data: Mapping[str, object]) -> Mapping[str, object]: class TestGeminiContents: def test_reads_system_instruction_and_text_parts(self): - assert _build_lakera_inspection_messages(_gemini_request()) == GEMINI_INSPECTION_MESSAGES + assert ( + _build_lakera_inspection_messages(_gemini_request(), call_type="pass_through_endpoint") + == GEMINI_INSPECTION_MESSAGES + ) def test_reads_a_single_contents_object_with_a_single_parts_object(self): data: Final = {"contents": {"role": "user", "parts": {"text": "hello world"}}} - assert _build_lakera_inspection_messages(data) == [{"role": "user", "content": "hello world"}] + assert _build_lakera_inspection_messages(data, call_type="pass_through_endpoint") == [ + {"role": "user", "content": "hello world"} + ] def test_reads_snake_case_system_instruction_and_joins_text_parts(self): data: Final = { "system_instruction": {"parts": [{"text": "a"}, {"text": "b"}]}, "contents": [{"parts": [{"text": "c"}]}, {"role": "function", "parts": [{"text": "d"}]}], } - assert _build_lakera_inspection_messages(data) == [ + assert _build_lakera_inspection_messages(data, call_type="pass_through_endpoint") == [ {"role": "system", "content": "a\nb"}, {"role": "user", "content": "c"}, {"role": "tool", "content": "d"}, @@ -1789,7 +1802,41 @@ class TestGeminiContents: @pytest.mark.parametrize("chat_field", ["messages", "input"]) def test_a_chat_field_next_to_contents_does_not_hide_the_contents(self, chat_field: str): data: Final = {chat_field: [], "contents": [{"parts": [{"text": "card"}]}]} - assert _build_lakera_inspection_messages(data) == [{"role": "user", "content": "card"}] + assert _build_lakera_inspection_messages(data, call_type="pass_through_endpoint") == [ + {"role": "user", "content": "card"} + ] + + @pytest.mark.parametrize("call_type", ["agenerate_content", "agenerate_content_stream"]) + def test_reads_contents_on_the_google_generate_content_routes(self, call_type: str): + assert _build_lakera_inspection_messages(_gemini_request(), call_type=call_type) == GEMINI_INSPECTION_MESSAGES + + @pytest.mark.parametrize("call_type", ["acompletion", "aresponses", "anthropic_messages", None]) + def test_contents_on_another_call_type_are_not_read(self, call_type: str | None): + data: Final = {**_gemini_request(), "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_contents_is_still_masked(self): + data: Final[dict[str, object]] = {"messages": [{"role": "user", "content": "a@b.com hi"}], "contents": []} + sent: Final = await _run_gemini_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"}] + + @pytest.mark.asyncio + async def test_monitor_still_masks_responses_instructions_next_to_contents(self): + data: Final[dict[str, object]] = {"instructions": "a@b.com is the contact", "input": "hi", "contents": []} + await _run_gemini_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_a_body_with_messages_and_contents_is_blocked_rather_than_masked(self):