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 ]