From 612ca2cbc6e9a76bebbafabfb118705598627693 Mon Sep 17 00:00:00 2001 From: Airia Open Source Date: Fri, 11 Sep 2026 20:04:51 +0000 Subject: [PATCH] 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 or non-2xx 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, or a rewrite of the wrong shape, 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. --- .../guardrail_hooks/airia/__init__.py | 32 ++ .../guardrails/guardrail_hooks/airia/airia.py | 170 ++++++++ litellm/types/guardrails.py | 1 + .../proxy/guardrails/guardrail_hooks/airia.py | 24 ++ ruff-strict.toml | 1 + .../guardrails/guardrail_hooks/test_airia.py | 368 ++++++++++++++++++ 6 files changed, 596 insertions(+) create mode 100644 litellm/proxy/guardrails/guardrail_hooks/airia/__init__.py create mode 100644 litellm/proxy/guardrails/guardrail_hooks/airia/airia.py create mode 100644 litellm/types/proxy/guardrails/guardrail_hooks/airia.py create mode 100644 tests/test_litellm/proxy/guardrails/guardrail_hooks/test_airia.py diff --git a/litellm/proxy/guardrails/guardrail_hooks/airia/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/airia/__init__.py new file mode 100644 index 00000000000..3ae1e33f813 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/airia/__init__.py @@ -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.get("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 = { # mutable-ok: module-level registry, built once and never mutated + SupportedGuardrailIntegrations.AIRIA.value: initialize_guardrail, +} + +guardrail_class_registry: Final = { # mutable-ok: module-level registry, built once and never mutated + SupportedGuardrailIntegrations.AIRIA.value: AiriaGuardrail, +} diff --git a/litellm/proxy/guardrails/guardrail_hooks/airia/airia.py b/litellm/proxy/guardrails/guardrail_hooks/airia/airia.py new file mode 100644 index 00000000000..5fc76cb85eb --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/airia/airia.py @@ -0,0 +1,170 @@ +import os +import uuid +from typing import ( + TYPE_CHECKING, + 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, +) +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.custom_httpx.http_handler import ( + get_async_httpx_client, + httpxSpecialProvider, +) +from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.utils import GenericGuardrailAPIInputs + +if TYPE_CHECKING: + from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel + +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) + + +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__ + **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 = get_async_httpx_client( + llm_provider=httpxSpecialProvider.GuardrailCallback, + params={"timeout": httpx.Timeout(timeout=resolved_timeout, connect=5.0)}, # mutable-ok: one-shot params + ) + + 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 + + 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__( + supported_event_hooks=supported_event_hooks or [*SUPPORTED_EVENT_HOOKS], # mutable-ok: base needs a list + **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( + f"{self.api_base}{GUARDRAIL_PATH}", + json={ # mutable-ok: httpx needs a plain dict; built once and sent + "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}"}, # mutable-ok: one-shot HTTP headers + ) + response.raise_for_status() + body: Final = response.json() + 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: + raise GuardrailRaisedException( + guardrail_name=self.guardrail_name, + message=body.get("blocked_reason") or DEFAULT_BLOCKED_MESSAGE, + blocked_content=True, + ) + + if action == ACTION_INTERVENED: + return self._rewritten(body, inputs) + + if action != ACTION_NONE: + raise self._blocked() + + return inputs + + def _rewritten( + self, + body: dict[str, object], # mutable-ok: response.json() returns a plain dict + inputs: GenericGuardrailAPIInputs, + ) -> GenericGuardrailAPIInputs: + texts: Final = body.get("texts") + structured_messages: Final = body.get("structured_messages") + if texts is None and structured_messages is None: + raise self._blocked() + if not (texts is None or isinstance(texts, list)): + raise self._blocked() + if not (structured_messages is None or isinstance(structured_messages, list)): + raise self._blocked() + + rewritten: Final[GenericGuardrailAPIInputs] = {**inputs} # mutable-ok: fresh copy; caller's object untouched + if isinstance(texts, list): + rewritten["texts"] = texts + if isinstance(structured_messages, list): + rewritten["structured_messages"] = structured_messages + return rewritten + + def _blocked(self) -> GuardrailRaisedException: + return GuardrailRaisedException( + guardrail_name=self.guardrail_name, + message=DEFAULT_BLOCKED_MESSAGE, + blocked_content=True, + ) + + @staticmethod + def get_config_model() -> type["GuardrailConfigModel"] | None: + from litellm.types.proxy.guardrails.guardrail_hooks.airia import ( + AiriaGuardrailConfigModel, + ) + + return AiriaGuardrailConfigModel diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 02dee40f2a3..b51bb48b59b 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -137,6 +137,7 @@ class SupportedGuardrailIntegrations(Enum): COMPRESR = "compresr" STRAIKER = "straiker" ALICE = "alice" + AIRIA = "airia" class Role(Enum): diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/airia.py b/litellm/types/proxy/guardrails/guardrail_hooks/airia.py new file mode 100644 index 00000000000..ec48f49d593 --- /dev/null +++ b/litellm/types/proxy/guardrails/guardrail_hooks/airia.py @@ -0,0 +1,24 @@ +from pydantic import Field + +from .base import GuardrailConfigModel + + +class AiriaGuardrailConfigModel(GuardrailConfigModel): + 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" diff --git a/ruff-strict.toml b/ruff-strict.toml index ae092bdde7d..5026d206af8 100644 --- a/ruff-strict.toml +++ b/ruff-strict.toml @@ -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"] [lint.mccabe] max-complexity = 15 diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_airia.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_airia.py new file mode 100644 index 00000000000..50d6a97ad00 --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_airia.py @@ -0,0 +1,368 @@ +import uuid +from typing import Literal, cast +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 +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) -> AiriaGuardrail: + """Build a guardrail; tests replace `async_handler.post`, so no socket is ever opened.""" + kwargs.setdefault("api_base", API_BASE) + kwargs.setdefault("api_key", API_KEY) + kwargs.setdefault("guardrail_name", "airia-guard") + return AiriaGuardrail(**kwargs) + + +def _mock_response(payload: dict, 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 _inputs(**overrides: object) -> GenericGuardrailAPIInputs: + return cast( + GenericGuardrailAPIInputs, {"texts": ["hello"], "model": "gpt-4o", **overrides} + ) # cast-ok: test fixture + + +def test_airia_guardrail_config(monkeypatch: pytest.MonkeyPatch): + """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 registered[0].async_handler.timeout.read == 45.0 + + +def test_during_call_is_not_a_supported_event_hook(): + """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): + """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): + 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(): + guardrail = _make_guardrail(api_base=f"{API_BASE}/") + + assert guardrail.api_base == API_BASE + + +def test_custom_timeout_from_kwargs(): + guardrail = _make_guardrail(timeout=45.0) + + assert guardrail.async_handler.timeout.read == 45.0 + assert guardrail.async_handler.timeout.connect == 5.0 + + +def test_timeout_falls_back_to_environment(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("AIRIA_TIMEOUT", "30") + + guardrail = _make_guardrail() + + assert guardrail.async_handler.timeout.read == 30.0 + + +@pytest.mark.asyncio +async def test_request_payload_and_auth_header(): + guardrail = _make_guardrail() + guardrail.async_handler.post = AsyncMock(return_value=_mock_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 = guardrail.async_handler.post.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"]): + guardrail = _make_guardrail() + guardrail.async_handler.post = AsyncMock(return_value=_mock_response({"action": "NONE"})) + + await guardrail.apply_guardrail(inputs=_inputs(), request_data={}, input_type=input_type) + + assert guardrail.async_handler.post.call_args.kwargs["json"]["input_type"] == input_type + + +@pytest.mark.asyncio +async def test_call_id_falls_back_to_a_generated_uuid(): + guardrail = _make_guardrail() + guardrail.async_handler.post = AsyncMock(return_value=_mock_response({"action": "NONE"})) + + await guardrail.apply_guardrail(inputs=_inputs(), request_data={}, input_type="request") + + call_id = guardrail.async_handler.post.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(): + guardrail = _make_guardrail() + guardrail.async_handler.post = AsyncMock(return_value=_mock_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(): + guardrail = _make_guardrail() + guardrail.async_handler.post = AsyncMock( + return_value=_mock_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(): + guardrail = _make_guardrail() + guardrail.async_handler.post = AsyncMock(return_value=_mock_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(): + """Substituting only one field would leave the other still carrying the unredacted text.""" + guardrail = _make_guardrail() + guardrail.async_handler.post = AsyncMock( + return_value=_mock_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(): + """An omitted field means "not rewritten", not "rewritten to empty".""" + guardrail = _make_guardrail() + guardrail.async_handler.post = AsyncMock( + return_value=_mock_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(): + """A rewrite this version cannot apply must not let the original text through.""" + guardrail = _make_guardrail() + guardrail.async_handler.post = AsyncMock( + return_value=_mock_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(): + guardrail = _make_guardrail() + guardrail.async_handler.post = AsyncMock( + return_value=_mock_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(): + """An intervention the proxy cannot apply must not let the original content through.""" + guardrail = _make_guardrail() + guardrail.async_handler.post = AsyncMock( + return_value=_mock_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 +async def test_intervened_returns_a_copy_and_leaves_the_caller_object_untouched(): + guardrail = _make_guardrail() + guardrail.async_handler.post = AsyncMock( + return_value=_mock_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"] + + +def test_streaming_moderates_the_whole_response_then_emits_it_redacted(): + """block_only (the default) drops rewrites; per-chunk rounds can underflow when an entity spans them.""" + guardrail = _make_guardrail() + + assert guardrail.streaming_transform_mode == "incremental_diff" + assert guardrail.streaming_end_of_stream_only is True + + +def test_config_model_is_exposed(): + 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(): + """An action this version cannot interpret must not become an allow.""" + guardrail = _make_guardrail() + guardrail.async_handler.post = AsyncMock(return_value=_mock_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(): + """A response with no action at all is not an allow either.""" + guardrail = _make_guardrail() + guardrail.async_handler.post = AsyncMock(return_value=_mock_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(): + """Still refused, but blocked_content stays False: could not evaluate is not a verdict.""" + guardrail = _make_guardrail() + guardrail.async_handler.post = AsyncMock(side_effect=httpx.ConnectError("connection refused")) + + 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(): + guardrail = _make_guardrail() + response = _mock_response({}, status_code=503) + response.raise_for_status.side_effect = httpx.HTTPStatusError("503", request=MagicMock(), response=MagicMock()) + guardrail.async_handler.post = AsyncMock(return_value=response) + + with pytest.raises(GuardrailRaisedException) as excinfo: + await guardrail.apply_guardrail(inputs=_inputs(), request_data={}, input_type="request") + + assert excinfo.value.blocked_content is False