mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-24 00:52:24 +00:00
Merge pull request #42354 from BerriAI/litellm_fix_video_key_guardrails_lit_6685
* 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> * chore(ui): regenerate api types for video call types Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test: skip avideo_generation in azure sdk client exhaustive check Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * 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> * 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> * 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> --------- Co-authored-by: yucheng <yucheng@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yucheng-berri
GitHub
commit
12f79308cf
10 changed files with 269 additions and 108 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 = { # mutable-ok: discover_guardrail_translation_mappings only accepts isinstance(mappings, dict)
|
||||
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):
|
||||
async def process_input_messages(
|
||||
self,
|
||||
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]: # 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=texts, model=model)
|
||||
if isinstance(model, str)
|
||||
else GenericGuardrailAPIInputs(texts=texts)
|
||||
)
|
||||
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")
|
||||
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,
|
||||
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 contract
|
||||
) -> 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,21 @@ 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:
|
||||
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,
|
||||
|
|
|
|||
69
tests/e2e/guardrails/test_key_guardrail_video_e2e.py
Normal file
69
tests/e2e/guardrails/test_key_guardrail_video_e2e.py
Normal file
|
|
@ -0,0 +1,69 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
from e2e_config import unique_marker
|
||||
from e2e_http import Success, UnknownApiError
|
||||
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"
|
||||
|
||||
|
||||
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)
|
||||
|
||||
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}")
|
||||
|
|
@ -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 ----------
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@
|
|||
|
||||
import logging
|
||||
from types import SimpleNamespace
|
||||
from typing import Final
|
||||
from typing import TYPE_CHECKING, Final, Literal
|
||||
|
||||
import pytest
|
||||
|
||||
|
|
@ -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,
|
||||
|
|
@ -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):
|
||||
|
|
@ -61,6 +64,18 @@ class RecordingGuardrail(CustomGuardrail):
|
|||
return {"texts": inputs.get("texts", [])}
|
||||
|
||||
|
||||
class RewritingGuardrail(RecordingGuardrail):
|
||||
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):
|
||||
"""Test translation handler that simply echoes input/output."""
|
||||
|
||||
|
|
@ -115,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
|
||||
|
|
@ -130,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
|
||||
|
||||
|
|
@ -169,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
|
||||
|
||||
|
|
@ -201,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:
|
||||
|
|
@ -229,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
|
||||
|
|
@ -248,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
|
||||
|
||||
|
|
@ -299,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
|
||||
|
||||
|
|
@ -331,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:
|
||||
|
|
@ -360,6 +341,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):
|
||||
|
|
@ -424,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(
|
||||
|
|
@ -493,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
|
||||
|
|
@ -546,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"]
|
||||
|
|
@ -592,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(),
|
||||
)
|
||||
|
||||
|
|
@ -622,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(),
|
||||
)
|
||||
|
||||
|
|
@ -734,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):
|
||||
|
|
@ -770,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):
|
||||
|
|
@ -789,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]}
|
||||
|
||||
|
|
@ -1538,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
|
||||
|
|
@ -1685,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(
|
||||
|
|
@ -1985,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
|
||||
|
|
@ -2011,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
|
||||
|
|
@ -2076,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 = []
|
||||
|
|
@ -2407,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
|
||||
]
|
||||
|
|
|
|||
2
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
2
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -25835,7 +25835,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 */
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue