scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-15 04:19:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 121

feat(guardrails): add Airia guardrail integration (#40784)

Browse source 
* feat(guardrails): add Airia guardrail integration

Adds `airia` as a built-in guardrail provider. The hook sends the prompt under `pre_call` and the
model's response under `post_call` to the Airia AI Gateway, which answers allow, block, or redacted
content; the only configuration is `api_base`, `api_key`, and an optional `timeout`.

`during_call` is not offered: it runs concurrently with the model call, so a block could land after
the prompt has already reached the provider.

A BLOCKED verdict, or any action this version does not recognise, raises with
`blocked_content=True`; a transport error, a non-2xx, or a response body that is not a JSON object
raises with `blocked_content=False`, so callers can tell "could not evaluate" from "evaluated and
blocked" while both still fail closed. On GUARDRAIL_INTERVENED a copy of the inputs is returned with
every rewritten field substituted; an intervention carrying no applicable rewrite, a rewrite of the
wrong shape, or a rewrite whose item count does not match what was sent (which would misalign the
positional write-back downstream) blocks instead of letting the original through.

Streamed responses are moderated whole and then emitted redacted: the hook opts into the unified
hook's `incremental_diff` mode (the default `block_only` drops rewrites) with
`streaming_end_of_stream_only`, so a redaction can never span transform rounds and underflow. Covered
by a test that drives the real streaming iterator hook end to end and checks the assembled client
output, not just the guardrail's own config attributes.

The HTTP client is dependency-injected (an `async_handler` constructor parameter, defaulting to a
real client) rather than mutated post-construction, so tests inject a fake client instead of
monkeypatching an instance method.

Known limitation, not fixable from this guardrail alone: LiteLLM's unified guardrail framework
yields a tool-call chunk to the client immediately during `incremental_diff` streaming, inspecting
the assembled tool calls against the guardrail's verdict only once the stream ends
(`_run_incremental_transform_stream` in `unified_guardrail.py`). A blocked tool call can therefore
reach an executing client before the terminal error arrives. Switching to `block_only` with
`streaming_buffer_until_moderated` does not close this without a regression: that flag has no effect
under `incremental_diff` (the mode branches away before it is read), and `block_only` itself replays
the original, unredacted chunks on release, silently dropping any REDACT verdict on streamed text.
Relatedly, LiteLLM's OpenAI chat completions guardrail translation skips `post_call` guardrails
entirely for a response with no text content, so a tool-call-only response bypasses this guardrail
(and every other post-call guardrail) even outside streaming. Both are shared-framework behavior,
not specific to Airia; fixing either is out of scope for an isolated single-guardrail contribution.

* fix(guardrails): fail closed on an Airia rewrite for a field that was not sent

Bugbot finding on #40784. An INTERVENED verdict carrying `structured_messages: []` with
`texts` omitted passed `_rewritten` unchanged: an empty list has the same length as an
omitted field, so it counted as an applied rewrite while the `texts` that actually carried
the content went through untouched. post_call sends only `texts`, so this was reachable.

A rewrite for a field the request did not send is now a mismatch, same as a wrong length,
and the verdict blocks. Regression test covers both the empty and the non-empty case.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

* fix(guardrails): fail closed on streamed non-chat routes the Airia rewrite cannot reach

Veria finding on #40784. incremental_diff, the mode that carries a rewrite into the stream,
exists only for OpenAI chat completions; UnifiedLLMGuardrails sends every other streamed route
(/v1/responses, /v1/messages) to block_only, which streams the original text live and drops
rewrites. So with post_call on, a streamed Responses or Anthropic call could receive text that
Airia then blocked or redacted.

Two changes, both confined to this guardrail:

- streaming_buffer_until_moderated is set, so on the block_only fallback the framework holds
  every chunk until the end-of-stream verdict instead of releasing it live. The flag is never
  read on the incremental_diff path, so streamed chat completions are unchanged.
- A REDACT verdict on a streamed response whose route cannot carry a rewrite is enforced as a
  block, since block_only would otherwise replay the original. The route check mirrors
  _resolve_transform_call_type so the two cannot disagree; an unresolvable route fails closed.
  Request-side rewrites and non-streamed calls are untouched.

The end-to-end test drives the real UnifiedLLMGuardrails iterator over an Anthropic
/v1/messages SSE stream and asserts not one original chunk is released for either verdict.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

* test(guardrails): keep the Airia stream test's pytest.raises block to one statement

CI's test-tree ruff run flags PT012 on the non-chat streaming test: the `async for` that
collected released chunks sat inside `pytest.raises`, so a raise from the loop body would be
indistinguishable from the one the test is about. The loop now lives in a `_drain_into` helper
and the `with` block awaits it alone. Same assertion, same partial-leak capture.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

* refactor(guardrails): keep the Airia guardrail within the per-rule basedpyright ceilings

scripts/type_check_gate.py now gates every rule on its merge-base count (#40193), and the
guardrail added sixteen strict-rule errors it previously absorbed under the global budget.

- The parsed gateway response goes through `_as_object` once, so `action`, `blocked_reason`
  and the rewrite fields are read from `dict[str, object]` rather than `dict[Unknown, Unknown]`;
  `blocked_reason` is only used as the message when it is a non-empty string.
- `get_config_model` returns the concrete `type[AiriaGuardrailConfigModel]`, and the config model
  carries the same pyright ignore as its siblings for the unparameterised base.
- The remaining untyped upstream members (the decorator, the client factory, `super().__init__`'s
  kwargs, the handler's `post`) are ignored per line with the reason, as the Wingback guardrail
  does. The registry reads `guardrail["guardrail_name"]`, a required TypedDict key.
- Tests reach the injected mock and the handler timeout through typed helpers instead of the
  real handler's partially typed attributes.

No behaviour change: 46/46 tests, ruff default and strict, and both lint gates pass.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

* fix(guardrails): clear the Airia guardrail's type-discipline gate violations

Drop the mutable-ok comments that suppress nothing (LIT013), give the two
deferred-import noqa lines a reason (LIT003), and annotate the parsed JSON
object as Mapping instead of dict (LIT001).

---------

Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
Co-authored-by: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com>
This commit is contained in:
airia-open-source 2026-10-14 18:45:47 -04:00 • committed by GitHub
parent d03923e6a9
commit 220bbfc0ce
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 929 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

32
litellm/proxy/guardrails/guardrail_hooks/airia/__init__.py Normal file
View file

@ -0,0 +1,32 @@
from typing import TYPE_CHECKING, Final
from litellm.proxy.guardrails.guardrail_hooks.airia.airia import AiriaGuardrail
from litellm.types.guardrails import SupportedGuardrailIntegrations
if TYPE_CHECKING:
from litellm.types.guardrails import Guardrail, LitellmParams
def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail") -> AiriaGuardrail:
import litellm
_airia_callback: Final = AiriaGuardrail(
api_base=litellm_params.api_base,
api_key=litellm_params.api_key,
timeout=litellm_params.timeout,
guardrail_name=guardrail["guardrail_name"],
event_hook=litellm_params.mode,
default_on=litellm_params.default_on,
)
litellm.logging_callback_manager.add_litellm_callback(_airia_callback)
return _airia_callback
guardrail_initializer_registry: Final = {
SupportedGuardrailIntegrations.AIRIA.value: initialize_guardrail,
}
guardrail_class_registry: Final = {
SupportedGuardrailIntegrations.AIRIA.value: AiriaGuardrail,
}

246
litellm/proxy/guardrails/guardrail_hooks/airia/airia.py Normal file
View file

@ -0,0 +1,246 @@
import os
import uuid
from collections.abc import Mapping
from typing import (
Any, # noqa: TID251 # the only type CustomGuardrail.__init__ accepts for its open-ended kwargs
Final,
Literal,
)
import httpx
from litellm._logging import verbose_proxy_logger
from litellm.exceptions import GuardrailRaisedException
from litellm.integrations.custom_guardrail import (
CustomGuardrail,
log_guardrail_information, # pyright: ignore[reportUnknownVariableType] # the decorator is untyped upstream
)
from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.custom_httpx.http_handler import (
AsyncHTTPHandler,
get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # the factory's params are untyped upstream
httpxSpecialProvider,
)
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.proxy.guardrails.guardrail_hooks.airia import AiriaGuardrailConfigModel
from litellm.types.utils import GenericGuardrailAPIInputs
GUARDRAIL_PATH: Final = "/v1/guardrails/litellm"
ACTION_NONE: Final = "NONE"
ACTION_BLOCKED: Final = "BLOCKED"
ACTION_INTERVENED: Final = "GUARDRAIL_INTERVENED"
DEFAULT_BLOCKED_MESSAGE: Final = "Blocked by your organization's content policy."
SUPPORTED_EVENT_HOOKS: Final = (GuardrailEventHooks.pre_call, GuardrailEventHooks.post_call)
def _stream_cannot_carry_a_rewrite(request_data: Mapping[str, object] | None) -> bool:
"""Whether this is a streamed response on a route the framework can only block, not rewrite.
``incremental_diff``, the mode that writes a rewrite into the stream, exists only for OpenAI chat
completions: ``UnifiedLLMGuardrails._resolve_transform_call_type`` sends every other streamed route
to ``block_only``, which drops the rewrite and replays the buffered original. On those routes a
REDACT verdict can only be honoured as a block. The route check mirrors that resolver so the two
cannot disagree, and a route it cannot resolve fails closed the same way the resolver does.
"""
if request_data is None or not _is_streamed(request_data):
return False
route: Final = _request_route(request_data)
call_types: Final = get_call_types_for_route(route) if route else None
if not call_types:
return True
# Deferred like the resolver's own import: the handler package imports the proxy.
from litellm.llms import load_guardrail_translation_mappings # noqa: PLC0415 # circular import
from litellm.llms.openai.chat.guardrail_translation.handler import ( # noqa: PLC0415 # circular import
OpenAIChatCompletionsHandler,
)
handler: Final[object] = load_guardrail_translation_mappings().get(call_types[0])
return not (isinstance(handler, type) and issubclass(handler, OpenAIChatCompletionsHandler))
def _is_streamed(request_data: Mapping[str, object]) -> bool:
if request_data.get("stream") is True:
return True
proxy_request: Final = _as_object(request_data.get("proxy_server_request"))
body: Final = _as_object(proxy_request.get("body")) if proxy_request is not None else None
return body is not None and body.get("stream") is True
def _request_route(request_data: Mapping[str, object]) -> str | None:
"""The route the proxy stamps as ``user_api_key_request_route``, in whichever metadata bucket the
route uses: ``metadata`` for chat completions, ``litellm_metadata`` for Responses and Anthropic."""
for bucket_name in ("metadata", "litellm_metadata"):
bucket = _as_object(request_data.get(bucket_name))
route = bucket.get("user_api_key_request_route") if bucket is not None else None
if isinstance(route, str) and route:
return route
return None
def _as_object(value: object) -> Mapping[str, object] | None:
"""`value` as a JSON object, or None. States the key type once so the checker does not see
`dict[Unknown, Unknown]` at every read."""
return value if isinstance(value, dict) else None # pyright: ignore[reportUnknownVariableType] # isinstance narrows to dict[Unknown, Unknown]; JSON keys are str
class AiriaGuardrail(CustomGuardrail):
"""Evaluates prompts and responses against your Airia guardrail policy."""
def __init__(
self,
api_base: str | None = None,
api_key: str | None = None,
timeout: float | None = None,
supported_event_hooks: list[GuardrailEventHooks] | None = None, # mutable-ok: matches CustomGuardrail.__init__
async_handler: AsyncHTTPHandler | None = None, # dependency-injected in tests; defaults to a real client
**kwargs: Any, # kwargs-ok: passed straight through to CustomGuardrail.__init__
) -> None:
resolved_timeout: Final = timeout or float(os.getenv("AIRIA_TIMEOUT", "10"))
self.async_handler = async_handler or get_async_httpx_client(
llm_provider=httpxSpecialProvider.GuardrailCallback,
params={"timeout": httpx.Timeout(timeout=resolved_timeout, connect=5.0)},
)
self.api_base = (api_base or os.getenv("AIRIA_GATEWAY_URL", "")).rstrip("/")
self.api_key = api_key or os.getenv("AIRIA_API_KEY")
self.streaming_transform_mode: Final[Literal["block_only", "incremental_diff"]] = "incremental_diff"
self.streaming_end_of_stream_only: Final = True
# incremental_diff exists only for streamed chat completions. Every other streamed route falls
# back to block_only, where the framework releases the original text live unless told to hold
# it; holding it keeps the text behind the end-of-stream verdict, and apply_guardrail turns a
# REDACT that block_only cannot deliver into a block. Never read on the incremental_diff path.
self.streaming_buffer_until_moderated: Final = True
if not self.api_base:
raise ValueError("AiriaGuardrail requires api_base, or the AIRIA_GATEWAY_URL environment variable.")
if not self.api_key:
raise ValueError("AiriaGuardrail requires api_key, or the AIRIA_API_KEY environment variable.")
self.optional_params = kwargs
super().__init__( # pyright: ignore[reportUnknownMemberType] # parent takes untyped **kwargs
supported_event_hooks=supported_event_hooks or [*SUPPORTED_EVENT_HOOKS],
**kwargs,
)
verbose_proxy_logger.info("AiriaGuardrail initialized with gateway: %s", self.api_base)
@log_guardrail_information
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict[str, object], # mutable-ok: the base class declares this parameter as a dict
input_type: Literal["request", "response"],
logging_obj: "LiteLLMLoggingObj | None" = None,
) -> GenericGuardrailAPIInputs:
call_id: Final = (
logging_obj.litellm_call_id
if logging_obj
else (request_data.get("litellm_call_id") if request_data else None)
) or str(uuid.uuid4())
try:
response: Final = await self.async_handler.post( # pyright: ignore[reportUnknownMemberType] # the handler's `data` param is an untyped dict upstream
f"{self.api_base}{GUARDRAIL_PATH}",
json={
"input_type": input_type,
"texts": inputs.get("texts") or (),
"images": inputs.get("images") or (),
"structured_messages": inputs.get("structured_messages") or (),
"tools": inputs.get("tools") or (),
"tool_calls": inputs.get("tool_calls") or (),
"model": inputs.get("model"),
"litellm_call_id": call_id,
},
headers={"Authorization": f"Bearer {self.api_key}"},
)
response.raise_for_status()
raw_body: Final[object] = response.json()
body: Final = _as_object(raw_body)
if body is None:
raise TypeError(f"expected a JSON object in the response body, got {type(raw_body).__name__}")
except Exception as error:
verbose_proxy_logger.error(
"Airia guardrail could not evaluate the request (litellm_call_id=%s, input_type=%s): %s",
call_id,
input_type,
error,
)
raise GuardrailRaisedException(
guardrail_name=self.guardrail_name,
message=f"Airia guardrail could not evaluate the request: {error}",
blocked_content=False,
) from error
action: Final = body.get("action")
if action == ACTION_BLOCKED:
reason: Final = body.get("blocked_reason")
raise GuardrailRaisedException(
guardrail_name=self.guardrail_name,
message=reason if isinstance(reason, str) and reason else DEFAULT_BLOCKED_MESSAGE,
blocked_content=True,
)
if action == ACTION_INTERVENED:
if input_type == "response" and _stream_cannot_carry_a_rewrite(request_data):
raise self._blocked()
return self._rewritten(body, inputs)
if action != ACTION_NONE:
raise self._blocked()
return inputs
def _rewritten(
self,
body: Mapping[str, object],
inputs: GenericGuardrailAPIInputs,
) -> GenericGuardrailAPIInputs:
texts: Final = body.get("texts")
structured_messages: Final = body.get("structured_messages")
if not self._is_applicable_rewrite(texts, inputs.get("texts")):
raise self._blocked()
if not self._is_applicable_rewrite(structured_messages, inputs.get("structured_messages")):
raise self._blocked()
# An intervention that rewrites nothing would hand the original content through under a
# verdict that said it must not pass.
if texts is None and structured_messages is None:
raise self._blocked()
rewritten: Final[GenericGuardrailAPIInputs] = {**inputs}
if isinstance(texts, list):
rewritten["texts"] = texts
if isinstance(structured_messages, list):
rewritten["structured_messages"] = structured_messages
return rewritten
@staticmethod
def _is_applicable_rewrite(rewrite: object, original: object) -> bool:
"""A rewrite is safe to apply positionally only when it is absent (this field was not
touched) or a list the same length as what was sent. A different count would misalign
the positional write-back downstream, dropping content or leaving other fields unchanged.
A rewrite for a field that was never sent is a mismatch too, not a no-op: post_call sends
only `texts`, so `structured_messages: []` would otherwise count as an applied rewrite
while the `texts` that carried the content passed through untouched.
"""
if rewrite is None:
return True
if not isinstance(original, list) or not isinstance(rewrite, list):
return False
return len(rewrite) == len(original) # pyright: ignore[reportUnknownArgumentType] # only the counts matter; the items are the gateway's
def _blocked(self) -> GuardrailRaisedException:
return GuardrailRaisedException(
guardrail_name=self.guardrail_name,
message=DEFAULT_BLOCKED_MESSAGE,
blocked_content=True,
)
@staticmethod
def get_config_model() -> type[AiriaGuardrailConfigModel] | None:
return AiriaGuardrailConfigModel

1
litellm/types/guardrails.py
View file

@ -156,6 +156,7 @@ class SupportedGuardrailIntegrations(Enum):
AGENT_365 = "agent_365"
LLM_SHIELD_PROXY = "llm_shield_proxy"
CONDUCT = "conduct"
AIRIA = "airia"
WINGBACK = "wingback"
DECISION_MODEL = "decision_model"

24
litellm/types/proxy/guardrails/guardrail_hooks/airia.py Normal file
View file

@ -0,0 +1,24 @@
from pydantic import Field
from .base import GuardrailConfigModel
class AiriaGuardrailConfigModel(GuardrailConfigModel): # pyright: ignore[reportMissingTypeArgument] # parameterizing adds an optional_params field to the UI form
api_base: str | None = Field(
default=None,
description="Base URL of the Airia AI Gateway, e.g. https://gateway.airia.ai. If not "
"provided, the `AIRIA_GATEWAY_URL` environment variable is checked.",
)
api_key: str | None = Field(
default=None,
description="An Airia API key. If not provided, the `AIRIA_API_KEY` environment variable is checked.",
)
timeout: float | None = Field(
default=None,
description="Request timeout in seconds. If not provided, the `AIRIA_TIMEOUT` environment "
"variable is checked, defaulting to 10.",
)
@staticmethod
def ui_friendly_name() -> str:
return "Airia Guardrail"

1
ruff-strict.toml
View file

@ -30,6 +30,7 @@ external = [
# grows over time; typing it concretely (`object`) broke that forwarding call outright —
# basedpyright turned every named param into a reportArgumentType error. Any is correct here.
"litellm/proxy/guardrails/guardrail_hooks/alice/alice.py" = ["ANN401"]
"litellm/proxy/guardrails/guardrail_hooks/airia/airia.py" = ["ANN401"]
# Same reason: `**kwargs` forwards verbatim to CustomGuardrail.__init__, and the lifecycle
# hook signatures inherit `Any` for `response` from CustomLogger, so narrowing them here
# would break the override rather than describe it.

625
tests/unit/proxy/guardrails/guardrail_hooks/test_airia.py Normal file
View file

@ -0,0 +1,625 @@
import json
import uuid
from collections.abc import AsyncGenerator
from typing import Final, Literal
from unittest.mock import AsyncMock, MagicMock
import httpx
import pytest
import litellm
from litellm.exceptions import GuardrailRaisedException
from litellm.proxy.guardrails.guardrail_hooks.airia.airia import AiriaGuardrail
from litellm.proxy.guardrails.init_guardrails import (
init_guardrails_v2, # pyright: ignore[reportUnknownVariableType] # untyped upstream
)
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.proxy.guardrails.guardrail_hooks.airia import AiriaGuardrailConfigModel
from litellm.types.utils import GenericGuardrailAPIInputs
API_BASE = "https://gateway.airia.ai"
API_KEY = "ak-test-key"
def _make_guardrail(**kwargs: object) -> AiriaGuardrail:
"""Build a guardrail; construction defaults still exercise the real HTTP client."""
kwargs.setdefault("api_base", API_BASE)
kwargs.setdefault("api_key", API_KEY)
kwargs.setdefault("guardrail_name", "airia-guard")
return AiriaGuardrail(**kwargs)
def _make_guardrail_with_response(payload: object, status_code: int = 200, **kwargs: object) -> AiriaGuardrail:
"""Build a guardrail with a fake HTTP client injected at construction, per CLAUDE.md's
dependency-injection rule: no test mutates `async_handler` after the fact."""
handler = MagicMock()
handler.post = AsyncMock(return_value=_mock_response(payload, status_code))
return _make_guardrail(async_handler=handler, **kwargs)
def _mock_response(payload: object, status_code: int = 200) -> MagicMock:
response = MagicMock()
response.status_code = status_code
response.json.return_value = payload
response.raise_for_status = MagicMock()
return response
def _as_mock(value: object) -> MagicMock:
"""Forget the declared type, which is the real handler's, and recover the mock that was injected."""
assert isinstance(value, MagicMock)
return value
def _post_mock(guardrail: AiriaGuardrail) -> AsyncMock:
"""The `post` the test injected, typed as the mock it is so call inspection stays checked."""
post = _as_mock(guardrail.async_handler).post
assert isinstance(post, AsyncMock)
return post
def _timeout(guardrail: AiriaGuardrail) -> httpx.Timeout:
timeout = guardrail.async_handler.timeout # pyright: ignore[reportUnknownMemberType] # the handler's attribute is partially untyped
assert isinstance(timeout, httpx.Timeout)
return timeout
def _inputs(**overrides: object) -> GenericGuardrailAPIInputs:
return GenericGuardrailAPIInputs(**{"texts": ["hello"], "model": "gpt-4o", **overrides})
def test_airia_guardrail_config(monkeypatch: pytest.MonkeyPatch) -> None:
"""The guardrail registers through init_guardrails_v2 under the `airia` key."""
monkeypatch.setattr(litellm, "guardrail_name_config_map", {})
monkeypatch.setattr(litellm, "callbacks", [])
monkeypatch.setenv("AIRIA_GATEWAY_URL", API_BASE)
monkeypatch.setenv("AIRIA_API_KEY", API_KEY)
init_guardrails_v2(
all_guardrails=[
{
"guardrail_name": "airia-guard",
"litellm_params": {
"guardrail": "airia",
"mode": "pre_call",
"default_on": True,
"timeout": 45.0,
},
}
],
config_file_path="",
)
registered = [c for c in litellm.callbacks if isinstance(c, AiriaGuardrail)]
assert len(registered) == 1
assert registered[0].guardrail_name == "airia-guard"
assert registered[0].default_on is True
assert registered[0].event_hook == "pre_call"
assert _timeout(registered[0]).read == 45.0
def test_during_call_is_not_a_supported_event_hook() -> None:
"""during_call runs concurrently with the upstream call, so a block can land too late."""
hooks = _make_guardrail().supported_event_hooks
assert hooks is not None
assert GuardrailEventHooks.during_call not in hooks
assert list(hooks) == [GuardrailEventHooks.pre_call, GuardrailEventHooks.post_call]
@pytest.mark.parametrize("missing", ["api_base", "api_key"])
def test_missing_credentials_raise_at_construction(missing: str, monkeypatch: pytest.MonkeyPatch) -> None:
"""Fail at startup rather than on the first request."""
monkeypatch.delenv("AIRIA_GATEWAY_URL", raising=False)
monkeypatch.delenv("AIRIA_API_KEY", raising=False)
kwargs = {"api_base": API_BASE, "api_key": API_KEY}
kwargs[missing] = None
with pytest.raises(ValueError, match="AiriaGuardrail requires"):
_make_guardrail(**kwargs)
def test_credentials_fall_back_to_environment(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("AIRIA_GATEWAY_URL", API_BASE)
monkeypatch.setenv("AIRIA_API_KEY", API_KEY)
guardrail = _make_guardrail(api_base=None, api_key=None)
assert guardrail.api_base == API_BASE
assert guardrail.api_key == API_KEY
def test_trailing_slash_is_stripped_from_api_base() -> None:
guardrail = _make_guardrail(api_base=f"{API_BASE}/")
assert guardrail.api_base == API_BASE
def test_custom_timeout_from_kwargs() -> None:
guardrail = _make_guardrail(timeout=45.0)
assert _timeout(guardrail).read == 45.0
assert _timeout(guardrail).connect == 5.0
def test_timeout_falls_back_to_environment(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("AIRIA_TIMEOUT", "30")
guardrail = _make_guardrail()
assert _timeout(guardrail).read == 30.0
@pytest.mark.asyncio
async def test_request_payload_and_auth_header() -> None:
guardrail = _make_guardrail_with_response({"action": "NONE"})
inputs = _inputs(
images=["data:image/png;base64,AAAA"],
structured_messages=[{"role": "user", "content": "hello"}],
tools=[{"type": "function", "function": {"name": "f"}}],
tool_calls=[{"id": "c1", "function": {"name": "f", "arguments": "{}"}}],
)
await guardrail.apply_guardrail(inputs=inputs, request_data={"litellm_call_id": "call-123"}, input_type="request")
call = _post_mock(guardrail).call_args
assert call.args[0] == f"{API_BASE}/v1/guardrails/litellm"
assert call.kwargs["headers"] == {"Authorization": f"Bearer {API_KEY}"}
payload = call.kwargs["json"]
assert payload["input_type"] == "request"
assert payload["texts"] == ["hello"]
assert payload["structured_messages"] == [{"role": "user", "content": "hello"}]
assert payload["images"] == ["data:image/png;base64,AAAA"]
assert payload["tools"] == [{"type": "function", "function": {"name": "f"}}]
assert payload["tool_calls"] == [{"id": "c1", "function": {"name": "f", "arguments": "{}"}}]
assert payload["model"] == "gpt-4o"
assert payload["litellm_call_id"] == "call-123"
@pytest.mark.asyncio
@pytest.mark.parametrize("input_type", ["request", "response"])
async def test_input_type_is_forwarded_verbatim(input_type: Literal["request", "response"]) -> None:
guardrail = _make_guardrail_with_response({"action": "NONE"})
await guardrail.apply_guardrail(inputs=_inputs(), request_data={}, input_type=input_type)
assert _post_mock(guardrail).call_args.kwargs["json"]["input_type"] == input_type
@pytest.mark.asyncio
async def test_call_id_falls_back_to_a_generated_uuid() -> None:
guardrail = _make_guardrail_with_response({"action": "NONE"})
await guardrail.apply_guardrail(inputs=_inputs(), request_data={}, input_type="request")
call_id = _post_mock(guardrail).call_args.kwargs["json"]["litellm_call_id"]
assert uuid.UUID(call_id).version == 4
@pytest.mark.asyncio
async def test_action_none_returns_inputs_unchanged() -> None:
guardrail = _make_guardrail_with_response({"action": "NONE"})
inputs = _inputs(structured_messages=[{"role": "user", "content": "hello"}])
result = await guardrail.apply_guardrail(inputs=inputs, request_data={}, input_type="request")
assert result["texts"] == ["hello"]
assert result["structured_messages"] == [{"role": "user", "content": "hello"}]
@pytest.mark.asyncio
async def test_action_blocked_raises_with_blocked_content_set() -> None:
guardrail = _make_guardrail_with_response({"action": "BLOCKED", "blocked_reason": "Contains a secret"})
with pytest.raises(GuardrailRaisedException) as excinfo:
await guardrail.apply_guardrail(inputs=_inputs(), request_data={}, input_type="request")
assert excinfo.value.blocked_content is True
assert "Contains a secret" in str(excinfo.value)
@pytest.mark.asyncio
async def test_action_blocked_without_a_reason_uses_the_default_message() -> None:
guardrail = _make_guardrail_with_response({"action": "BLOCKED"})
with pytest.raises(GuardrailRaisedException) as excinfo:
await guardrail.apply_guardrail(inputs=_inputs(), request_data={}, input_type="request")
assert excinfo.value.blocked_content is True
assert "Blocked by your organization's content policy." in str(excinfo.value)
@pytest.mark.asyncio
async def test_intervened_replaces_both_text_bearing_fields() -> None:
"""Substituting only one field would leave the other still carrying the unredacted text."""
guardrail = _make_guardrail_with_response(
{
"action": "GUARDRAIL_INTERVENED",
"texts": ["my email is [EmailAddress1]"],
"structured_messages": [{"role": "user", "content": "my email is [EmailAddress1]"}],
}
)
inputs = _inputs(
texts=["my email is a@b.com"],
structured_messages=[{"role": "user", "content": "my email is a@b.com"}],
)
result = await guardrail.apply_guardrail(inputs=inputs, request_data={}, input_type="request")
assert result["texts"] == ["my email is [EmailAddress1]"]
assert result["structured_messages"] == [{"role": "user", "content": "my email is [EmailAddress1]"}]
assert "a@b.com" not in str(result)
@pytest.mark.asyncio
async def test_intervened_leaves_a_field_alone_when_airia_omits_it() -> None:
"""An omitted field means "not rewritten", not "rewritten to empty"."""
guardrail = _make_guardrail_with_response({"action": "GUARDRAIL_INTERVENED", "texts": ["redacted"]})
inputs = _inputs(texts=["original"], structured_messages=[{"role": "user", "content": "keep me"}])
result = await guardrail.apply_guardrail(inputs=inputs, request_data={}, input_type="request")
assert result["texts"] == ["redacted"]
assert result["structured_messages"] == [{"role": "user", "content": "keep me"}]
@pytest.mark.asyncio
async def test_intervened_with_a_non_list_field_is_treated_as_a_block() -> None:
"""A rewrite this version cannot apply must not let the original text through."""
guardrail = _make_guardrail_with_response({"action": "GUARDRAIL_INTERVENED", "texts": "not a list"})
with pytest.raises(GuardrailRaisedException) as excinfo:
await guardrail.apply_guardrail(inputs=_inputs(texts=["secret"]), request_data={}, input_type="request")
assert excinfo.value.blocked_content is True
@pytest.mark.asyncio
async def test_intervened_with_a_non_list_structured_messages_is_treated_as_a_block() -> None:
guardrail = _make_guardrail_with_response(
{"action": "GUARDRAIL_INTERVENED", "structured_messages": {"role": "user"}}
)
with pytest.raises(GuardrailRaisedException) as excinfo:
await guardrail.apply_guardrail(inputs=_inputs(), request_data={}, input_type="request")
assert excinfo.value.blocked_content is True
@pytest.mark.asyncio
async def test_intervened_without_any_rewritten_field_is_treated_as_a_block() -> None:
"""An intervention the proxy cannot apply must not let the original content through."""
guardrail = _make_guardrail_with_response({"action": "GUARDRAIL_INTERVENED", "tool_calls": []})
with pytest.raises(GuardrailRaisedException) as excinfo:
await guardrail.apply_guardrail(inputs=_inputs(texts=["secret"]), request_data={}, input_type="request")
assert excinfo.value.blocked_content is True
@pytest.mark.asyncio
@pytest.mark.parametrize("body", [["not", "an", "object"], "just a string", 42, None])
async def test_a_non_object_response_body_fails_closed_but_is_not_reported_as_a_verdict(body: object) -> None:
"""response.json() is not guaranteed to return a dict; body.get(...) on a list or scalar
would otherwise raise AttributeError and surface as an internal error, not a clean refusal."""
guardrail = _make_guardrail_with_response(body)
with pytest.raises(GuardrailRaisedException) as excinfo:
await guardrail.apply_guardrail(inputs=_inputs(), request_data={}, input_type="request")
assert excinfo.value.blocked_content is False
assert "could not evaluate" in str(excinfo.value)
@pytest.mark.asyncio
async def test_intervened_with_a_texts_rewrite_of_the_wrong_length_is_treated_as_a_block() -> None:
"""A rewrite with a different item count than what was sent would misalign the positional
write-back downstream; accepting it could drop content or leave other items unredacted."""
guardrail = _make_guardrail_with_response({"action": "GUARDRAIL_INTERVENED", "texts": ["only one redacted item"]})
with pytest.raises(GuardrailRaisedException) as excinfo:
await guardrail.apply_guardrail(
inputs=_inputs(texts=["first secret", "second secret"]), request_data={}, input_type="request"
)
assert excinfo.value.blocked_content is True
@pytest.mark.asyncio
@pytest.mark.parametrize("structured_messages", [[], [{"role": "user", "content": "[Redacted]"}]])
async def test_intervened_rewrite_for_a_field_that_was_not_sent_is_treated_as_a_block(
structured_messages: list[object],
) -> None:
"""post_call sends only `texts`. A `structured_messages` rewrite then has nothing to replace, and an
empty one must not count as "applied" while the `texts` that carried the content pass through."""
guardrail = _make_guardrail_with_response(
{"action": "GUARDRAIL_INTERVENED", "structured_messages": structured_messages}
)
with pytest.raises(GuardrailRaisedException) as excinfo:
await guardrail.apply_guardrail(inputs=_inputs(texts=["secret"]), request_data={}, input_type="response")
assert excinfo.value.blocked_content is True
@pytest.mark.asyncio
async def test_intervened_returns_a_copy_and_leaves_the_caller_object_untouched() -> None:
guardrail = _make_guardrail_with_response({"action": "GUARDRAIL_INTERVENED", "texts": ["redacted"]})
inputs = _inputs(texts=["original"])
result = await guardrail.apply_guardrail(inputs=inputs, request_data={}, input_type="request")
assert result["texts"] == ["redacted"]
assert inputs["texts"] == ["original"]
@pytest.mark.asyncio
async def test_streaming_redaction_actually_reaches_the_client() -> None:
"""block_only (the framework default) drops guardrail rewrites on the streaming path, so a
REDACT verdict would silently leak the original text. Runs the real UnifiedLLMGuardrails
streaming iterator end to end (not just this guardrail's config attributes) and checks what
the client would actually receive.
"""
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
UnifiedLLMGuardrails,
)
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
guardrail: Final = _make_guardrail_with_response(
{"action": "GUARDRAIL_INTERVENED", "texts": ["Contact me at [EmailAddress1]"]},
event_hook="post_call",
default_on=True,
)
async def source_stream() -> AsyncGenerator[ModelResponseStream, None]:
for content in ("Contact ", "me at ", "ada@example.com"):
yield ModelResponseStream(
model="gpt-4o-mini", choices=[StreamingChoices(index=0, delta=Delta(content=content))]
)
yield ModelResponseStream(
model="gpt-4o-mini",
choices=[StreamingChoices(index=0, delta=Delta(content=None), finish_reason="stop")],
)
assembled = ""
async for chunk in UnifiedLLMGuardrails().async_post_call_streaming_iterator_hook( # pyright: ignore[reportUnknownMemberType] # annotates request_data as a bare dict
user_api_key_dict=UserAPIKeyAuth(api_key="test", request_route="/chat/completions"),
response=source_stream(),
request_data={ # mutable-ok: one-shot request_data the framework requires as a plain dict
"messages": [{"role": "user", "content": "share your contact"}],
"stream": True,
},
guardrail_to_apply=guardrail,
):
for choice in chunk.choices:
if choice.delta.content:
assembled += choice.delta.content
assert "[EmailAddress1]" in assembled
assert "ada@example.com" not in assembled
_post_mock(guardrail).assert_called_once()
def _streamed_request(route: str | None) -> dict[str, object]:
"""A proxy request_data as the streaming hook sees it: `stream` set, and the route the proxy
stamped into the metadata bucket (`litellm_metadata` for the non-chat routes under test)."""
request_data: dict[str, object] = {"stream": True} # mutable-ok: built once per test
if route is not None:
request_data["litellm_metadata"] = {"user_api_key_request_route": route}
return request_data
@pytest.mark.asyncio
@pytest.mark.parametrize("route", ["/v1/messages", "/v1/responses", "/responses", "/unknown/route", None])
async def test_streamed_redaction_on_a_route_without_stream_rewrites_is_enforced_as_a_block(
route: str | None,
) -> None:
"""incremental_diff exists only for chat completions; every other streamed route runs block_only,
which drops rewrites and replays the original text. There a REDACT verdict must block instead."""
guardrail = _make_guardrail_with_response({"action": "GUARDRAIL_INTERVENED", "texts": ["[Redacted]"]})
with pytest.raises(GuardrailRaisedException) as excinfo:
await guardrail.apply_guardrail(
inputs=_inputs(texts=["secret"]), request_data=_streamed_request(route), input_type="response"
)
assert excinfo.value.blocked_content is True
@pytest.mark.asyncio
@pytest.mark.parametrize("route", ["/chat/completions", "/v1/chat/completions"])
async def test_streamed_redaction_on_chat_completions_is_still_delivered(route: str) -> None:
guardrail = _make_guardrail_with_response({"action": "GUARDRAIL_INTERVENED", "texts": ["[Redacted]"]})
request_data: dict[str, object] = {"stream": True, "metadata": {"user_api_key_request_route": route}}
result = await guardrail.apply_guardrail(
inputs=_inputs(texts=["secret"]), request_data=request_data, input_type="response"
)
assert result["texts"] == ["[Redacted]"]
@pytest.mark.asyncio
async def test_unstreamed_redaction_on_a_non_chat_route_is_delivered() -> None:
"""Only the stream cannot carry a rewrite; a non-streamed Responses call is rewritten as usual."""
guardrail = _make_guardrail_with_response({"action": "GUARDRAIL_INTERVENED", "texts": ["[Redacted]"]})
request_data: dict[str, object] = {"litellm_metadata": {"user_api_key_request_route": "/v1/responses"}}
result = await guardrail.apply_guardrail(
inputs=_inputs(texts=["secret"]), request_data=request_data, input_type="response"
)
assert result["texts"] == ["[Redacted]"]
@pytest.mark.asyncio
async def test_streamed_request_side_redaction_is_unaffected() -> None:
"""A pre_call rewrite lands in the prompt before the model is called; streaming is irrelevant there."""
guardrail = _make_guardrail_with_response({"action": "GUARDRAIL_INTERVENED", "texts": ["[Redacted]"]})
result = await guardrail.apply_guardrail(
inputs=_inputs(texts=["secret"]), request_data=_streamed_request("/v1/messages"), input_type="request"
)
assert result["texts"] == ["[Redacted]"]
@pytest.mark.asyncio
async def test_stream_flag_is_also_read_from_the_proxied_request_body() -> None:
guardrail = _make_guardrail_with_response({"action": "GUARDRAIL_INTERVENED", "texts": ["[Redacted]"]})
request_data: dict[str, object] = {
"proxy_server_request": {"body": {"stream": True}},
"litellm_metadata": {"user_api_key_request_route": "/v1/messages"},
}
with pytest.raises(GuardrailRaisedException) as excinfo:
await guardrail.apply_guardrail(
inputs=_inputs(texts=["secret"]), request_data=request_data, input_type="response"
)
assert excinfo.value.blocked_content is True
def _anthropic_sse(event_type: str, data: dict[str, object]) -> bytes:
return f"event: {event_type}\ndata: {json.dumps(data)}\n\n".encode()
async def _anthropic_messages_stream() -> AsyncGenerator[bytes, None]:
"""A complete Anthropic /v1/messages SSE stream, the shape the proxy relays on that route."""
yield _anthropic_sse(
"message_start",
{
"type": "message_start",
"message": {
"id": "msg_orig",
"type": "message",
"role": "assistant",
"model": "claude-3-5-sonnet",
"content": [],
"stop_reason": None,
"usage": {"input_tokens": 1, "output_tokens": 0},
},
},
)
yield _anthropic_sse(
"content_block_start",
{"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},
)
for text in ("Contact ", "me at ", "ada@example.com"):
yield _anthropic_sse(
"content_block_delta",
{"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": text}},
)
yield _anthropic_sse("content_block_stop", {"type": "content_block_stop", "index": 0})
yield _anthropic_sse(
"message_delta",
{
"type": "message_delta",
"delta": {"stop_reason": "end_turn", "stop_sequence": None},
"usage": {"output_tokens": 5},
},
)
yield _anthropic_sse("message_stop", {"type": "message_stop"})
async def _drain_into(stream: AsyncGenerator[object, None], sink: list[object]) -> None:
"""Forward every chunk the client would have received into `sink`, keeping what leaked before a
raise (a comprehension would drop it). Lives outside the test so `pytest.raises` wraps one await."""
async for chunk in stream:
sink.append(chunk) # noqa: PERF401 # the partial list is the point when the stream is refused
@pytest.mark.asyncio
@pytest.mark.parametrize(
"verdict",
[
{"action": "BLOCKED", "blocked_reason": "policy"},
{"action": "GUARDRAIL_INTERVENED", "texts": ["Contact me at [EmailAddress1]"]},
],
)
async def test_non_chat_streamed_route_releases_nothing_before_the_verdict(verdict: dict[str, object]) -> None:
"""Drives the real UnifiedLLMGuardrails iterator on /v1/messages, which incremental_diff does not
support, so the framework takes its block_only fallback. streaming_buffer_until_moderated holds
the text until the end-of-stream verdict; a block, and a redaction block_only could not deliver,
both refuse the response without one original chunk reaching the client."""
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
UnifiedLLMGuardrails,
)
guardrail: Final = _make_guardrail_with_response(verdict, event_hook="post_call", default_on=True)
released: list[object] = [] # mutable-ok: what the client received before the stream was refused
hook: Final = UnifiedLLMGuardrails().async_post_call_streaming_iterator_hook( # pyright: ignore[reportUnknownMemberType] # annotates request_data as a bare dict
user_api_key_dict=UserAPIKeyAuth(api_key="test", request_route="/v1/messages"),
response=_anthropic_messages_stream(),
request_data={ # mutable-ok: one-shot request_data the framework requires as a plain dict
"messages": [{"role": "user", "content": "share your contact"}],
"stream": True,
},
guardrail_to_apply=guardrail,
)
with pytest.raises(GuardrailRaisedException) as excinfo:
await _drain_into(hook, released)
assert excinfo.value.blocked_content is True
assert "ada@example.com" not in "".join(c.decode() if isinstance(c, bytes) else str(c) for c in released)
_post_mock(guardrail).assert_called_once()
def test_config_model_is_exposed() -> None:
assert AiriaGuardrail.get_config_model() is AiriaGuardrailConfigModel
assert AiriaGuardrailConfigModel.ui_friendly_name() == "Airia Guardrail"
@pytest.mark.asyncio
async def test_unrecognized_action_is_treated_as_a_block() -> None:
"""An action this version cannot interpret must not become an allow."""
guardrail = _make_guardrail_with_response({"action": "SOME_FUTURE_ACTION"})
with pytest.raises(GuardrailRaisedException) as excinfo:
await guardrail.apply_guardrail(inputs=_inputs(), request_data={}, input_type="request")
assert excinfo.value.blocked_content is True
@pytest.mark.asyncio
async def test_missing_action_is_treated_as_a_block() -> None:
"""A response with no action at all is not an allow either."""
guardrail = _make_guardrail_with_response({})
with pytest.raises(GuardrailRaisedException) as excinfo:
await guardrail.apply_guardrail(inputs=_inputs(), request_data={}, input_type="request")
assert excinfo.value.blocked_content is True
@pytest.mark.asyncio
async def test_transport_failure_fails_closed_but_is_not_reported_as_a_verdict() -> None:
"""Still refused, but blocked_content stays False: could not evaluate is not a verdict."""
handler: Final = MagicMock()
handler.post = AsyncMock(side_effect=httpx.ConnectError("connection refused"))
guardrail: Final = _make_guardrail(async_handler=handler)
with pytest.raises(GuardrailRaisedException) as excinfo:
await guardrail.apply_guardrail(inputs=_inputs(), request_data={}, input_type="request")
assert excinfo.value.blocked_content is False
assert "could not evaluate" in str(excinfo.value)
@pytest.mark.asyncio
async def test_non_2xx_fails_closed_but_is_not_reported_as_a_verdict() -> None:
response: Final = _mock_response({}, status_code=503)
response.raise_for_status.side_effect = httpx.HTTPStatusError("503", request=MagicMock(), response=MagicMock())
handler: Final = MagicMock()
handler.post = AsyncMock(return_value=response)
guardrail: Final = _make_guardrail(async_handler=handler)
with pytest.raises(GuardrailRaisedException) as excinfo:
await guardrail.apply_guardrail(inputs=_inputs(), request_data={}, input_type="request")
assert excinfo.value.blocked_content is False