scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-24 00:52:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 91

test(guardrails): gate the video e2e on a chat probe so a miss starts at most one paid job

Browse source 
Addresses Greptile review: typed RewritingGuardrail override, dropped routine docstrings, and the e2e waits for the key guardrail to sync via /chat/completions before its single /v1/videos call

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yucheng 2026-09-21 23:20:29 +00:00
parent 652bddfdc6
commit 2f0584cec6
4 changed files with 73 additions and 149 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm/llms/openai/videos/guardrail_translation/handler.py
View file

@ -10,8 +10,6 @@ if TYPE_CHECKING:
class OpenAIVideoGenerationHandler(BaseTranslation):
"""Scans the text `prompt` of video create, remix, edit and extension requests."""
async def process_input_messages(
self,
data: dict[str, object], # mutable-ok: BaseTranslation contract passes the proxy's request dict through

2
tests/e2e/guardrails/guardrails_client.py
View file

@ -270,8 +270,6 @@ class GuardrailsClient:
return self.proxy.generate_key(KeyGenerateBody(team_id=team_id, user_id="e2e-guardrails-user"))
def create_key_with_guardrails(self, resources: ResourceManager, guardrails: list[str]) -> str:
"""A key whose metadata.guardrails attaches the named guardrails to every
request made with it, the way an admin attaches one from the key page."""
key = self.proxy.generate_key(
KeyGenerateBody(user_id="e2e-guardrails-user", metadata=KeyMetadata(guardrails=guardrails))
)

58
tests/e2e/guardrails/test_key_guardrail_video_e2e.py
View file

@ -1,30 +1,17 @@
"""Live e2e: a guardrail attached to a virtual key (metadata.guardrails) must run
on POST /v1/videos, so a banned prompt is rejected before the provider is called
instead of quietly starting a paid video generation job (LIT-6685).
Uses a local litellm_content_filter (keyword match, no external service) so the
block is deterministic, and a real Vertex AI Veo deployment so the sad path proves
the provider was never reached.
"""
from __future__ import annotations
import time
import pytest
from e2e_config import unique_marker
from e2e_http import Success, UnknownApiError
from guardrails_client import GuardrailsClient
from guardrails_client import GuardrailsClient, poll_until_blocked
from lifecycle import ResourceManager
from models import LiteLLMParamsBody
pytestmark = pytest.mark.e2e
CHAT_MODEL = "gemini-2.5-flash"
VIDEO_BACKEND = "vertex_ai/veo-3.1-fast-generate-001"
GUARDRAIL_PROPAGATION_DEADLINE_SECONDS = 40.0
GUARDRAIL_PROPAGATION_POLL_INTERVAL_SECONDS = 5.0
def _video_prompt_with(banned_keyword: str) -> str:
return f"A short clip of a paper boat floating down a stream. {banned_keyword}"
@ -61,25 +48,22 @@ class TestKeyAttachedGuardrailOnVideos:
key = client.create_key_with_guardrails(resources, [guardrail_name])
model = _create_video_model(client, resources)
deadline = time.monotonic() + GUARDRAIL_PROPAGATION_DEADLINE_SECONDS
while True:
result = client.create_video(key, model, _video_prompt_with(banned))
match result:
case UnknownApiError(status_code=status, body=body):
assert status == 400, f"expected a 400 guardrail block, got {status}: {body[:300]}"
assert "content blocked" in body.lower() or banned in body, (
f"block response missing content-filter reason: {body[:300]}"
)
return
case Success(data=video) if time.monotonic() >= deadline:
pytest.fail(
f"key-attached guardrail {guardrail_name!r} was skipped on /v1/videos: "
f"the banned prompt reached the provider and started video job {video.id}"
)
case _ if time.monotonic() < deadline:
time.sleep(GUARDRAIL_PROPAGATION_POLL_INTERVAL_SECONDS)
case _:
pytest.fail(
f"key-attached guardrail never blocked the banned prompt within "
f"{GUARDRAIL_PROPAGATION_DEADLINE_SECONDS}s; got {result}"
)
synced = poll_until_blocked(lambda: client.chat(key, CHAT_MODEL, _video_prompt_with(banned)))
assert isinstance(synced, UnknownApiError) and synced.status_code == 400, (
f"key guardrail {guardrail_name!r} never synced to the proxy on /chat/completions: {synced}"
)
result = client.create_video(key, model, _video_prompt_with(banned))
match result:
case UnknownApiError(status_code=status, body=body):
assert status == 400, f"expected a 400 guardrail block, got {status}: {body[:300]}"
assert "content blocked" in body.lower() or banned in body, (
f"block response missing content-filter reason: {body[:300]}"
)
case Success(data=video):
pytest.fail(
f"key-attached guardrail {guardrail_name!r} was skipped on /v1/videos: "
f"the banned prompt reached the provider and started video job {video.id}"
)
case _:
pytest.fail(f"unexpected /v1/videos outcome for a banned prompt: {result}")

160
tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py
View file

@ -2,7 +2,7 @@
import logging
from types import SimpleNamespace
from typing import Final
from typing import TYPE_CHECKING, Final, Literal
import pytest
@ -41,7 +41,10 @@ from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrai
)
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.llms.openai import ResponsesAPIResponse
from litellm.types.utils import CallTypes, Delta, ModelResponseStream, StreamingChoices
from litellm.types.utils import CallTypes, Delta, GenericGuardrailAPIInputs, ModelResponseStream, StreamingChoices
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
class RecordingGuardrail(CustomGuardrail):
@ -62,11 +65,15 @@ class RecordingGuardrail(CustomGuardrail):
class RewritingGuardrail(RecordingGuardrail):
"""Records like RecordingGuardrail and hands back a visibly rewritten text."""
async def apply_guardrail(self, inputs, request_data, input_type, **kwargs):
recorded = await super().apply_guardrail(inputs, request_data, input_type, **kwargs)
return {"texts": [f"{text} [GUARDRAILED]" for text in recorded["texts"]]}
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict, # mutable-ok: CustomGuardrail.apply_guardrail contract
input_type: Literal["request", "response"],
logging_obj: "LiteLLMLoggingObj | None" = None,
) -> GenericGuardrailAPIInputs:
recorded: Final = await super().apply_guardrail(inputs, request_data, input_type, logging_obj=logging_obj)
return GenericGuardrailAPIInputs(texts=[f"{text} [GUARDRAILED]" for text in recorded["texts"]])
class _NoopTranslation(BaseTranslation):
@ -123,9 +130,7 @@ class TestUnifiedLLMGuardrails:
assert msgs[0]["content"] == "sys"
def test_effective_skip_respects_per_guardrail_over_global(self, monkeypatch):
monkeypatch.setattr(
litellm, "skip_system_message_in_guardrail", True, raising=False
)
monkeypatch.setattr(litellm, "skip_system_message_in_guardrail", True, raising=False)
class G:
skip_system_message_in_guardrail = False
@ -138,21 +143,15 @@ class TestUnifiedLLMGuardrails:
assert effective_skip_system_message_for_guardrail(G2()) is True
@pytest.mark.asyncio
async def test_openai_handler_skips_system_in_guardrail_inputs(
self, monkeypatch
):
monkeypatch.setattr(
litellm, "skip_system_message_in_guardrail", True, raising=False
)
async def test_openai_handler_skips_system_in_guardrail_inputs(self, monkeypatch):
monkeypatch.setattr(litellm, "skip_system_message_in_guardrail", True, raising=False)
captured = {}
class MockGuardrail:
skip_system_message_in_guardrail = None
async def apply_guardrail(
self, inputs, request_data, input_type, logging_obj=None
):
async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
captured["inputs"] = inputs
return inputs
@ -177,21 +176,15 @@ class TestUnifiedLLMGuardrails:
assert data["messages"][0]["content"] == "secret system"
@pytest.mark.asyncio
async def test_openai_handler_per_guardrail_skip_false_overrides_global(
self, monkeypatch
):
monkeypatch.setattr(
litellm, "skip_system_message_in_guardrail", True, raising=False
)
async def test_openai_handler_per_guardrail_skip_false_overrides_global(self, monkeypatch):
monkeypatch.setattr(litellm, "skip_system_message_in_guardrail", True, raising=False)
captured = {}
class MockGuardrail:
skip_system_message_in_guardrail = False
async def apply_guardrail(
self, inputs, request_data, input_type, logging_obj=None
):
async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
captured["inputs"] = inputs
return inputs
@ -209,10 +202,7 @@ class TestUnifiedLLMGuardrails:
)
assert "sys" in captured["inputs"]["texts"]
roles = {
m.get("role")
for m in (captured["inputs"].get("structured_messages") or [])
}
roles = {m.get("role") for m in (captured["inputs"].get("structured_messages") or [])}
assert "system" in roles
class TestSkipToolMessageForChatCompletions:
@ -237,12 +227,8 @@ class TestUnifiedLLMGuardrails:
assert all(m["role"] != "tool" for m in out)
assert msgs[2]["content"] == "tool result"
def test_effective_skip_tool_respects_per_guardrail_over_global(
self, monkeypatch
):
monkeypatch.setattr(
litellm, "skip_tool_message_in_guardrail", True, raising=False
)
def test_effective_skip_tool_respects_per_guardrail_over_global(self, monkeypatch):
monkeypatch.setattr(litellm, "skip_tool_message_in_guardrail", True, raising=False)
class G:
skip_tool_message_in_guardrail = False
@ -256,18 +242,14 @@ class TestUnifiedLLMGuardrails:
@pytest.mark.asyncio
async def test_openai_handler_skips_tool_in_guardrail_inputs(self, monkeypatch):
monkeypatch.setattr(
litellm, "skip_tool_message_in_guardrail", True, raising=False
)
monkeypatch.setattr(litellm, "skip_tool_message_in_guardrail", True, raising=False)
captured = {}
class MockGuardrail:
skip_tool_message_in_guardrail = None
async def apply_guardrail(
self, inputs, request_data, input_type, logging_obj=None
):
async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
captured["inputs"] = inputs
return inputs
@ -307,21 +289,15 @@ class TestUnifiedLLMGuardrails:
assert data["messages"][2]["content"] == "secret tool result"
@pytest.mark.asyncio
async def test_openai_handler_per_guardrail_skip_tool_false_overrides_global(
self, monkeypatch
):
monkeypatch.setattr(
litellm, "skip_tool_message_in_guardrail", True, raising=False
)
async def test_openai_handler_per_guardrail_skip_tool_false_overrides_global(self, monkeypatch):
monkeypatch.setattr(litellm, "skip_tool_message_in_guardrail", True, raising=False)
captured = {}
class MockGuardrail:
skip_tool_message_in_guardrail = False
async def apply_guardrail(
self, inputs, request_data, input_type, logging_obj=None
):
async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
captured["inputs"] = inputs
return inputs
@ -339,10 +315,7 @@ class TestUnifiedLLMGuardrails:
)
assert "tr" in captured["inputs"]["texts"]
roles = {
m.get("role")
for m in (captured["inputs"].get("structured_messages") or [])
}
roles = {m.get("role") for m in (captured["inputs"].get("structured_messages") or [])}
assert "tool" in roles
class TestAsyncPreCallHook:
@ -464,7 +437,9 @@ class TestUnifiedLLMGuardrails:
async def process_input_messages(self, data, guardrail_to_apply, litellm_logging_obj=None): # type: ignore[override]
return data
async def process_output_response(self, response, guardrail_to_apply, litellm_logging_obj=None, user_api_key_dict=None): # type: ignore[override]
async def process_output_response(
self, response, guardrail_to_apply, litellm_logging_obj=None, user_api_key_dict=None
): # type: ignore[override]
return response
async def process_output_streaming_response(
@ -533,9 +508,7 @@ class TestUnifiedLLMGuardrails:
response=mock_stream(),
request_data=request_data,
):
content = (
item.choices[0].delta.content if item.choices[0].delta else None
)
content = item.choices[0].delta.content if item.choices[0].delta else None
yielded_contents.append(content)
# Every chunk should have non-empty content
@ -586,23 +559,18 @@ class TestUnifiedLLMGuardrails:
],
)
@pytest.mark.asyncio
async def test_post_call_scans_output_on_every_registered_alias(
self, request_route: str
) -> None:
async def test_post_call_scans_output_on_every_registered_alias(self, request_route: str) -> None:
handler = UnifiedLLMGuardrails()
guardrail = RecordingGuardrail()
await handler.async_post_call_success_hook(
data={"guardrail_to_apply": guardrail, "model": "gpt-4o"},
user_api_key_dict=UserAPIKeyAuth(
api_key="test-key", request_route=request_route
),
user_api_key_dict=UserAPIKeyAuth(api_key="test-key", request_route=request_route),
response=self._responses_api_response(),
)
assert guardrail.apply_calls, (
f"guardrail never ran for request_route={request_route!r}; model "
f"output reached the client unscanned"
f"guardrail never ran for request_route={request_route!r}; model output reached the client unscanned"
)
assert guardrail.apply_calls[0]["input_type"] == "response"
assert guardrail.apply_calls[0]["inputs"]["texts"] == ["Paris"]
@ -632,18 +600,14 @@ class TestUnifiedLLMGuardrails:
assert CallTypes.responses in mappings
@pytest.mark.asyncio
async def test_unresolvable_route_skips_scanning_and_says_so(
self, caplog: pytest.LogCaptureFixture
) -> None:
async def test_unresolvable_route_skips_scanning_and_says_so(self, caplog: pytest.LogCaptureFixture) -> None:
handler = UnifiedLLMGuardrails()
guardrail = RecordingGuardrail()
with caplog.at_level(logging.WARNING):
result = await handler.async_post_call_success_hook(
data={"guardrail_to_apply": guardrail, "model": "gpt-4o"},
user_api_key_dict=UserAPIKeyAuth(
api_key="test-key", request_route="/cursor/chat/completions"
),
user_api_key_dict=UserAPIKeyAuth(api_key="test-key", request_route="/cursor/chat/completions"),
response=self._responses_api_response(),
)
@ -662,9 +626,7 @@ class TestUnifiedLLMGuardrails:
with caplog.at_level(logging.WARNING):
await handler.async_post_call_success_hook(
data={"guardrail_to_apply": guardrail, "model": "gpt-4o"},
user_api_key_dict=UserAPIKeyAuth(
api_key="test-key", request_route="/v1/chat/completions"
),
user_api_key_dict=UserAPIKeyAuth(api_key="test-key", request_route="/v1/chat/completions"),
response=self._responses_api_response(),
)
@ -774,15 +736,10 @@ class TestUnifiedLLMGuardrails:
assert guardrail.event_history == [GuardrailEventHooks.pre_call]
assert len(guardrail.apply_calls) == 1
assert guardrail.apply_calls[0]["input_type"] == "request"
assert (
"https://arxiv.org/pdf/2201.04234"
in guardrail.apply_calls[0]["inputs"]["texts"]
)
assert "https://arxiv.org/pdf/2201.04234" in guardrail.apply_calls[0]["inputs"]["texts"]
# Data should be returned with document intact
assert (
result["document"]["document_url"] == "https://arxiv.org/pdf/2201.04234"
)
assert result["document"]["document_url"] == "https://arxiv.org/pdf/2201.04234"
@pytest.mark.asyncio
async def test_moderation_hook_invokes_ocr_handler(self):
@ -810,10 +767,7 @@ class TestUnifiedLLMGuardrails:
assert guardrail.event_history == [GuardrailEventHooks.during_call]
assert len(guardrail.apply_calls) == 1
assert (
"https://example.com/scan.png"
in guardrail.apply_calls[0]["inputs"]["texts"]
)
assert "https://example.com/scan.png" in guardrail.apply_calls[0]["inputs"]["texts"]
@pytest.mark.asyncio
async def test_post_call_success_hook_guardrails_ocr_output(self):
@ -829,9 +783,7 @@ class TestUnifiedLLMGuardrails:
def should_run_guardrail(self, data, event_type): # type: ignore[override]
return True
async def apply_guardrail(
self, inputs, request_data, input_type, **kwargs
):
async def apply_guardrail(self, inputs, request_data, input_type, **kwargs):
texts = inputs.get("texts", [])
return {"texts": [t.replace("SECRET", "[REDACTED]") for t in texts]}
@ -1578,9 +1530,7 @@ class TestStreamingTransform:
# And the redacted text ("SECRET") reached the wire on some non-tool
# chunk (i.e. the text terminator).
transformed = "".join(
item.choices[0].delta.content or ""
for item in out
if item.choices and not item.choices[0].delta.tool_calls
item.choices[0].delta.content or "" for item in out if item.choices and not item.choices[0].delta.tool_calls
)
assert "SECRET" in transformed
assert "secret" not in transformed
@ -1725,7 +1675,9 @@ class TestStreamingTransform:
_stream_chunk("went home."),
ModelResponseStream(
choices=[
StreamingChoices(index=0, delta=Delta(content=None, role="assistant", tool_calls=None), finish_reason=None),
StreamingChoices(
index=0, delta=Delta(content=None, role="assistant", tool_calls=None), finish_reason=None
),
StreamingChoices(
index=1,
delta=Delta(
@ -2025,9 +1977,7 @@ class TestStreamingHttpErrorFrames:
guardrail = _EosHttpBlockingGuardrail()
chunks = _anthropic_message_chunks(["hello ", "world"])
out = await _drive_stream(
UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/messages"
)
out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/messages")
raw = b"".join(c for c in out if isinstance(c, bytes)).decode()
assert "hello " in raw
@ -2051,9 +2001,7 @@ class TestStreamingHttpErrorFrames:
},
]
out = await _drive_stream(
UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/responses"
)
out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/responses")
assert chunks[0] in out and chunks[1] in out
assert chunks[2] not in out
@ -2116,9 +2064,7 @@ class TestStreamingGuardrailInformationBucket:
for chunk in chunks:
yield chunk
user_api_key_dict = UserAPIKeyAuth(
api_key="test-key", user_id="user-1", request_route="/v1/chat/completions"
)
user_api_key_dict = UserAPIKeyAuth(api_key="test-key", user_id="user-1", request_route="/v1/chat/completions")
request_data = {"guardrail_to_apply": guardrail, "model": "gpt-4", "metadata": {}}
out = []
@ -2447,7 +2393,5 @@ class TestTranslationMappingsAreReadLive:
assert len(guardrail.apply_calls) == 1
assert not [
name
for name, value in vars(unified_module).items()
if isinstance(value, dict) and CallTypes.aocr in value
name for name, value in vars(unified_module).items() if isinstance(value, dict) and CallTypes.aocr in value
]