diff --git a/litellm/proxy/guardrails/guardrail_hooks/collieai/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/collieai/__init__.py index fd3060e637b..b8d18fd4166 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/collieai/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/collieai/__init__.py @@ -13,8 +13,6 @@ if TYPE_CHECKING: from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.types.guardrails import Guardrail, LitellmParams -# The package's own settings a config entry may carry beside LiteLLM's. The package validates their values and -# names a bad one at startup PACKAGE_SETTINGS: Final = frozenset( { "enforcement_mode", diff --git a/litellm/proxy/guardrails/guardrail_hooks/collieai/collieai.py b/litellm/proxy/guardrails/guardrail_hooks/collieai/collieai.py index 07c6f81e84b..267bf48adb1 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/collieai/collieai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/collieai/collieai.py @@ -7,6 +7,7 @@ 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 @@ -14,7 +15,7 @@ from litellm.types.proxy.guardrails.guardrail_hooks.collieai import CollieAiGuar PACKAGE_MODULE: Final = "collieai_guardrail.guardrail" MISSING_PACKAGE_MESSAGE: Final = ( - "collieai-litellm-guardrail>=1.0.0 is required for the CollieAi guardrail. " + "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"' ) @@ -42,9 +43,9 @@ def missing_package(import_error: ImportError | None) -> type[CustomGuardrail]: return MissingPackage -def load_base() -> type[CustomGuardrail]: +def load_base(import_module: Callable[[str], object] = importlib.import_module) -> type[CustomGuardrail]: try: - package: Final = importlib.import_module(PACKAGE_MODULE) + package: Final = import_module(PACKAGE_MODULE) except ImportError as import_error: return missing_package(import_error) return package_guardrail(package) or missing_package(None) diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/collieai.py b/litellm/types/proxy/guardrails/guardrail_hooks/collieai.py index 57b73cda82b..d6d02a1571c 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/collieai.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/collieai.py @@ -26,6 +26,8 @@ class CollieAiGuardrailConfigModelOptionalParams(LiteLLMBaseModel): class CollieAiGuardrailConfigModel(GuardrailConfigModel[CollieAiGuardrailConfigModelOptionalParams]): + model_config = ConfigDict(frozen=True) + api_key: str = Field( min_length=1, description="CollieAi project API key, starting with clai_. Env: COLLIEAI_API_KEY.", diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_collieai.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_collieai.py index c1b60664e15..7085772b94b 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_collieai.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_collieai.py @@ -18,8 +18,9 @@ from litellm.proxy.guardrails.guardrail_endpoints import get_provider_specific_p 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, - missing_package, + load_base, package_guardrail, ) from litellm.proxy.guardrails.guardrail_registry import InMemoryGuardrailHandler @@ -126,7 +127,9 @@ def test_points_at_app_collieai_io_unless_told_otherwise() -> None: def test_ui_form_defaults_are_what_the_package_gets() -> None: optional: Final = CollieAiGuardrailConfigModelOptionalParams() model: Final = CollieAiGuardrailConfigModel(api_key=API_KEY) - callback: Final = _init(_params(**model.model_dump(exclude={"api_key", "optional_params"}), **optional.model_dump())) + callback: Final = _init( + _params(**model.model_dump(exclude={"api_key", "optional_params"}), **optional.model_dump()) + ) assert callback.given["api_base"] == model.api_base assert callback.given["deadline_ms"] == optional.deadline_ms @@ -161,11 +164,26 @@ def test_only_a_current_package_s_guardrail_class_is_taken() -> 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 = ImportError("No module named 'collieai_guardrail'") + def import_module(name: str) -> object: + raise cause + with pytest.raises(ImportError, match="pip install") as raised: - missing_package(cause)(guardrail_name="collieai") + load_base(import_module)(guardrail_name="collieai") assert str(raised.value) == MISSING_PACKAGE_MESSAGE assert raised.value.__cause__ is cause 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 e5f48d1fd41..4c7985b3114 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 @@ -396,8 +396,6 @@ export const GUARDRAIL_PRESETS: Record = { collieai: { provider: "CollieAi", guardrailNameSuggestion: "CollieAi", - // CollieAi checks both halves of every call. Block mode also needs the package's readiness - // guard in the proxy config file, so a guardrail added here observes. mode: ["pre_call", "post_call"], defaultOn: true, }, 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 5ee6530b9f5..2f315fe69aa 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 @@ -101,3 +101,14 @@ describe("decision model presets", () => { } }); }); + +describe("collieai preset", () => { + it("opens the CollieAi card as an always-on guardrail on both halves of the call", () => { + expect(GUARDRAIL_PRESETS["collieai"]).toEqual({ + provider: "CollieAi", + guardrailNameSuggestion: "CollieAi", + mode: ["pre_call", "post_call"], + defaultOn: true, + }); + }); +});