diff --git a/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py b/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py index d24a12ed7e8..8d7b7e1b81c 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py +++ b/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py @@ -289,6 +289,7 @@ class OvalixGuardrail(CustomGuardrail): actor: str, session_id: str, target: CheckpointTarget, + oversize: bool = False, ) -> Mapping[str, object]: """Call the Ovalix Tracker checkpoint API and return the JSON response. @@ -313,6 +314,7 @@ class OvalixGuardrail(CustomGuardrail): "data": data, "tool": "LiteLLM", **routing, + **(MappingProxyType({"oversize": True}) if oversize else _NO_METADATA), } response: Final = await self._async_handler.post( f"{self._tracker_api_base}/tracking/beta/{route}", headers=self._tracker_headers, json=payload @@ -332,9 +334,12 @@ class OvalixGuardrail(CustomGuardrail): session_id: str, target: CheckpointTarget, escalation_reason: str, + oversize: bool = False, ) -> str | None: try: - resp: Final = await self._call_checkpoint(data_type, data, checkpoint_id, actor, session_id, target) + resp: Final = await self._call_checkpoint( + data_type, data, checkpoint_id, actor, session_id, target, oversize=oversize + ) except Exception as e: verbose_proxy_logger.exception("Ovalix checkpoint call failed: %s", e) raise GuardrailRaisedException( @@ -377,7 +382,14 @@ class OvalixGuardrail(CustomGuardrail): for part in sorted(file_parts, key=lambda p: p.message_index, reverse=True): data = await self._file_part_to_data(part) reason = await self._block_reason_for_item( - "FILE", data, checkpoint_id, actor, session_id, target, _FILE_BLOCK_ESCALATION_REASON + "FILE", + data, + checkpoint_id, + actor, + session_id, + target, + _FILE_BLOCK_ESCALATION_REASON, + oversize=part.oversize, ) if reason is not None: return reason diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_ovalix.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_ovalix.py index 5b39ed42ed1..2b6d3ac91b1 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_ovalix.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_ovalix.py @@ -21,6 +21,7 @@ from litellm.proxy.guardrails.guardrail_hooks.ovalix.ovalix import ( OvalixGuardrailMissingSecrets, ResolvedRouting, ) +from litellm.proxy.guardrails.guardrail_hooks.ovalix.ovalix_extraction import FilePart from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import GenericGuardrailAPIInputs @@ -994,6 +995,26 @@ async def test_file_checkpoint_call_routes_to_litellm_file_endpoint(): assert seen["url"] == "https://t/tracking/beta/file_checkpoint" +@pytest.mark.asyncio +@pytest.mark.parametrize("oversize", [True, False]) +async def test_file_checkpoint_flags_oversize_parts(oversize): + g = _static_guardrail() + part = FilePart(name="huge.bin", data=None, mime_hint=None, inline=oversize, oversize=oversize, message_index=0) + seen = {} + + async def _post(url, headers=None, json=None): + seen["last"] = json + r = MagicMock() + r.json.return_value = _ALLOW + r.raise_for_status = MagicMock() + return r + + with patch.object(g._async_handler, "post", new=_post): + await g._check_files_for_block([part], "file-1", "a", "s", CheckpointTarget("app-1", "request")) + assert seen["last"]["data"]["content"] is None + assert seen["last"].get("oversize", False) is oversize + + @pytest.mark.asyncio async def test_tool_call_block_raises(): g = _static_guardrail()