scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-08 22:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 78

fix(guardrails): block Prompt Security file modifications

Browse source
This commit is contained in:
David Abutbul 2026-08-25 14:28:53 +03:00
parent 3bba1f66a3
commit fa15acf355
4 changed files with 140 additions and 22 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

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

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

@ -93,6 +93,7 @@ class PromptSecurityGuardrail(CustomGuardrail):
check_tool_results: bool | None = None,
file_sanitization_timeout: float = _SANITIZE_FILE_FAIL_OPEN_TIMEOUT_SECONDS,
file_sanitization_fail_open: bool | None = None,
block_on_file_modify: bool | None = None,
**kwargs,
):
kwargs.setdefault("supported_event_hooks", list(self.get_supported_event_hooks()))
@ -124,6 +125,7 @@ class PromptSecurityGuardrail(CustomGuardrail):
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
self.block_on_file_modify = block_on_file_modify is not False
super().__init__(**kwargs)
@ -372,13 +374,7 @@ class PromptSecurityGuardrail(CustomGuardrail):
result = await self.sanitize_file_content(
file_data, filename, user_api_key_alias=user_api_key_alias
)
if result.get("action") == "block":
violations = result.get("violations", [])
raise HTTPException(
status_code=400,
detail=f"Image blocked by Prompt Security. Violations: {', '.join(violations)}",
)
self._raise_if_file_blocked(result, "Image")
except HTTPException:
raise
except Exception as e:
@ -408,7 +404,7 @@ class PromptSecurityGuardrail(CustomGuardrail):
file_data: bytes,
filename: str,
user_api_key_alias: str | None = None,
) -> dict:
) -> _SanitizeResult:
"""
Sanitize file content using Prompt Security API.
Returns: dict with keys 'action', 'content', 'metadata'
@ -528,6 +524,17 @@ class PromptSecurityGuardrail(CustomGuardrail):
raise HTTPException(status_code=408, detail="File sanitization timeout")
def _raise_if_file_blocked(self, sanitization_result: _SanitizeResult, resource_name: str) -> None:
action: Final = sanitization_result.get("action")
if action != "block" and not (action == "modify" and self.block_on_file_modify):
return
violations: Final = sanitization_result.get("violations", ())
raise HTTPException(
status_code=400,
detail=f"{resource_name} blocked by Prompt Security. Violations: {', '.join(violations)}",
)
async def _process_image_url_item(self, item: dict, user_api_key_alias: str | None) -> dict:
"""Process and sanitize image_url items."""
image_url_data: Final = item.get("image_url", {})
@ -547,13 +554,7 @@ class PromptSecurityGuardrail(CustomGuardrail):
file_data, filename, user_api_key_alias=user_api_key_alias
)
action: Final = sanitization_result.get("action")
if action == "block":
violations: Final = sanitization_result.get("violations", [])
raise HTTPException(
status_code=400,
detail=f"File blocked by Prompt Security. Violations: {', '.join(violations)}",
)
self._raise_if_file_blocked(sanitization_result, "File")
if action == "modify":
sanitized_content: Final = sanitization_result.get("content", "")
@ -615,13 +616,7 @@ class PromptSecurityGuardrail(CustomGuardrail):
file_data, filename, user_api_key_alias=user_api_key_alias
)
action: Final = sanitization_result.get("action")
if action == "block":
violations: Final = sanitization_result.get("violations", [])
raise HTTPException(
status_code=400,
detail=f"Document blocked by Prompt Security. Violations: {', '.join(violations)}",
)
self._raise_if_file_blocked(sanitization_result, "Document")
if action == "modify":
sanitized_content: Final = sanitization_result.get("content", "")

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

@ -16,6 +16,10 @@ class PromptSecurityGuardrailConfigModel(GuardrailConfigModel):
default=True,
description="Whether file sanitization timeouts allow the original file through instead of blocking the request.",
)
block_on_file_modify: bool = Field(
default=True,
description="Whether a file sanitization `modify` verdict blocks the request instead of replacing the file content.",
)
@staticmethod
def ui_friendly_name() -> str:

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

@ -31,6 +31,7 @@ def test_prompt_security_guard_config(monkeypatch: pytest.MonkeyPatch):
"mode": "during_call",
"default_on": True,
"file_sanitization_fail_open": False,
"block_on_file_modify": False,
},
}
],
@ -43,9 +44,11 @@ def test_prompt_security_guard_config(monkeypatch: pytest.MonkeyPatch):
assert registered[0].default_on is True
assert registered[0].event_hook == "during_call"
assert registered[0].file_sanitization_fail_open is False
assert registered[0].block_on_file_modify is False
config_model = registered[0].get_config_model()
assert config_model is not None
assert config_model().file_sanitization_fail_open is True
assert config_model().block_on_file_modify is True
def test_prompt_security_guard_config_no_api_key(monkeypatch: pytest.MonkeyPatch):
@ -379,6 +382,121 @@ async def test_file_sanitization(monkeypatch: pytest.MonkeyPatch):
assert result is not None
@pytest.mark.asyncio
async def test_file_sanitization_modify_blocks_by_default(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
)
csv_data = b"name,email\nAlice,alice@example.com\n"
item = {
"type": "file",
"file": {
"data": base64.b64encode(csv_data).decode(),
"mime_type": "text/csv",
},
}
upload_response = Response(
json={"jobId": "modify-job"},
status_code=200,
request=Request(method="POST", url="https://test.prompt.security/api/sanitizeFile"),
)
poll_response = Response(
json={
"status": "done",
"content": "name,email\nAlice,[REDACTED]\n",
"metadata": {"action": "modify", "violations": ["Sensitive Data"]},
},
status_code=200,
request=Request(method="GET", url="https://test.prompt.security/api/sanitizeFile"),
)
with patch.object(guardrail.async_handler, "post", AsyncMock(return_value=upload_response)):
with patch.object(guardrail.async_handler, "get", AsyncMock(return_value=poll_response)):
with pytest.raises(HTTPException) as exc_info:
await guardrail._process_document_item(item, None)
assert exc_info.value.status_code == 400
assert exc_info.value.detail == "Document blocked by Prompt Security. Violations: Sensitive Data"
@pytest.mark.asyncio
async def test_standalone_image_sanitization_modify_blocks_by_default(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
)
guardrail.poll_interval = 0
image_url = "data:image/png;base64," + base64.b64encode(b"image-content").decode()
upload_response = Response(
json={"jobId": "modify-image-job"},
status_code=200,
request=Request(method="POST", url="https://test.prompt.security/api/sanitizeFile"),
)
poll_response = Response(
json={
"status": "done",
"content": "Email: [REDACTED]",
"metadata": {"action": "modify", "violations": ["Sensitive Data"]},
},
status_code=200,
request=Request(method="GET", url="https://test.prompt.security/api/sanitizeFile"),
)
with patch.object(guardrail.async_handler, "post", AsyncMock(return_value=upload_response)):
with patch.object(guardrail.async_handler, "get", AsyncMock(return_value=poll_response)):
with pytest.raises(HTTPException) as exc_info:
await guardrail._process_standalone_images([image_url], None)
assert exc_info.value.status_code == 400
assert exc_info.value.detail == "Image blocked by Prompt Security. Violations: Sensitive Data"
@pytest.mark.asyncio
async def test_file_sanitization_modify_can_rewrite_when_blocking_disabled(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,
block_on_file_modify=False,
)
csv_data = b"name,email\nAlice,alice@example.com\n"
item = {
"type": "file",
"file": {
"data": base64.b64encode(csv_data).decode(),
"mime_type": "text/csv",
},
}
upload_response = Response(
json={"jobId": "modify-job"},
status_code=200,
request=Request(method="POST", url="https://test.prompt.security/api/sanitizeFile"),
)
poll_response = Response(
json={
"status": "done",
"content": "name,email\nAlice,[REDACTED]\n",
"metadata": {"action": "modify", "violations": ["Sensitive Data"]},
},
status_code=200,
request=Request(method="GET", url="https://test.prompt.security/api/sanitizeFile"),
)
with patch.object(guardrail.async_handler, "post", AsyncMock(return_value=upload_response)):
with patch.object(guardrail.async_handler, "get", AsyncMock(return_value=poll_response)):
result = await guardrail._process_document_item(item, None)
assert base64.b64decode(result["file"]["data"]) == b"name,email\nAlice,[REDACTED]\n"
@pytest.mark.asyncio
@pytest.mark.parametrize(
"timeout",