From 0dab7f9dcc870be7ad989464b05f21494c6bf5c7 Mon Sep 17 00:00:00 2001 From: Shalom Jamil Date: Mon, 20 Jul 2026 15:36:35 +0300 Subject: [PATCH 01/16] feat(ovalix): optional routing + file-checkpoint + routing-cache config --- .../guardrail_hooks/ovalix/__init__.py | 4 ++ .../guardrail_hooks/ovalix/ovalix.py | 49 +++++++++++----- .../guardrails/guardrail_hooks/ovalix.py | 8 +++ .../guardrails/guardrail_hooks/test_ovalix.py | 56 +++++++++++++++++++ 4 files changed, 102 insertions(+), 15 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/ovalix/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/ovalix/__init__.py index b73572e4ed5..7b7f782f45b 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/ovalix/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/ovalix/__init__.py @@ -19,6 +19,8 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" application_id = getattr(litellm_params, "application_id", None) pre_checkpoint_id = getattr(litellm_params, "pre_checkpoint_id", None) post_checkpoint_id = getattr(litellm_params, "post_checkpoint_id", None) + file_checkpoint_id = getattr(litellm_params, "file_checkpoint_id", None) + enable_routing_cache = getattr(litellm_params, "enable_routing_cache", None) _ovalix_callback = OvalixGuardrail( guardrail_name=guardrail.get("guardrail_name", ""), @@ -27,6 +29,8 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" application_id=application_id, pre_checkpoint_id=pre_checkpoint_id, post_checkpoint_id=post_checkpoint_id, + file_checkpoint_id=file_checkpoint_id, + enable_routing_cache=enable_routing_cache, event_hook=litellm_params.mode, default_on=litellm_params.default_on, ) diff --git a/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py b/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py index e92ac37ca77..acfbf176d0a 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py +++ b/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py @@ -7,7 +7,9 @@ post_call (model output) checkpoints with optional correction/blocking. import datetime import hashlib import os -from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Type +import re +from collections import OrderedDict +from typing import TYPE_CHECKING, Any, Dict, List, Literal, NamedTuple, Optional, Type import httpx @@ -32,6 +34,20 @@ BLOCKED_BY_OVALIX_FALLBACK_MESSAGE = "This message was blocked by Ovalix" BLOCKED_ACTION_TYPE = "block" +class ResolvedRouting(NamedTuple): + application_id: str + checkpoint_id_pre: str | None + checkpoint_id_post: str | None + checkpoint_id_pre_file: str | None + checkpoint_id_post_file: str | None + + +def _coerce_bool(value: bool | str) -> bool: + if isinstance(value, bool): + return value + return str(value).strip().lower() in ("1", "true", "yes", "on") + + class OvalixGuardrailMissingSecrets(Exception): """Raised when required Ovalix config (API base, key, application/checkpoint IDs) is missing.""" @@ -80,6 +96,8 @@ class OvalixGuardrail(CustomGuardrail): application_id: Optional[str] = None, pre_checkpoint_id: Optional[str] = None, post_checkpoint_id: Optional[str] = None, + file_checkpoint_id: str | None = None, + enable_routing_cache: bool | None = None, **kwargs: Any, ): self._tracker_api_base = tracker_api_base or os.environ.get("OVALIX_TRACKER_API_BASE") @@ -87,6 +105,16 @@ class OvalixGuardrail(CustomGuardrail): self._application_id = application_id or os.environ.get("OVALIX_APPLICATION_ID") self._pre_checkpoint_id = pre_checkpoint_id or os.environ.get("OVALIX_PRE_CHECKPOINT_ID") self._post_checkpoint_id = post_checkpoint_id or os.environ.get("OVALIX_POST_CHECKPOINT_ID") + self._file_checkpoint_id = file_checkpoint_id or os.environ.get("OVALIX_FILE_CHECKPOINT_ID") + env_enable_routing_cache = os.environ.get("OVALIX_ENABLE_ROUTING_CACHE") + resolved_enable_routing_cache = ( + enable_routing_cache if enable_routing_cache is not None else env_enable_routing_cache + ) + self._enable_routing_cache = ( + True if resolved_enable_routing_cache is None else _coerce_bool(resolved_enable_routing_cache) + ) + self._routing_cache: OrderedDict[str, tuple[float, ResolvedRouting]] = OrderedDict() + self._app_name_regex: re.Pattern[str] | None = None if "supported_event_hooks" not in kwargs: kwargs["supported_event_hooks"] = [] @@ -113,31 +141,22 @@ class OvalixGuardrail(CustomGuardrail): ) def _validate_config(self, supported_event_hooks: List[GuardrailEventHooks]) -> None: - """Ensure required secrets and checkpoint IDs are set; auto-add hooks when IDs are present.""" + """Ensure required Tracker secrets are set; an application_id requires a checkpoint. Auto-adds both hooks.""" errors: List[str] = [] if not self._tracker_api_base: errors.append("Tracker API base, set OVALIX_TRACKER_API_BASE or pass tracker_api_base") if not self._tracker_api_key: errors.append("Tracker API key, set OVALIX_TRACKER_API_KEY or pass tracker_api_key") - if not self._application_id: - errors.append("Application ID, set OVALIX_APPLICATION_ID or pass application_id") - if not self._pre_checkpoint_id and GuardrailEventHooks.pre_call in supported_event_hooks: - errors.append("Pre-checkpoint ID, set OVALIX_PRE_CHECKPOINT_ID or pass pre_checkpoint_id") - if not self._post_checkpoint_id and GuardrailEventHooks.post_call in supported_event_hooks: - errors.append("Post-checkpoint ID, set OVALIX_POST_CHECKPOINT_ID or pass post_checkpoint_id") - if not self._pre_checkpoint_id and not self._post_checkpoint_id: - errors.append( - "Pre-checkpoint ID or Post-checkpoint ID, set OVALIX_PRE_CHECKPOINT_ID or OVALIX_POST_CHECKPOINT_ID or pass pre_checkpoint_id or post_checkpoint_id" - ) + if self._application_id and not self._pre_checkpoint_id and not self._post_checkpoint_id: + errors.append("With application_id set, provide OVALIX_PRE_CHECKPOINT_ID and/or OVALIX_POST_CHECKPOINT_ID") if errors: raise OvalixGuardrailMissingSecrets("Missing Ovalix guardrail configuration errors: " + ". ".join(errors)) - # auto-add hooks when checkpoint IDs are present - if self._pre_checkpoint_id and GuardrailEventHooks.pre_call not in supported_event_hooks: + if GuardrailEventHooks.pre_call not in supported_event_hooks: supported_event_hooks.append(GuardrailEventHooks.pre_call) - if self._post_checkpoint_id and GuardrailEventHooks.post_call not in supported_event_hooks: + if GuardrailEventHooks.post_call not in supported_event_hooks: supported_event_hooks.append(GuardrailEventHooks.post_call) def _get_actor(self, data: dict) -> str: diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/ovalix.py b/litellm/types/proxy/guardrails/guardrail_hooks/ovalix.py index 7417d1a00c9..d83f4a5ecc6 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/ovalix.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/ovalix.py @@ -30,6 +30,14 @@ class OvalixGuardrailConfigModel(GuardrailConfigModel): default=None, description="Post-checkpoint ID for the Ovalix Tracker service.", ) + file_checkpoint_id: str | None = Field( + default=None, + description="File-checkpoint ID for the Ovalix Tracker service (falls back to pre/post).", + ) + enable_routing_cache: bool | None = Field( + default=None, + description="Cache discovery-mode routing resolution per api-key alias for 1 hour. Default on.", + ) @staticmethod def ui_friendly_name() -> str: 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 4160a835ca4..0bf456830f0 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_ovalix.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_ovalix.py @@ -74,6 +74,62 @@ def _guardrail_kwargs(): } +def test_discovery_mode_initializes_without_application_id(): + guardrail = OvalixGuardrail( + tracker_api_base="https://tracker.test", + tracker_api_key="key", + guardrail_name="ovalix-test", + event_hook="pre_call", + default_on=True, + ) + assert guardrail._application_id is None + assert guardrail._enable_routing_cache is True + + +def test_routing_cache_defaults_true_and_can_disable(): + on = OvalixGuardrail( + tracker_api_base="https://t", tracker_api_key="k", guardrail_name="o", event_hook="pre_call", default_on=True + ) + assert on._enable_routing_cache is True + off = OvalixGuardrail( + tracker_api_base="https://t", + tracker_api_key="k", + enable_routing_cache=False, + guardrail_name="o", + event_hook="pre_call", + default_on=True, + ) + assert off._enable_routing_cache is False + + +def test_new_config_fields_from_params(): + guardrail = OvalixGuardrail( + tracker_api_base="https://tracker.test", + tracker_api_key="key", + application_id="app-1", + pre_checkpoint_id="pre-1", + file_checkpoint_id="file-1", + enable_routing_cache=True, + guardrail_name="ovalix-test", + event_hook="pre_call", + default_on=True, + ) + assert guardrail._file_checkpoint_id == "file-1" + assert guardrail._enable_routing_cache is True + + +def test_static_mode_requires_a_checkpoint(): + with pytest.raises(OvalixGuardrailMissingSecrets): + OvalixGuardrail( + tracker_api_base="https://tracker.test", + tracker_api_key="key", + application_id="app-1", + guardrail_name="ovalix-test", + event_hook="pre_call", + default_on=True, + ) + + class TestOvalixGuardrailConfigModel: """Minimal config model tests: wiring only.""" From 3b7bc3911748b7d92360306a0b8e156b120926c7 Mon Sep 17 00:00:00 2001 From: Shalom Jamil Date: Mon, 20 Jul 2026 15:49:01 +0300 Subject: [PATCH 02/16] feat(ovalix): add file/tool extraction helpers --- .../ovalix/ovalix_extraction.py | 267 ++++++++++++++++++ .../guardrail_hooks/test_ovalix_extraction.py | 140 +++++++++ 2 files changed, 407 insertions(+) create mode 100644 litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix_extraction.py create mode 100644 tests/test_litellm/proxy/guardrails/guardrail_hooks/test_ovalix_extraction.py diff --git a/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix_extraction.py b/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix_extraction.py new file mode 100644 index 00000000000..f3f91a85a56 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix_extraction.py @@ -0,0 +1,267 @@ +import base64 +import json +import posixpath +import re +from collections.abc import Callable +from typing import Any, NamedTuple +from urllib.parse import unquote, urlparse + +_TOOL_NAME_MAX_LENGTH = 100 +_DEFAULT_TOOL_RESULT_NAME = "tool_result" + +_DATA_URL_RE = re.compile(r"^data:(?P[^;,]+)?(?P(?:;[^;,]+)*?)(?P;base64)?,", re.IGNORECASE) +_URLSAFE_TO_STANDARD_B64 = str.maketrans("-_", "+/") + + +class FilePart(NamedTuple): + name: str | None + data: bytes | None + mime_hint: str | None + inline: bool + oversize: bool + message_index: int + + +def _split_data_url(value: str) -> tuple[str | None, str | None]: + match = _DATA_URL_RE.match(value) + if not match: + if value.lower().startswith("data:"): + return None, None + return None, value + mime = match.group("mime") or None + if not match.group("b64"): + return mime, None + return mime, value[match.end() :] + + +def _decode_base64_with_limit(b64_payload: str, size_limit: int | None) -> tuple[bytes | None, bool]: + cleaned = "".join(b64_payload.split()) + if not cleaned: + return None, False + if size_limit is not None and (len(cleaned) * 3) // 4 - 2 > size_limit: + return None, True + data = None + try: + data = base64.b64decode(cleaned, validate=True) + except ValueError: + if "-" in cleaned or "_" in cleaned: + try: + data = base64.b64decode(cleaned.translate(_URLSAFE_TO_STANDARD_B64), validate=True) + except ValueError: + return None, False + else: + return None, False + if size_limit is not None and len(data) > size_limit: + return None, True + return (data, False) if data else (None, False) + + +def _name_from_url(url: str) -> str | None: + try: + return unquote(posixpath.basename(urlparse(url).path)) or None + except ValueError: + return None + + +def _part_from_file_block(block: dict[str, Any], size_limit: int | None, message_index: int) -> FilePart | None: + file_obj = block.get("file") + if not isinstance(file_obj, dict): + return None + name = file_obj.get("filename") or file_obj.get("file_id") or None + file_data = file_obj.get("file_data") + if isinstance(file_data, str) and file_data: + mime_hint, payload = _split_data_url(file_data) + data, oversize = _decode_base64_with_limit(payload, size_limit) if payload else (None, False) + if data is not None or oversize: + return FilePart(name, data, mime_hint, True, oversize, message_index) + return FilePart(name, None, None, False, False, message_index) + + +def _part_from_image_url_block(block: dict[str, Any], size_limit: int | None, message_index: int) -> FilePart | None: + image_url = block.get("image_url") + url = image_url.get("url") if isinstance(image_url, dict) else image_url + if not isinstance(url, str) or not url: + return None + if url.startswith("data:"): + mime_hint, payload = _split_data_url(url) + data, oversize = _decode_base64_with_limit(payload, size_limit) if payload else (None, False) + if data is not None or oversize: + return FilePart(None, data, mime_hint, True, oversize, message_index) + return None + return FilePart(_name_from_url(url), None, None, False, False, message_index) + + +def _part_from_input_file_block(block: dict[str, Any], size_limit: int | None, message_index: int) -> FilePart | None: + name = block.get("filename") or block.get("file_id") or None + file_data = block.get("file_data") + if isinstance(file_data, str) and file_data: + mime_hint, payload = _split_data_url(file_data) + data, oversize = _decode_base64_with_limit(payload, size_limit) if payload else (None, False) + if data is not None or oversize: + return FilePart(name, data, mime_hint, True, oversize, message_index) + file_url = block.get("file_url") + if isinstance(file_url, str) and file_url and not name: + name = _name_from_url(file_url) + return FilePart(name, None, None, False, False, message_index) + + +def _part_from_input_audio_block(block: dict[str, Any], size_limit: int | None, message_index: int) -> FilePart | None: + audio = block.get("input_audio") + if not isinstance(audio, dict): + return None + data_b64 = audio.get("data") + if not isinstance(data_b64, str) or not data_b64: + return None + name = f"audio.{audio.get('format') or 'bin'}" + data, oversize = _decode_base64_with_limit(data_b64, size_limit) + if data is None and not oversize: + return FilePart(name, None, None, False, False, message_index) + return FilePart(name, data, None, True, oversize, message_index) + + +_BLOCK_PARSERS: dict[str, Callable[[dict[str, Any], int | None, int], FilePart | None]] = { + "file": _part_from_file_block, + "image_url": _part_from_image_url_block, + "input_image": _part_from_image_url_block, + "input_file": _part_from_input_file_block, + "input_audio": _part_from_input_audio_block, +} + + +def extract_file_parts_from_messages( + structured_messages: list[dict[str, Any]] | None, size_limit: int | None = None +) -> list[FilePart]: + parts: list[FilePart] = [] + for message_index, message in enumerate(structured_messages or []): + if not isinstance(message, dict): + continue + content = message.get("content") + if not isinstance(content, list): + continue + for block in content: + if not isinstance(block, dict): + continue + block_type = block.get("type") + if not isinstance(block_type, str): + continue + parser = _BLOCK_PARSERS.get(block_type) + if parser is None: + continue + try: + part = parser(block, size_limit, message_index) + except (TypeError, ValueError, AttributeError, KeyError): + continue + if part is not None and (part.inline or part.name): + parts.append(part) + return parts + + +def extract_file_parts_from_images(images: list[str] | None, size_limit: int | None = None) -> list[FilePart]: + parts: list[FilePart] = [] + for index, value in enumerate(images or []): + if not isinstance(value, str) or not value: + continue + if value.startswith(("http://", "https://")): + name = _name_from_url(value) + if name: + parts.append(FilePart(name, None, None, False, False, index)) + continue + mime_hint, payload = _split_data_url(value) + data, oversize = _decode_base64_with_limit(payload, size_limit) if payload else (None, False) + if data is not None or oversize: + parts.append(FilePart(None, data, mime_hint, True, oversize, index)) + return parts + + +def make_tool_data(name: str, content: str | None, tool_input: dict[str, Any] | None = None) -> dict[str, Any]: + action_name = str(name) if str(name).strip() else _DEFAULT_TOOL_RESULT_NAME + tool_name = action_name[:_TOOL_NAME_MAX_LENGTH] + if not tool_name.strip(): + tool_name = _DEFAULT_TOOL_RESULT_NAME + return {"content": content, "tool_name": tool_name, "action_name": action_name, "tool_input": tool_input or {}} + + +def _tool_call_field(tool_call: Any, key: str) -> Any: + if isinstance(tool_call, dict): + return tool_call.get(key) + return getattr(tool_call, key, None) + + +def tool_call_to_tool_data(tool_call: Any) -> dict[str, Any] | None: + function = _tool_call_field(tool_call, "function") + name = function.get("name") if isinstance(function, dict) else getattr(function, "name", None) + if not name or not str(name).strip(): + return None + raw_arguments = function.get("arguments") if isinstance(function, dict) else getattr(function, "arguments", None) + tool_input: dict[str, Any] = {} + if isinstance(raw_arguments, str): + content = raw_arguments + if content: + try: + parsed = json.loads(content) + if isinstance(parsed, dict): + tool_input = parsed + except (ValueError, TypeError): + tool_input = {} + elif raw_arguments is None: + content = "" + elif isinstance(raw_arguments, dict): + try: + content = json.dumps(raw_arguments) + except (TypeError, ValueError): + content = str(raw_arguments) + tool_input = raw_arguments + else: + try: + content = json.dumps(raw_arguments) + except (TypeError, ValueError): + content = str(raw_arguments) + return make_tool_data(name, content, tool_input) + + +def _extract_tool_content(content: Any) -> str | None: + if isinstance(content, list): + parts: list[str] = [] + for block in content: + if isinstance(block, dict): + text = block.get("text") + if isinstance(text, str) and text: + parts.append(text) + elif isinstance(block, str): + parts.append(block) + content = "\n".join(parts) + elif isinstance(content, dict): + try: + content = json.dumps(content) + except (TypeError, ValueError): + content = str(content) + if not isinstance(content, str) or not content.strip(): + return None + return content + + +def extract_tool_results(structured_messages: list[dict[str, Any]] | None) -> list[tuple[str, str, str | None]]: + id_to_name: dict[str, str] = {} + results: list[tuple[str, str, str | None]] = [] + for message in structured_messages or []: + if not isinstance(message, dict): + continue + role = message.get("role") + if role == "assistant": + for tool_call in message.get("tool_calls") or []: + if not isinstance(tool_call, dict): + continue + call_id = tool_call.get("id") + function = tool_call.get("function") + name = function.get("name") if isinstance(function, dict) else None + if isinstance(call_id, str) and call_id and name and str(name).strip(): + id_to_name[call_id] = name + elif role == "tool": + content = _extract_tool_content(message.get("content")) + if content is None: + continue + tool_call_id = message.get("tool_call_id") + resolved_name = id_to_name.get(tool_call_id) if isinstance(tool_call_id, str) else None + name = resolved_name or _DEFAULT_TOOL_RESULT_NAME + results.append((name, content, tool_call_id)) + return results diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_ovalix_extraction.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_ovalix_extraction.py new file mode 100644 index 00000000000..f8862640e75 --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_ovalix_extraction.py @@ -0,0 +1,140 @@ +import base64 + +from litellm.proxy.guardrails.guardrail_hooks.ovalix.ovalix_extraction import ( + extract_file_parts_from_images, + extract_file_parts_from_messages, + extract_tool_results, + make_tool_data, + tool_call_to_tool_data, +) + + +def _b64(raw: bytes) -> str: + return base64.b64encode(raw).decode() + + +def test_file_block_data_url_decoded(): + msgs = [ + { + "role": "user", + "content": [ + {"type": "file", "file": {"filename": "a.txt", "file_data": f"data:text/plain;base64,{_b64(b'hi')}"}} + ], + } + ] + parts = extract_file_parts_from_messages(msgs, size_limit=1000) + assert len(parts) == 1 and parts[0].data == b"hi" and parts[0].name == "a.txt" and parts[0].inline + + +def test_image_url_reference_tracked_by_name(): + msgs = [{"role": "user", "content": [{"type": "image_url", "image_url": {"url": "https://x.test/pic.png"}}]}] + parts = extract_file_parts_from_messages(msgs, size_limit=1000) + assert len(parts) == 1 and parts[0].data is None and parts[0].inline is False and parts[0].name == "pic.png" + + +def test_oversize_file_flagged_no_data(): + big = _b64(b"x" * 100) + msgs = [ + { + "role": "user", + "content": [ + { + "type": "file", + "file": {"filename": "big.bin", "file_data": f"data:application/octet-stream;base64,{big}"}, + } + ], + } + ] + parts = extract_file_parts_from_messages(msgs, size_limit=10) + assert len(parts) == 1 and parts[0].oversize is True and parts[0].data is None + + +def test_images_field_data_url(): + parts = extract_file_parts_from_images([f"data:image/png;base64,{_b64(b'png')}"], size_limit=1000) + assert len(parts) == 1 and parts[0].data == b"png" and parts[0].inline + + +def test_tool_call_to_tool_data_parses_arguments(): + td = tool_call_to_tool_data( + {"id": "c1", "type": "function", "function": {"name": "get_weather", "arguments": '{"city": "TLV"}'}} + ) + assert ( + td["content"] == '{"city": "TLV"}' + and td["tool_name"] == "get_weather" + and td["action_name"] == "get_weather" + and td["tool_input"] == {"city": "TLV"} + ) + + +def test_tool_call_malformed_dropped(): + assert tool_call_to_tool_data({"id": "c1", "type": "function", "function": {"name": ""}}) is None + assert tool_call_to_tool_data({"id": "c1"}) is None + + +def test_extract_tool_results_correlates_name(): + msgs = [ + {"role": "assistant", "tool_calls": [{"id": "c1", "function": {"name": "get_weather"}}]}, + {"role": "tool", "tool_call_id": "c1", "content": "sunny"}, + {"role": "tool", "tool_call_id": "unknown", "content": "orphan"}, + ] + results = extract_tool_results(msgs) + assert ("get_weather", "sunny", "c1") in results + assert ("tool_result", "orphan", "unknown") in results + + +def test_make_tool_data_truncates_and_defaults_name(): + td = make_tool_data(" ", "content") + assert td["tool_name"] == "tool_result" and td["action_name"] == "tool_result" + long = "x" * 200 + td2 = make_tool_data(long, "c") + assert len(td2["tool_name"]) == 100 and td2["action_name"] == long + + +def test_unhashable_block_type_skipped_without_raising(): + msgs = [{"role": "user", "content": [{"type": ["file"], "file": {"filename": "a.txt"}}]}] + parts = extract_file_parts_from_messages(msgs, size_limit=1000) + assert parts == [] + + +def test_unhashable_tool_call_id_skipped_without_raising(): + msgs = [ + {"role": "assistant", "tool_calls": [{"id": ["c1"], "function": {"name": "get_weather"}}]}, + {"role": "tool", "tool_call_id": ["c1"], "content": "sunny"}, + ] + results = extract_tool_results(msgs) + assert results == [("tool_result", "sunny", ["c1"])] + + +def test_extract_tool_results_list_form_content(): + msgs = [{"role": "tool", "tool_call_id": "c1", "content": [{"type": "text", "text": "part1"}, "part2"]}] + results = extract_tool_results(msgs) + assert ("tool_result", "part1\npart2", "c1") in results + + +def test_extract_tool_results_dict_form_content(): + msgs = [{"role": "tool", "tool_call_id": "c1", "content": {"city": "TLV"}}] + results = extract_tool_results(msgs) + assert ("tool_result", '{"city": "TLV"}', "c1") in results + + +def test_images_field_oversize_flagged_no_data(): + big = _b64(b"x" * 100) + parts = extract_file_parts_from_images([f"data:image/png;base64,{big}"], size_limit=10) + assert len(parts) == 1 and parts[0].oversize is True and parts[0].data is None + + +class _StubFunction: + def __init__(self, name, arguments): + self.name = name + self.arguments = arguments + + +class _StubToolCall: + def __init__(self, function): + self.function = function + + +def test_tool_call_to_tool_data_accepts_object_style_tool_call(): + tool_call = _StubToolCall(_StubFunction("get_weather", '{"city": "TLV"}')) + td = tool_call_to_tool_data(tool_call) + assert td["tool_name"] == "get_weather" and td["tool_input"] == {"city": "TLV"} From c023ec98b1ef1b16aede383fe7c98c2c373aad8a Mon Sep 17 00:00:00 2001 From: Shalom Jamil Date: Mon, 20 Jul 2026 16:07:16 +0300 Subject: [PATCH 03/16] feat(ovalix): regex-based discovery resolution with LRU+TTL cache --- .../guardrail_hooks/ovalix/ovalix.py | 106 ++++++++ .../guardrails/guardrail_hooks/test_ovalix.py | 246 ++++++++++++------ 2 files changed, 278 insertions(+), 74 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py b/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py index acfbf176d0a..b1dc368b8ee 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py +++ b/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py @@ -8,6 +8,7 @@ import datetime import hashlib import os import re +import time from collections import OrderedDict from typing import TYPE_CHECKING, Any, Dict, List, Literal, NamedTuple, Optional, Type @@ -32,6 +33,8 @@ if TYPE_CHECKING: BLOCKED_BY_OVALIX_FALLBACK_MESSAGE = "This message was blocked by Ovalix" BLOCKED_ACTION_TYPE = "block" +_ROUTING_CACHE_TTL_SECONDS = 3600 +_ROUTING_CACHE_MAX_SIZE = 1000 class ResolvedRouting(NamedTuple): @@ -302,6 +305,109 @@ class OvalixGuardrail(CustomGuardrail): return modified["content"] return None + def _get_key_alias(self, request_data: dict) -> str | None: + metadata = {**(request_data.get("metadata") or {}), **(request_data.get("litellm_metadata") or {})} + return metadata.get("user_api_key_alias") or metadata.get("user_api_key_key_alias") + + async def _get_app_name_regex(self) -> re.Pattern[str]: + if self._app_name_regex is not None: + return self._app_name_regex + url = f"{self._tracker_api_base}/tracking/custom_application/litellm_app_name_regex" + try: + response = await self._async_handler.get(url, headers=dict(self._tracker_headers)) + response.raise_for_status() + compiled = re.compile(response.json()["regex"]) + except Exception as e: + verbose_proxy_logger.exception("Ovalix app-name regex fetch failed: %s", e) + raise GuardrailRaisedException( + guardrail_name=self.guardrail_name, + message=f"Ovalix guardrail error: app-name regex fetch failed: {e!s}", + should_wrap_with_default_message=False, + ) from e + self._app_name_regex = compiled + return compiled + + def _extract_application_name(self, alias: str, regex: re.Pattern[str]) -> str | None: + match = regex.search(alias) + if not match: + return None + name = (match.group(1) if match.groups() else match.group(0)).strip() + return name or None + + def _routing_cache_get(self, name: str) -> ResolvedRouting | None: + entry = self._routing_cache.get(name) + if entry is None: + return None + stored_at, routing = entry + if time.monotonic() - stored_at >= _ROUTING_CACHE_TTL_SECONDS: + del self._routing_cache[name] + return None + self._routing_cache.move_to_end(name) + return routing + + def _routing_cache_put(self, name: str, routing: ResolvedRouting) -> None: + self._routing_cache[name] = (time.monotonic(), routing) + self._routing_cache.move_to_end(name) + while len(self._routing_cache) > _ROUTING_CACHE_MAX_SIZE: + self._routing_cache.popitem(last=False) + + async def _resolve_routing(self, request_data: dict) -> ResolvedRouting: + if self._application_id: + return ResolvedRouting( + self._application_id, + self._pre_checkpoint_id, + self._post_checkpoint_id, + self._file_checkpoint_id, + None, + ) + alias = self._get_key_alias(request_data) + if not alias: + raise GuardrailRaisedException( + guardrail_name=self.guardrail_name, + message="Ovalix guardrail error: no application_id configured and no user_api_key_alias to resolve by", + should_wrap_with_default_message=False, + ) + regex = await self._get_app_name_regex() + name = self._extract_application_name(alias, regex) + if not name: + raise GuardrailRaisedException( + guardrail_name=self.guardrail_name, + message="Ovalix guardrail error: could not extract an application name from the api key alias", + should_wrap_with_default_message=False, + ) + if self._enable_routing_cache: + cached = self._routing_cache_get(name) + if cached is not None: + return cached + routing = await self._resolve_via_tracker(name) + if self._enable_routing_cache: + self._routing_cache_put(name, routing) + return routing + + async def _resolve_via_tracker(self, application_name: str) -> ResolvedRouting: + url = f"{self._tracker_api_base}/tracking/custom_application/resolve_litellm_application" + try: + response = await self._async_handler.post( + url, headers=dict(self._tracker_headers), json={"application_name": application_name} + ) + response.raise_for_status() + body = response.json() + routing = ResolvedRouting( + application_id=str(body["application_id"]), + checkpoint_id_pre=body.get("checkpoint_id_pre"), + checkpoint_id_post=body.get("checkpoint_id_post"), + checkpoint_id_pre_file=body.get("checkpoint_id_pre_file"), + checkpoint_id_post_file=body.get("checkpoint_id_post_file"), + ) + except Exception as e: + verbose_proxy_logger.exception("Ovalix routing resolution failed: %s", e) + raise GuardrailRaisedException( + guardrail_name=self.guardrail_name, + message=f"Ovalix guardrail error: routing resolution failed: {e!s}", + should_wrap_with_default_message=False, + ) from e + return routing + @staticmethod def get_config_model() -> Optional[Type["GuardrailConfigModel"]]: from litellm.types.proxy.guardrails.guardrail_hooks.ovalix import ( 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 0bf456830f0..d6a8f3bf2e1 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_ovalix.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_ovalix.py @@ -15,9 +15,18 @@ from litellm.proxy.guardrails.guardrail_hooks.ovalix.ovalix import ( OvalixGuardrail, OvalixGuardrailBlockedException, OvalixGuardrailMissingSecrets, + ResolvedRouting, ) from litellm.types.utils import GenericGuardrailAPIInputs + +@pytest.fixture(autouse=True) +def _clear_ovalix_env(monkeypatch): + for key in list(os.environ.keys()): + if key.startswith("OVALIX_"): + monkeypatch.delenv(key, raising=False) + + # Example Tracker responses (as returned by the checkpoint API) TRACKER_RESPONSE_ALLOW = { "action_type": "allow", @@ -217,9 +226,7 @@ class TestOvalixGuardrail: mock_response.json.return_value = TRACKER_RESPONSE_ALLOW mock_response.raise_for_status = MagicMock() - with patch.object( - guardrail._async_handler, "post", new_callable=AsyncMock - ) as mock_post: + with patch.object(guardrail._async_handler, "post", new_callable=AsyncMock) as mock_post: mock_post.return_value = mock_response result = await guardrail._call_checkpoint( content="hello", @@ -231,9 +238,7 @@ class TestOvalixGuardrail: assert result == TRACKER_RESPONSE_ALLOW mock_post.assert_called_once() call_args = mock_post.call_args - assert call_args.args[0] == ( - "https://tracker.test/tracking/custom_application/checkpoint" - ) + assert call_args.args[0] == ("https://tracker.test/tracking/custom_application/checkpoint") body = call_args.kwargs["json"] assert body["application_id"] == "app-1" assert body["checkpoint_id"] == "pre-1" @@ -263,9 +268,7 @@ class TestOvalixGuardrail: mock_response.json.return_value = TRACKER_RESPONSE_ALLOW mock_response.raise_for_status = MagicMock() - with patch.object( - guardrail._async_handler, "post", new_callable=AsyncMock - ) as mock_post: + with patch.object(guardrail._async_handler, "post", new_callable=AsyncMock) as mock_post: mock_post.return_value = mock_response result = await guardrail.apply_guardrail( inputs=inputs, @@ -289,9 +292,7 @@ class TestOvalixGuardrail: try: guardrail = OvalixGuardrail(**_guardrail_kwargs()) inputs = GenericGuardrailAPIInputs( - structured_messages=[ - {"role": "user", "content": "Hello, my name is David."} - ], + structured_messages=[{"role": "user", "content": "Hello, my name is David."}], texts=["Hello, my name is David."], ) request_data = {} @@ -300,9 +301,7 @@ class TestOvalixGuardrail: mock_response.json.return_value = TRACKER_RESPONSE_ANONYMIZE mock_response.raise_for_status = MagicMock() - with patch.object( - guardrail._async_handler, "post", new_callable=AsyncMock - ) as mock_post: + with patch.object(guardrail._async_handler, "post", new_callable=AsyncMock) as mock_post: mock_post.return_value = mock_response result = await guardrail.apply_guardrail( inputs=inputs, @@ -335,9 +334,7 @@ class TestOvalixGuardrail: mock_response.json.return_value = TRACKER_RESPONSE_BLOCK mock_response.raise_for_status = MagicMock() - with patch.object( - guardrail._async_handler, "post", new_callable=AsyncMock - ) as mock_post: + with patch.object(guardrail._async_handler, "post", new_callable=AsyncMock) as mock_post: mock_post.return_value = mock_response with pytest.raises(OvalixGuardrailBlockedException) as exc_info: await guardrail.apply_guardrail( @@ -382,9 +379,7 @@ class TestOvalixGuardrail: resp.raise_for_status = MagicMock() return resp - with patch.object( - guardrail._async_handler, "post", new_callable=AsyncMock - ) as mock_post: + with patch.object(guardrail._async_handler, "post", new_callable=AsyncMock) as mock_post: mock_post.side_effect = side_effect result = await guardrail.apply_guardrail( inputs=inputs, @@ -411,9 +406,7 @@ class TestOvalixGuardrail: try: guardrail = OvalixGuardrail(**_guardrail_kwargs()) inputs = GenericGuardrailAPIInputs( - structured_messages=[ - {"role": "assistant", "content": "Safe assistant reply"} - ], + structured_messages=[{"role": "assistant", "content": "Safe assistant reply"}], texts=["Safe assistant reply"], ) request_data = {} @@ -422,9 +415,7 @@ class TestOvalixGuardrail: mock_response.json.return_value = TRACKER_RESPONSE_ALLOW mock_response.raise_for_status = MagicMock() - with patch.object( - guardrail._async_handler, "post", new_callable=AsyncMock - ) as mock_post: + with patch.object(guardrail._async_handler, "post", new_callable=AsyncMock) as mock_post: mock_post.return_value = mock_response result = await guardrail.apply_guardrail( inputs=inputs, @@ -454,9 +445,7 @@ class TestOvalixGuardrail: mock_response.json.return_value = TRACKER_RESPONSE_BLOCK mock_response.raise_for_status = MagicMock() - with patch.object( - guardrail._async_handler, "post", new_callable=AsyncMock - ) as mock_post: + with patch.object(guardrail._async_handler, "post", new_callable=AsyncMock) as mock_post: mock_post.return_value = mock_response with pytest.raises(OvalixGuardrailBlockedException) as exc_info: await guardrail.apply_guardrail( @@ -471,9 +460,7 @@ class TestOvalixGuardrail: assert mock_post.call_count == 1 @pytest.mark.asyncio - async def test_apply_guardrail_request_missing_modified_data_uses_original_content( - self, guardrail_with_env - ): + async def test_apply_guardrail_request_missing_modified_data_uses_original_content(self, guardrail_with_env): """When Tracker response has no modified_data.content, original content is used.""" guardrail = guardrail_with_env inputs = GenericGuardrailAPIInputs( @@ -492,9 +479,7 @@ class TestOvalixGuardrail: } mock_response.raise_for_status = MagicMock() - with patch.object( - guardrail._async_handler, "post", new_callable=AsyncMock - ) as mock_post: + with patch.object(guardrail._async_handler, "post", new_callable=AsyncMock) as mock_post: mock_post.return_value = mock_response result = await guardrail.apply_guardrail( inputs=inputs, @@ -507,9 +492,7 @@ class TestOvalixGuardrail: assert mock_post.call_count == 1 @pytest.mark.asyncio - async def test_apply_guardrail_tracker_http_error_raises_guardrail_exception( - self, guardrail_with_env - ): + async def test_apply_guardrail_tracker_http_error_raises_guardrail_exception(self, guardrail_with_env): """When Tracker returns HTTP error (e.g. 400), GuardrailRaisedException is raised.""" guardrail = guardrail_with_env inputs = GenericGuardrailAPIInputs( @@ -525,9 +508,7 @@ class TestOvalixGuardrail: response=MagicMock(status_code=400), ) - with patch.object( - guardrail._async_handler, "post", new_callable=AsyncMock - ) as mock_post: + with patch.object(guardrail._async_handler, "post", new_callable=AsyncMock) as mock_post: mock_post.return_value = mock_response with pytest.raises(GuardrailRaisedException): await guardrail.apply_guardrail( @@ -580,9 +561,7 @@ class TestOvalixGuardrail: inputs = GenericGuardrailAPIInputs(structured_messages=[], texts=[]) request_data = {} - with patch.object( - guardrail._async_handler, "post", new_callable=AsyncMock - ) as mock_post: + with patch.object(guardrail._async_handler, "post", new_callable=AsyncMock) as mock_post: result = await guardrail.apply_guardrail( inputs=inputs, request_data=request_data, @@ -603,22 +582,9 @@ class TestOvalixGuardrail: os.environ[k] = v try: guardrail = OvalixGuardrail(**_guardrail_kwargs()) - assert ( - guardrail._get_actor( - {"metadata": {"user_api_key_user_email": "a@b.com"}} - ) - == "a@b.com" - ) - assert ( - guardrail._get_actor({"metadata": {"user_api_key_user_id": "uid-1"}}) - == "uid-1" - ) - assert ( - guardrail._get_actor( - {"litellm_metadata": {"user_api_key_user_id": "uid-2"}} - ) - == "uid-2" - ) + assert guardrail._get_actor({"metadata": {"user_api_key_user_email": "a@b.com"}}) == "a@b.com" + assert guardrail._get_actor({"metadata": {"user_api_key_user_id": "uid-1"}}) == "uid-1" + assert guardrail._get_actor({"litellm_metadata": {"user_api_key_user_id": "uid-2"}}) == "uid-2" assert guardrail._get_actor({}) == "unknown" finally: for k in _ovalix_env(): @@ -656,9 +622,7 @@ class TestOvalixGuardrail: assert session_id_1 == session_id_2 assert "app-1" in session_id_1 - def test_block_current_message_raises_ovalix_blocked_exception( - self, guardrail_with_env - ): + def test_block_current_message_raises_ovalix_blocked_exception(self, guardrail_with_env): """_block_current_message raises OvalixGuardrailBlockedException with status_code 400.""" guardrail = guardrail_with_env with pytest.raises(OvalixGuardrailBlockedException) as exc_info: @@ -670,16 +634,11 @@ class TestOvalixGuardrail: """_get_trackers_corrected_message returns modified_data.content or None.""" guardrail = guardrail_with_env assert ( - guardrail._get_trackers_corrected_message( - {"modified_data": {"content": "corrected text"}} - ) + guardrail._get_trackers_corrected_message({"modified_data": {"content": "corrected text"}}) == "corrected text" ) assert guardrail._get_trackers_corrected_message({"modified_data": {}}) is None - assert ( - guardrail._get_trackers_corrected_message({"modified_data": "not-a-dict"}) - is None - ) + assert guardrail._get_trackers_corrected_message({"modified_data": "not-a-dict"}) is None @pytest.mark.asyncio async def test_apply_guardrail_response_no_texts_returns_unchanged(self): @@ -691,9 +650,7 @@ class TestOvalixGuardrail: inputs = GenericGuardrailAPIInputs() request_data = {} - with patch.object( - guardrail._async_handler, "post", new_callable=AsyncMock - ) as mock_post: + with patch.object(guardrail._async_handler, "post", new_callable=AsyncMock) as mock_post: result = await guardrail.apply_guardrail( inputs=inputs, request_data=request_data, @@ -707,3 +664,144 @@ class TestOvalixGuardrail: for k in _ovalix_env(): if k in os.environ: del os.environ[k] + + +_REGEX = r"^\s*\[([^\]]+)\]" +_ROUTING_BODY = { + "application_id": "app-9", + "checkpoint_id_pre": "pre", + "checkpoint_id_post": "post", + "checkpoint_id_pre_file": None, + "checkpoint_id_post_file": None, +} + + +def _discovery_guardrail(enable_cache=True): + return OvalixGuardrail( + tracker_api_base="https://tracker.test", + tracker_api_key="key", + enable_routing_cache=enable_cache, + guardrail_name="ovalix-test", + event_hook="pre_call", + default_on=True, + ) + + +def _alias_request_data(alias="[Weather App] prod"): + return {"metadata": {"user_api_key_alias": alias, "user_api_key_user_email": "u@x.com"}} + + +def _mock_handler(g, routing=None): + get_resp = MagicMock() + get_resp.json.return_value = {"regex": _REGEX} + get_resp.raise_for_status = MagicMock() + post_resp = MagicMock() + post_resp.json.return_value = routing or _ROUTING_BODY + post_resp.raise_for_status = MagicMock() + g._async_handler.get = AsyncMock(return_value=get_resp) + g._async_handler.post = AsyncMock(return_value=post_resp) + return g._async_handler.get, g._async_handler.post + + +@pytest.mark.asyncio +async def test_static_mode_uses_config_routing(): + g = OvalixGuardrail( + tracker_api_base="https://t", + tracker_api_key="k", + application_id="app-1", + pre_checkpoint_id="pre-1", + post_checkpoint_id="post-1", + file_checkpoint_id="file-1", + guardrail_name="o", + event_hook="pre_call", + default_on=True, + ) + routing = await g._resolve_routing({}) + assert routing == ResolvedRouting("app-1", "pre-1", "post-1", "file-1", None) + + +@pytest.mark.asyncio +async def test_discovery_extracts_name_and_resolves(): + g = _discovery_guardrail(enable_cache=False) + mock_get, mock_post = _mock_handler(g) + routing = await g._resolve_routing(_alias_request_data("[Weather App] prod")) + assert routing.application_id == "app-9" + assert mock_post.call_args.args[0].endswith("/tracking/custom_application/resolve_litellm_application") + assert mock_post.call_args.kwargs["json"] == {"application_name": "Weather App"} + assert mock_get.call_args.args[0].endswith("/tracking/custom_application/litellm_app_name_regex") + + +@pytest.mark.asyncio +async def test_regex_fetched_once_even_with_cache_off(): + g = _discovery_guardrail(enable_cache=False) + mock_get, mock_post = _mock_handler(g) + await g._resolve_routing(_alias_request_data()) + await g._resolve_routing(_alias_request_data()) + assert mock_get.call_count == 1 + assert mock_post.call_count == 2 + + +@pytest.mark.asyncio +async def test_no_bracket_alias_fails_closed(): + g = _discovery_guardrail() + _mock_handler(g) + with pytest.raises(GuardrailRaisedException): + await g._resolve_routing(_alias_request_data("no brackets here")) + + +@pytest.mark.asyncio +async def test_discovery_missing_alias_raises(): + g = _discovery_guardrail() + _mock_handler(g) + with pytest.raises(GuardrailRaisedException): + await g._resolve_routing({"metadata": {"user_api_key_user_email": "u@x.com"}}) + + +@pytest.mark.asyncio +async def test_regex_fetch_failure_raises_guardrail_exception(): + g = _discovery_guardrail(enable_cache=False) + g._async_handler.get = AsyncMock(side_effect=httpx.ConnectError("boom")) + with pytest.raises(GuardrailRaisedException): + await g._resolve_routing(_alias_request_data()) + + +@pytest.mark.asyncio +async def test_resolve_endpoint_failure_raises_guardrail_exception(): + g = _discovery_guardrail(enable_cache=False) + _mock_handler(g) + g._async_handler.post = AsyncMock(side_effect=httpx.ConnectError("boom")) + with pytest.raises(GuardrailRaisedException): + await g._resolve_routing(_alias_request_data()) + + +@pytest.mark.asyncio +async def test_resolve_missing_application_id_raises_guardrail_exception(): + g = _discovery_guardrail(enable_cache=False) + _mock_handler(g, routing={"checkpoint_id_pre": "pre", "checkpoint_id_post": "post"}) + with pytest.raises(GuardrailRaisedException): + await g._resolve_routing(_alias_request_data()) + + +@pytest.mark.asyncio +async def test_routing_cache_hit_and_ttl_expiry(monkeypatch): + g = _discovery_guardrail(enable_cache=True) + mock_get, mock_post = _mock_handler(g) + clock = [1000.0] + monkeypatch.setattr("litellm.proxy.guardrails.guardrail_hooks.ovalix.ovalix.time.monotonic", lambda: clock[0]) + await g._resolve_routing(_alias_request_data()) + clock[0] = 1000.0 + 3599 + await g._resolve_routing(_alias_request_data()) + assert mock_post.call_count == 1 + clock[0] = 1000.0 + 3601 + await g._resolve_routing(_alias_request_data()) + assert mock_post.call_count == 2 + + +@pytest.mark.asyncio +async def test_routing_cache_lru_eviction(monkeypatch): + monkeypatch.setattr("litellm.proxy.guardrails.guardrail_hooks.ovalix.ovalix._ROUTING_CACHE_MAX_SIZE", 2) + g = _discovery_guardrail(enable_cache=True) + _mock_handler(g) + for name in ("[App A] x", "[App B] x", "[App C] x"): + await g._resolve_routing(_alias_request_data(name)) + assert "App A" not in g._routing_cache and len(g._routing_cache) == 2 From fa0d3ea14cabdb9b2a68dc93144d9c049d4e7331 Mon Sep 17 00:00:00 2001 From: Shalom Jamil Date: Mon, 20 Jul 2026 16:34:51 +0300 Subject: [PATCH 04/16] feat(ovalix): scan files, tool calls, tool results via custom_application checkpoint --- .../guardrail_hooks/ovalix/ovalix.py | 270 ++++++++++++++---- .../guardrails/guardrail_hooks/test_ovalix.py | 222 +++++++++++++- 2 files changed, 424 insertions(+), 68 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py b/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py index b1dc368b8ee..756ac150314 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py +++ b/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py @@ -4,13 +4,17 @@ Use Ovalix Guardrails for your LLM calls. Supports pre_call (user input) and post_call (model output) checkpoints with optional correction/blocking. """ +import asyncio +import base64 import datetime +import gzip import hashlib +import mimetypes import os import re import time from collections import OrderedDict -from typing import TYPE_CHECKING, Any, Dict, List, Literal, NamedTuple, Optional, Type +from typing import TYPE_CHECKING, Any, List, Literal, NamedTuple, Optional, Type import httpx @@ -24,6 +28,14 @@ from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, ) +from litellm.proxy.guardrails.guardrail_hooks.ovalix.ovalix_extraction import ( + FilePart, + extract_file_parts_from_images, + extract_file_parts_from_messages, + extract_tool_results, + make_tool_data, + tool_call_to_tool_data, +) from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import GenericGuardrailAPIInputs @@ -33,8 +45,23 @@ if TYPE_CHECKING: BLOCKED_BY_OVALIX_FALLBACK_MESSAGE = "This message was blocked by Ovalix" BLOCKED_ACTION_TYPE = "block" +_MODIFY_ACTION_TYPES = ("anonymize", "sanitize") _ROUTING_CACHE_TTL_SECONDS = 3600 _ROUTING_CACHE_MAX_SIZE = 1000 +_DEFAULT_FILE_SIZE_LIMIT = 64 * 1024 * 1024 +_FILE_BLOCK_ESCALATION_REASON = ( + "This message was blocked by Ovalix because file content anonymization isn't possible via LiteLLM" +) +_TOOL_BLOCK_ESCALATION_REASON = ( + "This message was blocked by Ovalix because tool call anonymization isn't possible via LiteLLM" +) +_TOOL_RESULT_BLOCK_ESCALATION_REASON = ( + "This message was blocked by Ovalix because tool result anonymization isn't possible via LiteLLM" +) + + +def _encode_file_wire_format(raw: bytes) -> str: + return base64.b64encode(gzip.compress(raw)).decode() class ResolvedRouting(NamedTuple): @@ -183,36 +210,98 @@ class OvalixGuardrail(CustomGuardrail): def _get_session_id(self, data: dict) -> str: """Return a unique identifier for the chat/session (actor + date + application_id).""" - actor_hash = self._get_tracker_actor_id(data) - today = datetime.datetime.now(datetime.timezone.utc).strftime("%Y-%m-%d") - return f"{actor_hash}_{today}_{self._application_id}" + return self._get_session_id_for_application(data, self._application_id) async def _call_checkpoint( self, - content: str, + data_type: str, + data: dict[str, Any], checkpoint_id: str, actor: str, session_id: str, - ) -> Dict[str, Any]: + application_id: str, + ) -> dict[str, Any]: """Call the Ovalix Tracker checkpoint API and return the JSON response.""" - application_id = self._application_id if not application_id or not checkpoint_id: raise ValueError("Ovalix: application_id or checkpoint_id not resolved") url = f"{self._tracker_api_base}/tracking/custom_application/checkpoint" - headers = dict(self._tracker_headers) payload = { "application_id": application_id, "checkpoint_id": checkpoint_id, "actor": actor, "session_id": session_id, - "data_type": "TEXT", - "data": {"content": content}, + "data_type": data_type, + "data": data, + "tool": "LiteLLM", } - response = await self._async_handler.post(url, headers=headers, json=payload) + response = await self._async_handler.post(url, headers=dict(self._tracker_headers), json=payload) response.raise_for_status() return response.json() + def _verdict(self, resp: dict[str, Any]) -> tuple[str, str | None]: + return (resp.get("action_type") or "").lower(), self._get_trackers_corrected_message(resp) + + async def _block_reason_for_item( + self, + data_type: str, + data: dict[str, Any], + checkpoint_id: str, + actor: str, + session_id: str, + application_id: str, + escalation_reason: str, + ) -> str | None: + try: + resp = await self._call_checkpoint(data_type, data, checkpoint_id, actor, session_id, application_id) + except Exception as e: + verbose_proxy_logger.exception("Ovalix checkpoint call failed: %s", e) + raise GuardrailRaisedException( + guardrail_name=self.guardrail_name, + message=f"Ovalix guardrail error: {e!s}", + should_wrap_with_default_message=False, + ) from e + action, corrected = self._verdict(resp) + if action == BLOCKED_ACTION_TYPE: + return corrected or BLOCKED_BY_OVALIX_FALLBACK_MESSAGE + if action in _MODIFY_ACTION_TYPES: + return escalation_reason + return None + + async def _check_items_block_only( + self, + items: list[tuple[str, dict[str, Any]]], + checkpoint_id: str, + actor: str, + session_id: str, + application_id: str, + escalation_reason: str, + ) -> str | None: + for data_type, data in items: + reason = await self._block_reason_for_item( + data_type, data, checkpoint_id, actor, session_id, application_id, escalation_reason + ) + if reason is not None: + return reason + return None + + async def _check_files_for_block( + self, + file_parts: list[FilePart], + checkpoint_id: str, + actor: str, + session_id: str, + application_id: str, + ) -> str | None: + 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, application_id, _FILE_BLOCK_ESCALATION_REASON + ) + if reason is not None: + return reason + return None + @log_guardrail_information async def apply_guardrail( self, @@ -221,74 +310,129 @@ class OvalixGuardrail(CustomGuardrail): input_type: Literal["request", "response"], logging_obj: Optional[Any] = None, ) -> GenericGuardrailAPIInputs: - """ - Apply Ovalix guardrail to the given inputs (request or response text). + routing = await self._resolve_routing(request_data) + actor = self._get_tracker_actor_id(request_data) + session_id = self._get_session_id_for_application(request_data, routing.application_id) + is_response = input_type == "response" - Used by the unified guardrail flow and the /apply_guardrail API. - For "request", uses the pre-checkpoint; for "response", uses the post-checkpoint. + prompt_checkpoint = routing.checkpoint_id_post if is_response else routing.checkpoint_id_pre + file_checkpoint = ( + routing.checkpoint_id_post_file if is_response else routing.checkpoint_id_pre_file + ) or prompt_checkpoint + if not prompt_checkpoint: + raise GuardrailRaisedException( + guardrail_name=self.guardrail_name, + message="Ovalix guardrail error: no checkpoint resolved for input_type", + should_wrap_with_default_message=False, + ) - Args: - inputs: Guardrail API inputs (e.g. texts to check). - request_data: Full request payload (messages, metadata, response). - input_type: "request" (pre_call) or "response" (post_call). - logging_obj: Optional logging context. + structured_messages = inputs.get("structured_messages") or [] + file_parts = ( + extract_file_parts_from_images(inputs.get("images"), size_limit=_DEFAULT_FILE_SIZE_LIMIT) + if is_response + else extract_file_parts_from_messages(structured_messages, size_limit=_DEFAULT_FILE_SIZE_LIMIT) + ) + file_block = await self._check_files_for_block( + file_parts, file_checkpoint, actor, session_id, routing.application_id + ) + if file_block is not None: + self._block_current_message(file_block) - Returns: - Updated inputs (e.g. with replaced/corrected texts, or unchanged). - """ - if not self._pre_checkpoint_id and not self._post_checkpoint_id: - return inputs + tool_call_items = [ + ("TOOL", td) for td in (tool_call_to_tool_data(tc) for tc in (inputs.get("tool_calls") or [])) if td + ] + tool_block = await self._check_items_block_only( + tool_call_items, + prompt_checkpoint, + actor, + session_id, + routing.application_id, + _TOOL_BLOCK_ESCALATION_REASON, + ) + if tool_block is not None: + self._block_current_message(tool_block) + + tool_results = extract_tool_results(structured_messages) + tool_result_items = [("TOOL", make_tool_data(name, content)) for name, content, _ in tool_results] + tool_result_block = await self._check_items_block_only( + tool_result_items, + prompt_checkpoint, + actor, + session_id, + routing.application_id, + _TOOL_RESULT_BLOCK_ESCALATION_REASON, + ) + if tool_result_block is not None: + self._block_current_message(tool_result_block) - tracker_actor_id = self._get_tracker_actor_id(request_data) - session_id = self._get_session_id(request_data) texts = inputs.get("texts") or [] if not texts or not isinstance(texts, list): return inputs + skip_contents = {content for _, content, _ in tool_results} + output_texts = await self._check_texts( + texts, prompt_checkpoint, actor, session_id, routing.application_id, skip_contents + ) + if output_texts is None: + return inputs + return {**inputs, "texts": output_texts} - if input_type == "response": - if not self._post_checkpoint_id: - return inputs - corrected_llm_responses = await self._generate_post_guardrail_llm_texts( - texts, tracker_actor_id, session_id, self._post_checkpoint_id - ) - return {**inputs, "texts": corrected_llm_responses} + async def _file_part_to_data(self, part: FilePart) -> dict[str, Any]: + extension = mimetypes.guess_extension(part.mime_hint) if part.mime_hint else None + name = part.name or (f"file{extension}" if extension else "file") + content = ( + await asyncio.get_event_loop().run_in_executor(None, _encode_file_wire_format, part.data) + if part.data + else None + ) + return {"name": name, "content": content} - if self._pre_checkpoint_id: - post_guardrail_texts = await self._generate_post_guardrail_llm_texts( - texts, tracker_actor_id, session_id, self._pre_checkpoint_id - ) - return {**inputs, "texts": post_guardrail_texts} - return inputs - - async def _generate_post_guardrail_llm_texts( - self, texts: List[str], actor: str, session_id: str, checkpoint_id: str - ) -> List[str]: - """Generate post-guardrail LLM responses for the given LLM responses.""" - post_guardrail_texts: List[str] = [] - - is_first_response = True - for llm_response in reversed(texts): + async def _check_texts( + self, + texts: list[str], + checkpoint_id: str, + actor: str, + session_id: str, + application_id: str, + skip_contents: set[str], + ) -> list[str] | None: + output = list(texts) + changed = False + count = len(texts) + for reversed_index in range(count): + original_index = count - 1 - reversed_index + is_newest = reversed_index == 0 + content = texts[original_index] + if content in skip_contents: + continue try: - resp = await self._call_checkpoint(llm_response, checkpoint_id, actor, session_id) + resp = await self._call_checkpoint( + "TEXT", {"content": content}, checkpoint_id, actor, session_id, application_id + ) except Exception as e: - verbose_proxy_logger.exception("Ovalix apply_guardrail checkpoint call failed: %s", e) + verbose_proxy_logger.exception("Ovalix checkpoint call failed: %s", e) raise GuardrailRaisedException( guardrail_name=self.guardrail_name, message=f"Ovalix guardrail error: {e!s}", should_wrap_with_default_message=False, ) from e + action, corrected = self._verdict(resp) + if action == BLOCKED_ACTION_TYPE: + block_message = corrected or BLOCKED_BY_OVALIX_FALLBACK_MESSAGE + if is_newest: + self._block_current_message(block_message) + if output[original_index] != block_message: + changed = True + output[original_index] = block_message + continue + if action in _MODIFY_ACTION_TYPES and corrected is not None and corrected != content: + changed = True + output[original_index] = corrected + return output if changed else None - action_type = (resp.get("action_type") or "").lower() - blocking_message = self._get_trackers_corrected_message(resp) or BLOCKED_BY_OVALIX_FALLBACK_MESSAGE - if action_type == BLOCKED_ACTION_TYPE and is_first_response: - self._block_current_message(blocking_message) - elif action_type == BLOCKED_ACTION_TYPE: - post_guardrail_texts.insert(0, blocking_message) - else: - corrected_text = self._get_trackers_corrected_message(resp) or llm_response - post_guardrail_texts.insert(0, corrected_text) - is_first_response = False - return post_guardrail_texts + def _get_session_id_for_application(self, data: dict, application_id: str | None) -> str: + actor_hash = self._get_tracker_actor_id(data) + today = datetime.datetime.now(datetime.timezone.utc).strftime("%Y-%m-%d") + return f"{actor_hash}_{today}_{application_id}" def _block_current_message(self, blocking_message: str) -> None: """Raise OvalixGuardrailBlockedException with the given message (no default wrapper).""" @@ -358,7 +502,7 @@ class OvalixGuardrail(CustomGuardrail): self._pre_checkpoint_id, self._post_checkpoint_id, self._file_checkpoint_id, - None, + self._file_checkpoint_id, ) alias = self._get_key_alias(request_data) if not alias: 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 d6a8f3bf2e1..4c656896c3d 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_ovalix.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_ovalix.py @@ -3,6 +3,8 @@ Unit tests for Ovalix guardrail: config resolution and apply_guardrail behavior with mocked Tracker service responses (allow, anonymize, block). """ +import base64 +import gzip import os from typing import Any, List from unittest.mock import AsyncMock, MagicMock, patch @@ -217,7 +219,7 @@ class TestOvalixGuardrail: @pytest.mark.asyncio async def test_call_checkpoint_sends_correct_payload_and_returns_json(self): - """_call_checkpoint POSTs to tracker with application_id, checkpoint_id, actor, session_id, data.""" + """_call_checkpoint POSTs to tracker with application_id, checkpoint_id, actor, session_id, data, tool.""" for k, v in _ovalix_env().items(): os.environ[k] = v try: @@ -229,10 +231,12 @@ class TestOvalixGuardrail: with patch.object(guardrail._async_handler, "post", new_callable=AsyncMock) as mock_post: mock_post.return_value = mock_response result = await guardrail._call_checkpoint( - content="hello", + data_type="TEXT", + data={"content": "hello"}, checkpoint_id="pre-1", actor="a1b2c3d4", session_id="session-1", + application_id="app-1", ) assert result == TRACKER_RESPONSE_ALLOW @@ -246,6 +250,7 @@ class TestOvalixGuardrail: assert body["session_id"] == "session-1" assert body["data_type"] == "TEXT" assert body["data"] == {"content": "hello"} + assert body["tool"] == "LiteLLM" finally: for k in _ovalix_env(): if k in os.environ: @@ -400,7 +405,7 @@ class TestOvalixGuardrail: @pytest.mark.asyncio async def test_apply_guardrail_response_allow_returns_inputs(self): - """When input_type is response and Tracker allows, apply_guardrail returns inputs with texts updated from Tracker.""" + """When input_type is response and Tracker allows, apply_guardrail leaves texts unchanged (allow never rewrites).""" for k, v in _ovalix_env().items(): os.environ[k] = v try: @@ -424,7 +429,7 @@ class TestOvalixGuardrail: logging_obj=None, ) - assert result.get("texts") == ["how are you?"] + assert result.get("texts") == ["Safe assistant reply"] assert mock_post.call_count == 1 finally: for k in _ovalix_env(): @@ -717,7 +722,7 @@ async def test_static_mode_uses_config_routing(): default_on=True, ) routing = await g._resolve_routing({}) - assert routing == ResolvedRouting("app-1", "pre-1", "post-1", "file-1", None) + assert routing == ResolvedRouting("app-1", "pre-1", "post-1", "file-1", "file-1") @pytest.mark.asyncio @@ -805,3 +810,210 @@ async def test_routing_cache_lru_eviction(monkeypatch): for name in ("[App A] x", "[App B] x", "[App C] x"): await g._resolve_routing(_alias_request_data(name)) assert "App A" not in g._routing_cache and len(g._routing_cache) == 2 + + +_ALLOW = {"action_type": "allow", "modified_data": {"content": "x"}} +_BLOCK = {"action_type": "block", "modified_data": {"content": "stop-reason"}} +_ANON = {"action_type": "anonymize", "modified_data": {"content": "redacted"}} + + +def _static_guardrail(): + return OvalixGuardrail( + tracker_api_base="https://t", + tracker_api_key="k", + application_id="app-1", + pre_checkpoint_id="pre-1", + post_checkpoint_id="post-1", + file_checkpoint_id="file-1", + guardrail_name="o", + event_hook="pre_call", + default_on=True, + ) + + +def _post_returning(mapping_fn): + resp_factory = mapping_fn + + async def _post(url, headers=None, json=None): + r = MagicMock() + r.json.return_value = resp_factory(json) + r.raise_for_status = MagicMock() + return r + + return _post + + +@pytest.mark.asyncio +async def test_file_block_raises(): + g = _static_guardrail() + data_url = "data:text/plain;base64," + base64.b64encode(b"secret").decode() + inputs = GenericGuardrailAPIInputs( + texts=["hi"], + structured_messages=[ + {"role": "user", "content": [{"type": "file", "file": {"filename": "s.txt", "file_data": data_url}}]} + ], + ) + with patch.object( + g._async_handler, "post", new=_post_returning(lambda body: _BLOCK if body["data_type"] == "FILE" else _ALLOW) + ): + with pytest.raises(OvalixGuardrailBlockedException) as exc: + await g.apply_guardrail(inputs=inputs, request_data={}, input_type="request", logging_obj=None) + assert "stop-reason" in str(exc.value.message) + + +@pytest.mark.asyncio +async def test_file_uses_file_checkpoint_and_gzip_wire(): + g = _static_guardrail() + data_url = "data:text/plain;base64," + base64.b64encode(b"secret").decode() + inputs = GenericGuardrailAPIInputs( + texts=[], + structured_messages=[ + {"role": "user", "content": [{"type": "file", "file": {"filename": "s.txt", "file_data": data_url}}]} + ], + ) + 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.apply_guardrail(inputs=inputs, request_data={}, input_type="request", logging_obj=None) + assert seen["last"]["data_type"] == "FILE" + assert seen["last"]["checkpoint_id"] == "file-1" + assert gzip.decompress(base64.b64decode(seen["last"]["data"]["content"])) == b"secret" + + +@pytest.mark.asyncio +async def test_response_side_file_uses_file_checkpoint(): + g = _static_guardrail() + data_url = "data:image/png;base64," + base64.b64encode(b"img").decode() + inputs = GenericGuardrailAPIInputs(texts=[], images=[data_url]) + 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.apply_guardrail(inputs=inputs, request_data={}, input_type="response", logging_obj=None) + assert seen["last"]["data_type"] == "FILE" + assert seen["last"]["checkpoint_id"] == "file-1" + + +@pytest.mark.asyncio +async def test_tool_call_block_raises(): + g = _static_guardrail() + inputs = GenericGuardrailAPIInputs( + texts=[], tool_calls=[{"id": "c1", "type": "function", "function": {"name": "exfil", "arguments": "{}"}}] + ) + with patch.object( + g._async_handler, "post", new=_post_returning(lambda body: _BLOCK if body["data_type"] == "TOOL" else _ALLOW) + ): + with pytest.raises(OvalixGuardrailBlockedException): + await g.apply_guardrail(inputs=inputs, request_data={}, input_type="response", logging_obj=None) + + +@pytest.mark.asyncio +async def test_tool_call_anonymize_escalates_to_block(): + g = _static_guardrail() + inputs = GenericGuardrailAPIInputs( + texts=[], tool_calls=[{"id": "c1", "type": "function", "function": {"name": "exfil", "arguments": "{}"}}] + ) + with patch.object( + g._async_handler, "post", new=_post_returning(lambda body: _ANON if body["data_type"] == "TOOL" else _ALLOW) + ): + with pytest.raises(OvalixGuardrailBlockedException) as exc: + await g.apply_guardrail(inputs=inputs, request_data={}, input_type="response", logging_obj=None) + assert "tool call anonymization isn't possible" in str(exc.value.message) + + +@pytest.mark.asyncio +async def test_tool_result_block_raises(): + g = _static_guardrail() + inputs = GenericGuardrailAPIInputs( + texts=["sunny"], + structured_messages=[ + {"role": "assistant", "tool_calls": [{"id": "c1", "function": {"name": "get_weather"}}]}, + {"role": "tool", "tool_call_id": "c1", "content": "sunny"}, + ], + ) + with patch.object( + g._async_handler, "post", new=_post_returning(lambda body: _BLOCK if body["data_type"] == "TOOL" else _ALLOW) + ): + with pytest.raises(OvalixGuardrailBlockedException): + await g.apply_guardrail(inputs=inputs, request_data={}, input_type="request", logging_obj=None) + + +@pytest.mark.asyncio +async def test_newest_text_block_raises_older_anonymized(): + g = _static_guardrail() + inputs = GenericGuardrailAPIInputs(texts=["old", "new"]) + + def _map(body): + if body["data_type"] != "TEXT": + return _ALLOW + return _BLOCK if body["data"]["content"] == "new" else _ALLOW + + with patch.object(g._async_handler, "post", new=_post_returning(_map)): + with pytest.raises(OvalixGuardrailBlockedException): + await g.apply_guardrail(inputs=inputs, request_data={}, input_type="request", logging_obj=None) + + +@pytest.mark.asyncio +async def test_older_text_anonymized_returns_modified_texts(): + g = _static_guardrail() + inputs = GenericGuardrailAPIInputs(texts=["make me anon", "safe newest"]) + + def _map(body): + if body["data_type"] != "TEXT": + return _ALLOW + return ( + {"action_type": "anonymize", "modified_data": {"content": "ANON"}} + if body["data"]["content"] == "make me anon" + else _ALLOW + ) + + with patch.object(g._async_handler, "post", new=_post_returning(_map)): + result = await g.apply_guardrail(inputs=inputs, request_data={}, input_type="request", logging_obj=None) + assert result["texts"] == ["ANON", "safe newest"] + + +@pytest.mark.asyncio +async def test_all_allow_passes_through(): + g = _static_guardrail() + inputs = GenericGuardrailAPIInputs(texts=["hi"]) + with patch.object(g._async_handler, "post", new=_post_returning(lambda body: _ALLOW)): + result = await g.apply_guardrail(inputs=inputs, request_data={}, input_type="request", logging_obj=None) + assert result["texts"] == ["hi"] + + +@pytest.mark.asyncio +async def test_tool_result_content_skipped_on_text_path(): + g = _static_guardrail() + inputs = GenericGuardrailAPIInputs( + texts=["sunny"], + structured_messages=[ + {"role": "assistant", "tool_calls": [{"id": "c1", "function": {"name": "get_weather"}}]}, + {"role": "tool", "tool_call_id": "c1", "content": "sunny"}, + ], + ) + text_calls = [] + + async def _post(url, headers=None, json=None): + if json["data_type"] == "TEXT": + text_calls.append(json["data"]["content"]) + 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.apply_guardrail(inputs=inputs, request_data={}, input_type="request", logging_obj=None) + assert "sunny" not in text_calls From 8b7a8bf73b96347bbbbea869bb9526881b285dc2 Mon Sep 17 00:00:00 2001 From: Shalom Jamil Date: Tue, 21 Jul 2026 12:46:22 +0300 Subject: [PATCH 05/16] fix(ovalix): send the raw user identifier (empty when absent), not its hash, as the tracker actor --- .../guardrail_hooks/ovalix/ovalix.py | 6 +-- .../guardrails/guardrail_hooks/test_ovalix.py | 41 ++++++++++++++++++- 2 files changed, 43 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py b/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py index 756ac150314..74d4e291666 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py +++ b/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py @@ -196,12 +196,12 @@ class OvalixGuardrail(CustomGuardrail): return metadata["user_api_key_user_email"] if metadata.get("user_api_key_user_id"): return metadata["user_api_key_user_id"] - return "unknown" + return "" def _get_tracker_actor_id(self, data: dict) -> str: """Normalize the actor string into a short, stable id for Tracker API payloads.""" # NOTE: this hash is purely for normalization — it collapses an arbitrary actor - # string (email, user id, or "unknown") into a compact, fixed-length, consistent + # string (email, user id, or empty) into a compact, fixed-length, consistent # key. It is not a privacy/security measure and the actor value is not sensitive, # so a plain SHA-256 (truncated) is sufficient; no salting/KDF is needed here. actor_id = self._get_actor(data).encode() @@ -311,7 +311,7 @@ class OvalixGuardrail(CustomGuardrail): logging_obj: Optional[Any] = None, ) -> GenericGuardrailAPIInputs: routing = await self._resolve_routing(request_data) - actor = self._get_tracker_actor_id(request_data) + actor = self._get_actor(request_data) session_id = self._get_session_id_for_application(request_data, routing.application_id) is_response = input_type == "response" 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 4c656896c3d..1342dd5cbe7 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_ovalix.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_ovalix.py @@ -590,7 +590,7 @@ class TestOvalixGuardrail: assert guardrail._get_actor({"metadata": {"user_api_key_user_email": "a@b.com"}}) == "a@b.com" assert guardrail._get_actor({"metadata": {"user_api_key_user_id": "uid-1"}}) == "uid-1" assert guardrail._get_actor({"litellm_metadata": {"user_api_key_user_id": "uid-2"}}) == "uid-2" - assert guardrail._get_actor({}) == "unknown" + assert guardrail._get_actor({}) == "" finally: for k in _ovalix_env(): if k in os.environ: @@ -994,6 +994,45 @@ async def test_all_allow_passes_through(): assert result["texts"] == ["hi"] +@pytest.mark.asyncio +async def test_actor_sent_to_tracker_is_raw_identifier_not_hash(): + g = _static_guardrail() + request_data = {"metadata": {"user_api_key_user_email": "user@example.com"}} + inputs = GenericGuardrailAPIInputs(texts=["hi"]) + 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.apply_guardrail(inputs=inputs, request_data=request_data, input_type="request", logging_obj=None) + assert seen["last"]["actor"] == "user@example.com" + assert seen["last"]["session_id"] != "user@example.com" + assert g._get_tracker_actor_id(request_data) in seen["last"]["session_id"] + + +@pytest.mark.asyncio +async def test_empty_user_sends_empty_actor_matching_reference(): + g = _static_guardrail() + inputs = GenericGuardrailAPIInputs(texts=["hi"]) + 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.apply_guardrail(inputs=inputs, request_data={}, input_type="request", logging_obj=None) + assert seen["last"]["actor"] == "" + + @pytest.mark.asyncio async def test_tool_result_content_skipped_on_text_path(): g = _static_guardrail() From 05c5669ba9cc3948992a56bb196546c6e6b79971 Mon Sep 17 00:00:00 2001 From: Shalom Jamil Date: Thu, 23 Jul 2026 11:06:53 +0300 Subject: [PATCH 06/16] addressing PR comments --- .../guardrail_hooks/ovalix/ovalix.py | 16 +- .../guardrails/guardrail_hooks/test_ovalix.py | 142 +++++++++++++- .../guardrail_hooks/test_ovalix_extraction.py | 177 ++++++++++++++++++ 3 files changed, 323 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py b/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py index 74d4e291666..0254c24a118 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py +++ b/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py @@ -171,7 +171,7 @@ class OvalixGuardrail(CustomGuardrail): ) def _validate_config(self, supported_event_hooks: List[GuardrailEventHooks]) -> None: - """Ensure required Tracker secrets are set; an application_id requires a checkpoint. Auto-adds both hooks.""" + """Ensure required Tracker secrets are set; register the pre/post hooks this config can serve (both in discovery mode; only configured-checkpoint directions in static mode).""" errors: List[str] = [] if not self._tracker_api_base: @@ -184,9 +184,11 @@ class OvalixGuardrail(CustomGuardrail): if errors: raise OvalixGuardrailMissingSecrets("Missing Ovalix guardrail configuration errors: " + ". ".join(errors)) - if GuardrailEventHooks.pre_call not in supported_event_hooks: + supports_pre = not self._application_id or bool(self._pre_checkpoint_id) + supports_post = not self._application_id or bool(self._post_checkpoint_id) + if supports_pre and GuardrailEventHooks.pre_call not in supported_event_hooks: supported_event_hooks.append(GuardrailEventHooks.pre_call) - if GuardrailEventHooks.post_call not in supported_event_hooks: + if supports_post and GuardrailEventHooks.post_call not in supported_event_hooks: supported_event_hooks.append(GuardrailEventHooks.post_call) def _get_actor(self, data: dict) -> str: @@ -368,10 +370,7 @@ class OvalixGuardrail(CustomGuardrail): texts = inputs.get("texts") or [] if not texts or not isinstance(texts, list): return inputs - skip_contents = {content for _, content, _ in tool_results} - output_texts = await self._check_texts( - texts, prompt_checkpoint, actor, session_id, routing.application_id, skip_contents - ) + output_texts = await self._check_texts(texts, prompt_checkpoint, actor, session_id, routing.application_id) if output_texts is None: return inputs return {**inputs, "texts": output_texts} @@ -393,7 +392,6 @@ class OvalixGuardrail(CustomGuardrail): actor: str, session_id: str, application_id: str, - skip_contents: set[str], ) -> list[str] | None: output = list(texts) changed = False @@ -402,8 +400,6 @@ class OvalixGuardrail(CustomGuardrail): original_index = count - 1 - reversed_index is_newest = reversed_index == 0 content = texts[original_index] - if content in skip_contents: - continue try: resp = await self._call_checkpoint( "TEXT", {"content": content}, checkpoint_id, actor, session_id, application_id 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 1342dd5cbe7..40715eb94b4 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_ovalix.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_ovalix.py @@ -19,6 +19,7 @@ from litellm.proxy.guardrails.guardrail_hooks.ovalix.ovalix import ( OvalixGuardrailMissingSecrets, ResolvedRouting, ) +from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import GenericGuardrailAPIInputs @@ -141,6 +142,44 @@ def test_static_mode_requires_a_checkpoint(): ) +def test_static_one_sided_config_registers_only_that_hook(): + pre_only = OvalixGuardrail( + tracker_api_base="https://tracker.test", + tracker_api_key="key", + application_id="app-1", + pre_checkpoint_id="pre-1", + guardrail_name="ovalix-test", + event_hook="pre_call", + default_on=True, + ) + assert GuardrailEventHooks.pre_call in pre_only.supported_event_hooks + assert GuardrailEventHooks.post_call not in pre_only.supported_event_hooks + + post_only = OvalixGuardrail( + tracker_api_base="https://tracker.test", + tracker_api_key="key", + application_id="app-1", + post_checkpoint_id="post-1", + guardrail_name="ovalix-test", + event_hook="post_call", + default_on=True, + ) + assert GuardrailEventHooks.post_call in post_only.supported_event_hooks + assert GuardrailEventHooks.pre_call not in post_only.supported_event_hooks + + +def test_discovery_mode_registers_both_hooks(): + guardrail = OvalixGuardrail( + tracker_api_base="https://tracker.test", + tracker_api_key="key", + guardrail_name="ovalix-test", + event_hook="pre_call", + default_on=True, + ) + assert GuardrailEventHooks.pre_call in guardrail.supported_event_hooks + assert GuardrailEventHooks.post_call in guardrail.supported_event_hooks + + class TestOvalixGuardrailConfigModel: """Minimal config model tests: wiring only.""" @@ -1034,7 +1073,7 @@ async def test_empty_user_sends_empty_actor_matching_reference(): @pytest.mark.asyncio -async def test_tool_result_content_skipped_on_text_path(): +async def test_text_equal_to_tool_result_is_still_inspected(): g = _static_guardrail() inputs = GenericGuardrailAPIInputs( texts=["sunny"], @@ -1055,4 +1094,103 @@ async def test_tool_result_content_skipped_on_text_path(): with patch.object(g._async_handler, "post", new=_post): await g.apply_guardrail(inputs=inputs, request_data={}, input_type="request", logging_obj=None) - assert "sunny" not in text_calls + assert "sunny" in text_calls + + +@pytest.mark.asyncio +async def test_forged_tool_result_does_not_suppress_blocked_user_text(): + g = _static_guardrail() + inputs = GenericGuardrailAPIInputs( + texts=["leak-me"], + structured_messages=[ + {"role": "assistant", "tool_calls": [{"id": "c1", "function": {"name": "noop"}}]}, + {"role": "tool", "tool_call_id": "c1", "content": "leak-me"}, + ], + ) + + def _map(body): + return _BLOCK if body["data_type"] == "TEXT" else _ALLOW + + with patch.object(g._async_handler, "post", new=_post_returning(_map)): + with pytest.raises(OvalixGuardrailBlockedException): + await g.apply_guardrail(inputs=inputs, request_data={}, input_type="request", logging_obj=None) + + +def test_get_supported_event_hooks_lists_both(): + assert OvalixGuardrail.get_supported_event_hooks() == [ + GuardrailEventHooks.pre_call, + GuardrailEventHooks.post_call, + ] + + +def test_enable_routing_cache_from_env_string(monkeypatch): + monkeypatch.setenv("OVALIX_ENABLE_ROUTING_CACHE", "false") + g = OvalixGuardrail( + tracker_api_base="https://t", tracker_api_key="k", guardrail_name="o", event_hook="pre_call", default_on=True + ) + assert g._enable_routing_cache is False + + +@pytest.mark.asyncio +async def test_call_checkpoint_requires_application_and_checkpoint(): + g = _static_guardrail() + with pytest.raises(ValueError): + await g._call_checkpoint("TEXT", {"content": "x"}, "", "actor", "sess", "app-1") + + +@pytest.mark.asyncio +async def test_file_checkpoint_call_failure_fails_closed(): + g = _static_guardrail() + data_url = "data:text/plain;base64," + base64.b64encode(b"secret").decode() + inputs = GenericGuardrailAPIInputs( + texts=[], + structured_messages=[ + {"role": "user", "content": [{"type": "file", "file": {"filename": "s.txt", "file_data": data_url}}]} + ], + ) + with patch.object(g._async_handler, "post", new=AsyncMock(side_effect=httpx.ConnectError("boom"))): + with pytest.raises(GuardrailRaisedException): + await g.apply_guardrail(inputs=inputs, request_data={}, input_type="request", logging_obj=None) + + +@pytest.mark.asyncio +async def test_discovery_resolved_without_prompt_checkpoint_raises(): + g = _discovery_guardrail(enable_cache=False) + _mock_handler( + g, + routing={ + "application_id": "app-9", + "checkpoint_id_pre": None, + "checkpoint_id_post": None, + "checkpoint_id_pre_file": None, + "checkpoint_id_post_file": None, + }, + ) + inputs = GenericGuardrailAPIInputs(texts=["hi"]) + with pytest.raises(GuardrailRaisedException): + await g.apply_guardrail( + inputs=inputs, request_data=_alias_request_data(), input_type="request", logging_obj=None + ) + + +def test_initialize_guardrail_wires_new_params(monkeypatch): + import litellm + from litellm.proxy.guardrails.guardrail_hooks.ovalix import initialize_guardrail + + monkeypatch.setattr(litellm.logging_callback_manager, "add_litellm_callback", lambda callback: None) + + class _Params: + tracker_api_base = "https://t" + tracker_api_key = "k" + application_id = "app-1" + pre_checkpoint_id = "pre-1" + post_checkpoint_id = "post-1" + file_checkpoint_id = "file-1" + enable_routing_cache = False + mode = "pre_call" + default_on = True + + guardrail = initialize_guardrail(_Params(), {"guardrail_name": "ovalix"}) + assert guardrail._file_checkpoint_id == "file-1" + assert guardrail._enable_routing_cache is False + assert guardrail.guardrail_name == "ovalix" diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_ovalix_extraction.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_ovalix_extraction.py index f8862640e75..c3e5017b1aa 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_ovalix_extraction.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_ovalix_extraction.py @@ -138,3 +138,180 @@ def test_tool_call_to_tool_data_accepts_object_style_tool_call(): tool_call = _StubToolCall(_StubFunction("get_weather", '{"city": "TLV"}')) td = tool_call_to_tool_data(tool_call) assert td["tool_name"] == "get_weather" and td["tool_input"] == {"city": "TLV"} + + +def _msgs(block): + return [{"role": "user", "content": [block]}] + + +def test_image_url_block_data_url_decoded_from_messages(): + block = {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{_b64(b'png')}"}} + parts = extract_file_parts_from_messages(_msgs(block), size_limit=1000) + assert len(parts) == 1 and parts[0].data == b"png" and parts[0].inline and parts[0].name is None + + +def test_image_url_block_non_string_url_skipped(): + block = {"type": "image_url", "image_url": {"url": 123}} + assert extract_file_parts_from_messages(_msgs(block), size_limit=1000) == [] + + +def test_input_image_block_data_url_decoded(): + block = {"type": "input_image", "image_url": f"data:image/png;base64,{_b64(b'img')}"} + parts = extract_file_parts_from_messages(_msgs(block), size_limit=1000) + assert len(parts) == 1 and parts[0].data == b"img" and parts[0].inline + + +def test_file_block_non_dict_file_skipped(): + block = {"type": "file", "file": "not-a-dict"} + assert extract_file_parts_from_messages(_msgs(block), size_limit=1000) == [] + + +def test_file_block_reference_without_bytes_is_name_only(): + block = {"type": "file", "file": {"file_id": "file-abc"}} + parts = extract_file_parts_from_messages(_msgs(block), size_limit=1000) + assert len(parts) == 1 and parts[0].name == "file-abc" and parts[0].data is None and parts[0].inline is False + + +def test_file_block_urlsafe_base64_decoded_via_fallback(): + raw = b"\xff\xff\xfe" # encodes with url-unsafe chars '+'/'/' in standard b64 + urlsafe = base64.urlsafe_b64encode(raw).decode() + assert "-" in urlsafe or "_" in urlsafe + block = { + "type": "file", + "file": {"filename": "b.bin", "file_data": f"data:application/octet-stream;base64,{urlsafe}"}, + } + parts = extract_file_parts_from_messages(_msgs(block), size_limit=1000) + assert len(parts) == 1 and parts[0].data == raw + + +def test_file_block_data_url_without_base64_marker_has_no_bytes(): + block = {"type": "file", "file": {"filename": "n.txt", "file_data": "data:text/plain,hello"}} + parts = extract_file_parts_from_messages(_msgs(block), size_limit=1000) + assert len(parts) == 1 and parts[0].data is None and parts[0].inline is False and parts[0].name == "n.txt" + + +def test_input_file_block_data_url_decoded(): + block = {"type": "input_file", "filename": "doc.pdf", "file_data": f"data:application/pdf;base64,{_b64(b'pdf')}"} + parts = extract_file_parts_from_messages(_msgs(block), size_limit=1000) + assert len(parts) == 1 and parts[0].data == b"pdf" and parts[0].name == "doc.pdf" and parts[0].inline + + +def test_input_file_block_file_url_reference_name_only(): + block = {"type": "input_file", "file_url": "https://x.test/report.csv"} + parts = extract_file_parts_from_messages(_msgs(block), size_limit=1000) + assert len(parts) == 1 and parts[0].name == "report.csv" and parts[0].data is None and parts[0].inline is False + + +def test_input_audio_block_decoded(): + block = {"type": "input_audio", "input_audio": {"data": _b64(b"wav"), "format": "wav"}} + parts = extract_file_parts_from_messages(_msgs(block), size_limit=1000) + assert len(parts) == 1 and parts[0].data == b"wav" and parts[0].name == "audio.wav" and parts[0].inline + + +def test_input_audio_block_undecodable_is_name_only(): + block = {"type": "input_audio", "input_audio": {"data": "!!!not-base64!!!"}} + parts = extract_file_parts_from_messages(_msgs(block), size_limit=1000) + assert len(parts) == 1 and parts[0].name == "audio.bin" and parts[0].data is None and parts[0].inline is False + + +def test_input_audio_block_non_dict_skipped(): + block = {"type": "input_audio", "input_audio": "nope"} + assert extract_file_parts_from_messages(_msgs(block), size_limit=1000) == [] + + +def test_tool_call_dict_arguments_serialized_and_parsed(): + td = tool_call_to_tool_data({"function": {"name": "f", "arguments": {"a": 1}}}) + assert td["content"] == '{"a": 1}' and td["tool_input"] == {"a": 1} + + +def test_tool_call_none_arguments_yields_empty_content(): + td = tool_call_to_tool_data({"function": {"name": "f", "arguments": None}}) + assert td["content"] == "" and td["tool_input"] == {} + + +def test_tool_call_non_string_non_dict_arguments_serialized(): + td = tool_call_to_tool_data({"function": {"name": "f", "arguments": [1, 2]}}) + assert td["content"] == "[1, 2]" and td["tool_input"] == {} + + +def test_tool_call_invalid_json_string_arguments_kept_as_content(): + td = tool_call_to_tool_data({"function": {"name": "f", "arguments": "{not json"}}) + assert td["content"] == "{not json" and td["tool_input"] == {} + + +def test_tool_result_with_non_string_tool_call_id_uses_default_name(): + msgs = [{"role": "tool", "tool_call_id": ["c1"], "content": "orphan"}] + results = extract_tool_results(msgs) + assert results == [("tool_result", "orphan", ["c1"])] + + +def test_images_field_http_url_is_name_only_reference(): + parts = extract_file_parts_from_images(["https://x.test/pic.png"], size_limit=1000) + assert len(parts) == 1 and parts[0].name == "pic.png" and parts[0].data is None and parts[0].inline is False + + +def test_messages_skip_non_dict_and_unknown_blocks(): + msgs = [ + "not-a-message", + { + "role": "user", + "content": [ + "bare-string-block", + {"type": "text", "text": "hi"}, + {"type": "file", "file": {"filename": "a.txt", "file_data": f"data:text/plain;base64,{_b64(b'x')}"}}, + ], + }, + ] + parts = extract_file_parts_from_messages(msgs, size_limit=1000) + assert len(parts) == 1 and parts[0].name == "a.txt" and parts[0].data == b"x" + + +def test_image_url_block_invalid_data_url_returns_no_part(): + block = {"type": "image_url", "image_url": {"url": "data:image/png;base64,%%%invalid%%%"}} + assert extract_file_parts_from_messages(_msgs(block), size_limit=1000) == [] + + +def test_tool_message_with_empty_content_is_skipped(): + msgs = [{"role": "tool", "tool_call_id": "c1", "content": " "}] + assert extract_tool_results(msgs) == [] + + +def test_extract_tool_results_skips_non_dict_messages_and_tool_calls(): + msgs = [ + "junk", + {"role": "assistant", "tool_calls": ["not-a-dict", {"id": "c1", "function": {"name": "f"}}]}, + {"role": "tool", "tool_call_id": "c1", "content": "ok"}, + ] + assert extract_tool_results(msgs) == [("f", "ok", "c1")] + + +def test_malformed_data_url_yields_no_bytes(): + block = {"type": "file", "file": {"filename": "x.bin", "file_data": "data:garbage-no-comma"}} + parts = extract_file_parts_from_messages(_msgs(block), size_limit=1000) + assert len(parts) == 1 and parts[0].data is None and parts[0].inline is False and parts[0].name == "x.bin" + + +def test_input_audio_block_without_data_skipped(): + block = {"type": "input_audio", "input_audio": {"format": "wav"}} + assert extract_file_parts_from_messages(_msgs(block), size_limit=1000) == [] + + +def test_images_field_non_string_entries_skipped(): + assert extract_file_parts_from_images([123, None, ""], size_limit=1000) == [] + + +def test_make_tool_data_whitespace_after_truncation_defaults_name(): + name = " " * 100 + "x" + td = make_tool_data(name, "c") + assert td["tool_name"] == "tool_result" and td["action_name"] == name + + +def test_raw_base64_file_data_without_data_url_prefix_decoded(): + block = {"type": "file", "file": {"filename": "a.bin", "file_data": _b64(b"rawbytes")}} + parts = extract_file_parts_from_messages(_msgs(block), size_limit=1000) + assert len(parts) == 1 and parts[0].data == b"rawbytes" and parts[0].mime_hint is None + + +def test_images_field_raw_base64_without_data_url_prefix_decoded(): + parts = extract_file_parts_from_images([_b64(b"rawimg")], size_limit=1000) + assert len(parts) == 1 and parts[0].data == b"rawimg" From f93552cb1a559db961c269c11f00ee9bc9272c64 Mon Sep 17 00:00:00 2001 From: Shalom Jamil Date: Fri, 24 Jul 2026 06:05:56 +0300 Subject: [PATCH 07/16] more tests --- .../guardrail_hooks/test_ovalix_extraction.py | 45 +++++++++++++++++++ 1 file changed, 45 insertions(+) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_ovalix_extraction.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_ovalix_extraction.py index c3e5017b1aa..25d67d8dda9 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_ovalix_extraction.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_ovalix_extraction.py @@ -315,3 +315,48 @@ def test_raw_base64_file_data_without_data_url_prefix_decoded(): def test_images_field_raw_base64_without_data_url_prefix_decoded(): parts = extract_file_parts_from_images([_b64(b"rawimg")], size_limit=1000) assert len(parts) == 1 and parts[0].data == b"rawimg" + + +def test_whitespace_only_base64_payload_yields_no_bytes(): + block = {"type": "file", "file": {"filename": "n.txt", "file_data": "data:text/plain;base64, "}} + parts = extract_file_parts_from_messages(_msgs(block), size_limit=1000) + assert len(parts) == 1 and parts[0].data is None and parts[0].inline is False and parts[0].name == "n.txt" + + +def test_base64_invalid_even_after_urlsafe_translate_yields_no_bytes(): + block = {"type": "file", "file": {"filename": "b.bin", "file_data": "_@@@"}} + parts = extract_file_parts_from_messages(_msgs(block), size_limit=1000) + assert len(parts) == 1 and parts[0].data is None and parts[0].inline is False and parts[0].name == "b.bin" + + +def test_file_oversize_detected_after_decode_when_estimate_passes(): + block = {"type": "file", "file": {"filename": "s.bin", "file_data": _b64(b"abc")}} + parts = extract_file_parts_from_messages(_msgs(block), size_limit=2) + assert len(parts) == 1 and parts[0].oversize is True and parts[0].data is None + + +def test_image_http_url_that_urlparse_rejects_is_dropped(): + parts = extract_file_parts_from_images(["http://["], size_limit=1000) + assert parts == [] + + +def test_message_content_not_a_list_is_skipped(): + parts = extract_file_parts_from_messages( + [{"role": "user", "content": "just a plain string prompt"}], size_limit=1000 + ) + assert parts == [] + + +def test_tool_call_dict_arguments_non_serializable_falls_back_to_str(): + td = tool_call_to_tool_data({"function": {"name": "f", "arguments": {"a": {1, 2}}}}) + assert isinstance(td["content"], str) and td["tool_input"] == {"a": {1, 2}} + + +def test_tool_call_non_serializable_other_arguments_falls_back_to_str(): + td = tool_call_to_tool_data({"function": {"name": "f", "arguments": {1, 2}}}) + assert isinstance(td["content"], str) and td["tool_input"] == {} + + +def test_tool_result_dict_content_non_serializable_falls_back_to_str(): + results = extract_tool_results([{"role": "tool", "tool_call_id": "c1", "content": {"x": {1, 2}}}]) + assert len(results) == 1 and isinstance(results[0][1], str) and results[0][1].strip() From a33c6247e6b4539c4f393519ed3645e33a35ebe6 Mon Sep 17 00:00:00 2001 From: Shalom Jamil Date: Sun, 2 Aug 2026 14:12:32 +0300 Subject: [PATCH 08/16] name resolve failed policy flag --- .../guardrail_hooks/ovalix/__init__.py | 2 + .../guardrail_hooks/ovalix/ovalix.py | 89 +++--- .../guardrails/guardrail_hooks/ovalix.py | 7 + .../guardrails/guardrail_hooks/test_ovalix.py | 256 ++++++++++++++++++ 4 files changed, 322 insertions(+), 32 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/ovalix/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/ovalix/__init__.py index 7b7f782f45b..de557f7ff6e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/ovalix/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/ovalix/__init__.py @@ -21,6 +21,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" post_checkpoint_id = getattr(litellm_params, "post_checkpoint_id", None) file_checkpoint_id = getattr(litellm_params, "file_checkpoint_id", None) enable_routing_cache = getattr(litellm_params, "enable_routing_cache", None) + fail_if_no_application = getattr(litellm_params, "fail_if_no_application", None) _ovalix_callback = OvalixGuardrail( guardrail_name=guardrail.get("guardrail_name", ""), @@ -31,6 +32,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" post_checkpoint_id=post_checkpoint_id, file_checkpoint_id=file_checkpoint_id, enable_routing_cache=enable_routing_cache, + fail_if_no_application=fail_if_no_application, event_hook=litellm_params.mode, default_on=litellm_params.default_on, ) diff --git a/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py b/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py index 0254c24a118..b75fab49ad0 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py +++ b/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py @@ -46,7 +46,9 @@ if TYPE_CHECKING: BLOCKED_BY_OVALIX_FALLBACK_MESSAGE = "This message was blocked by Ovalix" BLOCKED_ACTION_TYPE = "block" _MODIFY_ACTION_TYPES = ("anonymize", "sanitize") +_APPLICATION_NOT_FOUND_STATUS = 404 _ROUTING_CACHE_TTL_SECONDS = 3600 +_ROUTING_CACHE_NEGATIVE_TTL_SECONDS = 300 _ROUTING_CACHE_MAX_SIZE = 1000 _DEFAULT_FILE_SIZE_LIMIT = 64 * 1024 * 1024 _FILE_BLOCK_ESCALATION_REASON = ( @@ -128,6 +130,7 @@ class OvalixGuardrail(CustomGuardrail): post_checkpoint_id: Optional[str] = None, file_checkpoint_id: str | None = None, enable_routing_cache: bool | None = None, + fail_if_no_application: bool | None = None, **kwargs: Any, ): self._tracker_api_base = tracker_api_base or os.environ.get("OVALIX_TRACKER_API_BASE") @@ -143,7 +146,14 @@ class OvalixGuardrail(CustomGuardrail): self._enable_routing_cache = ( True if resolved_enable_routing_cache is None else _coerce_bool(resolved_enable_routing_cache) ) - self._routing_cache: OrderedDict[str, tuple[float, ResolvedRouting]] = OrderedDict() + env_fail_if_no_application = os.environ.get("OVALIX_FAIL_IF_NO_APPLICATION") + resolved_fail_if_no_application = ( + fail_if_no_application if fail_if_no_application is not None else env_fail_if_no_application + ) + self._fail_if_no_application = ( + True if resolved_fail_if_no_application is None else _coerce_bool(resolved_fail_if_no_application) + ) + self._routing_cache: OrderedDict[str, tuple[float, ResolvedRouting | None]] = OrderedDict() self._app_name_regex: re.Pattern[str] | None = None if "supported_event_hooks" not in kwargs: @@ -313,6 +323,8 @@ class OvalixGuardrail(CustomGuardrail): logging_obj: Optional[Any] = None, ) -> GenericGuardrailAPIInputs: routing = await self._resolve_routing(request_data) + if routing is None: + return inputs actor = self._get_actor(request_data) session_id = self._get_session_id_for_application(request_data, routing.application_id) is_response = input_type == "response" @@ -474,24 +486,45 @@ class OvalixGuardrail(CustomGuardrail): name = (match.group(1) if match.groups() else match.group(0)).strip() return name or None - def _routing_cache_get(self, name: str) -> ResolvedRouting | None: + def _routing_cache_get(self, name: str) -> tuple[bool, ResolvedRouting | None]: entry = self._routing_cache.get(name) if entry is None: - return None - stored_at, routing = entry - if time.monotonic() - stored_at >= _ROUTING_CACHE_TTL_SECONDS: + return False, None + expires_at, routing = entry + if time.monotonic() >= expires_at: del self._routing_cache[name] - return None + return False, None self._routing_cache.move_to_end(name) - return routing + return True, routing - def _routing_cache_put(self, name: str, routing: ResolvedRouting) -> None: - self._routing_cache[name] = (time.monotonic(), routing) + def _routing_cache_put(self, name: str, routing: ResolvedRouting | None) -> None: + ttl = _ROUTING_CACHE_TTL_SECONDS if routing is not None else _ROUTING_CACHE_NEGATIVE_TTL_SECONDS + self._routing_cache[name] = (time.monotonic() + ttl, routing) self._routing_cache.move_to_end(name) while len(self._routing_cache) > _ROUTING_CACHE_MAX_SIZE: self._routing_cache.popitem(last=False) - async def _resolve_routing(self, request_data: dict) -> ResolvedRouting: + def _no_application(self, reason: str) -> None: + if self._fail_if_no_application: + raise GuardrailRaisedException( + guardrail_name=self.guardrail_name, + message=f"Ovalix guardrail error: {reason}", + should_wrap_with_default_message=False, + ) + verbose_proxy_logger.warning( + "Ovalix guardrail passing the call through unguarded (fail_if_no_application=false): %s", reason + ) + return None + + def _routing_error(self, error: Exception) -> GuardrailRaisedException: + verbose_proxy_logger.exception("Ovalix routing resolution failed: %s", error) + return GuardrailRaisedException( + guardrail_name=self.guardrail_name, + message=f"Ovalix guardrail error: routing resolution failed: {error!s}", + should_wrap_with_default_message=False, + ) + + async def _resolve_routing(self, request_data: dict) -> ResolvedRouting | None: if self._application_id: return ResolvedRouting( self._application_id, @@ -502,29 +535,23 @@ class OvalixGuardrail(CustomGuardrail): ) alias = self._get_key_alias(request_data) if not alias: - raise GuardrailRaisedException( - guardrail_name=self.guardrail_name, - message="Ovalix guardrail error: no application_id configured and no user_api_key_alias to resolve by", - should_wrap_with_default_message=False, - ) + return self._no_application("no application_id configured and no user_api_key_alias to resolve by") regex = await self._get_app_name_regex() name = self._extract_application_name(alias, regex) if not name: - raise GuardrailRaisedException( - guardrail_name=self.guardrail_name, - message="Ovalix guardrail error: could not extract an application name from the api key alias", - should_wrap_with_default_message=False, - ) + return self._no_application("could not extract an application name from the api key alias") if self._enable_routing_cache: - cached = self._routing_cache_get(name) - if cached is not None: - return cached + hit, cached = self._routing_cache_get(name) + if hit: + return cached if cached is not None else self._no_application(f"application '{name}' was not found") routing = await self._resolve_via_tracker(name) if self._enable_routing_cache: self._routing_cache_put(name, routing) + if routing is None: + return self._no_application(f"application '{name}' was not found") return routing - async def _resolve_via_tracker(self, application_name: str) -> ResolvedRouting: + async def _resolve_via_tracker(self, application_name: str) -> ResolvedRouting | None: url = f"{self._tracker_api_base}/tracking/custom_application/resolve_litellm_application" try: response = await self._async_handler.post( @@ -532,21 +559,19 @@ class OvalixGuardrail(CustomGuardrail): ) response.raise_for_status() body = response.json() - routing = ResolvedRouting( + return ResolvedRouting( application_id=str(body["application_id"]), checkpoint_id_pre=body.get("checkpoint_id_pre"), checkpoint_id_post=body.get("checkpoint_id_post"), checkpoint_id_pre_file=body.get("checkpoint_id_pre_file"), checkpoint_id_post_file=body.get("checkpoint_id_post_file"), ) + except httpx.HTTPStatusError as e: + if e.response.status_code == _APPLICATION_NOT_FOUND_STATUS: + return None + raise self._routing_error(e) from e except Exception as e: - verbose_proxy_logger.exception("Ovalix routing resolution failed: %s", e) - raise GuardrailRaisedException( - guardrail_name=self.guardrail_name, - message=f"Ovalix guardrail error: routing resolution failed: {e!s}", - should_wrap_with_default_message=False, - ) from e - return routing + raise self._routing_error(e) from e @staticmethod def get_config_model() -> Optional[Type["GuardrailConfigModel"]]: diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/ovalix.py b/litellm/types/proxy/guardrails/guardrail_hooks/ovalix.py index d83f4a5ecc6..00af16bf5c0 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/ovalix.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/ovalix.py @@ -38,6 +38,13 @@ class OvalixGuardrailConfigModel(GuardrailConfigModel): default=None, description="Cache discovery-mode routing resolution per api-key alias for 1 hour. Default on.", ) + fail_if_no_application: bool | None = Field( + default=None, + description=( + "Fail the call when no application is configured and none is discovered. Default on; " + "set false to let such calls through unguarded instead." + ), + ) @staticmethod def ui_friendly_name() -> str: 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 40715eb94b4..42fd6c19ad0 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_ovalix.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_ovalix.py @@ -1194,3 +1194,259 @@ def test_initialize_guardrail_wires_new_params(monkeypatch): assert guardrail._file_checkpoint_id == "file-1" assert guardrail._enable_routing_cache is False assert guardrail.guardrail_name == "ovalix" + + +def test_fail_if_no_application_defaults_true(): + g = OvalixGuardrail( + tracker_api_base="https://t", tracker_api_key="k", guardrail_name="o", event_hook="pre_call", default_on=True + ) + assert g._fail_if_no_application is True + + +def test_fail_if_no_application_from_param(): + g = OvalixGuardrail( + tracker_api_base="https://t", + tracker_api_key="k", + fail_if_no_application=False, + guardrail_name="o", + event_hook="pre_call", + default_on=True, + ) + assert g._fail_if_no_application is False + + +def test_fail_if_no_application_from_env_string(monkeypatch): + monkeypatch.setenv("OVALIX_FAIL_IF_NO_APPLICATION", "false") + g = OvalixGuardrail( + tracker_api_base="https://t", tracker_api_key="k", guardrail_name="o", event_hook="pre_call", default_on=True + ) + assert g._fail_if_no_application is False + + +def test_explicit_param_beats_env_for_fail_if_no_application(monkeypatch): + monkeypatch.setenv("OVALIX_FAIL_IF_NO_APPLICATION", "false") + g = OvalixGuardrail( + tracker_api_base="https://t", + tracker_api_key="k", + fail_if_no_application=True, + guardrail_name="o", + event_hook="pre_call", + default_on=True, + ) + assert g._fail_if_no_application is True + + +def test_initialize_guardrail_wires_fail_if_no_application(monkeypatch): + import litellm + from litellm.proxy.guardrails.guardrail_hooks.ovalix import initialize_guardrail + + monkeypatch.setattr(litellm.logging_callback_manager, "add_litellm_callback", lambda callback: None) + + class _Params: + tracker_api_base = "https://t" + tracker_api_key = "k" + application_id = None + pre_checkpoint_id = None + post_checkpoint_id = None + file_checkpoint_id = None + enable_routing_cache = None + fail_if_no_application = False + mode = "pre_call" + default_on = True + + guardrail = initialize_guardrail(_Params(), {"guardrail_name": "ovalix"}) + assert guardrail._fail_if_no_application is False + + +def test_config_model_declares_fail_if_no_application(): + from litellm.types.proxy.guardrails.guardrail_hooks.ovalix import OvalixGuardrailConfigModel + + assert OvalixGuardrailConfigModel.model_fields["fail_if_no_application"].default is None + + +def _fail_open_discovery_guardrail(enable_cache=False): + return OvalixGuardrail( + tracker_api_base="https://tracker.test", + tracker_api_key="key", + enable_routing_cache=enable_cache, + fail_if_no_application=False, + guardrail_name="ovalix-test", + event_hook="pre_call", + default_on=True, + ) + + +def _http_status_error(status_code): + request = httpx.Request("POST", "https://tracker.test/x") + response = httpx.Response(status_code, request=request) + return httpx.HTTPStatusError("boom", request=request, response=response) + + +def _mock_handler_resolve_status(g, status_code): + get_resp = MagicMock() + get_resp.json.return_value = {"regex": _REGEX} + get_resp.raise_for_status = MagicMock() + post_resp = MagicMock() + post_resp.raise_for_status = MagicMock(side_effect=_http_status_error(status_code)) + g._async_handler.get = AsyncMock(return_value=get_resp) + g._async_handler.post = AsyncMock(return_value=post_resp) + return g._async_handler.post + + +@pytest.mark.asyncio +async def test_fail_open_missing_alias_returns_none(): + g = _fail_open_discovery_guardrail() + _mock_handler(g) + assert await g._resolve_routing({"metadata": {"user_api_key_user_email": "u@x.com"}}) is None + + +@pytest.mark.asyncio +async def test_fail_open_unparseable_alias_returns_none(): + g = _fail_open_discovery_guardrail() + _mock_handler(g) + assert await g._resolve_routing(_alias_request_data("no brackets here")) is None + + +@pytest.mark.asyncio +async def test_fail_open_tracker_404_returns_none(): + g = _fail_open_discovery_guardrail() + _mock_handler_resolve_status(g, 404) + assert await g._resolve_routing(_alias_request_data("[Ghost App] prod")) is None + + +@pytest.mark.asyncio +async def test_fail_closed_tracker_404_still_raises(): + g = _discovery_guardrail(enable_cache=False) + _mock_handler_resolve_status(g, 404) + with pytest.raises(GuardrailRaisedException): + await g._resolve_routing(_alias_request_data("[Ghost App] prod")) + + +@pytest.mark.asyncio +async def test_fail_open_tracker_500_still_raises(): + g = _fail_open_discovery_guardrail() + _mock_handler_resolve_status(g, 500) + with pytest.raises(GuardrailRaisedException): + await g._resolve_routing(_alias_request_data("[Weather App] prod")) + + +@pytest.mark.asyncio +async def test_fail_open_tracker_unreachable_still_raises(): + g = _fail_open_discovery_guardrail() + _mock_handler(g) + g._async_handler.post = AsyncMock(side_effect=httpx.ConnectError("boom")) + with pytest.raises(GuardrailRaisedException): + await g._resolve_routing(_alias_request_data()) + + +@pytest.mark.asyncio +async def test_fail_open_regex_fetch_failure_still_raises(): + g = _fail_open_discovery_guardrail() + g._async_handler.get = AsyncMock(side_effect=httpx.ConnectError("boom")) + with pytest.raises(GuardrailRaisedException): + await g._resolve_routing(_alias_request_data()) + + +@pytest.mark.asyncio +async def test_fail_open_no_application_makes_no_checkpoint_call(): + g = _fail_open_discovery_guardrail() + _mock_handler_resolve_status(g, 404) + checkpoint = AsyncMock() + with patch.object(g, "_call_checkpoint", new=checkpoint): + inputs = GenericGuardrailAPIInputs(texts=["hello"]) + result = await g.apply_guardrail( + inputs=inputs, + request_data=_alias_request_data("[Ghost App] prod"), + input_type="request", + logging_obj=None, + ) + assert result == inputs + checkpoint.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_fail_open_resolved_application_without_checkpoint_still_raises(): + g = _fail_open_discovery_guardrail() + _mock_handler( + g, + routing={ + "application_id": "app-9", + "checkpoint_id_pre": None, + "checkpoint_id_post": None, + "checkpoint_id_pre_file": None, + "checkpoint_id_post_file": None, + }, + ) + with pytest.raises(GuardrailRaisedException): + await g.apply_guardrail( + inputs=GenericGuardrailAPIInputs(texts=["hi"]), + request_data=_alias_request_data(), + input_type="request", + logging_obj=None, + ) + + +@pytest.mark.asyncio +async def test_fail_open_still_guards_when_application_resolves(): + g = _fail_open_discovery_guardrail() + _mock_handler(g) + checkpoint = AsyncMock(return_value=_BLOCK) + with patch.object(g, "_call_checkpoint", new=checkpoint): + with pytest.raises(OvalixGuardrailBlockedException): + await g.apply_guardrail( + inputs=GenericGuardrailAPIInputs(texts=["hi"]), + request_data=_alias_request_data(), + input_type="request", + logging_obj=None, + ) + + +@pytest.mark.asyncio +async def test_negative_routing_cache_hit_and_expiry(monkeypatch): + g = _fail_open_discovery_guardrail(enable_cache=True) + mock_post = _mock_handler_resolve_status(g, 404) + clock = [1000.0] + monkeypatch.setattr("litellm.proxy.guardrails.guardrail_hooks.ovalix.ovalix.time.monotonic", lambda: clock[0]) + + assert await g._resolve_routing(_alias_request_data("[Ghost App] prod")) is None + clock[0] = 1000.0 + 299 + assert await g._resolve_routing(_alias_request_data("[Ghost App] prod")) is None + assert mock_post.call_count == 1 + + clock[0] = 1000.0 + 301 + assert await g._resolve_routing(_alias_request_data("[Ghost App] prod")) is None + assert mock_post.call_count == 2 + + +@pytest.mark.asyncio +async def test_negative_cache_expires_sooner_than_positive(monkeypatch): + g = _discovery_guardrail(enable_cache=True) + _mock_handler(g) + clock = [1000.0] + monkeypatch.setattr("litellm.proxy.guardrails.guardrail_hooks.ovalix.ovalix.time.monotonic", lambda: clock[0]) + await g._resolve_routing(_alias_request_data("[Weather App] prod")) + clock[0] = 1000.0 + 301 + hit, cached = g._routing_cache_get("Weather App") + assert hit is True and cached is not None + + +@pytest.mark.asyncio +async def test_negative_result_not_cached_when_cache_disabled(): + g = _fail_open_discovery_guardrail(enable_cache=False) + mock_post = _mock_handler_resolve_status(g, 404) + await g._resolve_routing(_alias_request_data("[Ghost App] prod")) + await g._resolve_routing(_alias_request_data("[Ghost App] prod")) + assert mock_post.call_count == 2 + assert len(g._routing_cache) == 0 + + +@pytest.mark.asyncio +async def test_cached_404_still_raises_when_failing_closed(monkeypatch): + g = _discovery_guardrail(enable_cache=True) + mock_post = _mock_handler_resolve_status(g, 404) + clock = [1000.0] + monkeypatch.setattr("litellm.proxy.guardrails.guardrail_hooks.ovalix.ovalix.time.monotonic", lambda: clock[0]) + for _ in range(2): + with pytest.raises(GuardrailRaisedException): + await g._resolve_routing(_alias_request_data("[Ghost App] prod")) + assert mock_post.call_count == 1 From 51855a20f524e23f798c68cdcc13811660fbe93b Mon Sep 17 00:00:00 2001 From: Shalom Jamil Date: Mon, 3 Aug 2026 12:32:33 +0300 Subject: [PATCH 09/16] fixing PR comments --- .../guardrail_hooks/ovalix/ovalix.py | 127 +++++--- .../ovalix/ovalix_extraction.py | 302 +++++++++++------- .../guardrails/guardrail_hooks/test_ovalix.py | 131 +++++++- .../guardrail_hooks/test_ovalix_extraction.py | 67 +++- 4 files changed, 435 insertions(+), 192 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py b/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py index 65c068162e9..b22b8a41dae 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py +++ b/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py @@ -14,7 +14,9 @@ import os import re import time from collections import OrderedDict -from typing import TYPE_CHECKING, Any, List, Literal, NamedTuple, Optional, Type +from collections.abc import Mapping, Sequence +from types import MappingProxyType +from typing import TYPE_CHECKING, Any, Literal, NamedTuple import httpx @@ -35,6 +37,7 @@ from litellm.proxy.guardrails.guardrail_hooks.ovalix.ovalix_extraction import ( extract_tool_results, make_tool_data, tool_call_to_tool_data, + tool_result_text_indices, ) from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import GenericGuardrailAPIInputs @@ -51,6 +54,7 @@ _ROUTING_CACHE_TTL_SECONDS = 3600 _ROUTING_CACHE_NEGATIVE_TTL_SECONDS = 300 _ROUTING_CACHE_MAX_SIZE = 1000 _DEFAULT_FILE_SIZE_LIMIT = 64 * 1024 * 1024 +_NO_METADATA: Mapping[str, Any] = MappingProxyType({}) _FILE_BLOCK_ESCALATION_REASON = ( "This message was blocked by Ovalix because file content anonymization isn't possible via LiteLLM" ) @@ -159,12 +163,16 @@ class OvalixGuardrail(CustomGuardrail): self._validate_config(kwargs["supported_event_hooks"]) - self._tracker_headers = httpx.Headers( - { - "Authorization": f"Bearer {self._tracker_api_key}", - "Content-Type": "application/json", - }, - encoding="utf-8", + self._tracker_headers = dict( + httpx.Headers( + MappingProxyType( + { + "Authorization": f"Bearer {self._tracker_api_key}", + "Content-Type": "application/json", + } + ), + encoding="utf-8", + ) ) self._async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) @@ -180,14 +188,18 @@ class OvalixGuardrail(CustomGuardrail): def _validate_config(self, supported_event_hooks: list[GuardrailEventHooks]) -> None: """Ensure required Tracker secrets are set; register the pre/post hooks this config can serve (both in discovery mode; only configured-checkpoint directions in static mode).""" - errors: list[str] = [] - - if not self._tracker_api_base: - errors.append("Tracker API base, set OVALIX_TRACKER_API_BASE or pass tracker_api_base") - if not self._tracker_api_key: - errors.append("Tracker API key, set OVALIX_TRACKER_API_KEY or pass tracker_api_key") - if self._application_id and not self._pre_checkpoint_id and not self._post_checkpoint_id: - errors.append("With application_id set, provide OVALIX_PRE_CHECKPOINT_ID and/or OVALIX_POST_CHECKPOINT_ID") + errors = tuple( + message + for present, message in ( + (not self._tracker_api_base, "Tracker API base, set OVALIX_TRACKER_API_BASE or pass tracker_api_base"), + (not self._tracker_api_key, "Tracker API key, set OVALIX_TRACKER_API_KEY or pass tracker_api_key"), + ( + bool(self._application_id) and not self._pre_checkpoint_id and not self._post_checkpoint_id, + "With application_id set, provide OVALIX_PRE_CHECKPOINT_ID and/or OVALIX_POST_CHECKPOINT_ID", + ), + ) + if present + ) if errors: raise OvalixGuardrailMissingSecrets("Missing Ovalix guardrail configuration errors: " + ". ".join(errors)) @@ -199,16 +211,16 @@ class OvalixGuardrail(CustomGuardrail): if supports_post and GuardrailEventHooks.post_call not in supported_event_hooks: supported_event_hooks.append(GuardrailEventHooks.post_call) - def _get_actor(self, data: dict) -> str: + def _get_actor(self, data: Mapping[str, Any]) -> str: """Return a stable actor identifier from request metadata (e.g. user email or id).""" - metadata = data.get("metadata") or data.get("litellm_metadata") or {} + metadata = data.get("metadata") or data.get("litellm_metadata") or _NO_METADATA if metadata.get("user_api_key_user_email"): return metadata["user_api_key_user_email"] if metadata.get("user_api_key_user_id"): return metadata["user_api_key_user_id"] return "" - def _get_tracker_actor_id(self, data: dict) -> str: + def _get_tracker_actor_id(self, data: Mapping[str, Any]) -> str: """Normalize the actor string into a short, stable id for Tracker API payloads.""" # NOTE: this hash is purely for normalization — it collapses an arbitrary actor # string (email, user id, or empty) into a compact, fixed-length, consistent @@ -218,24 +230,28 @@ class OvalixGuardrail(CustomGuardrail): normalized_actor_id = hashlib.sha256(actor_id).hexdigest()[:8] return normalized_actor_id - def _get_session_id(self, data: dict) -> str: + def _get_session_id(self, data: Mapping[str, Any]) -> str: """Return a unique identifier for the chat/session (actor + date + application_id).""" return self._get_session_id_for_application(data, self._application_id) async def _call_checkpoint( self, data_type: str, - data: dict[str, Any], + data: Mapping[str, Any], checkpoint_id: str, actor: str, session_id: str, application_id: str, - ) -> dict[str, Any]: + ) -> Mapping[str, Any]: """Call the Ovalix Tracker checkpoint API and return the JSON response.""" if not application_id or not checkpoint_id: raise ValueError("Ovalix: application_id or checkpoint_id not resolved") - url = f"{self._tracker_api_base}/tracking/custom_application/checkpoint" + url = ( + f"{self._tracker_api_base}/tracking/litellm/file_checkpoint" + if data_type == "FILE" + else f"{self._tracker_api_base}/tracking/custom_application/checkpoint" + ) payload = { "application_id": application_id, "checkpoint_id": checkpoint_id, @@ -245,17 +261,17 @@ class OvalixGuardrail(CustomGuardrail): "data": data, "tool": "LiteLLM", } - response = await self._async_handler.post(url, headers=dict(self._tracker_headers), json=payload) + response = await self._async_handler.post(url, headers=self._tracker_headers, json=payload) response.raise_for_status() return response.json() - def _verdict(self, resp: dict[str, Any]) -> tuple[str, str | None]: + def _verdict(self, resp: Mapping[str, Any]) -> tuple[str, str | None]: return (resp.get("action_type") or "").lower(), self._get_trackers_corrected_message(resp) async def _block_reason_for_item( self, data_type: str, - data: dict[str, Any], + data: Mapping[str, Any], checkpoint_id: str, actor: str, session_id: str, @@ -280,7 +296,7 @@ class OvalixGuardrail(CustomGuardrail): async def _check_items_block_only( self, - items: list[tuple[str, dict[str, Any]]], + items: Sequence[tuple[str, Mapping[str, Any]]], checkpoint_id: str, actor: str, session_id: str, @@ -297,7 +313,7 @@ class OvalixGuardrail(CustomGuardrail): async def _check_files_for_block( self, - file_parts: list[FilePart], + file_parts: Sequence[FilePart], checkpoint_id: str, actor: str, session_id: str, @@ -316,7 +332,7 @@ class OvalixGuardrail(CustomGuardrail): async def apply_guardrail( self, inputs: GenericGuardrailAPIInputs, - request_data: dict, + request_data: Mapping[str, Any], input_type: Literal["request", "response"], logging_obj: Any | None = None, ) -> GenericGuardrailAPIInputs: @@ -338,7 +354,7 @@ class OvalixGuardrail(CustomGuardrail): should_wrap_with_default_message=False, ) - structured_messages = inputs.get("structured_messages") or [] + structured_messages = inputs.get("structured_messages") or () file_parts = ( extract_file_parts_from_images(inputs.get("images"), size_limit=_DEFAULT_FILE_SIZE_LIMIT) if is_response @@ -350,9 +366,9 @@ class OvalixGuardrail(CustomGuardrail): if file_block is not None: self._block_current_message(file_block) - tool_call_items = [ - ("TOOL", td) for td in (tool_call_to_tool_data(tc) for tc in (inputs.get("tool_calls") or [])) if td - ] + tool_call_items = tuple( + ("TOOL", td) for td in (tool_call_to_tool_data(tc) for tc in (inputs.get("tool_calls") or ())) if td + ) tool_block = await self._check_items_block_only( tool_call_items, prompt_checkpoint, @@ -365,7 +381,7 @@ class OvalixGuardrail(CustomGuardrail): self._block_current_message(tool_block) tool_results = extract_tool_results(structured_messages) - tool_result_items = [("TOOL", make_tool_data(name, content)) for name, content, _ in tool_results] + tool_result_items = tuple(("TOOL", make_tool_data(name, content)) for name, content, _ in tool_results) tool_result_block = await self._check_items_block_only( tool_result_items, prompt_checkpoint, @@ -377,15 +393,22 @@ class OvalixGuardrail(CustomGuardrail): if tool_result_block is not None: self._block_current_message(tool_result_block) - texts = inputs.get("texts") or [] + texts = inputs.get("texts") or () if not texts or not isinstance(texts, list): return inputs - output_texts = await self._check_texts(texts, prompt_checkpoint, actor, session_id, routing.application_id) + output_texts = await self._check_texts( + texts, + prompt_checkpoint, + actor, + session_id, + routing.application_id, + tool_result_text_indices(structured_messages, texts), + ) if output_texts is None: return inputs return {**inputs, "texts": output_texts} - async def _file_part_to_data(self, part: FilePart) -> dict[str, Any]: + async def _file_part_to_data(self, part: FilePart) -> Mapping[str, Any]: extension = mimetypes.guess_extension(part.mime_hint) if part.mime_hint else None name = part.name or (f"file{extension}" if extension else "file") content = ( @@ -397,17 +420,20 @@ class OvalixGuardrail(CustomGuardrail): async def _check_texts( self, - texts: list[str], + texts: Sequence[str], checkpoint_id: str, actor: str, session_id: str, application_id: str, + skip_indices: frozenset[int], ) -> list[str] | None: output = list(texts) changed = False count = len(texts) for reversed_index in range(count): original_index = count - 1 - reversed_index + if original_index in skip_indices: + continue is_newest = reversed_index == 0 content = texts[original_index] try: @@ -435,7 +461,7 @@ class OvalixGuardrail(CustomGuardrail): output[original_index] = corrected return output if changed else None - def _get_session_id_for_application(self, data: dict, application_id: str | None) -> str: + def _get_session_id_for_application(self, data: Mapping[str, Any], application_id: str | None) -> str: actor_hash = self._get_tracker_actor_id(data) today = datetime.datetime.now(datetime.timezone.utc).strftime("%Y-%m-%d") return f"{actor_hash}_{today}_{application_id}" @@ -448,23 +474,29 @@ class OvalixGuardrail(CustomGuardrail): should_wrap_with_default_message=False, ) - def _get_trackers_corrected_message(self, resp: dict) -> str | None: + def _get_trackers_corrected_message(self, resp: Mapping[str, Any]) -> str | None: """Extract corrected/blocking message content from Tracker checkpoint response.""" modified = resp.get("modified_data") if isinstance(modified, dict) and "content" in modified: return modified["content"] return None - def _get_key_alias(self, request_data: dict) -> str | None: - metadata = {**(request_data.get("metadata") or {}), **(request_data.get("litellm_metadata") or {})} - return metadata.get("user_api_key_alias") or metadata.get("user_api_key_key_alias") + def _get_key_alias(self, request_data: Mapping[str, Any]) -> str | None: + litellm_metadata = request_data.get("litellm_metadata") or _NO_METADATA + metadata = request_data.get("metadata") or _NO_METADATA + + def _merged(key: str) -> object: + return litellm_metadata.get(key) if key in litellm_metadata else metadata.get(key) + + alias = _merged("user_api_key_alias") or _merged("user_api_key_key_alias") + return alias if isinstance(alias, str) else None async def _get_app_name_regex(self) -> re.Pattern[str]: if self._app_name_regex is not None: return self._app_name_regex - url = f"{self._tracker_api_base}/tracking/custom_application/litellm_app_name_regex" + url = f"{self._tracker_api_base}/tracking/litellm/app_name_regex" try: - response = await self._async_handler.get(url, headers=dict(self._tracker_headers)) + response = await self._async_handler.get(url, headers=self._tracker_headers) response.raise_for_status() compiled = re.compile(response.json()["regex"]) except Exception as e: @@ -512,7 +544,6 @@ class OvalixGuardrail(CustomGuardrail): verbose_proxy_logger.warning( "Ovalix guardrail passing the call through unguarded (fail_if_no_application=false): %s", reason ) - return None def _routing_error(self, error: Exception) -> GuardrailRaisedException: verbose_proxy_logger.exception("Ovalix routing resolution failed: %s", error) @@ -522,7 +553,7 @@ class OvalixGuardrail(CustomGuardrail): should_wrap_with_default_message=False, ) - async def _resolve_routing(self, request_data: dict) -> ResolvedRouting | None: + async def _resolve_routing(self, request_data: Mapping[str, Any]) -> ResolvedRouting | None: if self._application_id: return ResolvedRouting( self._application_id, @@ -550,10 +581,10 @@ class OvalixGuardrail(CustomGuardrail): return routing async def _resolve_via_tracker(self, application_name: str) -> ResolvedRouting | None: - url = f"{self._tracker_api_base}/tracking/custom_application/resolve_litellm_application" + url = f"{self._tracker_api_base}/tracking/litellm/resolve_application" try: response = await self._async_handler.post( - url, headers=dict(self._tracker_headers), json={"application_name": application_name} + url, headers=self._tracker_headers, json={"application_name": application_name} ) response.raise_for_status() body = response.json() diff --git a/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix_extraction.py b/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix_extraction.py index f3f91a85a56..c614daf6182 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix_extraction.py +++ b/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix_extraction.py @@ -2,12 +2,14 @@ import base64 import json import posixpath import re -from collections.abc import Callable -from typing import Any, NamedTuple +from collections.abc import Callable, Iterator, Mapping, Sequence +from types import MappingProxyType +from typing import NamedTuple from urllib.parse import unquote, urlparse _TOOL_NAME_MAX_LENGTH = 100 _DEFAULT_TOOL_RESULT_NAME = "tool_result" +_NO_TOOL_INPUT: Mapping[str, object] = MappingProxyType({}) _DATA_URL_RE = re.compile(r"^data:(?P[^;,]+)?(?P(?:;[^;,]+)*?)(?P;base64)?,", re.IGNORECASE) _URLSAFE_TO_STANDARD_B64 = str.maketrans("-_", "+/") @@ -63,7 +65,7 @@ def _name_from_url(url: str) -> str | None: return None -def _part_from_file_block(block: dict[str, Any], size_limit: int | None, message_index: int) -> FilePart | None: +def _part_from_file_block(block: Mapping[str, object], size_limit: int | None, message_index: int) -> FilePart | None: file_obj = block.get("file") if not isinstance(file_obj, dict): return None @@ -77,7 +79,9 @@ def _part_from_file_block(block: dict[str, Any], size_limit: int | None, message return FilePart(name, None, None, False, False, message_index) -def _part_from_image_url_block(block: dict[str, Any], size_limit: int | None, message_index: int) -> FilePart | None: +def _part_from_image_url_block( + block: Mapping[str, object], size_limit: int | None, message_index: int +) -> FilePart | None: image_url = block.get("image_url") url = image_url.get("url") if isinstance(image_url, dict) else image_url if not isinstance(url, str) or not url: @@ -91,7 +95,9 @@ def _part_from_image_url_block(block: dict[str, Any], size_limit: int | None, me return FilePart(_name_from_url(url), None, None, False, False, message_index) -def _part_from_input_file_block(block: dict[str, Any], size_limit: int | None, message_index: int) -> FilePart | None: +def _part_from_input_file_block( + block: Mapping[str, object], size_limit: int | None, message_index: int +) -> FilePart | None: name = block.get("filename") or block.get("file_id") or None file_data = block.get("file_data") if isinstance(file_data, str) and file_data: @@ -105,7 +111,9 @@ def _part_from_input_file_block(block: dict[str, Any], size_limit: int | None, m return FilePart(name, None, None, False, False, message_index) -def _part_from_input_audio_block(block: dict[str, Any], size_limit: int | None, message_index: int) -> FilePart | None: +def _part_from_input_audio_block( + block: Mapping[str, object], size_limit: int | None, message_index: int +) -> FilePart | None: audio = block.get("input_audio") if not isinstance(audio, dict): return None @@ -119,149 +127,215 @@ def _part_from_input_audio_block(block: dict[str, Any], size_limit: int | None, return FilePart(name, data, None, True, oversize, message_index) -_BLOCK_PARSERS: dict[str, Callable[[dict[str, Any], int | None, int], FilePart | None]] = { - "file": _part_from_file_block, - "image_url": _part_from_image_url_block, - "input_image": _part_from_image_url_block, - "input_file": _part_from_input_file_block, - "input_audio": _part_from_input_audio_block, -} +_BLOCK_PARSERS: Mapping[str, Callable[[Mapping[str, object], int | None, int], FilePart | None]] = MappingProxyType( + { + "file": _part_from_file_block, + "image_url": _part_from_image_url_block, + "input_image": _part_from_image_url_block, + "input_file": _part_from_input_file_block, + "input_audio": _part_from_input_audio_block, + } +) + + +def _file_parts_of_message( + message: Mapping[str, object], size_limit: int | None, message_index: int +) -> Iterator[FilePart]: + content = message.get("content") + if not isinstance(content, list): + return + for block in content: + if not isinstance(block, Mapping): + continue + block_type = block.get("type") + if not isinstance(block_type, str): + continue + parser = _BLOCK_PARSERS.get(block_type) + if parser is None: + continue + try: + part = parser(block, size_limit, message_index) + except (TypeError, ValueError, AttributeError, KeyError): + continue + if part is not None and (part.inline or part.name): + yield part def extract_file_parts_from_messages( - structured_messages: list[dict[str, Any]] | None, size_limit: int | None = None -) -> list[FilePart]: - parts: list[FilePart] = [] - for message_index, message in enumerate(structured_messages or []): - if not isinstance(message, dict): - continue - content = message.get("content") - if not isinstance(content, list): - continue - for block in content: - if not isinstance(block, dict): - continue - block_type = block.get("type") - if not isinstance(block_type, str): - continue - parser = _BLOCK_PARSERS.get(block_type) - if parser is None: - continue - try: - part = parser(block, size_limit, message_index) - except (TypeError, ValueError, AttributeError, KeyError): - continue - if part is not None and (part.inline or part.name): - parts.append(part) - return parts + structured_messages: Sequence[Mapping[str, object]] | None, size_limit: int | None = None +) -> tuple[FilePart, ...]: + return tuple( + part + for message_index, message in enumerate(structured_messages or ()) + if isinstance(message, Mapping) + for part in _file_parts_of_message(message, size_limit, message_index) + ) -def extract_file_parts_from_images(images: list[str] | None, size_limit: int | None = None) -> list[FilePart]: - parts: list[FilePart] = [] - for index, value in enumerate(images or []): - if not isinstance(value, str) or not value: - continue - if value.startswith(("http://", "https://")): - name = _name_from_url(value) - if name: - parts.append(FilePart(name, None, None, False, False, index)) - continue - mime_hint, payload = _split_data_url(value) - data, oversize = _decode_base64_with_limit(payload, size_limit) if payload else (None, False) - if data is not None or oversize: - parts.append(FilePart(None, data, mime_hint, True, oversize, index)) - return parts +def _file_part_of_image(value: str, size_limit: int | None, index: int) -> FilePart | None: + if value.startswith(("http://", "https://")): + name = _name_from_url(value) + return FilePart(name, None, None, False, False, index) if name else None + mime_hint, payload = _split_data_url(value) + data, oversize = _decode_base64_with_limit(payload, size_limit) if payload else (None, False) + if data is None and not oversize: + return None + return FilePart(None, data, mime_hint, True, oversize, index) -def make_tool_data(name: str, content: str | None, tool_input: dict[str, Any] | None = None) -> dict[str, Any]: +def extract_file_parts_from_images(images: Sequence[str] | None, size_limit: int | None = None) -> tuple[FilePart, ...]: + candidates = ( + _file_part_of_image(value, size_limit, index) + for index, value in enumerate(images or ()) + if isinstance(value, str) and value + ) + return tuple(part for part in candidates if part is not None) + + +def make_tool_data( + name: str, content: str | None, tool_input: Mapping[str, object] | None = None +) -> Mapping[str, object]: action_name = str(name) if str(name).strip() else _DEFAULT_TOOL_RESULT_NAME tool_name = action_name[:_TOOL_NAME_MAX_LENGTH] if not tool_name.strip(): tool_name = _DEFAULT_TOOL_RESULT_NAME - return {"content": content, "tool_name": tool_name, "action_name": action_name, "tool_input": tool_input or {}} + return { + "content": content, + "tool_name": tool_name, + "action_name": action_name, + "tool_input": dict(tool_input or ()), + } -def _tool_call_field(tool_call: Any, key: str) -> Any: +def _tool_call_field(tool_call: object, key: str) -> object: if isinstance(tool_call, dict): return tool_call.get(key) return getattr(tool_call, key, None) -def tool_call_to_tool_data(tool_call: Any) -> dict[str, Any] | None: +def _json_or_str(value: object) -> str: + try: + return json.dumps(value) + except (TypeError, ValueError): + return str(value) + + +def _parsed_tool_input(raw_arguments: str) -> Mapping[str, object]: + if not raw_arguments: + return _NO_TOOL_INPUT + try: + parsed = json.loads(raw_arguments) + except (ValueError, TypeError): + return _NO_TOOL_INPUT + return parsed if isinstance(parsed, dict) else _NO_TOOL_INPUT + + +def _tool_content_and_input(raw_arguments: object) -> tuple[str, Mapping[str, object]]: + if isinstance(raw_arguments, str): + return raw_arguments, _parsed_tool_input(raw_arguments) + if raw_arguments is None: + return "", _NO_TOOL_INPUT + if isinstance(raw_arguments, dict): + return _json_or_str(raw_arguments), raw_arguments + return _json_or_str(raw_arguments), _NO_TOOL_INPUT + + +def tool_call_to_tool_data(tool_call: object) -> Mapping[str, object] | None: function = _tool_call_field(tool_call, "function") name = function.get("name") if isinstance(function, dict) else getattr(function, "name", None) if not name or not str(name).strip(): return None raw_arguments = function.get("arguments") if isinstance(function, dict) else getattr(function, "arguments", None) - tool_input: dict[str, Any] = {} - if isinstance(raw_arguments, str): - content = raw_arguments - if content: - try: - parsed = json.loads(content) - if isinstance(parsed, dict): - tool_input = parsed - except (ValueError, TypeError): - tool_input = {} - elif raw_arguments is None: - content = "" - elif isinstance(raw_arguments, dict): - try: - content = json.dumps(raw_arguments) - except (TypeError, ValueError): - content = str(raw_arguments) - tool_input = raw_arguments - else: - try: - content = json.dumps(raw_arguments) - except (TypeError, ValueError): - content = str(raw_arguments) + content, tool_input = _tool_content_and_input(raw_arguments) return make_tool_data(name, content, tool_input) -def _extract_tool_content(content: Any) -> str | None: +def _tool_content_blocks(content: Sequence[object]) -> Iterator[str]: + for block in content: + if isinstance(block, Mapping): + text = block.get("text") + if isinstance(text, str) and text: + yield text + elif isinstance(block, str): + yield block + + +def _extract_tool_content(content: object) -> str | None: if isinstance(content, list): - parts: list[str] = [] - for block in content: - if isinstance(block, dict): - text = block.get("text") - if isinstance(text, str) and text: - parts.append(text) - elif isinstance(block, str): - parts.append(block) - content = "\n".join(parts) + content = "\n".join(_tool_content_blocks(content)) elif isinstance(content, dict): - try: - content = json.dumps(content) - except (TypeError, ValueError): - content = str(content) + content = _json_or_str(content) if not isinstance(content, str) or not content.strip(): return None return content -def extract_tool_results(structured_messages: list[dict[str, Any]] | None) -> list[tuple[str, str, str | None]]: - id_to_name: dict[str, str] = {} - results: list[tuple[str, str, str | None]] = [] - for message in structured_messages or []: - if not isinstance(message, dict): +def _declared_names_for_call(message: Mapping[str, object], call_id: str) -> Iterator[str]: + for tool_call in message.get("tool_calls") or (): + if not isinstance(tool_call, Mapping) or tool_call.get("id") != call_id: continue - role = message.get("role") - if role == "assistant": - for tool_call in message.get("tool_calls") or []: - if not isinstance(tool_call, dict): - continue - call_id = tool_call.get("id") - function = tool_call.get("function") - name = function.get("name") if isinstance(function, dict) else None - if isinstance(call_id, str) and call_id and name and str(name).strip(): - id_to_name[call_id] = name - elif role == "tool": + function = tool_call.get("function") + name = function.get("name") if isinstance(function, Mapping) else None + if name and str(name).strip(): + yield name + + +def _resolve_tool_name(messages: Sequence[Mapping[str, object]], tool_index: int, tool_call_id: object) -> str: + if not isinstance(tool_call_id, str) or not tool_call_id: + return _DEFAULT_TOOL_RESULT_NAME + declared = tuple( + name + for message in messages[:tool_index] + if isinstance(message, Mapping) and message.get("role") == "assistant" + for name in _declared_names_for_call(message, tool_call_id) + ) + return declared[-1] if declared else _DEFAULT_TOOL_RESULT_NAME + + +def extract_tool_results( + structured_messages: Sequence[Mapping[str, object]] | None, +) -> tuple[tuple[str, str, str | None], ...]: + messages = tuple(structured_messages or ()) + + def _results() -> Iterator[tuple[str, str, str | None]]: + for index, message in enumerate(messages): + if not isinstance(message, Mapping) or message.get("role") != "tool": + continue content = _extract_tool_content(message.get("content")) if content is None: continue tool_call_id = message.get("tool_call_id") - resolved_name = id_to_name.get(tool_call_id) if isinstance(tool_call_id, str) else None - name = resolved_name or _DEFAULT_TOOL_RESULT_NAME - results.append((name, content, tool_call_id)) - return results + yield _resolve_tool_name(messages, index, tool_call_id), content, tool_call_id + + return tuple(_results()) + + +def _message_text_origins(structured_messages: Sequence[Mapping[str, object]] | None) -> Iterator[tuple[str, bool]]: + for message in structured_messages or (): + if not isinstance(message, Mapping): + continue + content = message.get("content") + from_tool_result = message.get("role") == "tool" and _extract_tool_content(content) is not None + if isinstance(content, str): + yield content, from_tool_result + elif isinstance(content, list): + for block in content: + if isinstance(block, Mapping) and block.get("text") is not None: + yield block["text"], from_tool_result + + +def tool_result_text_indices( + structured_messages: Sequence[Mapping[str, object]] | None, texts: Sequence[str] +) -> frozenset[int]: + """Positions in ``texts`` that hold content already submitted under the TOOL policy. + + The chat-completions guardrail flow builds ``texts`` and ``structured_messages`` from the + same message list, so tool-role content lands in both and would otherwise be checked twice. + Other surfaces (e.g. Anthropic messages) build ``texts`` from a differently shaped payload, + so the mapping is only trusted when replaying it reproduces ``texts`` exactly; anything else + falls back to checking every text. + """ + origins = tuple(_message_text_origins(structured_messages)) + if tuple(text for text, _ in origins) != tuple(texts): + return frozenset() + return frozenset(index for index, (_, from_tool_result) in enumerate(origins) if from_tool_result) 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 42fd6c19ad0..83f01091038 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_ovalix.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_ovalix.py @@ -5,6 +5,7 @@ with mocked Tracker service responses (allow, anonymize, block). import base64 import gzip +import json as json_lib import os from typing import Any, List from unittest.mock import AsyncMock, MagicMock, patch @@ -770,9 +771,9 @@ async def test_discovery_extracts_name_and_resolves(): mock_get, mock_post = _mock_handler(g) routing = await g._resolve_routing(_alias_request_data("[Weather App] prod")) assert routing.application_id == "app-9" - assert mock_post.call_args.args[0].endswith("/tracking/custom_application/resolve_litellm_application") + assert mock_post.call_args.args[0].endswith("/tracking/litellm/resolve_application") assert mock_post.call_args.kwargs["json"] == {"application_name": "Weather App"} - assert mock_get.call_args.args[0].endswith("/tracking/custom_application/litellm_app_name_regex") + assert mock_get.call_args.args[0].endswith("/tracking/litellm/app_name_regex") @pytest.mark.asyncio @@ -946,6 +947,23 @@ async def test_response_side_file_uses_file_checkpoint(): assert seen["last"]["checkpoint_id"] == "file-1" +@pytest.mark.asyncio +async def test_file_checkpoint_call_routes_to_litellm_file_endpoint(): + g = _static_guardrail() + seen = {} + + async def _post(url, headers=None, json=None): + seen["url"] = url + 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._call_checkpoint("FILE", {"name": "f.txt", "content": "x"}, "file-1", "a", "s", "app-1") + assert seen["url"] == "https://t/tracking/litellm/file_checkpoint" + + @pytest.mark.asyncio async def test_tool_call_block_raises(): g = _static_guardrail() @@ -1072,8 +1090,70 @@ async def test_empty_user_sends_empty_actor_matching_reference(): assert seen["last"]["actor"] == "" +def _recording_post(mapping_fn=lambda body: _ALLOW): + calls = [] + + async def _post(url, headers=None, json=None): + calls.append((json["data_type"], json["data"].get("content"))) + r = MagicMock() + r.json.return_value = mapping_fn(json) + r.raise_for_status = MagicMock() + return r + + return _post, calls + + @pytest.mark.asyncio -async def test_text_equal_to_tool_result_is_still_inspected(): +async def test_every_checkpoint_payload_is_json_serializable(): + """httpx json-encodes the checkpoint body, so a non-dict mapping in it would 500 at runtime.""" + g = _static_guardrail() + data_url = "data:text/plain;base64," + base64.b64encode(b"secret").decode() + inputs = GenericGuardrailAPIInputs( + texts=["hello", "sunny"], + tool_calls=[{"id": "c1", "type": "function", "function": {"name": "noop", "arguments": None}}], + structured_messages=[ + {"role": "user", "content": [{"type": "file", "file": {"filename": "s.txt", "file_data": data_url}}]}, + {"role": "assistant", "tool_calls": [{"id": "c2", "function": {"name": "get_weather"}}]}, + {"role": "tool", "tool_call_id": "c2", "content": "sunny"}, + ], + ) + encoded = [] + + async def _post(url, headers=None, json=None): + encoded.append(json_lib.dumps(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.apply_guardrail(inputs=inputs, request_data={}, input_type="request", logging_obj=None) + payloads = [json_lib.loads(body) for body in encoded] + assert sorted(p["data_type"] for p in payloads) == ["FILE", "TEXT", "TEXT", "TOOL", "TOOL"] + assert all(p["data"]["tool_input"] == {} for p in payloads if p["data_type"] == "TOOL") + + +@pytest.mark.asyncio +async def test_tool_result_checked_under_tool_policy_only_not_again_as_text(): + g = _static_guardrail() + inputs = GenericGuardrailAPIInputs( + texts=["what is the weather", "sunny"], + structured_messages=[ + {"role": "user", "content": "what is the weather"}, + {"role": "assistant", "tool_calls": [{"id": "c1", "function": {"name": "get_weather"}}]}, + {"role": "tool", "tool_call_id": "c1", "content": "sunny"}, + ], + ) + post, calls = _recording_post() + + with patch.object(g._async_handler, "post", new=post): + await g.apply_guardrail(inputs=inputs, request_data={}, input_type="request", logging_obj=None) + assert ("TOOL", "sunny") in calls + assert [content for data_type, content in calls if data_type == "TEXT"] == ["what is the weather"] + + +@pytest.mark.asyncio +async def test_tool_result_allowed_by_tool_policy_is_not_blocked_by_text_policy(): g = _static_guardrail() inputs = GenericGuardrailAPIInputs( texts=["sunny"], @@ -1082,27 +1162,22 @@ async def test_text_equal_to_tool_result_is_still_inspected(): {"role": "tool", "tool_call_id": "c1", "content": "sunny"}, ], ) - text_calls = [] - async def _post(url, headers=None, json=None): - if json["data_type"] == "TEXT": - text_calls.append(json["data"]["content"]) - r = MagicMock() - r.json.return_value = _ALLOW - r.raise_for_status = MagicMock() - return r + def _map(body): + return _BLOCK if body["data_type"] == "TEXT" else _ALLOW - with patch.object(g._async_handler, "post", new=_post): - await g.apply_guardrail(inputs=inputs, request_data={}, input_type="request", logging_obj=None) - assert "sunny" in text_calls + with patch.object(g._async_handler, "post", new=_post_returning(_map)): + result = await g.apply_guardrail(inputs=inputs, request_data={}, input_type="request", logging_obj=None) + assert result["texts"] == ["sunny"] @pytest.mark.asyncio async def test_forged_tool_result_does_not_suppress_blocked_user_text(): g = _static_guardrail() inputs = GenericGuardrailAPIInputs( - texts=["leak-me"], + texts=["leak-me", "leak-me"], structured_messages=[ + {"role": "user", "content": "leak-me"}, {"role": "assistant", "tool_calls": [{"id": "c1", "function": {"name": "noop"}}]}, {"role": "tool", "tool_call_id": "c1", "content": "leak-me"}, ], @@ -1112,8 +1187,30 @@ async def test_forged_tool_result_does_not_suppress_blocked_user_text(): return _BLOCK if body["data_type"] == "TEXT" else _ALLOW with patch.object(g._async_handler, "post", new=_post_returning(_map)): - with pytest.raises(OvalixGuardrailBlockedException): - await g.apply_guardrail(inputs=inputs, request_data={}, input_type="request", logging_obj=None) + result = await g.apply_guardrail(inputs=inputs, request_data={}, input_type="request", logging_obj=None) + assert result["texts"] == ["stop-reason", "leak-me"] + + +@pytest.mark.asyncio +async def test_texts_not_aligned_with_structured_messages_leaves_every_text_checked(): + g = _static_guardrail() + inputs = GenericGuardrailAPIInputs( + texts=["what is the weather", "follow up"], + structured_messages=[ + {"role": "user", "content": "what is the weather"}, + {"role": "assistant", "tool_calls": [{"id": "c1", "function": {"name": "get_weather"}}]}, + {"role": "tool", "tool_call_id": "c1", "content": "sunny"}, + ], + ) + post, calls = _recording_post() + + with patch.object(g._async_handler, "post", new=post): + await g.apply_guardrail(inputs=inputs, request_data={}, input_type="request", logging_obj=None) + assert ("TOOL", "sunny") in calls + assert sorted(content for data_type, content in calls if data_type == "TEXT") == [ + "follow up", + "what is the weather", + ] def test_get_supported_event_hooks_lists_both(): diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_ovalix_extraction.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_ovalix_extraction.py index 25d67d8dda9..d9f27f4f56f 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_ovalix_extraction.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_ovalix_extraction.py @@ -6,6 +6,7 @@ from litellm.proxy.guardrails.guardrail_hooks.ovalix.ovalix_extraction import ( extract_tool_results, make_tool_data, tool_call_to_tool_data, + tool_result_text_indices, ) @@ -93,7 +94,7 @@ def test_make_tool_data_truncates_and_defaults_name(): def test_unhashable_block_type_skipped_without_raising(): msgs = [{"role": "user", "content": [{"type": ["file"], "file": {"filename": "a.txt"}}]}] parts = extract_file_parts_from_messages(msgs, size_limit=1000) - assert parts == [] + assert parts == () def test_unhashable_tool_call_id_skipped_without_raising(): @@ -102,7 +103,7 @@ def test_unhashable_tool_call_id_skipped_without_raising(): {"role": "tool", "tool_call_id": ["c1"], "content": "sunny"}, ] results = extract_tool_results(msgs) - assert results == [("tool_result", "sunny", ["c1"])] + assert results == (("tool_result", "sunny", ["c1"]),) def test_extract_tool_results_list_form_content(): @@ -152,7 +153,7 @@ def test_image_url_block_data_url_decoded_from_messages(): def test_image_url_block_non_string_url_skipped(): block = {"type": "image_url", "image_url": {"url": 123}} - assert extract_file_parts_from_messages(_msgs(block), size_limit=1000) == [] + assert extract_file_parts_from_messages(_msgs(block), size_limit=1000) == () def test_input_image_block_data_url_decoded(): @@ -163,7 +164,7 @@ def test_input_image_block_data_url_decoded(): def test_file_block_non_dict_file_skipped(): block = {"type": "file", "file": "not-a-dict"} - assert extract_file_parts_from_messages(_msgs(block), size_limit=1000) == [] + assert extract_file_parts_from_messages(_msgs(block), size_limit=1000) == () def test_file_block_reference_without_bytes_is_name_only(): @@ -216,7 +217,7 @@ def test_input_audio_block_undecodable_is_name_only(): def test_input_audio_block_non_dict_skipped(): block = {"type": "input_audio", "input_audio": "nope"} - assert extract_file_parts_from_messages(_msgs(block), size_limit=1000) == [] + assert extract_file_parts_from_messages(_msgs(block), size_limit=1000) == () def test_tool_call_dict_arguments_serialized_and_parsed(): @@ -242,7 +243,7 @@ def test_tool_call_invalid_json_string_arguments_kept_as_content(): def test_tool_result_with_non_string_tool_call_id_uses_default_name(): msgs = [{"role": "tool", "tool_call_id": ["c1"], "content": "orphan"}] results = extract_tool_results(msgs) - assert results == [("tool_result", "orphan", ["c1"])] + assert results == (("tool_result", "orphan", ["c1"]),) def test_images_field_http_url_is_name_only_reference(): @@ -268,12 +269,12 @@ def test_messages_skip_non_dict_and_unknown_blocks(): def test_image_url_block_invalid_data_url_returns_no_part(): block = {"type": "image_url", "image_url": {"url": "data:image/png;base64,%%%invalid%%%"}} - assert extract_file_parts_from_messages(_msgs(block), size_limit=1000) == [] + assert extract_file_parts_from_messages(_msgs(block), size_limit=1000) == () def test_tool_message_with_empty_content_is_skipped(): msgs = [{"role": "tool", "tool_call_id": "c1", "content": " "}] - assert extract_tool_results(msgs) == [] + assert extract_tool_results(msgs) == () def test_extract_tool_results_skips_non_dict_messages_and_tool_calls(): @@ -282,7 +283,7 @@ def test_extract_tool_results_skips_non_dict_messages_and_tool_calls(): {"role": "assistant", "tool_calls": ["not-a-dict", {"id": "c1", "function": {"name": "f"}}]}, {"role": "tool", "tool_call_id": "c1", "content": "ok"}, ] - assert extract_tool_results(msgs) == [("f", "ok", "c1")] + assert extract_tool_results(msgs) == (("f", "ok", "c1"),) def test_malformed_data_url_yields_no_bytes(): @@ -293,11 +294,11 @@ def test_malformed_data_url_yields_no_bytes(): def test_input_audio_block_without_data_skipped(): block = {"type": "input_audio", "input_audio": {"format": "wav"}} - assert extract_file_parts_from_messages(_msgs(block), size_limit=1000) == [] + assert extract_file_parts_from_messages(_msgs(block), size_limit=1000) == () def test_images_field_non_string_entries_skipped(): - assert extract_file_parts_from_images([123, None, ""], size_limit=1000) == [] + assert extract_file_parts_from_images([123, None, ""], size_limit=1000) == () def test_make_tool_data_whitespace_after_truncation_defaults_name(): @@ -337,14 +338,14 @@ def test_file_oversize_detected_after_decode_when_estimate_passes(): def test_image_http_url_that_urlparse_rejects_is_dropped(): parts = extract_file_parts_from_images(["http://["], size_limit=1000) - assert parts == [] + assert parts == () def test_message_content_not_a_list_is_skipped(): parts = extract_file_parts_from_messages( [{"role": "user", "content": "just a plain string prompt"}], size_limit=1000 ) - assert parts == [] + assert parts == () def test_tool_call_dict_arguments_non_serializable_falls_back_to_str(): @@ -360,3 +361,43 @@ def test_tool_call_non_serializable_other_arguments_falls_back_to_str(): def test_tool_result_dict_content_non_serializable_falls_back_to_str(): results = extract_tool_results([{"role": "tool", "tool_call_id": "c1", "content": {"x": {1, 2}}}]) assert len(results) == 1 and isinstance(results[0][1], str) and results[0][1].strip() + + +def test_tool_result_text_indices_marks_only_tool_role_positions(): + messages = [ + {"role": "user", "content": "ask"}, + {"role": "assistant", "tool_calls": [{"id": "c1", "function": {"name": "f"}}]}, + {"role": "tool", "tool_call_id": "c1", "content": "result"}, + ] + assert tool_result_text_indices(messages, ["ask", "result"]) == frozenset({1}) + + +def test_tool_result_text_indices_covers_every_text_block_of_a_tool_message(): + messages = [ + {"role": "user", "content": "ask"}, + { + "role": "tool", + "tool_call_id": "c1", + "content": [{"type": "text", "text": "a"}, {"type": "text", "text": "b"}], + }, + ] + assert tool_result_text_indices(messages, ["ask", "a", "b"]) == frozenset({1, 2}) + + +def test_tool_result_text_indices_empty_when_texts_do_not_replay_messages(): + messages = [ + {"role": "user", "content": "ask"}, + {"role": "tool", "tool_call_id": "c1", "content": "result"}, + ] + assert tool_result_text_indices(messages, ["result"]) == frozenset() + assert tool_result_text_indices(messages, ["ask", "tampered"]) == frozenset() + + +def test_tool_result_text_indices_skips_blank_tool_content_never_submitted_as_tool(): + messages = [{"role": "tool", "tool_call_id": "c1", "content": " "}] + assert extract_tool_results(messages) == () + assert tool_result_text_indices(messages, [" "]) == frozenset() + + +def test_tool_result_text_indices_empty_without_structured_messages(): + assert tool_result_text_indices(None, ["ask"]) == frozenset() From eea475a65e4f4090f54e2938f0e89238758554b9 Mon Sep 17 00:00:00 2001 From: Shalom Jamil Date: Mon, 3 Aug 2026 16:03:24 +0300 Subject: [PATCH 10/16] feat(ovalix): call tracker /beta endpoints for discovery + file scanning, send x-api-key --- litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py | 7 ++++--- .../proxy/guardrails/guardrail_hooks/test_ovalix.py | 6 +++--- 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py b/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py index b22b8a41dae..7284e126d10 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py +++ b/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py @@ -168,6 +168,7 @@ class OvalixGuardrail(CustomGuardrail): MappingProxyType( { "Authorization": f"Bearer {self._tracker_api_key}", + "x-api-key": self._tracker_api_key or "", "Content-Type": "application/json", } ), @@ -248,7 +249,7 @@ class OvalixGuardrail(CustomGuardrail): raise ValueError("Ovalix: application_id or checkpoint_id not resolved") url = ( - f"{self._tracker_api_base}/tracking/litellm/file_checkpoint" + f"{self._tracker_api_base}/tracking/beta/file_checkpoint" if data_type == "FILE" else f"{self._tracker_api_base}/tracking/custom_application/checkpoint" ) @@ -494,7 +495,7 @@ class OvalixGuardrail(CustomGuardrail): async def _get_app_name_regex(self) -> re.Pattern[str]: if self._app_name_regex is not None: return self._app_name_regex - url = f"{self._tracker_api_base}/tracking/litellm/app_name_regex" + url = f"{self._tracker_api_base}/tracking/beta/app_name_regex" try: response = await self._async_handler.get(url, headers=self._tracker_headers) response.raise_for_status() @@ -581,7 +582,7 @@ class OvalixGuardrail(CustomGuardrail): return routing async def _resolve_via_tracker(self, application_name: str) -> ResolvedRouting | None: - url = f"{self._tracker_api_base}/tracking/litellm/resolve_application" + url = f"{self._tracker_api_base}/tracking/beta/resolve_application" try: response = await self._async_handler.post( url, headers=self._tracker_headers, json={"application_name": application_name} 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 83f01091038..8f47097e17b 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_ovalix.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_ovalix.py @@ -771,9 +771,9 @@ async def test_discovery_extracts_name_and_resolves(): mock_get, mock_post = _mock_handler(g) routing = await g._resolve_routing(_alias_request_data("[Weather App] prod")) assert routing.application_id == "app-9" - assert mock_post.call_args.args[0].endswith("/tracking/litellm/resolve_application") + assert mock_post.call_args.args[0].endswith("/tracking/beta/resolve_application") assert mock_post.call_args.kwargs["json"] == {"application_name": "Weather App"} - assert mock_get.call_args.args[0].endswith("/tracking/litellm/app_name_regex") + assert mock_get.call_args.args[0].endswith("/tracking/beta/app_name_regex") @pytest.mark.asyncio @@ -961,7 +961,7 @@ async def test_file_checkpoint_call_routes_to_litellm_file_endpoint(): with patch.object(g._async_handler, "post", new=_post): await g._call_checkpoint("FILE", {"name": "f.txt", "content": "x"}, "file-1", "a", "s", "app-1") - assert seen["url"] == "https://t/tracking/litellm/file_checkpoint" + assert seen["url"] == "https://t/tracking/beta/file_checkpoint" @pytest.mark.asyncio From af9436f44ae742813cb5c7e352f8ce55fd4dbe88 Mon Sep 17 00:00:00 2001 From: Shalom Jamil Date: Tue, 4 Aug 2026 09:54:11 +0300 Subject: [PATCH 11/16] CR comment fix --- .../guardrail_hooks/ovalix/ovalix.py | 30 +++++++- .../guardrails/guardrail_hooks/test_ovalix.py | 69 +++++++++++++++---- 2 files changed, 85 insertions(+), 14 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py b/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py index 7284e126d10..ed7b8fa3618 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py +++ b/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py @@ -77,6 +77,17 @@ class ResolvedRouting(NamedTuple): checkpoint_id_pre_file: str | None checkpoint_id_post_file: str | None + @property + def has_any_checkpoint(self) -> bool: + return any( + ( + self.checkpoint_id_pre, + self.checkpoint_id_post, + self.checkpoint_id_pre_file, + self.checkpoint_id_post_file, + ) + ) + def _coerce_bool(value: bool | str) -> bool: if isinstance(value, bool): @@ -348,12 +359,19 @@ class OvalixGuardrail(CustomGuardrail): file_checkpoint = ( routing.checkpoint_id_post_file if is_response else routing.checkpoint_id_pre_file ) or prompt_checkpoint - if not prompt_checkpoint: + if not routing.has_any_checkpoint: raise GuardrailRaisedException( guardrail_name=self.guardrail_name, - message="Ovalix guardrail error: no checkpoint resolved for input_type", + message=f"Ovalix guardrail error: application {routing.application_id} has no checkpoints configured", should_wrap_with_default_message=False, ) + if not file_checkpoint: + verbose_proxy_logger.debug( + "Ovalix guardrail: application %s has no %s checkpoint, leaving this direction uninspected", + routing.application_id, + input_type, + ) + return inputs structured_messages = inputs.get("structured_messages") or () file_parts = ( @@ -367,6 +385,14 @@ class OvalixGuardrail(CustomGuardrail): if file_block is not None: self._block_current_message(file_block) + if not prompt_checkpoint: + verbose_proxy_logger.debug( + "Ovalix guardrail: application %s has only a %s file checkpoint, skipping text and tool inspection", + routing.application_id, + input_type, + ) + return inputs + tool_call_items = tuple( ("TOOL", td) for td in (tool_call_to_tool_data(tc) for tc in (inputs.get("tool_calls") or ())) if td ) 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 8f47097e17b..92835546452 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_ovalix.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_ovalix.py @@ -736,6 +736,25 @@ def _alias_request_data(alias="[Weather App] prod"): return {"metadata": {"user_api_key_alias": alias, "user_api_key_user_email": "u@x.com"}} +def _routing_body(pre, post, pre_file, post_file, application_id="app-9"): + return { + "application_id": application_id, + "checkpoint_id_pre": pre, + "checkpoint_id_post": post, + "checkpoint_id_pre_file": pre_file, + "checkpoint_id_post_file": post_file, + } + + +def _checkpoint_bodies(mock_post): + """Bodies of the tracker checkpoint calls only, excluding regex/resolve traffic.""" + return [ + c.kwargs["json"] + for c in mock_post.call_args_list + if c.args and c.args[0].endswith(("/checkpoint", "/file_checkpoint")) + ] + + def _mock_handler(g, routing=None): get_resp = MagicMock() get_resp.json.return_value = {"regex": _REGEX} @@ -1251,23 +1270,49 @@ async def test_file_checkpoint_call_failure_fails_closed(): @pytest.mark.asyncio -async def test_discovery_resolved_without_prompt_checkpoint_raises(): +async def test_discovery_resolved_without_any_checkpoint_raises(): + """An application with no checkpoints in any direction is a tracker misconfiguration, so fail closed.""" g = _discovery_guardrail(enable_cache=False) - _mock_handler( - g, - routing={ - "application_id": "app-9", - "checkpoint_id_pre": None, - "checkpoint_id_post": None, - "checkpoint_id_pre_file": None, - "checkpoint_id_post_file": None, - }, - ) + _mock_handler(g, routing=_routing_body(None, None, None, None)) inputs = GenericGuardrailAPIInputs(texts=["hi"]) - with pytest.raises(GuardrailRaisedException): + with pytest.raises(GuardrailRaisedException) as exc: await g.apply_guardrail( inputs=inputs, request_data=_alias_request_data(), input_type="request", logging_obj=None ) + assert "no checkpoints configured" in str(exc.value.message) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("input_type, inspected", [("request", ["pre-9"]), ("response", [])]) +async def test_one_sided_discovery_inspects_configured_direction_only(input_type, inspected): + """Discovery registers both hooks speculatively, so a direction the app does not inspect must pass through.""" + g = _discovery_guardrail(enable_cache=False) + _, mock_post = _mock_handler(g, routing=_routing_body("pre-9", None, None, None)) + inputs = GenericGuardrailAPIInputs(texts=["hi"]) + result = await g.apply_guardrail( + inputs=inputs, request_data=_alias_request_data(), input_type=input_type, logging_obj=None + ) + assert result["texts"] == ["hi"] + assert [b["checkpoint_id"] for b in _checkpoint_bodies(mock_post)] == inspected + + +@pytest.mark.asyncio +async def test_file_only_checkpoint_inspects_files_and_skips_text(): + """A direction with just a file checkpoint still scans files; text and tools need a prompt checkpoint.""" + g = _discovery_guardrail(enable_cache=False) + _, mock_post = _mock_handler(g, routing=_routing_body(None, None, "pre-file-9", None)) + data_url = "data:text/plain;base64," + base64.b64encode(b"secret").decode() + inputs = GenericGuardrailAPIInputs( + texts=["hi"], + structured_messages=[ + {"role": "user", "content": [{"type": "file", "file": {"filename": "s.txt", "file_data": data_url}}]} + ], + ) + result = await g.apply_guardrail( + inputs=inputs, request_data=_alias_request_data(), input_type="request", logging_obj=None + ) + assert result["texts"] == ["hi"] + assert [(b["data_type"], b["checkpoint_id"]) for b in _checkpoint_bodies(mock_post)] == [("FILE", "pre-file-9")] def test_initialize_guardrail_wires_new_params(monkeypatch): From e9e1fc6bcbc639a21813c8bebeacc378f9aa62f3 Mon Sep 17 00:00:00 2001 From: Shalom Jamil Date: Tue, 4 Aug 2026 15:52:42 +0300 Subject: [PATCH 12/16] application name over application ID --- .../guardrail_hooks/ovalix/ovalix.py | 87 ++++++--- .../guardrails/guardrail_hooks/test_ovalix.py | 175 ++++++++++++++++-- 2 files changed, 227 insertions(+), 35 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py b/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py index ed7b8fa3618..597108df521 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py +++ b/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py @@ -89,6 +89,21 @@ class ResolvedRouting(NamedTuple): ) +class CheckpointTarget(NamedTuple): + """How a checkpoint call addresses the application. + + When application_name is set the call sends the name and the direction, and the tracker resolves + and evaluates in one request; resolution can create the application, and ids from a separate + resolve call may name one the tracker's config has not caught up with, which it reports as an + uninspected allow rather than an error. application_id is what the session id groups on, and is + what gets sent when the deployment pins the application in config and reads no alias. + """ + + application_id: str + input_type: str + application_name: str | None = None + + def _coerce_bool(value: bool | str) -> bool: if isinstance(value, bool): return value @@ -253,27 +268,35 @@ class OvalixGuardrail(CustomGuardrail): checkpoint_id: str, actor: str, session_id: str, - application_id: str, + target: CheckpointTarget, ) -> Mapping[str, Any]: - """Call the Ovalix Tracker checkpoint API and return the JSON response.""" - if not application_id or not checkpoint_id: + """Call the Ovalix Tracker checkpoint API and return the JSON response. + + Both routes live on the tracker's /beta litellm router, which accepts the api key this + guardrail already sends. The name and id routing forms are mutually exclusive, so exactly one + reaches the wire; the name form also sends the direction, which is what the tracker selects + the pre or post checkpoint by. + """ + if not target.application_name and (not target.application_id or not checkpoint_id): raise ValueError("Ovalix: application_id or checkpoint_id not resolved") - url = ( - f"{self._tracker_api_base}/tracking/beta/file_checkpoint" - if data_type == "FILE" - else f"{self._tracker_api_base}/tracking/custom_application/checkpoint" + route = "file_checkpoint" if data_type == "FILE" else "checkpoint" + routing = ( + {"application_name": target.application_name, "input_type": target.input_type} + if target.application_name + else {"application_id": target.application_id, "checkpoint_id": checkpoint_id} ) payload = { - "application_id": application_id, - "checkpoint_id": checkpoint_id, "actor": actor, "session_id": session_id, "data_type": data_type, "data": data, "tool": "LiteLLM", + **routing, } - response = await self._async_handler.post(url, headers=self._tracker_headers, json=payload) + response = await self._async_handler.post( + f"{self._tracker_api_base}/tracking/beta/{route}", headers=self._tracker_headers, json=payload + ) response.raise_for_status() return response.json() @@ -287,11 +310,11 @@ class OvalixGuardrail(CustomGuardrail): checkpoint_id: str, actor: str, session_id: str, - application_id: str, + target: CheckpointTarget, escalation_reason: str, ) -> str | None: try: - resp = await self._call_checkpoint(data_type, data, checkpoint_id, actor, session_id, application_id) + resp = await self._call_checkpoint(data_type, data, checkpoint_id, actor, session_id, target) except Exception as e: verbose_proxy_logger.exception("Ovalix checkpoint call failed: %s", e) raise GuardrailRaisedException( @@ -312,12 +335,12 @@ class OvalixGuardrail(CustomGuardrail): checkpoint_id: str, actor: str, session_id: str, - application_id: str, + target: CheckpointTarget, escalation_reason: str, ) -> str | None: for data_type, data in items: reason = await self._block_reason_for_item( - data_type, data, checkpoint_id, actor, session_id, application_id, escalation_reason + data_type, data, checkpoint_id, actor, session_id, target, escalation_reason ) if reason is not None: return reason @@ -329,12 +352,12 @@ class OvalixGuardrail(CustomGuardrail): checkpoint_id: str, actor: str, session_id: str, - application_id: str, + target: CheckpointTarget, ) -> str | None: 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, application_id, _FILE_BLOCK_ESCALATION_REASON + "FILE", data, checkpoint_id, actor, session_id, target, _FILE_BLOCK_ESCALATION_REASON ) if reason is not None: return reason @@ -354,6 +377,11 @@ class OvalixGuardrail(CustomGuardrail): actor = self._get_actor(request_data) session_id = self._get_session_id_for_application(request_data, routing.application_id) is_response = input_type == "response" + target = CheckpointTarget( + application_id=routing.application_id, + input_type=input_type, + application_name=await self._checkpoint_routing_name(request_data), + ) prompt_checkpoint = routing.checkpoint_id_post if is_response else routing.checkpoint_id_pre file_checkpoint = ( @@ -379,9 +407,7 @@ class OvalixGuardrail(CustomGuardrail): if is_response else extract_file_parts_from_messages(structured_messages, size_limit=_DEFAULT_FILE_SIZE_LIMIT) ) - file_block = await self._check_files_for_block( - file_parts, file_checkpoint, actor, session_id, routing.application_id - ) + file_block = await self._check_files_for_block(file_parts, file_checkpoint, actor, session_id, target) if file_block is not None: self._block_current_message(file_block) @@ -401,7 +427,7 @@ class OvalixGuardrail(CustomGuardrail): prompt_checkpoint, actor, session_id, - routing.application_id, + target, _TOOL_BLOCK_ESCALATION_REASON, ) if tool_block is not None: @@ -414,7 +440,7 @@ class OvalixGuardrail(CustomGuardrail): prompt_checkpoint, actor, session_id, - routing.application_id, + target, _TOOL_RESULT_BLOCK_ESCALATION_REASON, ) if tool_result_block is not None: @@ -428,7 +454,7 @@ class OvalixGuardrail(CustomGuardrail): prompt_checkpoint, actor, session_id, - routing.application_id, + target, tool_result_text_indices(structured_messages, texts), ) if output_texts is None: @@ -451,7 +477,7 @@ class OvalixGuardrail(CustomGuardrail): checkpoint_id: str, actor: str, session_id: str, - application_id: str, + target: CheckpointTarget, skip_indices: frozenset[int], ) -> list[str] | None: output = list(texts) @@ -465,7 +491,7 @@ class OvalixGuardrail(CustomGuardrail): content = texts[original_index] try: resp = await self._call_checkpoint( - "TEXT", {"content": content}, checkpoint_id, actor, session_id, application_id + "TEXT", {"content": content}, checkpoint_id, actor, session_id, target ) except Exception as e: verbose_proxy_logger.exception("Ovalix checkpoint call failed: %s", e) @@ -580,6 +606,19 @@ class OvalixGuardrail(CustomGuardrail): should_wrap_with_default_message=False, ) + async def _checkpoint_routing_name(self, request_data: Mapping[str, Any]) -> str | None: + """The application name to route checkpoints by, or None to route by resolved ids. + + None when the deployment pins an application in config, or when no name can be read from the + api key alias. Only reached after _resolve_routing has already fetched and cached the regex. + """ + if self._application_id: + return None + alias = self._get_key_alias(request_data) + if not alias: + return None + return self._extract_application_name(alias, await self._get_app_name_regex()) + async def _resolve_routing(self, request_data: Mapping[str, Any]) -> ResolvedRouting | None: if self._application_id: return ResolvedRouting( 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 92835546452..3a96f3c4b18 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_ovalix.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_ovalix.py @@ -7,7 +7,7 @@ import base64 import gzip import json as json_lib import os -from typing import Any, List +import re from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -15,6 +15,7 @@ import pytest from litellm.exceptions import GuardrailRaisedException from litellm.proxy.guardrails.guardrail_hooks.ovalix.ovalix import ( + CheckpointTarget, OvalixGuardrail, OvalixGuardrailBlockedException, OvalixGuardrailMissingSecrets, @@ -276,13 +277,13 @@ class TestOvalixGuardrail: checkpoint_id="pre-1", actor="a1b2c3d4", session_id="session-1", - application_id="app-1", + target=CheckpointTarget("app-1", "request"), ) assert result == TRACKER_RESPONSE_ALLOW mock_post.assert_called_once() call_args = mock_post.call_args - assert call_args.args[0] == ("https://tracker.test/tracking/custom_application/checkpoint") + assert call_args.args[0] == ("https://tracker.test/tracking/beta/checkpoint") body = call_args.kwargs["json"] assert body["application_id"] == "app-1" assert body["checkpoint_id"] == "pre-1" @@ -746,15 +747,23 @@ def _routing_body(pre, post, pre_file, post_file, application_id="app-9"): } -def _checkpoint_bodies(mock_post): - """Bodies of the tracker checkpoint calls only, excluding regex/resolve traffic.""" +def _checkpoint_calls(mock_post): + """(body, url) of the tracker checkpoint calls only, excluding regex/resolve traffic.""" return [ - c.kwargs["json"] + (c.kwargs["json"], c.args[0]) for c in mock_post.call_args_list if c.args and c.args[0].endswith(("/checkpoint", "/file_checkpoint")) ] +def _checkpoint_bodies(mock_post): + return [body for body, _ in _checkpoint_calls(mock_post)] + + +def _route_of(url): + return url.rsplit("/", 1)[-1] + + def _mock_handler(g, routing=None): get_resp = MagicMock() get_resp.json.return_value = {"regex": _REGEX} @@ -979,7 +988,9 @@ async def test_file_checkpoint_call_routes_to_litellm_file_endpoint(): return r with patch.object(g._async_handler, "post", new=_post): - await g._call_checkpoint("FILE", {"name": "f.txt", "content": "x"}, "file-1", "a", "s", "app-1") + await g._call_checkpoint( + "FILE", {"name": "f.txt", "content": "x"}, "file-1", "a", "s", CheckpointTarget("app-1", "request") + ) assert seen["url"] == "https://t/tracking/beta/file_checkpoint" @@ -1251,7 +1262,7 @@ def test_enable_routing_cache_from_env_string(monkeypatch): async def test_call_checkpoint_requires_application_and_checkpoint(): g = _static_guardrail() with pytest.raises(ValueError): - await g._call_checkpoint("TEXT", {"content": "x"}, "", "actor", "sess", "app-1") + await g._call_checkpoint("TEXT", {"content": "x"}, "", "actor", "sess", CheckpointTarget("app-1", "request")) @pytest.mark.asyncio @@ -1283,7 +1294,7 @@ async def test_discovery_resolved_without_any_checkpoint_raises(): @pytest.mark.asyncio -@pytest.mark.parametrize("input_type, inspected", [("request", ["pre-9"]), ("response", [])]) +@pytest.mark.parametrize("input_type, inspected", [("request", ["request"]), ("response", [])]) async def test_one_sided_discovery_inspects_configured_direction_only(input_type, inspected): """Discovery registers both hooks speculatively, so a direction the app does not inspect must pass through.""" g = _discovery_guardrail(enable_cache=False) @@ -1293,7 +1304,7 @@ async def test_one_sided_discovery_inspects_configured_direction_only(input_type inputs=inputs, request_data=_alias_request_data(), input_type=input_type, logging_obj=None ) assert result["texts"] == ["hi"] - assert [b["checkpoint_id"] for b in _checkpoint_bodies(mock_post)] == inspected + assert [b["input_type"] for b in _checkpoint_bodies(mock_post)] == inspected @pytest.mark.asyncio @@ -1312,7 +1323,7 @@ async def test_file_only_checkpoint_inspects_files_and_skips_text(): inputs=inputs, request_data=_alias_request_data(), input_type="request", logging_obj=None ) assert result["texts"] == ["hi"] - assert [(b["data_type"], b["checkpoint_id"]) for b in _checkpoint_bodies(mock_post)] == [("FILE", "pre-file-9")] + assert [(b["data_type"], _route_of(c)) for b, c in _checkpoint_calls(mock_post)] == [("FILE", "file_checkpoint")] def test_initialize_guardrail_wires_new_params(monkeypatch): @@ -1592,3 +1603,145 @@ async def test_cached_404_still_raises_when_failing_closed(monkeypatch): with pytest.raises(GuardrailRaisedException): await g._resolve_routing(_alias_request_data("[Ghost App] prod")) assert mock_post.call_count == 1 + + +def _alias_guardrail(): + return OvalixGuardrail( + tracker_api_base="https://t", + tracker_api_key="k", + guardrail_name="o", + event_hook="pre_call", + default_on=True, + ) + + +def _capturing_post(response=None): + seen = {} + + async def _post(url, headers=None, json=None): + seen["url"] = url + seen["body"] = json + r = MagicMock() + r.json.return_value = response if response is not None else _ALLOW + r.raise_for_status = MagicMock() + return r + + return seen, _post + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "data_type,expected_route", + [("TEXT", "checkpoint"), ("TOOL", "checkpoint"), ("FILE", "file_checkpoint")], +) +async def test_call_checkpoint_by_name_sends_name_and_direction_not_ids(data_type, expected_route): + """Name routing must put application_name and input_type on the wire and no ids at all. + + The two forms are mutually exclusive server-side, so leaking an id alongside the name is rejected + """ + g = _alias_guardrail() + seen, post = _capturing_post() + + with patch.object(g._async_handler, "post", new=post): + await g._call_checkpoint( + data_type, + {"content": "x"}, + "unused-cp", + "actor", + "sess", + CheckpointTarget("app-9", "response", application_name="Weather App"), + ) + + assert seen["url"] == f"https://t/tracking/beta/{expected_route}" + assert seen["body"]["application_name"] == "Weather App" + assert seen["body"]["input_type"] == "response" + assert "application_id" not in seen["body"] + assert "checkpoint_id" not in seen["body"] + + +@pytest.mark.asyncio +async def test_call_checkpoint_by_ids_sends_ids_and_no_name_or_direction(): + g = _static_guardrail() + seen, post = _capturing_post() + + with patch.object(g._async_handler, "post", new=post): + await g._call_checkpoint( + "TEXT", {"content": "x"}, "pre-1", "actor", "sess", CheckpointTarget("app-1", "request") + ) + + assert seen["body"]["application_id"] == "app-1" + assert seen["body"]["checkpoint_id"] == "pre-1" + assert "application_name" not in seen["body"] + assert "input_type" not in seen["body"] + + +@pytest.mark.asyncio +async def test_call_checkpoint_by_name_does_not_require_a_checkpoint_id(): + """The tracker chooses the checkpoint under name routing, so an empty id must not be rejected.""" + g = _alias_guardrail() + seen, post = _capturing_post() + + with patch.object(g._async_handler, "post", new=post): + await g._call_checkpoint( + "TEXT", + {"content": "x"}, + "", + "actor", + "sess", + CheckpointTarget("", "request", application_name="Weather App"), + ) + + assert seen["body"]["application_name"] == "Weather App" + + +@pytest.mark.asyncio +async def test_checkpoint_routing_name_is_none_when_application_is_pinned(): + """A deployment that pins application_id reads no alias, so it must keep routing by ids.""" + g = _static_guardrail() + + assert await g._checkpoint_routing_name({"metadata": {"user_api_key_alias": "[Weather App] k"}}) is None + + +@pytest.mark.asyncio +async def test_checkpoint_routing_name_extracted_from_key_alias(): + g = _alias_guardrail() + with patch.object(g, "_get_app_name_regex", new=AsyncMock(return_value=re.compile(r"^\s*\[([^\]]+)\]"))): + name = await g._checkpoint_routing_name({"metadata": {"user_api_key_alias": "[Weather App] free text"}}) + + assert name == "Weather App" + + +@pytest.mark.asyncio +async def test_checkpoint_routing_name_is_none_without_an_alias(): + g = _alias_guardrail() + + assert await g._checkpoint_routing_name({"metadata": {}}) is None + + +@pytest.mark.asyncio +async def test_apply_guardrail_resolved_by_alias_routes_checkpoints_by_name(): + """End to end: an alias-resolved application sends its name per checkpoint, never the resolved id. + + This is what keeps a just-created application inspectable: the ids from resolution may name an + application the tracker's process-wide config has not picked up yet + """ + g = _alias_guardrail() + seen, post = _capturing_post() + + with ( + patch.object( + g, "_resolve_routing", new=AsyncMock(return_value=ResolvedRouting("app-9", "pre-9", "post-9", None, None)) + ), + patch.object(g, "_get_app_name_regex", new=AsyncMock(return_value=re.compile(r"^\s*\[([^\]]+)\]"))), + patch.object(g._async_handler, "post", new=post), + ): + await g.apply_guardrail( + GenericGuardrailAPIInputs(texts=["hello"]), + {"metadata": {"user_api_key_alias": "[Weather App] k", "user_api_key_user_email": "u@e.com"}}, + "request", + ) + + assert seen["url"] == "https://t/tracking/beta/checkpoint" + assert seen["body"]["application_name"] == "Weather App" + assert seen["body"]["input_type"] == "request" + assert "application_id" not in seen["body"] From 3641581823cf9894f625c3696dcd50d9dcfdb77c Mon Sep 17 00:00:00 2001 From: Shalom Jamil Date: Mon, 7 Sep 2026 10:11:12 +0300 Subject: [PATCH 13/16] fix(ovalix): scan tool calls carried on messages, tolerate optional regex groups Tool-call blocking only read the top-level tool_calls input. The Anthropic request path fills structured_messages but leaves that input empty, so calls made in prior assistant turns were never checkpointed. Merge both sources and dedupe on the tool payload, since the OpenAI path populates both and would otherwise scan every call twice. _extract_application_name called strip() on group(1) whenever the regex had any groups. An optional group that did not participate is None, which raised AttributeError instead of falling through to the no-application handling. --- .../guardrail_hooks/ovalix/ovalix.py | 15 ++++-- .../ovalix/ovalix_extraction.py | 24 +++++++++ .../guardrails/guardrail_hooks/test_ovalix.py | 50 ++++++++++++++++++- 3 files changed, 84 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py b/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py index ade2fba1156..5bf5e35f7c7 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py +++ b/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py @@ -34,9 +34,11 @@ from litellm.proxy.guardrails.guardrail_hooks.ovalix.ovalix_extraction import ( FilePart, extract_file_parts_from_images, extract_file_parts_from_messages, + extract_tool_calls_from_messages, extract_tool_results, make_tool_data, tool_call_to_tool_data, + tool_data_key, tool_result_text_indices, ) from litellm.types.guardrails import GuardrailEventHooks @@ -418,9 +420,14 @@ class OvalixGuardrail(CustomGuardrail): ) return inputs - tool_call_items: Final = tuple( - ("TOOL", td) for td in (tool_call_to_tool_data(tc) for tc in (inputs.get("tool_calls") or ())) if td + tool_calls: Final = ( + *(inputs.get("tool_calls") or ()), + *extract_tool_calls_from_messages(structured_messages), ) + unique_tool_data: Final = { + tool_data_key(data): data for data in (tool_call_to_tool_data(tc) for tc in tool_calls) if data + } + tool_call_items: Final = tuple(("TOOL", data) for data in unique_tool_data.values()) tool_block: Final = await self._check_items_block_only( tool_call_items, prompt_checkpoint, @@ -562,8 +569,8 @@ class OvalixGuardrail(CustomGuardrail): match: Final = regex.search(alias) if not match: return None - name: Final = (match.group(1) if match.groups() else match.group(0)).strip() - return name or None + captured: Final = (match.group(1) if match.groups() else match.group(0)) or "" + return captured.strip() or None def _routing_cache_get(self, name: str) -> tuple[bool, ResolvedRouting | None]: entry: Final = self._routing_cache.get(name) diff --git a/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix_extraction.py b/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix_extraction.py index 6e0d5f895c8..099af7f2bb8 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix_extraction.py +++ b/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix_extraction.py @@ -261,6 +261,30 @@ def tool_call_to_tool_data(tool_call: object) -> Mapping[str, object] | None: return make_tool_data(name, content, tool_input) +def _message_tool_calls(message: Mapping[str, object]) -> Sequence[object]: + tool_calls: Final = message.get("tool_calls") + return tool_calls if isinstance(tool_calls, list) else () + + +def extract_tool_calls_from_messages(structured_messages: Sequence[object] | None) -> tuple[object, ...]: + """Tool calls declared on the messages themselves. + + Surfaces such as the Anthropic request path populate ``structured_messages`` but leave the + top-level ``tool_calls`` input empty, so calls made in prior assistant turns are only visible here. + """ + return tuple( + tool_call + for message in structured_messages or () + if isinstance(message, Mapping) + for tool_call in _message_tool_calls(message) + ) + + +def tool_data_key(tool_data: Mapping[str, object]) -> str: + """Stable identity for a tool payload, so a call reached from two sources is only scanned once.""" + return json.dumps(tool_data, sort_keys=True, default=str) + + def _tool_content_blocks(content: Sequence[object]) -> Iterator[str]: for block in content: if isinstance(block, Mapping): 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 3a96f3c4b18..928bdefa477 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_ovalix.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_ovalix.py @@ -1159,7 +1159,7 @@ async def test_every_checkpoint_payload_is_json_serializable(): with patch.object(g._async_handler, "post", new=_post): await g.apply_guardrail(inputs=inputs, request_data={}, input_type="request", logging_obj=None) payloads = [json_lib.loads(body) for body in encoded] - assert sorted(p["data_type"] for p in payloads) == ["FILE", "TEXT", "TEXT", "TOOL", "TOOL"] + assert sorted(p["data_type"] for p in payloads) == ["FILE", "TEXT", "TEXT", "TOOL", "TOOL", "TOOL"] assert all(p["data"]["tool_input"] == {} for p in payloads if p["data_type"] == "TOOL") @@ -1745,3 +1745,51 @@ async def test_apply_guardrail_resolved_by_alias_routes_checkpoints_by_name(): assert seen["body"]["application_name"] == "Weather App" assert seen["body"]["input_type"] == "request" assert "application_id" not in seen["body"] + + +@pytest.mark.asyncio +async def test_message_tool_calls_are_scanned_without_top_level_tool_calls(): + g = _static_guardrail() + inputs = GenericGuardrailAPIInputs( + texts=[], + structured_messages=[ + { + "role": "assistant", + "tool_calls": [{"id": "c1", "type": "function", "function": {"name": "exfil", "arguments": "{}"}}], + } + ], + ) + with patch.object( + g._async_handler, "post", new=_post_returning(lambda body: _BLOCK if body["data_type"] == "TOOL" else _ALLOW) + ): + with pytest.raises(OvalixGuardrailBlockedException): + await g.apply_guardrail(inputs=inputs, request_data={}, input_type="request", logging_obj=None) + + +@pytest.mark.asyncio +async def test_tool_call_present_in_both_sources_is_scanned_once(): + g = _static_guardrail() + call = {"id": "c1", "type": "function", "function": {"name": "get_weather", "arguments": '{"city": "sf"}'}} + inputs = GenericGuardrailAPIInputs( + texts=[], + tool_calls=[call], + structured_messages=[{"role": "assistant", "tool_calls": [call]}], + ) + tool_bodies = [] + + async def _post(url, headers=None, json=None): + if json["data_type"] == "TOOL": + tool_bodies.append(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.apply_guardrail(inputs=inputs, request_data={}, input_type="request", logging_obj=None) + assert len(tool_bodies) == 1 + + +def test_extract_application_name_tolerates_unmatched_optional_group(): + g = _static_guardrail() + assert g._extract_application_name("app-", re.compile(r"app-(\w+)?")) is None From b5ce6464e69d753b2db4588f5fd53965238589d5 Mon Sep 17 00:00:00 2001 From: Shalom Jamil Date: Mon, 7 Sep 2026 10:50:08 +0300 Subject: [PATCH 14/16] chore(ovalix): satisfy the tightened lint gates on the current base The base moved on and lowered the LIT002 ceiling, so bring the guardrail back under it. _validate_config now returns the event hooks it resolves instead of appending to a list the caller owns, which lets __init__ take supported_event_hooks as a real parameter rather than digging it out of kwargs and rebuilding the dict. The checkpoint routing forms and the tool dedup map are read-only, so they are MappingProxyType now; the payload itself stays a plain dict because httpx has to json-encode it. Also set match= on the _call_checkpoint raises assertion for PT011. --- .../guardrail_hooks/ovalix/ovalix.py | 34 +++++++++++-------- .../guardrails/guardrail_hooks/test_ovalix.py | 2 +- 2 files changed, 20 insertions(+), 16 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py b/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py index ee041faf19e..6a0643d8e2d 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py +++ b/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py @@ -163,6 +163,7 @@ class OvalixGuardrail(CustomGuardrail): file_checkpoint_id: str | None = None, enable_routing_cache: bool | None = None, fail_if_no_application: bool | None = None, + supported_event_hooks: list[GuardrailEventHooks] | None = None, **kwargs: Any, ): self._tracker_api_base = tracker_api_base or os.environ.get("OVALIX_TRACKER_API_BASE") @@ -188,9 +189,7 @@ class OvalixGuardrail(CustomGuardrail): self._routing_cache: OrderedDict[str, tuple[float, ResolvedRouting | None]] = OrderedDict() self._app_name_regex: re.Pattern[str] | None = None - supported_event_hooks: Final[list[GuardrailEventHooks]] = list(kwargs.get("supported_event_hooks") or ()) - - self._validate_config(supported_event_hooks) + event_hooks: Final = self._validated_event_hooks(supported_event_hooks or ()) self._tracker_headers = dict( httpx.Headers( @@ -207,7 +206,7 @@ class OvalixGuardrail(CustomGuardrail): self._async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) - super().__init__(**{**kwargs, "supported_event_hooks": supported_event_hooks}) + super().__init__(supported_event_hooks=event_hooks, **kwargs) verbose_proxy_logger.debug( "Ovalix Guardrail initialized: tracker=%s, application_id=%s, pre_checkpoint_id=%s, post_checkpoint_id=%s", self._tracker_api_base, @@ -216,8 +215,8 @@ class OvalixGuardrail(CustomGuardrail): self._post_checkpoint_id, ) - def _validate_config(self, supported_event_hooks: list[GuardrailEventHooks]) -> None: - """Ensure required Tracker secrets are set; register the pre/post hooks this config can serve (both in discovery mode; only configured-checkpoint directions in static mode).""" + def _validated_event_hooks(self, requested: Sequence[GuardrailEventHooks]) -> list[GuardrailEventHooks]: + """Ensure required Tracker secrets are set; return the pre/post hooks this config can serve (both in discovery mode; only configured-checkpoint directions in static mode).""" errors: Final = tuple( message for present, message in ( @@ -236,10 +235,15 @@ class OvalixGuardrail(CustomGuardrail): supports_pre: Final = not self._application_id or bool(self._pre_checkpoint_id) supports_post: Final = not self._application_id or bool(self._post_checkpoint_id) - if supports_pre and GuardrailEventHooks.pre_call not in supported_event_hooks: - supported_event_hooks.append(GuardrailEventHooks.pre_call) - if supports_post and GuardrailEventHooks.post_call not in supported_event_hooks: - supported_event_hooks.append(GuardrailEventHooks.post_call) + auto_added: Final = tuple( + hook + for supported, hook in ( + (supports_pre, GuardrailEventHooks.pre_call), + (supports_post, GuardrailEventHooks.post_call), + ) + if supported and hook not in requested + ) + return [*requested, *auto_added] def _get_actor(self, data: Mapping[str, Any]) -> str: """Return a stable actor identifier from request metadata (e.g. user email or id).""" @@ -285,9 +289,9 @@ class OvalixGuardrail(CustomGuardrail): route: Final = "file_checkpoint" if data_type == "FILE" else "checkpoint" routing: Final = ( - {"application_name": target.application_name, "input_type": target.input_type} + MappingProxyType({"application_name": target.application_name, "input_type": target.input_type}) if target.application_name - else {"application_id": target.application_id, "checkpoint_id": checkpoint_id} + else MappingProxyType({"application_id": target.application_id, "checkpoint_id": checkpoint_id}) ) payload: Final = { "actor": actor, @@ -426,9 +430,9 @@ class OvalixGuardrail(CustomGuardrail): *(inputs.get("tool_calls") or ()), *extract_tool_calls_from_messages(structured_messages), ) - unique_tool_data: Final = { - tool_data_key(data): data for data in (tool_call_to_tool_data(tc) for tc in tool_calls) if data - } + unique_tool_data: Final = MappingProxyType( + {tool_data_key(data): data for data in (tool_call_to_tool_data(tc) for tc in tool_calls) if data} + ) tool_call_items: Final = tuple(("TOOL", data) for data in unique_tool_data.values()) tool_block: Final = await self._check_items_block_only( tool_call_items, 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 928bdefa477..b0d66c39286 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_ovalix.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_ovalix.py @@ -1261,7 +1261,7 @@ def test_enable_routing_cache_from_env_string(monkeypatch): @pytest.mark.asyncio async def test_call_checkpoint_requires_application_and_checkpoint(): g = _static_guardrail() - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="application_id or checkpoint_id not resolved"): await g._call_checkpoint("TEXT", {"content": "x"}, "", "actor", "sess", CheckpointTarget("app-1", "request")) From e4a72ef0bc3efb793e4a450a03f7e94c00965d43 Mon Sep 17 00:00:00 2001 From: Shalom Jamil Date: Mon, 7 Sep 2026 10:59:21 +0300 Subject: [PATCH 15/16] chore(ovalix): regenerate the OpenAPI snapshot and UI schema types The guardrail config model gained file_checkpoint_id, enable_routing_cache and fail_if_no_application, so the lazy OpenAPI snapshot and the dashboard's schema.d.ts both went stale. Regenerated with the documented commands. --- litellm/proxy/_lazy_openapi_snapshot.json | 36 +++++++++++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 15 ++++++++ 2 files changed, 51 insertions(+) diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 91a97ad6544..ad6af60d466 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -11490,6 +11490,18 @@ "description": "If True, will not raise an exception when the guardrail is blocked. Useful for OpenWebUI where exceptions can end the chat flow.", "title": "Disable Exception On Block" }, + "enable_routing_cache": { + "anyOf": [ + { + "type": "boolean" + }, + { + "type": "null" + } + ], + "description": "Cache discovery-mode routing resolution per api-key alias for 1 hour. Default on.", + "title": "Enable Routing Cache" + }, "end_session_after_n_fails": { "anyOf": [ { @@ -11542,6 +11554,18 @@ "description": "Header names to forward from the client request to the guardrail (e.g. x-request-id). Only these headers' values are sent; others may be omitted or sent as [present]. Used by generic_guardrail_api (similar to MCP extra_headers).", "title": "Extra Headers" }, + "fail_if_no_application": { + "anyOf": [ + { + "type": "boolean" + }, + { + "type": "null" + } + ], + "description": "Fail the call when no application is configured and none is discovered. Default on; set false to let such calls through unguarded instead.", + "title": "Fail If No Application" + }, "fail_on_error": { "anyOf": [ { @@ -11555,6 +11579,18 @@ "description": "Whether to fail the request if the guardrail encounters an error. Implemented by guardrail='model_armor', 'generic_guardrail_api' and 'crowdstrike_aidr'. True (default) raises the error. False logs a critical error and lets the request proceed, so only a valid guardrail response can block or modify it.", "title": "Fail On Error" }, + "file_checkpoint_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "description": "File-checkpoint ID for the Ovalix Tracker service (falls back to pre/post).", + "title": "File Checkpoint Id" + }, "grounding_check": { "anyOf": [ { diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index b1534c19670..a4e2987f676 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -30611,6 +30611,11 @@ export interface components { * @default false */ disable_exception_on_block: boolean | null; + /** + * Enable Routing Cache + * @description Cache discovery-mode routing resolution per api-key alias for 1 hour. Default on. + */ + enable_routing_cache?: boolean | null; /** * End Session After N Fails * @description For /v1/realtime sessions: automatically close the session after this many guardrail violations. @@ -30632,12 +30637,22 @@ export interface components { * @description Header names to forward from the client request to the guardrail (e.g. x-request-id). Only these headers' values are sent; others may be omitted or sent as [present]. Used by generic_guardrail_api (similar to MCP extra_headers). */ extra_headers?: string[] | null; + /** + * Fail If No Application + * @description Fail the call when no application is configured and none is discovered. Default on; set false to let such calls through unguarded instead. + */ + fail_if_no_application?: boolean | null; /** * Fail On Error * @description Whether to fail the request if the guardrail encounters an error. Implemented by guardrail='model_armor', 'generic_guardrail_api' and 'crowdstrike_aidr'. True (default) raises the error. False logs a critical error and lets the request proceed, so only a valid guardrail response can block or modify it. * @default true */ fail_on_error: boolean | null; + /** + * File Checkpoint Id + * @description File-checkpoint ID for the Ovalix Tracker service (falls back to pre/post). + */ + file_checkpoint_id?: string | null; /** * Grounding Check * @description Enable grounding verification to ensure output is grounded in provided context. From f9d9aba4214f1516558064c13365505232a3780a Mon Sep 17 00:00:00 2001 From: Shalom Jamil Date: Mon, 7 Sep 2026 12:24:06 +0300 Subject: [PATCH 16/16] refactor(ovalix): type the guardrail payloads as object instead of Any Every Mapping[str, Any] in the hook is now Mapping[str, object], matching what ovalix_extraction.py already used. object forces the reads to narrow, so the places that were quietly trusting request metadata and tracker responses now check the type they claim to return. Two of those were real crashes waiting to happen. A non-string user email in the metadata used to come back from _get_actor and blow up on .encode(), and a non-string modified_data.content was returned as the corrected message despite the str | None annotation. Both are covered by tests now. apply_guardrail keeps Mapping since the base class declares dict there, and **kwargs stays Any because it forwards to CustomGuardrail like the other hooks. --- .../guardrail_hooks/ovalix/ovalix.py | 65 +++++++++++-------- .../guardrails/guardrail_hooks/test_ovalix.py | 11 ++++ 2 files changed, 48 insertions(+), 28 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py b/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py index 6a0643d8e2d..fc37411c296 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py +++ b/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py @@ -57,7 +57,7 @@ _ROUTING_CACHE_TTL_SECONDS: Final = 3600 _ROUTING_CACHE_NEGATIVE_TTL_SECONDS: Final = 300 _ROUTING_CACHE_MAX_SIZE: Final = 1000 _DEFAULT_FILE_SIZE_LIMIT: Final = 64 * 1024 * 1024 -_NO_METADATA: Final[Mapping[str, Any]] = MappingProxyType({}) +_NO_METADATA: Final[Mapping[str, object]] = MappingProxyType({}) _FILE_BLOCK_ESCALATION_REASON: Final = ( "This message was blocked by Ovalix because file content anonymization isn't possible via LiteLLM" ) @@ -107,6 +107,14 @@ class CheckpointTarget(NamedTuple): application_name: str | None = None +def _mapping_or_empty(value: object) -> Mapping[str, object]: + return value if isinstance(value, Mapping) else _NO_METADATA + + +def _str_or_none(value: object) -> str | None: + return value if isinstance(value, str) and value else None + + def _coerce_bool(value: bool | str) -> bool: if isinstance(value, bool): return value @@ -245,16 +253,16 @@ class OvalixGuardrail(CustomGuardrail): ) return [*requested, *auto_added] - def _get_actor(self, data: Mapping[str, Any]) -> str: + def _get_actor(self, data: Mapping[str, object]) -> str: """Return a stable actor identifier from request metadata (e.g. user email or id).""" - metadata: Final = data.get("metadata") or data.get("litellm_metadata") or _NO_METADATA - if metadata.get("user_api_key_user_email"): - return metadata["user_api_key_user_email"] - if metadata.get("user_api_key_user_id"): - return metadata["user_api_key_user_id"] - return "" + metadata: Final = _mapping_or_empty(data.get("metadata") or data.get("litellm_metadata")) + return ( + _str_or_none(metadata.get("user_api_key_user_email")) + or _str_or_none(metadata.get("user_api_key_user_id")) + or "" + ) - def _get_tracker_actor_id(self, data: Mapping[str, Any]) -> str: + def _get_tracker_actor_id(self, data: Mapping[str, object]) -> str: """Normalize the actor string into a short, stable id for Tracker API payloads.""" # NOTE: this hash is purely for normalization — it collapses an arbitrary actor # string (email, user id, or empty) into a compact, fixed-length, consistent @@ -264,19 +272,19 @@ class OvalixGuardrail(CustomGuardrail): normalized_actor_id: Final = hashlib.sha256(actor_id).hexdigest()[:8] return normalized_actor_id - def _get_session_id(self, data: Mapping[str, Any]) -> str: + def _get_session_id(self, data: Mapping[str, object]) -> str: """Return a unique identifier for the chat/session (actor + date + application_id).""" return self._get_session_id_for_application(data, self._application_id) async def _call_checkpoint( self, data_type: str, - data: Mapping[str, Any], + data: Mapping[str, object], checkpoint_id: str, actor: str, session_id: str, target: CheckpointTarget, - ) -> Mapping[str, Any]: + ) -> Mapping[str, object]: """Call the Ovalix Tracker checkpoint API and return the JSON response. Both routes live on the tracker's /beta litellm router, which accepts the api key this @@ -307,13 +315,13 @@ class OvalixGuardrail(CustomGuardrail): response.raise_for_status() return response.json() - def _verdict(self, resp: Mapping[str, Any]) -> tuple[str, str | None]: - return (resp.get("action_type") or "").lower(), self._get_trackers_corrected_message(resp) + def _verdict(self, resp: Mapping[str, object]) -> tuple[str, str | None]: + return (_str_or_none(resp.get("action_type")) or "").lower(), self._get_trackers_corrected_message(resp) async def _block_reason_for_item( self, data_type: str, - data: Mapping[str, Any], + data: Mapping[str, object], checkpoint_id: str, actor: str, session_id: str, @@ -338,7 +346,7 @@ class OvalixGuardrail(CustomGuardrail): async def _check_items_block_only( self, - items: Sequence[tuple[str, Mapping[str, Any]]], + items: Sequence[tuple[str, Mapping[str, object]]], checkpoint_id: str, actor: str, session_id: str, @@ -374,7 +382,7 @@ class OvalixGuardrail(CustomGuardrail): async def apply_guardrail( self, inputs: GenericGuardrailAPIInputs, - request_data: Mapping[str, Any], + request_data: Mapping[str, object], input_type: Literal["request", "response"], logging_obj: "LiteLLMLoggingObj | None" = None, ) -> GenericGuardrailAPIInputs: @@ -473,7 +481,7 @@ class OvalixGuardrail(CustomGuardrail): return inputs return {**inputs, "texts": output_texts} - async def _file_part_to_data(self, part: FilePart) -> Mapping[str, Any]: + async def _file_part_to_data(self, part: FilePart) -> Mapping[str, object]: extension: Final = mimetypes.guess_extension(part.mime_hint) if part.mime_hint else None name: Final = part.name or (f"file{extension}" if extension else "file") content: Final = ( @@ -523,7 +531,7 @@ class OvalixGuardrail(CustomGuardrail): output[original_index] = corrected return output if tuple(output) != original else None - def _get_session_id_for_application(self, data: Mapping[str, Any], application_id: str | None) -> str: + def _get_session_id_for_application(self, data: Mapping[str, object], application_id: str | None) -> str: actor_hash: Final = self._get_tracker_actor_id(data) today: Final = datetime.datetime.now(datetime.timezone.utc).strftime("%Y-%m-%d") return f"{actor_hash}_{today}_{application_id}" @@ -536,16 +544,17 @@ class OvalixGuardrail(CustomGuardrail): should_wrap_with_default_message=False, ) - def _get_trackers_corrected_message(self, resp: Mapping[str, Any]) -> str | None: + def _get_trackers_corrected_message(self, resp: Mapping[str, object]) -> str | None: """Extract corrected/blocking message content from Tracker checkpoint response.""" modified: Final = resp.get("modified_data") - if isinstance(modified, dict) and "content" in modified: - return modified["content"] - return None + if not isinstance(modified, Mapping): + return None + content: Final = modified.get("content") + return content if isinstance(content, str) else None - def _get_key_alias(self, request_data: Mapping[str, Any]) -> str | None: - litellm_metadata: Final = request_data.get("litellm_metadata") or _NO_METADATA - metadata: Final = request_data.get("metadata") or _NO_METADATA + def _get_key_alias(self, request_data: Mapping[str, object]) -> str | None: + litellm_metadata: Final = _mapping_or_empty(request_data.get("litellm_metadata")) + metadata: Final = _mapping_or_empty(request_data.get("metadata")) def _merged(key: str) -> object: return litellm_metadata.get(key) if key in litellm_metadata else metadata.get(key) @@ -615,7 +624,7 @@ class OvalixGuardrail(CustomGuardrail): should_wrap_with_default_message=False, ) - async def _checkpoint_routing_name(self, request_data: Mapping[str, Any]) -> str | None: + async def _checkpoint_routing_name(self, request_data: Mapping[str, object]) -> str | None: """The application name to route checkpoints by, or None to route by resolved ids. None when the deployment pins an application in config, or when no name can be read from the @@ -628,7 +637,7 @@ class OvalixGuardrail(CustomGuardrail): return None return self._extract_application_name(alias, await self._get_app_name_regex()) - async def _resolve_routing(self, request_data: Mapping[str, Any]) -> ResolvedRouting | None: + async def _resolve_routing(self, request_data: Mapping[str, object]) -> ResolvedRouting | None: if self._application_id: return ResolvedRouting( self._application_id, 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 b0d66c39286..5b39ed42ed1 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_ovalix.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_ovalix.py @@ -1793,3 +1793,14 @@ async def test_tool_call_present_in_both_sources_is_scanned_once(): def test_extract_application_name_tolerates_unmatched_optional_group(): g = _static_guardrail() assert g._extract_application_name("app-", re.compile(r"app-(\w+)?")) is None + + +def test_non_string_actor_metadata_does_not_crash_hashing(): + g = _static_guardrail() + assert g._get_tracker_actor_id({"metadata": {"user_api_key_user_email": 12345}}) == g._get_tracker_actor_id({}) + + +def test_non_string_corrected_content_is_ignored(): + g = _static_guardrail() + assert g._get_trackers_corrected_message({"modified_data": {"content": {"nested": "obj"}}}) is None + assert g._get_trackers_corrected_message({"modified_data": {"content": ""}}) == ""