From 0dab7f9dcc870be7ad989464b05f21494c6bf5c7 Mon Sep 17 00:00:00 2001 From: Shalom Jamil Date: Mon, 20 Jul 2026 15:36:35 +0300 Subject: [PATCH 01/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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"]