feat(guardrails): add the CollieAi guardrail integration (#49224)

* feat(guardrails): add the CollieAi guardrail integration

* 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

* 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

* 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

* 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

* fix(guardrails): reject CollieAi guardrails the package refuses and stop env-keyed rebuilds

A Garden create whose key or settings the package refuses now answers 400 and
saves nothing, instead of saving a guardrail that checks nothing. Every setting
on the entry now reaches the package, like the class-path route. DB guardrails
keyed with an os.environ/ reference are no longer rebuilt on every DB poll.
The Garden card says which calls CollieAi checks.

* fix(guardrails): load CollieAi stream_scope and import errors as setup errors, default the Garden card off

Pass stream_scope into the package constructor inside the setup-error wrap so a refused scope answers 400 on create and stops a config-file proxy at startup. Drop settings cleared to null so the package default applies. Say an installed package failed to load, with the error, instead of asking to install it, and keep any import-time exception from breaking the guardrail routes. Make api_key optional in the form since COLLIEAI_API_KEY is the fallback, and open the Garden card with Always On off like the other partner presets.

* fix(guardrails): leave CollieAi form fields empty and share one guardrail setup error

The Garden form no longer pre-fills api_base, deadline_ms or forward_conversation_id, so a
create saves only what the admin typed and the package keeps its own defaults. Setup failures
raise GuardrailSetupError from custom_guardrail, the missing-package hint covers the Docker
image, and the card copy no longer repeats the package's eligibility rules.

* test(guardrails): cover env-reference api_base in the DB poll comparison

---------

Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
collieai-dev 2026-11-12 02:05:12 +01:00 • committed by GitHub
parent 74d76fd758
commit 208a5863a5
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
16 changed files with 678 additions and 8 deletions

View file

@ -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.

View file

@ -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:

View file

@ -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,
}

View file

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

View file

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

View file

@ -163,6 +163,7 @@ class SupportedGuardrailIntegrations(Enum):
AIRIA = "airia"
WINGBACK = "wingback"
DECISION_MODEL = "decision_model"
COLLIEAI = "collieai"
class Role(Enum):

View file

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

View file

@ -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")

View file

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

View file

@ -0,0 +1,6 @@
<svg width="29" height="36" viewBox="0 0 29 36" fill="none" xmlns="http://www.w3.org/2000/svg">
<path d="M27.6737 5.65918H22.0257C16.9593 5.65918 12.8496 9.80562 12.8496 14.9173V15.2074C12.8496 19.8634 16.2629 23.7242 20.7017 24.3651C21.3981 24.4654 22.018 23.9327 22.018 23.2223V15.2074V14.9173V13.3807C25.3012 13.3807 28.0181 10.5083 28.4696 6.7788C28.5385 6.19196 28.1635 5.6669 27.666 5.6669L27.6737 5.65918Z" fill="#071E54"/>
<path d="M9.17608 9.1715C9.17608 4.50772 5.76279 0.654688 1.31633 0.0138033C0.627555 -0.0942977 0 0.446207 0 1.15659V9.1715V15.1987V23.2136C0 23.924 0.627555 24.4568 1.31633 24.3564C5.75514 23.7078 9.16842 19.8625 9.16842 15.2064V9.16378L9.17608 9.1715Z" fill="#071E54"/>
<path d="M12.795 23.499H9.23635C7.99655 23.499 7.37665 25.0124 8.25675 25.8927L10.5297 28.186C10.7976 28.4562 11.2415 28.4562 11.5093 28.186L13.7823 25.8927C14.6548 25.0124 14.0348 23.499 12.8027 23.499H12.795Z" fill="#071E54"/>
<path d="M11.3545 30.1963C11.8914 30.596 12.6123 30.8506 13.4092 30.8506V30.8584C14.1361 30.8583 14.7941 30.6414 15.3145 30.3018V31.8232C15.3145 34.132 13.4621 36.001 11.1738 36.001H11.1357V36H11.0977C8.80941 36 6.95708 34.1319 6.95703 31.8232V30.3018C7.47735 30.6414 8.13542 30.8583 8.8623 30.8584V30.8506C9.77297 30.8506 10.5847 30.5184 11.1357 30.0166L11.3545 30.1963Z" fill="#FF383C"/>
</svg>

After

Width:  |  Height:  |  Size: 1.3 KiB

View file

@ -393,4 +393,10 @@ export const GUARDRAIL_PRESETS: Record<string, GuardrailPreset> = {
mode: ["pre_call", "post_call"],
defaultOn: false,
},
collieai: {
provider: "Collieai",
guardrailNameSuggestion: "CollieAi",
mode: ["pre_call", "post_call"],
defaultOn: false,
},
};

View file

@ -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<string, string> = {
presidio: "microsoft_azure.svg",
@ -39,6 +46,7 @@ const EXPECTED_PARTNER_LOGO_FILES: Record<string, string> = {
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);
});
});

View file

@ -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];

View file

@ -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<string, string> = {
"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<string, string>;

View file

@ -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", () => {

View file

@ -36,6 +36,7 @@ const TREATMENT_BY_ASSET: Readonly<Record<string, LogoTreatment>> = {
"aiml_api.svg": "plate",
"akto.svg": "plate",
"aws.svg": "plate",
"collieai.svg": "plate",
"deepkeep.svg": "plate",
"fireworks.svg": "plate",
"llm_guard.png": "plate",