scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-03 07:23:33 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 166

fix(guardrails): tolerate pass-through bodies request setup hasn't normalized

Browse source
This commit is contained in:
yucheng 2026-11-02 17:27:38 -08:00
parent 1d956d84d1
commit de63702f08
8 changed files with 160 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py
View file

@ -365,7 +365,10 @@ class GenericGuardrailAPI(CustomGuardrail):
top_level_metadata: Final = request_data.get("metadata", {})
# Merge both sources, preferring litellm_metadata if both exist
metadata_dict: Final = {**top_level_metadata, **litellm_metadata}
metadata_dict: Final = {
**(top_level_metadata if isinstance(top_level_metadata, dict) else {}),
**(litellm_metadata if isinstance(litellm_metadata, dict) else {}),
}
if not metadata_dict:
return result_metadata

5
litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py
View file

@ -219,7 +219,10 @@ class lakeraAI_Moderation(CustomGuardrail):
}
)
elif "input" in data and isinstance(data["input"], list):
text = "\n".join(data["input"])
text = "\n".join(item for item in data["input"] if isinstance(item, str))
if not text:
verbose_proxy_logger.debug("Skipping lakera prompt injection, input has no text")
return
_json_data = json.dumps(
{
"input": text,

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

@ -781,6 +781,9 @@ class LakeraAIGuardrail(CustomGuardrail):
# Use a copy of original_messages so _mask_pii_in_messages does not mutate data["messages"]
post_call_messages: Final = list(copy.deepcopy(original_messages)) + response_messages
if not post_call_messages:
verbose_proxy_logger.debug("Lakera AI: not running guardrail. No inspectable text in data")
return response
# Call Lakera guardrail
lakera_guardrail_response, _ = await self.call_v2_guard(

10
litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py
View file

@ -21,6 +21,7 @@ from litellm.integrations.custom_guardrail import (
CustomGuardrail,
log_guardrail_information,
)
from litellm.litellm_core_utils.core_helpers import as_str_mapping
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
from litellm.llms.custom_httpx.http_handler import (
@ -63,6 +64,11 @@ class _Action(str, enum.Enum):
GUARDRAIL_INTERVENED = "GUARDRAIL_INTERVENED"
def _user_api_key_alias(request_data: Mapping[str, object], metadata_key: str) -> object:
"""The key alias recorded under ``metadata_key``, which a pass-through client can set to anything."""
return (as_str_mapping(request_data.get(metadata_key)) or {}).get("user_api_key_alias")
class NomaV2Guardrail(CustomGuardrail):
reads_stream_marks: ClassVar[bool] = True
"""Noma keeps one record per streamed reply when every scan carries ``stream`` and only the last one carries the
@ -336,8 +342,8 @@ class NomaV2Guardrail(CustomGuardrail):
# (ports v1 fallback from PR #16832).
if application_id is None:
application_id = self._get_non_empty_str(
request_data.get("litellm_metadata", {}).get("user_api_key_alias")
) or self._get_non_empty_str(request_data.get("metadata", {}).get("user_api_key_alias"))
_user_api_key_alias(request_data, "litellm_metadata")
) or self._get_non_empty_str(_user_api_key_alias(request_data, "metadata"))
try:
payload: Final = self._build_scan_payload(

3
litellm/proxy/litellm_pre_call_utils.py
View file

@ -4377,11 +4377,12 @@ async def add_guardrails_from_policy_engine(
all_tags: Final = _policy_match_tags(data, metadata_variable_name)
_team_alias: Final = user_api_key_dict.team_alias
_key_alias: Final = user_api_key_dict.key_alias
_model: Final = data.get("model")
context: Final = PolicyMatchContext.model_validate(
{
"team_alias": _team_alias if isinstance(_team_alias, str) else None,
"key_alias": _key_alias if isinstance(_key_alias, str) else None,
"model": data.get("model"),
"model": _model if isinstance(_model, str) else None,
"tags": all_tags,
}
)

83
tests/unit/proxy/guardrails/guardrail_hooks/test_guardrails_on_unexpected_request_fields.py Normal file
View file

@ -0,0 +1,83 @@
"""Guardrails on built-in pass-through routes read client JSON that request setup hasn't normalized."""
from typing import Final
from unittest.mock import AsyncMock, patch
import httpx
import pytest
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api.generic_guardrail_api import (
GenericGuardrailAPI,
)
from litellm.proxy.guardrails.guardrail_hooks.lakera_ai import lakeraAI_Moderation
from litellm.proxy.guardrails.guardrail_hooks.noma.noma_v2 import NomaV2Guardrail
NOT_AN_OBJECT: Final = (None, 0, False, "", [], "text", 7)
@pytest.mark.parametrize("value", NOT_AN_OBJECT)
def test_generic_guardrail_reads_metadata_past_a_litellm_metadata_that_is_not_an_object(value: object) -> None:
guardrail: Final = GenericGuardrailAPI(api_base="http://vendor.invalid", guardrail_name="generic")
metadata: Final = guardrail._extract_user_api_key_metadata(
{"metadata": {"user_api_key_hash": "hash"}, "litellm_metadata": value}
)
assert metadata["user_api_key_hash"] == "hash"
@pytest.mark.asyncio
@pytest.mark.parametrize("value", NOT_AN_OBJECT)
async def test_noma_v2_scans_a_request_whose_litellm_metadata_is_not_an_object(value: object) -> None:
guardrail: Final = NomaV2Guardrail(api_key="key", api_base="http://vendor.invalid", guardrail_name="noma")
verdict: Final = httpx.Response(
200, json={"action": "NONE"}, request=httpx.Request("POST", "http://vendor.invalid")
)
with patch.object(guardrail.async_handler, "post", new=AsyncMock(return_value=verdict)):
result: Final = await guardrail.apply_guardrail(
inputs={"texts": ["hello"]},
request_data={"metadata": {}, "litellm_metadata": value},
input_type="request",
)
assert result["texts"] == ["hello"]
@pytest.mark.asyncio
@pytest.mark.parametrize("token_arrays", [[1, 2, 3], [[1, 2], [3, 4]]])
async def test_legacy_lakera_has_no_text_to_scan_in_token_array_embeddings(token_arrays: list[object]) -> None:
guardrail: Final = lakeraAI_Moderation(api_key="key", api_base="http://vendor.invalid")
vendor_error: Final = httpx.HTTPStatusError(
"vendor called",
request=httpx.Request("POST", "http://vendor.invalid"),
response=httpx.Response(500, text="vendor called"),
)
with patch.object(guardrail.async_handler, "post", new=AsyncMock(side_effect=vendor_error)):
await guardrail._check(
data={"input": token_arrays},
user_api_key_dict=UserAPIKeyAuth(api_key="key"),
call_type="pass_through_endpoint",
)
@pytest.mark.asyncio
async def test_legacy_lakera_still_scans_string_inputs() -> None:
guardrail: Final = lakeraAI_Moderation(api_key="key", api_base="http://vendor.invalid")
vendor_error: Final = httpx.HTTPStatusError(
"vendor called",
request=httpx.Request("POST", "http://vendor.invalid"),
response=httpx.Response(500, text="vendor called"),
)
with (
patch.object(guardrail.async_handler, "post", new=AsyncMock(side_effect=vendor_error)),
pytest.raises(Exception, match="vendor called"),
):
await guardrail._check(
data={"input": ["hello"]},
user_api_key_dict=UserAPIKeyAuth(api_key="key"),
call_type="pass_through_endpoint",
)

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

@ -2413,6 +2413,35 @@ async def test_lakera_post_call_allows_clean_content():
assert result is llm_response
@pytest.mark.asyncio
async def test_post_call_with_no_messages_and_no_reply_text_returns_the_response():
guardrail: Final = LakeraAIGuardrail(api_key="test_key")
response: Final = {"candidates": [{"content": {"role": "model", "parts": [{"text": "hi"}]}}]}
rejected: Final = HTTPException(status_code=400, detail="either content, messages, or tools must be provided")
with patch.object(guardrail, "call_v2_guard", new=AsyncMock(side_effect=rejected)):
result: Final = await guardrail.async_post_call_success_hook(
data={"contents": [{"role": "user", "parts": [{"text": "hello"}]}]},
user_api_key_dict=UserAPIKeyAuth(api_key="test_key"),
response=response,
)
assert result == response
@pytest.mark.asyncio
async def test_post_call_with_messages_still_reaches_lakera():
guardrail: Final = LakeraAIGuardrail(api_key="test_key")
rejected: Final = HTTPException(status_code=400, detail="rejected")
with (
patch.object(guardrail, "call_v2_guard", new=AsyncMock(side_effect=rejected)),
pytest.raises(HTTPException),
):
await guardrail.async_post_call_success_hook(
data={"messages": [{"role": "user", "content": "hello"}]},
user_api_key_dict=UserAPIKeyAuth(api_key="test_key"),
response={},
)
@pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown")
@pytest.mark.asyncio
async def test_lakera_post_call_masks_pii_and_allows():

27
tests/unit/proxy/test_litellm_pre_call_utils.py
View file

@ -5327,6 +5327,33 @@ async def test_embedding_header_forwarding_without_model_group_config():
litellm.model_group_settings = original_model_group_settings
@pytest.mark.asyncio
@pytest.mark.parametrize("model", [0, False, [], {}, 7, True])
async def test_policy_engine_leaves_a_model_that_is_not_a_string_alone(
model: object, monkeypatch: pytest.MonkeyPatch
) -> None:
from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry
from litellm.proxy.policy_engine.policy_registry import get_policy_registry
from litellm.types.proxy.policy_engine import Policy, PolicyAttachment, PolicyGuardrails
policy_registry = get_policy_registry()
monkeypatch.setattr(
policy_registry, "_policies", {"baseline": Policy(guardrails=PolicyGuardrails(add=["pii_blocker"]))}
)
monkeypatch.setattr(policy_registry, "_initialized", True)
attachment_registry = get_attachment_registry()
monkeypatch.setattr(attachment_registry, "_attachments", [PolicyAttachment(policy="baseline", scope="*")])
monkeypatch.setattr(attachment_registry, "_initialized", True)
data = {"model": model, "metadata": {}}
await add_guardrails_from_policy_engine(
data=data, metadata_variable_name="metadata", user_api_key_dict=UserAPIKeyAuth(api_key="test-key")
)
assert data["model"] == model
assert data["metadata"]["guardrails"] == ["pii_blocker"]
@pytest.mark.asyncio
async def test_add_guardrails_from_policy_engine():
"""