diff --git a/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py b/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py index 809d5e0fb31..178f89814db 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py +++ b/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py @@ -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,7 @@ 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, **kwargs, ): kwargs.setdefault("supported_event_hooks", list(self.get_supported_event_hooks())) @@ -108,6 +121,7 @@ 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 super().__init__(**kwargs) @@ -397,6 +411,31 @@ 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: + 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 diff --git a/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py b/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py index 26beaa78a46..49c4e3147ec 100644 --- a/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py @@ -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 @@ -374,6 +374,62 @@ 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"), +) +async def test_file_sanitization_fails_open_on_request_timeout(monkeypatch: pytest.MonkeyPatch, timeout: Exception): + 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) + + with patch.object(guardrail.async_handler, "post", AsyncMock(side_effect=timeout)): + result = await guardrail.sanitize_file_content(b"file-content", "document.pdf") + + assert result == { + "action": "allow", + "content": None, + "metadata": {}, + "violations": (), + } + + +@pytest.mark.asyncio +async def test_file_sanitization_fails_open_on_overall_timeout(monkeypatch: pytest.MonkeyPatch): + 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, + ) + + async def hanging_post(*args, **kwargs): + await asyncio.sleep(60) + raise AssertionError("sanitization request should have been cancelled") + + with patch.object(guardrail.async_handler, "post", side_effect=hanging_post): + 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 +600,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",