mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-10 22:41:41 +00:00
* feat(guardrails): honor Lakera v2 skip-message flags and add advisory (inject_system_message) mode Squashed rebase of bugfix/lakera-v2-skip-system-tool-messages onto latest litellm_internal_staging (900+ commits ahead; a commit-by-commit rebase hit repeated conflicts against the same files across earlier review-round commits, so the branch's cumulative diff was reapplied in one pass instead). Adds skip_system_message_in_guardrail/skip_tool_message_in_guardrail support to Lakera v2, a third on_flagged: "inject_system_message" advisory mode, and the associated masking-safety-guard hardening (multimodal content, non- maskable message fields, combined messages+input, and structured Responses- API input in advisory delivery) found across this PR's review rounds. * fix(guardrails): don't let one invalid guardrail config crash proxy boot init_guardrails_v2 had no try/except around initialize_guardrail, so a guardrail whose litellm_params fail validation at construction time (for example Lakera's on_flagged=inject_system_message combined with mode=during_call, or a malformed advisory_system_message template) raised uncaught and crashed the entire proxy at startup, taking down every other, correctly-configured guardrail in the list. Catch ValueError/TypeError per guardrail, log a warning, and skip it, matching the same pattern already used for the DB-driven guardrail-creation path in guardrail_endpoints.py. * fix(guardrails): preserve message fields and mask PII before advising in Lakera v2 Mask-in-place degraded to a hard block for any message carrying a field beyond role/content (tool_call_id, tool_calls, name, cache_control), for a message excluded by skip_system_message_in_guardrail/skip_tool_message_in_guardrail, or for a message with no inspectable text, since it rewrote data["messages"] wholesale from a synthetic role/content-only list built for the Lakera API call. That made masking effectively unusable for any real tool-calling conversation and made the skip flags flip every PII-only violation to a hard block instead of masking just the in-scope text. Replace the wholesale rewrite with a scope-index merge, reusing the same merge_guardrailed_scoped_messages helper the OpenAI/Anthropic guardrail translation handlers already use for this: patch content in place on a copy of each original message actually sent to Lakera, and leave every skipped/no-text/out-of-scope message untouched at its original position. This also fixes on_flagged="inject_system_message" (advisory mode) shipping raw unmasked PII to the model: a PII-only violation is now masked the same way regardless of on_flagged, and the advisory note is reserved for flags masking can't resolve on its own. Addresses maintainer-reported regressions on BerriAI/litellm#34940. * fix(guardrails): satisfy new lint gates for the masking/advisory fix Parameterize the write-back helper's dict param and suppress the two new lint rules that landed on the base while this branch was in flight: TQ008 (patching an internal collaborator) for two pre-existing tests unrelated to this change, and LIT001 for a param that genuinely needs to mutate the caller's request dict in place. * fix(guardrails): normalize role casing in Lakera v2 masking scope, log skipped guardrails louder Greptile finding: the masking scope helper compared roles case-sensitively while filter_messages_by_skip_flags (used to build what's actually sent to Lakera) normalizes casing, so an uppercase-cased "System"/"TOOL" role survived the scope filter but was excluded from the inspected list. The resulting length mismatch raised inside the strict positional zip, turning a maskable PII-only violation into an unhandled request failure. Lowercase the role comparison to match. Also, per veria-ai's finding that a skipped invalid guardrail now fails open: log it at error level with an explicit note that the proxy is starting without that guardrail, so it's not mistaken for routine info. * fix(guardrails): mask maskable PII in mixed violations before advising in Lakera v2 on_flagged="inject_system_message" only masked when a violation was PII-only; a mixed violation (PII plus a non-PII flag like prompt injection) fell straight through to the advisory branch with the raw PII still in place, in both async_pre_call_hook and async_moderation_hook. Mask whatever Lakera returned location data for before appending or logging the advisory, so a mixed violation never ships raw PII just because something else was also flagged. Also degrade to blocking, same as block mode already does, when nothing can be safely masked at all (multimodal content, or messages combined with a Responses API input field) instead of showing an advisory note next to raw, unredacted content. Widened call_v2_guard/_mask_pii_in_messages/the write-back helper's message parameters from list to Sequence to match what's actually passed through from _filter_skipped_messages, instead of duplicating list(...) casts at every call site. * fix(guardrails): don't hard-block advisory mode for non-PII flags on non-maskable input Bugbot finding: gating the entire inject_system_message branch on is_multimodal_input hard-blocked every flagged request on Responses instructions, combined messages+input, or multimodal content, including a prompt-injection-only violation with no PII at all. Masking safety only matters when there's actual PII to mask; a violation with no PII needs no masking, so the advisory should still be delivered normally. Only degrade to blocking when the breakdown actually contains a PII detection and masking isn't safely possible. Otherwise, mask whatever's maskable (if any) and deliver the advisory as before. * fix(guardrails): require payload and breakdown for Lakera v2 advisory mode Advisory mode's mixed-violation masking safety net can only redact detected PII when Lakera's response carries both the breakdown (to detect a PII hit at all) and payload (the location data to mask by). payload=False or breakdown=False alongside on_flagged='inject_system_message' silently forwarded raw PII next to the advisory note. Reject that combination at construction and hot-reload time instead. * fix(guardrails): skip_system_message_in_guardrail must not force-block Lakera masking _has_responses_instructions treated any non-empty data["instructions"] as unsafe to mask regardless of skip_system_message_in_guardrail, even though that flag excludes the instructions-derived synthetic system message from what Lakera ever inspects. PII detected purely in the maskable non-system content was force-blocked instead of masked. Also fixes pre-existing LIT010 (missing Final) violations in _has_responses_instructions, _breakdown_has_pii_violation, and async_post_call_success_hook that the rebase's lowered budget ceiling now flags. * chore: retrigger CI (GitHub Actions runner-acquisition failure on prior push) * fix(guardrails): address maintainer review findings on Lakera v2 advisory mode - Gate advisory_system_message template validation on on_flagged= 'inject_system_message', since block/monitor mode never reads it. - Allow on_flagged='inject_system_message' with mode='during_call' at construction/hot-reload instead of rejecting it; async_moderation_hook already degrades gracefully (masks if possible, else logs a warning). - reinitialize_guardrail now restores the previous live instance when the new config fails to initialize, instead of leaving the guardrail deleted entirely with nothing enforcing it. - PATCH /guardrails/{id} rolls back the DB write and returns 422 when the in-memory sync rejects the new config, instead of persisting a config that never actually took effect and returning 200. - Qualifire now rejects on_flagged values it doesn't implement (only Lakera should accept 'inject_system_message'; LitellmParams flattens the field across every guardrail config mixin). * fix(tests): satisfy lint gates and update collateral test for advisory-mode fixes - Add match= to a too-broad pytest.raises(ValueError), and suppress the new TQ008 mocker.patch findings (same pattern already used by sibling scenarios in this test). - test_init_guardrails_v2_skips_invalid_guardrail_instead_of_crashing_boot used mode='during_call' + on_flagged='inject_system_message' as its invalid-config example; that combination is now accepted, so swap in the payload/breakdown-missing case and add a test confirming during_call advisory mode constructs successfully. * docs(CLAUDE.md): auto-capture review learnings without being asked This session found three real bugs a human maintainer caught after eight rounds of bot review and live-proxy verification all missed them. Add a standing instruction to write learnings.md entries the moment a root cause is understood, in both the repo-wide file and any relevant skill's own file, instead of relying on being asked. * feat(guardrails): add scan_raw_request flag so YAML order can't change enforcement Maintainer finding on BerriAI/litellm#34940: guardrails for the same hook run sequentially over one shared, progressively-mutated request dict, so declaring a masking guardrail before a blocking one hides the violation from it (200 vs 400 depending purely on YAML order). scan_raw_request opts a guardrail into always evaluating a snapshot taken before any guardrail in the hook ran, regardless of its declared position. Same contract as run_in_parallel: block-only, its own mutations discarded. Verified live: real proxy, real Gemini call, two custom guardrails (a redactor then a blocker). Same request, same declared order -- without the flag the blocker never sees the raw secret (200); with it, the blocker correctly rejects before any provider call (400). * fix(guardrails): harden scan_raw_request against review findings - Use safe_deep_copy instead of a bare deepcopy for the raw-request snapshot; request payloads commonly carry unpicklable objects (e.g. an otel span in metadata), which previously raised on every guarded request when tracing was enabled (Bugbot, High). - Only compute the snapshot when a guardrail actually opted in, and take it before _maybe_execute_pipelines runs, so a pipeline-mutated payload can't hide a violation from a scan_raw_request guardrail outside the pipeline (veria-ai). - Log a warning when a scan_raw_request guardrail returns a modified payload, since that mutation is discarded and the combination is otherwise silently exploitable for a masking-capable integration misconfigured this way (veria-ai). * chore(openapi): regenerate lazy snapshot and dashboard schema types The lazy OpenAPI snapshot (litellm/proxy/_lazy_openapi_snapshot.json) and the derived dashboard schema.d.ts had drifted stale relative to the guardrail config model changes across this PR's rounds (advisory mode, scan_raw_request, and upstream additions picked up by rebasing). Regenerated via the CI's own documented fix: uv run python -m litellm.proxy._lazy_openapi_snapshot npm run gen:api (via make check) * chore(openapi): pick up cache_hit_filter field after rebase * fix(guardrails): stop scan_raw_request warning from firing on every call _process_guardrail_callback always returns a dict once a guardrail runs (mark_pre_call_hook_ran unconditionally stamps bookkeeping metadata), so comparing the result to non-None warned on every request even when the guardrail never touched the payload. Compare against a bookkeeping-only baseline instead, so only an actual content mutation triggers the warning. * fix(guardrails): make scan_raw_request snapshots independent of safe_memory_mode safe_deep_copy can return the original object under litellm.safe_memory_mode, or alias a per-key reference on copy failure. Under that mode, the scan_raw_request comparison baseline aliased raw_request_snapshot (and therefore the live request), letting mark_pre_call_hook_ran write a premature execution marker that a deployment-level guardrail sharing the same name would read as "already ran" and skip. Also affected the feature's core isolation guarantee: input_data itself could alias the live request under the same mode. Replace every scan_raw_request snapshot with _independent_snapshot, which never returns an alias, only a genuine copy or None. * fix(guardrails): gate during_call mixed-violation masking behind an actual PII check The during_call branch for a mixed violation under on_flagged=inject_system_message unconditionally masked and reassigned data["messages"], even for a pure prompt-injection violation with zero PII, unlike async_pre_call_hook which already gates the same call behind _breakdown_has_pii_violation. The unconditional reassignment touched shared request state during a hook documented as racing with the concurrent LLM dispatch, for no reason when there was nothing to mask. * fix(guardrails): stop scan_raw_request from silently no-op'ing on real requests _independent_snapshot did one whole-dict copy.deepcopy and returned None on any failure. Every real proxy request carries data["litellm_logging_obj"] (a Logging instance nesting a live OTel span with a real lock) by the time pre_call_hook runs, which can never be deep-copied, so the snapshot failed on every real request and silently fell back to the live, unisolated data with no warning -- defeating the entire feature in production while every existing test (none of which set litellm_logging_obj) kept passing. Rework the helper to deep-copy each top-level key independently, falling back to the original reference only for the specific key that fails, same crash tolerance as safe_deep_copy's own per-key fallback. It never returns None now; only the keys scan_raw_request actually depends on (messages/ input, metadata/litellm_metadata) need to be genuinely independent. * fix(guardrails): block during_call when PII can't be safely masked Greptile finding (P1, security): async_moderation_hook's inject_system_message branch had no equivalent to async_pre_call_hook's degrade-to-blocking case for a PII violation on input that can't be safely masked (e.g. combined messages+input). It fell through to the advisory no-op branch and let raw, unredacted PII reach the model with no protection at all. Raising still blocks the response from reaching the caller even though during_call races with the LLM dispatch, the same mechanism on_flagged="block" already relies on for this hook, so add the same block-instead-of-advisory branch pre_call already has. * chore(lint): fix LIT002 ceiling after rebase merge conflict resolution * fix(lint): suppress genuine LIT002 hits instead of padding the ceiling My earlier rebase conflict resolution for type-discipline-budget.json's LIT002 limit was too low, then overcorrected by padding it well above the actual measured count. Root-caused instead: _independent_snapshot and the PATCH-endpoint rollback path legitimately construct plain, mutable request-payload/config dicts (matching this file's existing precedent for the same shape), so suppress those four sites with `# mutable-ok:` rather than reshaping code that must stay a plain dict by contract. Set the limit to the exact current measured total; the small remaining gap vs upstream's own committed ceiling is pre-existing drift in litellm_internal_staging itself (its own tree already measures over its committed limit), not attributable to this PR. * fix(guardrails): stamp live request when a scan_raw_request guardrail runs _run_sequential_guardrail_callback and _run_parallel_pre_call_guardrails only called mark_pre_call_hook_ran on throwaway snapshot copies for a scan_raw_request guardrail, never on the live request returned to the caller. A later async_pre_call_deployment_hook (router-level guardrail re-check) reads that marker on live kwargs to decide whether to skip re-running the same guardrail; since it was never stamped there, the guardrail ran a second time on live data, doubling the external call and re-applying whatever scan_raw_request's contract says should be discarded. * fix(guardrails): revalidate Qualifire's on_flagged on live config reload on_flagged was validated only in __init__. The base CustomGuardrail.update_in_memory_litellm_params is a generic setattr loop with no revalidation, so a live config update (PUT /guardrails/{id}, no restart) could setattr on_flagged="inject_system_message" onto a running instance, bypassing the constructor's rejection -- silently blocking every flagged request under an "advisory" label. Mirrors LakeraAIGuardrail's own update_in_memory_litellm_params override added earlier in this PR. * fix(guardrails): honor scan_raw_request for pipeline-managed guardrails A scan_raw_request=True guardrail that is itself a pipeline step never saw raw_request_snapshot: PipelineExecutor.execute_steps had no way to receive it, and pipeline-managed guardrails are fully excluded from the normal sequential/parallel loops that implement the flag. Such a guardrail silently evaluated whatever an earlier pass_data step in the same pipeline had already rewritten, defeating the flag for pipeline-managed guardrails. Moves the snapshot helper (renamed independent_snapshot) from proxy/utils.py to litellm_core_utils/core_helpers.py so pipeline_executor.py can use the same independent-copy logic without a circular import, threads raw_request_snapshot through _maybe_execute_pipelines and PipelineExecutor.execute_steps/_run_step, and discards a scan_raw_request step's returned data the same way the sequential/parallel loops already do. * chore(openapi): pick up upstream drift after rebase onto litellm_internal_staging * fix(guardrails): stop attempting PII masking during during_call in Lakera v2 Greptile finding (P1, security): during_call runs concurrently with the LLM dispatch. In the common path, the provider call already binds its messages kwarg before this guardrail's coroutine gets a chance to run, let alone before its own network round trip to Lakera completes -- masking here can never reliably reach the outgoing request, and _apply_redacted_messages_back_ preserving_fields reassigns to a new list object rather than mutating in place, so even winning the race wouldn't help. This affected both the PII-only and mixed-violation masking branches, all added in this same PR. Remove masking from async_moderation_hook entirely and let PII violations fall through to the normal on_flagged branching: block under "block" or "inject_system_message" (extending the existing multimodal-only block to cover every PII case, since masking is proven non-functional regardless of input shape), log-and-allow under "monitor" -- consistent with how every other violation type in this hook is already handled. --------- Co-authored-by: Deepanshu <deepanshu.lulla@alpha-sense.com>
811 lines
28 KiB
Python
811 lines
28 KiB
Python
"""
|
|
Regression tests for guardrail-coverage gaps.
|
|
|
|
Each test confirms that a previously-bypassable input shape now triggers
|
|
inspection by the relevant guardrail hook:
|
|
|
|
- VERIA-11: multimodal list-format ``content`` is inspected (no longer
|
|
silently skipped because of an ``isinstance(content, str)`` check).
|
|
- fniVO9-F: Responses-API ``data["input"]`` is inspected (no longer
|
|
silently skipped because the hook only looked at ``data["messages"]``).
|
|
- yVS0wMDO: Aim's post-call hook inspects every choice when ``n>1``,
|
|
not just ``choices[0]``.
|
|
"""
|
|
|
|
from typing import Any, Dict
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import pytest
|
|
from httpx import Request, Response
|
|
|
|
from litellm import DualCache
|
|
from litellm.proxy._types import UserAPIKeyAuth
|
|
from litellm.types.utils import Choices, Message, ModelResponse
|
|
|
|
|
|
@pytest.fixture
|
|
def user_api_key():
|
|
return UserAPIKeyAuth(api_key="hashed", user_id="u", key_alias=None)
|
|
|
|
|
|
# ── Aim ───────────────────────────────────────────────────────────────────────
|
|
|
|
|
|
def _aim_no_action_response() -> Response:
|
|
return Response(
|
|
status_code=200,
|
|
json={"required_action": None},
|
|
request=Request("POST", "https://api.aim.security/fw/v1/analyze"),
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_aim_inspects_multimodal_list_content(user_api_key, monkeypatch):
|
|
monkeypatch.setenv("AIM_API_KEY", "hs-aim-key")
|
|
from litellm.proxy.guardrails.guardrail_hooks.aim.aim import AimGuardrail
|
|
|
|
guard = AimGuardrail()
|
|
sent_payload: Dict[str, Any] = {}
|
|
|
|
async def capture(url, headers, json):
|
|
sent_payload.update(json)
|
|
return _aim_no_action_response()
|
|
|
|
with patch.object(guard.async_handler, "post", side_effect=capture):
|
|
await guard.async_pre_call_hook(
|
|
user_api_key_dict=user_api_key,
|
|
cache=DualCache(),
|
|
data={
|
|
"messages": [
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{"type": "text", "text": "secret payload"},
|
|
{"type": "image_url", "image_url": {"url": "..."}},
|
|
],
|
|
}
|
|
]
|
|
},
|
|
call_type="acompletion",
|
|
)
|
|
|
|
# The multimodal text part must be visible to Aim.
|
|
assert sent_payload["messages"] == [{"role": "user", "content": "secret payload"}]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_aim_inspects_responses_api_input(user_api_key, monkeypatch):
|
|
monkeypatch.setenv("AIM_API_KEY", "hs-aim-key")
|
|
from litellm.proxy.guardrails.guardrail_hooks.aim.aim import AimGuardrail
|
|
|
|
guard = AimGuardrail()
|
|
sent_payload: Dict[str, Any] = {}
|
|
|
|
async def capture(url, headers, json):
|
|
sent_payload.update(json)
|
|
return _aim_no_action_response()
|
|
|
|
with patch.object(guard.async_handler, "post", side_effect=capture):
|
|
await guard.async_pre_call_hook(
|
|
user_api_key_dict=user_api_key,
|
|
cache=DualCache(),
|
|
data={"input": "responses-api content"},
|
|
call_type="acompletion",
|
|
)
|
|
|
|
assert sent_payload["messages"] == [
|
|
{"role": "user", "content": "responses-api content"}
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_aim_post_call_inspects_all_choices(user_api_key, monkeypatch):
|
|
"""yVS0wMDO: ``n>1`` no longer bypasses Aim by hiding violations in
|
|
``choices[1+]``."""
|
|
monkeypatch.setenv("AIM_API_KEY", "hs-aim-key")
|
|
from litellm.proxy.guardrails.guardrail_hooks.aim.aim import AimGuardrail
|
|
|
|
guard = AimGuardrail()
|
|
inspected_outputs = []
|
|
|
|
async def capture(request_data, output, hook, key_alias):
|
|
inspected_outputs.append(output)
|
|
return {"redacted_output": output}
|
|
|
|
response = ModelResponse(
|
|
choices=[
|
|
Choices(index=0, message=Message(role="assistant", content="first")),
|
|
Choices(index=1, message=Message(role="assistant", content="second")),
|
|
Choices(index=2, message=Message(role="assistant", content="third")),
|
|
]
|
|
)
|
|
|
|
with patch.object(guard, "call_aim_guardrail_on_output", side_effect=capture):
|
|
await guard.async_post_call_success_hook(
|
|
data={"messages": [{"role": "user", "content": "hi"}]},
|
|
user_api_key_dict=user_api_key,
|
|
response=response,
|
|
)
|
|
|
|
# ``asyncio.gather`` is used for parallelism, so order of inspection is
|
|
# not guaranteed.
|
|
assert sorted(inspected_outputs) == ["first", "second", "third"]
|
|
|
|
|
|
# ── Lakera v2 ─────────────────────────────────────────────────────────────────
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_lakera_v2_inspects_responses_api_input(user_api_key, monkeypatch):
|
|
monkeypatch.setenv("LAKERA_API_KEY", "lk-test")
|
|
from litellm.proxy.guardrails.guardrail_hooks.lakera_ai_v2 import (
|
|
LakeraAIGuardrail,
|
|
)
|
|
|
|
guard = LakeraAIGuardrail(api_key="lk-test", on_flagged="monitor")
|
|
|
|
seen_messages = []
|
|
|
|
async def fake_call_v2_guard(messages, request_data, event_type):
|
|
seen_messages.append(messages)
|
|
return {"flagged": False}, {}
|
|
|
|
with patch.object(guard, "call_v2_guard", side_effect=fake_call_v2_guard):
|
|
await guard.async_pre_call_hook(
|
|
user_api_key_dict=user_api_key,
|
|
cache=DualCache(),
|
|
data={"input": "responses-api content"},
|
|
call_type="responses",
|
|
)
|
|
|
|
assert seen_messages == [({"role": "user", "content": "responses-api content"},)]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_lakera_v2_responses_api_input_redacted_writeback(
|
|
user_api_key, monkeypatch
|
|
):
|
|
"""Greptile P1: when input arrives via Responses-API ``data["input"]``
|
|
(string) and Lakera flags PII, the redacted content must be written
|
|
back to ``data["input"]`` — the Responses-API backend reads from
|
|
``input``, so writing only to ``messages`` would let unredacted PII
|
|
reach the LLM."""
|
|
monkeypatch.setenv("LAKERA_API_KEY", "lk-test")
|
|
from litellm.proxy.guardrails.guardrail_hooks.lakera_ai_v2 import (
|
|
LakeraAIGuardrail,
|
|
)
|
|
|
|
guard = LakeraAIGuardrail(api_key="lk-test", on_flagged="block")
|
|
|
|
async def fake_call_v2_guard(messages, request_data, event_type):
|
|
return ({"flagged": True, "payload": []}, {"EMAIL": 1})
|
|
|
|
def fake_mask(messages, lakera_response, masked_entity_count):
|
|
return [{"role": "user", "content": "[REDACTED EMAIL]"}]
|
|
|
|
with (
|
|
patch.object(guard, "call_v2_guard", side_effect=fake_call_v2_guard),
|
|
patch.object(guard, "_is_only_pii_violation", return_value=True),
|
|
patch.object(guard, "_mask_pii_in_messages", side_effect=fake_mask),
|
|
):
|
|
data = {"input": "user@example.com leaked"}
|
|
await guard.async_pre_call_hook(
|
|
user_api_key_dict=user_api_key,
|
|
cache=DualCache(),
|
|
data=data,
|
|
call_type="responses",
|
|
)
|
|
|
|
assert data["input"] == "[REDACTED EMAIL]"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_aim_responses_api_input_anonymize_writeback(user_api_key, monkeypatch):
|
|
"""Greptile P1: Aim's anonymize action must redact ``data["input"]``
|
|
for Responses-API requests, not just ``data["messages"]``."""
|
|
monkeypatch.setenv("AIM_API_KEY", "hs-aim-key")
|
|
from litellm.proxy.guardrails.guardrail_hooks.aim.aim import AimGuardrail
|
|
|
|
guard = AimGuardrail()
|
|
|
|
aim_response_body = {
|
|
"required_action": {"action_type": "anonymize_action"},
|
|
"redacted_chat": {
|
|
"all_redacted_messages": [
|
|
{"role": "user", "content": "[REDACTED] anonymised"}
|
|
]
|
|
},
|
|
}
|
|
|
|
async def capture(url, headers, json):
|
|
return Response(
|
|
status_code=200,
|
|
json=aim_response_body,
|
|
request=Request("POST", "https://api.aim.security/fw/v1/analyze"),
|
|
)
|
|
|
|
with patch.object(guard.async_handler, "post", side_effect=capture):
|
|
data = {"input": "user@example.com leaked"}
|
|
await guard.async_pre_call_hook(
|
|
user_api_key_dict=user_api_key,
|
|
cache=DualCache(),
|
|
data=data,
|
|
call_type="responses",
|
|
)
|
|
|
|
assert data["input"] == "[REDACTED] anonymised"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_lakera_v2_multimodal_pii_degrades_to_block(user_api_key, monkeypatch):
|
|
"""Mask-in-place uses Lakera offsets and cannot preserve image/audio
|
|
parts of multimodal input. When PII is detected on a multimodal
|
|
request, the hook must raise the block exception instead of silently
|
|
flattening ``data["messages"]`` to text-only."""
|
|
monkeypatch.setenv("LAKERA_API_KEY", "lk-test")
|
|
from fastapi import HTTPException
|
|
|
|
from litellm.proxy.guardrails.guardrail_hooks.lakera_ai_v2 import (
|
|
LakeraAIGuardrail,
|
|
)
|
|
|
|
guard = LakeraAIGuardrail(api_key="lk-test", on_flagged="block")
|
|
|
|
async def fake_call_v2_guard(messages, request_data, event_type):
|
|
return (
|
|
{
|
|
"flagged": True,
|
|
"payload": [{"detector_type": "pii/email", "start": 0, "end": 5}],
|
|
},
|
|
{"EMAIL": 1},
|
|
)
|
|
|
|
with (
|
|
patch.object(guard, "call_v2_guard", side_effect=fake_call_v2_guard),
|
|
patch.object(guard, "_is_only_pii_violation", return_value=True),
|
|
patch.object(
|
|
guard,
|
|
"_get_http_exception_for_blocked_guardrail",
|
|
return_value=HTTPException(status_code=400, detail="blocked"),
|
|
),
|
|
):
|
|
with pytest.raises(HTTPException):
|
|
await guard.async_pre_call_hook(
|
|
user_api_key_dict=user_api_key,
|
|
cache=DualCache(),
|
|
data={
|
|
"messages": [
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{"type": "text", "text": "leak"},
|
|
{"type": "image_url", "image_url": {"url": "..."}},
|
|
],
|
|
}
|
|
]
|
|
},
|
|
call_type="acompletion",
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_lakera_v2_inspects_multimodal_list_content(user_api_key, monkeypatch):
|
|
monkeypatch.setenv("LAKERA_API_KEY", "lk-test")
|
|
from litellm.proxy.guardrails.guardrail_hooks.lakera_ai_v2 import (
|
|
LakeraAIGuardrail,
|
|
)
|
|
|
|
guard = LakeraAIGuardrail(api_key="lk-test", on_flagged="monitor")
|
|
seen_messages = []
|
|
|
|
async def fake_call_v2_guard(messages, request_data, event_type):
|
|
seen_messages.append(messages)
|
|
return {"flagged": False}, {}
|
|
|
|
with patch.object(guard, "call_v2_guard", side_effect=fake_call_v2_guard):
|
|
await guard.async_pre_call_hook(
|
|
user_api_key_dict=user_api_key,
|
|
cache=DualCache(),
|
|
data={
|
|
"messages": [
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{"type": "text", "text": "AKIAEXAMPLE"},
|
|
{"type": "image_url", "image_url": {"url": "..."}},
|
|
],
|
|
}
|
|
]
|
|
},
|
|
call_type="acompletion",
|
|
)
|
|
|
|
assert seen_messages == [({"role": "user", "content": "AKIAEXAMPLE"},)]
|
|
|
|
|
|
# ── Lasso ─────────────────────────────────────────────────────────────────────
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_lasso_multimodal_falls_back_to_classify(user_api_key, monkeypatch):
|
|
"""Lasso's classifix (mask) endpoint returns text that overwrites
|
|
``data["messages"]``. For multimodal input that would silently strip
|
|
image parts — the hook must use the classify endpoint instead and
|
|
leave the original payload intact."""
|
|
monkeypatch.setenv("LASSO_API_KEY", "ls-test")
|
|
from litellm.proxy.guardrails.guardrail_hooks.lasso.lasso import LassoGuardrail
|
|
|
|
guard = LassoGuardrail(lasso_api_key="ls-test", mask=True)
|
|
|
|
masking_called = False
|
|
classify_called = False
|
|
|
|
async def fake_masking(data, cache, message_type, messages):
|
|
nonlocal masking_called
|
|
masking_called = True
|
|
return data
|
|
|
|
async def fake_classification(data, cache, message_type, messages):
|
|
nonlocal classify_called
|
|
classify_called = True
|
|
return data
|
|
|
|
with (
|
|
patch.object(guard, "_handle_masking", side_effect=fake_masking),
|
|
patch.object(guard, "_handle_classification", side_effect=fake_classification),
|
|
):
|
|
await guard._run_lasso_guardrail(
|
|
data={
|
|
"messages": [
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{"type": "text", "text": "hello"},
|
|
{"type": "image_url", "image_url": {"url": "..."}},
|
|
],
|
|
}
|
|
]
|
|
},
|
|
cache=DualCache(),
|
|
message_type="PROMPT",
|
|
)
|
|
|
|
assert classify_called is True
|
|
assert masking_called is False
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_lasso_inspects_responses_api_input(user_api_key, monkeypatch):
|
|
monkeypatch.setenv("LASSO_API_KEY", "ls-test")
|
|
from litellm.proxy.guardrails.guardrail_hooks.lasso.lasso import LassoGuardrail
|
|
|
|
guard = LassoGuardrail(lasso_api_key="ls-test")
|
|
|
|
seen_messages = []
|
|
|
|
async def fake_handle_classification(data, cache, message_type, messages):
|
|
seen_messages.append(messages)
|
|
return data
|
|
|
|
with patch.object(
|
|
guard, "_handle_classification", side_effect=fake_handle_classification
|
|
):
|
|
await guard._run_lasso_guardrail(
|
|
data={"input": "responses-api content"},
|
|
cache=DualCache(),
|
|
message_type="PROMPT",
|
|
)
|
|
|
|
assert seen_messages == [[{"role": "user", "content": "responses-api content"}]]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_lasso_masking_writes_back_responses_api_input(user_api_key, monkeypatch):
|
|
"""Krrish blocker: Lasso classifix masking must update ``data["input"]``
|
|
for Responses-API requests, not only ``data["messages"]``."""
|
|
monkeypatch.setenv("LASSO_API_KEY", "ls-test")
|
|
from litellm.proxy.guardrails.guardrail_hooks.lasso.lasso import LassoGuardrail
|
|
|
|
guard = LassoGuardrail(lasso_api_key="ls-test", mask=True)
|
|
lasso_response = {
|
|
"violations_detected": True,
|
|
"deputies": {"pii": True},
|
|
"findings": {"pii": [{"action": "AUTO_MASKING"}]},
|
|
"messages": [{"role": "user", "content": "[REDACTED]"}],
|
|
}
|
|
|
|
async def fake_call_lasso_api(headers, payload, api_url=None):
|
|
return lasso_response
|
|
|
|
data = {"input": "user@example.com leaked"}
|
|
|
|
with patch.object(guard, "_call_lasso_api", side_effect=fake_call_lasso_api):
|
|
await guard._run_lasso_guardrail(
|
|
data=data,
|
|
cache=DualCache(),
|
|
message_type="PROMPT",
|
|
)
|
|
|
|
assert data["input"] == "[REDACTED]"
|
|
|
|
|
|
# ── Banned Keywords ───────────────────────────────────────────────────────────
|
|
|
|
|
|
def test_banned_keywords_blocks_multimodal_content(monkeypatch):
|
|
"""VERIA-11: a banned word hidden in a multimodal text part is now caught.
|
|
|
|
Uses ``acompletion`` — the value the proxy ingress actually passes
|
|
for ``/v1/chat/completions``. Asserting against the literal sync
|
|
``"completion"`` would pass even if the hook's call-type gate were
|
|
misaligned with the runtime, so the test wouldn't catch regressions.
|
|
"""
|
|
monkeypatch.setattr("litellm.banned_keywords_list", ["forbidden"], raising=False)
|
|
from enterprise.enterprise_hooks.banned_keywords import _ENTERPRISE_BannedKeywords
|
|
from fastapi import HTTPException
|
|
|
|
guard = _ENTERPRISE_BannedKeywords()
|
|
|
|
async def _run():
|
|
await guard.async_pre_call_hook(
|
|
user_api_key_dict=UserAPIKeyAuth(api_key="hashed", user_id="u"),
|
|
cache=DualCache(),
|
|
data={
|
|
"messages": [
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{"type": "text", "text": "forbidden word here"},
|
|
{"type": "image_url", "image_url": {"url": "..."}},
|
|
],
|
|
}
|
|
]
|
|
},
|
|
call_type="acompletion",
|
|
)
|
|
|
|
import asyncio
|
|
|
|
with pytest.raises(HTTPException) as exc:
|
|
asyncio.run(_run())
|
|
assert "forbidden" in str(exc.value.detail).lower()
|
|
|
|
|
|
def test_banned_keywords_blocks_responses_api_input(monkeypatch):
|
|
monkeypatch.setattr("litellm.banned_keywords_list", ["forbidden"], raising=False)
|
|
from enterprise.enterprise_hooks.banned_keywords import _ENTERPRISE_BannedKeywords
|
|
from fastapi import HTTPException
|
|
|
|
guard = _ENTERPRISE_BannedKeywords()
|
|
|
|
async def _run():
|
|
await guard.async_pre_call_hook(
|
|
user_api_key_dict=UserAPIKeyAuth(api_key="hashed", user_id="u"),
|
|
cache=DualCache(),
|
|
data={"input": "this contains forbidden content"},
|
|
call_type="aresponses",
|
|
)
|
|
|
|
import asyncio
|
|
|
|
with pytest.raises(HTTPException):
|
|
asyncio.run(_run())
|
|
|
|
|
|
@pytest.mark.parametrize("call_type", ["completion", "acompletion", "aresponses"])
|
|
def test_banned_keywords_fires_on_text_content_call_types(monkeypatch, call_type):
|
|
"""Locks the call-type gate to the runtime ``route_type`` values the
|
|
proxy actually emits — pinning a regression where the hook had
|
|
``call_type == "completion"`` and silently no-op'd both
|
|
``acompletion`` (chat completions) and ``aresponses`` (Responses API).
|
|
"""
|
|
monkeypatch.setattr("litellm.banned_keywords_list", ["forbidden"], raising=False)
|
|
from enterprise.enterprise_hooks.banned_keywords import _ENTERPRISE_BannedKeywords
|
|
from fastapi import HTTPException
|
|
|
|
guard = _ENTERPRISE_BannedKeywords()
|
|
|
|
import asyncio
|
|
|
|
with pytest.raises(HTTPException):
|
|
asyncio.run(
|
|
guard.async_pre_call_hook(
|
|
user_api_key_dict=UserAPIKeyAuth(api_key="hashed", user_id="u"),
|
|
cache=DualCache(),
|
|
data={
|
|
"messages": [{"role": "user", "content": "forbidden text"}],
|
|
"input": "forbidden text",
|
|
},
|
|
call_type=call_type,
|
|
)
|
|
)
|
|
|
|
|
|
def test_banned_keywords_skips_non_text_call_types(monkeypatch):
|
|
"""Embedding / moderation / audio paths don't carry chat text and
|
|
aren't in the text-guardrail scope. They must not trigger the hook
|
|
even when the request body otherwise looks like a chat payload.
|
|
"""
|
|
monkeypatch.setattr("litellm.banned_keywords_list", ["forbidden"], raising=False)
|
|
from enterprise.enterprise_hooks.banned_keywords import _ENTERPRISE_BannedKeywords
|
|
|
|
guard = _ENTERPRISE_BannedKeywords()
|
|
|
|
import asyncio
|
|
|
|
for call_type in ("aembedding", "amoderation", "aspeech", "atranscription"):
|
|
# Should return without raising, even though the data carries the banned word.
|
|
asyncio.run(
|
|
guard.async_pre_call_hook(
|
|
user_api_key_dict=UserAPIKeyAuth(api_key="hashed", user_id="u"),
|
|
cache=DualCache(),
|
|
data={"input": "forbidden text"},
|
|
call_type=call_type,
|
|
)
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_banned_keywords_post_call_checks_all_choices(monkeypatch, user_api_key):
|
|
"""Krrish blocker: ``n>1`` responses must not bypass post-call checks by
|
|
placing the banned text in ``choices[1+]``."""
|
|
monkeypatch.setattr("litellm.banned_keywords_list", ["forbidden"], raising=False)
|
|
from enterprise.enterprise_hooks.banned_keywords import _ENTERPRISE_BannedKeywords
|
|
from fastapi import HTTPException
|
|
|
|
guard = _ENTERPRISE_BannedKeywords()
|
|
response = ModelResponse(
|
|
choices=[
|
|
Choices(index=0, message=Message(role="assistant", content="clean")),
|
|
Choices(index=1, message=Message(role="assistant", content="forbidden")),
|
|
]
|
|
)
|
|
|
|
with pytest.raises(HTTPException) as exc:
|
|
await guard.async_post_call_success_hook(
|
|
data={},
|
|
user_api_key_dict=user_api_key,
|
|
response=response,
|
|
)
|
|
|
|
assert "forbidden" in str(exc.value.detail).lower()
|
|
|
|
|
|
# ── Azure Content Safety ──────────────────────────────────────────────────────
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
"call_type, data",
|
|
[
|
|
(
|
|
"acompletion",
|
|
{
|
|
"messages": [
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{"type": "text", "text": "scan me"},
|
|
{"type": "image_url", "image_url": {"url": "..."}},
|
|
],
|
|
}
|
|
]
|
|
},
|
|
),
|
|
("aresponses", {"input": "scan me"}),
|
|
],
|
|
)
|
|
async def test_azure_content_safety_pre_call_fires_on_runtime_call_types(
|
|
user_api_key, call_type, data
|
|
):
|
|
"""The proxy ingress passes ``route_type`` straight through as
|
|
``call_type`` — ``acompletion`` for chat completions and
|
|
``aresponses`` for the Responses API. The hook must inspect text
|
|
fragments under both, not only the literal ``"completion"`` string
|
|
used by some SDK callers."""
|
|
from litellm.proxy.hooks.azure_content_safety import _PROXY_AzureContentSafety
|
|
|
|
guard = _PROXY_AzureContentSafety.__new__(_PROXY_AzureContentSafety)
|
|
seen = []
|
|
|
|
async def fake_test_violation(content, source=None):
|
|
seen.append((content, source))
|
|
|
|
guard.test_violation = fake_test_violation
|
|
await guard.async_pre_call_hook(
|
|
user_api_key_dict=user_api_key,
|
|
cache=DualCache(),
|
|
data=data,
|
|
call_type=call_type,
|
|
)
|
|
assert ("scan me", "input") in seen
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_azure_content_safety_post_call_checks_all_choices(user_api_key):
|
|
"""Krrish blocker: ``n>1`` responses must not bypass Azure Content Safety
|
|
by placing the unsafe text in ``choices[1+]``."""
|
|
from fastapi import HTTPException
|
|
from litellm.proxy.hooks.azure_content_safety import _PROXY_AzureContentSafety
|
|
|
|
guard = _PROXY_AzureContentSafety.__new__(_PROXY_AzureContentSafety)
|
|
seen_outputs = []
|
|
|
|
async def fake_test_violation(content, source=None):
|
|
seen_outputs.append((content, source))
|
|
if "unsafe" in content:
|
|
raise HTTPException(status_code=400, detail={"error": "unsafe"})
|
|
|
|
guard.test_violation = fake_test_violation
|
|
response = ModelResponse(
|
|
choices=[
|
|
Choices(index=0, message=Message(role="assistant", content="clean")),
|
|
Choices(index=1, message=Message(role="assistant", content="unsafe")),
|
|
Choices(index=2, message=Message(role="assistant", content="later")),
|
|
]
|
|
)
|
|
|
|
with pytest.raises(HTTPException):
|
|
await guard.async_post_call_success_hook(
|
|
data={},
|
|
user_api_key_dict=user_api_key,
|
|
response=response,
|
|
)
|
|
|
|
assert seen_outputs == [("clean", "output"), ("unsafe", "output")]
|
|
|
|
|
|
# ── Secret Detection ──────────────────────────────────────────────────────────
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_secret_detection_redacts_multimodal_text_parts(user_api_key):
|
|
from litellm_enterprise.enterprise_callbacks.secret_detection import (
|
|
_ENTERPRISE_SecretDetection,
|
|
)
|
|
|
|
guard = _ENTERPRISE_SecretDetection()
|
|
data = {
|
|
"messages": [
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{
|
|
"type": "text",
|
|
"text": "AKIAIOSFODNN7EXAMPLE is the key",
|
|
},
|
|
{"type": "image_url", "image_url": {"url": "..."}},
|
|
],
|
|
}
|
|
]
|
|
}
|
|
|
|
await guard.async_pre_call_hook(
|
|
user_api_key_dict=user_api_key,
|
|
cache=DualCache(),
|
|
data=data,
|
|
call_type="completion",
|
|
)
|
|
|
|
parts = data["messages"][0]["content"]
|
|
assert "AKIAIOSFODNN7EXAMPLE" not in parts[0]["text"]
|
|
assert "[REDACTED]" in parts[0]["text"]
|
|
# Non-text part is preserved untouched.
|
|
assert parts[1] == {"type": "image_url", "image_url": {"url": "..."}}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_secret_detection_redacts_responses_api_input(user_api_key):
|
|
from litellm_enterprise.enterprise_callbacks.secret_detection import (
|
|
_ENTERPRISE_SecretDetection,
|
|
)
|
|
|
|
guard = _ENTERPRISE_SecretDetection()
|
|
data = {"input": "leak: AKIAIOSFODNN7EXAMPLE"}
|
|
|
|
await guard.async_pre_call_hook(
|
|
user_api_key_dict=user_api_key,
|
|
cache=DualCache(),
|
|
data=data,
|
|
call_type="moderation",
|
|
)
|
|
|
|
assert "AKIAIOSFODNN7EXAMPLE" not in data["input"]
|
|
assert "[REDACTED]" in data["input"]
|
|
|
|
|
|
# ── OpenAI Moderation ─────────────────────────────────────────────────────────
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_openai_moderation_inspects_multimodal_content(monkeypatch, user_api_key):
|
|
"""The aggregated text passed to ``llm_router.amoderation`` must include
|
|
list-format text parts and Responses-API input — without this, multimodal
|
|
content silently passed moderation."""
|
|
from enterprise.enterprise_hooks.openai_moderation import (
|
|
_ENTERPRISE_OpenAI_Moderation,
|
|
)
|
|
|
|
guard = _ENTERPRISE_OpenAI_Moderation()
|
|
|
|
seen_inputs = []
|
|
|
|
class FakeModeration:
|
|
results = [type("R", (), {"flagged": False})()]
|
|
|
|
async def fake_amoderation(model, input):
|
|
seen_inputs.append(input)
|
|
return FakeModeration()
|
|
|
|
fake_router = MagicMock()
|
|
fake_router.amoderation = AsyncMock(side_effect=fake_amoderation)
|
|
|
|
monkeypatch.setattr(
|
|
"litellm.proxy.proxy_server.llm_router", fake_router, raising=False
|
|
)
|
|
|
|
await guard.async_moderation_hook(
|
|
data={
|
|
"messages": [
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{"type": "text", "text": "alpha "},
|
|
{"type": "image_url", "image_url": {"url": "..."}},
|
|
{"type": "text", "text": "beta"},
|
|
],
|
|
}
|
|
]
|
|
},
|
|
user_api_key_dict=user_api_key,
|
|
call_type="acompletion",
|
|
)
|
|
|
|
assert seen_inputs == ["alpha beta"]
|
|
|
|
|
|
# ── Google Text Moderation ────────────────────────────────────────────────────
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_google_text_moderation_inspects_multimodal_content(user_api_key):
|
|
"""The text passed to Google's moderation client must include list-format
|
|
text parts."""
|
|
from enterprise.enterprise_hooks.google_text_moderation import (
|
|
_ENTERPRISE_GoogleTextModeration,
|
|
)
|
|
|
|
guard = _ENTERPRISE_GoogleTextModeration.__new__(_ENTERPRISE_GoogleTextModeration)
|
|
seen_documents = []
|
|
|
|
def fake_language_document(content, type_):
|
|
seen_documents.append(content)
|
|
return MagicMock()
|
|
|
|
fake_response = MagicMock()
|
|
fake_response.moderation_categories = []
|
|
|
|
guard.language_document = fake_language_document
|
|
guard.moderate_text_request = MagicMock(return_value=MagicMock())
|
|
guard.document_type = MagicMock()
|
|
guard.client = MagicMock()
|
|
guard.client.moderate_text = MagicMock(return_value=fake_response)
|
|
|
|
await guard.async_moderation_hook(
|
|
data={
|
|
"messages": [
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{"type": "text", "text": "hello "},
|
|
{"type": "image_url", "image_url": {"url": "..."}},
|
|
{"type": "text", "text": "world"},
|
|
],
|
|
}
|
|
]
|
|
},
|
|
user_api_key_dict=user_api_key,
|
|
call_type="acompletion",
|
|
)
|
|
|
|
assert seen_documents == ["hello world"]
|