mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-27 06:13:07 +00:00
feat(guardrails): add native Reco guardrail provider (pre_call), fail closed by default (#48243)
* feat(guardrails): add native Reco guardrail provider (pre_call) Registers guardrail: reco as a first-class vendor integration, built on the existing generic guardrail API wire contract rather than reimplementing the HTTP client and payload mapping. Maps reco_tenant_id to the X-Reco-Tenant-Id header and hardcodes fail_open/fail_on_error since those should not be customer-configurable. v1 only supports mode: pre_call; post_call and streaming redaction already work against the existing generic pilot and are deferred. Also overrides apply_guardrail to re-raise blocked-request errors with the guardrail's own configured name: the base class hardcodes a module-level name in that exception, which would otherwise mislabel every blocked Reco request in logs. * feat(guardrails): configurable Reco fail modes Expose unreachable_fallback and fail_on_error on the Reco guardrail and keep fail-open as the default. Values under optional_params win. Top-level litellm_params values apply only when set explicitly, since their defaults are fail-closed and would otherwise flip the default Also type RecoGuardrail's forwarded kwargs with an Unpack'd TypedDict instead of silencing ANN003 Signed-off-by: MickeyShnaiderman-RecoLabs <mickeys@reco.ai> Co-authored-by: Cursor <cursoragent@cursor.com> * chore(guardrails): align Reco with main's CI gates Upstream moved unit tests from tests/test_litellm to tests/unit, so test_reco.py was invoked by no CI job. Move it next to the other guardrail tests Replace RecoGuardrail's forwarded kwargs with explicit keyword params, read optional_params through a pydantic adapter instead of untyped dict/getattr access, freeze RecoOptionalParams, and drop mutable-ok markers that suppress nothing, to stay within main's type-discipline and basedpyright budgets Signed-off-by: MickeyShnaiderman-RecoLabs <mickeys@reco.ai> Co-authored-by: Cursor <cursoragent@cursor.com> * feat(guardrails): fail closed by default in the Reco guardrail Matches GenericGuardrailAPI and every other guardrail here: a Reco error or outage now blocks the request unless the admin sets fail_on_error: false. Tests cover the default, explicit fail-open, and top-level vs optional_params precedence Co-authored-by: RecoLabs-YaronBratspiess <208601538+RecoLabs-YaronBratspiess@users.noreply.github.com> Co-authored-by: MickeyShnaiderman-RecoLabs <77006830+MickeyShnaiderman-RecoLabs@users.noreply.github.com> Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(guardrails): honor explicit fail-open on Reco when optional_params is a model model_dump(exclude_unset=True) so pydantic defaults in optional_params no longer override an explicit top-level fail_on_error. Also drop the duplicate log_guardrail_information decorator, call the parent by keyword, freeze RecoConfigModel, and type the tests Co-authored-by: RecoLabs-YaronBratspiess <208601538+RecoLabs-YaronBratspiess@users.noreply.github.com> Co-authored-by: MickeyShnaiderman-RecoLabs <77006830+MickeyShnaiderman-RecoLabs@users.noreply.github.com> Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(guardrails): keep Reco config model assignable and the hook decorator the gate requires frozen=True on RecoConfigModel froze the shared LitellmParams mixin, which the guardrail registry assigns resolved secrets on, so every guardrail failed to load. Dropped it and added a test that LitellmParams stays assignable. Restored log_guardrail_information on the Reco override since tests/code_coverage_tests/check_guardrail_apply_decorator.py requires it on every hook. Test params typed as Mapping for LIT001 Co-authored-by: RecoLabs-YaronBratspiess <208601538+RecoLabs-YaronBratspiess@users.noreply.github.com> Co-authored-by: MickeyShnaiderman-RecoLabs <77006830+MickeyShnaiderman-RecoLabs@users.noreply.github.com> Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style(guardrails): format the Reco test file Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Signed-off-by: MickeyShnaiderman-RecoLabs <mickeys@reco.ai> Co-authored-by: Yaron Bratspiess <yaronb@reco.ai> Co-authored-by: MickeyShnaiderman-RecoLabs <mickeys@reco.ai> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: joshua <joshua@berri.ai> Co-authored-by: MickeyShnaiderman-RecoLabs <77006830+MickeyShnaiderman-RecoLabs@users.noreply.github.com> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
172aa72ad2
commit
1418c1990f
5 changed files with 494 additions and 0 deletions
72
litellm/proxy/guardrails/guardrail_hooks/reco/__init__.py
Normal file
72
litellm/proxy/guardrails/guardrail_hooks/reco/__init__.py
Normal file
|
|
@ -0,0 +1,72 @@
|
|||
from collections.abc import Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from pydantic import BaseModel, TypeAdapter
|
||||
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.guardrails import SupportedGuardrailIntegrations
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.reco import RecoOptionalParams
|
||||
|
||||
from .reco import RecoGuardrail
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.types.guardrails import Guardrail, LitellmParams
|
||||
|
||||
_CONFIG_MAPPING: Final = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
def _optional_params_mapping(optional_params: object) -> Mapping[str, object]:
|
||||
"""optional_params arrives as a pydantic object from YAML, or as a plain dict when reconciled from the DB."""
|
||||
if optional_params is None:
|
||||
return MappingProxyType({})
|
||||
raw: Final = (
|
||||
optional_params.model_dump(exclude_unset=True) if isinstance(optional_params, BaseModel) else optional_params
|
||||
)
|
||||
return _CONFIG_MAPPING.validate_python(raw)
|
||||
|
||||
|
||||
def _first_configured(field_name: str, sources: tuple[Mapping[str, object], ...]) -> object:
|
||||
value: Final = next((source[field_name] for source in sources if source.get(field_name) is not None), None)
|
||||
if isinstance(value, str) and value.startswith("os.environ/"):
|
||||
return get_secret_str(value)
|
||||
return value
|
||||
|
||||
|
||||
def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail") -> RecoGuardrail:
|
||||
import litellm
|
||||
|
||||
# Top-level fields only count when set explicitly: their defaults are fail-closed and would flip Reco's.
|
||||
explicit_top_level: Final = _CONFIG_MAPPING.validate_python(
|
||||
litellm_params.model_dump(include=litellm_params.model_fields_set & RecoOptionalParams.model_fields.keys())
|
||||
)
|
||||
sources: Final = (_optional_params_mapping(litellm_params.optional_params), explicit_top_level)
|
||||
configured: Final = (
|
||||
(field_name, _first_configured(field_name, sources)) for field_name in RecoOptionalParams.model_fields
|
||||
)
|
||||
reco_params: Final = RecoOptionalParams.model_validate(
|
||||
{field_name: value for field_name, value in configured if value is not None}
|
||||
)
|
||||
|
||||
_reco_callback: Final = RecoGuardrail(
|
||||
guardrail_name=guardrail["guardrail_name"],
|
||||
reco_tenant_id=reco_params.reco_tenant_id,
|
||||
api_base=reco_params.api_base,
|
||||
unreachable_fallback=reco_params.unreachable_fallback,
|
||||
fail_on_error=reco_params.fail_on_error,
|
||||
event_hook=litellm_params.mode,
|
||||
default_on=bool(litellm_params.default_on),
|
||||
)
|
||||
litellm.logging_callback_manager.add_litellm_callback(_reco_callback)
|
||||
|
||||
return _reco_callback
|
||||
|
||||
|
||||
guardrail_initializer_registry: Final = {
|
||||
SupportedGuardrailIntegrations.RECO.value: initialize_guardrail,
|
||||
}
|
||||
|
||||
|
||||
guardrail_class_registry: Final = {
|
||||
SupportedGuardrailIntegrations.RECO.value: RecoGuardrail,
|
||||
}
|
||||
84
litellm/proxy/guardrails/guardrail_hooks/reco/reco.py
Normal file
84
litellm/proxy/guardrails/guardrail_hooks/reco/reco.py
Normal file
|
|
@ -0,0 +1,84 @@
|
|||
from collections.abc import Mapping, Sequence
|
||||
from typing import TYPE_CHECKING, Final, Literal, Optional
|
||||
|
||||
from litellm.exceptions import GuardrailRaisedException
|
||||
from litellm.integrations.custom_guardrail import (
|
||||
log_guardrail_information, # pyright: ignore[reportUnknownVariableType] # decorator is untyped
|
||||
)
|
||||
from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api.generic_guardrail_api import (
|
||||
GenericGuardrailAPI,
|
||||
)
|
||||
from litellm.types.guardrails import GuardrailEventHooks, Mode
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.reco import validate_reco_tenant_id
|
||||
from litellm.types.utils import GenericGuardrailAPIInputs
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.reco import RecoOptionalParams
|
||||
|
||||
|
||||
class RecoGuardrail(GenericGuardrailAPI):
|
||||
"""Reco guardrail integration for LiteLLM, built on the Generic Guardrail API wire contract."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
reco_tenant_id: str | None,
|
||||
api_base: str | None,
|
||||
headers: Mapping[str, str] | None = None,
|
||||
unreachable_fallback: Literal["fail_closed", "fail_open"] = "fail_closed",
|
||||
fail_on_error: bool = True,
|
||||
guardrail_name: str | None = None,
|
||||
event_hook: str | Sequence[str] | Mode | None = None,
|
||||
default_on: bool = False,
|
||||
) -> None:
|
||||
if not reco_tenant_id:
|
||||
raise ValueError("reco_tenant_id is required for the Reco guardrail")
|
||||
if not api_base:
|
||||
raise ValueError("api_base is required for the Reco guardrail")
|
||||
|
||||
validated_tenant_id: Final = validate_reco_tenant_id(reco_tenant_id)
|
||||
|
||||
super().__init__( # pyright: ignore[reportUnknownMemberType] # parent takes untyped extra_headers and **kwargs
|
||||
headers={**(headers or {}), "X-Reco-Tenant-Id": validated_tenant_id},
|
||||
api_base=api_base,
|
||||
unreachable_fallback=unreachable_fallback,
|
||||
fail_on_error=fail_on_error,
|
||||
guardrail_name=guardrail_name,
|
||||
event_hook=event_hook,
|
||||
default_on=default_on,
|
||||
)
|
||||
|
||||
@log_guardrail_information
|
||||
async def apply_guardrail(
|
||||
self,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
request_data: dict[str, object], # mutable-ok: must match GenericGuardrailAPI.apply_guardrail's own param type
|
||||
input_type: Literal["request", "response"],
|
||||
logging_obj: Optional["LiteLLMLoggingObj"] = None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
# GenericGuardrailAPI raises with a hardcoded module-level guardrail_name on BLOCKED,
|
||||
# not the guardrail's own configured name. Re-raise with the right one instead of
|
||||
# patching the shared base class.
|
||||
try:
|
||||
return await super().apply_guardrail( # pyright: ignore[reportUnknownMemberType] # parent annotates request_data as a bare dict
|
||||
inputs=inputs, request_data=request_data, input_type=input_type, logging_obj=logging_obj
|
||||
)
|
||||
except GuardrailRaisedException as e:
|
||||
raise GuardrailRaisedException(
|
||||
guardrail_name=self.guardrail_name,
|
||||
message=e.message,
|
||||
should_wrap_with_default_message=False,
|
||||
status_code=e.status_code,
|
||||
blocked_content=e.blocked_content,
|
||||
) from e
|
||||
|
||||
@staticmethod
|
||||
def get_config_model() -> type["GuardrailConfigModel[RecoOptionalParams]"] | None:
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.reco import RecoConfigModel
|
||||
|
||||
return RecoConfigModel
|
||||
|
||||
@classmethod
|
||||
def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]: # mutable-ok: fixed by base class
|
||||
return [GuardrailEventHooks.pre_call]
|
||||
|
|
@ -59,6 +59,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.reco import (
|
||||
RecoConfigModel,
|
||||
)
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.repelloai import (
|
||||
RepelloAIGuardrailConfigModel,
|
||||
)
|
||||
|
|
@ -153,6 +156,7 @@ class SupportedGuardrailIntegrations(Enum):
|
|||
TYPESAFE = "typesafe"
|
||||
STRAIKER = "straiker"
|
||||
ALICE = "alice"
|
||||
RECO = "reco"
|
||||
AGENT_365 = "agent_365"
|
||||
LLM_SHIELD_PROXY = "llm_shield_proxy"
|
||||
CONDUCT = "conduct"
|
||||
|
|
@ -1398,6 +1402,7 @@ class LitellmParams( # pyright: ignore[reportIncompatibleVariableOverride] # o
|
|||
QostodianNexusConfigModel,
|
||||
VigilGuardGuardrailConfigModel,
|
||||
SingulrGuardrailConfigModel,
|
||||
RecoConfigModel,
|
||||
Agent365GuardrailConfigModel,
|
||||
WingbackGuardrailConfigModel,
|
||||
):
|
||||
|
|
|
|||
48
litellm/types/proxy/guardrails/guardrail_hooks/reco.py
Normal file
48
litellm/types/proxy/guardrails/guardrail_hooks/reco.py
Normal file
|
|
@ -0,0 +1,48 @@
|
|||
import uuid
|
||||
from typing import Literal
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, field_validator
|
||||
|
||||
from .base import GuardrailConfigModel
|
||||
|
||||
|
||||
def validate_reco_tenant_id(value: str) -> str:
|
||||
try:
|
||||
uuid.UUID(value)
|
||||
except ValueError as e:
|
||||
raise ValueError(f"reco_tenant_id must be a valid UUID, got {value!r}") from e
|
||||
return value
|
||||
|
||||
|
||||
class RecoOptionalParams(BaseModel):
|
||||
"""Configuration parameters for the Reco guardrail"""
|
||||
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
reco_tenant_id: str = Field(
|
||||
description="Tenant identifier for the Reco account, as a UUID. Sent as the X-Reco-Tenant-Id header on every guardrail request.",
|
||||
)
|
||||
api_base: str = Field(
|
||||
description="Base URL for the Reco guardrail API. Reco's endpoint is per-region and per-silo, so this must be set explicitly rather than defaulted.",
|
||||
)
|
||||
unreachable_fallback: Literal["fail_closed", "fail_open"] = Field(
|
||||
default="fail_closed",
|
||||
description="Behavior when Reco is unreachable. 'fail_closed' (default) blocks the request; 'fail_open' lets it through uninspected. Only applies while fail_on_error is true.",
|
||||
)
|
||||
fail_on_error: bool = Field(
|
||||
default=True,
|
||||
description="Behavior on any Reco error, including unreachability. True (default) blocks the request; False lets it through uninspected.",
|
||||
)
|
||||
|
||||
@field_validator("reco_tenant_id")
|
||||
@classmethod
|
||||
def _check_reco_tenant_id(cls, v: str) -> str:
|
||||
return validate_reco_tenant_id(v)
|
||||
|
||||
|
||||
class RecoConfigModel(GuardrailConfigModel[RecoOptionalParams]):
|
||||
"""Configuration parameters for the Reco guardrail"""
|
||||
|
||||
@staticmethod
|
||||
def ui_friendly_name() -> str:
|
||||
return "Reco"
|
||||
285
tests/unit/proxy/guardrails/guardrail_hooks/test_reco.py
Normal file
285
tests/unit/proxy/guardrails/guardrail_hooks/test_reco.py
Normal file
|
|
@ -0,0 +1,285 @@
|
|||
"""
|
||||
Tests for the Reco guardrail integration.
|
||||
"""
|
||||
|
||||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from pydantic import ValidationError
|
||||
|
||||
from litellm.exceptions import GuardrailRaisedException
|
||||
from litellm.proxy.guardrails.guardrail_hooks.reco import (
|
||||
RecoGuardrail,
|
||||
initialize_guardrail,
|
||||
)
|
||||
from litellm.types.guardrails import GuardrailEventHooks, LitellmParams
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.reco import (
|
||||
RecoConfigModel,
|
||||
RecoOptionalParams,
|
||||
)
|
||||
from litellm.types.utils import GenericGuardrailAPIInputs
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def reco_guardrail() -> RecoGuardrail:
|
||||
return RecoGuardrail(
|
||||
reco_tenant_id="11111111-1111-1111-1111-111111111111",
|
||||
api_base="https://edge1.us.reco.ai",
|
||||
guardrail_name="my-reco-guardrail",
|
||||
event_hook="pre_call",
|
||||
default_on=True,
|
||||
)
|
||||
|
||||
|
||||
class TestRecoGuardrailConfiguration:
|
||||
def test_tenant_id_forwarded_as_header(self, reco_guardrail: RecoGuardrail) -> None:
|
||||
assert reco_guardrail.headers == {"X-Reco-Tenant-Id": "11111111-1111-1111-1111-111111111111"}
|
||||
|
||||
def test_tenant_header_merged_with_existing_headers(self) -> None:
|
||||
guardrail: Final = RecoGuardrail(
|
||||
reco_tenant_id="11111111-1111-1111-1111-111111111111",
|
||||
api_base="https://edge1.us.reco.ai",
|
||||
headers={"X-Custom": "value"},
|
||||
)
|
||||
assert guardrail.headers == {
|
||||
"X-Custom": "value",
|
||||
"X-Reco-Tenant-Id": "11111111-1111-1111-1111-111111111111",
|
||||
}
|
||||
|
||||
def test_missing_reco_tenant_id_raises(self) -> None:
|
||||
with pytest.raises(ValueError, match="reco_tenant_id"):
|
||||
RecoGuardrail(reco_tenant_id=None, api_base="https://edge1.us.reco.ai")
|
||||
|
||||
def test_missing_api_base_raises(self) -> None:
|
||||
with pytest.raises(ValueError, match="api_base"):
|
||||
RecoGuardrail(reco_tenant_id="11111111-1111-1111-1111-111111111111", api_base=None)
|
||||
|
||||
def test_non_uuid_reco_tenant_id_raises(self) -> None:
|
||||
with pytest.raises(ValueError, match="UUID"):
|
||||
RecoGuardrail(reco_tenant_id="tenant-123", api_base="https://edge1.us.reco.ai")
|
||||
|
||||
def test_optional_params_rejects_non_uuid_tenant_id(self) -> None:
|
||||
with pytest.raises(ValidationError, match="UUID"):
|
||||
RecoOptionalParams(reco_tenant_id="tenant-123", api_base="https://edge1.us.reco.ai")
|
||||
|
||||
def test_get_supported_event_hooks_returns_pre_call_only(self) -> None:
|
||||
assert RecoGuardrail.get_supported_event_hooks() == [GuardrailEventHooks.pre_call]
|
||||
|
||||
def test_get_config_model_returns_reco_config_model(self) -> None:
|
||||
assert RecoGuardrail.get_config_model() is RecoConfigModel
|
||||
|
||||
def test_config_model_ui_friendly_name(self) -> None:
|
||||
assert RecoConfigModel.ui_friendly_name() == "Reco"
|
||||
|
||||
def test_litellm_params_stay_assignable_for_registry_secret_resolution(self) -> None:
|
||||
litellm_params: Final = LitellmParams(guardrail="reco", mode="pre_call")
|
||||
|
||||
litellm_params.api_base = "https://edge1.us.reco.ai"
|
||||
|
||||
assert litellm_params.api_base == "https://edge1.us.reco.ai"
|
||||
|
||||
|
||||
class TestRecoGuardrailInitializer:
|
||||
def test_initialize_guardrail_wires_optional_params(self) -> None:
|
||||
litellm_params: Final = LitellmParams(
|
||||
guardrail="reco",
|
||||
mode="pre_call",
|
||||
optional_params={
|
||||
"reco_tenant_id": "22222222-2222-2222-2222-222222222222",
|
||||
"api_base": "https://edge2.eu.reco.ai",
|
||||
},
|
||||
)
|
||||
|
||||
guardrail: Final = initialize_guardrail(litellm_params, {"guardrail_name": "prod-reco"})
|
||||
|
||||
assert isinstance(guardrail, RecoGuardrail)
|
||||
assert guardrail.guardrail_name == "prod-reco"
|
||||
assert guardrail.headers == {"X-Reco-Tenant-Id": "22222222-2222-2222-2222-222222222222"}
|
||||
assert guardrail.api_base == "https://edge2.eu.reco.ai/beta/litellm_basic_guardrail_api"
|
||||
|
||||
def test_initialize_guardrail_reads_plain_dict_optional_params(self) -> None:
|
||||
"""Guardrails reconciled from the DB can carry `optional_params` as a
|
||||
plain dict rather than a pydantic object; attribute-only access would
|
||||
silently read None for every field."""
|
||||
litellm_params: Final = LitellmParams.model_construct(
|
||||
guardrail="reco",
|
||||
optional_params={
|
||||
"reco_tenant_id": "33333333-3333-3333-3333-333333333333",
|
||||
"api_base": "https://edge3.eu.reco.ai",
|
||||
},
|
||||
mode="pre_call",
|
||||
default_on=True,
|
||||
)
|
||||
|
||||
guardrail: Final = initialize_guardrail(litellm_params, {"guardrail_name": "prod-reco"})
|
||||
|
||||
assert guardrail.headers == {"X-Reco-Tenant-Id": "33333333-3333-3333-3333-333333333333"}
|
||||
assert guardrail.api_base == "https://edge3.eu.reco.ai/beta/litellm_basic_guardrail_api"
|
||||
|
||||
def test_initialize_guardrail_resolves_env_var_reference(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("RECO_TEST_TENANT_ID", "44444444-4444-4444-4444-444444444444")
|
||||
litellm_params: Final = LitellmParams(
|
||||
guardrail="reco",
|
||||
optional_params={
|
||||
"reco_tenant_id": "os.environ/RECO_TEST_TENANT_ID",
|
||||
"api_base": "https://edge4.eu.reco.ai",
|
||||
},
|
||||
mode="pre_call",
|
||||
default_on=True,
|
||||
)
|
||||
|
||||
guardrail: Final = initialize_guardrail(litellm_params, {"guardrail_name": "prod-reco"})
|
||||
|
||||
assert guardrail.headers == {"X-Reco-Tenant-Id": "44444444-4444-4444-4444-444444444444"}
|
||||
|
||||
|
||||
def _litellm_params(top_level: Mapping[str, object], nested: Mapping[str, object]) -> LitellmParams:
|
||||
return LitellmParams(
|
||||
guardrail="reco",
|
||||
mode="pre_call",
|
||||
optional_params={
|
||||
"reco_tenant_id": "22222222-2222-2222-2222-222222222222",
|
||||
"api_base": "https://edge2.eu.reco.ai",
|
||||
**nested,
|
||||
},
|
||||
**top_level,
|
||||
)
|
||||
|
||||
|
||||
async def _apply_with_reco_unreachable(guardrail: RecoGuardrail) -> GenericGuardrailAPIInputs:
|
||||
with patch.object(guardrail.async_handler, "post", side_effect=httpx.ConnectError("connection refused")):
|
||||
return await guardrail.apply_guardrail(
|
||||
inputs={"texts": ["hello"]},
|
||||
request_data={"metadata": {}},
|
||||
input_type="request",
|
||||
)
|
||||
|
||||
|
||||
class TestRecoGuardrailFailureModes:
|
||||
@pytest.mark.asyncio
|
||||
async def test_unreachable_reco_blocks_request_by_default(self, reco_guardrail: RecoGuardrail) -> None:
|
||||
with pytest.raises(Exception, match="connection refused"):
|
||||
await _apply_with_reco_unreachable(reco_guardrail)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fail_open_lets_request_through_when_reco_unreachable(self) -> None:
|
||||
guardrail: Final = RecoGuardrail(
|
||||
reco_tenant_id="11111111-1111-1111-1111-111111111111",
|
||||
api_base="https://edge1.us.reco.ai",
|
||||
fail_on_error=False,
|
||||
)
|
||||
|
||||
result: Final = await _apply_with_reco_unreachable(guardrail)
|
||||
|
||||
assert result["texts"] == ["hello"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fail_closed_blocks_request_when_reco_unreachable(self) -> None:
|
||||
guardrail: Final = RecoGuardrail(
|
||||
reco_tenant_id="11111111-1111-1111-1111-111111111111",
|
||||
api_base="https://edge1.us.reco.ai",
|
||||
unreachable_fallback="fail_closed",
|
||||
fail_on_error=True,
|
||||
)
|
||||
|
||||
with pytest.raises(Exception, match="connection refused"):
|
||||
await _apply_with_reco_unreachable(guardrail)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("top_level", "nested"),
|
||||
[
|
||||
pytest.param({}, {}, id="defaults"),
|
||||
pytest.param({}, {"unreachable_fallback": "fail_closed", "fail_on_error": True}, id="optional-params"),
|
||||
pytest.param({"unreachable_fallback": "fail_closed", "fail_on_error": True}, {}, id="top-level"),
|
||||
],
|
||||
)
|
||||
async def test_initialize_guardrail_fails_closed_when_configured(
|
||||
self, top_level: Mapping[str, object], nested: Mapping[str, object]
|
||||
) -> None:
|
||||
guardrail: Final = initialize_guardrail(_litellm_params(top_level, nested), {"guardrail_name": "prod-reco"})
|
||||
|
||||
with pytest.raises(Exception, match="connection refused"):
|
||||
await _apply_with_reco_unreachable(guardrail)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("top_level", "nested"),
|
||||
[
|
||||
pytest.param({}, {"fail_on_error": False}, id="optional-params"),
|
||||
pytest.param({"fail_on_error": False}, {}, id="top-level"),
|
||||
pytest.param(
|
||||
{"unreachable_fallback": "fail_closed", "fail_on_error": True},
|
||||
{"fail_on_error": False},
|
||||
id="optional-params-wins-over-top-level",
|
||||
),
|
||||
],
|
||||
)
|
||||
async def test_initialize_guardrail_fails_open_when_configured(
|
||||
self, top_level: Mapping[str, object], nested: Mapping[str, object]
|
||||
) -> None:
|
||||
guardrail: Final = initialize_guardrail(_litellm_params(top_level, nested), {"guardrail_name": "prod-reco"})
|
||||
|
||||
result: Final = await _apply_with_reco_unreachable(guardrail)
|
||||
|
||||
assert result["texts"] == ["hello"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_initialize_guardrail_keeps_top_level_fail_open_when_nested_params_are_a_model(self) -> None:
|
||||
nested: Final = RecoOptionalParams(
|
||||
reco_tenant_id="22222222-2222-2222-2222-222222222222", api_base="https://edge2.eu.reco.ai"
|
||||
)
|
||||
litellm_params: Final = LitellmParams.model_construct(
|
||||
guardrail="reco", mode="pre_call", optional_params=nested, fail_on_error=False
|
||||
)
|
||||
|
||||
guardrail: Final = initialize_guardrail(litellm_params, {"guardrail_name": "prod-reco"})
|
||||
result: Final = await _apply_with_reco_unreachable(guardrail)
|
||||
|
||||
assert result["texts"] == ["hello"]
|
||||
|
||||
|
||||
class TestRecoGuardrailBlocking:
|
||||
@pytest.mark.asyncio
|
||||
async def test_action_blocked_raises_exception_with_configured_guardrail_name(
|
||||
self, reco_guardrail: RecoGuardrail
|
||||
) -> None:
|
||||
mock_response: Final = MagicMock()
|
||||
mock_response.json.return_value = {
|
||||
"action": "BLOCKED",
|
||||
"blocked_reason": "Sensitive data detected",
|
||||
}
|
||||
mock_response.raise_for_status = MagicMock()
|
||||
|
||||
with patch.object(reco_guardrail.async_handler, "post", return_value=mock_response) as mock_post:
|
||||
with pytest.raises(GuardrailRaisedException) as exc_info:
|
||||
await reco_guardrail.apply_guardrail(
|
||||
inputs={"texts": ["some sensitive prompt"]},
|
||||
request_data={"metadata": {}},
|
||||
input_type="request",
|
||||
)
|
||||
|
||||
assert exc_info.value.guardrail_name == "my-reco-guardrail"
|
||||
assert str(exc_info.value) == "Sensitive data detected"
|
||||
assert exc_info.value.blocked_content is True
|
||||
|
||||
_, call_kwargs = mock_post.call_args
|
||||
assert call_kwargs["headers"]["X-Reco-Tenant-Id"] == "11111111-1111-1111-1111-111111111111"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_action_none_allows_content(self, reco_guardrail: RecoGuardrail) -> None:
|
||||
mock_response: Final = MagicMock()
|
||||
mock_response.json.return_value = {"action": "NONE", "texts": ["hello"]}
|
||||
mock_response.raise_for_status = MagicMock()
|
||||
|
||||
with patch.object(reco_guardrail.async_handler, "post", return_value=mock_response):
|
||||
result: Final = await reco_guardrail.apply_guardrail(
|
||||
inputs={"texts": ["hello"]},
|
||||
request_data={"metadata": {}},
|
||||
input_type="request",
|
||||
)
|
||||
|
||||
assert result["texts"] == ["hello"]
|
||||
Loading…
Add table
Reference in a new issue