mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-08 03:08:45 +00:00
feat(guardrails): resolve default_on from env var via os.environ/ prefix
Allow guardrail configs to set default_on to an os.environ/<NAME> reference that resolves to a bool via get_secret_bool, mirroring how api_key and api_base already accept os.environ/ references in the same loader. The use case is keeping a single central config.yaml for the litellm proxy across kubernetes environments while letting each environment drive whether a guardrail is default_on through its own env var, instead of forking the config per deployment. The field's type is unchanged (Optional[bool]); plain true/false values pass through untouched. Resolution only triggers on the os.environ/ prefix, so a quoted "true" still coerces normally and a typo like ture still fails loudly at startup rather than silently disabling the guardrail. An unset referenced variable falls back to False, matching LitellmParams' existing default.
This commit is contained in:
parent
1f9323792c
commit
85c751c6b5
2 changed files with 99 additions and 1 deletions
|
|
@ -22,7 +22,7 @@ from litellm.proxy.guardrails.guardrail_hooks.grayswan import (
|
|||
from litellm.proxy.types_utils.utils import get_instance_fn
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
from litellm.repositories.table_repositories import GuardrailsRepository
|
||||
from litellm.secret_managers.main import get_secret
|
||||
from litellm.secret_managers.main import get_secret, get_secret_bool
|
||||
from litellm.types.guardrails import (
|
||||
Guardrail,
|
||||
GuardrailEventHooks,
|
||||
|
|
@ -444,6 +444,14 @@ class InMemoryGuardrailHandler:
|
|||
verbose_proxy_logger.debug("litellm_params= %s", litellm_params_data)
|
||||
|
||||
if isinstance(litellm_params_data, dict):
|
||||
default_on_raw = litellm_params_data.get("default_on")
|
||||
if isinstance(default_on_raw, str) and default_on_raw.startswith(
|
||||
"os.environ/"
|
||||
):
|
||||
litellm_params_data = {
|
||||
**litellm_params_data,
|
||||
"default_on": get_secret_bool(default_on_raw),
|
||||
}
|
||||
litellm_params = LitellmParams(**litellm_params_data)
|
||||
else:
|
||||
litellm_params = litellm_params_data
|
||||
|
|
|
|||
|
|
@ -349,3 +349,93 @@ def test_repeated_db_sync_does_not_accumulate_runner_instances():
|
|||
finally:
|
||||
for cb_list, snapshot in zip(lists, snapshots):
|
||||
cb_list[:] = snapshot
|
||||
|
||||
|
||||
def _guardrail_with_default_on(raw_default_on) -> Guardrail:
|
||||
return Guardrail(
|
||||
guardrail_id="default-on-env-var",
|
||||
guardrail_name="default-on-env-var-guard",
|
||||
litellm_params={
|
||||
"guardrail": "default_on_env_var_fake",
|
||||
"mode": "pre_call",
|
||||
"default_on": raw_default_on,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _init_handler_with_fake_initializer(captured, default_on_value):
|
||||
from unittest.mock import patch
|
||||
|
||||
from litellm.proxy.guardrails import guardrail_registry as registry_module
|
||||
|
||||
def fake_initializer(params, guardrail):
|
||||
captured["default_on"] = params.default_on
|
||||
return CustomGuardrail(
|
||||
guardrail_name=guardrail.get("guardrail_name", ""),
|
||||
event_hook=params.mode,
|
||||
default_on=params.default_on,
|
||||
)
|
||||
|
||||
with patch.dict(
|
||||
registry_module.guardrail_initializer_registry,
|
||||
{"default_on_env_var_fake": fake_initializer},
|
||||
):
|
||||
handler = registry_module.InMemoryGuardrailHandler()
|
||||
handler.initialize_guardrail(
|
||||
guardrail=_guardrail_with_default_on(default_on_value), source="config"
|
||||
)
|
||||
return handler
|
||||
|
||||
|
||||
def test_initialize_guardrail_resolves_default_on_env_var_true(monkeypatch):
|
||||
captured: dict = {}
|
||||
monkeypatch.setenv("LITELLM_TEST_DEFAULT_ON_TRUE", "true")
|
||||
handler = _init_handler_with_fake_initializer(
|
||||
captured, "os.environ/LITELLM_TEST_DEFAULT_ON_TRUE"
|
||||
)
|
||||
|
||||
assert captured["default_on"] is True
|
||||
callback = handler.guardrail_id_to_custom_guardrail["default-on-env-var"]
|
||||
assert isinstance(callback, CustomGuardrail)
|
||||
assert callback.default_on is True
|
||||
|
||||
|
||||
def test_initialize_guardrail_resolves_default_on_env_var_false(monkeypatch):
|
||||
captured: dict = {}
|
||||
monkeypatch.setenv("LITELLM_TEST_DEFAULT_ON_FALSE", "false")
|
||||
handler = _init_handler_with_fake_initializer(
|
||||
captured, "os.environ/LITELLM_TEST_DEFAULT_ON_FALSE"
|
||||
)
|
||||
|
||||
assert captured["default_on"] is False
|
||||
callback = handler.guardrail_id_to_custom_guardrail["default-on-env-var"]
|
||||
assert callback.default_on is False
|
||||
|
||||
|
||||
def test_initialize_guardrail_unsets_default_on_env_var_falls_back_false(monkeypatch):
|
||||
captured: dict = {}
|
||||
monkeypatch.delenv("LITELLM_TEST_DEFAULT_ON_MISSING", raising=False)
|
||||
handler = _init_handler_with_fake_initializer(
|
||||
captured, "os.environ/LITELLM_TEST_DEFAULT_ON_MISSING"
|
||||
)
|
||||
|
||||
assert captured["default_on"] is False
|
||||
callback = handler.guardrail_id_to_custom_guardrail["default-on-env-var"]
|
||||
assert callback.default_on is False
|
||||
|
||||
|
||||
def test_initialize_guardrail_non_prefixed_bool_string_not_env_resolved(monkeypatch):
|
||||
monkeypatch.setenv("true", "false")
|
||||
captured: dict = {}
|
||||
_init_handler_with_fake_initializer(captured, "true")
|
||||
|
||||
assert captured["default_on"] is True
|
||||
|
||||
|
||||
def test_initialize_guardrail_plain_bool_default_on_still_works():
|
||||
captured: dict = {}
|
||||
handler = _init_handler_with_fake_initializer(captured, True)
|
||||
|
||||
assert captured["default_on"] is True
|
||||
callback = handler.guardrail_id_to_custom_guardrail["default-on-env-var"]
|
||||
assert callback.default_on is True
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue