Merge branch 'litellm_lakera_v2_gemini_contents' into litellm_passthrough_applied_guardrails_header

litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py

@@ -3,7 +3,7 @@ import json
 import os
 from collections.abc import Mapping, Sequence
 from datetime import datetime
-from itertools import chain
+from itertools import chain, product
 from string import Formatter
 from types import MappingProxyType
 from typing import Final, Literal
@@ -232,6 +232,20 @@ def _unmaskable_reason(
     return None
 
 
+_GEMINI_TOOL_PART_FIELDS: Final = MappingProxyType(
+    {
+        "functionCall": "args",
+        "function_call": "args",
+        "functionResponse": "response",
+        "function_response": "response",
+        "executableCode": "code",
+        "executable_code": "code",
+        "codeExecutionResult": "output",
+        "code_execution_result": "output",
+    }
+)
+
+
 def _gemini_items(value: object) -> Sequence[object]:
     """A repeated Gemini field as a sequence: a JSON array, or the single object Google also accepts in its place."""
     single: Final = as_str_mapping(value)
@@ -247,13 +261,19 @@ def _gemini_text(content: object) -> str:
     return "\n".join(text for text in texts if isinstance(text, str) and text)
 
 
-def _gemini_tool_results(content: object) -> str:
-    """The results one Gemini ``Content`` hands back to the model in its function response parts, as JSON."""
+def _gemini_tool_text(content: object) -> str:
+    """What one Gemini ``Content`` carries in its function call, function response and code execution parts, with
+    structured values as JSON."""
     parts: Final = _gemini_items((as_str_mapping(content) or {}).get("parts"))
-    part_fields: Final = (as_str_mapping(part) or {} for part in parts)
-    responses: Final = (fields.get("functionResponse") or fields.get("function_response") for fields in part_fields)
-    results: Final = ((as_str_mapping(response) or {}).get("response") for response in responses)
-    return "\n".join(json.dumps(result, default=str, ensure_ascii=False) for result in results if result)
+    values: Final = (
+        (as_str_mapping((as_str_mapping(part) or {}).get(part_name)) or {}).get(field)
+        for part, (part_name, field) in product(parts, _GEMINI_TOOL_PART_FIELDS.items())
+    )
+    return "\n".join(
+        value if isinstance(value, str) else json.dumps(value, default=str, ensure_ascii=False)
+        for value in values
+        if value
+    )
 
 
 def _has_gemini_contents(data: object, call_type: str | None) -> bool:
@@ -274,15 +294,15 @@ def _has_gemini_contents(data: object, call_type: str | None) -> bool:
 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``, followed by a tool message with the function results it carries. Turns with no
-    text are dropped."""
+    ``model`` read as ``assistant``, followed by a tool message with the tool parts it carries. Turns with no text
+    are dropped."""
     if not _has_gemini_contents(data, call_type):
         return []
     system_text: Final = _gemini_text(data.get("systemInstruction") or data.get("system_instruction"))
     turns: Final = chain.from_iterable(
         (
             {"role": _GEMINI_ROLES.get(str(content.get("role") or "user"), "user"), "content": _gemini_text(content)},
-            {"role": "tool", "content": _gemini_tool_results(content)},
+            {"role": "tool", "content": _gemini_tool_text(content)},
         )
         for content in map(as_str_mapping, _gemini_items(data.get("contents")))
         if content is not None

tests/integration/authorization/test_lakera_v2_gemini_contents.py

@@ -177,13 +177,17 @@ def test_card_in_gemini_body_is_blocked_before_the_provider(rig: GeminiLakeraRig
     assert rig.upstream.drain() == ()
 
 
-def test_card_in_a_gemini_function_response_is_blocked_before_the_provider(rig: GeminiLakeraRig) -> None:
-    body: Final = {
-        "contents": [
-            {"role": "model", "parts": [{"functionCall": {"name": "lookup", "args": {}}}]},
-            {"role": "user", "parts": [{"functionResponse": {"name": "lookup", "response": {"result": PII_PROMPT}}}]},
-        ]
-    }
+@pytest.mark.parametrize(
+    "part",
+    [
+        {"functionCall": {"name": "lookup", "args": {"note": PII_PROMPT}}},
+        {"functionResponse": {"name": "lookup", "response": {"result": PII_PROMPT}}},
+    ],
+)
+def test_card_in_a_gemini_tool_part_is_blocked_before_the_provider(
+    rig: GeminiLakeraRig, part: dict[str, JsonValue]
+) -> None:
+    body: Final = {"contents": [{"role": "user", "parts": [part]}]}
 
     response: Final = _send(rig, "block", body)
 

tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py

@@ -1799,7 +1799,6 @@ class TestGeminiContents:
                     "role": "user",
                     "parts": [{"text": "hi"}, {field: {"name": "lookup", "response": {"card": "4111"}}}],
                 },
-                {"role": "model", "parts": [{"functionCall": {"name": "lookup", "args": {"card": "5555"}}}]},
             ]
         }
 
@@ -1808,6 +1807,22 @@ class TestGeminiContents:
             {"role": "tool", "content": '{"card": "4111"}'},
         ]
 
+    @pytest.mark.parametrize(
+        ("part", "expected"),
+        [
+            ({"functionCall": {"name": "lookup", "args": {"card": "4111"}}}, '{"card": "4111"}'),
+            ({"function_call": {"name": "lookup", "args": {"card": "4111"}}}, '{"card": "4111"}'),
+            ({"executableCode": {"language": "PYTHON", "code": "print(4111)"}}, "print(4111)"),
+            ({"codeExecutionResult": {"outcome": "OUTCOME_OK", "output": "4111"}}, "4111"),
+        ],
+    )
+    def test_gemini_tool_part_is_read_as_a_tool_turn(self, part: dict[str, object], expected: str):
+        data: Final = {"contents": [{"role": "model", "parts": [part]}]}
+
+        assert _build_lakera_inspection_messages(data, call_type="pass_through_endpoint") == [
+            {"role": "tool", "content": expected},
+        ]
+
     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, call_type="pass_through_endpoint") == [