diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 3346b0c9ff8..e09cceb501c 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -13498,6 +13498,18 @@ "description": "If True, will not raise an exception when the guardrail is blocked. Useful for OpenWebUI where exceptions can end the chat flow.", "title": "Disable Exception On Block" }, + "enable_thinking": { + "anyOf": [ + { + "type": "boolean" + }, + { + "type": "null" + } + ], + "description": "Whether MLS should render the chat template in thinking mode. Defaults to False.", + "title": "Enable Thinking" + }, "end_session_after_n_fails": { "anyOf": [ { @@ -13668,6 +13680,18 @@ "description": "Enable hallucination detection to detect factual inaccuracies.", "title": "Hallucinations Check" }, + "hazard_threshold": { + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + } + ], + "description": "Block the request when the hazard_prompt probe scores strictly above this value. Defaults to 0.703, the threshold MLS reports for that probe. Note the probe also responds to instruction-style phrasing such as \"repeat this back verbatim\", so raise this if benign traffic is being blocked.", + "title": "Hazard Threshold" + }, "include_evidence": { "anyOf": [ { @@ -14306,6 +14330,24 @@ "description": "Optional per-entity minimum confidence scores for Presidio detections. Entities below the threshold are ignored.", "title": "Presidio Score Thresholds" }, + "probes": { + "anyOf": [ + { + "items": { + "type": "string" + }, + "type": "array" + }, + { + "type": "string" + }, + { + "type": "null" + } + ], + "description": "Which classifier probes to run: a list of probe names, or \"all\". Defaults to [\"hazard_prompt\"] - the only probe whose score this guardrail enforces. An unknown probe name makes MLS return 404.", + "title": "Probes" + }, "project_id": { "anyOf": [ { diff --git a/litellm/proxy/guardrails/guardrail_hooks/realmlabs/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/realmlabs/__init__.py new file mode 100644 index 00000000000..4d06c5af79e --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/realmlabs/__init__.py @@ -0,0 +1,72 @@ +from types import MappingProxyType +from typing import TYPE_CHECKING, Final + +from pydantic import BaseModel + +from litellm.types.guardrails import GuardrailEventHooks, Mode, SupportedGuardrailIntegrations +from litellm.types.proxy.guardrails.guardrail_hooks.realmlabs import RealmLabsGuardrailOptionalParams + +from .realmlabs import RealmLabsGuardrail + +if TYPE_CHECKING: + from litellm.types.guardrails import Guardrail, LitellmParams + +__all__ = ("RealmLabsGuardrail",) + + +def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail") -> RealmLabsGuardrail: + """Build the guardrail from its ``config.yaml`` entry and register it; found via the registries below.""" + import litellm + + settings: Final = _resolved_params(litellm_params) + _realmlabs_callback: Final = RealmLabsGuardrail( + api_key=litellm_params.api_key, + api_base=litellm_params.api_base, + probes=settings.probes, + hazard_threshold=settings.hazard_threshold, + block_on_error=settings.block_on_error, + enable_thinking=settings.enable_thinking, + timeout=settings.timeout, + guardrail_name=guardrail["guardrail_name"], + event_hook=_coerce_event_hook(litellm_params.mode), + default_on=litellm_params.default_on or False, + ) + litellm.logging_callback_manager.add_litellm_callback( # pyright: ignore[reportUnknownMemberType] # callback manager is untyped + _realmlabs_callback + ) + return _realmlabs_callback + + +guardrail_initializer_registry: Final = { + SupportedGuardrailIntegrations.REALMLABS.value: initialize_guardrail, +} + +guardrail_class_registry: Final = { + SupportedGuardrailIntegrations.REALMLABS.value: RealmLabsGuardrail, +} + + +def _coerce_event_hook( + mode: str | list[str] | Mode, # mutable-ok: mirrors LitellmParams.mode +) -> GuardrailEventHooks | list[GuardrailEventHooks] | Mode: # mutable-ok: mirrors CustomGuardrail's event_hook + """Convert the ``mode`` strings from ``config.yaml`` into the enum values ``CustomGuardrail`` expects.""" + if isinstance(mode, Mode): + return mode + if isinstance(mode, list): + return [GuardrailEventHooks(item) for item in mode] + return GuardrailEventHooks(mode) + + +def _resolved_params(litellm_params: "LitellmParams") -> RealmLabsGuardrailOptionalParams: + """Resolve explicit nested values, then top-level values, then RealmLabs defaults. + + Excluding unset fields prevents another guardrail's parsed defaults from overriding RealmLabs settings. + Null nested values fall back to the top level; false, zero, and empty lists remain explicit overrides. + """ + top_level: Final = RealmLabsGuardrailOptionalParams.model_validate(litellm_params.model_dump(exclude_unset=True)) + nested: Final = litellm_params.optional_params + if not isinstance(nested, BaseModel): + return top_level + return RealmLabsGuardrailOptionalParams.model_validate( + MappingProxyType({**top_level.model_dump(), **nested.model_dump(exclude_unset=True, exclude_none=True)}) + ) diff --git a/litellm/proxy/guardrails/guardrail_hooks/realmlabs/realmlabs.py b/litellm/proxy/guardrails/guardrail_hooks/realmlabs/realmlabs.py new file mode 100644 index 00000000000..9ccc142f013 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/realmlabs/realmlabs.py @@ -0,0 +1,208 @@ +"""RealmLabs MLS request hazard guardrail.""" + +from __future__ import annotations + +import json +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from typing import TYPE_CHECKING, Final, Literal + +from httpx import HTTPError +from pydantic import TypeAdapter, ValidationError + +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] # decorator is untyped in custom_guardrail +) +from litellm.llms.custom_httpx.http_handler import ( + get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # helper is untyped in http_handler + httpxSpecialProvider, +) +from litellm.secret_managers.main import get_secret_str +from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.proxy.guardrails.guardrail_hooks.realmlabs import ( + RealmLabsChatMessage, + RealmLabsGuardrailConfigModel, + RealmLabsGuardrailRequest, + RealmLabsGuardrailResponse, +) + +if TYPE_CHECKING: + from httpx import Response as HttpxResponse + + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + from litellm.types.guardrails import Mode + from litellm.types.utils import GenericGuardrailAPIInputs + +_DEFAULT_API_BASE: Final = "https://mls.realmlabs.ai" +_GUARDRAIL_ENDPOINT: Final = "/guardrail" +_HAZARD_PROBE: Final = "hazard_prompt" +_DEFAULT_HAZARD_THRESHOLD: Final = 0.703 +_DEFAULT_TIMEOUT: Final = 15.0 + +_RESPONSE_ADAPTER: Final = TypeAdapter(RealmLabsGuardrailResponse) + + +class RealmLabsMissingCredentials(Exception): + """Raised at startup when no MLS API key is configured.""" + + +@dataclass(frozen=True, slots=True) +class _InvalidResponse: + reason: str + + +class RealmLabsGuardrail(CustomGuardrail): + """Blocks hazardous prompts using the RealmLabs MLS endpoint.""" + + def __init__( + self, + api_key: str | None = None, + api_base: str | None = None, + probes: Sequence[str] | str | None = None, + hazard_threshold: float | None = None, + block_on_error: bool | None = None, + enable_thinking: bool | None = None, + timeout: float | None = None, + guardrail_name: str | None = None, + event_hook: ( # mutable-ok: same type as CustomGuardrail + GuardrailEventHooks | list[GuardrailEventHooks] | Mode | None + ) = None, + default_on: bool = False, + ) -> None: + """Resolve each setting from its argument, then the ``REALMLABS_*`` env vars, then the module defaults.""" + self.api_key = api_key or get_secret_str("REALMLABS_API_KEY") + if not self.api_key: + raise RealmLabsMissingCredentials( + "RealmLabs API key is required. Set REALMLABS_API_KEY in the environment " + "or pass api_key in the guardrail config." + ) + + self.api_base = (api_base or get_secret_str("REALMLABS_API_BASE") or _DEFAULT_API_BASE).rstrip("/") + self.probes: Sequence[str] | str = (_HAZARD_PROBE,) if probes is None else probes + self.hazard_threshold = _DEFAULT_HAZARD_THRESHOLD if hazard_threshold is None else hazard_threshold + self.block_on_error = False if block_on_error is None else block_on_error + self.enable_thinking = False if enable_thinking is None else enable_thinking + self.timeout = _DEFAULT_TIMEOUT if timeout is None else timeout + self.async_handler: AsyncHTTPHandler = get_async_httpx_client( + llm_provider=httpxSpecialProvider.GuardrailCallback, + ) + super().__init__( # pyright: ignore[reportUnknownMemberType] # CustomGuardrail.__init__ is untyped + guardrail_name=guardrail_name, + supported_event_hooks=self.get_supported_event_hooks(), + event_hook=event_hook, + default_on=default_on, + ) + + @staticmethod + def get_config_model() -> type[RealmLabsGuardrailConfigModel] | None: + """Config model the admin UI uses to render and validate this guardrail's settings.""" + return RealmLabsGuardrailConfigModel + + @classmethod + def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]: # mutable-ok: base class returns a list + """Inspect requests on ``pre_call``.""" + return [GuardrailEventHooks.pre_call] + + @staticmethod + def _hazard_score(response: RealmLabsGuardrailResponse) -> float | None: + """Score of the first ``hazard_prompt`` verdict, unless MLS reports a role mismatch.""" + for result in response["results"]: + if result["probe"] == _HAZARD_PROBE: + return None if result.get("role_mismatch") else result["prob"] + return None + + def _build_request(self, messages: Sequence[Mapping[str, object]]) -> RealmLabsGuardrailRequest: + return RealmLabsGuardrailRequest( + messages=messages, + probes=self.probes, + pii=False, + enable_thinking=self.enable_thinking, + ) + + def _parse_response(self, content: bytes) -> RealmLabsGuardrailResponse | _InvalidResponse: + try: + result: Final = _RESPONSE_ADAPTER.validate_json(content) + except ValidationError as exc: + return _InvalidResponse(exc.json(include_input=False, include_context=False, include_url=False)) + + return result + + async def _call_mls( + self, messages: Sequence[Mapping[str, object]] + ) -> RealmLabsGuardrailResponse | _InvalidResponse: + endpoint: Final = f"{self.api_base}{_GUARDRAIL_ENDPOINT}" + verbose_proxy_logger.debug( + "RealmLabs MLS: %s msgs=%d probes=%s", + endpoint, + len(messages), + self.probes, + ) + response: Final[HttpxResponse] = await self.async_handler.post( # pyright: ignore[reportUnknownMemberType] # AsyncHTTPHandler.post is untyped + url=endpoint, + headers={ + "Authorization": f"Bearer {self.api_key}", + "Content-Type": "application/json", + }, + content=json.dumps(self._build_request(messages)), + timeout=self.timeout, + ) + response.raise_for_status() + return self._parse_response(response.content) + + def _handle_mls_error(self, inputs: GenericGuardrailAPIInputs, reason: str) -> GenericGuardrailAPIInputs: + verbose_proxy_logger.error("RealmLabs MLS error: %s", reason) + if self.block_on_error: + raise GuardrailRaisedException( + guardrail_name=self.guardrail_name, + message=f"RealmLabs MLS error (block_on_error=True): {reason}", + ) + return inputs + + @log_guardrail_information + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: Mapping[str, object], + input_type: Literal["request", "response"], + logging_obj: LiteLLMLoggingObj | None = None, + ) -> GenericGuardrailAPIInputs: + """Screen requests for hazardous content through LiteLLM's unified guardrail layer.""" + if input_type != "request": + return inputs + texts: Final = tuple(inputs.get("texts") or ()) + messages: Final = tuple(inputs.get("structured_messages") or ()) or tuple( + RealmLabsChatMessage(role="user", content=text) for text in texts + ) + if not messages: + return inputs + + try: + result: Final = await self._call_mls(messages) + except (HTTPError, TypeError, ValueError) as exc: + return self._handle_mls_error(inputs, str(exc)) + + if isinstance(result, _InvalidResponse): + return self._handle_mls_error(inputs, f"Invalid RealmLabs guardrail response: {result.reason}") + + hazard_score: Final = self._hazard_score(result) + if hazard_score is not None and hazard_score > self.hazard_threshold: + verbose_proxy_logger.warning( + "RealmLabs MLS blocked request: %s=%s > %s", + _HAZARD_PROBE, + hazard_score, + self.hazard_threshold, + ) + raise GuardrailRaisedException( + guardrail_name=self.guardrail_name, + message=( + f"Blocked by RealmLabs {_HAZARD_PROBE} probe: " + f"score={hazard_score} exceeds threshold={self.hazard_threshold}" + ), + blocked_content=True, + ) + + return inputs diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 46026c12d24..ed416e8718c 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -53,6 +53,9 @@ from litellm.types.proxy.guardrails.guardrail_hooks.qohash import ( from litellm.types.proxy.guardrails.guardrail_hooks.qualifire import ( QualifireGuardrailConfigModel, ) +from litellm.types.proxy.guardrails.guardrail_hooks.realmlabs import ( + RealmLabsGuardrailConfigModel, +) from litellm.types.proxy.guardrails.guardrail_hooks.repelloai import ( RepelloAIGuardrailConfigModel, ) @@ -120,6 +123,7 @@ class SupportedGuardrailIntegrations(Enum): MCP_SECURITY = "mcp_security" ONYX = "onyx" PROMPTGUARD = "promptguard" + REALMLABS = "realmlabs" XECGUARD = "xecguard" PROMPT_SECURITY = "prompt_security" GENERIC_GUARDRAIL_API = "generic_guardrail_api" @@ -1187,6 +1191,7 @@ class LitellmParams( # pyright: ignore[reportIncompatibleVariableOverride] # o GraySwanGuardrailConfigModel, NomaGuardrailConfigModel, PromptGuardConfigModel, + RealmLabsGuardrailConfigModel, XecGuardConfigModel, ToolPermissionGuardrailConfigModel, ZscalerAIGuardConfigModel, diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/realmlabs.py b/litellm/types/proxy/guardrails/guardrail_hooks/realmlabs.py new file mode 100644 index 00000000000..9c7c8072b33 --- /dev/null +++ b/litellm/types/proxy/guardrails/guardrail_hooks/realmlabs.py @@ -0,0 +1,122 @@ +from collections.abc import Mapping, Sequence +from typing import Annotated + +from pydantic import BaseModel, Field +from typing_extensions import ReadOnly, Required, TypedDict + +from .base import GuardrailConfigModel + + +class RealmLabsChatMessage(TypedDict): + """A plain-text chat turn sent to MLS.""" + + role: ReadOnly[str] + content: ReadOnly[str] + + +class RealmLabsGuardrailRequest(TypedDict): + messages: ReadOnly[Sequence[Mapping[str, object]]] + probes: ReadOnly[Sequence[str] | str] + pii: ReadOnly[bool] + enable_thinking: ReadOnly[bool] + + +class RealmLabsProbeResult(TypedDict, total=False): + """One classifier probe verdict. + + ``prob`` is compared against the guardrail's own ``hazard_threshold``, not the ``threshold`` MLS reports. + """ + + probe: ReadOnly[Required[Annotated[str, Field(strict=True, min_length=1)]]] + prob: ReadOnly[Required[Annotated[float, Field(strict=True, ge=0, le=1, allow_inf_nan=False)]]] + threshold: ReadOnly[float | None] + decision: ReadOnly[bool | None] + role_mismatch: ReadOnly[Annotated[bool, Field(strict=True)] | None] + + +class RealmLabsGuardrailResponse(TypedDict, total=False): + """Response body of ``POST {api_base}/guardrail``. MLS is stateless, so it carries no turn id.""" + + results: ReadOnly[Required[Sequence[RealmLabsProbeResult]]] + focal_role: ReadOnly[str | None] + + +class RealmLabsGuardrailOptionalParams(BaseModel): + """Nested tuning settings; explicitly supplied non-null values override the top-level settings.""" + + probes: Sequence[str] | str | None = Field( + default=None, + description="Classifier probes to run: a list of names or 'all'. Overrides top-level probes when supplied.", + ) + hazard_threshold: float | None = Field( + default=None, + description="Block hazard scores strictly above this value. Overrides top-level hazard_threshold when supplied.", + ) + block_on_error: bool | None = Field( + default=None, + description="Whether to block when MLS fails. Overrides top-level block_on_error when supplied.", + ) + + enable_thinking: bool | None = Field( + default=False, + description="Whether MLS should render the chat template in thinking mode. Defaults to False.", + ) + + timeout: float | None = Field( + default=15.0, + description="Timeout in seconds for the MLS request. Defaults to 15.", + ) + + +class RealmLabsGuardrailConfigModel(GuardrailConfigModel[RealmLabsGuardrailOptionalParams]): + """Settings accepted under ``litellm_params`` for ``guardrail: realmlabs``.""" + + api_key: str | None = Field( + default=None, + description=( + "API key for the RealmLabs MLS guardrail endpoint, sent as a bearer token. " + "If not provided, the REALMLABS_API_KEY environment variable is used." + ), + ) + api_base: str | None = Field( + default=None, + description=( + "Base URL of the RealmLabs MLS deployment. The /guardrail path is " + "appended automatically. Defaults to https://mls.realmlabs.ai, and falls " + "back to the REALMLABS_API_BASE environment variable." + ), + ) + probes: list[str] | str | None = Field( # mutable-ok: the admin UI renders only list[...] fields as array inputs + default=None, + description=( + 'Which classifier probes to run: a list of probe names, or "all". ' + 'Defaults to ["hazard_prompt"] - the only probe whose score this ' + "guardrail enforces. An unknown probe name makes MLS return 404." + ), + ) + hazard_threshold: float | None = Field( + default=None, + description=( + "Block the request when the hazard_prompt probe scores strictly above this " + "value. Defaults to 0.703, the threshold MLS reports for that probe. Note " + 'the probe also responds to instruction-style phrasing such as "repeat this ' + 'back verbatim", so raise this if benign traffic is being blocked.' + ), + ) + block_on_error: bool | None = Field( + default=None, + description=( + "Whether to block the request when MLS is unreachable or returns an " + "unreadable response. Defaults to False (fail open), so an MLS outage does " + "not take the gateway down with it. Set to True to fail closed." + ), + ) + enable_thinking: bool | None = Field( + default=None, + description="Whether MLS should render the chat template in thinking mode. Defaults to False.", + ) + + @staticmethod + def ui_friendly_name() -> str: + """Name the admin UI shows for this guardrail.""" + return "RealmLabs MLS" diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/realmlabs/__init__.py b/tests/unit/proxy/guardrails/guardrail_hooks/realmlabs/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/realmlabs/test_realmlabs.py b/tests/unit/proxy/guardrails/guardrail_hooks/realmlabs/test_realmlabs.py new file mode 100644 index 00000000000..a5ac4c6de2e --- /dev/null +++ b/tests/unit/proxy/guardrails/guardrail_hooks/realmlabs/test_realmlabs.py @@ -0,0 +1,559 @@ +import json +from collections.abc import Mapping +from http import HTTPStatus +from typing import Final, Literal + +import httpx +import pytest +import respx +from pydantic import TypeAdapter + +import litellm +from litellm.exceptions import GuardrailRaisedException +from litellm.proxy.guardrails.guardrail_hooks.realmlabs import guardrail_initializer_registry +from litellm.proxy.guardrails.guardrail_hooks.realmlabs.realmlabs import ( + RealmLabsGuardrail, + RealmLabsMissingCredentials, +) +from litellm.types.guardrails import GuardrailEventHooks, LitellmParams +from litellm.types.llms.openai import AllMessageValues +from litellm.types.proxy.guardrails.guardrail_hooks.realmlabs import RealmLabsGuardrailOptionalParams +from litellm.types.utils import GenericGuardrailAPIInputs + +_API_KEY = "mls_gr_test" +_API_BASE = "https://mls.example.test" +_URL = f"{_API_BASE}/guardrail" +_DEFAULT_THRESHOLD = 0.703 + + +_JSON_OBJECT = TypeAdapter(dict[str, object]) + + +@pytest.fixture(autouse=True) +def fresh_httpx_client(monkeypatch: pytest.MonkeyPatch) -> None: + """Route MLS calls through a fresh httpx client that ``respx`` can intercept, and clear ``REALMLABS_*`` env vars.""" + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", None) + monkeypatch.delenv("REALMLABS_API_KEY", raising=False) + monkeypatch.delenv("REALMLABS_API_BASE", raising=False) + + +def _configured_guardrail(settings: Mapping[str, object]) -> RealmLabsGuardrail: + """Build a guardrail through the real config model and initializer, keeping each test's settings visible.""" + + params: Final = LitellmParams.model_validate( + { + "guardrail": "realmlabs", + "mode": "pre_call", + "api_key": _API_KEY, + "api_base": _API_BASE, + **settings, + } + ) + return guardrail_initializer_registry["realmlabs"]( + params, {"guardrail_name": "rl-config", "litellm_params": params} + ) + + +def _guardrail( + hazard_threshold: float | None = None, + block_on_error: bool | None = None, + event_hook: GuardrailEventHooks = GuardrailEventHooks.pre_call, +) -> RealmLabsGuardrail: + """Guardrail pointed at the fake MLS URL; unset arguments keep the guardrail's defaults.""" + + return RealmLabsGuardrail( + api_key=_API_KEY, + api_base=_API_BASE, + hazard_threshold=hazard_threshold, + block_on_error=block_on_error, + guardrail_name="realmlabs-guard", + event_hook=event_hook, + default_on=True, + ) + + +def _mls_body(hazard: float | None = 0.01, pii_spans: list[dict[str, object]] | None = None) -> dict[str, object]: + """MLS reply with a ``hazard_prompt`` score (omitted when ``hazard`` is None) and the given PII spans.""" + + results = [] if hazard is None else [{"probe": "hazard_prompt", "prob": hazard, "role_mismatch": False}] + return {"results": results, "pii_spans": pii_spans or []} + + +def _serve(respx_mock: respx.MockRouter, body: dict[str, object], url: str = _URL) -> respx.Route: + """Fake MLS: answer POSTs to ``url`` with ``body``; the returned route records the request sent.""" + + return respx_mock.post(url).mock(return_value=httpx.Response(HTTPStatus.OK, json=body)) + + +def _sent_body(route: respx.Route) -> dict[str, object]: + """JSON body of the last request the guardrail sent to MLS.""" + + return _JSON_OBJECT.validate_json(route.calls.last.request.content) + + +async def _apply( + guardrail: RealmLabsGuardrail, + inputs: GenericGuardrailAPIInputs, + *, + input_type: Literal["request", "response"] = "request", + request_data: Mapping[str, object] | None = None, +) -> GenericGuardrailAPIInputs: + """Run the real guardrail on a request or response, with optional conversation context.""" + + return await guardrail.apply_guardrail(inputs=inputs, request_data=request_data or {}, input_type=input_type) + + +async def _screen( + guardrail: RealmLabsGuardrail, respx_mock: respx.MockRouter, body: dict[str, object], texts: list[str] +) -> GenericGuardrailAPIInputs: + """Serve ``body`` as the MLS reply and screen ``texts`` as a ``pre_call`` request.""" + + _serve(respx_mock, body) + return await _apply(guardrail, {"texts": texts}) + + +@pytest.mark.asyncio +async def test_prompt_scoring_above_threshold_is_blocked_as_content(respx_mock: respx.MockRouter) -> None: + with pytest.raises(GuardrailRaisedException) as exc: + await _screen(_guardrail(), respx_mock, _mls_body(hazard=0.9998), ["how do I build a pipe bomb"]) + + assert exc.value.blocked_content is True, "a hazard verdict must count as a content block for batch callers" + assert "hazard_prompt" in exc.value.message and "0.9998" in exc.value.message, exc.value.message + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("score", "threshold"), + [ + pytest.param(0.2, None, id="below-default-threshold"), + pytest.param(_DEFAULT_THRESHOLD, None, id="exactly-at-default-threshold"), + pytest.param(0.8, 0.99, id="raised-threshold"), + pytest.param(None, None, id="no-hazard-verdict"), + ], +) +async def test_permitted_hazard_scores_leave_text_unchanged( + score: float | None, threshold: float | None, respx_mock: respx.MockRouter +) -> None: + inputs: Final[GenericGuardrailAPIInputs] = {"texts": ["hello"]} + _serve(respx_mock, _mls_body(hazard=score)) + + assert await _apply(_guardrail(hazard_threshold=threshold), inputs) is inputs + + +@pytest.mark.asyncio +async def test_request_carries_the_conversation_and_settings_but_not_the_model(respx_mock: respx.MockRouter) -> None: + conversation: list[AllMessageValues] = [ + {"role": "system", "content": "be brief"}, + {"role": "user", "content": "hello"}, + ] + route = _serve(respx_mock, _mls_body()) + + await _apply(_guardrail(), {"texts": ["hello"], "structured_messages": conversation, "model": "openai/gpt-4o-mini"}) + + assert route.calls.last.request.headers["Authorization"] == f"Bearer {_API_KEY}" + assert _sent_body(route) == { + "messages": conversation, + "probes": ["hazard_prompt"], + "pii": False, + "enable_thinking": False, + } + + +@pytest.mark.asyncio +async def test_plain_texts_are_sent_as_user_turns(respx_mock: respx.MockRouter) -> None: + route = _serve(respx_mock, _mls_body()) + + await _apply(_guardrail(), {"texts": ["a", "b"]}) + + assert _sent_body(route) == { + "messages": [{"role": "user", "content": "a"}, {"role": "user", "content": "b"}], + "probes": ["hazard_prompt"], + "pii": False, + "enable_thinking": False, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("input_type", ["request"]) +async def test_empty_inputs_are_returned_without_calling_mls( + input_type: Literal["request", "response"], respx_mock: respx.MockRouter +) -> None: + route = _serve(respx_mock, _mls_body()) + inputs: GenericGuardrailAPIInputs = {"texts": []} + + assert await _apply(_guardrail(), inputs, input_type=input_type) is inputs + assert route.call_count == 0, "empty input must not be billed as an MLS call" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("input_type", ["request"]) +@pytest.mark.parametrize("block_on_error", [None, True], ids=["default-fail-open", "fail-closed"]) +@pytest.mark.parametrize( + ("mls_reply", "reason"), + [ + pytest.param(httpx.ConnectError("connection refused"), "connection refused", id="connection-error"), + pytest.param( + httpx.Response(HTTPStatus.SERVICE_UNAVAILABLE), + f"{HTTPStatus.SERVICE_UNAVAILABLE.value} {HTTPStatus.SERVICE_UNAVAILABLE.phrase}", + id="server-error", + ), + pytest.param( + httpx.Response(HTTPStatus.OK, json={"results": "not a list"}), + "Invalid RealmLabs guardrail response", + id="invalid-body", + ), + ], +) +async def test_mls_failures_follow_the_error_policy( + mls_reply: httpx.Response | httpx.HTTPError, + reason: str, + block_on_error: bool | None, + input_type: Literal["request", "response"], + respx_mock: respx.MockRouter, + caplog: pytest.LogCaptureFixture, +) -> None: + route: Final = respx_mock.post(_URL) + if isinstance(mls_reply, httpx.HTTPError): + route.mock(side_effect=mls_reply) + else: + route.mock(return_value=mls_reply) + inputs: Final[GenericGuardrailAPIInputs] = {"texts": ["hello"]} + guardrail: Final = _guardrail(block_on_error=block_on_error) + + if block_on_error: + with pytest.raises(GuardrailRaisedException) as exc: + await _apply(guardrail, inputs, input_type=input_type) + assert exc.value.blocked_content is False + assert reason in exc.value.message + else: + assert await _apply(guardrail, inputs, input_type=input_type) is inputs + assert reason in caplog.text + + +@pytest.mark.asyncio +@pytest.mark.parametrize("input_type", ["request"]) +@pytest.mark.parametrize("block_on_error", [False, True]) +@pytest.mark.parametrize( + "body", + [ + pytest.param({}, id="empty-object"), + pytest.param({"choices": [{"message": {"content": "hello"}}]}, id="chat-envelope"), + pytest.param({"pii_spans": []}, id="missing-results"), + pytest.param( + {"results": [{"probe": "hazard_prompt", "probability": 0.99}], "pii_spans": []}, id="renamed-score" + ), + pytest.param({"results": [{"prob": 0.99}], "pii_spans": []}, id="missing-probe"), + pytest.param( + {"results": [{"probe": "hazard_prompt", "prob": 0.99, "role_mismatch": "true"}], "pii_spans": []}, + id="invalid-role-mismatch", + ), + ], +) +async def test_incomplete_verdicts_follow_the_error_policy( + body: dict[str, object], + block_on_error: bool, + input_type: Literal["request", "response"], + respx_mock: respx.MockRouter, + caplog: pytest.LogCaptureFixture, +) -> None: + _serve(respx_mock, body) + inputs: Final[GenericGuardrailAPIInputs] = {"texts": ["Hello Alex."]} + guardrail: Final = _guardrail(block_on_error=block_on_error) + + if block_on_error: + with pytest.raises(GuardrailRaisedException) as exc: + await _apply(guardrail, inputs, input_type=input_type) + assert exc.value.blocked_content is False, exc.value.message + assert "Invalid RealmLabs guardrail response" in exc.value.message, exc.value.message + else: + assert await _apply(guardrail, inputs, input_type=input_type) is inputs + assert "Invalid RealmLabs guardrail response" in caplog.text, caplog.text + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "prob", + [ + pytest.param(None, id="null"), + pytest.param("0.99", id="string"), + pytest.param(True, id="boolean"), + pytest.param(-0.1, id="negative"), + pytest.param(1.1, id="above-one"), + pytest.param(float("nan"), id="nan"), + pytest.param(float("inf"), id="infinity"), + ], +) +async def test_invalid_probabilities_cannot_bypass_fail_closed(prob: object, respx_mock: respx.MockRouter) -> None: + body: Final = {"results": [{"probe": "hazard_prompt", "prob": prob}], "pii_spans": []} + respx_mock.post(_URL).mock(return_value=httpx.Response(HTTPStatus.OK, content=json.dumps(body))) + + with pytest.raises(GuardrailRaisedException) as exc: + await _apply(_guardrail(block_on_error=True), {"texts": ["hello"]}) + + assert exc.value.blocked_content is False, exc.value.message + assert "Invalid RealmLabs guardrail response" in exc.value.message, exc.value.message + + +@pytest.mark.asyncio +@pytest.mark.parametrize("input_type", ["request"]) +async def test_empty_verdict_arrays_are_valid_in_fail_closed_mode( + input_type: Literal["request", "response"], respx_mock: respx.MockRouter +) -> None: + _serve(respx_mock, {"results": [], "pii_spans": []}) + inputs: Final[GenericGuardrailAPIInputs] = {"texts": ["hello"]} + + assert await _apply(_guardrail(block_on_error=True), inputs, input_type=input_type) is inputs + + +@pytest.mark.asyncio +@pytest.mark.parametrize("hazard", [False, True]) +async def test_additional_response_fields_preserve_policy_enforcement( + hazard: bool, respx_mock: respx.MockRouter +) -> None: + _serve( + respx_mock, + { + "future_metadata": {"version": 2}, + "results": [{"probe": "hazard_prompt", "prob": 1.0 if hazard else 0.0, "future_field": [1, 2]}], + "pii_spans": [], + }, + ) + guardrail: Final = _guardrail(hazard_threshold=0.5, block_on_error=True) + inputs: Final[GenericGuardrailAPIInputs] = {"texts": ["Hello Alex."]} + + if hazard: + with pytest.raises(GuardrailRaisedException) as exc: + await _apply(guardrail, inputs) + assert exc.value.blocked_content is True, exc.value.message + assert "hazard_prompt" in exc.value.message, exc.value.message + else: + assert await _apply(guardrail, inputs) is inputs + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("api_key", "api_base", "expected_key", "expected_base"), + [ + pytest.param(None, None, "env_key", "https://env.example.test", id="environment-fallback"), + pytest.param(_API_KEY, f"{_API_BASE}/", _API_KEY, _API_BASE, id="config-overrides-environment"), + ], +) +async def test_credentials_resolve_config_before_environment( + api_key: str | None, + api_base: str | None, + expected_key: str, + expected_base: str, + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +) -> None: + monkeypatch.setenv("REALMLABS_API_KEY", "env_key") + monkeypatch.setenv("REALMLABS_API_BASE", "https://env.example.test/") + route: Final = _serve(respx_mock, _mls_body(), f"{expected_base}/guardrail") + + await _apply(RealmLabsGuardrail(api_key=api_key, api_base=api_base), {"texts": ["hello"]}) + + assert route.call_count == 1 + assert route.calls.last.request.headers["Authorization"] == f"Bearer {expected_key}" + + +def test_missing_api_key_is_rejected_at_startup() -> None: + with pytest.raises(RealmLabsMissingCredentials): + RealmLabsGuardrail() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("settings", "expected_probes", "expected_thinking"), + [ + pytest.param( + {"probes": ["hazard_prompt", "dispute"], "enable_thinking": True}, + ["hazard_prompt", "dispute"], + True, + id="top-level", + ), + pytest.param( + {"optional_params": {"probes": "all", "enable_thinking": True}}, + "all", + True, + id="nested", + ), + pytest.param( + { + "probes": ["hazard_prompt"], + "enable_thinking": True, + "optional_params": {"probes": [], "enable_thinking": False}, + }, + [], + False, + id="nested-false-and-empty-list-win", + ), + pytest.param( + { + "probes": "all", + "enable_thinking": True, + "optional_params": {"probes": None, "enable_thinking": None}, + }, + "all", + True, + id="nested-null-falls-back", + ), + pytest.param( + {"probes": "all", "enable_thinking": True, "optional_params": {}}, + "all", + True, + id="empty-options-keep-top-level", + ), + ], +) +async def test_config_yaml_settings_reach_mls( + settings: dict[str, object], + expected_probes: list[str] | str, + expected_thinking: bool, + respx_mock: respx.MockRouter, +) -> None: + guardrail: Final = _configured_guardrail(settings) + route: Final = _serve(respx_mock, _mls_body()) + + await _apply(guardrail, {"texts": ["hello"]}) + + assert route.calls.last.request.headers["Authorization"] == f"Bearer {_API_KEY}" + assert _sent_body(route) == { + "messages": [{"role": "user", "content": "hello"}], + "probes": expected_probes, + "pii": False, + "enable_thinking": expected_thinking, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "settings", + [ + pytest.param({"hazard_threshold": 0.0}, id="top-level-zero"), + pytest.param({"optional_params": {"hazard_threshold": 0.0}}, id="nested-zero"), + pytest.param({"hazard_threshold": 0.99, "optional_params": {"hazard_threshold": 0.0}}, id="nested-zero-wins"), + ], +) +async def test_configured_zero_hazard_threshold_blocks_a_positive_score( + settings: dict[str, object], respx_mock: respx.MockRouter +) -> None: + guardrail: Final = _configured_guardrail(settings) + _serve(respx_mock, _mls_body(hazard=0.01)) + + with pytest.raises(GuardrailRaisedException) as exc: + await _apply(guardrail, {"texts": ["hello"]}) + + assert exc.value.blocked_content is True + assert "threshold=0.0" in exc.value.message, exc.value.message + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "settings", + [ + pytest.param({"hazard_threshold": 0.99, "optional_params": {}}, id="omitted-nested"), + pytest.param({"hazard_threshold": 0.99, "optional_params": {"hazard_threshold": None}}, id="null-nested"), + pytest.param({"hazard_threshold": 0.0, "optional_params": {"hazard_threshold": 0.99}}, id="nested-wins"), + ], +) +async def test_configured_higher_hazard_threshold_allows_the_request( + settings: dict[str, object], respx_mock: respx.MockRouter +) -> None: + guardrail: Final = _configured_guardrail(settings) + _serve(respx_mock, _mls_body(hazard=0.9)) + inputs: Final[GenericGuardrailAPIInputs] = {"texts": ["hello"]} + + assert await _apply(guardrail, inputs) is inputs + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "settings", + [ + pytest.param({"block_on_error": True}, id="top-level"), + pytest.param({"optional_params": {"block_on_error": True}}, id="nested"), + pytest.param({"block_on_error": False, "optional_params": {"block_on_error": True}}, id="nested-wins"), + pytest.param({"block_on_error": True, "optional_params": {"block_on_error": None}}, id="null-nested"), + ], +) +async def test_configured_fail_closed_blocks_an_mls_outage( + settings: dict[str, object], respx_mock: respx.MockRouter +) -> None: + guardrail: Final = _configured_guardrail(settings) + respx_mock.post(_URL).mock(side_effect=httpx.ConnectError("connection refused")) + + with pytest.raises(GuardrailRaisedException) as exc: + await _apply(guardrail, {"texts": ["hello"]}) + + assert exc.value.blocked_content is False + assert "block_on_error=True" in exc.value.message, exc.value.message + + +@pytest.mark.asyncio +async def test_nested_fail_open_overrides_top_level_fail_closed(respx_mock: respx.MockRouter) -> None: + guardrail: Final = _configured_guardrail({"block_on_error": True, "optional_params": {"block_on_error": False}}) + respx_mock.post(_URL).mock(side_effect=httpx.ConnectError("connection refused")) + inputs: Final[GenericGuardrailAPIInputs] = {"texts": ["hello"]} + + assert await _apply(guardrail, inputs) is inputs + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("settings", "expected_timeout"), + [ + pytest.param({"optional_params": None}, None, id="default-no-options"), + pytest.param({"optional_params": {}}, None, id="default-empty-options"), + pytest.param({"optional_params": {"enable_thinking": True}}, None, id="default-thinking-enabled"), + pytest.param({"optional_params": {"timeout": None}}, None, id="default-null-timeout"), + pytest.param({"timeout": 2.0, "optional_params": None}, 2.0, id="top-level-no-options"), + pytest.param({"timeout": 2.0, "optional_params": {}}, 2.0, id="top-level-empty-options"), + pytest.param( + {"timeout": 2.0, "optional_params": {"enable_thinking": True}}, 2.0, id="top-level-thinking-enabled" + ), + pytest.param({"timeout": 2.0, "optional_params": {"timeout": None}}, 2.0, id="top-level-null-timeout"), + pytest.param({"timeout": 0.0, "optional_params": None}, 0.0, id="zero-no-options"), + pytest.param({"timeout": 0.0, "optional_params": {}}, 0.0, id="zero-empty-options"), + pytest.param({"timeout": 0.0, "optional_params": {"enable_thinking": True}}, 0.0, id="zero-thinking-enabled"), + pytest.param({"timeout": 0.0, "optional_params": {"timeout": None}}, 0.0, id="zero-null-timeout"), + pytest.param({"timeout": 2.0, "optional_params": {"timeout": 3.0}}, 3.0, id="nested-overrides-top-level"), + pytest.param({"timeout": 2.0, "optional_params": {"timeout": 10.0}}, 10.0, id="explicit-nested-ten-seconds"), + ], +) +async def test_configured_timeout_reaches_mls( + settings: dict[str, object], expected_timeout: float | None, respx_mock: respx.MockRouter +) -> None: + guardrail: Final = _configured_guardrail(settings) + route: Final = _serve(respx_mock, _mls_body()) + + await _apply(guardrail, {"texts": ["hello"]}) + + expected: Final = RealmLabsGuardrailOptionalParams().timeout if expected_timeout is None else expected_timeout + extensions: Final = _JSON_OBJECT.validate_python(route.calls.last.request.extensions) + timeouts: Final = _JSON_OBJECT.validate_python(extensions["timeout"]) + assert timeouts["read"] == expected + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mismatch_fields", [{}, {"role_mismatch": None}], ids=["omitted", "null"]) +async def test_missing_role_mismatch_still_enforces_request_hazard( + mismatch_fields: dict[str, object], respx_mock: respx.MockRouter +) -> None: + _serve(respx_mock, {"results": [{"probe": "hazard_prompt", "prob": 0.99, **mismatch_fields}], "pii_spans": []}) + + with pytest.raises(GuardrailRaisedException) as exc: + await _apply(_guardrail(), {"texts": ["a hazardous request"]}) + + assert exc.value.blocked_content is True + assert "hazard_prompt" in exc.value.message, exc.value.message + + +@pytest.mark.asyncio +async def test_response_text_is_returned_without_calling_mls(respx_mock: respx.MockRouter) -> None: + route: Final = _serve(respx_mock, _mls_body()) + inputs: Final[GenericGuardrailAPIInputs] = {"texts": ["hello"]} + + assert await _apply(_guardrail(), inputs, input_type="response") is inputs + assert route.call_count == 0 diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index cc0b081b670..b6b7beaf86f 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -36234,6 +36234,11 @@ export interface components { * @default false */ disable_exception_on_block: boolean | null; + /** + * Enable Thinking + * @description Whether MLS should render the chat template in thinking mode. Defaults to False. + */ + enable_thinking?: boolean | null; /** * End Session After N Fails * @description For /v1/realtime sessions: automatically close the session after this many guardrail violations. @@ -36306,6 +36311,11 @@ export interface components { * @description Enable hallucination detection to detect factual inaccuracies. */ hallucinations_check?: boolean | null; + /** + * Hazard Threshold + * @description Block the request when the hazard_prompt probe scores strictly above this value. Defaults to 0.703, the threshold MLS reports for that probe. Note the probe also responds to instruction-style phrasing such as "repeat this back verbatim", so raise this if benign traffic is being blocked. + */ + hazard_threshold?: number | null; /** * Include Evidence * @description Include detailed evidence payloads in responses (sets `plr_evidence` header). @@ -36558,6 +36568,11 @@ export interface components { presidio_score_thresholds?: { [key: string]: number; } | null; + /** + * Probes + * @description Which classifier probes to run: a list of probe names, or "all". Defaults to ["hazard_prompt"] - the only probe whose score this guardrail enforces. An unknown probe name makes MLS return 404. + */ + probes?: string[] | string | null; /** * Project Id * @description Project ID for the Lakera AI project