From 9ae2fe2ea4095c9d2baa3d1613756a7ea7f4535f Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 21 Sep 2026 22:23:11 +0000 Subject: [PATCH 1/6] 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> --- .../videos/guardrail_translation/__init__.py | 23 +++++ .../videos/guardrail_translation/handler.py | 48 +++++++++++ litellm/types/utils.py | 2 + tests/e2e/coverage_registry/guardrail.yaml | 1 + tests/e2e/guardrails/guardrails_client.py | 24 +++++- .../test_key_guardrail_video_e2e.py | 85 +++++++++++++++++++ tests/e2e/models.py | 15 ++++ .../test_unified_guardrail.py | 42 ++++++++- 8 files changed, 237 insertions(+), 3 deletions(-) create mode 100644 litellm/llms/openai/videos/guardrail_translation/__init__.py create mode 100644 litellm/llms/openai/videos/guardrail_translation/handler.py create mode 100644 tests/e2e/guardrails/test_key_guardrail_video_e2e.py 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): From 8ffca3bd196c6f5d8fbbd50c26855bc269b68c69 Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 21 Sep 2026 22:27:44 +0000 Subject: [PATCH 2/6] chore(ui): regenerate api types for video call types Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 0b09f3654dc..7f0b692049c 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -25794,7 +25794,7 @@ export interface components { * CallTypes * @enum {string} */ - CallTypes: "embedding" | "aembedding" | "completion" | "acompletion" | "atext_completion" | "text_completion" | "image_generation" | "aimage_generation" | "image_edit" | "aimage_edit" | "moderation" | "amoderation" | "atranscription" | "transcription" | "aspeech" | "speech" | "rerank" | "arerank" | "search" | "asearch" | "_arealtime" | "_aresponses_websocket" | "create_batch" | "acreate_batch" | "aretrieve_batch" | "retrieve_batch" | "acancel_batch" | "cancel_batch" | "pass_through_endpoint" | "anthropic_messages" | "aanthropic_messages" | "get_assistants" | "aget_assistants" | "create_assistants" | "acreate_assistants" | "delete_assistant" | "adelete_assistant" | "acreate_thread" | "create_thread" | "aget_thread" | "get_thread" | "a_add_message" | "add_message" | "aget_messages" | "get_messages" | "arun_thread" | "run_thread" | "arun_thread_stream" | "run_thread_stream" | "afile_retrieve" | "file_retrieve" | "afile_delete" | "file_delete" | "afile_list" | "file_list" | "acreate_file" | "create_file" | "afile_content" | "file_content" | "create_fine_tuning_job" | "acreate_fine_tuning_job" | "create_video" | "acreate_video" | "avideo_retrieve" | "video_retrieve" | "avideo_content" | "video_content" | "video_remix" | "avideo_remix" | "video_list" | "avideo_list" | "video_retrieve_job" | "avideo_retrieve_job" | "video_delete" | "avideo_delete" | "video_create_character" | "avideo_create_character" | "video_get_character" | "avideo_get_character" | "video_edit" | "avideo_edit" | "video_extension" | "avideo_extension" | "vector_store_file_create" | "avector_store_file_create" | "vector_store_file_list" | "avector_store_file_list" | "vector_store_file_retrieve" | "avector_store_file_retrieve" | "vector_store_file_content" | "avector_store_file_content" | "vector_store_file_update" | "avector_store_file_update" | "vector_store_file_delete" | "avector_store_file_delete" | "vector_store_create" | "avector_store_create" | "vector_store_search" | "avector_store_search" | "ingest" | "aingest" | "query" | "aquery" | "create_interaction" | "acreate_interaction" | "create_container" | "acreate_container" | "list_containers" | "alist_containers" | "retrieve_container" | "aretrieve_container" | "delete_container" | "adelete_container" | "list_container_files" | "alist_container_files" | "upload_container_file" | "aupload_container_file" | "create_sandbox" | "acreate_sandbox" | "delete_sandbox" | "adelete_sandbox" | "run_code" | "arun_code" | "code_interpreter_tool" | "acode_interpreter_tool" | "acancel_fine_tuning_job" | "cancel_fine_tuning_job" | "alist_fine_tuning_jobs" | "list_fine_tuning_jobs" | "aretrieve_fine_tuning_job" | "retrieve_fine_tuning_job" | "responses" | "aresponses" | "alist_input_items" | "llm_passthrough_route" | "allm_passthrough_route" | "generate_content" | "agenerate_content" | "generate_content_stream" | "agenerate_content_stream" | "ocr" | "aocr" | "call_mcp_tool" | "list_mcp_tools" | "asend_message" | "send_message" | "acreate_skill"; + CallTypes: "embedding" | "aembedding" | "completion" | "acompletion" | "atext_completion" | "text_completion" | "image_generation" | "aimage_generation" | "image_edit" | "aimage_edit" | "moderation" | "amoderation" | "atranscription" | "transcription" | "aspeech" | "speech" | "rerank" | "arerank" | "search" | "asearch" | "_arealtime" | "_aresponses_websocket" | "create_batch" | "acreate_batch" | "aretrieve_batch" | "retrieve_batch" | "acancel_batch" | "cancel_batch" | "pass_through_endpoint" | "anthropic_messages" | "aanthropic_messages" | "get_assistants" | "aget_assistants" | "create_assistants" | "acreate_assistants" | "delete_assistant" | "adelete_assistant" | "acreate_thread" | "create_thread" | "aget_thread" | "get_thread" | "a_add_message" | "add_message" | "aget_messages" | "get_messages" | "arun_thread" | "run_thread" | "arun_thread_stream" | "run_thread_stream" | "afile_retrieve" | "file_retrieve" | "afile_delete" | "file_delete" | "afile_list" | "file_list" | "acreate_file" | "create_file" | "afile_content" | "file_content" | "create_fine_tuning_job" | "acreate_fine_tuning_job" | "create_video" | "acreate_video" | "video_generation" | "avideo_generation" | "avideo_retrieve" | "video_retrieve" | "avideo_content" | "video_content" | "video_remix" | "avideo_remix" | "video_list" | "avideo_list" | "video_retrieve_job" | "avideo_retrieve_job" | "video_delete" | "avideo_delete" | "video_create_character" | "avideo_create_character" | "video_get_character" | "avideo_get_character" | "video_edit" | "avideo_edit" | "video_extension" | "avideo_extension" | "vector_store_file_create" | "avector_store_file_create" | "vector_store_file_list" | "avector_store_file_list" | "vector_store_file_retrieve" | "avector_store_file_retrieve" | "vector_store_file_content" | "avector_store_file_content" | "vector_store_file_update" | "avector_store_file_update" | "vector_store_file_delete" | "avector_store_file_delete" | "vector_store_create" | "avector_store_create" | "vector_store_search" | "avector_store_search" | "ingest" | "aingest" | "query" | "aquery" | "create_interaction" | "acreate_interaction" | "create_container" | "acreate_container" | "list_containers" | "alist_containers" | "retrieve_container" | "aretrieve_container" | "delete_container" | "adelete_container" | "list_container_files" | "alist_container_files" | "upload_container_file" | "aupload_container_file" | "create_sandbox" | "acreate_sandbox" | "delete_sandbox" | "adelete_sandbox" | "run_code" | "arun_code" | "code_interpreter_tool" | "acode_interpreter_tool" | "acancel_fine_tuning_job" | "cancel_fine_tuning_job" | "alist_fine_tuning_jobs" | "list_fine_tuning_jobs" | "aretrieve_fine_tuning_job" | "retrieve_fine_tuning_job" | "responses" | "aresponses" | "alist_input_items" | "llm_passthrough_route" | "allm_passthrough_route" | "generate_content" | "agenerate_content" | "generate_content_stream" | "agenerate_content_stream" | "ocr" | "aocr" | "call_mcp_tool" | "list_mcp_tools" | "asend_message" | "send_message" | "acreate_skill"; /** CallbackDelete */ CallbackDelete: { /** Callback Name */ From a44befa8c529465c331cecd373f1f38d810927ee Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 21 Sep 2026 22:33:09 +0000 Subject: [PATCH 3/6] test: skip avideo_generation in azure sdk client exhaustive check Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/llms/azure/test_azure_common_utils.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/test_litellm/llms/azure/test_azure_common_utils.py b/tests/test_litellm/llms/azure/test_azure_common_utils.py index 83ec85f1176..caf941ebd19 100644 --- a/tests/test_litellm/llms/azure/test_azure_common_utils.py +++ b/tests/test_litellm/llms/azure/test_azure_common_utils.py @@ -597,7 +597,8 @@ async def test_ensure_initialize_azure_sdk_client_always_used(call_type): "litellm.files.main.azure_files_instance.initialize_azure_sdk_client" ) elif ( - call_type == CallTypes.avideo_content + call_type == CallTypes.avideo_generation + or call_type == CallTypes.avideo_content or call_type == CallTypes.avideo_list or call_type == CallTypes.avideo_remix or call_type == CallTypes.avideo_create_character From dfc5ef70b4259bd5f189af1f04d11a9d9ffb8a41 Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 21 Sep 2026 22:44:05 +0000 Subject: [PATCH 4/6] test(e2e): retry a leaked video job until the guardrail sync deadline Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/guardrails/test_key_guardrail_video_e2e.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/e2e/guardrails/test_key_guardrail_video_e2e.py b/tests/e2e/guardrails/test_key_guardrail_video_e2e.py index 8c70ea5b46f..100e9af5a5f 100644 --- a/tests/e2e/guardrails/test_key_guardrail_video_e2e.py +++ b/tests/e2e/guardrails/test_key_guardrail_video_e2e.py @@ -71,7 +71,7 @@ class TestKeyAttachedGuardrailOnVideos: f"block response missing content-filter reason: {body[:300]}" ) return - case Success(data=video): + 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}" From 652bddfdc60a1369ae5416f92f957e49d38864df Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 21 Sep 2026 22:55:01 +0000 Subject: [PATCH 5/6] refactor(guardrails): satisfy the type-discipline gate in the video handler Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../videos/guardrail_translation/__init__.py | 4 ++-- .../videos/guardrail_translation/handler.py | 16 +++++++++------- 2 files changed, 11 insertions(+), 9 deletions(-) diff --git a/litellm/llms/openai/videos/guardrail_translation/__init__.py b/litellm/llms/openai/videos/guardrail_translation/__init__.py index 6540ba994e4..7bd869612d6 100644 --- a/litellm/llms/openai/videos/guardrail_translation/__init__.py +++ b/litellm/llms/openai/videos/guardrail_translation/__init__.py @@ -7,7 +7,7 @@ from litellm.llms.openai.videos.guardrail_translation.handler import ( ) from litellm.types.utils import CallTypes -guardrail_translation_mappings: Final = { +guardrail_translation_mappings: Final = { # mutable-ok: discover_guardrail_translation_mappings only accepts isinstance(mappings, dict) CallTypes.video_generation: OpenAIVideoGenerationHandler, CallTypes.avideo_generation: OpenAIVideoGenerationHandler, CallTypes.create_video: OpenAIVideoGenerationHandler, @@ -20,4 +20,4 @@ guardrail_translation_mappings: Final = { CallTypes.avideo_extension: OpenAIVideoGenerationHandler, } -__all__ = ["OpenAIVideoGenerationHandler", "guardrail_translation_mappings"] +__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 index 74fafb4bcfe..3735094c87b 100644 --- a/litellm/llms/openai/videos/guardrail_translation/handler.py +++ b/litellm/llms/openai/videos/guardrail_translation/handler.py @@ -14,19 +14,20 @@ class OpenAIVideoGenerationHandler(BaseTranslation): async def process_input_messages( self, - data: dict[str, object], + data: dict[str, object], # mutable-ok: BaseTranslation contract passes the proxy's request dict through guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: "LiteLLMLoggingObj | None" = None, - ) -> dict[str, object]: + ) -> dict[str, object]: # mutable-ok: BaseTranslation contract returns the proxy's request dict prompt: Final = data.get("prompt") if not isinstance(prompt, str): return data model: Final = data.get("model") + texts: Final = [prompt] # mutable-ok: GenericGuardrailAPIInputs.texts is declared list[str] inputs: Final = ( - GenericGuardrailAPIInputs(texts=[prompt], model=model) + GenericGuardrailAPIInputs(texts=texts, model=model) if isinstance(model, str) - else GenericGuardrailAPIInputs(texts=[prompt]) + else GenericGuardrailAPIInputs(texts=texts) ) guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail( # pyright: ignore[reportUnknownMemberType] # request_data is a bare dict inputs=inputs, @@ -34,8 +35,9 @@ class OpenAIVideoGenerationHandler(BaseTranslation): 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} + guardrailed_texts: Final = guardrailed_inputs.get("texts") + guardrailed_prompt: Final = guardrailed_texts[0] if guardrailed_texts else prompt + return {**data, "prompt": guardrailed_prompt} # mutable-ok: BaseTranslation contract returns a dict async def process_output_response( self, @@ -43,6 +45,6 @@ class OpenAIVideoGenerationHandler(BaseTranslation): guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: "LiteLLMLoggingObj | None" = None, user_api_key_dict: "UserAPIKeyAuth | None" = None, - request_data: dict[str, object] | None = None, + request_data: dict[str, object] | None = None, # mutable-ok: BaseTranslation contract ) -> object: return response From 2f0584cec6cd6e6ce445db254ff9473b0d01d306 Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 21 Sep 2026 23:20:29 +0000 Subject: [PATCH 6/6] test(guardrails): gate the video e2e on a chat probe so a miss starts at most one paid job 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> --- .../videos/guardrail_translation/handler.py | 2 - tests/e2e/guardrails/guardrails_client.py | 2 - .../test_key_guardrail_video_e2e.py | 58 +++---- .../test_unified_guardrail.py | 160 ++++++------------ 4 files changed, 73 insertions(+), 149 deletions(-) diff --git a/litellm/llms/openai/videos/guardrail_translation/handler.py b/litellm/llms/openai/videos/guardrail_translation/handler.py index 3735094c87b..49a8d05100c 100644 --- a/litellm/llms/openai/videos/guardrail_translation/handler.py +++ b/litellm/llms/openai/videos/guardrail_translation/handler.py @@ -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 diff --git a/tests/e2e/guardrails/guardrails_client.py b/tests/e2e/guardrails/guardrails_client.py index 11efaf0296e..3ceac737399 100644 --- a/tests/e2e/guardrails/guardrails_client.py +++ b/tests/e2e/guardrails/guardrails_client.py @@ -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)) ) diff --git a/tests/e2e/guardrails/test_key_guardrail_video_e2e.py b/tests/e2e/guardrails/test_key_guardrail_video_e2e.py index 100e9af5a5f..5f318e141a5 100644 --- a/tests/e2e/guardrails/test_key_guardrail_video_e2e.py +++ b/tests/e2e/guardrails/test_key_guardrail_video_e2e.py @@ -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}") 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 138b6d739f7..c90f88ec110 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 @@ -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 ]