mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-01 07:03:02 +00:00
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 5ffc268585, 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 <shalom@ovalix.io>
This commit is contained in:
yucheng-berri
GitHub
parent
901b7f762f
commit
0f4d5c6ad0
9 changed files with 2298 additions and 199 deletions
|
|
@ -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": [
|
||||
{
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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<mime>[^;,]+)?(?P<params>(?:;[^;,]+)*?)(?P<b64>;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)
|
||||
|
|
@ -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:
|
||||
|
|
|
|||
338
tests/integration/observability/test_ovalix_guardrail.py
Normal file
338
tests/integration/observability/test_ovalix_guardrail.py
Normal file
|
|
@ -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)
|
||||
File diff suppressed because it is too large
Load diff
|
|
@ -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()
|
||||
5
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
5
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -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.
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue