From b5bd7ff6820cb2e169618c4167508ac5268d441a Mon Sep 17 00:00:00 2001 From: CollieAi Date: Thu, 29 Oct 2026 20:39:34 +0100 Subject: [PATCH 1/5] feat(guardrails): add the CollieAi guardrail integration --- .../guardrail_hooks/collieai/__init__.py | 62 +++++ .../guardrail_hooks/collieai/collieai.py | 70 ++++++ litellm/types/guardrails.py | 1 + .../guardrails/guardrail_hooks/collieai.py | 40 +++ .../guardrail_hooks/test_collieai.py | 230 ++++++++++++++++++ .../public/assets/logos/collieai.svg | 6 + .../_components/guardrail_garden_configs.ts | 8 + .../_components/guardrail_garden_data.test.ts | 1 + .../_components/guardrail_garden_data.ts | 10 + .../_components/guardrail_info_helpers.tsx | 3 + 10 files changed, 431 insertions(+) create mode 100644 litellm/proxy/guardrails/guardrail_hooks/collieai/__init__.py create mode 100644 litellm/proxy/guardrails/guardrail_hooks/collieai/collieai.py create mode 100644 litellm/types/proxy/guardrails/guardrail_hooks/collieai.py create mode 100644 tests/unit/proxy/guardrails/guardrail_hooks/test_collieai.py create mode 100644 ui/litellm-dashboard/public/assets/logos/collieai.svg 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..fd3060e637b --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/collieai/__init__.py @@ -0,0 +1,62 @@ +from __future__ import annotations + +from collections.abc import Callable, Mapping +from types import MappingProxyType +from typing import TYPE_CHECKING, Final + +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 + +# 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", + "deadline_ms", + "network_margin_ms", + "max_body_bytes", + "forward_conversation_id", + "summary_interval_s", + "synthetic_interval_s", + } +) +_NO_EXTRAS: Final[Mapping[str, object]] = MappingProxyType({}) + + +def package_settings(extras: Mapping[str, object]) -> Mapping[str, object]: + return MappingProxyType({key: value for key, value in extras.items() if key in PACKAGE_SETTINGS}) + + +def initialize_guardrail( + litellm_params: LitellmParams, + guardrail: Guardrail, + guardrail_cls: Callable[..., CustomGuardrail] = CollieAiGuardrail, +) -> CustomGuardrail: + import litellm + + extras: Final = litellm_params.model_extra or _NO_EXTRAS + _callback: Final = guardrail_cls( + api_base=litellm_params.api_base or DEFAULT_API_BASE, + api_key=litellm_params.api_key, + guardrail_name=guardrail["guardrail_name"], + event_hook=litellm_params.mode, + default_on=litellm_params.default_on, + **package_settings(extras), + ) + 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..07c6f81e84b --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/collieai/collieai.py @@ -0,0 +1,70 @@ +"""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 typing import Final + +from litellm.integrations.custom_guardrail import CustomGuardrail +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 is required for the CollieAi guardrail. " + 'Install it with: pip install "collieai-litellm-guardrail>=1.0.0"' +) + + +def has_readiness_guard(package: object) -> bool: + """1.0.0 is the first release with the readiness guard, `require_guardrail`, that the package's block mode needs""" + return getattr(package, "require_guardrail", None) is not None + + +def package_guardrail(package: object) -> type[CustomGuardrail] | None: + """The package's guardrail class, or None for a package too old or not shaped as one""" + 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 missing_package(import_error: ImportError | None) -> type[CustomGuardrail]: + """A stand-in that fails where the proxy builds the guardrail, naming what to install""" + + class MissingPackage(CustomGuardrail): + def __init__(self, **kwargs: object) -> None: # kwargs-ok: mirrors the package constructor, only raises + raise ImportError(MISSING_PACKAGE_MESSAGE) from import_error + + return MissingPackage + + +def load_base() -> type[CustomGuardrail]: + try: + package: Final = importlib.import_module(PACKAGE_MODULE) + except ImportError as import_error: + return missing_package(import_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 + + +__all__ = ( + "MISSING_PACKAGE_MESSAGE", + "PACKAGE_MODULE", + "CollieAiGuardrail", + "has_readiness_guard", + "load_base", + "missing_package", + "package_guardrail", +) diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 0803485aa02..67d9050ac9a 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..57b73cda82b --- /dev/null +++ b/litellm/types/proxy/guardrails/guardrail_hooks/collieai.py @@ -0,0 +1,40 @@ +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=2500, + gt=0, + description="Each phase's whole deadline in milliseconds: connection, send, reply and parse.", + ) + forward_conversation_id: bool | None = Field( + default=False, + description="Forward the caller's x-collieai-conversation-id header, and only that header.", + ) + + +class CollieAiGuardrailConfigModel(GuardrailConfigModel[CollieAiGuardrailConfigModelOptionalParams]): + api_key: str = Field( + min_length=1, + description="CollieAi project API key, starting with clai_. Env: COLLIEAI_API_KEY.", + ) + api_base: str | None = Field( + default=DEFAULT_API_BASE, + description="CollieAi's LiteLLM integration base URL.", + ) + + @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..c1b60664e15 --- /dev/null +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_collieai.py @@ -0,0 +1,230 @@ +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 +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.guardrails.guardrail_endpoints import 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, + has_readiness_guard, + missing_package, + 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-mini", "messages": [{"role": "user", "content": "hello"}]} + + +class _RecordingGuardrail(CustomGuardrail): + """Stand-in with the ``collieai_guardrail.guardrail.CollieAIGuardrail`` constructor: it keeps what it was given""" + + 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, + **settings: object, + ) -> None: + super().__init__(guardrail_name=guardrail_name, default_on=bool(default_on)) + self.given = { + "api_base": api_base, + "api_key": api_key, + "guardrail_name": guardrail_name, + "event_hook": event_hook, + "default_on": default_on, + **settings, + } + + +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_the_package_settings_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", + ) + ) + + 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, + "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, + } + 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"} + + +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())) + + assert callback.given["api_base"] == model.api_base + assert callback.given["deadline_ms"] == optional.deadline_ms + assert callback.given["forward_conversation_id"] == optional.forward_conversation_id + + +@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 True + assert fields["api_base"]["default_value"] == DEFAULT_API_BASE + assert set(fields["optional_params"]["fields"]) == {"deadline_ms", "forward_conversation_id"} + + +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_without_the_package_building_the_guardrail_names_what_to_install() -> None: + cause: Final = ImportError("No module named 'collieai_guardrail'") + + with pytest.raises(ImportError, match="pip install") as raised: + missing_package(cause)(guardrail_name="collieai") + + assert str(raised.value) == MISSING_PACKAGE_MESSAGE + 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(ImportError, match="pip install"): + InMemoryGuardrailHandler().initialize_guardrail(_guardrail(_params())) + + 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/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..e5f48d1fd41 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,12 @@ export const GUARDRAIL_PRESETS: Record = { mode: ["pre_call", "post_call"], defaultOn: false, }, + 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 20725e73e37..5ee6530b9f5 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 @@ -39,6 +39,7 @@ const EXPECTED_PARTNER_LOGO_FILES: Record = { conduct: "conduct.png", wingback: "wingback.png", airia: "airia.svg", + collieai: "collieai.svg", }; describe("guardrail_garden_data logos", () => { 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 7ef4f159717..aebfadfad2f 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 @@ -577,6 +577,16 @@ export const PARTNER_GUARDRAIL_CARDS: GuardrailCardInfo[] = [ tags: ["Security", "Prompt Injection", "PII", "Secrets", "Redaction"], providerKey: "Airia", }, + { + id: "collieai", + name: "CollieAi", + description: + "CollieAi checks prompts and model responses 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, set in the proxy config file, refuses the calls the policy blocks.", + 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 6c4bb8c4313..cf10e03d81c 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 @@ -7,6 +7,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"; @@ -92,6 +93,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 @@ -327,6 +329,7 @@ export const guardrailLogoMap = { "Conduct Guard": conductLogo.src, Wingback: wingbackLogo.src, "Airia Guardrail": airiaLogo.src, + CollieAi: collieaiLogo.src, "Decision Model": litellmLogo.src, } satisfies Record; From 348b82fd33758e111db6ffdb5bda8f5b5a4a6eba Mon Sep 17 00:00:00 2001 From: CollieAi Date: Sat, 31 Oct 2026 10:43:06 +0100 Subject: [PATCH 2/5] 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 --- .../guardrail_hooks/collieai/__init__.py | 2 -- .../guardrail_hooks/collieai/collieai.py | 7 +++--- .../guardrails/guardrail_hooks/collieai.py | 2 ++ .../guardrail_hooks/test_collieai.py | 24 ++++++++++++++++--- .../_components/guardrail_garden_configs.ts | 2 -- .../_components/guardrail_garden_data.test.ts | 11 +++++++++ 6 files changed, 38 insertions(+), 10 deletions(-) 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, + }); + }); +}); From 39e4bbc19179dc8388ea0789a50541fa5a69cd03 Mon Sep 17 00:00:00 2001 From: CollieAi Date: Wed, 4 Nov 2026 21:57:48 +0100 Subject: [PATCH 3/5] fix(guardrails): use the generated Collieai provider key in the dashboard The dashboard builds provider keys from the proxy's provider list, so collieai becomes Collieai. The static CollieAi key came first in the provider map, which left saved CollieAi guardrails without their display name and logo. The map, the Garden preset and the card now use Collieai, and a test covers the lookup The preset expectation is a named value now, which keeps the dashboard under its inline object lint budget --- .../_components/guardrail_garden_configs.ts | 2 +- .../_components/guardrail_garden_data.test.ts | 25 ++++++++++++++++--- .../_components/guardrail_garden_data.ts | 2 +- .../_components/guardrail_info_helpers.tsx | 2 +- 4 files changed, 25 insertions(+), 6 deletions(-) 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 4c7985b3114..ca3101224ae 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 @@ -394,7 +394,7 @@ export const GUARDRAIL_PRESETS: Record = { defaultOn: false, }, collieai: { - provider: "CollieAi", + provider: "Collieai", guardrailNameSuggestion: "CollieAi", 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 2f315fe69aa..061e0fefb40 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", @@ -104,11 +111,23 @@ 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", + const expectedPreset = { + provider: "Collieai", guardrailNameSuggestion: "CollieAi", mode: ["pre_call", "post_call"], defaultOn: true, - }); + }; + 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 aebfadfad2f..3ffe2ba2d1f 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 @@ -585,7 +585,7 @@ export const PARTNER_GUARDRAIL_CARDS: GuardrailCardInfo[] = [ category: "partner", logo: guardrailLogoMap["CollieAi"], tags: ["Security", "Prompt Injection", "PII", "Policy", "Monitoring"], - providerKey: "CollieAi", + providerKey: "Collieai", }, ]; 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 cf10e03d81c..6675f3121b8 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 @@ -93,7 +93,7 @@ export const guardrail_provider_map: Record = { "LLM Shield Proxy": "llm_shield_proxy", Conduct: "conduct", DecisionModel: "decision_model", - CollieAi: "collieai", + Collieai: "collieai", }; // Search aliases matched alongside the provider label in the add-guardrail provider combobox From 462b8f795ad8a9fd83fd21542dde4a4b8db601ba Mon Sep 17 00:00:00 2001 From: CollieAi Date: Wed, 4 Nov 2026 21:57:48 +0100 Subject: [PATCH 4/5] refactor(guardrails): drop the CollieAi wrapper's function docstrings AGENTS.md keeps comments to complex logic and tool input, and these only restated what each function returns --- litellm/proxy/guardrails/guardrail_hooks/collieai/collieai.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/collieai/collieai.py b/litellm/proxy/guardrails/guardrail_hooks/collieai/collieai.py index 267bf48adb1..5b77aa6985d 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/collieai/collieai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/collieai/collieai.py @@ -21,12 +21,10 @@ MISSING_PACKAGE_MESSAGE: Final = ( def has_readiness_guard(package: object) -> bool: - """1.0.0 is the first release with the readiness guard, `require_guardrail`, that the package's block mode needs""" return getattr(package, "require_guardrail", None) is not None def package_guardrail(package: object) -> type[CustomGuardrail] | None: - """The package's guardrail class, or None for a package too old or not shaped as one""" guardrail: Final[object] = getattr(package, "CollieAIGuardrail", None) if not has_readiness_guard(package) or not isinstance(guardrail, type): return None @@ -34,8 +32,6 @@ def package_guardrail(package: object) -> type[CustomGuardrail] | None: def missing_package(import_error: ImportError | None) -> type[CustomGuardrail]: - """A stand-in that fails where the proxy builds the guardrail, naming what to install""" - class MissingPackage(CustomGuardrail): def __init__(self, **kwargs: object) -> None: # kwargs-ok: mirrors the package constructor, only raises raise ImportError(MISSING_PACKAGE_MESSAGE) from import_error From 740f7b56bf34ddd79883026cc835781d1379806e Mon Sep 17 00:00:00 2001 From: CollieAi Date: Thu, 5 Nov 2026 09:26:12 +0100 Subject: [PATCH 5/5] fix(guardrails): declare CollieAi's phases, read COLLIEAI_API_KEY and plate its logo The dashboard offered every mode for CollieAi, and a guardrail saved with a mode the package refuses was reported as created but never loaded. The wrapper now declares pre_call and post_call, so the form offers only those api_key falls back to COLLIEAI_API_KEY as its description says, and the navy logo gets a plate in dark mode like DeepKeep's The tests now also cover an entry's own name, phase and default_on, the key fallback, and the modes the dashboard offers --- .../guardrail_hooks/collieai/__init__.py | 3 +- .../guardrail_hooks/collieai/collieai.py | 5 +++ .../guardrail_hooks/test_collieai.py | 36 +++++++++++++++++-- .../src/lib/logoTreatments.test.ts | 1 + .../src/lib/logoTreatments.ts | 1 + 5 files changed, 43 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/collieai/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/collieai/__init__.py index b8d18fd4166..7f08ce6f15c 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/collieai/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/collieai/__init__.py @@ -37,11 +37,12 @@ def initialize_guardrail( guardrail_cls: Callable[..., CustomGuardrail] = CollieAiGuardrail, ) -> CustomGuardrail: import litellm + from litellm.secret_managers.main import get_secret_str extras: Final = litellm_params.model_extra or _NO_EXTRAS _callback: Final = guardrail_cls( api_base=litellm_params.api_base or DEFAULT_API_BASE, - api_key=litellm_params.api_key, + 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, diff --git a/litellm/proxy/guardrails/guardrail_hooks/collieai/collieai.py b/litellm/proxy/guardrails/guardrail_hooks/collieai/collieai.py index 5b77aa6985d..e8ae264db8a 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/collieai/collieai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/collieai/collieai.py @@ -11,6 +11,7 @@ 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" @@ -55,6 +56,10 @@ class CollieAiGuardrail(_Base): 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", diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_collieai.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_collieai.py index 7085772b94b..687709c93c5 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_collieai.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_collieai.py @@ -14,7 +14,7 @@ from litellm.caching.dual_cache import DualCache from litellm.exceptions import GuardrailRaisedException from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.proxy._types import UserAPIKeyAuth -from litellm.proxy.guardrails.guardrail_endpoints import get_provider_specific_params +from litellm.proxy.guardrails.guardrail_endpoints import 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, @@ -34,7 +34,7 @@ from litellm.types.proxy.guardrails.guardrail_hooks.collieai import ( PACKAGE_INSTALLED: Final = importlib.util.find_spec("collieai_guardrail") is not None API_KEY: Final = "clai_" + "0" * 40 -CHAT_REQUEST: Final = {"model": "gpt-5-mini", "messages": [{"role": "user", "content": "hello"}]} +CHAT_REQUEST: Final = {"model": "gpt-5.4-mini", "messages": [{"role": "user", "content": "hello"}]} class _RecordingGuardrail(CustomGuardrail): @@ -124,6 +124,31 @@ def test_points_at_app_collieai_io_unless_told_otherwise() -> None: assert set(callback.given) == {"api_base", "api_key", "guardrail_name", "event_hook", "default_on"} +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_ui_form_defaults_are_what_the_package_gets() -> None: optional: Final = CollieAiGuardrailConfigModelOptionalParams() model: Final = CollieAiGuardrailConfigModel(api_key=API_KEY) @@ -148,6 +173,13 @@ async def test_ui_offers_collieai_fields_without_the_package() -> None: assert set(fields["optional_params"]["fields"]) == {"deadline_ms", "forward_conversation_id"} +@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())) 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 fc19c5c364d..4a63b34ce64 100644 --- a/ui/litellm-dashboard/src/lib/logoTreatments.ts +++ b/ui/litellm-dashboard/src/lib/logoTreatments.ts @@ -35,6 +35,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",