Merge branch 'litellm_lakera_v2_gemini_contents' into litellm_passthrough_applied_guardrails_header

litellm/llms/pass_through/guardrail_translation/handler.py

@@ -63,25 +63,21 @@ _REPLY_NESTED_FIELDS: Final = MappingProxyType(
     }
 )
 _REPLY_MODEL_OUTPUT_FIELDS: Final = frozenset(("args", "arguments", "input", "parameters"))
-_REPLY_MARKER_FIELDS: Final = frozenset(
-    (
-        "candidates",
-        "choices",
-        "finish_reason",
-        "object",
-        "output",
-        "predictions",
-        "stop_reason",
-        "usage",
-        "usageMetadata",
-    )
-)
+_REPLY_LIST_FIELDS: Final = ("candidates", "choices", "output", "predictions")
 _IMAGE_REPLY_FIELDS: Final = frozenset(("created", "data"))
 _NO_FIELDS: Final[frozenset[str]] = frozenset()
 
 
 def _is_provider_reply(response: Mapping[str, object]) -> bool:
-    return bool(_REPLY_MARKER_FIELDS & response.keys()) or _IMAGE_REPLY_FIELDS.issubset(response)
+    """Whether a JSON response has the shape of a provider's reply, by the type of the field that names it, since a
+    custom API's JSON can use the same field names."""
+    return (
+        any(isinstance(response.get(field), list) for field in _REPLY_LIST_FIELDS)
+        or isinstance(response.get("object"), str)
+        or (response.get("type") == "message" and "stop_reason" in response)
+        or ("finish_reason" in response and "message" in response)
+        or _IMAGE_REPLY_FIELDS.issubset(response)
+    )
 
 
 def _reply_field_position(name: str, skipped: frozenset[str] | None) -> frozenset[str] | None:

litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py

@@ -1,4 +1,5 @@
 import copy
+import json
 import os
 from collections.abc import Mapping, Sequence
 from datetime import datetime
@@ -17,7 +18,6 @@ from litellm.integrations.custom_guardrail import (
     CustomGuardrail,
 )
 from litellm.litellm_core_utils.core_helpers import as_str_mapping
-from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
 from litellm.llms.base_llm.guardrail_translation.utils import (
     effective_skip_system_message_for_guardrail,
     effective_skip_tool_message_for_guardrail,
@@ -253,7 +253,7 @@ def _gemini_tool_results(content: object) -> str:
     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(safe_dumps(result) for result in results if result)
+    return "\n".join(json.dumps(result, default=str, ensure_ascii=False) for result in results if result)
 
 
 def _has_gemini_contents(data: object, call_type: str | None) -> bool:

tests/unit/llms/pass_through/guardrail_translation/test_handler.py

@@ -286,6 +286,9 @@ async def test_post_call_scan_reads_a_field_named_like_the_envelope_in_what_the_
             "4111111111111111",
         ),
         ({"choices": [{"message": {"tool_calls": [{"function": {"arguments": {"n": 5551234567}}}]}}]}, "5551234567"),
+        ({"usage": {}, "customer": {"phone": 5551234567}}, "5551234567"),
+        ({"output": {"card": 4111111111111111}}, "4111111111111111"),
+        ({"type": "message", "phone": 5551234567}, "5551234567"),
     ],
 )
 async def test_post_call_scan_reads_an_integer_in_what_the_model_or_api_wrote(response, expected):
@@ -311,6 +314,7 @@ async def test_post_call_scan_reads_a_field_name_in_what_the_model_or_api_wrote(
         {"created": 1793600123, "data": [{"url": "https://x.example/a.png"}]},
         {"id": "batch_1", "object": "batch", "created_at": 1793600123, "request_counts": {"total": 5551234567}},
         {"predictions": [{"embeddings": {"values": [0.0123456789]}, "total": 5551234567}]},
+        {"finish_reason": "COMPLETE", "message": {"role": "assistant"}, "usage": {"total": 5551234567}},
     ],
 )
 async def test_post_call_scan_leaves_out_the_numbers_of_a_provider_reply(response):

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

@@ -1782,6 +1782,15 @@ class TestGeminiContents:
             == GEMINI_INSPECTION_MESSAGES
         )
 
+    def test_gemini_function_response_reaches_inspection_exactly_as_the_provider_gets_it(self):
+        result: Final = {"output": "4111 1111 1111 1111", "\x00output": "clean", "name": "José"}
+        data: Final = {"contents": [{"parts": [{"functionResponse": {"name": "lookup", "response": result}}]}]}
+
+        messages: Final = _build_lakera_inspection_messages(data, call_type="pass_through_endpoint")
+
+        assert json.loads(messages[0]["content"]) == result
+        assert "José" in messages[0]["content"]
+
     @pytest.mark.parametrize("field", ["functionResponse", "function_response"])
     def test_gemini_function_response_is_read_as_a_tool_turn(self, field: str):
         data: Final = {