From ca0213ae52b145c45bf4ab94a33b95ef51f19f71 Mon Sep 17 00:00:00 2001 From: yucheng Date: Sat, 31 Oct 2026 11:26:05 -0700 Subject: [PATCH] fix(guardrails): keep reading Gemini contents when a messages or input field sits next to them --- .../guardrail_hooks/lakera_ai_v2.py | 7 +---- .../guardrail_hooks/test_lakera_ai_v2.py | 27 ++++++++++--------- 2 files changed, 15 insertions(+), 19 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py index 9c098125b37..234cac7847a 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py @@ -237,11 +237,8 @@ def _gemini_text(content: object) -> str: 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``. A body that also carries - ``messages`` or ``input`` is a chat or Responses body and is read as one.""" + ``contents`` and ``systemInstruction`` instead of ``messages``.""" fields: Final = as_str_mapping(data) or {} - if "messages" in fields or "input" in fields: - return False return isinstance(fields.get("contents"), list | dict) or isinstance( fields.get("systemInstruction") or fields.get("system_instruction"), dict ) @@ -251,8 +248,6 @@ def _gemini_inspection_messages(data: Mapping[str, object]) -> Sequence[Mapping[ """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.""" - if not _has_gemini_contents(data): - 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)} 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 1fa12a6d332..0d98e4ace64 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 @@ -1787,21 +1787,22 @@ class TestGeminiContents: ] @pytest.mark.parametrize("chat_field", ["messages", "input"]) - def test_a_chat_body_with_a_stray_contents_field_is_read_as_chat(self, chat_field: str): - data: Final = { - chat_field: [{"role": "user", "content": "hello"}], - "contents": [{"parts": [{"text": "stray"}]}], - "system_instruction": {"parts": [{"text": "stray"}]}, - } - assert _build_lakera_inspection_messages(data) == [{"role": "user", "content": "hello"}] + 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"}] @pytest.mark.asyncio - async def test_a_chat_body_with_a_stray_contents_field_is_still_masked(self): - data: Final = {"messages": [{"role": "user", "content": "alice@example.com"}], "contents": []} - await _run_gemini_pre_call( - LakeraAIGuardrail(api_key="test_key", on_flagged="mask"), PII_ONLY_LAKERA_RESPONSE, data - ) - assert data["messages"] == [{"role": "user", "content": "[MASKED EMAIL]@example.com"}] + async def test_a_body_with_messages_and_contents_is_blocked_rather_than_masked(self): + data: Final = { + "messages": [{"role": "user", "content": "alice@example.com"}], + "contents": [{"parts": [{"text": "hello"}]}], + } + with pytest.raises(HTTPException) as blocked: + await _run_gemini_pre_call( + LakeraAIGuardrail(api_key="test_key", on_flagged="mask"), PII_ONLY_LAKERA_RESPONSE, data + ) + assert blocked.value.status_code == 400 + assert data["messages"] == [{"role": "user", "content": "alice@example.com"}] @pytest.mark.asyncio async def test_clean_body_reaches_lakera_and_is_forwarded_unchanged(self):