fix(proxy): guard lakera_prompt_injection callback_specific_params against non-dict

Addresses review feedback: forwarding callback_settings as callback_specific_params
(so DatadogCostManagementLogger receives cost_tag_keys) exposed the
lakera_prompt_injection branch, which did lakeraAI_Moderation(**callback_specific_params
["lakera_prompt_injection"]) with no type guard. A config like
`callback_settings: {lakera_prompt_injection: "any-string"}` then hit `**"any-string"`
-> TypeError: argument after ** must be a mapping, not str.

Guard the lakera branch with isinstance(dict), matching the existing presidio and
datadog_cost_management branches (non-dict values fall back to {}). Add a regression
test asserting initialize_callbacks_on_proxy ignores a non-dict value instead of crashing.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Hedi Daoud 2026-06-08 13:06:41 +02:00
parent 26273eff1d
commit bab6a8f208
No known key found for this signature in database
2 changed files with 55 additions and 1 deletions

View file

@ -166,7 +166,12 @@ def initialize_callbacks_on_proxy( # noqa: PLR0915
)
init_params = {}
if "lakera_prompt_injection" in callback_specific_params:
if (
"lakera_prompt_injection" in callback_specific_params
and isinstance(
callback_specific_params["lakera_prompt_injection"], dict
)
):
init_params = callback_specific_params["lakera_prompt_injection"]
lakera_moderations_object = lakeraAI_Moderation(**init_params)
imported_list.append(lakera_moderations_object)

View file

@ -309,3 +309,52 @@ def test_encrypt_callback_vars_only_encrypts_credential_fields(monkeypatch):
assert cv["langfuse_host"] == "https://cloud.langfuse.com"
assert cv["langsmith_project"] == "my-proj"
assert cv["langsmith_base_url"] == "https://smith.example"
def test_initialize_callbacks_on_proxy_lakera_ignores_non_dict_callback_settings(
monkeypatch,
):
"""Regression: a non-dict value under callback_settings.lakera_prompt_injection
must not crash initialize_callbacks_on_proxy.
Forwarding callback_settings as callback_specific_params (so callbacks like
DatadogCostManagementLogger receive their init params) exposes the lakera
branch, which previously did lakeraAI_Moderation(**callback_specific_params[
"lakera_prompt_injection"]) with no isinstance(dict) guard. For a config like
{"lakera_prompt_injection": "x"} that is `**"x"` -> TypeError: argument after
** must be a mapping, not str. The branch now guards on isinstance(dict),
matching the presidio / datadog_cost_management branches.
"""
captured = {}
class _DummyLakera:
def __init__(self, **kwargs):
captured["kwargs"] = kwargs
monkeypatch.setitem(
sys.modules,
"litellm.proxy.proxy_server",
SimpleNamespace(prisma_client=None),
)
monkeypatch.setattr(
"litellm.proxy.guardrails.guardrail_hooks.lakera_ai.lakeraAI_Moderation",
_DummyLakera,
)
original_callbacks = (
list(litellm.callbacks) if isinstance(litellm.callbacks, list) else []
)
litellm.callbacks = []
try:
# A non-dict value must be ignored (init_params stays {}), not **-unpacked.
initialize_callbacks_on_proxy(
value=["lakera_prompt_injection"],
premium_user=False,
config_file_path=".",
litellm_settings={},
callback_specific_params={"lakera_prompt_injection": "any-string"},
)
assert captured["kwargs"] == {}
assert any(isinstance(c, _DummyLakera) for c in litellm.callbacks)
finally:
litellm.callbacks = original_callbacks