diff --git a/litellm/llms/openai/videos/guardrail_translation/__init__.py b/litellm/llms/openai/videos/guardrail_translation/__init__.py new file mode 100644 index 00000000000..6540ba994e4 --- /dev/null +++ b/litellm/llms/openai/videos/guardrail_translation/__init__.py @@ -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"] diff --git a/litellm/llms/openai/videos/guardrail_translation/handler.py b/litellm/llms/openai/videos/guardrail_translation/handler.py new file mode 100644 index 00000000000..74fafb4bcfe --- /dev/null +++ b/litellm/llms/openai/videos/guardrail_translation/handler.py @@ -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 diff --git a/litellm/types/utils.py b/litellm/types/utils.py index e1d43b7fccb..e23f329ee83 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -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" diff --git a/tests/e2e/coverage_registry/guardrail.yaml b/tests/e2e/coverage_registry/guardrail.yaml index 81832bebf49..86eb44f6cb1 100644 --- a/tests/e2e/coverage_registry/guardrail.yaml +++ b/tests/e2e/coverage_registry/guardrail.yaml @@ -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"} diff --git a/tests/e2e/guardrails/guardrails_client.py b/tests/e2e/guardrails/guardrails_client.py index ed112a79b9b..11efaf0296e 100644 --- a/tests/e2e/guardrails/guardrails_client.py +++ b/tests/e2e/guardrails/guardrails_client.py @@ -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, diff --git a/tests/e2e/guardrails/test_key_guardrail_video_e2e.py b/tests/e2e/guardrails/test_key_guardrail_video_e2e.py new file mode 100644 index 00000000000..8c70ea5b46f --- /dev/null +++ b/tests/e2e/guardrails/test_key_guardrail_video_e2e.py @@ -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}" + ) diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 355329585fb..3eb18cdf899 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -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 ---------- diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py index d1d22d0d7c2..138b6d739f7 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py @@ -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):