Merge 7cbf47cf47 into ec3f8183c3

litellm/proxy/guardrails/guardrail_hooks/reco/__init__.py

@@ -0,0 +1,35 @@
+from typing import TYPE_CHECKING, Final
+
+from litellm.types.guardrails import SupportedGuardrailIntegrations
+
+from .reco import RecoGuardrail
+
+if TYPE_CHECKING:
+    from litellm.types.guardrails import Guardrail, LitellmParams
+
+
+def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"):
+    import litellm
+
+    optional_params: Final = getattr(litellm_params, "optional_params", None)
+
+    _reco_callback: Final = RecoGuardrail(
+        guardrail_name=guardrail.get("guardrail_name", ""),
+        reco_tenant_id=getattr(optional_params, "reco_tenant_id", None),
+        api_base=getattr(optional_params, "api_base", None),
+        event_hook=litellm_params.mode,
+        default_on=litellm_params.default_on,
+    )
+    litellm.logging_callback_manager.add_litellm_callback(_reco_callback)
+
+    return _reco_callback
+
+
+guardrail_initializer_registry: Final = {  # mutable-ok: guardrail auto-discovery requires an actual dict instance
+    SupportedGuardrailIntegrations.RECO.value: initialize_guardrail,
+}
+
+
+guardrail_class_registry: Final = {  # mutable-ok: guardrail auto-discovery requires an actual dict instance
+    SupportedGuardrailIntegrations.RECO.value: RecoGuardrail,
+}

litellm/proxy/guardrails/guardrail_hooks/reco/reco.py

@@ -0,0 +1,72 @@
+from collections.abc import Mapping
+from typing import TYPE_CHECKING, Literal, Optional
+
+from litellm.exceptions import GuardrailRaisedException
+from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api.generic_guardrail_api import (
+    GenericGuardrailAPI,
+)
+from litellm.types.guardrails import GuardrailEventHooks
+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
+
+
+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,
+        **kwargs,  # noqa: ANN003  # forwards guardrail_name/event_hook/default_on to CustomGuardrail, whose types are narrower than LitellmParams' own field types
+    ) -> 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 = validate_reco_tenant_id(reco_tenant_id)
+        base_headers = headers or {}  # mutable-ok: normalizes headers to a dict
+        merged_headers = {**base_headers, "X-Reco-Tenant-Id": validated_tenant_id}  # mutable-ok: required by base class
+
+        super().__init__(
+            headers=merged_headers,
+            api_base=api_base,
+            unreachable_fallback="fail_open",
+            fail_on_error=False,
+            **kwargs,
+        )
+
+    async def apply_guardrail(
+        self,
+        inputs: GenericGuardrailAPIInputs,
+        request_data: dict,  # 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(inputs, request_data, input_type, 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,
+            ) from e
+
+    @staticmethod
+    def get_config_model() -> type["GuardrailConfigModel"] | 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]  # mutable-ok: fixed by base class

litellm/types/guardrails.py

@@ -136,6 +136,7 @@ class SupportedGuardrailIntegrations(Enum):
     HEADROOM = "headroom"
     COMPRESR = "compresr"
     STRAIKER = "straiker"
+    RECO = "reco"
 
 
 class Role(Enum):

litellm/types/proxy/guardrails/guardrail_hooks/reco.py

@@ -0,0 +1,37 @@
+import uuid
+
+from pydantic import BaseModel, 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"""
+
+    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.",
+    )
+
+    @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"

tests/test_litellm/proxy/guardrails/guardrail_hooks/test_reco.py

@@ -0,0 +1,134 @@
+"""
+Tests for the Reco guardrail integration.
+"""
+
+from unittest.mock import MagicMock, patch
+
+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,
+)
+
+
+@pytest.fixture
+def reco_guardrail():
+    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):
+        assert reco_guardrail.headers == {"X-Reco-Tenant-Id": "11111111-1111-1111-1111-111111111111"}
+
+    def test_tenant_header_merged_with_existing_headers(self):
+        guardrail = 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_unreachable_fallback_and_fail_on_error_are_hardcoded(self, reco_guardrail):
+        assert reco_guardrail.unreachable_fallback == "fail_open"
+        assert reco_guardrail.fail_on_error is False
+
+    def test_missing_reco_tenant_id_raises(self):
+        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):
+        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):
+        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):
+        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):
+        assert RecoGuardrail.get_supported_event_hooks() == [GuardrailEventHooks.pre_call]
+
+    def test_get_config_model_returns_reco_config_model(self):
+        assert RecoGuardrail.get_config_model() is RecoConfigModel
+
+    def test_config_model_ui_friendly_name(self):
+        assert RecoConfigModel.ui_friendly_name() == "Reco"
+
+
+class TestRecoGuardrailInitializer:
+    def test_initialize_guardrail_wires_optional_params(self):
+        litellm_params = LitellmParams(
+            guardrail="reco",
+            mode="pre_call",
+            optional_params={
+                "reco_tenant_id": "22222222-2222-2222-2222-222222222222",
+                "api_base": "https://edge2.eu.reco.ai",
+            },
+        )
+
+        guardrail = 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"
+
+
+class TestRecoGuardrailBlocking:
+    @pytest.mark.asyncio
+    async def test_action_blocked_raises_exception_with_configured_guardrail_name(self, reco_guardrail):
+        mock_response = 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"
+
+        _, 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):
+        mock_response = 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 = await reco_guardrail.apply_guardrail(
+                inputs={"texts": ["hello"]},
+                request_data={"metadata": {}},
+                input_type="request",
+            )
+
+        assert result["texts"] == ["hello"]