refactor(guardrails): tighten the CollieAi wrapper and its tests

The loader takes its importer as a parameter, so a test covers the fallback for a package that is too old, and the install hint names the Python the package needs

The config model is frozen, the Garden preset's phases and default_on are pinned by a test, and comments the code already says are gone

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",

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)

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.",

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

ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_configs.ts

@@ -396,8 +396,6 @@ export const GUARDRAIL_PRESETS: Record<string, GuardrailPreset> = {
   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,
   },

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,
+    });
+  });
+});