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..71b37e7d25e --- /dev/null +++ b/litellm/llms/openai/videos/guardrail_translation/__init__.py @@ -0,0 +1,22 @@ +"""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.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..8ea6f3cf944 --- /dev/null +++ b/litellm/llms/openai/videos/guardrail_translation/handler.py @@ -0,0 +1,55 @@ +""" +OpenAI Video Generation Handler for Unified Guardrails + +Scans the `prompt` of video create / remix / edit / extension requests. The output is a +video job or binary content, so there is no text to guardrail on the response side. +""" + +from typing import TYPE_CHECKING, Final + +from litellm._logging import verbose_proxy_logger +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): + async def process_input_messages( + self, + data: dict[str, object], # mutable-ok: BaseTranslation signature; guardrails write call ids into request_data + guardrail_to_apply: "CustomGuardrail", + litellm_logging_obj: "LiteLLMLoggingObj | None" = None, + ) -> dict[str, object]: # mutable-ok: BaseTranslation signature + prompt: Final = data.get("prompt") + if not isinstance(prompt, str): + verbose_proxy_logger.debug("OpenAI Video Generation: no string prompt in request data, skipping guardrail") + 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( + 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, # mutable-ok: BaseTranslation signature + ) -> object: + return response diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 5783a39b30c..b8d3769970e 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -437,6 +437,7 @@ class CallTypes(str, Enum): ######################################################### create_video = "create_video" acreate_video = "acreate_video" + avideo_generation = "avideo_generation" avideo_retrieve = "avideo_retrieve" video_retrieve = "video_retrieve" avideo_content = "avideo_content" 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 8cad1c634a9..e34b181f2e7 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 @@ -634,6 +634,58 @@ class TestUnifiedLLMGuardrails: f"key-scoped users get 403 on this alias" ) + class TestVideoGuardrailE2E: + """UnifiedLLMGuardrails must scan the prompt on the video routes the proxy exposes.""" + + @pytest.fixture(autouse=True) + def _use_real_mappings(self): + unified_module.endpoint_guardrail_translation_mappings = load_guardrail_translation_mappings() + yield + unified_module.endpoint_guardrail_translation_mappings = None + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "call_type", + [ + "avideo_generation", + CallTypes.acreate_video.value, + CallTypes.avideo_remix.value, + CallTypes.avideo_edit.value, + CallTypes.avideo_extension.value, + ], + ) + async def test_pre_call_hook_scans_video_prompt(self, call_type: str) -> None: + class RewritingGuardrail(RecordingGuardrail): + async def apply_guardrail(self, inputs, request_data, input_type, **kwargs): + await super().apply_guardrail(inputs, request_data, input_type, **kwargs) + return {"texts": [f"{text} [GUARDRAILED]" for text in inputs.get("texts", [])]} + + handler = UnifiedLLMGuardrails() + guardrail = RewritingGuardrail() + data = { + "guardrail_to_apply": guardrail, + "model": "vertex_ai/veo-3.1-fast-generate-001", + "prompt": "A dog surfing", + "seconds": "4", + } + + result = await handler.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test-key", request_route="/v1/videos"), + cache=DualCache(), + data=data, + call_type=call_type, + ) + + assert guardrail.event_history == [GuardrailEventHooks.pre_call] + assert guardrail.apply_calls == [ + { + "inputs": {"texts": ["A dog surfing"], "model": "vertex_ai/veo-3.1-fast-generate-001"}, + "input_type": "request", + } + ] + assert result["prompt"] == "A dog surfing [GUARDRAILED]" + assert result["seconds"] == "4" + class TestOCRGuardrailE2E: """End-to-end tests: UnifiedLLMGuardrails -> OCRHandler."""