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

fix(prompt_security): poll only on queued statuses, keep 500 for terminal or missing status

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yucheng 2026-09-14 22:01:39 +00:00
parent fb60f80c80
commit 3a3075b8a2
2 changed files with 33 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -27,6 +27,7 @@ if TYPE_CHECKING:
_SANITIZE_FILE_FAIL_OPEN_TIMEOUT_SECONDS: Final = 30.0
_SANITIZE_FILE_QUEUED_STATUSES: Final = frozenset({"created", "in progress"})
class PromptSecurityGuardrailMissingSecrets(Exception):
@ -513,6 +514,9 @@ class PromptSecurityGuardrail(CustomGuardrail):
"violations": result.get("metadata", {}).get("violations", []),
}
if status not in _SANITIZE_FILE_QUEUED_STATUSES:
raise HTTPException(status_code=500, detail=f"Unexpected sanitization status: {status}")
verbose_proxy_logger.debug(
"Prompt Security Guardrail: File sanitization status=%s for jobId=%s (attempt %d/%d)",
status,

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

@ -560,6 +560,35 @@ async def test_file_sanitization_never_finishing_job_times_out(monkeypatch: pyte
assert exc_info.value.detail == "File sanitization timeout"
@pytest.mark.asyncio
@pytest.mark.parametrize("poll_body", [{"status": "failed"}, {}])
async def test_file_sanitization_terminal_failure_does_not_fail_open(monkeypatch: pytest.MonkeyPatch, poll_body):
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
upload_response = Response(
json={"jobId": "failed-job"},
status_code=200,
request=Request(method="POST", url="https://test.prompt.security/api/sanitizeFile"),
)
poll_response = Response(
json=poll_body,
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)) as poll_mock:
with pytest.raises(HTTPException) as exc_info:
await guardrail.sanitize_file_content(b"file-content", "document.pdf")
assert poll_mock.await_count == 1
assert exc_info.value.status_code == 500
assert exc_info.value.detail == f"Unexpected sanitization status: {poll_body.get('status')}"
@pytest.mark.asyncio
@pytest.mark.parametrize(
"timeout",