mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
feat(guardrails): add payload, applicability, dispatch and logging controls to generic_guardrail_api
Guardrail providers that are observe-only, that ignore parts of the payload, or that only care about some traffic currently have no way to say so: LiteLLM sends the whole payload for every call type, blocks the request until the guardrail answers, and appends a logging entry on every single call. On long agent sessions that means the entire transcript is re-sent every turn, synchronously, on the user's critical path, and one spend-log row accumulates thousands of guardrail entries. Seven additive options, all defaulting to today's behavior: - fire_and_forget dispatches the call as a background task and lets the request proceed, which makes the guardrail observe-only (BLOCKED and GUARDRAIL_INTERVENED can no longer be honored, and it says so at init). fire_and_forget_max_inflight bounds the queue so a slow endpoint cannot pile up tasks; in-flight tasks are strong-referenced so the event loop cannot collect them mid-call. Streaming is forced to end-of-stream observation, so a stream dispatches one call instead of one per sampled chunk. - send_images, exclude_payload_fields, max_messages and max_text_chars declare what the endpoint receives. send_images=false also replaces inline image data in structured_messages with a marker, so the base64 payload really does stay in LiteLLM. max_messages windows the flat texts list as well as structured_messages, since the translation handlers build texts from the whole conversation and windowing one without the other would leave payload size proportional to session length. input_type and litellm_call_id cannot be excluded. - strip_patterns removes configured regex matches from text before the payload is sent, touching only string text fields, never JSON structure, tool schemas, ids or metadata. - skip_if_system_prompt_matches and skip_if_first_role_in skip a call entirely, anchored to the system/developer message so pasted user content cannot trigger them. Both read the request body, which the caller controls, so a caller that knows the configured value can exempt itself; that is stated in the field descriptions and warned about at init. skip_if_key_alias_in and skip_if_team_id_in are the alternative that holds against the caller, matching on what authentication established. The message-based decision is replayed on the paired response through the shared logging object, with a bounded call-id cache for paths that have none; the identity filters need no correlation, since both hooks see the same metadata. - run_only_on_call_types and skip_call_types skip whole call types (embeddings, image, audio, rerank, moderation) before anything is sent. An unresolvable call type still runs, so a misconfiguration cannot silently blind the guardrail. - guardrail_information_scope picks how often a StandardLoggingGuardrailInformation entry is recorded: per_call keeps today's behavior, per_session records only the first call of a session, off records nothing on success. Session dedup needs a resolvable session id and falls back to per_call without one, and the session set is bounded by size and TTL. Suppression rides the log_guardrail_information contract via a new named entry point, suppress_guardrail_information_record, and is applied only after a normal return, so blocks and guardrail failures still record under every scope. Payload shaping is lossy, so each shaped payload carries the loss with it and the guardrail refuses to write back text the endpoint never saw in full: a stripped, truncated or windowed block keeps the caller's original content, and images or tools that were not sent cannot be replaced. Shaping runs on the dumped JSON because pydantic validates structured_messages content lazily into a ValidatorIterator that is unsafe to walk or re-iterate. The message-mapping helpers live in the shared guardrails content utils, next to the existing walkers, so the Chat Completions and Responses content shapes are handled in one place.
This commit is contained in:
Itay Ovadia
parent
abddd64285
commit
c3ad8ac34e
11 changed files with 2463 additions and 39 deletions
|
|
@ -65,6 +65,20 @@ _guardrail_self_recorded: Final[contextvars.ContextVar[bool]] = contextvars.Cont
|
|||
)
|
||||
|
||||
|
||||
def suppress_guardrail_information_record() -> None:
|
||||
"""Skip the automatic ``StandardLoggingGuardrailInformation`` entry for this call.
|
||||
|
||||
``log_guardrail_information`` clears the flag before invoking the guardrail and
|
||||
skips its own append when the flag is set afterwards, so this is the supported
|
||||
way for a guardrail to opt one invocation out of recording (for instance a
|
||||
guardrail configured to record once per session rather than per call).
|
||||
|
||||
Call it only on a success path: the decorator's exception branch reads the same
|
||||
flag, so setting it before a raise also drops the error entry.
|
||||
"""
|
||||
_guardrail_self_recorded.set(True)
|
||||
|
||||
|
||||
def _strict_guardrail_modes_enabled() -> bool:
|
||||
"""Whether guardrail-mode validation raises (default) or logs a warning.
|
||||
|
||||
|
|
|
|||
|
|
@ -8,8 +8,10 @@ skip the other shapes — these helpers normalise that so every hook sees
|
|||
every text fragment.
|
||||
"""
|
||||
|
||||
from collections.abc import Callable, Iterator
|
||||
from typing import Any, Final
|
||||
from collections.abc import Callable, Collection, Iterator, Mapping, Sequence
|
||||
from typing import Any, Final, TypeAlias
|
||||
|
||||
from pydantic import JsonValue
|
||||
|
||||
# Call types whose body carries free-form chat / prompt text that
|
||||
# text-content guardrails (banned keywords, content moderation, secret
|
||||
|
|
@ -35,6 +37,12 @@ def is_text_content_call_type(call_type: str) -> bool:
|
|||
|
||||
TEXT_PART_TYPES: Final[frozenset[str]] = frozenset({"text", "input_text", "output_text"})
|
||||
|
||||
# A text rewrite applied to one fragment at a time.
|
||||
TextTransform: TypeAlias = Callable[[str], str]
|
||||
|
||||
# Content part types whose payload is image data (Chat Completions and Responses API).
|
||||
IMAGE_PART_TYPES: Final[frozenset[str]] = frozenset({"image_url", "input_image"})
|
||||
|
||||
# Responses-API item types whose ``output`` field carries user/tool text
|
||||
# that guardrails should inspect. ``function_call_output`` is the
|
||||
# built-in shape; ``custom_tool_call_output`` is the custom-tool
|
||||
|
|
@ -104,6 +112,111 @@ def iter_message_text(data: dict[str, Any]) -> Iterator[str]:
|
|||
yield from _iter_text_parts_in_content(message.get("content"))
|
||||
|
||||
|
||||
def iter_role_text(messages: Sequence[Mapping[str, object]], roles: Collection[str]) -> Iterator[str]:
|
||||
"""Yield every text fragment carried by messages whose role is in ``roles``.
|
||||
|
||||
Lets a hook anchor a decision to a specific role (e.g. the system /
|
||||
developer instructions) instead of the whole transcript, which would match
|
||||
on any user-pasted content.
|
||||
"""
|
||||
for message in messages:
|
||||
# The annotation is the contract; the values arrive from provider
|
||||
# translation handlers, so a stray non-message entry must not raise.
|
||||
if not isinstance(message, Mapping): # pyright: ignore[reportUnnecessaryIsInstance] # untyped upstream data
|
||||
continue
|
||||
if message.get("role") in roles:
|
||||
yield from _iter_text_parts_in_content(message.get("content"))
|
||||
|
||||
|
||||
def map_content_text(content: JsonValue, transform: TextTransform) -> JsonValue:
|
||||
"""Return ``content`` with every text fragment replaced by ``transform(fragment)``.
|
||||
|
||||
Pure counterpart of ``walk_user_text``: string content, bare strings in a
|
||||
content list, and ``{"type": "text", "text": ...}`` parts are transformed;
|
||||
every other part (images, audio, tool payloads) is passed through by
|
||||
reference so the content structure cannot change shape.
|
||||
"""
|
||||
if isinstance(content, str):
|
||||
return transform(content) if content else content
|
||||
if not isinstance(content, list):
|
||||
return content
|
||||
return [_mapped_text_part(part, transform) for part in content] # mutable-ok: JSON content is a list
|
||||
|
||||
|
||||
def _mapped_text_part(part: JsonValue, transform: TextTransform) -> JsonValue:
|
||||
if isinstance(part, str):
|
||||
# A bare string in a content list is itself a text fragment.
|
||||
return transform(part) if part else part
|
||||
if not isinstance(part, dict) or part.get("type") not in TEXT_PART_TYPES:
|
||||
return part
|
||||
text: Final = part.get("text")
|
||||
if not isinstance(text, str) or not text:
|
||||
return part
|
||||
return {**part, "text": transform(text)} # mutable-ok: JSON parts are dicts
|
||||
|
||||
|
||||
def map_content_image_urls(content: JsonValue, transform: TextTransform) -> JsonValue:
|
||||
"""Return ``content`` with every image payload replaced by ``transform(url)``.
|
||||
|
||||
Covers both image part shapes: Chat Completions ``{"type": "image_url",
|
||||
"image_url": {"url": ...}}`` (and its bare-string variant) and the
|
||||
Responses-API ``{"type": "input_image", "image_url": ...}``. Everything else
|
||||
is passed through by reference.
|
||||
"""
|
||||
if not isinstance(content, list):
|
||||
return content
|
||||
return [_mapped_image_part(part, transform) for part in content] # mutable-ok: JSON content is a list
|
||||
|
||||
|
||||
def _mapped_image_part(part: JsonValue, transform: TextTransform) -> JsonValue:
|
||||
if not isinstance(part, dict) or part.get("type") not in IMAGE_PART_TYPES:
|
||||
return part
|
||||
image_url: Final = part.get("image_url")
|
||||
if isinstance(image_url, str):
|
||||
return {**part, "image_url": transform(image_url)} # mutable-ok: JSON parts are dicts
|
||||
if not isinstance(image_url, dict):
|
||||
return part
|
||||
url: Final = image_url.get("url")
|
||||
if not isinstance(url, str):
|
||||
return part
|
||||
nested: Final = {**image_url, "url": transform(url)} # mutable-ok: JSON parts are dicts
|
||||
return {**part, "image_url": nested} # mutable-ok: JSON parts are dicts
|
||||
|
||||
|
||||
def map_messages_image_urls(messages: JsonValue, transform: TextTransform) -> JsonValue:
|
||||
"""Return a new message list with every image payload transformed."""
|
||||
return _mapped_messages(messages, map_content_image_urls, transform)
|
||||
|
||||
|
||||
def map_messages_text(messages: JsonValue, transform: TextTransform) -> JsonValue:
|
||||
"""Return a new message list with every ``content`` text fragment transformed.
|
||||
|
||||
Roles, ids, tool calls, tool schemas and any other key are copied through
|
||||
untouched, so only text can differ from the input.
|
||||
"""
|
||||
return _mapped_messages(messages, map_content_text, transform)
|
||||
|
||||
|
||||
def _mapped_messages(
|
||||
messages: JsonValue,
|
||||
map_content: Callable[[JsonValue, TextTransform], JsonValue],
|
||||
transform: TextTransform,
|
||||
) -> JsonValue:
|
||||
if not isinstance(messages, list):
|
||||
return messages
|
||||
return [_mapped_message(message, map_content, transform) for message in messages] # mutable-ok: JSON is a list
|
||||
|
||||
|
||||
def _mapped_message(
|
||||
message: JsonValue,
|
||||
map_content: Callable[[JsonValue, TextTransform], JsonValue],
|
||||
transform: TextTransform,
|
||||
) -> JsonValue:
|
||||
if not isinstance(message, dict) or "content" not in message:
|
||||
return message
|
||||
return {**message, "content": map_content(message["content"], transform)} # mutable-ok: JSON messages are dicts
|
||||
|
||||
|
||||
def walk_user_text(data: dict[str, Any], visit: Callable[[str], str]) -> int:
|
||||
"""Rewrite every text fragment in place via ``visit``.
|
||||
|
||||
|
|
@ -113,31 +226,13 @@ def walk_user_text(data: dict[str, Any], visit: Callable[[str], str]) -> int:
|
|||
"""
|
||||
visited = 0
|
||||
|
||||
def _rewrite_content(content: Any) -> Any:
|
||||
def _counting_visit(text: str) -> str:
|
||||
nonlocal visited
|
||||
if isinstance(content, str):
|
||||
if content:
|
||||
visited += 1
|
||||
return visit(content)
|
||||
return content
|
||||
if isinstance(content, list):
|
||||
new_parts: Final[list[Any]] = []
|
||||
for part in content:
|
||||
if isinstance(part, str) and part:
|
||||
visited += 1
|
||||
new_parts.append(visit(part))
|
||||
elif (
|
||||
isinstance(part, dict)
|
||||
and part.get("type") in TEXT_PART_TYPES
|
||||
and isinstance(part.get("text"), str)
|
||||
and part["text"]
|
||||
):
|
||||
visited += 1
|
||||
new_parts.append({**part, "text": visit(part["text"])})
|
||||
else:
|
||||
new_parts.append(part)
|
||||
return new_parts
|
||||
return content
|
||||
visited += 1
|
||||
return visit(text)
|
||||
|
||||
def _rewrite_content(content: Any) -> Any:
|
||||
return map_content_text(content, _counting_visit)
|
||||
|
||||
messages: Final = data.get("messages")
|
||||
if isinstance(messages, list):
|
||||
|
|
|
|||
|
|
@ -39,6 +39,22 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"
|
|||
streaming_end_of_stream_only=_get_config_value(litellm_params, optional_params, "streaming_end_of_stream_only"),
|
||||
streaming_sampling_rate=_get_config_value(litellm_params, optional_params, "streaming_sampling_rate"),
|
||||
streaming_transform_mode=_get_config_value(litellm_params, optional_params, "streaming_transform_mode"),
|
||||
fire_and_forget=_get_config_value(litellm_params, optional_params, "fire_and_forget"),
|
||||
fire_and_forget_max_inflight=_get_config_value(litellm_params, optional_params, "fire_and_forget_max_inflight"),
|
||||
send_images=_get_config_value(litellm_params, optional_params, "send_images"),
|
||||
exclude_payload_fields=_get_config_value(litellm_params, optional_params, "exclude_payload_fields"),
|
||||
max_messages=_get_config_value(litellm_params, optional_params, "max_messages"),
|
||||
max_text_chars=_get_config_value(litellm_params, optional_params, "max_text_chars"),
|
||||
strip_patterns=_get_config_value(litellm_params, optional_params, "strip_patterns"),
|
||||
skip_if_system_prompt_matches=_get_config_value(
|
||||
litellm_params, optional_params, "skip_if_system_prompt_matches"
|
||||
),
|
||||
skip_if_first_role_in=_get_config_value(litellm_params, optional_params, "skip_if_first_role_in"),
|
||||
skip_if_key_alias_in=_get_config_value(litellm_params, optional_params, "skip_if_key_alias_in"),
|
||||
skip_if_team_id_in=_get_config_value(litellm_params, optional_params, "skip_if_team_id_in"),
|
||||
run_only_on_call_types=_get_config_value(litellm_params, optional_params, "run_only_on_call_types"),
|
||||
skip_call_types=_get_config_value(litellm_params, optional_params, "skip_call_types"),
|
||||
guardrail_information_scope=_get_config_value(litellm_params, optional_params, "guardrail_information_scope"),
|
||||
)
|
||||
|
||||
litellm.logging_callback_manager.add_litellm_callback(_generic_guardrail_api_callback)
|
||||
|
|
|
|||
|
|
@ -0,0 +1,85 @@
|
|||
"""Off-critical-path dispatch for the Generic Guardrail API (``fire_and_forget``).
|
||||
|
||||
The event loop only holds a weak reference to a task created with
|
||||
``asyncio.create_task``, so a task with no other referrer can be garbage
|
||||
collected mid-flight and silently cancelled. Every dispatched task is therefore
|
||||
kept in a strong-ref set until it completes.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from collections.abc import Awaitable, Callable
|
||||
from typing import Final
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
||||
DEFAULT_FIRE_AND_FORGET_MAX_INFLIGHT: Final = 100
|
||||
|
||||
# Log one warning per this many drops, so a saturated endpoint cannot turn into a
|
||||
# log flood of its own.
|
||||
_DROP_LOG_INTERVAL: Final = 100
|
||||
|
||||
|
||||
class BackgroundDispatcher:
|
||||
"""Runs guardrail calls detached from the request, with a bounded queue.
|
||||
|
||||
Async dispatch decouples the request rate from the endpoint's throughput, so
|
||||
a slow endpoint would otherwise pile up tasks without limit. Once
|
||||
``max_inflight`` tasks are outstanding, further calls are dropped and
|
||||
counted rather than queued.
|
||||
"""
|
||||
|
||||
def __init__(self, *, guardrail_name: str | None, max_inflight: int) -> None:
|
||||
if max_inflight < 1:
|
||||
raise ValueError(f"fire_and_forget_max_inflight must be >= 1 (got {max_inflight})")
|
||||
self._guardrail_name: Final = guardrail_name
|
||||
self._max_inflight: Final = max_inflight
|
||||
# Strong refs to in-flight tasks; entries are discarded on completion.
|
||||
self._pending: set[asyncio.Task[None]] = set() # mutable-ok: tasks are added/removed as they complete
|
||||
self._dropped: int = 0
|
||||
|
||||
@property
|
||||
def pending_count(self) -> int:
|
||||
return len(self._pending)
|
||||
|
||||
@property
|
||||
def dropped_count(self) -> int:
|
||||
return self._dropped
|
||||
|
||||
def dispatch(self, run: Callable[[], Awaitable[None]], *, context: str) -> bool:
|
||||
"""Start ``run`` detached. Returns False when dropped for backpressure."""
|
||||
if len(self._pending) >= self._max_inflight:
|
||||
self._dropped += 1
|
||||
if self._dropped % _DROP_LOG_INTERVAL == 1:
|
||||
verbose_proxy_logger.warning(
|
||||
"Generic Guardrail API (%s, fire_and_forget): dropped %d call(s) so far; "
|
||||
"%d already in flight (fire_and_forget_max_inflight=%d). %s",
|
||||
self._guardrail_name,
|
||||
self._dropped,
|
||||
len(self._pending),
|
||||
self._max_inflight,
|
||||
context,
|
||||
)
|
||||
return False
|
||||
|
||||
task: Final = asyncio.create_task(self._guarded(run, context=context))
|
||||
self._pending.add(task)
|
||||
task.add_done_callback(self._pending.discard)
|
||||
return True
|
||||
|
||||
async def _guarded(self, run: Callable[[], Awaitable[None]], *, context: str) -> None:
|
||||
"""Never raise: a detached task's exception has no frame to surface in."""
|
||||
try:
|
||||
await run()
|
||||
except Exception as e: # noqa: BLE001 # a detached task has no frame to raise into
|
||||
verbose_proxy_logger.warning(
|
||||
"Generic Guardrail API (%s, fire_and_forget) call failed. %s: %s",
|
||||
self._guardrail_name,
|
||||
context,
|
||||
e,
|
||||
)
|
||||
|
||||
async def wait_for_pending(self) -> None:
|
||||
"""Await every in-flight task. Used by tests and best-effort teardown."""
|
||||
pending: Final = tuple(self._pending)
|
||||
if pending:
|
||||
await asyncio.gather(*pending, return_exceptions=True)
|
||||
|
|
@ -7,22 +7,27 @@
|
|||
|
||||
import fnmatch
|
||||
import os
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, Optional
|
||||
|
||||
import httpx
|
||||
from pydantic import JsonValue
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm._version import version as litellm_version
|
||||
from litellm.exceptions import GuardrailRaisedException, Timeout
|
||||
from litellm.integrations.custom_guardrail import (
|
||||
CustomGuardrail,
|
||||
get_session_id_from_request_data,
|
||||
log_guardrail_information,
|
||||
suppress_guardrail_information_record,
|
||||
)
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
AsyncHTTPHandler,
|
||||
get_async_httpx_client,
|
||||
httpxSpecialProvider,
|
||||
)
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
from litellm.types.guardrails import GuardrailEventHooks, Mode
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import (
|
||||
GenericGuardrailAPIMetadata,
|
||||
GenericGuardrailAPIRequest,
|
||||
|
|
@ -31,6 +36,32 @@ from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import
|
|||
)
|
||||
from litellm.types.utils import GenericGuardrailAPIInputs
|
||||
|
||||
from .background_dispatch import (
|
||||
DEFAULT_FIRE_AND_FORGET_MAX_INFLIGHT,
|
||||
BackgroundDispatcher,
|
||||
)
|
||||
from .payload_policy import (
|
||||
PayloadLoss,
|
||||
PayloadPolicy,
|
||||
compile_patterns,
|
||||
merge_guardrailed_texts,
|
||||
resolve_exclude_fields,
|
||||
shape_payload,
|
||||
)
|
||||
from .record_scope import (
|
||||
DEFAULT_GUARDRAIL_INFORMATION_SCOPE,
|
||||
GuardrailInformationScope,
|
||||
RecordScope,
|
||||
)
|
||||
from .request_filters import (
|
||||
SkipDecisionStore,
|
||||
SkipPolicy,
|
||||
call_type_allowed,
|
||||
identity_matches_skip,
|
||||
request_matches_skip,
|
||||
validate_call_types,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel
|
||||
|
|
@ -148,6 +179,42 @@ def _extract_inbound_headers(
|
|||
return None
|
||||
|
||||
|
||||
def _resolve_call_type(
|
||||
request_data: Mapping[str, object],
|
||||
logging_obj: Optional["LiteLLMLoggingObj"],
|
||||
) -> str | None:
|
||||
"""Resolve the call type of the current request.
|
||||
|
||||
The proxy passes the route type straight through to ``function_setup``, which
|
||||
stamps it on the logging object before the pre-call hook runs and keeps it
|
||||
for the post-call hook, so both sides of a call resolve the same value.
|
||||
"""
|
||||
from_logging: Final = getattr(logging_obj, "call_type", None) if logging_obj else None
|
||||
if isinstance(from_logging, str) and from_logging:
|
||||
return from_logging
|
||||
from_request: Final = request_data.get("call_type")
|
||||
return from_request if isinstance(from_request, str) and from_request else None
|
||||
|
||||
|
||||
def _passthrough_inputs(inputs: GenericGuardrailAPIInputs) -> GenericGuardrailAPIInputs:
|
||||
"""Return the inputs untouched (same value identities), as action=NONE."""
|
||||
return GenericGuardrailAPIInputs(**inputs)
|
||||
|
||||
|
||||
def _has_request_side_hook(event_hook: str | Sequence[str] | Mode | None) -> bool:
|
||||
"""Whether the configured mode(s) include a hook that sees the request.
|
||||
|
||||
Tag-based ``Mode`` config is resolved per request, so it is treated as having
|
||||
one rather than emitting a warning that may not apply.
|
||||
"""
|
||||
request_side: Final = frozenset({GuardrailEventHooks.pre_call.value, GuardrailEventHooks.during_call.value})
|
||||
if event_hook is None or isinstance(event_hook, Mode):
|
||||
return True
|
||||
if isinstance(event_hook, str):
|
||||
return event_hook in request_side
|
||||
return any(hook in request_side for hook in event_hook)
|
||||
|
||||
|
||||
class GenericGuardrailAPI(CustomGuardrail):
|
||||
"""
|
||||
Generic Guardrail API integration for LiteLLM.
|
||||
|
|
@ -182,9 +249,27 @@ class GenericGuardrailAPI(CustomGuardrail):
|
|||
streaming_end_of_stream_only: bool | None = None,
|
||||
streaming_sampling_rate: int | None = None,
|
||||
streaming_transform_mode: Literal["block_only", "incremental_diff"] | None = None,
|
||||
fire_and_forget: bool | None = None,
|
||||
fire_and_forget_max_inflight: int | None = None,
|
||||
send_images: bool | None = None,
|
||||
exclude_payload_fields: Sequence[str] | None = None,
|
||||
max_messages: int | None = None,
|
||||
max_text_chars: int | None = None,
|
||||
strip_patterns: Sequence[str] | None = None,
|
||||
skip_if_system_prompt_matches: Sequence[str] | None = None,
|
||||
skip_if_first_role_in: Sequence[str] | None = None,
|
||||
skip_if_key_alias_in: Sequence[str] | None = None,
|
||||
skip_if_team_id_in: Sequence[str] | None = None,
|
||||
run_only_on_call_types: Sequence[str] | None = None,
|
||||
skip_call_types: Sequence[str] | None = None,
|
||||
guardrail_information_scope: GuardrailInformationScope | None = None,
|
||||
async_handler: AsyncHTTPHandler | None = None,
|
||||
dispatcher: BackgroundDispatcher | None = None,
|
||||
**kwargs,
|
||||
):
|
||||
self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback)
|
||||
self.async_handler = async_handler or get_async_httpx_client(
|
||||
llm_provider=httpxSpecialProvider.GuardrailCallback
|
||||
)
|
||||
self.headers = headers or {}
|
||||
self.extra_headers = extra_headers or []
|
||||
|
||||
|
|
@ -229,6 +314,93 @@ class GenericGuardrailAPI(CustomGuardrail):
|
|||
"block_only" if streaming_transform_mode is None else streaming_transform_mode
|
||||
)
|
||||
|
||||
configured_name: Final = kwargs.get("guardrail_name")
|
||||
|
||||
self.fire_and_forget: bool = False if fire_and_forget is None else fire_and_forget
|
||||
self._dispatcher: Final = dispatcher or BackgroundDispatcher(
|
||||
guardrail_name=configured_name,
|
||||
max_inflight=(
|
||||
DEFAULT_FIRE_AND_FORGET_MAX_INFLIGHT
|
||||
if fire_and_forget_max_inflight is None
|
||||
else fire_and_forget_max_inflight
|
||||
),
|
||||
)
|
||||
|
||||
self._payload_policy: Final = PayloadPolicy(
|
||||
send_images=True if send_images is None else send_images,
|
||||
exclude_fields=resolve_exclude_fields(exclude_payload_fields, guardrail_name=configured_name),
|
||||
max_messages=max_messages,
|
||||
max_text_chars=max_text_chars,
|
||||
strip_patterns=compile_patterns(strip_patterns, option_name="strip_patterns"),
|
||||
)
|
||||
|
||||
self._skip_policy: Final = SkipPolicy(
|
||||
system_prompt_patterns=compile_patterns(
|
||||
skip_if_system_prompt_matches, option_name="skip_if_system_prompt_matches"
|
||||
),
|
||||
first_role_in=frozenset(skip_if_first_role_in or ()),
|
||||
key_aliases=frozenset(skip_if_key_alias_in or ()),
|
||||
team_ids=frozenset(skip_if_team_id_in or ()),
|
||||
run_only_on_call_types=(
|
||||
validate_call_types(
|
||||
run_only_on_call_types,
|
||||
option_name="run_only_on_call_types",
|
||||
guardrail_name=configured_name,
|
||||
)
|
||||
if run_only_on_call_types
|
||||
else None
|
||||
),
|
||||
skip_call_types=validate_call_types(
|
||||
skip_call_types, option_name="skip_call_types", guardrail_name=configured_name
|
||||
),
|
||||
)
|
||||
self._skip_store: Final = SkipDecisionStore(guardrail_name=configured_name)
|
||||
|
||||
self._record_scope: Final = RecordScope(
|
||||
DEFAULT_GUARDRAIL_INFORMATION_SCOPE if guardrail_information_scope is None else guardrail_information_scope
|
||||
)
|
||||
|
||||
if self.fire_and_forget:
|
||||
# Nothing awaits the response, so a block cannot be honored and the
|
||||
# stream can only be observed. Say so at boot rather than surprising
|
||||
# the operator with a guardrail that never blocks.
|
||||
verbose_proxy_logger.warning(
|
||||
"Generic Guardrail API (%s): fire_and_forget=True makes this guardrail observe-only. "
|
||||
"action=BLOCKED and action=GUARDRAIL_INTERVENED are ignored, and "
|
||||
"fail_on_error=%s / unreachable_fallback=%s cannot block the request. "
|
||||
"Streaming is forced to end-of-stream observation.",
|
||||
configured_name,
|
||||
self.fail_on_error,
|
||||
self.unreachable_fallback,
|
||||
)
|
||||
self.streaming_end_of_stream_only = True
|
||||
|
||||
if self._skip_policy.filters_requests:
|
||||
verbose_proxy_logger.warning(
|
||||
"Generic Guardrail API (%s): skip_if_system_prompt_matches / skip_if_first_role_in match on the "
|
||||
"request body, which the caller controls, so a caller that knows the configured value can exempt "
|
||||
"itself from this guardrail. Use skip_if_key_alias_in / skip_if_team_id_in when the exemption must "
|
||||
"hold against the caller.",
|
||||
configured_name,
|
||||
)
|
||||
|
||||
if self._skip_policy.filters_requests and not _has_request_side_hook(kwargs.get("event_hook")):
|
||||
verbose_proxy_logger.warning(
|
||||
"Generic Guardrail API (%s): skip_if_system_prompt_matches / skip_if_first_role_in need a "
|
||||
"request-side hook (pre_call or during_call) to decide anything. mode=%s only sees "
|
||||
"responses, so nothing will be skipped.",
|
||||
configured_name,
|
||||
kwargs.get("event_hook"),
|
||||
)
|
||||
|
||||
if self._skip_policy.run_only_on_call_types is not None and self._skip_policy.skip_call_types:
|
||||
verbose_proxy_logger.warning(
|
||||
"Generic Guardrail API (%s): both run_only_on_call_types and skip_call_types are set. "
|
||||
"The allowlist wins; skip_call_types=%s is ignored.",
|
||||
configured_name,
|
||||
sorted(self._skip_policy.skip_call_types),
|
||||
)
|
||||
|
||||
# Set supported event hooks
|
||||
kwargs.setdefault("supported_event_hooks", list(self.get_supported_event_hooks()))
|
||||
|
||||
|
|
@ -303,9 +475,7 @@ class GenericGuardrailAPI(CustomGuardrail):
|
|||
exc_info=error,
|
||||
)
|
||||
# Keep flow going - treat as action=NONE (no modifications)
|
||||
return_inputs: Final[GenericGuardrailAPIInputs] = {}
|
||||
return_inputs.update(inputs)
|
||||
return return_inputs
|
||||
return _passthrough_inputs(inputs)
|
||||
|
||||
def _build_request_headers(self) -> dict:
|
||||
"""Build HTTP headers for the guardrail API request."""
|
||||
|
|
@ -321,16 +491,24 @@ class GenericGuardrailAPI(CustomGuardrail):
|
|||
images: Any,
|
||||
tools: Any,
|
||||
guardrail_response: GenericGuardrailAPIResponse,
|
||||
loss: PayloadLoss,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
# Action is NONE or no modifications needed
|
||||
# Action is NONE or no modifications needed. A component the guardrail
|
||||
# never received in full (payload shaping) keeps the caller's original
|
||||
# value: it cannot rewrite what it could not see.
|
||||
return_inputs: Final = GenericGuardrailAPIInputs(texts=texts)
|
||||
if guardrail_response.texts:
|
||||
return_inputs["texts"] = guardrail_response.texts
|
||||
if guardrail_response.images:
|
||||
return_inputs["texts"] = merge_guardrailed_texts(
|
||||
original=texts,
|
||||
returned=guardrail_response.texts,
|
||||
loss=loss,
|
||||
guardrail_name=getattr(self, "guardrail_name", None),
|
||||
)
|
||||
if guardrail_response.images and not loss.images_omitted:
|
||||
return_inputs["images"] = guardrail_response.images
|
||||
elif images:
|
||||
return_inputs["images"] = images
|
||||
if guardrail_response.tools:
|
||||
if guardrail_response.tools and not loss.tools_omitted:
|
||||
return_inputs["tools"] = guardrail_response.tools
|
||||
elif tools:
|
||||
return_inputs["tools"] = tools
|
||||
|
|
@ -359,6 +537,54 @@ class GenericGuardrailAPI(CustomGuardrail):
|
|||
verbose_proxy_logger.error("Generic Guardrail API: failed to make request: %s", str(error))
|
||||
raise Exception(f"Generic Guardrail API failed: {error}")
|
||||
|
||||
def _should_skip_out_of_scope_request(
|
||||
self,
|
||||
*,
|
||||
input_type: Literal["request", "response"],
|
||||
structured_messages: Sequence[Mapping[str, object]] | None,
|
||||
request_data: Mapping[str, object],
|
||||
logging_obj: Optional["LiteLLMLoggingObj"],
|
||||
) -> bool:
|
||||
"""Whether this call is out of scope per skip_if_* (Feature 4).
|
||||
|
||||
Only the request carries the system prompt, so the response side replays
|
||||
the decision the request side recorded instead of re-deciding.
|
||||
"""
|
||||
if not self._skip_policy.filters_requests:
|
||||
return False
|
||||
|
||||
call_id: Final = (getattr(logging_obj, "litellm_call_id", None) if logging_obj else None) or request_data.get(
|
||||
"litellm_call_id"
|
||||
)
|
||||
|
||||
if input_type == "response":
|
||||
return self._skip_store.consume(logging_obj=logging_obj, call_id=call_id)
|
||||
|
||||
if not request_matches_skip(self._skip_policy, structured_messages):
|
||||
return False
|
||||
|
||||
self._skip_store.record(logging_obj=logging_obj, call_id=call_id)
|
||||
return True
|
||||
|
||||
def _dispatch_background_post(
|
||||
self,
|
||||
*,
|
||||
payload: Mapping[str, JsonValue],
|
||||
headers: Mapping[str, str],
|
||||
input_type: Literal["request", "response"],
|
||||
logging_obj: Optional["LiteLLMLoggingObj"],
|
||||
) -> None:
|
||||
context: Final = (
|
||||
f"input_type={input_type} "
|
||||
f"litellm_call_id={getattr(logging_obj, 'litellm_call_id', None) if logging_obj else None}"
|
||||
)
|
||||
|
||||
async def _post() -> None:
|
||||
response: Final = await self.async_handler.post(url=self.api_base, json=payload, headers=headers)
|
||||
response.raise_for_status()
|
||||
|
||||
self._dispatcher.dispatch(_post, context=context)
|
||||
|
||||
@log_guardrail_information
|
||||
async def apply_guardrail(
|
||||
self,
|
||||
|
|
@ -366,6 +592,29 @@ class GenericGuardrailAPI(CustomGuardrail):
|
|||
request_data: dict,
|
||||
input_type: Literal["request", "response"],
|
||||
logging_obj: Optional["LiteLLMLoggingObj"] = None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
"""Run the guardrail, then decide whether this call records a log entry.
|
||||
|
||||
The suppression flag is set only once the call has returned normally, so
|
||||
a block or a guardrail failure still records under every scope: the
|
||||
decorator's exception branch reads the same flag.
|
||||
"""
|
||||
result: Final = await self._apply_guardrail(
|
||||
inputs=inputs,
|
||||
request_data=request_data,
|
||||
input_type=input_type,
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
if self._record_scope.should_suppress(get_session_id_from_request_data(request_data or {})):
|
||||
suppress_guardrail_information_record()
|
||||
return result
|
||||
|
||||
async def _apply_guardrail(
|
||||
self,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
request_data: dict,
|
||||
input_type: Literal["request", "response"],
|
||||
logging_obj: Optional["LiteLLMLoggingObj"] = None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
"""
|
||||
Apply the Generic Guardrail API to the given inputs.
|
||||
|
|
@ -401,6 +650,40 @@ class GenericGuardrailAPI(CustomGuardrail):
|
|||
if request_data is None:
|
||||
request_data = {}
|
||||
|
||||
# Extract user API key metadata (also the basis of the identity filters below)
|
||||
user_metadata: Final = self._extract_user_api_key_metadata(request_data)
|
||||
|
||||
call_type: Final = _resolve_call_type(request_data=request_data, logging_obj=logging_obj)
|
||||
if not call_type_allowed(self._skip_policy, call_type):
|
||||
verbose_proxy_logger.debug(
|
||||
"Generic Guardrail API: skipping call_type=%s (input_type=%s) per call-type filter",
|
||||
call_type,
|
||||
input_type,
|
||||
)
|
||||
return _passthrough_inputs(inputs)
|
||||
|
||||
if identity_matches_skip(self._skip_policy, user_metadata):
|
||||
verbose_proxy_logger.debug(
|
||||
"Generic Guardrail API: skipping out-of-scope caller (input_type=%s, key_alias=%s, team_id=%s)",
|
||||
input_type,
|
||||
user_metadata.get("user_api_key_alias"),
|
||||
user_metadata.get("user_api_key_team_id"),
|
||||
)
|
||||
return _passthrough_inputs(inputs)
|
||||
|
||||
if self._should_skip_out_of_scope_request(
|
||||
input_type=input_type,
|
||||
structured_messages=structured_messages,
|
||||
request_data=request_data,
|
||||
logging_obj=logging_obj,
|
||||
):
|
||||
verbose_proxy_logger.debug(
|
||||
"Generic Guardrail API: skipping out-of-scope request (input_type=%s, litellm_call_id=%s)",
|
||||
input_type,
|
||||
getattr(logging_obj, "litellm_call_id", None) if logging_obj else None,
|
||||
)
|
||||
return _passthrough_inputs(inputs)
|
||||
|
||||
request_body: Final = request_data.get("body") or {}
|
||||
|
||||
# Merge additional provider specific params from config and dynamic params
|
||||
|
|
@ -411,8 +694,6 @@ class GenericGuardrailAPI(CustomGuardrail):
|
|||
if dynamic_params:
|
||||
additional_params.update(dynamic_params)
|
||||
|
||||
# Extract user API key metadata
|
||||
user_metadata: Final = self._extract_user_api_key_metadata(request_data)
|
||||
extra_allowlist = {h.lower() for h in self.extra_headers if isinstance(h, str)} if self.extra_headers else None
|
||||
inbound_headers: Final = _extract_inbound_headers(
|
||||
request_data=request_data,
|
||||
|
|
@ -440,11 +721,25 @@ class GenericGuardrailAPI(CustomGuardrail):
|
|||
|
||||
headers: Final = self._build_request_headers()
|
||||
|
||||
# Shape the payload (send_images / exclude_payload_fields / max_messages /
|
||||
# max_text_chars / strip_patterns) once, so the awaited and the
|
||||
# fire_and_forget paths send exactly the same bytes.
|
||||
# mode="json" ensures all iterables are converted to lists.
|
||||
payload, payload_loss = shape_payload(guardrail_request, self._payload_policy)
|
||||
|
||||
if self.fire_and_forget:
|
||||
self._dispatch_background_post(
|
||||
payload=payload,
|
||||
headers=headers,
|
||||
input_type=input_type,
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
return _passthrough_inputs(inputs)
|
||||
|
||||
# Make the API request
|
||||
# Use mode="json" to ensure all iterables are converted to lists
|
||||
response: Final = await self.async_handler.post(
|
||||
url=self.api_base,
|
||||
json=guardrail_request.model_dump(mode="json"),
|
||||
json=payload,
|
||||
headers=headers,
|
||||
)
|
||||
|
||||
|
|
@ -471,6 +766,7 @@ class GenericGuardrailAPI(CustomGuardrail):
|
|||
images=images,
|
||||
tools=tools,
|
||||
guardrail_response=guardrail_response,
|
||||
loss=payload_loss,
|
||||
)
|
||||
|
||||
except GuardrailRaisedException:
|
||||
|
|
|
|||
|
|
@ -0,0 +1,224 @@
|
|||
"""Payload shaping for the Generic Guardrail API.
|
||||
|
||||
Lets the config declare what the guardrail endpoint actually receives:
|
||||
``send_images``, ``exclude_payload_fields``, ``max_messages``,
|
||||
``max_text_chars`` and ``strip_patterns``.
|
||||
|
||||
Shaping is lossy by design, so every shaped payload carries a ``PayloadLoss``
|
||||
describing what the endpoint could not see. The guardrail refuses to write back
|
||||
mutations for those components: a provider that never received a text block (or
|
||||
received a stripped / truncated one) must not be able to replace the caller's
|
||||
original content with it.
|
||||
"""
|
||||
|
||||
import re
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from functools import reduce
|
||||
from typing import Final
|
||||
|
||||
from pydantic import JsonValue
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.proxy.guardrails._content_utils import (
|
||||
map_messages_image_urls,
|
||||
map_messages_text,
|
||||
)
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import (
|
||||
GenericGuardrailAPIRequest,
|
||||
)
|
||||
|
||||
# Routing / correlation keys the endpoint needs to interpret any payload at all.
|
||||
PROTECTED_PAYLOAD_FIELDS: Final = frozenset({"input_type", "litellm_call_id"})
|
||||
|
||||
# Ceiling on substitutions per pattern per text, so a pathological pattern on a
|
||||
# huge transcript cannot stall the request path.
|
||||
MAX_STRIP_SUBSTITUTIONS: Final = 64
|
||||
|
||||
# Stands in for image data when send_images=False, mirroring the "[present]"
|
||||
# convention used for headers whose value is not forwarded.
|
||||
IMAGE_OMITTED_PLACEHOLDER: Final = "[omitted]"
|
||||
|
||||
_IMAGES_FIELD: Final = frozenset({"images"})
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class PayloadPolicy:
|
||||
"""Resolved payload-shaping configuration."""
|
||||
|
||||
send_images: bool = True
|
||||
exclude_fields: frozenset[str] = frozenset()
|
||||
max_messages: int | None = None
|
||||
max_text_chars: int | None = None
|
||||
strip_patterns: tuple[re.Pattern[str], ...] = ()
|
||||
|
||||
@property
|
||||
def shapes_text(self) -> bool:
|
||||
return self.max_text_chars is not None or bool(self.strip_patterns)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class PayloadLoss:
|
||||
"""What the guardrail endpoint did not get to see."""
|
||||
|
||||
altered_text_indices: frozenset[int] = frozenset()
|
||||
images_omitted: bool = False
|
||||
tools_omitted: bool = False
|
||||
|
||||
|
||||
def resolve_exclude_fields(raw: Sequence[str] | None, *, guardrail_name: str | None) -> frozenset[str]:
|
||||
"""Validate ``exclude_payload_fields`` against the request model.
|
||||
|
||||
Unknown keys and routing-critical keys are dropped with a warning rather
|
||||
than raising, so a stale config never takes the proxy down.
|
||||
"""
|
||||
if not raw:
|
||||
return frozenset()
|
||||
|
||||
known: Final = frozenset(GenericGuardrailAPIRequest.model_fields)
|
||||
unknown: Final = tuple(field for field in raw if field not in known)
|
||||
if unknown:
|
||||
verbose_proxy_logger.warning(
|
||||
"Generic Guardrail API (%s): exclude_payload_fields contains unknown field(s) %s; ignoring them. "
|
||||
"Known fields: %s",
|
||||
guardrail_name,
|
||||
unknown,
|
||||
sorted(known),
|
||||
)
|
||||
|
||||
protected: Final = tuple(field for field in raw if field in PROTECTED_PAYLOAD_FIELDS)
|
||||
if protected:
|
||||
verbose_proxy_logger.warning(
|
||||
"Generic Guardrail API (%s): exclude_payload_fields cannot drop routing-critical field(s) %s; "
|
||||
"they will still be sent.",
|
||||
guardrail_name,
|
||||
protected,
|
||||
)
|
||||
|
||||
return frozenset(field for field in raw if field in known and field not in PROTECTED_PAYLOAD_FIELDS)
|
||||
|
||||
|
||||
def compile_patterns(raw: Sequence[str] | None, *, option_name: str) -> tuple[re.Pattern[str], ...]:
|
||||
"""Compile config-supplied regexes, raising on an invalid pattern.
|
||||
|
||||
A bad regex is a config bug: failing at init surfaces it at proxy boot
|
||||
instead of silently disabling the option on the request path.
|
||||
"""
|
||||
if not raw:
|
||||
return ()
|
||||
try:
|
||||
return tuple(re.compile(pattern) for pattern in raw)
|
||||
except re.error as e:
|
||||
raise ValueError(f"{option_name} contains an invalid regex: {e}") from e
|
||||
|
||||
|
||||
def _strip_text(text: str, patterns: tuple[re.Pattern[str], ...]) -> str:
|
||||
return reduce(lambda acc, pattern: pattern.sub("", acc, count=MAX_STRIP_SUBSTITUTIONS), patterns, text)
|
||||
|
||||
|
||||
def _shape_text(text: str, policy: PayloadPolicy) -> str:
|
||||
stripped: Final = _strip_text(text, policy.strip_patterns)
|
||||
return stripped[: policy.max_text_chars] if policy.max_text_chars is not None else stripped
|
||||
|
||||
|
||||
def _shape_texts(texts: Sequence[str], policy: PayloadPolicy) -> tuple[JsonValue, frozenset[int]]:
|
||||
"""Window and rewrite the flat text list, and report which indices changed.
|
||||
|
||||
``max_messages`` has to bound ``texts`` as well as ``structured_messages``:
|
||||
the translation handlers populate the flat list from the whole conversation,
|
||||
so windowing only the structured form would leave payload size proportional
|
||||
to session length, which is the thing the option exists to stop. Windowing
|
||||
shifts every position, so the caller treats the whole list as unusable for
|
||||
write-back (see ``merge_guardrailed_texts``).
|
||||
"""
|
||||
windowed: Final = texts[-policy.max_messages :] if policy.max_messages is not None else texts
|
||||
shaped: Final = [_shape_text(text, policy) for text in windowed] # mutable-ok: JSON texts is an array
|
||||
if len(windowed) != len(texts):
|
||||
return shaped, frozenset(range(len(texts)))
|
||||
return shaped, frozenset(index for index, text in enumerate(texts) if text != shaped[index])
|
||||
|
||||
|
||||
def _shape_messages(messages: JsonValue, policy: PayloadPolicy) -> JsonValue:
|
||||
if not isinstance(messages, list) or not messages:
|
||||
return None
|
||||
windowed: Final = messages[-policy.max_messages :] if policy.max_messages is not None else messages
|
||||
texted: Final = (
|
||||
map_messages_text(windowed, lambda text: _shape_text(text, policy)) if policy.shapes_text else windowed
|
||||
)
|
||||
# send_images=False has to cover the inline image parts too, otherwise the
|
||||
# base64 payload still leaves LiteLLM inside structured_messages. The part is
|
||||
# kept so the guardrail still knows an image was there.
|
||||
if policy.send_images:
|
||||
return texted
|
||||
return map_messages_image_urls(texted, lambda _: IMAGE_OMITTED_PLACEHOLDER)
|
||||
|
||||
|
||||
def shape_payload(
|
||||
request: GenericGuardrailAPIRequest,
|
||||
policy: PayloadPolicy,
|
||||
) -> tuple[Mapping[str, JsonValue], PayloadLoss]:
|
||||
"""Return the JSON payload to POST plus the loss it represents.
|
||||
|
||||
Shaping runs on the dumped payload rather than on the model: pydantic
|
||||
validates ``structured_messages`` list content lazily into a
|
||||
``ValidatorIterator``, which is neither a list to walk nor safe to iterate
|
||||
twice, so the plain JSON form is the only reliable input here.
|
||||
"""
|
||||
exclude: Final = policy.exclude_fields if policy.send_images else policy.exclude_fields | _IMAGES_FIELD
|
||||
excluded: Final = set(exclude) or None # mutable-ok: model_dump(exclude=...) takes a set
|
||||
dumped: Final = request.model_dump(mode="json", exclude=excluded)
|
||||
|
||||
texts: Final = request.texts
|
||||
texts_sent: Final = bool(texts) and "texts" not in exclude
|
||||
shaped_texts, altered_indices = _shape_texts(texts or (), policy) if texts_sent else (None, frozenset[int]())
|
||||
shaped_messages: Final = _shape_messages(dumped.get("structured_messages"), policy)
|
||||
|
||||
payload: Final = { # mutable-ok: the POST body is a JSON object
|
||||
**dumped,
|
||||
**({"texts": shaped_texts} if texts_sent else {}), # mutable-ok: JSON object
|
||||
**({"structured_messages": shaped_messages} if shaped_messages is not None else {}), # mutable-ok: JSON
|
||||
}
|
||||
loss: Final = PayloadLoss(
|
||||
# An excluded text block was never seen at all, so no index may be rewritten.
|
||||
altered_text_indices=(altered_indices if texts_sent else frozenset(range(len(texts or ())))),
|
||||
images_omitted="images" in exclude,
|
||||
tools_omitted="tools" in exclude,
|
||||
)
|
||||
return payload, loss
|
||||
|
||||
|
||||
def merge_guardrailed_texts(
|
||||
*,
|
||||
original: list[str], # mutable-ok: the framework's write-back contract
|
||||
returned: list[str], # mutable-ok: the framework's write-back contract
|
||||
loss: PayloadLoss,
|
||||
guardrail_name: str | None,
|
||||
) -> list[str]: # mutable-ok: GenericGuardrailAPIInputs["texts"] is list[str]
|
||||
"""Accept guardrail-returned text only where the endpoint saw the original.
|
||||
|
||||
Indices the payload altered (stripped / truncated) or never carried keep the
|
||||
caller's text, so lossy shaping can never splice a mangled prompt back into
|
||||
the request. Length mismatch breaks index alignment, so with any loss the
|
||||
whole rewrite is refused.
|
||||
"""
|
||||
if not loss.altered_text_indices:
|
||||
return returned
|
||||
if len(returned) != len(original):
|
||||
verbose_proxy_logger.warning(
|
||||
"Generic Guardrail API (%s): ignoring returned texts. The payload was shaped "
|
||||
"(strip_patterns / max_text_chars / excluded texts) and the guardrail returned "
|
||||
"%d text(s) for %d sent, so they cannot be aligned.",
|
||||
guardrail_name,
|
||||
len(returned),
|
||||
len(original),
|
||||
)
|
||||
return original
|
||||
verbose_proxy_logger.warning(
|
||||
"Generic Guardrail API (%s): keeping original text for shaped index(es) %s; the guardrail "
|
||||
"never saw the full content at those positions.",
|
||||
guardrail_name,
|
||||
sorted(loss.altered_text_indices),
|
||||
)
|
||||
return [ # mutable-ok: write-back contract
|
||||
original[index] if index in loss.altered_text_indices else text for index, text in enumerate(returned)
|
||||
]
|
||||
|
|
@ -0,0 +1,64 @@
|
|||
"""How often this guardrail records a spend-log / OTEL entry.
|
||||
|
||||
``@log_guardrail_information`` appends a ``StandardLoggingGuardrailInformation``
|
||||
entry on every invocation, with no cap and no dedup, so a long-lived session
|
||||
accumulates one entry per guardrail call in the same request metadata. This
|
||||
module decides, per call, whether that entry should be recorded at all.
|
||||
|
||||
Suppression rides the decorator's existing contract: it clears
|
||||
``_guardrail_self_recorded`` before calling the guardrail and skips its own
|
||||
append if the flag is set afterwards. Setting the flag only on a normal return
|
||||
(never before a raise) keeps blocks and guardrail failures recorded under every
|
||||
scope, since the decorator's exception branch honors the same flag.
|
||||
"""
|
||||
|
||||
from typing import Final, Literal
|
||||
|
||||
from litellm.caching.in_memory_cache import InMemoryCache
|
||||
|
||||
GuardrailInformationScope = Literal["per_call", "per_session", "off"]
|
||||
|
||||
DEFAULT_GUARDRAIL_INFORMATION_SCOPE: Final[GuardrailInformationScope] = "per_call"
|
||||
|
||||
# Bounded so a long-running proxy cannot accumulate session keys forever.
|
||||
_SESSION_CACHE_MAX_ENTRIES: Final = 100_000
|
||||
_SESSION_CACHE_TTL_SECONDS: Final = 3600
|
||||
|
||||
|
||||
class RecordScope:
|
||||
"""Tracks which sessions have already recorded an entry for one guardrail."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
scope: GuardrailInformationScope,
|
||||
*,
|
||||
cache: InMemoryCache | None = None,
|
||||
) -> None:
|
||||
self._scope: Final = scope
|
||||
self._recorded_sessions: Final = cache or InMemoryCache(
|
||||
max_size_in_memory=_SESSION_CACHE_MAX_ENTRIES,
|
||||
default_ttl=_SESSION_CACHE_TTL_SECONDS,
|
||||
)
|
||||
|
||||
@property
|
||||
def scope(self) -> GuardrailInformationScope:
|
||||
return self._scope
|
||||
|
||||
def should_suppress(self, session_id: str | None) -> bool:
|
||||
"""Whether this successful call should skip its logging entry.
|
||||
|
||||
``per_session`` claims the session on its first call, so the calls that
|
||||
follow it suppress. A session id is required: without one there is
|
||||
nothing to dedup against, so the call records as it would under
|
||||
``per_call`` rather than silently dropping every entry.
|
||||
"""
|
||||
if self._scope == "per_call":
|
||||
return False
|
||||
if self._scope == "off":
|
||||
return True
|
||||
if session_id is None:
|
||||
return False
|
||||
if self._recorded_sessions.get_cache(session_id) is True:
|
||||
return True
|
||||
self._recorded_sessions.set_cache(session_id, True, ttl=_SESSION_CACHE_TTL_SECONDS)
|
||||
return False
|
||||
|
|
@ -0,0 +1,187 @@
|
|||
"""Applicability filters for the Generic Guardrail API.
|
||||
|
||||
Two independent ways to keep a request away from the guardrail endpoint
|
||||
entirely:
|
||||
|
||||
- by call type (``run_only_on_call_types`` / ``skip_call_types``): every hook
|
||||
resolves its own call type, so request and response skip symmetrically with no
|
||||
correlation needed.
|
||||
- by request content (``skip_if_system_prompt_matches`` /
|
||||
``skip_if_first_role_in``): only the request side can decide, so the decision
|
||||
is recorded and replayed on the paired response via ``SkipDecisionStore``.
|
||||
"""
|
||||
|
||||
import re
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Final, Optional
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.caching.in_memory_cache import InMemoryCache
|
||||
from litellm.proxy.guardrails._content_utils import iter_role_text
|
||||
from litellm.types.utils import CallTypes
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
|
||||
# Roles carrying the operator's instructions. Matching these instead of the whole
|
||||
# transcript avoids skipping a request because a user pasted a matching string.
|
||||
INSTRUCTION_ROLES: Final = frozenset({"system", "developer"})
|
||||
|
||||
# Fallback correlation cache, used only when the request and response hooks do
|
||||
# not share a logging object. Bounded and short-lived: a response that never
|
||||
# fires (stream abort, upstream error) must not leak an entry.
|
||||
_SKIP_CACHE_MAX_ENTRIES: Final = 1000
|
||||
_SKIP_CACHE_TTL_SECONDS: Final = 300
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class SkipPolicy:
|
||||
"""Resolved applicability configuration."""
|
||||
|
||||
system_prompt_patterns: tuple[re.Pattern[str], ...] = ()
|
||||
first_role_in: frozenset[str] = frozenset()
|
||||
key_aliases: frozenset[str] = frozenset()
|
||||
team_ids: frozenset[str] = frozenset()
|
||||
run_only_on_call_types: frozenset[str] | None = None
|
||||
skip_call_types: frozenset[str] = frozenset()
|
||||
|
||||
@property
|
||||
def filters_requests(self) -> bool:
|
||||
"""Whether any message-based filter is configured.
|
||||
|
||||
These read the request body, so only the request hook can decide and the
|
||||
decision has to be replayed on the response.
|
||||
"""
|
||||
return bool(self.system_prompt_patterns) or bool(self.first_role_in)
|
||||
|
||||
@property
|
||||
def filters_identities(self) -> bool:
|
||||
"""Whether any caller-identity filter is configured.
|
||||
|
||||
These read what authentication established, which both hooks see, so each
|
||||
side decides for itself and no correlation is needed.
|
||||
"""
|
||||
return bool(self.key_aliases) or bool(self.team_ids)
|
||||
|
||||
|
||||
def validate_call_types(
|
||||
raw: Sequence[str] | None,
|
||||
*,
|
||||
option_name: str,
|
||||
guardrail_name: str | None,
|
||||
) -> frozenset[str]:
|
||||
"""Validate configured call types against ``CallTypes``, warning on unknowns."""
|
||||
if not raw:
|
||||
return frozenset()
|
||||
known: Final = frozenset(call_type.value for call_type in CallTypes)
|
||||
unknown: Final = tuple(value for value in raw if value not in known)
|
||||
if unknown:
|
||||
verbose_proxy_logger.warning(
|
||||
"Generic Guardrail API (%s): %s contains unrecognized call type(s) %s. "
|
||||
"Values must be CallTypes names (e.g. acompletion, aembedding, anthropic_messages).",
|
||||
guardrail_name,
|
||||
option_name,
|
||||
unknown,
|
||||
)
|
||||
return frozenset(raw)
|
||||
|
||||
|
||||
def call_type_allowed(policy: SkipPolicy, call_type: str | None) -> bool:
|
||||
"""Whether the guardrail runs for this call type.
|
||||
|
||||
An unresolved call type runs, so a filter misconfiguration can never
|
||||
silently blind the guardrail. The allowlist takes precedence over the
|
||||
denylist.
|
||||
"""
|
||||
if call_type is None:
|
||||
return True
|
||||
if policy.run_only_on_call_types is not None:
|
||||
return call_type in policy.run_only_on_call_types
|
||||
return call_type not in policy.skip_call_types
|
||||
|
||||
|
||||
def identity_matches_skip(policy: SkipPolicy, metadata: Mapping[str, object]) -> bool:
|
||||
"""Whether the calling key or team is out of scope for the guardrail.
|
||||
|
||||
Matched against the metadata the auth layer produced, not the request body,
|
||||
so a caller cannot exempt itself by choosing what it sends.
|
||||
"""
|
||||
if not policy.filters_identities:
|
||||
return False
|
||||
key_alias: Final = metadata.get("user_api_key_alias")
|
||||
if policy.key_aliases and isinstance(key_alias, str) and key_alias in policy.key_aliases:
|
||||
return True
|
||||
team_id: Final = metadata.get("user_api_key_team_id")
|
||||
return bool(policy.team_ids and isinstance(team_id, str) and team_id in policy.team_ids)
|
||||
|
||||
|
||||
def request_matches_skip(
|
||||
policy: SkipPolicy,
|
||||
structured_messages: Sequence[Mapping[str, object]] | None,
|
||||
) -> bool:
|
||||
"""Whether this request is out of scope for the guardrail."""
|
||||
if not policy.filters_requests or not structured_messages:
|
||||
return False
|
||||
|
||||
first_role: Final = structured_messages[0].get("role")
|
||||
if policy.first_role_in and isinstance(first_role, str) and first_role in policy.first_role_in:
|
||||
return True
|
||||
|
||||
if not policy.system_prompt_patterns:
|
||||
return False
|
||||
|
||||
return any(
|
||||
pattern.search(text)
|
||||
for text in iter_role_text(structured_messages, INSTRUCTION_ROLES)
|
||||
for pattern in policy.system_prompt_patterns
|
||||
)
|
||||
|
||||
|
||||
class SkipDecisionStore:
|
||||
"""Carries a request-side skip decision to the paired response hook.
|
||||
|
||||
Primary channel is the shared logging object, which is exact and dies with
|
||||
the call. The bounded cache keyed by ``litellm_call_id`` covers the paths
|
||||
where no logging object reaches the hooks; entries are evicted on read.
|
||||
"""
|
||||
|
||||
def __init__(self, *, guardrail_name: str | None, cache: InMemoryCache | None = None) -> None:
|
||||
self._detail_key: Final = f"_generic_guardrail_skip::{guardrail_name}"
|
||||
self._cache: Final = cache or InMemoryCache(
|
||||
max_size_in_memory=_SKIP_CACHE_MAX_ENTRIES,
|
||||
default_ttl=_SKIP_CACHE_TTL_SECONDS,
|
||||
)
|
||||
|
||||
def _cache_key(self, call_id: str) -> str:
|
||||
return f"{self._detail_key}::{call_id}"
|
||||
|
||||
@staticmethod
|
||||
def _call_details(
|
||||
logging_obj: Optional["LiteLLMLoggingObj"],
|
||||
) -> dict[str, bool] | None: # mutable-ok: this is the logging object's live dict, written through on purpose
|
||||
"""The shared write-through channel, or None when this call has no logging object."""
|
||||
details: Final[object] = getattr(logging_obj, "model_call_details", None) if logging_obj else None
|
||||
# model_call_details is a free-form dict; only our own bool marker is read
|
||||
# or written here, so the narrower value type is accurate for this use.
|
||||
return details if isinstance(details, dict) else None
|
||||
|
||||
def record(self, *, logging_obj: Optional["LiteLLMLoggingObj"], call_id: str | None) -> None:
|
||||
details: Final = self._call_details(logging_obj)
|
||||
if details is not None:
|
||||
details[self._detail_key] = True
|
||||
return
|
||||
if call_id:
|
||||
self._cache.set_cache(self._cache_key(call_id), True, ttl=_SKIP_CACHE_TTL_SECONDS)
|
||||
|
||||
def consume(self, *, logging_obj: Optional["LiteLLMLoggingObj"], call_id: str | None) -> bool:
|
||||
details: Final = self._call_details(logging_obj)
|
||||
if details is not None and details.get(self._detail_key) is True:
|
||||
return True
|
||||
if not call_id:
|
||||
return False
|
||||
key: Final = self._cache_key(call_id)
|
||||
if self._cache.get_cache(key) is not True:
|
||||
return False
|
||||
self._cache.delete_cache(key)
|
||||
return True
|
||||
|
|
@ -102,6 +102,171 @@ class GenericGuardrailAPIOptionalParams(BaseModel):
|
|||
),
|
||||
)
|
||||
|
||||
fire_and_forget: bool | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"If True, the guardrail HTTP call is dispatched as a background task and the "
|
||||
"request proceeds immediately without awaiting the response. Applies to every "
|
||||
"mode (pre_call, post_call, during_call), adding ~0 latency. Because the "
|
||||
"response is never awaited, the guardrail is observe-only: action=BLOCKED and "
|
||||
"action=GUARDRAIL_INTERVENED are ignored. Also forces "
|
||||
"streaming_end_of_stream_only=True so a stream dispatches one background call "
|
||||
"instead of one per sampled chunk. Defaults to False in "
|
||||
"GenericGuardrailAPI.__init__ when None."
|
||||
),
|
||||
)
|
||||
|
||||
fire_and_forget_max_inflight: int | None = Field(
|
||||
default=None,
|
||||
ge=1,
|
||||
description=(
|
||||
"Maximum number of fire_and_forget calls in flight at once. Async dispatch "
|
||||
"decouples the request rate from the guardrail endpoint's throughput, so a slow "
|
||||
"endpoint would otherwise pile up tasks without bound. Calls beyond this limit "
|
||||
"are dropped and counted, with a rate-limited warning. Ignored unless "
|
||||
"fire_and_forget is True. Defaults to 100 in GenericGuardrailAPI.__init__ when None."
|
||||
),
|
||||
)
|
||||
|
||||
send_images: bool | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"If False, base64 image data is omitted from the guardrail request even when the "
|
||||
"LLM request contains images. Large payload saver for guardrails that only "
|
||||
"inspect text. Images the guardrail never received cannot be rewritten by its "
|
||||
"response. Defaults to True in GenericGuardrailAPI.__init__ when None."
|
||||
),
|
||||
)
|
||||
|
||||
exclude_payload_fields: tuple[str, ...] | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Top-level GenericGuardrailAPIRequest keys to omit from the payload (e.g. "
|
||||
"['images','texts','request_headers']), for providers that do not consume them. "
|
||||
"Unknown keys and the routing-critical keys input_type and litellm_call_id are "
|
||||
"ignored with a warning at init. A component that is not sent cannot be "
|
||||
"rewritten by the guardrail response."
|
||||
),
|
||||
)
|
||||
|
||||
max_messages: int | None = Field(
|
||||
default=None,
|
||||
ge=1,
|
||||
description=(
|
||||
"If set, only the last N entries of structured_messages and the last N text "
|
||||
"blocks are sent. Bounds payload size when the full conversation is re-sent "
|
||||
"every turn. Note the system prompt and early context fall out of the window "
|
||||
"once the session exceeds N. LOSSY: windowing shifts text positions, so the "
|
||||
"guardrail can no longer rewrite text on this call (action=BLOCKED still "
|
||||
"applies); leave it unset on a guardrail that masks or redacts."
|
||||
),
|
||||
)
|
||||
|
||||
max_text_chars: int | None = Field(
|
||||
default=None,
|
||||
ge=1,
|
||||
description=(
|
||||
"If set, each individual text block is truncated to this many characters before "
|
||||
"sending, for providers that only need a prefix. LOSSY: truncated text blocks "
|
||||
"keep their original content on write-back, so the guardrail cannot rewrite them."
|
||||
),
|
||||
)
|
||||
|
||||
strip_patterns: tuple[str, ...] | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Regex patterns applied to text content (texts[] and the text of each "
|
||||
"structured_messages entry) before the guardrail request is sent. Matches are "
|
||||
"removed. Intended to drop volatile boilerplate the provider does not need. "
|
||||
"Applied only to string text fields, never to JSON structure, tool schemas, ids "
|
||||
"or metadata. LOSSY: stripped content cannot be inspected by the guardrail, and "
|
||||
"stripped text blocks keep their original content on write-back. An invalid "
|
||||
"regex raises at init."
|
||||
),
|
||||
)
|
||||
|
||||
skip_if_system_prompt_matches: tuple[str, ...] | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Regex patterns matched against the request's system message (role=system or "
|
||||
"developer). On a match the guardrail is skipped for this call: no request is "
|
||||
"sent, and the paired response is skipped too, so both telemetry and enforcement "
|
||||
"are suppressed for matched requests. Matching the system message only (not "
|
||||
"arbitrary user text) avoids false positives from pasted content. Requires a "
|
||||
"request-side hook (pre_call or during_call) in the mode list. "
|
||||
"TRUST BOUNDARY: the system message comes from the request body, so any caller "
|
||||
"who knows the configured pattern can add a system message and exempt itself "
|
||||
"from this guardrail. Treat this as traffic scoping, not as enforcement, and "
|
||||
"prefer skip_if_key_alias_in / skip_if_team_id_in when the exemption has to "
|
||||
"hold against the caller."
|
||||
),
|
||||
)
|
||||
|
||||
skip_if_first_role_in: tuple[str, ...] | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"If the first message's role is in this list (e.g. ['developer']), skip the "
|
||||
"guardrail for the call, with the same request/response semantics as, and the "
|
||||
"same trust boundary as, skip_if_system_prompt_matches: the caller chooses the "
|
||||
"roles it sends."
|
||||
),
|
||||
)
|
||||
|
||||
skip_if_key_alias_in: tuple[str, ...] | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"If the calling virtual key's alias is in this list, skip the guardrail for the "
|
||||
"call. Unlike the message-based filters this reads what authentication "
|
||||
"established, so a caller cannot exempt itself by changing its request body. "
|
||||
"Same request/response semantics: neither side is sent."
|
||||
),
|
||||
)
|
||||
|
||||
skip_if_team_id_in: tuple[str, ...] | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"If the calling key's team id is in this list, skip the guardrail for the call. "
|
||||
"Admin-controlled like skip_if_key_alias_in, and matched on the same "
|
||||
"authenticated metadata."
|
||||
),
|
||||
)
|
||||
|
||||
run_only_on_call_types: tuple[str, ...] | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"If set, the guardrail runs ONLY for these call types (e.g. "
|
||||
"['completion','acompletion','anthropic_messages','responses','aresponses']). All "
|
||||
"other call types, including embeddings, image generation, audio, rerank and "
|
||||
"moderation, are skipped before any request is sent. Allowlist; takes precedence "
|
||||
"over skip_call_types. Values are CallTypes names. An unresolvable call type runs."
|
||||
),
|
||||
)
|
||||
|
||||
skip_call_types: tuple[str, ...] | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"If set (and run_only_on_call_types is not), the guardrail is skipped for these "
|
||||
"call types (e.g. ['embedding','aembedding','image_generation','transcription',"
|
||||
"'speech','rerank','moderation']). Denylist. Values are CallTypes names."
|
||||
),
|
||||
)
|
||||
|
||||
guardrail_information_scope: Literal["per_call", "per_session", "off"] | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"How often this guardrail records a StandardLoggingGuardrailInformation entry "
|
||||
"into request metadata (spend logs / OTEL). Every invocation records one today, "
|
||||
"with no cap and no dedup, so a long agent session accumulates one entry per "
|
||||
"guardrail call in the same row. 'per_call' (default) keeps that behavior. "
|
||||
"'per_session' records only the first call of a session, and needs a resolvable "
|
||||
"session id (litellm_session_id or metadata.session_id); without one it behaves "
|
||||
"as 'per_call' rather than dropping every entry. 'off' records nothing on "
|
||||
"success. Blocks and guardrail failures are recorded under every scope. Dedup "
|
||||
"is per proxy process, so a session spread across workers records once per "
|
||||
"worker. Defaults to 'per_call' in GenericGuardrailAPI.__init__ when None."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class GenericGuardrailAPIConfigModel(
|
||||
GuardrailConfigModel[GenericGuardrailAPIOptionalParams],
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -5,6 +5,10 @@ from litellm.proxy.guardrails._content_utils import (
|
|||
build_inspection_messages,
|
||||
has_non_string_content,
|
||||
iter_message_text,
|
||||
iter_role_text,
|
||||
map_content_text,
|
||||
map_messages_image_urls,
|
||||
map_messages_text,
|
||||
walk_user_text,
|
||||
)
|
||||
|
||||
|
|
@ -563,3 +567,83 @@ def test_build_inspection_messages_custom_tool_call_output():
|
|||
}
|
||||
msgs = build_inspection_messages(data)
|
||||
assert any("custom-tool-leak" in m["content"] for m in msgs)
|
||||
|
||||
|
||||
# ── map_content_text / map_messages_text / map_messages_image_urls ───────────────
|
||||
|
||||
|
||||
def test_map_content_text_transforms_string_and_text_parts_only():
|
||||
content = [
|
||||
{"type": "text", "text": "keep me"},
|
||||
{"type": "image_url", "image_url": {"url": "data:image/png;base64,AAA"}},
|
||||
"bare fragment",
|
||||
]
|
||||
assert map_content_text(content, lambda t: t.upper()) == [
|
||||
{"type": "text", "text": "KEEP ME"},
|
||||
{"type": "image_url", "image_url": {"url": "data:image/png;base64,AAA"}},
|
||||
"BARE FRAGMENT",
|
||||
]
|
||||
|
||||
|
||||
def test_map_content_text_does_not_mutate_the_input():
|
||||
content = [{"type": "text", "text": "original"}]
|
||||
map_content_text(content, lambda t: "changed")
|
||||
assert content == [{"type": "text", "text": "original"}]
|
||||
|
||||
|
||||
def test_map_messages_text_preserves_every_other_key():
|
||||
messages = [
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "answer",
|
||||
"tool_calls": [{"id": "call_1", "function": {"name": "run", "arguments": "{}"}}],
|
||||
}
|
||||
]
|
||||
mapped = map_messages_text(messages, lambda t: t.upper())
|
||||
assert mapped[0]["role"] == "assistant"
|
||||
assert mapped[0]["content"] == "ANSWER"
|
||||
assert mapped[0]["tool_calls"] == messages[0]["tool_calls"]
|
||||
|
||||
|
||||
def test_map_messages_image_urls_covers_both_part_shapes():
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "describe"},
|
||||
{"type": "image_url", "image_url": {"url": "data:image/png;base64,AAA", "detail": "low"}},
|
||||
{"type": "image_url", "image_url": "data:image/png;base64,BBB"},
|
||||
{"type": "input_image", "image_url": "data:image/png;base64,CCC"},
|
||||
],
|
||||
}
|
||||
]
|
||||
mapped = map_messages_image_urls(messages, lambda _: "[omitted]")
|
||||
content = mapped[0]["content"]
|
||||
assert content[0] == {"type": "text", "text": "describe"}
|
||||
assert content[1] == {"type": "image_url", "image_url": {"url": "[omitted]", "detail": "low"}}
|
||||
assert content[2] == {"type": "image_url", "image_url": "[omitted]"}
|
||||
assert content[3] == {"type": "input_image", "image_url": "[omitted]"}
|
||||
|
||||
|
||||
def test_map_messages_image_urls_leaves_string_content_alone():
|
||||
messages = [{"role": "user", "content": "just text"}]
|
||||
assert map_messages_image_urls(messages, lambda _: "[omitted]") == messages
|
||||
|
||||
|
||||
# ── iter_role_text ──────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_iter_role_text_selects_only_the_requested_roles():
|
||||
messages = [
|
||||
{"role": "system", "content": "system rules"},
|
||||
{"role": "developer", "content": [{"type": "text", "text": "developer rules"}]},
|
||||
{"role": "user", "content": "user text"},
|
||||
]
|
||||
assert list(iter_role_text(messages, {"system", "developer"})) == [
|
||||
"system rules",
|
||||
"developer rules",
|
||||
]
|
||||
|
||||
|
||||
def test_iter_role_text_ignores_non_mapping_entries():
|
||||
assert list(iter_role_text(["not a message", {"role": "system", "content": "ok"}], {"system"})) == ["ok"]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue