scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-06 08:16:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 79

fix(guardrails): configure Prompt Security file timeout policy (#38083)

Browse source 
* fix(guardrails): fail open on Prompt Security file timeouts

* fix(guardrails): configure Prompt Security timeout policy
This commit is contained in:
davida-ps 2026-08-31 22:05:57 +03:00 • committed by GitHub
parent c09fa5b712
commit 9f67a58198
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 148 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
litellm/proxy/guardrails/guardrail_hooks/prompt_security/__init__.py
View file

@ -20,6 +20,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"
guardrail_name=guardrail.get("guardrail_name", ""),
event_hook=litellm_params.mode,
default_on=litellm_params.default_on,
file_sanitization_fail_open=getattr(litellm_params, "file_sanitization_fail_open", None),
)
litellm.logging_callback_manager.add_litellm_callback(_prompt_security_callback)

49
litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py
View file

@ -4,10 +4,12 @@ import os
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Final, Literal, Optional
import httpx
from fastapi import HTTPException
from typing_extensions import ReadOnly, TypedDict
from litellm._logging import verbose_proxy_logger
from litellm.exceptions import Timeout as LiteLLMTimeout
from litellm.integrations.custom_guardrail import (
CustomGuardrail,
log_guardrail_information,
@ -24,6 +26,9 @@ if TYPE_CHECKING:
from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel
_SANITIZE_FILE_FAIL_OPEN_TIMEOUT_SECONDS: Final = 30.0
class PromptSecurityGuardrailMissingSecrets(Exception):
pass
@ -63,6 +68,13 @@ class _SanitizeStatusResponse(TypedDict, total=False):
metadata: ReadOnly[_SanitizeMetadata]
class _SanitizeResult(TypedDict):
action: ReadOnly[str]
content: ReadOnly[str | None]
metadata: ReadOnly[_SanitizeMetadata]
violations: ReadOnly[Sequence[str]]
class PromptSecurityGuardrail(CustomGuardrail):
@classmethod
def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]:
@ -79,6 +91,8 @@ class PromptSecurityGuardrail(CustomGuardrail):
user: str | None = None,
system_prompt: str | None = None,
check_tool_results: bool | None = None,
file_sanitization_timeout: float = _SANITIZE_FILE_FAIL_OPEN_TIMEOUT_SECONDS,
file_sanitization_fail_open: bool | None = None,
**kwargs,
):
kwargs.setdefault("supported_event_hooks", list(self.get_supported_event_hooks()))
@ -108,6 +122,8 @@ class PromptSecurityGuardrail(CustomGuardrail):
# Configuration for file sanitization
self.max_poll_attempts = 30 # Maximum number of polling attempts
self.poll_interval = 2 # Seconds between polling attempts
self.file_sanitization_timeout = file_sanitization_timeout
self.file_sanitization_fail_open = file_sanitization_fail_open is not False
super().__init__(**kwargs)
@ -397,6 +413,39 @@ class PromptSecurityGuardrail(CustomGuardrail):
Sanitize file content using Prompt Security API.
Returns: dict with keys 'action', 'content', 'metadata'
"""
try:
return await asyncio.wait_for(
self._sanitize_file_content(file_data, filename, user_api_key_alias),
timeout=self.file_sanitization_timeout,
)
except (asyncio.TimeoutError, httpx.TimeoutException, LiteLLMTimeout) as exc:
if not self.file_sanitization_fail_open:
verbose_proxy_logger.error(
"Prompt Security Guardrail: file sanitization for %s timed out with %s; failing closed",
filename,
type(exc).__name__,
)
raise HTTPException(status_code=408, detail="File sanitization timeout") from exc
verbose_proxy_logger.error(
"Prompt Security Guardrail: file sanitization for %s timed out with %s; failing open",
filename,
type(exc).__name__,
)
fail_open_result: Final[_SanitizeResult] = {
"action": "allow",
"content": None,
"metadata": {},
"violations": (),
}
return fail_open_result
async def _sanitize_file_content(
self,
file_data: bytes,
filename: str,
user_api_key_alias: str | None,
) -> _SanitizeResult:
headers: Final = {"APP-ID": self.api_key}
if user_api_key_alias:
headers["X-LiteLLM-Key-Alias"] = user_api_key_alias

4
litellm/types/proxy/guardrails/guardrail_hooks/prompt_security.py
View file

@ -12,6 +12,10 @@ class PromptSecurityGuardrailConfigModel(GuardrailConfigModel):
default=None,
description="The API base for the Prompt Security guardrail. If not provided, the `PROMPT_SECURITY_API_BASE` environment variable is used.",
)
file_sanitization_fail_open: bool = Field(
default=True,
description="Whether file sanitization timeouts allow the original file through instead of blocking the request.",
)
@staticmethod
def ui_friendly_name() -> str:

103
tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py
View file

@ -1,16 +1,16 @@
from fastapi.exceptions import HTTPException
from unittest.mock import patch, AsyncMock
from httpx import Response, Request
import asyncio
import base64
from unittest.mock import AsyncMock, patch
import pytest
from litellm.proxy.guardrails.guardrail_hooks.prompt_security.prompt_security import (
PromptSecurityGuardrailMissingSecrets,
PromptSecurityGuardrail,
)
from fastapi.exceptions import HTTPException
from httpx import ReadTimeout, Request, Response
import litellm
from litellm.proxy.guardrails.guardrail_hooks.prompt_security.prompt_security import (
PromptSecurityGuardrail,
PromptSecurityGuardrailMissingSecrets,
)
from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2
@ -30,6 +30,7 @@ def test_prompt_security_guard_config(monkeypatch: pytest.MonkeyPatch):
"guardrail": "prompt_security",
"mode": "during_call",
"default_on": True,
"file_sanitization_fail_open": False,
},
}
],
@ -41,6 +42,10 @@ def test_prompt_security_guard_config(monkeypatch: pytest.MonkeyPatch):
assert registered[0].guardrail_name == "prompt_security"
assert registered[0].default_on is True
assert registered[0].event_hook == "during_call"
assert registered[0].file_sanitization_fail_open is False
config_model = registered[0].get_config_model()
assert config_model is not None
assert config_model().file_sanitization_fail_open is True
def test_prompt_security_guard_config_no_api_key(monkeypatch: pytest.MonkeyPatch):
@ -374,6 +379,86 @@ async def test_file_sanitization(monkeypatch: pytest.MonkeyPatch):
assert result is not None
@pytest.mark.asyncio
@pytest.mark.parametrize(
"timeout",
(
litellm.Timeout(
message="Prompt Security upload timed out",
model="default-model-name",
llm_provider="litellm-httpx-handler",
),
ReadTimeout(
"Prompt Security poll timed out",
request=Request(method="GET", url="https://test.prompt.security/api/sanitizeFile"),
),
),
ids=("litellm", "httpx"),
)
@pytest.mark.parametrize("fail_open", (True, False), ids=("fail-open", "fail-closed"))
async def test_file_sanitization_request_timeout_policy(
monkeypatch: pytest.MonkeyPatch, timeout: Exception, fail_open: bool
):
monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key")
monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security")
guardrail = PromptSecurityGuardrail(
guardrail_name="test-guard",
event_hook="pre_call",
default_on=True,
file_sanitization_fail_open=fail_open,
)
with patch.object(guardrail.async_handler, "post", AsyncMock(side_effect=timeout)):
if not fail_open:
with pytest.raises(HTTPException) as exc_info:
await guardrail.sanitize_file_content(b"file-content", "document.pdf")
assert exc_info.value.status_code == 408
assert exc_info.value.detail == "File sanitization timeout"
return
result = await guardrail.sanitize_file_content(b"file-content", "document.pdf")
assert result == {
"action": "allow",
"content": None,
"metadata": {},
"violations": (),
}
@pytest.mark.asyncio
@pytest.mark.parametrize("fail_open", (True, False), ids=("fail-open", "fail-closed"))
async def test_file_sanitization_overall_timeout_policy(monkeypatch: pytest.MonkeyPatch, fail_open: bool):
monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key")
monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security")
guardrail = PromptSecurityGuardrail(
guardrail_name="test-guard",
event_hook="pre_call",
default_on=True,
file_sanitization_timeout=0.01,
file_sanitization_fail_open=fail_open,
)
async def hanging_post(*_args: object, **_kwargs: object) -> None:
await asyncio.sleep(60)
raise AssertionError("sanitization request should have been cancelled")
with patch.object(guardrail.async_handler, "post", side_effect=hanging_post):
if not fail_open:
with pytest.raises(HTTPException) as exc_info:
await guardrail.sanitize_file_content(b"file-content", "document.pdf")
assert exc_info.value.status_code == 408
assert exc_info.value.detail == "File sanitization timeout"
return
result = await guardrail.sanitize_file_content(b"file-content", "document.pdf")
assert result["action"] == "allow"
assert result["content"] is None
@pytest.mark.asyncio
async def test_file_sanitization_block(monkeypatch: pytest.MonkeyPatch):
"""Test that file sanitization blocks malicious files"""
@ -544,7 +629,7 @@ async def test_role_filtering(monkeypatch: pytest.MonkeyPatch):
return mock_response
with patch.object(guardrail.async_handler, "post", side_effect=mock_post):
result = await guardrail.apply_guardrail(
await guardrail.apply_guardrail(
inputs=inputs,
request_data=request_data,
input_type="request",