From 0f4d5c6ad04669846e761dba4dac934c262c1fa7 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Sat, 31 Oct 2026 17:01:12 -0700 Subject: [PATCH] feat(ovalix)!: scan files, tool calls and tool results (#49309) * feat(ovalix): file and tool scanning, app discovery from the key alias Mirror of BerriAI/litellm#34132 at 5ffc268585d, rebuilt on main. * fix(ovalix): keep inspecting requests whose key has no user The tracker allows a checkpoint call with an empty actor without inspecting it, so a key with no user (or a /v1/messages or /v1/responses request that carries its own metadata object) got no Ovalix enforcement. Fall back to "unknown" as the guardrail did before, so every request is inspected. * refactor(ovalix): check texts without mutation and test discovery through tracker requests - _check_texts builds its result from per-text verdicts checked newest first by _check_text - discovery and routing cache unit tests answer the tracker routes instead of replacing private methods - integration test adds streaming cells and a tracker outage during a concurrent burst * fix(ovalix): resolve tool result names in one pass Each tool result rescanned every earlier message for its call id, so a request with thousands of tool results spent seconds in the event loop. Declarations are now indexed once per request and each result takes the nearest earlier one by bisection * test(ovalix): annotate new test bindings and build the burst plan with product * fix(ovalix): drop app discovery from the key alias A config with no application_id no longer boots in discovery mode, where the [Name] prefix of the key alias picked the Ovalix application and so its policy. Anyone who could rename a key could move it to a weaker app, or to an unknown name the tracker created with its default policy. The guardrail now needs application_id at startup again, as on main, and every checkpoint call goes by the configured application and checkpoint ids. Removes enable_routing_cache and fail_if_no_application, the app_name_regex and resolve_application tracker calls, and the per-alias routing cache. Restores main's startup check that a requested pre_call or post_call hook has its checkpoint id. --------- Co-authored-by: Shalom Jamil --- litellm/proxy/_lazy_openapi_snapshot.json | 12 + .../guardrail_hooks/ovalix/__init__.py | 2 + .../guardrail_hooks/ovalix/ovalix.py | 443 ++++++--- .../ovalix/ovalix_extraction.py | 396 ++++++++ .../guardrails/guardrail_hooks/ovalix.py | 4 + .../observability/test_ovalix_guardrail.py | 338 +++++++ .../guardrails/guardrail_hooks/test_ovalix.py | 854 ++++++++++++++++-- .../guardrail_hooks/test_ovalix_extraction.py | 443 +++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 5 + 9 files changed, 2298 insertions(+), 199 deletions(-) create mode 100644 litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix_extraction.py create mode 100644 tests/integration/observability/test_ovalix_guardrail.py create mode 100644 tests/unit/proxy/guardrails/guardrail_hooks/test_ovalix_extraction.py diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index d6d1636909b..5a80583b207 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -18222,6 +18222,18 @@ "description": "Whether to fail the request if the guardrail encounters an error. Implemented by guardrail='model_armor', 'generic_guardrail_api' and 'crowdstrike_aidr'. True (default) raises the error. False logs a critical error and lets the request proceed, so only a valid guardrail response can block or modify it.", "title": "Fail On Error" }, + "file_checkpoint_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "description": "File-checkpoint ID for the Ovalix Tracker service (falls back to pre/post).", + "title": "File Checkpoint Id" + }, "file_guardrail_timeout": { "anyOf": [ { diff --git a/litellm/proxy/guardrails/guardrail_hooks/ovalix/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/ovalix/__init__.py index 0c651864bbb..e6efc0d3592 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/ovalix/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/ovalix/__init__.py @@ -19,6 +19,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" application_id: Final = getattr(litellm_params, "application_id", None) pre_checkpoint_id: Final = getattr(litellm_params, "pre_checkpoint_id", None) post_checkpoint_id: Final = getattr(litellm_params, "post_checkpoint_id", None) + file_checkpoint_id: Final = getattr(litellm_params, "file_checkpoint_id", None) _ovalix_callback: Final = OvalixGuardrail( guardrail_name=guardrail.get("guardrail_name", ""), @@ -27,6 +28,7 @@ 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, event_hook=litellm_params.mode, default_on=litellm_params.default_on, timeout=litellm_params.timeout, diff --git a/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py b/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py index 8409a801c0f..f4e9cdaa2f4 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py +++ b/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py @@ -4,13 +4,19 @@ 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 -from typing import TYPE_CHECKING, Any, Final, Literal +from collections.abc import Mapping, Sequence +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, Literal import httpx -from typing_extensions import ReadOnly, TypedDict, Unpack +from typing_extensions import TypedDict, Unpack from litellm._logging import verbose_proxy_logger from litellm.exceptions import GuardrailRaisedException @@ -18,10 +24,23 @@ from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, ) +from litellm.litellm_core_utils.core_helpers import as_str_mapping 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_calls_from_messages, + extract_tool_results, + make_tool_data, + str_or_none, + tool_call_to_tool_data, + tool_data_key, + tool_result_text_indices, +) from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import GenericGuardrailAPIInputs @@ -32,13 +51,34 @@ if TYPE_CHECKING: BLOCKED_BY_OVALIX_FALLBACK_MESSAGE: Final = "This message was blocked by Ovalix" BLOCKED_ACTION_TYPE: Final = "block" +_MODIFY_ACTION_TYPES: Final = ("anonymize", "sanitize") +_DEFAULT_FILE_SIZE_LIMIT: Final = 64 * 1024 * 1024 +_NO_METADATA: Final[Mapping[str, object]] = MappingProxyType({}) +_FILE_BLOCK_ESCALATION_REASON: Final = ( + "This message was blocked by Ovalix because file content anonymization isn't possible via LiteLLM" +) +_UNDECODABLE_FILE_BLOCK_REASON: Final = ( + "This message was blocked by Ovalix because an attached file couldn't be decoded for inspection" +) +_TOOL_BLOCK_ESCALATION_REASON: Final = ( + "This message was blocked by Ovalix because tool call anonymization isn't possible via LiteLLM" +) +_TOOL_RESULT_BLOCK_ESCALATION_REASON: Final = ( + "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() + + +def _mapping_or_empty(value: object) -> Mapping[str, object]: + return as_str_mapping(value) or _NO_METADATA class _CustomGuardrailOptions(TypedDict, total=False, extra_items=object): """Base-class constructor options this guardrail forwards untouched to CustomGuardrail.""" - supported_event_hooks: ReadOnly[list[GuardrailEventHooks]] - class OvalixGuardrailMissingSecrets(Exception): """Raised when required Ovalix config (API base, key, application/checkpoint IDs) is missing.""" @@ -87,6 +127,8 @@ class OvalixGuardrail(CustomGuardrail): application_id: str | None = None, pre_checkpoint_id: str | None = None, post_checkpoint_id: str | None = None, + file_checkpoint_id: str | None = None, + supported_event_hooks: list[GuardrailEventHooks] | None = None, **kwargs: Unpack[_CustomGuardrailOptions], ): self._tracker_api_base = tracker_api_base or os.environ.get("OVALIX_TRACKER_API_BASE") @@ -94,23 +136,26 @@ 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") - supported_event_hooks: Final = kwargs.get("supported_event_hooks", []) + event_hooks: Final = self._validated_event_hooks(supported_event_hooks or ()) - self._validate_config(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}", + "x-api-key": self._tracker_api_key or "", + "Content-Type": "application/json", + } + ), + encoding="utf-8", + ) ) self._async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) - forwarded: Final[_CustomGuardrailOptions] = {**kwargs, "supported_event_hooks": supported_event_hooks} - super().__init__(**forwarded) + super().__init__(supported_event_hooks=event_hooks, **kwargs) verbose_proxy_logger.debug( "Ovalix Guardrail initialized: tracker=%s, application_id=%s, pre_checkpoint_id=%s, post_checkpoint_id=%s", self._tracker_api_base, @@ -119,44 +164,53 @@ class OvalixGuardrail(CustomGuardrail): self._post_checkpoint_id, ) - 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.""" - errors: Final[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" + def _validated_event_hooks(self, requested: Sequence[GuardrailEventHooks]) -> list[GuardrailEventHooks]: + """Ensure required secrets and checkpoint IDs are set; return the requested hooks plus one per configured checkpoint.""" + errors: Final = tuple( + message + for missing, 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"), + (not self._application_id, "Application ID, set OVALIX_APPLICATION_ID or pass application_id"), + ( + not self._pre_checkpoint_id and GuardrailEventHooks.pre_call in requested, + "Pre-checkpoint ID, set OVALIX_PRE_CHECKPOINT_ID or pass pre_checkpoint_id", + ), + ( + not self._post_checkpoint_id and GuardrailEventHooks.post_call in requested, + "Post-checkpoint ID, set OVALIX_POST_CHECKPOINT_ID or pass post_checkpoint_id", + ), + ( + not self._pre_checkpoint_id and not self._post_checkpoint_id, + "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 missing + ) 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: - supported_event_hooks.append(GuardrailEventHooks.pre_call) - if self._post_checkpoint_id and GuardrailEventHooks.post_call not in supported_event_hooks: - supported_event_hooks.append(GuardrailEventHooks.post_call) + auto_added: Final = tuple( + hook + for configured, hook in ( + (self._pre_checkpoint_id, GuardrailEventHooks.pre_call), + (self._post_checkpoint_id, GuardrailEventHooks.post_call), + ) + if configured and hook not in requested + ) + return [*requested, *auto_added] - def _get_actor(self, data: dict) -> str: + def _get_actor(self, data: Mapping[str, object]) -> str: """Return a stable actor identifier from request metadata (e.g. user email or id).""" - metadata: Final = data.get("metadata") or data.get("litellm_metadata") or {} - 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 "unknown" + metadata: Final = _mapping_or_empty(data.get("metadata") or data.get("litellm_metadata")) + return ( + str_or_none(metadata.get("user_api_key_user_email")) + or str_or_none(metadata.get("user_api_key_user_id")) + or "unknown" + ) - def _get_tracker_actor_id(self, data: dict) -> str: + def _get_tracker_actor_id(self, data: Mapping[str, object]) -> str: """Normalize the actor string into a short, stable id for Tracker API payloads.""" # NOTE: this hash is purely for normalization — it collapses an arbitrary actor # string (email, user id, or "unknown") into a compact, fixed-length, consistent @@ -166,7 +220,7 @@ class OvalixGuardrail(CustomGuardrail): normalized_actor_id: Final = 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, object]) -> str: """Return a unique identifier for the chat/session (actor + date + application_id).""" actor_hash: Final = self._get_tracker_actor_id(data) today: Final = datetime.datetime.now(datetime.timezone.utc).strftime("%Y-%m-%d") @@ -174,106 +228,254 @@ class OvalixGuardrail(CustomGuardrail): async def _call_checkpoint( self, - content: str, + data_type: str, + data: Mapping[str, object], checkpoint_id: str, actor: str, session_id: str, - ) -> dict[str, Any]: - """Call the Ovalix Tracker checkpoint API and return the JSON response.""" + oversize: bool = False, + ) -> Mapping[str, object]: + """Call the Ovalix Tracker checkpoint API and return the JSON response. + + Both routes live on the tracker's /beta litellm router, which accepts the api key this + guardrail already sends. + """ application_id: Final = self._application_id if not application_id or not checkpoint_id: raise ValueError("Ovalix: application_id or checkpoint_id not resolved") - url: Final = f"{self._tracker_api_base}/tracking/custom_application/checkpoint" - headers: Final = dict(self._tracker_headers) + route: Final = "file_checkpoint" if data_type == "FILE" else "checkpoint" payload: Final = { "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", + **(MappingProxyType({"oversize": True}) if oversize else _NO_METADATA), } - response: Final = await self._async_handler.post(url, headers=headers, json=payload, timeout=self.timeout) + response: Final = await self._async_handler.post( + f"{self._tracker_api_base}/tracking/beta/{route}", + headers=self._tracker_headers, + json=payload, + timeout=self.timeout, + ) response.raise_for_status() return response.json() + def _verdict(self, resp: Mapping[str, object]) -> tuple[str, str | None]: + return (str_or_none(resp.get("action_type")) or "").lower(), self._get_trackers_corrected_message(resp) + + async def _block_reason_for_item( + self, + data_type: str, + data: Mapping[str, object], + checkpoint_id: str, + actor: str, + session_id: str, + escalation_reason: str, + oversize: bool = False, + ) -> str | None: + try: + resp: Final = await self._call_checkpoint( + data_type, data, checkpoint_id, actor, session_id, oversize=oversize + ) + 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}", + 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: Sequence[tuple[str, Mapping[str, object]]], + checkpoint_id: str, + actor: str, + session_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, escalation_reason + ) + if reason is not None: + return reason + return None + + async def _check_files_for_block( + self, + file_parts: Sequence[FilePart], + checkpoint_id: str, + actor: str, + session_id: str, + ) -> str | None: + if any(part.undecodable for part in file_parts): + return _UNDECODABLE_FILE_BLOCK_REASON + 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, + _FILE_BLOCK_ESCALATION_REASON, + oversize=part.oversize, + ) + if reason is not None: + return reason + return None + @log_guardrail_information async def apply_guardrail( self, inputs: GenericGuardrailAPIInputs, - request_data: dict, + request_data: Mapping[str, object], input_type: Literal["request", "response"], logging_obj: "LiteLLMLoggingObj | None" = None, ) -> GenericGuardrailAPIInputs: - """ - Apply Ovalix guardrail to the given inputs (request or response text). - - Used by the unified guardrail flow and the /apply_guardrail API. - For "request", uses the pre-checkpoint; for "response", uses the post-checkpoint. - - 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. - - 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 - - tracker_actor_id: Final = self._get_tracker_actor_id(request_data) + actor: Final = self._get_actor(request_data) session_id: Final = self._get_session_id(request_data) - texts: Final = inputs.get("texts") or [] - if not texts or not isinstance(texts, list): + is_response: Final = input_type == "response" + + prompt_checkpoint: Final = self._post_checkpoint_id if is_response else self._pre_checkpoint_id + file_checkpoint: Final = self._file_checkpoint_id or prompt_checkpoint + if not file_checkpoint: + verbose_proxy_logger.debug( + "Ovalix guardrail: application %s has no %s checkpoint, leaving this direction uninspected", + self._application_id, + input_type, + ) return inputs - if input_type == "response": - if not self._post_checkpoint_id: - return inputs - corrected_llm_responses: Final = await self._generate_post_guardrail_llm_texts( - texts, tracker_actor_id, session_id, self._post_checkpoint_id + structured_messages: Final = inputs.get("structured_messages") or () + file_parts: Final = ( + 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: Final = await self._check_files_for_block(file_parts, file_checkpoint, actor, session_id) + 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", + self._application_id, + input_type, ) - return {**inputs, "texts": corrected_llm_responses} + return inputs - if self._pre_checkpoint_id: - post_guardrail_texts: Final = 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 + tool_calls: Final = ( + *(inputs.get("tool_calls") or ()), + *extract_tool_calls_from_messages(structured_messages), + ) + unique_tool_data: Final = MappingProxyType( + {tool_data_key(data): data for data in (tool_call_to_tool_data(tc) for tc in tool_calls) if data} + ) + tool_call_items: Final = tuple(("TOOL", data) for data in unique_tool_data.values()) + tool_block: Final = await self._check_items_block_only( + tool_call_items, + prompt_checkpoint, + actor, + session_id, + _TOOL_BLOCK_ESCALATION_REASON, + ) + if tool_block is not None: + self._block_current_message(tool_block) - 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: Final[list[str]] = [] + tool_results: Final = extract_tool_results(structured_messages) + tool_result_items: Final = tuple(("TOOL", make_tool_data(name, content)) for name, content, _ in tool_results) + tool_result_block: Final = await self._check_items_block_only( + tool_result_items, + prompt_checkpoint, + actor, + session_id, + _TOOL_RESULT_BLOCK_ESCALATION_REASON, + ) + if tool_result_block is not None: + self._block_current_message(tool_result_block) - is_first_response = True - for llm_response in reversed(texts): - try: - resp = await self._call_checkpoint(llm_response, checkpoint_id, actor, session_id) - except Exception as e: - verbose_proxy_logger.exception("Ovalix apply_guardrail checkpoint call failed: %s", e) - raise GuardrailRaisedException( - guardrail_name=self.guardrail_name, - message=f"Ovalix guardrail error: {e}", - should_wrap_with_default_message=False, - ) from e + texts: Final = inputs.get("texts") or () + if not texts: + return inputs + output_texts: Final = await self._check_texts( + texts, + prompt_checkpoint, + actor, + session_id, + tool_result_text_indices(structured_messages, texts), + ) + if output_texts is None: + return inputs + return {**inputs, "texts": output_texts} - 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 + async def _file_part_to_data(self, part: FilePart) -> Mapping[str, object]: + extension: Final = mimetypes.guess_extension(part.mime_hint) if part.mime_hint else None + name: Final = part.name or (f"file{extension}" if extension else "file") + content: Final = ( + 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} + + async def _check_texts( + self, + texts: Sequence[str], + checkpoint_id: str, + actor: str, + session_id: str, + skip_indices: frozenset[int], + ) -> list[str] | None: + newest_first: Final = tuple(index for index in reversed(range(len(texts))) if index not in skip_indices) + checked: Final = MappingProxyType( + { + index: await self._check_text( + texts[index], checkpoint_id, actor, session_id, is_newest=index == len(texts) - 1 + ) + for index in newest_first + } + ) + output: Final = [checked.get(index, text) for index, text in enumerate(texts)] + return output if output != list(texts) else None + + async def _check_text( + self, + content: str, + checkpoint_id: str, + actor: str, + session_id: str, + is_newest: bool, + ) -> str: + """Return the text that replaces `content` after Ovalix checks it; a blocked newest text raises.""" + try: + resp: Final = await self._call_checkpoint("TEXT", {"content": content}, checkpoint_id, actor, session_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}", + should_wrap_with_default_message=False, + ) from e + action, corrected = self._verdict(resp) + if action == BLOCKED_ACTION_TYPE: + block_message: Final = corrected or BLOCKED_BY_OVALIX_FALLBACK_MESSAGE + if is_newest: + self._block_current_message(block_message) + return block_message + if action in _MODIFY_ACTION_TYPES and corrected is not None: + return corrected + return content def _block_current_message(self, blocking_message: str) -> None: """Raise OvalixGuardrailBlockedException with the given message (no default wrapper).""" @@ -283,12 +485,13 @@ 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, object]) -> str | None: """Extract corrected/blocking message content from Tracker checkpoint response.""" - modified: Final = resp.get("modified_data") - if isinstance(modified, dict) and "content" in modified: - return modified["content"] - return None + modified: Final = as_str_mapping(resp.get("modified_data")) + if modified is None: + return None + content: Final = modified.get("content") + return content if isinstance(content, str) else None @staticmethod def get_config_model() -> type["GuardrailConfigModel"] | None: 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..ddc6485b7d5 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix_extraction.py @@ -0,0 +1,396 @@ +import base64 +import json +import posixpath +import re +from bisect import bisect_left +from collections.abc import Callable, Iterator, Mapping, Sequence +from itertools import chain, groupby +from operator import attrgetter +from types import MappingProxyType +from typing import Final, NamedTuple +from urllib.parse import unquote, urlparse + +from pydantic import TypeAdapter, ValidationError + +from litellm.litellm_core_utils.core_helpers import as_str_mapping + +_TOOL_NAME_MAX_LENGTH: Final = 100 +_DEFAULT_TOOL_RESULT_NAME: Final = "tool_result" +_NO_TOOL_INPUT: Final[Mapping[str, object]] = MappingProxyType({}) + +_DATA_URL_RE: Final = re.compile(r"^data:(?P[^;,]+)?(?P(?:;[^;,]+)*?)(?P;base64)?,", re.IGNORECASE) +_URLSAFE_TO_STANDARD_B64: Final = str.maketrans("-_", "+/") +_OBJECT_LIST: Final = TypeAdapter(list[object]) + + +def str_or_none(value: object) -> str | None: + return value if isinstance(value, str) and value else None + + +def _list_items(value: object) -> Sequence[object]: + try: + return _OBJECT_LIST.validate_python(value, strict=True) + except ValidationError: + return () + + +class FilePart(NamedTuple): + name: str | None + data: bytes | None + mime_hint: str | None + inline: bool + oversize: bool + message_index: int + undecodable: bool = False + + +def _split_data_url(value: str) -> tuple[str | None, str | None]: + match: Final = _DATA_URL_RE.match(value) + if not match: + if value.lower().startswith("data:"): + return None, None + return None, value + mime: Final = match.group("mime") or None + if not match.group("b64"): + return mime, None + return mime, value[match.end() :] + + +def _b64decode_or_none(payload: str) -> bytes | None: + try: + return base64.b64decode(payload, validate=True) + except ValueError: + return None + + +def _decoded_base64(cleaned: str) -> bytes | None: + strict: Final = _b64decode_or_none(cleaned) + if strict is not None or ("-" not in cleaned and "_" not in cleaned): + return strict + return _b64decode_or_none(cleaned.translate(_URLSAFE_TO_STANDARD_B64)) + + +def _inline_part(name: str | None, value: str, size_limit: int | None, message_index: int) -> FilePart | None: + mime_hint, payload = _split_data_url(value) + if payload is None: + return FilePart(name, None, mime_hint, True, False, message_index, undecodable=True) + cleaned: Final = "".join(payload.split()) + if not cleaned: + return None + if size_limit is not None and (len(cleaned) * 3) // 4 - 2 > size_limit: + return FilePart(name, None, mime_hint, True, True, message_index) + data: Final = _decoded_base64(cleaned) + if data is None: + return FilePart(name, None, mime_hint, True, False, message_index, undecodable=True) + if size_limit is not None and len(data) > size_limit: + return FilePart(name, None, mime_hint, True, True, message_index) + return FilePart(name, data, mime_hint, True, False, message_index) + + +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: Mapping[str, object], size_limit: int | None, message_index: int) -> FilePart | None: + file_obj: Final = as_str_mapping(block.get("file")) + if file_obj is None: + return None + name: Final = str_or_none(file_obj.get("filename")) or str_or_none(file_obj.get("file_id")) + file_data: Final = str_or_none(file_obj.get("file_data")) + inline: Final = _inline_part(name, file_data, size_limit, message_index) if file_data else None + return inline or FilePart(name, None, None, False, False, message_index) + + +def _part_from_image_url_block( + block: Mapping[str, object], size_limit: int | None, message_index: int +) -> FilePart | None: + image_url: Final = block.get("image_url") + image_url_fields: Final = as_str_mapping(image_url) + url: Final = image_url_fields.get("url") if image_url_fields is not None else image_url + if not isinstance(url, str) or not url: + return None + if url.startswith("data:"): + return _inline_part(None, url, size_limit, message_index) + return FilePart(_name_from_url(url), None, None, False, False, message_index) + + +def _part_from_input_file_block( + block: Mapping[str, object], size_limit: int | None, message_index: int +) -> FilePart | None: + declared_name: Final = str_or_none(block.get("filename")) or str_or_none(block.get("file_id")) + file_data: Final = str_or_none(block.get("file_data")) + inline: Final = _inline_part(declared_name, file_data, size_limit, message_index) if file_data else None + if inline is not None: + return inline + file_url: Final = block.get("file_url") + name: Final = ( + _name_from_url(file_url) if isinstance(file_url, str) and file_url and not declared_name else declared_name + ) + return FilePart(name, None, None, False, False, message_index) + + +def _part_from_input_audio_block( + block: Mapping[str, object], size_limit: int | None, message_index: int +) -> FilePart | None: + audio: Final = as_str_mapping(block.get("input_audio")) + if audio is None: + return None + data_b64: Final = audio.get("data") + if not isinstance(data_b64, str) or not data_b64: + return None + name: Final = f"audio.{audio.get('format') or 'bin'}" + return _inline_part(name, data_b64, size_limit, message_index) or FilePart( + name, None, None, False, False, message_index + ) + + +_BLOCK_PARSERS: Final[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]: + for block in map(as_str_mapping, _list_items(message.get("content"))): + if block is None: + 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 _file_parts_of_messages(structured_messages: Sequence[object], size_limit: int | None) -> Iterator[FilePart]: + for message_index, message in enumerate(map(as_str_mapping, structured_messages)): + if message is not None: + yield from _file_parts_of_message(message, size_limit, message_index) + + +def extract_file_parts_from_messages( + structured_messages: Sequence[object] | None, size_limit: int | None = None +) -> tuple[FilePart, ...]: + return tuple(_file_parts_of_messages(structured_messages or (), size_limit)) + + +def _file_part_of_image(value: str, size_limit: int | None, index: int) -> FilePart | None: + if value.startswith(("http://", "https://")): + name: Final = _name_from_url(value) + return FilePart(name, None, None, False, False, index) if name else None + return _inline_part(None, value, size_limit, index) + + +def extract_file_parts_from_images( + images: Sequence[object] | None, size_limit: int | None = None +) -> tuple[FilePart, ...]: + candidates: Final = ( + _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: Final = str(name) if str(name).strip() else _DEFAULT_TOOL_RESULT_NAME + truncated: Final = action_name[:_TOOL_NAME_MAX_LENGTH] + tool_name: Final = truncated if truncated.strip() else _DEFAULT_TOOL_RESULT_NAME + return { + "content": content, + "tool_name": tool_name, + "action_name": action_name, + "tool_input": dict(tool_input or ()), + } + + +def _tool_call_field(tool_call: object, key: str) -> object: + fields: Final = as_str_mapping(tool_call) + return fields.get(key) if fields is not None else getattr(tool_call, key, 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: Final = json.loads(raw_arguments) + except (ValueError, TypeError): + return _NO_TOOL_INPUT + return as_str_mapping(parsed) or _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 + arguments: Final = as_str_mapping(raw_arguments) + return _json_or_str(raw_arguments), arguments if arguments is not None else _NO_TOOL_INPUT + + +def tool_call_to_tool_data(tool_call: object) -> Mapping[str, object] | None: + function: Final = _tool_call_field(tool_call, "function") + name: Final = _tool_call_field(function, "name") + if not name or not str(name).strip(): + return None + content, tool_input = _tool_content_and_input(_tool_call_field(function, "arguments")) + return make_tool_data(str(name), content, tool_input) + + +def _message_tool_calls(message: object) -> Sequence[object]: + fields: Final = as_str_mapping(message) + return _list_items(fields.get("tool_calls")) if fields is not None else () + + +def extract_tool_calls_from_messages(structured_messages: Sequence[object] | None) -> tuple[object, ...]: + """Tool calls declared on the messages themselves. + + Surfaces such as the Anthropic request path populate ``structured_messages`` but leave the + top-level ``tool_calls`` input empty, so calls made in prior assistant turns are only visible here. + """ + return tuple(chain.from_iterable(_message_tool_calls(message) for message in structured_messages or ())) + + +def tool_data_key(tool_data: Mapping[str, object]) -> str: + """Stable identity for a tool payload, so a call reached from two sources is only scanned once.""" + return json.dumps(tool_data, sort_keys=True, default=str) + + +def _tool_content_block_text(block: object) -> str | None: + if isinstance(block, str): + return block + fields: Final = as_str_mapping(block) + return str_or_none(fields.get("text")) if fields is not None else None + + +def _tool_content_text(content: object) -> str: + if isinstance(content, str): + return content + fields: Final = as_str_mapping(content) + if fields is not None: + return _json_or_str(fields) + texts: Final = (_tool_content_block_text(block) for block in _list_items(content)) + return "\n".join(text for text in texts if text is not None) + + +def _extract_tool_content(content: object) -> str | None: + text: Final = _tool_content_text(content) + return text if text.strip() else None + + +class _Declaration(NamedTuple): + call_id: str + message_index: int + name: str + + +def _declarations_in(message_index: int, message: object) -> Iterator[_Declaration]: + fields: Final = as_str_mapping(message) + if fields is None or fields.get("role") != "assistant": + return + for tool_call in map(as_str_mapping, _list_items(fields.get("tool_calls"))): + if tool_call is None: + continue + call_id = tool_call.get("id") + function = as_str_mapping(tool_call.get("function")) + name = str_or_none(function.get("name")) if function is not None else None + if isinstance(call_id, str) and call_id and name and name.strip(): + yield _Declaration(call_id, message_index, name) + + +def _declarations_by_call_id(messages: Sequence[object]) -> Mapping[str, tuple[_Declaration, ...]]: + """Every named declaration of each tool call id, in message order, built in one pass.""" + declarations: Final = chain.from_iterable( + _declarations_in(index, message) for index, message in enumerate(messages) + ) + ordered: Final = sorted(declarations, key=attrgetter("call_id", "message_index")) + return MappingProxyType({call_id: tuple(group) for call_id, group in groupby(ordered, key=attrgetter("call_id"))}) + + +def _resolve_tool_name( + declarations: Mapping[str, tuple[_Declaration, ...]], tool_index: int, tool_call_id: object +) -> str: + """Name from the nearest declaration of ``tool_call_id`` before ``tool_index``.""" + if not isinstance(tool_call_id, str) or not tool_call_id: + return _DEFAULT_TOOL_RESULT_NAME + declared: Final = declarations.get(tool_call_id, ()) + earlier: Final = bisect_left(declared, tool_index, key=attrgetter("message_index")) + return declared[earlier - 1].name if earlier else _DEFAULT_TOOL_RESULT_NAME + + +def extract_tool_results( + structured_messages: Sequence[object] | None, +) -> tuple[tuple[str, str, object], ...]: + messages: Final = tuple(structured_messages or ()) + declarations: Final = _declarations_by_call_id(messages) + + def _results() -> Iterator[tuple[str, str, object]]: + for index, message in enumerate(map(as_str_mapping, messages)): + if message is None 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") + yield _resolve_tool_name(declarations, index, tool_call_id), content, tool_call_id + + return tuple(_results()) + + +def _block_texts(content: object) -> Iterator[object]: + blocks: Final = (as_str_mapping(block) for block in _list_items(content)) + return (block.get("text") for block in blocks if block is not None and block.get("text") is not None) + + +def _message_text_origins(structured_messages: Sequence[object] | None) -> Iterator[tuple[object, bool]]: + for message in map(as_str_mapping, structured_messages or ()): + if message is None: + 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 + else: + yield from ((text, from_tool_result) for text in _block_texts(content)) + + +def tool_result_text_indices(structured_messages: Sequence[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: Final = 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/litellm/types/proxy/guardrails/guardrail_hooks/ovalix.py b/litellm/types/proxy/guardrails/guardrail_hooks/ovalix.py index 0ec353fa945..7613a2016bb 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/ovalix.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/ovalix.py @@ -28,6 +28,10 @@ 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).", + ) @staticmethod def ui_friendly_name() -> str: diff --git a/tests/integration/observability/test_ovalix_guardrail.py b/tests/integration/observability/test_ovalix_guardrail.py new file mode 100644 index 00000000000..d7c9e5ac449 --- /dev/null +++ b/tests/integration/observability/test_ovalix_guardrail.py @@ -0,0 +1,338 @@ +"""Ovalix guardrail driven through a real proxy. + +The Ovalix tracker is the only guardrail double: an owned wire peer that answers the `/tracking/beta` routes +the way the live tracker does, including allowing a checkpoint call whose `actor` is empty without inspecting +it. The provider is a second owned peer. The proxy, Postgres and Redis run for real with two workers. +""" + +import json +import threading +import uuid +from collections.abc import Iterator +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from itertools import product +from pathlib import Path +from typing import Final +from urllib.parse import urlsplit + +import httpx +import pytest +import yaml +from integration._support.client import Gateway, JsonValue, gateway_from_environment +from integration._support.process import owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server + +pytestmark: Final = pytest.mark.shares_worker_state + +API_KEY: Final = "synthetic-ovalix-key" +APP_ID: Final = "synthetic-ovalix-app" +OTHER_APP: Final = "Synthetic Other App" +PRE_CHECKPOINT: Final = "synthetic-pre-checkpoint" +POST_CHECKPOINT: Final = "synthetic-post-checkpoint" +KEYWORD: Final = "Project Falcon" +NAME: Final = "David" +DENIAL: Final = "synthetic ovalix denial" +REPLY: Final = "synthetic provider reply" +PINNED: Final = "ovalix-pinned" +ENDPOINTS: Final = ("/v1/chat/completions", "/v1/messages", "/v1/responses") +CHECKPOINT: Final = "/tracking/beta/checkpoint" +OUTAGE_AFTER: Final = 6 + + +class _Outage: + """Takes the tracker double down once it has answered a set number of checkpoint calls.""" + + def __init__(self) -> None: + self._lock: Final = threading.Lock() + self._remaining: int | None = None + + def arm(self, answered: int) -> None: + with self._lock: + self._remaining = answered + + def clear(self) -> None: + with self._lock: + self._remaining = None + + def admits(self, path: str) -> bool: + with self._lock: + if self._remaining is None: + return True + if self._remaining == 0: + return False + if path == CHECKPOINT: + self._remaining -= 1 + return True + + +OUTAGE: Final = _Outage() + + +def _checkpoint_verdict(body: dict[str, object]) -> dict[str, object]: + data: Final = body["data"] + content: Final = str(data.get("content", "")) if isinstance(data, dict) else "" + if not body["actor"]: + return {"action_type": "allow"} + if KEYWORD in content: + return {"action_type": "block", "modified_data": {"content": DENIAL}} + if NAME in content: + return {"action_type": "anonymize", "modified_data": {"content": content.replace(NAME, "{Name}")}} + return {"action_type": "allow"} + + +def _tracker(request: Request) -> Reply: + if request.headers.get("x-api-key") != API_KEY: + return Reply(status=401, body=b'{"detail": "missing api key"}') + path: Final = urlsplit(request.target).path + if not OUTAGE.admits(path): + return Reply(status=503, body=b'{"detail": "tracker unavailable"}') + if path != CHECKPOINT: + return Reply(status=404, body=b'{"detail": "not found"}') + return Reply(body=json.dumps(_checkpoint_verdict(json.loads(request.body))).encode()) + + +def _responses_reply(marker: str) -> dict[str, object]: + return { + "id": "resp_" + marker, + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-4o-mini", + "output": [ + { + "type": "message", + "id": "msg_" + marker, + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": REPLY, "annotations": []}], + } + ], + "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + } + + +def _chat_reply(marker: str) -> dict[str, object]: + return { + "id": "chatcmpl-" + marker, + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [{"index": 0, "message": {"role": "assistant", "content": REPLY}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + } + + +def _provider(request: Request) -> Reply: + marker: Final = uuid.uuid4().hex + is_responses: Final = urlsplit(request.target).path.endswith("/responses") + return Reply(body=json.dumps(_responses_reply(marker) if is_responses else _chat_reply(marker)).encode()) + + +def _alias() -> str: + return f"[{OTHER_APP}] integration key {uuid.uuid4().hex}" + + +def _ovalix(name: str, tracker: str, **routing: str) -> dict[str, object]: + return { + "guardrail_name": name, + "litellm_params": { + "guardrail": "ovalix", + "mode": "pre_call", + "default_on": False, + "tracker_api_base": tracker, + "tracker_api_key": API_KEY, + **routing, + }, + } + + +@dataclass(frozen=True, slots=True) +class Keys: + no_user: str + user_id: str + with_user_id: str + email: str + with_email: str + + +@dataclass(frozen=True, slots=True) +class Rig: + proxy: Gateway + tracker: Wire + provider: Wire + model: str + keys: Keys + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + root: Final = tmp_path_factory.mktemp("ovalix") + with ( + gateway_from_environment() as gateway, + wire_server(_tracker) as tracker, + wire_server(_provider) as provider, + ): + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["litellm_settings"]["cache"] = False + config["guardrails"] = [ + _ovalix( + PINNED, + tracker.url, + application_id=APP_ID, + pre_checkpoint_id=PRE_CHECKPOINT, + post_checkpoint_id=POST_CHECKPOINT, + ), + ] + path: Final = root / "ovalix.yaml" + path.write_text(yaml.safe_dump(config)) + with ( + owned_proxy_process(gateway, root, {}, config=path, workers=2) as owned, + owned.gateway.scenario() as scenario, + ): + model: Final = scenario.model( + model="openai/gpt-4o-mini", api_base=provider.url + "/v1", api_key="synthetic-openai-key" + ) + user_id: Final = scenario.user() + email: Final = f"ovalix-{uuid.uuid4().hex}@example.com" + keys: Final = Keys( + no_user=scenario.key(key_alias=_alias()), + user_id=user_id, + with_user_id=scenario.key(key_alias=_alias(), user_id=user_id), + email=email, + with_email=scenario.key(key_alias=_alias(), user_id=scenario.user(user_email=email)), + ) + yield Rig(owned.gateway, tracker, provider, model, keys) + + +def _post(rig: Rig, endpoint: str, key: str, text: str, extra: dict[str, JsonValue] | None = None) -> httpx.Response: + prompt: Final[dict[str, JsonValue]] = ( + {"input": text} + if endpoint == "/v1/responses" + else {"max_tokens": 64, "messages": [{"role": "user", "content": text}]} + if endpoint == "/v1/messages" + else {"messages": [{"role": "user", "content": text}]} + ) + return rig.proxy.request( + "POST", endpoint, {"model": rig.model, "guardrails": [PINNED], **prompt, **(extra or {})}, key=key + ) + + +def _send(rig: Rig, endpoint: str, key: str, text: str, extra: dict[str, JsonValue] | None = None) -> httpx.Response: + rig.tracker.drain() + rig.provider.drain() + return _post(rig, endpoint, key, text, extra) + + +def _checkpoint_actors(rig: Rig) -> list[object]: + return [ + json.loads(request.body)["actor"] + for request in rig.tracker.drain() + if urlsplit(request.target).path == CHECKPOINT + ] + + +@pytest.mark.parametrize("stream", (False, True)) +@pytest.mark.parametrize("endpoint", ENDPOINTS) +def test_a_prompt_from_a_key_with_no_user_is_inspected(rig: Rig, endpoint: str, stream: bool) -> None: + response: Final = _send( + rig, endpoint, rig.keys.no_user, f"Tell me about {KEYWORD} {uuid.uuid4()}", {"stream": stream} + ) + + assert response.status_code == 400 and DENIAL in response.text, response.text + assert rig.provider.drain() == () + assert _checkpoint_actors(rig) == ["unknown"] + + +@pytest.mark.parametrize( + ("endpoint", "client_metadata"), + (("/v1/messages", {"user_id": "client-chosen"}), ("/v1/responses", {"client": "chosen"})), +) +def test_a_client_metadata_object_does_not_switch_inspection_off( + rig: Rig, endpoint: str, client_metadata: dict[str, JsonValue] +) -> None: + response: Final = _send( + rig, + endpoint, + rig.keys.with_user_id, + f"Tell me about {KEYWORD} {uuid.uuid4()}", + {"metadata": client_metadata}, + ) + + assert response.status_code == 400 and DENIAL in response.text, response.text + assert rig.provider.drain() == () + assert _checkpoint_actors(rig) == ["unknown"] + + +def test_a_name_from_a_key_with_no_user_reaches_the_provider_anonymized(rig: Rig) -> None: + marker: Final = uuid.uuid4().hex + response: Final = _send(rig, "/v1/chat/completions", rig.keys.no_user, f"My name is {NAME} {marker}") + + assert response.status_code == 200, response.text + sent: Final = [json.loads(request.body)["messages"][-1]["content"] for request in rig.provider.drain()] + assert sent == [f"My name is {{Name}} {marker}"] + + +def test_the_actor_is_the_key_users_email_then_their_id(rig: Rig) -> None: + by_id: Final = _send(rig, "/v1/chat/completions", rig.keys.with_user_id, f"hello {uuid.uuid4()}") + assert by_id.status_code == 200, by_id.text + assert _checkpoint_actors(rig) == [rig.keys.user_id] + + by_email: Final = _send(rig, "/v1/chat/completions", rig.keys.with_email, f"hello {uuid.uuid4()}") + assert by_email.status_code == 200, by_email.text + assert _checkpoint_actors(rig) == [rig.keys.email] + + +@pytest.mark.parametrize("endpoint", ENDPOINTS) +def test_a_key_alias_naming_another_app_does_not_change_the_application_checked(rig: Rig, endpoint: str) -> None: + response: Final = _send(rig, endpoint, rig.keys.with_email, f"Tell me about {KEYWORD} {uuid.uuid4()}") + + assert response.status_code == 400 and DENIAL in response.text, response.text + assert rig.provider.drain() == () + calls: Final = [(urlsplit(request.target).path, json.loads(request.body)) for request in rig.tracker.drain()] + assert [ + (path, body["application_id"], body["checkpoint_id"], "application_name" in body) for path, body in calls + ] == [(CHECKPOINT, APP_ID, PRE_CHECKPOINT, False)] + + +def _burst(rig: Rig) -> tuple[httpx.Response, ...]: + plan: Final = tuple(product(ENDPOINTS, (False, True), range(4))) + with ThreadPoolExecutor(max_workers=12) as pool: + return tuple( + pool.map( + lambda cell: _post( + rig, + cell[0], + rig.keys.no_user, + f"Tell me about {KEYWORD} {uuid.uuid4()}", + {"stream": cell[1]}, + ), + plan, + ) + ) + + +def test_a_tracker_outage_mid_burst_fails_closed_and_recovers(rig: Rig) -> None: + rig.tracker.drain() + rig.provider.drain() + OUTAGE.arm(OUTAGE_AFTER) + try: + during: Final = _burst(rig) + finally: + OUTAGE.clear() + + denied: Final = [response for response in during if response.status_code == 400 and DENIAL in response.text] + failed: Final = [response for response in during if response.status_code >= 400 and DENIAL not in response.text] + assert len(denied) == OUTAGE_AFTER, [response.text for response in during] + assert len(failed) == len(during) - OUTAGE_AFTER, [response.text for response in during] + assert all("Ovalix guardrail error" in response.text for response in failed), [response.text for response in failed] + assert rig.provider.drain() == () + assert set(_checkpoint_actors(rig)) == {"unknown"} + + after: Final = _burst(rig) + + assert all(response.status_code == 400 and DENIAL in response.text for response in after), [ + response.text for response in after + ] + assert rig.provider.drain() == () + assert _checkpoint_actors(rig) == ["unknown"] * len(after) diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_ovalix.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_ovalix.py index 4160a835ca4..d3553206394 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_ovalix.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_ovalix.py @@ -3,8 +3,11 @@ Unit tests for Ovalix guardrail: config resolution and apply_guardrail behavior 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 typing import Final from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -16,8 +19,18 @@ from litellm.proxy.guardrails.guardrail_hooks.ovalix.ovalix import ( OvalixGuardrailBlockedException, OvalixGuardrailMissingSecrets, ) +from litellm.proxy.guardrails.guardrail_hooks.ovalix.ovalix_extraction import FilePart +from litellm.types.guardrails import GuardrailEventHooks 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", @@ -74,6 +87,44 @@ def _guardrail_kwargs(): } +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, + ) + + +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 + + class TestOvalixGuardrailConfigModel: """Minimal config model tests: wiring only.""" @@ -152,7 +203,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: @@ -161,12 +212,11 @@ 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", + data_type="TEXT", + data={"content": "hello"}, checkpoint_id="pre-1", actor="a1b2c3d4", session_id="session-1", @@ -175,9 +225,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/beta/checkpoint") body = call_args.kwargs["json"] assert body["application_id"] == "app-1" assert body["checkpoint_id"] == "pre-1" @@ -185,6 +233,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: @@ -207,9 +256,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, @@ -233,9 +280,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 = {} @@ -244,9 +289,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, @@ -279,9 +322,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( @@ -326,9 +367,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, @@ -349,15 +388,13 @@ 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: 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 = {} @@ -366,9 +403,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, @@ -377,7 +412,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(): @@ -398,9 +433,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( @@ -415,9 +448,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( @@ -436,9 +467,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, @@ -451,9 +480,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( @@ -469,9 +496,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( @@ -524,9 +549,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, @@ -547,23 +570,14 @@ 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" + assert ( + guardrail._get_actor({"metadata": {"user_api_key_user_email": "", "user_api_key_user_id": ""}}) + == "unknown" + ) finally: for k in _ovalix_env(): if k in os.environ: @@ -600,9 +614,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: @@ -614,16 +626,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): @@ -635,9 +642,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, @@ -651,3 +656,694 @@ class TestOvalixGuardrail: for k in _ovalix_env(): if k in os.environ: del os.environ[k] + + +def _checkpoint_calls(mock_post): + """(body, url) of every tracker call the guardrail made.""" + return [(c.kwargs["json"], c.args[0]) for c in mock_post.call_args_list] + + +def _route_of(url): + return url.rsplit("/", 1)[-1] + + +def _mock_tracker(g): + post_resp = MagicMock() + post_resp.json.return_value = TRACKER_RESPONSE_ALLOW + post_resp.raise_for_status = MagicMock() + g._async_handler.get = AsyncMock() + g._async_handler.post = AsyncMock(return_value=post_resp) + return g._async_handler.get, g._async_handler.post + + +_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, timeout=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, timeout=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, timeout=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_file_checkpoint_call_routes_to_litellm_file_endpoint(): + g = _static_guardrail() + seen = {} + + async def _post(url, headers=None, json=None, timeout=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") + assert seen["url"] == "https://t/tracking/beta/file_checkpoint" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("oversize", [True, False]) +async def test_file_checkpoint_flags_oversize_parts(oversize): + g = _static_guardrail() + part = FilePart(name="huge.bin", data=None, mime_hint=None, inline=oversize, oversize=oversize, message_index=0) + seen = {} + + async def _post(url, headers=None, json=None, timeout=None): + seen["last"] = json + r = MagicMock() + r.json.return_value = _ALLOW + r.raise_for_status = MagicMock() + return r + + with patch.object(g._async_handler, "post", new=_post): + await g._check_files_for_block([part], "file-1", "a", "s") + assert seen["last"]["data"]["content"] is None + assert seen["last"].get("oversize", False) is oversize + + +@pytest.mark.asyncio +async def test_tool_call_block_raises(): + g = _static_guardrail() + 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_texts_are_checked_newest_first_and_a_blocked_newest_stops_the_scan(): + g: Final = _static_guardrail() + post, calls = _recording_post(lambda body: _BLOCK if body["data"].get("content") == "new" else _ALLOW) + + with patch.object(g._async_handler, "post", new=post): + with pytest.raises(OvalixGuardrailBlockedException): + await g.apply_guardrail( + inputs=GenericGuardrailAPIInputs(texts=["old", "middle", "new"]), + request_data={}, + input_type="request", + logging_obj=None, + ) + + assert [content for data_type, content in calls if data_type == "TEXT"] == ["new"] + + +@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_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, timeout=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 +@pytest.mark.parametrize( + "request_data", + [ + {}, + {"metadata": {"user_id": "client-chosen"}, "litellm_metadata": {"user_api_key_user_id": "uid-1"}}, + ], +) +async def test_request_with_no_user_field_sends_unknown_actor(request_data): + g = _static_guardrail() + inputs = GenericGuardrailAPIInputs(texts=["hi"]) + seen = {} + + async def _post(url, headers=None, json=None, timeout=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"] == "unknown" + + +def _recording_post(mapping_fn=lambda body: _ALLOW): + calls = [] + + async def _post(url, headers=None, json=None, timeout=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_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, timeout=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", "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"], + structured_messages=[ + {"role": "assistant", "tool_calls": [{"id": "c1", "function": {"name": "get_weather"}}]}, + {"role": "tool", "tool_call_id": "c1", "content": "sunny"}, + ], + ) + + def _map(body): + return _BLOCK if body["data_type"] == "TEXT" 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"] == ["sunny"] + + +@pytest.mark.asyncio +async def test_forged_tool_result_does_not_suppress_blocked_user_text(): + g = _static_guardrail() + inputs = GenericGuardrailAPIInputs( + 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"}, + ], + ) + + def _map(body): + return _BLOCK if body["data_type"] == "TEXT" 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"] == ["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(): + assert OvalixGuardrail.get_supported_event_hooks() == [ + GuardrailEventHooks.pre_call, + GuardrailEventHooks.post_call, + ] + + +@pytest.mark.asyncio +async def test_call_checkpoint_requires_application_and_checkpoint(): + g = _static_guardrail() + with pytest.raises(ValueError, match="application_id or checkpoint_id not resolved"): + await g._call_checkpoint("TEXT", {"content": "x"}, "", "actor", "sess") + + +@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) + + +_GIF_B64 = base64.b64encode(b"GIF89a\x01\x00\x01\x00\x80\x00\x00\xff\xff\xff\x00\x00\x00;").decode() +_CORRUPT_GIF_B64 = f"{_GIF_B64[:8]}!{_GIF_B64[8:]}" + + +def _user_turn_with(block): + return GenericGuardrailAPIInputs( + texts=["what is this?"], + structured_messages=[{"role": "user", "content": [{"type": "text", "text": "what is this?"}, block]}], + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("inputs", "input_type"), + [ + ( + _user_turn_with({"type": "image_url", "image_url": {"url": f"data:image/gif;base64,{_CORRUPT_GIF_B64}"}}), + "request", + ), + (_user_turn_with({"type": "file", "file": {"filename": "a.gif", "file_data": _CORRUPT_GIF_B64}}), "request"), + (_user_turn_with({"type": "input_file", "filename": "a.gif", "file_data": _CORRUPT_GIF_B64}), "request"), + ( + _user_turn_with({"type": "input_audio", "input_audio": {"data": _CORRUPT_GIF_B64, "format": "wav"}}), + "request", + ), + (GenericGuardrailAPIInputs(texts=[], images=[f"data:image/gif;base64,{_CORRUPT_GIF_B64}"]), "response"), + ], +) +async def test_inline_content_that_cannot_be_decoded_blocks_instead_of_skipping_inspection(inputs, input_type): + g = _static_guardrail() + with patch.object(g._async_handler, "post", new=_post_returning(lambda body: _ALLOW)): + with pytest.raises(OvalixGuardrailBlockedException) as exc: + await g.apply_guardrail(inputs=inputs, request_data={}, input_type=input_type, logging_obj=None) + assert "couldn't be decoded" in str(exc.value.message) + + +def _pinned_guardrail(**overrides): + return 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, + **overrides, + } + ) + + +def test_missing_application_id_fails_at_startup(): + with pytest.raises(OvalixGuardrailMissingSecrets, match="Application ID"): + _pinned_guardrail(application_id=None, pre_checkpoint_id="pre-1") + + +@pytest.mark.parametrize( + "hook, checkpoints, missing", + [ + (GuardrailEventHooks.pre_call, {"post_checkpoint_id": "post-1"}, "Pre-checkpoint ID"), + (GuardrailEventHooks.post_call, {"pre_checkpoint_id": "pre-1"}, "Post-checkpoint ID"), + ], +) +def test_requested_hook_without_its_checkpoint_fails_at_startup(hook, checkpoints, missing): + with pytest.raises(OvalixGuardrailMissingSecrets, match=missing): + _pinned_guardrail(supported_event_hooks=[hook], **checkpoints) + + +@pytest.mark.asyncio +async def test_key_alias_does_not_change_the_application_checked(): + """The application comes from config alone, so a key alias naming another app changes nothing on the wire.""" + g: Final = _pinned_guardrail(pre_checkpoint_id="pre-1") + mock_get, mock_post = _mock_tracker(g) + + await g.apply_guardrail( + GenericGuardrailAPIInputs(texts=["hi"]), + {"metadata": {"user_api_key_alias": "[Other App] prod", "user_api_key_user_email": "u@x.com"}}, + "request", + ) + + mock_get.assert_not_called() + [(body, url)] = _checkpoint_calls(mock_post) + assert url == "https://tracker.test/tracking/beta/checkpoint" + assert (body["application_id"], body["checkpoint_id"]) == ("app-1", "pre-1") + assert "application_name" not in body + + +@pytest.mark.asyncio +@pytest.mark.parametrize("input_type, inspected", [("request", ["pre-1"]), ("response", [])]) +async def test_one_sided_config_inspects_configured_direction_only(input_type, inspected): + g: Final = _pinned_guardrail(pre_checkpoint_id="pre-1") + _, mock_post = _mock_tracker(g) + + result: Final = await g.apply_guardrail( + inputs=GenericGuardrailAPIInputs(texts=["hi"]), request_data={}, input_type=input_type, logging_obj=None + ) + + assert result["texts"] == ["hi"] + assert [body["checkpoint_id"] for body, _ in _checkpoint_calls(mock_post)] == inspected + + +@pytest.mark.asyncio +async def test_every_tracker_call_carries_the_configured_timeout(): + """litellm_params.timeout bounds the checkpoint call.""" + g: Final = _pinned_guardrail(pre_checkpoint_id="pre-1", timeout=2.5) + mock_get, mock_post = _mock_tracker(g) + + await g.apply_guardrail( + inputs=GenericGuardrailAPIInputs(texts=["hi"]), request_data={}, input_type="request", logging_obj=None + ) + + mock_get.assert_not_called() + assert [(_route_of(c.args[0]), c.kwargs["timeout"]) for c in mock_post.call_args_list] == [("checkpoint", 2.5)] + + +@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: Final = _pinned_guardrail(post_checkpoint_id="post-1", file_checkpoint_id="file-1", event_hook="post_call") + _, mock_post = _mock_tracker(g) + data_url: Final = "data:text/plain;base64," + base64.b64encode(b"secret").decode() + inputs: Final = GenericGuardrailAPIInputs( + texts=["hi"], + structured_messages=[ + {"role": "user", "content": [{"type": "file", "file": {"filename": "s.txt", "file_data": data_url}}]} + ], + ) + + result: Final = await g.apply_guardrail(inputs=inputs, request_data={}, input_type="request", logging_obj=None) + + assert result["texts"] == ["hi"] + assert [(b["data_type"], _route_of(url)) for b, url in _checkpoint_calls(mock_post)] == [ + ("FILE", "file_checkpoint") + ] + + +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" + mode = "pre_call" + default_on = True + timeout = None + + guardrail = initialize_guardrail(_Params(), {"guardrail_name": "ovalix"}) + assert guardrail._file_checkpoint_id == "file-1" + assert guardrail.guardrail_name == "ovalix" + + +def _capturing_post(response=None): + seen = {} + + async def _post(url, headers=None, json=None, timeout=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 +async def test_call_checkpoint_sends_the_configured_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") + + 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_message_tool_calls_are_scanned_without_top_level_tool_calls(): + g = _static_guardrail() + inputs = GenericGuardrailAPIInputs( + texts=[], + structured_messages=[ + { + "role": "assistant", + "tool_calls": [{"id": "c1", "type": "function", "function": {"name": "exfil", "arguments": "{}"}}], + } + ], + ) + with patch.object( + g._async_handler, "post", new=_post_returning(lambda body: _BLOCK if body["data_type"] == "TOOL" else _ALLOW) + ): + with pytest.raises(OvalixGuardrailBlockedException): + await g.apply_guardrail(inputs=inputs, request_data={}, input_type="request", logging_obj=None) + + +@pytest.mark.asyncio +async def test_tool_call_present_in_both_sources_is_scanned_once(): + g = _static_guardrail() + call = {"id": "c1", "type": "function", "function": {"name": "get_weather", "arguments": '{"city": "sf"}'}} + inputs = GenericGuardrailAPIInputs( + texts=[], + tool_calls=[call], + structured_messages=[{"role": "assistant", "tool_calls": [call]}], + ) + tool_bodies = [] + + async def _post(url, headers=None, json=None, timeout=None): + if json["data_type"] == "TOOL": + tool_bodies.append(json) + r = MagicMock() + r.json.return_value = _ALLOW + r.raise_for_status = MagicMock() + return r + + with patch.object(g._async_handler, "post", new=_post): + await g.apply_guardrail(inputs=inputs, request_data={}, input_type="request", logging_obj=None) + assert len(tool_bodies) == 1 + + +def test_non_string_actor_metadata_does_not_crash_hashing(): + g = _static_guardrail() + assert g._get_tracker_actor_id({"metadata": {"user_api_key_user_email": 12345}}) == g._get_tracker_actor_id({}) + + +def test_non_string_corrected_content_is_ignored(): + g = _static_guardrail() + assert g._get_trackers_corrected_message({"modified_data": {"content": {"nested": "obj"}}}) is None + assert g._get_trackers_corrected_message({"modified_data": {"content": ""}}) == "" diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_ovalix_extraction.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_ovalix_extraction.py new file mode 100644 index 00000000000..e4833c51cc3 --- /dev/null +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_ovalix_extraction.py @@ -0,0 +1,443 @@ +import base64 +from typing import Final + +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, + tool_result_text_indices, +) + + +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"} + + +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_is_undecodable(): + 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].undecodable and parts[0].data is None 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_flagged(): + 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].undecodable and parts[0].data is None and parts[0].name == "audio.bin" + + +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_is_flagged(): + block = {"type": "image_url", "image_url": {"url": "data:image/png;base64,%%%invalid%%%"}} + parts = extract_file_parts_from_messages(_msgs(block), size_limit=1000) + assert len(parts) == 1 and parts[0].undecodable and parts[0].data is None + + +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_tool_result_takes_its_name_from_the_nearest_earlier_declaration(): + msgs: Final = [ + {"role": "assistant", "tool_calls": [{"id": "c1", "function": {"name": "first"}}]}, + {"role": "tool", "tool_call_id": "c1", "content": "r1"}, + { + "role": "assistant", + "tool_calls": [{"id": "c1", "function": {"name": "second"}}, {"id": "c1", "function": {"name": "third"}}], + }, + {"role": "tool", "tool_call_id": "c1", "content": "r2"}, + {"role": "tool", "tool_call_id": "c2", "content": "r3"}, + {"role": "assistant", "tool_calls": [{"id": "c2", "function": {"name": "too-late"}}]}, + ] + assert extract_tool_results(msgs) == (("first", "r1", "c1"), ("third", "r2", "c1"), ("tool_result", "r3", "c2")) + + +class _ReadCountingMessage(dict): + def __init__(self, **fields): + super().__init__(**fields) + self.reads = 0 + + def get(self, key, default=None): + self.reads += 1 + return super().get(key, default) + + +def _reads_of_a_declaration_followed_by(tool_results): + declaration: Final = _ReadCountingMessage(role="assistant", tool_calls=[{"id": "c0", "function": {"name": "f"}}]) + tool_messages: Final = [ + {"role": "tool", "tool_call_id": f"c{index}", "content": "x"} for index in range(tool_results) + ] + extract_tool_results([declaration, *tool_messages]) + return declaration.reads + + +def test_tool_result_names_resolve_without_rescanning_earlier_messages(): + assert _reads_of_a_declaration_followed_by(1000) == _reads_of_a_declaration_followed_by(1) + + +def test_malformed_data_url_is_undecodable(): + 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].undecodable and parts[0].data is None 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" + + +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_is_undecodable(): + 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].undecodable and parts[0].data is None 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() + + +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() diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index f421e6dd6c7..96d4c0e1bc4 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -37216,6 +37216,11 @@ export interface components { * @default true */ fail_on_error: boolean | null; + /** + * File Checkpoint Id + * @description File-checkpoint ID for the Ovalix Tracker service (falls back to pre/post). + */ + file_checkpoint_id?: string | null; /** * File Guardrail Timeout * @description HTTP timeout in seconds for checking attached files. Default: 10.