scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-12 08:42:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 207

fix(guardrails): read only the admin opt-out settings for a global Lakera guardrail on a Gemini body
Some checks are pending
Cost map pins recheck / rerun-cost-map-guard (push) Waiting to run
Details
LiteLLM Rust / rust-lint (push) Waiting to run
Details
LiteLLM Rust / rust-test (push) Waiting to run
Details
LiteLLM Rust / rust-wheel (push) Waiting to run
Details

Browse source 
The opt-out check reran should_run_guardrail on metadata alone, so a streaming-only or tag-scoped global
guardrail looked opted out and skipped the Gemini body. It now reads the key's and team's opt-out fields.
This commit is contained in:
yucheng 2026-11-02 18:05:07 -08:00
parent a00a45541f
commit 6ef3c0e056
2 changed files with 47 additions and 12 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

22
litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py
View file

@ -298,18 +298,16 @@ def _opted_out_of_global_guardrail(guardrail: "LakeraAIGuardrail", user_api_key_
"""True if the key's or team's admin metadata opts out of this ``default_on`` guardrail. A pass-through body has no
admin metadata for ``should_run_guardrail`` to read, so it is read from the key here."""
admin_fields: Final = user_api_key_dict.model_dump(include={"metadata", "team_metadata"})
return (
guardrail.default_on is True
and guardrail.should_run_guardrail( # pyright: ignore[reportUnknownMemberType] # takes an untyped data dict
data={
"metadata": {
"user_api_key_metadata": admin_fields.get("metadata"),
"user_api_key_team_metadata": admin_fields.get("team_metadata"),
}
},
event_type=GuardrailEventHooks.pre_call,
)
is not True
admin_data: Final = {
"metadata": {
"user_api_key_metadata": admin_fields.get("metadata"),
"user_api_key_team_metadata": admin_fields.get("team_metadata"),
}
}
return guardrail.default_on is True and (
guardrail.get_disable_global_guardrail(admin_data) # pyright: ignore[reportUnknownMemberType] # takes an untyped data dict
is True
or guardrail.guardrail_name in guardrail.get_opted_out_global_guardrails_from_metadata(admin_data) # pyright: ignore[reportUnknownMemberType] # takes an untyped data dict
)

37
tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py
View file

@ -1940,6 +1940,43 @@ class TestGeminiContents:
call_type="pass_through_endpoint",
)
@pytest.mark.asyncio
@pytest.mark.parametrize("call_type", ["pass_through_endpoint", "agenerate_content_stream"])
async def test_streaming_only_global_guardrail_reads_a_gemini_stream(self, call_type: str):
guardrail: Final = LakeraAIGuardrail(
api_key="test_key", guardrail_name="lakera-global", default_on=True, stream_scope="streaming"
)
with (
patch.object(guardrail, "call_v2_guard", new=AsyncMock(return_value=(GEMINI_FLAGGED, {}))),
pytest.raises(HTTPException),
):
await guardrail.async_pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(api_key="test_key"),
cache=DualCache(),
data={**_gemini_request(), "stream": True},
call_type=call_type,
)
@pytest.mark.asyncio
async def test_tag_scoped_global_guardrail_reads_a_gemini_body_with_its_tag(self):
guardrail: Final = LakeraAIGuardrail(
api_key="test_key",
guardrail_name="lakera-global",
default_on=True,
event_hook=Mode(tags={"prod": "pre_call"}, default="during_call"),
)
with (
patch.object(guardrail, "call_v2_guard", new=AsyncMock(return_value=(GEMINI_FLAGGED, {}))),
patch("litellm.proxy.proxy_server.premium_user", True),
pytest.raises(HTTPException, match="Violated"),
):
await guardrail.async_pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(api_key="test_key"),
cache=DualCache(),
data={**_gemini_request(), "metadata": {"tags": ["prod"]}},
call_type="pass_through_endpoint",
)
@pytest.mark.parametrize("chat_field", ["messages", "input"])
def test_a_pass_through_chat_body_with_stray_contents_is_read_as_chat(self, chat_field: str):
data: Final = {chat_field: [USER_MSG], "contents": [{"parts": [{"text": "stray"}]}]}