diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index ae469a104f8..fafbdb1f95a 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -386,6 +386,10 @@ def _guardrails_listed_in(metadata: object) -> list[str | dict[str, DynamicGuard return metadata.get("guardrails") or [] +class GuardrailSetupError(Exception): + pass + + class CustomGuardrail(CustomLogger): mcp_block_behavior: Literal["tool_result", "raise"] | None = None # If True, during_call runs async_moderation_hook instead of the unified apply_guardrail path. diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index eaece1fe4c6..f6591ca00d5 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -18,7 +18,7 @@ from pydantic import BaseModel, TypeAdapter, ValidationError from litellm._logging import verbose_proxy_logger from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH -from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.integrations.custom_guardrail import CustomGuardrail, GuardrailSetupError from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy._types import AUDIT_ACTIONS, LitellmTableNames, LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth @@ -540,7 +540,7 @@ async def create_guardrail( verbose_proxy_logger.info( "Immediate sync: Successfully initialized guardrail '%s' (ID: %s)", guardrail_name, guardrail_id ) - except (ValueError, TypeError, CustomCodeCompilationError) as init_error: + except (ValueError, TypeError, CustomCodeCompilationError, GuardrailSetupError) as init_error: # Configuration error — roll back the DB write so the guardrail isn't orphaned if prisma_client is not None: try: diff --git a/litellm/proxy/guardrails/guardrail_hooks/collieai/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/collieai/__init__.py new file mode 100644 index 00000000000..c0ea6335741 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/collieai/__init__.py @@ -0,0 +1,55 @@ +from __future__ import annotations + +from collections.abc import Callable, Mapping +from types import MappingProxyType +from typing import TYPE_CHECKING, Final + +from litellm.integrations.custom_guardrail import GuardrailSetupError +from litellm.types.guardrails import SupportedGuardrailIntegrations +from litellm.types.proxy.guardrails.guardrail_hooks.collieai import DEFAULT_API_BASE + +from .collieai import CollieAiGuardrail + +if TYPE_CHECKING: + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.types.guardrails import Guardrail, LitellmParams + +_NO_EXTRAS: Final[Mapping[str, object]] = MappingProxyType({}) + + +def initialize_guardrail( + litellm_params: LitellmParams, + guardrail: Guardrail, + guardrail_cls: Callable[..., CustomGuardrail] = CollieAiGuardrail, +) -> CustomGuardrail: + import litellm + from litellm.secret_managers.main import get_secret_str + + extras: Final = MappingProxyType( + {key: value for key, value in (litellm_params.model_extra or _NO_EXTRAS).items() if value is not None} + ) + try: + _callback: Final = guardrail_cls( + **{ + **extras, + "api_base": litellm_params.api_base or DEFAULT_API_BASE, + "api_key": litellm_params.api_key or get_secret_str("COLLIEAI_API_KEY"), + "guardrail_name": guardrail["guardrail_name"], + "event_hook": litellm_params.mode, + "default_on": litellm_params.default_on, + "stream_scope": litellm_params.stream_scope, + } + ) + except Exception as setup_error: + raise GuardrailSetupError(str(setup_error)) from setup_error + litellm.logging_callback_manager.add_litellm_callback(_callback) + return _callback + + +guardrail_initializer_registry: Final = { + SupportedGuardrailIntegrations.COLLIEAI.value: initialize_guardrail, +} + +guardrail_class_registry: Final = { + SupportedGuardrailIntegrations.COLLIEAI.value: CollieAiGuardrail, +} diff --git a/litellm/proxy/guardrails/guardrail_hooks/collieai/collieai.py b/litellm/proxy/guardrails/guardrail_hooks/collieai/collieai.py new file mode 100644 index 00000000000..5ea1d147f91 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/collieai/collieai.py @@ -0,0 +1,94 @@ +"""CollieAi as a LiteLLM guardrail, backed by the ``collieai-litellm-guardrail`` PyPI package. + +Install: ``pip install "collieai-litellm-guardrail>=1.0.0"`` +Source: https://github.com/CollieAi/collieai-litellm-guardrail +""" + +from __future__ import annotations + +import importlib +from collections.abc import Callable +from typing import Final + +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.proxy.guardrails.guardrail_hooks.collieai import CollieAiGuardrailConfigModel + +PACKAGE_MODULE: Final = "collieai_guardrail.guardrail" +MISSING_PACKAGE_MESSAGE: Final = ( + "collieai-litellm-guardrail>=1.0.0 (Python 3.11 or newer) is required for the CollieAi guardrail. " + 'Install it with: pip install "collieai-litellm-guardrail>=1.0.0". ' + "The LiteLLM Docker image does not include it, " + "see https://docs.collieai.io/integrations/litellm to add it to an image." +) + + +def has_readiness_guard(package: object) -> bool: + return getattr(package, "require_guardrail", None) is not None + + +def package_guardrail(package: object) -> type[CustomGuardrail] | None: + guardrail: Final[object] = getattr(package, "CollieAIGuardrail", None) + if not has_readiness_guard(package) or not isinstance(guardrail, type): + return None + return guardrail if issubclass(guardrail, CustomGuardrail) else None + + +def unavailable_package(message: str, cause: BaseException | None) -> type[CustomGuardrail]: + class UnavailablePackage(CustomGuardrail): + def __init__(self, **kwargs: object) -> None: # kwargs-ok: mirrors the package constructor, only raises + raise ImportError(message) from cause + + return UnavailablePackage + + +def missing_package(import_error: ImportError | None) -> type[CustomGuardrail]: + return unavailable_package(MISSING_PACKAGE_MESSAGE, import_error) + + +def broken_package(load_error: Exception) -> type[CustomGuardrail]: + return unavailable_package( + "collieai-litellm-guardrail is installed but failed to load with this LiteLLM version " + f"({type(load_error).__name__}: {load_error}). " + 'Upgrade it with: pip install -U "collieai-litellm-guardrail>=1.0.0"', + load_error, + ) + + +def is_the_package_missing(import_error: ImportError) -> bool: + package_root: Final = PACKAGE_MODULE.partition(".")[0] + return isinstance(import_error, ModuleNotFoundError) and (import_error.name or "").partition(".")[0] == package_root + + +def load_base(import_module: Callable[[str], object] = importlib.import_module) -> type[CustomGuardrail]: + try: + package: Final = import_module(PACKAGE_MODULE) + except ImportError as import_error: + return missing_package(import_error) if is_the_package_missing(import_error) else broken_package(import_error) + except Exception as load_error: # noqa: BLE001 # a broken third-party install must not break the guardrail routes + return broken_package(load_error) + return package_guardrail(package) or missing_package(None) + + +_Base: Final = load_base() + + +class CollieAiGuardrail(_Base): + @staticmethod + def get_config_model() -> type[CollieAiGuardrailConfigModel]: + return CollieAiGuardrailConfigModel + + @classmethod + def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]: # mutable-ok: CustomGuardrail's signature + return [GuardrailEventHooks.pre_call, GuardrailEventHooks.post_call] + + +__all__ = ( + "MISSING_PACKAGE_MESSAGE", + "PACKAGE_MODULE", + "CollieAiGuardrail", + "has_readiness_guard", + "load_base", + "missing_package", + "package_guardrail", +) diff --git a/litellm/proxy/guardrails/guardrail_registry.py b/litellm/proxy/guardrails/guardrail_registry.py index cafd79626ef..dd4d8f49010 100644 --- a/litellm/proxy/guardrails/guardrail_registry.py +++ b/litellm/proxy/guardrails/guardrail_registry.py @@ -685,6 +685,24 @@ def _configure_callback_scoping( _ParamsT = TypeVar("_ParamsT", bound=BaseModel) +SECRET_REFERENCE_PARAMS: Final = ("api_key", "api_base") + + +def resolve_secret_reference(value: str | None) -> str | None: + return str(get_secret(value)) if value and value.startswith("os.environ/") else value + + +def _with_resolved_secret_references(params: Mapping[str, object]) -> Mapping[str, object]: + return { + **params, + **{ + key: resolve_secret_reference(value) + for key in SECRET_REFERENCE_PARAMS + if isinstance(value := params.get(key), str) + }, + } + + def parse_tolerant_litellm_params( litellm_params_data: Mapping[str, object], guardrail_name: str, @@ -788,11 +806,8 @@ class InMemoryGuardrailHandler: ) litellm_params.category_thresholds = lakera_category_thresholds - if litellm_params.api_key and litellm_params.api_key.startswith("os.environ/"): - litellm_params.api_key = str(get_secret(litellm_params.api_key)) - - if litellm_params.api_base and litellm_params.api_base.startswith("os.environ/"): - litellm_params.api_base = str(get_secret(litellm_params.api_base)) + litellm_params.api_key = resolve_secret_reference(litellm_params.api_key) + litellm_params.api_base = resolve_secret_reference(litellm_params.api_base) guardrail_type: Final = litellm_params.guardrail @@ -1068,7 +1083,9 @@ class InMemoryGuardrailHandler: return params.model_dump() if isinstance(params, dict): try: - return parse_tolerant_litellm_params(with_tolerated_stream_scope(params), guardrail_name).model_dump() + return _with_resolved_secret_references( + parse_tolerant_litellm_params(with_tolerated_stream_scope(params), guardrail_name).model_dump() + ) except ValidationError as e: verbose_proxy_logger.warning( "Could not normalize guardrail litellm_params for comparison; treating the guardrail as changed. Error: %s", diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index e311f2ee5f9..8765db30ed5 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -163,6 +163,7 @@ class SupportedGuardrailIntegrations(Enum): AIRIA = "airia" WINGBACK = "wingback" DECISION_MODEL = "decision_model" + COLLIEAI = "collieai" class Role(Enum): diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/collieai.py b/litellm/types/proxy/guardrails/guardrail_hooks/collieai.py new file mode 100644 index 00000000000..784bbfc9d0d --- /dev/null +++ b/litellm/types/proxy/guardrails/guardrail_hooks/collieai.py @@ -0,0 +1,47 @@ +from __future__ import annotations + +from typing import Final + +from pydantic import ConfigDict, Field + +from litellm.types.llms.base import LiteLLMBaseModel + +from .base import GuardrailConfigModel + +DEFAULT_API_BASE: Final = "https://app.collieai.io/v1/integrations/litellm" + + +class CollieAiGuardrailConfigModelOptionalParams(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + deadline_ms: int | None = Field( + default=None, + gt=0, + description=( + "Each phase's whole deadline in milliseconds: connection, send, reply and parse. " + "Leave empty for the package default." + ), + ) + forward_conversation_id: bool | None = Field( + default=None, + description=( + "Forward the caller's x-collieai-conversation-id header, and only that header. Leave empty to keep it off." + ), + ) + + +class CollieAiGuardrailConfigModel(GuardrailConfigModel[CollieAiGuardrailConfigModelOptionalParams]): + model_config = ConfigDict(frozen=True) + + api_key: str | None = Field( + default=None, + description="CollieAi project API key, starting with clai_. Leave empty to use COLLIEAI_API_KEY.", + ) + api_base: str | None = Field( + default=None, + description="CollieAi's LiteLLM integration base URL. Leave empty for CollieAi's hosted service.", + ) + + @staticmethod + def ui_friendly_name() -> str: + return "CollieAi" diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_collieai.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_collieai.py new file mode 100644 index 00000000000..40629355a57 --- /dev/null +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_collieai.py @@ -0,0 +1,362 @@ +from __future__ import annotations + +import importlib.util +from types import ModuleType, SimpleNamespace +from typing import Final + +import httpx +import pytest +import respx +from fastapi import HTTPException + +import litellm +from litellm.caching.dual_cache import DualCache +from litellm.exceptions import GuardrailRaisedException +from litellm.integrations.custom_guardrail import CustomGuardrail, GuardrailSetupError +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.guardrails.guardrail_endpoints import ( + CreateGuardrailRequest, + create_guardrail, + get_guardrail_ui_settings, + get_provider_specific_params, +) +from litellm.proxy.guardrails.guardrail_hooks.collieai import CollieAiGuardrail, initialize_guardrail +from litellm.proxy.guardrails.guardrail_hooks.collieai.collieai import ( + MISSING_PACKAGE_MESSAGE, + PACKAGE_MODULE, + has_readiness_guard, + load_base, + package_guardrail, +) +from litellm.proxy.guardrails.guardrail_registry import InMemoryGuardrailHandler +from litellm.proxy.utils import ProxyLogging +from litellm.types.guardrails import Guardrail, LitellmParams +from litellm.types.proxy.guardrails.guardrail_hooks.collieai import ( + DEFAULT_API_BASE, + CollieAiGuardrailConfigModel, + CollieAiGuardrailConfigModelOptionalParams, +) + +PACKAGE_INSTALLED: Final = importlib.util.find_spec("collieai_guardrail") is not None +API_KEY: Final = "clai_" + "0" * 40 +CHAT_REQUEST: Final = {"model": "gpt-5.4-mini", "messages": [{"role": "user", "content": "hello"}]} + + +class _RecordingGuardrail(CustomGuardrail): + def __init__( + self, + *, + api_base: str | None = None, + api_key: str | None = None, + guardrail_name: str | None = None, + event_hook: object = None, + default_on: bool | None = None, + stream_scope: object = None, + **settings: object, + ) -> None: + super().__init__(guardrail_name=guardrail_name, default_on=bool(default_on), stream_scope=stream_scope) + self.given = { + "api_base": api_base, + "api_key": api_key, + "guardrail_name": guardrail_name, + "event_hook": event_hook, + "default_on": default_on, + "stream_scope": stream_scope, + **settings, + } + + +class _RefusesStreamScope(_RecordingGuardrail): + def apply_stream_scope(self, stream_scope: object) -> None: + if stream_scope not in (None, "both"): + raise RuntimeError("enforcement_mode block needs stream_scope unset or both") + super().apply_stream_scope(stream_scope) + + +def _params(**extras: object) -> LitellmParams: + return LitellmParams( + guardrail="collieai", mode=["pre_call", "post_call"], default_on=True, api_key=API_KEY, **extras + ) + + +def _guardrail(litellm_params: LitellmParams) -> Guardrail: + return Guardrail(guardrail_name="collieai", litellm_params=litellm_params) + + +def _init(litellm_params: LitellmParams) -> _RecordingGuardrail: + callback: Final = initialize_guardrail( + litellm_params, _guardrail(litellm_params), guardrail_cls=_RecordingGuardrail + ) + assert isinstance(callback, _RecordingGuardrail) + return callback + + +@pytest.fixture(autouse=True) +def _isolate_callbacks(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "callbacks", []) + + +def test_hands_the_entry_and_every_setting_to_the_package() -> None: + callback: Final = _init( + _params( + api_base="https://collieai.internal.example/v1/integrations/litellm", + enforcement_mode="block", + deadline_ms=2000, + network_margin_ms=300, + max_body_bytes=1024, + forward_conversation_id=True, + summary_interval_s=30, + synthetic_interval_s=0, + workspace_id="not-a-collieai-setting", + stream_scope="non_streaming", + ) + ) + + assert callback.given == { + "api_base": "https://collieai.internal.example/v1/integrations/litellm", + "api_key": API_KEY, + "guardrail_name": "collieai", + "event_hook": ["pre_call", "post_call"], + "default_on": True, + "stream_scope": "non_streaming", + "enforcement_mode": "block", + "deadline_ms": 2000, + "network_margin_ms": 300, + "max_body_bytes": 1024, + "forward_conversation_id": True, + "summary_interval_s": 30, + "synthetic_interval_s": 0, + "workspace_id": "not-a-collieai-setting", + } + assert litellm.callbacks == [callback] + + +def test_points_at_app_collieai_io_unless_told_otherwise() -> None: + callback: Final = _init(_params()) + + # The hosted service's base URL, from https://docs.collieai.io/integrations/litellm (2026-10-29) + assert callback.given["api_base"] == DEFAULT_API_BASE == "https://app.collieai.io/v1/integrations/litellm" + assert set(callback.given) == {"api_base", "api_key", "guardrail_name", "event_hook", "default_on", "stream_scope"} + + +def test_a_cleared_optional_setting_leaves_the_package_its_default() -> None: + callback: Final = _init(_params(deadline_ms=None, forward_conversation_id=None, network_margin_ms=300)) + + assert "deadline_ms" not in callback.given + assert "forward_conversation_id" not in callback.given + assert callback.given["network_margin_ms"] == 300 + + +def test_a_stream_scope_the_package_refuses_is_a_setup_error_and_adds_no_callback() -> None: + litellm_params: Final = _params(enforcement_mode="block", stream_scope="streaming") + + with pytest.raises(GuardrailSetupError, match="stream_scope"): + initialize_guardrail(litellm_params, _guardrail(litellm_params), guardrail_cls=_RefusesStreamScope) + + assert litellm.callbacks == [] + + +def test_hands_each_entry_its_own_name_phase_and_default_on() -> None: + litellm_params: Final = LitellmParams(guardrail="collieai", mode="pre_call", default_on=False, api_key=API_KEY) + callback: Final = initialize_guardrail( + litellm_params, + Guardrail(guardrail_name="collieai-prompts", litellm_params=litellm_params), + guardrail_cls=_RecordingGuardrail, + ) + + assert isinstance(callback, _RecordingGuardrail) + assert (callback.given["guardrail_name"], callback.given["event_hook"], callback.given["default_on"]) == ( + "collieai-prompts", + "pre_call", + False, + ) + + +def test_takes_the_key_from_collieai_api_key_only_when_the_entry_has_none(monkeypatch: pytest.MonkeyPatch) -> None: + from_environment: Final = "clai_" + "1" * 40 + monkeypatch.setenv("COLLIEAI_API_KEY", from_environment) + without_key: Final = LitellmParams(guardrail="collieai", mode=["pre_call", "post_call"], default_on=True) + + assert _init(without_key).given["api_key"] == from_environment + assert _init(_params()).given["api_key"] == API_KEY + + +def test_an_untouched_ui_form_leaves_every_setting_to_the_package() -> None: + model: Final = CollieAiGuardrailConfigModel(api_key=API_KEY) + optional: Final = CollieAiGuardrailConfigModelOptionalParams() + callback: Final = _init( + _params(**model.model_dump(exclude={"api_key", "optional_params"}), **optional.model_dump()) + ) + + assert callback.given["api_base"] == DEFAULT_API_BASE + assert set(callback.given) == {"api_base", "api_key", "guardrail_name", "event_hook", "default_on", "stream_scope"} + + +@pytest.mark.asyncio +async def test_ui_offers_collieai_fields_without_the_package() -> None: + assert CollieAiGuardrail.get_config_model() is CollieAiGuardrailConfigModel + + fields: Final = (await get_provider_specific_params())["collieai"] + + assert fields["ui_friendly_name"] == "CollieAi" + assert fields["api_key"]["required"] is False + assert set(fields["optional_params"]["fields"]) == {"deadline_ms", "forward_conversation_id"} + assert [name for name in ("api_key", "api_base") if "default_value" in fields[name]] == [] + assert [name for name, field in fields["optional_params"]["fields"].items() if "default_value" in field] == [] + + +@pytest.mark.asyncio +async def test_ui_offers_only_the_two_phases_collieai_checks() -> None: + settings: Final = await get_guardrail_ui_settings() + + assert settings.supported_modes_by_provider["collieai"] == ["pre_call", "post_call"] + + +def test_a_package_without_the_readiness_guard_is_too_old() -> None: + assert not has_readiness_guard(ModuleType("collieai_guardrail.guardrail")) + assert has_readiness_guard(SimpleNamespace(require_guardrail=object())) + + +def test_only_a_current_package_s_guardrail_class_is_taken() -> None: + guard: Final = object() + + assert package_guardrail(SimpleNamespace(require_guardrail=guard, CollieAIGuardrail=_RecordingGuardrail)) is ( + _RecordingGuardrail + ) + assert package_guardrail(SimpleNamespace(CollieAIGuardrail=_RecordingGuardrail)) is None + assert package_guardrail(SimpleNamespace(require_guardrail=guard, CollieAIGuardrail=object())) is None + assert package_guardrail(SimpleNamespace(require_guardrail=guard, CollieAIGuardrail=ValueError)) is None + + +def test_builds_on_the_package_only_when_it_is_current() -> None: + current: Final = SimpleNamespace(require_guardrail=object(), CollieAIGuardrail=_RecordingGuardrail) + too_old: Final = SimpleNamespace(CollieAIGuardrail=_RecordingGuardrail) + + assert load_base({PACKAGE_MODULE: current}.__getitem__) is _RecordingGuardrail + with pytest.raises(ImportError, match="pip install") as raised: + load_base({PACKAGE_MODULE: too_old}.__getitem__)(guardrail_name="collieai") + + assert str(raised.value) == MISSING_PACKAGE_MESSAGE + assert raised.value.__cause__ is None + + +def test_without_the_package_building_the_guardrail_names_what_to_install() -> None: + cause: Final = ModuleNotFoundError("No module named 'collieai_guardrail'", name="collieai_guardrail") + + def import_module(name: str) -> object: + raise cause + + with pytest.raises(ImportError, match="pip install") as raised: + load_base(import_module)(guardrail_name="collieai") + + assert str(raised.value) == MISSING_PACKAGE_MESSAGE + assert raised.value.__cause__ is cause + + +@pytest.mark.parametrize( + "cause", + [ + ModuleNotFoundError("No module named 'litellm.gone'", name="litellm.gone"), + ImportError("cannot import name 'Moved' from 'litellm.types'", name="litellm.types"), + ImportError("cannot import name 'Observer' from 'collieai_guardrail.observer'", name="collieai_guardrail.observer"), + AttributeError("module 'litellm' has no attribute 'moved'"), + ], +) +def test_an_installed_package_that_fails_to_load_says_so_instead_of_asking_to_install_it(cause: Exception) -> None: + def import_module(name: str) -> object: + raise cause + + with pytest.raises(ImportError, match="installed but failed to load") as raised: + load_base(import_module)(guardrail_name="collieai") + + assert str(cause) in str(raised.value) + assert MISSING_PACKAGE_MESSAGE not in str(raised.value) + assert raised.value.__cause__ is cause + + +@pytest.mark.skipif(PACKAGE_INSTALLED, reason="exercises the missing-package fallback") +def test_missing_package_fails_at_config_load_with_install_hint() -> None: + with pytest.raises(GuardrailSetupError, match="pip install") as raised: + InMemoryGuardrailHandler().initialize_guardrail(_guardrail(_params())) + + assert not isinstance(raised.value, (ValueError, TypeError)) + assert litellm.callbacks == [] + + +@pytest.mark.asyncio +async def test_creating_a_guardrail_the_package_refuses_answers_400_and_removes_the_row(mocker) -> None: + from unittest.mock import AsyncMock + + litellm_params: Final = { + "guardrail": "collieai", + "mode": ["pre_call", "post_call"], + "default_on": True, + "api_key": "sk-not-a-collieai-key", + } + row: Final = {"guardrail_name": "collieai-refused", "litellm_params": litellm_params, "guardrail_info": {}} + registry: Final = mocker.Mock(add_guardrail_to_db=AsyncMock(return_value={**row, "guardrail_id": "refused"})) + mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", registry) + mocker.patch("litellm.proxy.proxy_server.prisma_client", mocker.Mock()) + mocker.patch("litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", InMemoryGuardrailHandler()) + delete_row: Final = mocker.patch("litellm.proxy.guardrails.guardrail_endpoints._delete_guardrail_row", AsyncMock()) + + with pytest.raises(HTTPException) as refused: + await create_guardrail( + CreateGuardrailRequest(guardrail=row), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + assert refused.value.status_code == 400 + delete_row.assert_awaited_once() + assert litellm.callbacks == [] + + +def _load(monkeypatch: pytest.MonkeyPatch, **extras: object) -> CustomGuardrail: + """The guardrail as the proxy loads `guardrail: collieai` from its config, with no synthetic check at start. Block + mode also needs the package's readiness guard among the proxy's callbacks, which the config file loads first""" + from collieai_guardrail.guardrail import require_guardrail + + block: Final = extras.get("enforcement_mode") == "block" + monkeypatch.setattr(litellm, "callbacks", [require_guardrail] if block else []) + InMemoryGuardrailHandler().initialize_guardrail(_guardrail(_params(synthetic_interval_s=0, **extras))) + loaded: Final = [callback for callback in litellm.callbacks if isinstance(callback, CollieAiGuardrail)] + assert len(loaded) == 1 + return loaded[0] + + +async def _pre_call() -> dict[str, object]: + data: Final = {**CHAT_REQUEST, "metadata": {}} + await ProxyLogging(user_api_key_cache=DualCache()).pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(request_route="/v1/chat/completions"), data=data, call_type="acompletion" + ) + return data + + +@pytest.mark.skipif(not PACKAGE_INSTALLED, reason="needs collieai-litellm-guardrail") +@pytest.mark.asyncio +@respx.mock +async def test_block_refuses_the_call_when_collieai_is_unreachable(monkeypatch: pytest.MonkeyPatch) -> None: + route: Final = respx.post(url__startswith=DEFAULT_API_BASE).mock(side_effect=httpx.ConnectError("refused")) + _load(monkeypatch, enforcement_mode="block") + + with pytest.raises((GuardrailRaisedException, HTTPException)) as refused: + await _pre_call() + + assert route.called + assert getattr(refused.value, "status_code", None) == 503 + + +@pytest.mark.skipif(not PACKAGE_INSTALLED, reason="needs collieai-litellm-guardrail") +@pytest.mark.asyncio +@respx.mock +async def test_monitor_lets_the_call_go_on_and_records_why(monkeypatch: pytest.MonkeyPatch) -> None: + route: Final = respx.post(url__startswith=DEFAULT_API_BASE).mock(side_effect=httpx.ConnectError("refused")) + _load(monkeypatch) + + data: Final = await _pre_call() + + assert route.called + metadata: Final = data["metadata"] + assert isinstance(metadata, dict) + [record] = metadata["standard_logging_guardrail_information"] + assert (record["guardrail_provider"], record["guardrail_status"]) == ("collieai", "guardrail_failed_to_respond") diff --git a/tests/unit/proxy/guardrails/test_guardrail_registry.py b/tests/unit/proxy/guardrails/test_guardrail_registry.py index 7eabc24e090..bf328264bab 100644 --- a/tests/unit/proxy/guardrails/test_guardrail_registry.py +++ b/tests/unit/proxy/guardrails/test_guardrail_registry.py @@ -649,6 +649,38 @@ def test_unchanged_db_params_do_not_register_as_changed(): assert handler._has_guardrail_params_changed(gid, new) is False +@pytest.mark.parametrize("field", ("api_key", "api_base")) +def test_db_poll_does_not_rebuild_a_guardrail_whose_secret_field_is_an_env_reference(monkeypatch, field): + from litellm.proxy.guardrails import guardrail_registry as registry_module + + built_with: Final[list[str | None]] = [] + + def _initializer(litellm_params, guardrail): + built_with.append(getattr(litellm_params, field)) + return CustomGuardrail( + guardrail_name=guardrail["guardrail_name"], event_hook=GuardrailEventHooks.pre_call, default_on=True + ) + + monkeypatch.setitem(registry_module.guardrail_initializer_registry, "env_ref_test", _initializer) + monkeypatch.setenv("ENV_REF_TEST_VALUE", "resolved-value") + row: Final = { + "guardrail_id": "env-ref", + "guardrail_name": "env-ref", + "litellm_params": { + "guardrail": "env_ref_test", + "mode": "pre_call", + field: "os.environ/ENV_REF_TEST_VALUE", + }, + } + handler: Final = InMemoryGuardrailHandler() + handler.initialize_guardrail(guardrail=dict(row), source="db") + + handler.sync_guardrail_from_db(dict(row)) + handler.sync_guardrail_from_db(dict(row)) + + assert built_with == ["resolved-value"] + + def test_db_poll_does_not_reinitialize_config_guardrail_without_default_on(): handler = InMemoryGuardrailHandler() guardrail_id: Final = "config-default-on-guardrail" diff --git a/ui/litellm-dashboard/public/assets/logos/collieai.svg b/ui/litellm-dashboard/public/assets/logos/collieai.svg new file mode 100644 index 00000000000..dff072d0c30 --- /dev/null +++ b/ui/litellm-dashboard/public/assets/logos/collieai.svg @@ -0,0 +1,6 @@ + + + + + + diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_configs.ts b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_configs.ts index e532b5bfd74..7a9060c5d02 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_configs.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_configs.ts @@ -393,4 +393,10 @@ export const GUARDRAIL_PRESETS: Record = { mode: ["pre_call", "post_call"], defaultOn: false, }, + collieai: { + provider: "Collieai", + guardrailNameSuggestion: "CollieAi", + mode: ["pre_call", "post_call"], + defaultOn: false, + }, }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.test.ts index 20725e73e37..9839db91376 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.test.ts @@ -6,6 +6,13 @@ import { LITELLM_CONTENT_FILTER_CARDS, PARTNER_GUARDRAIL_CARDS, } from "./guardrail_garden_data"; +import { + getGuardrailLogoAndName, + getGuardrailProviders, + guardrailLogoMap, + populateGuardrailProviderMap, + populateGuardrailProviders, +} from "./guardrail_info_helpers"; const EXPECTED_PARTNER_LOGO_FILES: Record = { presidio: "microsoft_azure.svg", @@ -39,6 +46,7 @@ const EXPECTED_PARTNER_LOGO_FILES: Record = { conduct: "conduct.png", wingback: "wingback.png", airia: "airia.svg", + collieai: "collieai.svg", }; describe("guardrail_garden_data logos", () => { @@ -100,3 +108,26 @@ describe("decision model presets", () => { } }); }); + +describe("collieai preset", () => { + it("opens the CollieAi card on both halves of the call, off until a request or Always On asks for it", () => { + const expectedPreset = { + provider: "Collieai", + guardrailNameSuggestion: "CollieAi", + mode: ["pre_call", "post_call"], + defaultOn: false, + }; + expect(GUARDRAIL_PRESETS["collieai"]).toEqual(expectedPreset); + }); + + it("uses the provider key the proxy's provider list generates, so saved guardrails keep their name and logo", () => { + const providerParams = { collieai: { ui_friendly_name: "CollieAi" } }; + populateGuardrailProviders(providerParams); + populateGuardrailProviderMap(providerParams); + const provider = GUARDRAIL_PRESETS["collieai"].provider; + + expect(getGuardrailLogoAndName("collieai")).toEqual({ logo: guardrailLogoMap.CollieAi, displayName: "CollieAi" }); + expect(Object.keys(getGuardrailProviders())).toContain(provider); + expect(PARTNER_GUARDRAIL_CARDS.find((card) => card.id === "collieai")?.providerKey).toBe(provider); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.ts b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.ts index dc6acffa13c..9cb9521bb9d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.ts @@ -717,6 +717,16 @@ export const PARTNER_GUARDRAIL_CARDS: GuardrailCardInfo[] = [ tags: ["Security", "Prompt Injection", "PII", "Secrets", "Redaction"], providerKey: "Airia", }, + { + id: "collieai", + name: "CollieAi", + description: + "CollieAi checks chat completion text against your CollieAi policy before and after the model call and records each finding in LiteLLM's guardrail log. It observes by default. Block mode refuses the calls the policy blocks and every call it cannot check.", + category: "partner", + logo: guardrailLogoMap["CollieAi"], + tags: ["Security", "Prompt Injection", "PII", "Policy", "Monitoring"], + providerKey: "Collieai", + }, ]; export const ALL_CARDS = [...LITELLM_CONTENT_FILTER_CARDS, ...DECISION_MODEL_CARDS, ...PARTNER_GUARDRAIL_CARDS]; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx index d0253875735..bb773bbbb7f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx @@ -9,6 +9,7 @@ import conductLogo from "../../../../../public/assets/logos/conduct.png"; import aporiaLogo from "../../../../../public/assets/logos/aporia.png"; import bedrockLogo from "../../../../../public/assets/logos/bedrock.svg"; import catoNetworksLogo from "../../../../../public/assets/logos/cato_networks.svg"; +import collieaiLogo from "../../../../../public/assets/logos/collieai.svg"; import ciscoLogo from "../../../../../public/assets/logos/cisco.png"; import deepkeepLogo from "../../../../../public/assets/logos/deepkeep.svg"; import enkryptAiLogo from "../../../../../public/assets/logos/enkrypt_ai.avif"; @@ -94,6 +95,7 @@ export const guardrail_provider_map: Record = { "LLM Shield Proxy": "llm_shield_proxy", Conduct: "conduct", DecisionModel: "decision_model", + Collieai: "collieai", }; // Search aliases matched alongside the provider label in the add-guardrail provider combobox @@ -332,6 +334,7 @@ export const guardrailLogoMap = { "Conduct Guard": conductLogo.src, Wingback: wingbackLogo.src, "Airia Guardrail": airiaLogo.src, + CollieAi: collieaiLogo.src, "Decision Model": litellmLogo.src, } satisfies Record; diff --git a/ui/litellm-dashboard/src/lib/logoTreatments.test.ts b/ui/litellm-dashboard/src/lib/logoTreatments.test.ts index 828ea72a57f..6ae74738b4d 100644 --- a/ui/litellm-dashboard/src/lib/logoTreatments.test.ts +++ b/ui/litellm-dashboard/src/lib/logoTreatments.test.ts @@ -14,6 +14,7 @@ describe("logoTreatmentFor", () => { it("marks a multicolor dark mark for a plate instead of inversion", () => { expect(logoTreatmentFor("/ui/assets/logos/fireworks.svg")).toBe("plate"); + expect(logoTreatmentFor("/ui/assets/logos/collieai.svg")).toBe("plate"); }); it("plates a dark mark with a light knockout, which inversion would flatten away", () => { diff --git a/ui/litellm-dashboard/src/lib/logoTreatments.ts b/ui/litellm-dashboard/src/lib/logoTreatments.ts index 4cd6f65a613..f8b4ecc05b1 100644 --- a/ui/litellm-dashboard/src/lib/logoTreatments.ts +++ b/ui/litellm-dashboard/src/lib/logoTreatments.ts @@ -36,6 +36,7 @@ const TREATMENT_BY_ASSET: Readonly> = { "aiml_api.svg": "plate", "akto.svg": "plate", "aws.svg": "plate", + "collieai.svg": "plate", "deepkeep.svg": "plate", "fireworks.svg": "plate", "llm_guard.png": "plate",