mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-28 01:32:17 +00:00
fix(guardrails): scan video prompts for key-attached guardrails on /v1/videos
/v1/videos dispatches call_type avideo_generation, which CallTypes did not know and no guardrail translation handler covered, so the unified guardrail hook returned the request unscanned. Add the video call types and an OpenAI video guardrail translation package that scans the prompt for create, remix, edit and extension requests Resolves LIT-6685 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
parent
403b4be40e
commit
9ae2fe2ea4
8 changed files with 237 additions and 3 deletions
23
litellm/llms/openai/videos/guardrail_translation/__init__.py
Normal file
23
litellm/llms/openai/videos/guardrail_translation/__init__.py
Normal file
|
|
@ -0,0 +1,23 @@
|
|||
"""OpenAI Video Generation handler for Unified Guardrails."""
|
||||
|
||||
from typing import Final
|
||||
|
||||
from litellm.llms.openai.videos.guardrail_translation.handler import (
|
||||
OpenAIVideoGenerationHandler,
|
||||
)
|
||||
from litellm.types.utils import CallTypes
|
||||
|
||||
guardrail_translation_mappings: Final = {
|
||||
CallTypes.video_generation: OpenAIVideoGenerationHandler,
|
||||
CallTypes.avideo_generation: OpenAIVideoGenerationHandler,
|
||||
CallTypes.create_video: OpenAIVideoGenerationHandler,
|
||||
CallTypes.acreate_video: OpenAIVideoGenerationHandler,
|
||||
CallTypes.video_remix: OpenAIVideoGenerationHandler,
|
||||
CallTypes.avideo_remix: OpenAIVideoGenerationHandler,
|
||||
CallTypes.video_edit: OpenAIVideoGenerationHandler,
|
||||
CallTypes.avideo_edit: OpenAIVideoGenerationHandler,
|
||||
CallTypes.video_extension: OpenAIVideoGenerationHandler,
|
||||
CallTypes.avideo_extension: OpenAIVideoGenerationHandler,
|
||||
}
|
||||
|
||||
__all__ = ["OpenAIVideoGenerationHandler", "guardrail_translation_mappings"]
|
||||
48
litellm/llms/openai/videos/guardrail_translation/handler.py
Normal file
48
litellm/llms/openai/videos/guardrail_translation/handler.py
Normal file
|
|
@ -0,0 +1,48 @@
|
|||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
|
||||
from litellm.types.utils import GenericGuardrailAPIInputs
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
|
||||
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],
|
||||
guardrail_to_apply: "CustomGuardrail",
|
||||
litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
|
||||
) -> dict[str, object]:
|
||||
prompt: Final = data.get("prompt")
|
||||
if not isinstance(prompt, str):
|
||||
return data
|
||||
|
||||
model: Final = data.get("model")
|
||||
inputs: Final = (
|
||||
GenericGuardrailAPIInputs(texts=[prompt], model=model)
|
||||
if isinstance(model, str)
|
||||
else GenericGuardrailAPIInputs(texts=[prompt])
|
||||
)
|
||||
guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail( # pyright: ignore[reportUnknownMemberType] # request_data is a bare dict
|
||||
inputs=inputs,
|
||||
request_data=data,
|
||||
input_type="request",
|
||||
logging_obj=litellm_logging_obj,
|
||||
)
|
||||
guardrailed_texts: Final = guardrailed_inputs.get("texts", [])
|
||||
return {**data, "prompt": guardrailed_texts[0] if guardrailed_texts else prompt}
|
||||
|
||||
async def process_output_response(
|
||||
self,
|
||||
response: object,
|
||||
guardrail_to_apply: "CustomGuardrail",
|
||||
litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
|
||||
user_api_key_dict: "UserAPIKeyAuth | None" = None,
|
||||
request_data: dict[str, object] | None = None,
|
||||
) -> object:
|
||||
return response
|
||||
|
|
@ -459,6 +459,8 @@ class CallTypes(str, Enum):
|
|||
#########################################################
|
||||
create_video = "create_video"
|
||||
acreate_video = "acreate_video"
|
||||
video_generation = "video_generation"
|
||||
avideo_generation = "avideo_generation"
|
||||
avideo_retrieve = "avideo_retrieve"
|
||||
video_retrieve = "video_retrieve"
|
||||
avideo_content = "avideo_content"
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@
|
|||
- {id: guardrail.presidio.pre_call.logs_masked_entities, module: guardrail, tier: P0, hook_point: pre_call, assertions: [logs_masked_entities], exercised_on: [chat_completions], source: "guardrail_hooks/presidio.py", rationale: "A masking run must record itself on the spend log: the dashboard's guardrail panel renders the masked-entity counts and per-entity scores straight off metadata.guardrail_information, so a run that masks but records nothing leaves an operator unable to audit it"}
|
||||
- {id: guardrail.bedrock.pre_call.blocks, module: guardrail, tier: P0, hook_point: pre_call, assertions: [blocks], exercised_on: [chat_completions], source: "guardrail_hooks/bedrock_guardrails.py", rationale: "AWS content guardrail blocks harmful input"}
|
||||
- {id: guardrail.litellm_content_filter.pre_call.blocks, module: guardrail, tier: P0, hook_point: pre_call, assertions: [blocks], exercised_on: [chat_completions], source: "test_team_disable_global_guardrail_e2e.py", rationale: "Local content-filter default-on blocks banned keyword pre-call"}
|
||||
- {id: guardrail.litellm_content_filter.pre_call.blocks_video, module: guardrail, tier: P0, hook_point: pre_call, assertions: [blocks], exercised_on: [videos], source: "test_key_guardrail_video_e2e.py", fail_before_fix: proven, rationale: "A content-filter guardrail attached to a key (metadata.guardrails) blocks a banned prompt on POST /v1/videos before the provider is called; before the fix the route's call type was unknown to the unified guardrail hook and the prompt went to the provider unscanned (LIT-6685)"}
|
||||
- {id: guardrail.litellm_content_filter.pre_call.allows, module: guardrail, tier: P0, hook_point: pre_call, assertions: [allows], exercised_on: [chat_completions], source: "test_team_disable_global_guardrail_e2e.py", rationale: "Team disable_global_guardrails bypasses default-on content filter"}
|
||||
- {id: guardrail.litellm_content_filter.apply_endpoint.blocks, module: guardrail, tier: P0, hook_point: apply_endpoint, assertions: [blocks], exercised_on: [chat_completions], source: "guardrail_endpoints.py:apply_guardrail", rationale: "POST /guardrails/apply_guardrail blocks banned content for customers that call the apply surface directly"}
|
||||
- {id: guardrail.litellm_content_filter.apply_endpoint.allows, module: guardrail, tier: P0, hook_point: apply_endpoint, assertions: [allows], exercised_on: [chat_completions], source: "guardrail_endpoints.py:apply_guardrail", rationale: "POST /guardrails/apply_guardrail returns clean text for allowed input"}
|
||||
|
|
|
|||
|
|
@ -20,6 +20,7 @@ from models import (
|
|||
ChatResponse,
|
||||
ChatTool,
|
||||
KeyGenerateBody,
|
||||
KeyMetadata,
|
||||
LiteLLMParamsBody,
|
||||
TeamDeleteBody,
|
||||
TeamInfoParams,
|
||||
|
|
@ -27,6 +28,8 @@ from models import (
|
|||
TeamMetadata,
|
||||
TeamNewBody,
|
||||
TeamNewResponse,
|
||||
VideoCreateBody,
|
||||
VideoCreateResponse,
|
||||
)
|
||||
from proxy_client import ProxyClient
|
||||
from pydantic import BaseModel
|
||||
|
|
@ -151,12 +154,12 @@ class _ResponsesGuardrailBody(BaseModel):
|
|||
class GuardrailsClient:
|
||||
proxy: ProxyClient
|
||||
|
||||
def create_content_filter_guardrail(self, name: str, blocked_keyword: str) -> str:
|
||||
def create_content_filter_guardrail(self, name: str, blocked_keyword: str, *, default_on: bool = True) -> str:
|
||||
return self.register(
|
||||
name,
|
||||
ContentFilterParamsBody(
|
||||
mode="pre_call",
|
||||
default_on=True,
|
||||
default_on=default_on,
|
||||
blocked_words=[BlockedWordBody(keyword=blocked_keyword, action="BLOCK")],
|
||||
),
|
||||
)
|
||||
|
|
@ -266,6 +269,23 @@ class GuardrailsClient:
|
|||
def create_key_in_team(self, team_id: str) -> str:
|
||||
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))
|
||||
)
|
||||
resources.defer(lambda: self.proxy.delete_key(key))
|
||||
return key
|
||||
|
||||
def create_video(self, key: str, model: str, prompt: str) -> Result[VideoCreateResponse]:
|
||||
return self.proxy.transport.post(
|
||||
"/v1/videos",
|
||||
headers=self.proxy.transport.bearer(key),
|
||||
json=VideoCreateBody(model=model, prompt=prompt, seconds="4"),
|
||||
response_type=VideoCreateResponse,
|
||||
)
|
||||
|
||||
def chat(
|
||||
self,
|
||||
key: str,
|
||||
|
|
|
|||
85
tests/e2e/guardrails/test_key_guardrail_video_e2e.py
Normal file
85
tests/e2e/guardrails/test_key_guardrail_video_e2e.py
Normal file
|
|
@ -0,0 +1,85 @@
|
|||
"""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 lifecycle import ResourceManager
|
||||
from models import LiteLLMParamsBody
|
||||
|
||||
pytestmark = pytest.mark.e2e
|
||||
|
||||
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}"
|
||||
|
||||
|
||||
def _create_video_model(client: GuardrailsClient, resources: ResourceManager) -> str:
|
||||
model_name = f"e2e-guard-video-{unique_marker()}"
|
||||
model_id = client.proxy.create_model(
|
||||
model_name,
|
||||
LiteLLMParamsBody(
|
||||
model=VIDEO_BACKEND,
|
||||
vertex_project="os.environ/VERTEXAI_PROJECT",
|
||||
vertex_location="os.environ/VERTEXAI_LOCATION",
|
||||
vertex_credentials="os.environ/VERTEXAI_CREDENTIALS",
|
||||
),
|
||||
provider_live=True,
|
||||
)
|
||||
resources.defer(lambda: client.proxy.delete_model(model_id))
|
||||
return model_name
|
||||
|
||||
|
||||
class TestKeyAttachedGuardrailOnVideos:
|
||||
@pytest.mark.covers(
|
||||
"guardrail.litellm_content_filter.pre_call.blocks_video",
|
||||
exercised_on=["videos"],
|
||||
)
|
||||
def test_key_attached_content_filter_blocks_banned_video_prompt(
|
||||
self, client: GuardrailsClient, resources: ResourceManager
|
||||
) -> None:
|
||||
banned = unique_marker()
|
||||
guardrail_name = f"e2e-video-filter-{banned}"
|
||||
guardrail_id = client.create_content_filter_guardrail(guardrail_name, banned, default_on=False)
|
||||
resources.defer(lambda: client.delete_guardrail(guardrail_id))
|
||||
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):
|
||||
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}"
|
||||
)
|
||||
|
|
@ -60,6 +60,7 @@ class KeyMetadata(BaseModel):
|
|||
priority: str | None = None
|
||||
batch_enqueued_token_limit: int | None = None
|
||||
tag: str | None = None
|
||||
guardrails: list[str] | None = None
|
||||
|
||||
|
||||
class ObjectPermission(BaseModel):
|
||||
|
|
@ -697,6 +698,20 @@ class EmbedResponse(BaseModel):
|
|||
model: str | None = None
|
||||
|
||||
|
||||
# ---------- videos ----------
|
||||
|
||||
|
||||
class VideoCreateBody(BaseModel):
|
||||
model: str
|
||||
prompt: str
|
||||
seconds: str | None = None
|
||||
|
||||
|
||||
class VideoCreateResponse(BaseModel):
|
||||
id: str
|
||||
status: str | None = None
|
||||
|
||||
|
||||
# ---------- rerank ----------
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -13,7 +13,7 @@ from litellm.integrations.custom_guardrail import (
|
|||
log_guardrail_information,
|
||||
)
|
||||
from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route
|
||||
from litellm.llms import load_guardrail_translation_mappings
|
||||
from litellm.llms import discover_guardrail_translation_mappings, load_guardrail_translation_mappings
|
||||
from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
|
||||
from litellm.llms.base_llm.guardrail_translation.utils import (
|
||||
effective_skip_system_message_for_guardrail,
|
||||
|
|
@ -61,6 +61,14 @@ class RecordingGuardrail(CustomGuardrail):
|
|||
return {"texts": inputs.get("texts", [])}
|
||||
|
||||
|
||||
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"]]}
|
||||
|
||||
|
||||
class _NoopTranslation(BaseTranslation):
|
||||
"""Test translation handler that simply echoes input/output."""
|
||||
|
||||
|
|
@ -360,6 +368,38 @@ class TestUnifiedLLMGuardrails:
|
|||
|
||||
assert guardrail.event_history == [GuardrailEventHooks.pre_mcp_call]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"call_type",
|
||||
["avideo_generation", "acreate_video", "avideo_remix", "avideo_edit", "avideo_extension"],
|
||||
)
|
||||
async def test_video_routes_scan_prompt_and_keep_rewrite(self, monkeypatch, call_type: str) -> None:
|
||||
"""LIT-6685: /v1/videos dispatches call_type="avideo_generation", which the
|
||||
hook once swallowed as an unknown CallTypes value and returned unscanned.
|
||||
Runs against the discovered handler map so the video package must really exist."""
|
||||
_patch_translation_mappings(monkeypatch, discover_guardrail_translation_mappings())
|
||||
handler = UnifiedLLMGuardrails()
|
||||
guardrail = RewritingGuardrail()
|
||||
data = {
|
||||
"guardrail_to_apply": guardrail,
|
||||
"model": "veo-3.1-fast",
|
||||
"prompt": "a paper boat on a stream",
|
||||
"seconds": "4",
|
||||
}
|
||||
|
||||
result = await handler.async_pre_call_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="test-key"),
|
||||
cache=DualCache(),
|
||||
data=data,
|
||||
call_type=call_type,
|
||||
)
|
||||
|
||||
assert guardrail.event_history == [GuardrailEventHooks.pre_call]
|
||||
assert [call["inputs"]["texts"] for call in guardrail.apply_calls] == [["a paper boat on a stream"]]
|
||||
assert guardrail.apply_calls[0]["inputs"]["model"] == "veo-3.1-fast"
|
||||
assert result["prompt"] == "a paper boat on a stream [GUARDRAILED]"
|
||||
assert result["seconds"] == "4"
|
||||
|
||||
class TestAsyncModerationHook:
|
||||
@pytest.mark.asyncio
|
||||
async def test_uses_mcp_event_type(self):
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue