mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-10 22:41:41 +00:00
Address Greptile review: redact path, null decision, error context
- P1: Filter _extract_texts_from_messages to user-role messages only,
preventing system/assistant content from being injected into texts
- P1: Strengthen test_redact_updates_structured_messages assertion from
weak `in` check to strict equality, catching the injection bug
- P2: Use `result.get("decision") or "allow"` to handle explicit null
decision values (not just absent keys)
- P2: Wrap bare exception re-raise in GuardrailRaisedException so the
caller knows which guardrail failed (block_on_error=True path)
- P2: Add static Promptguard entry in guardrail_provider_map so the
preset works before populateGuardrailProviderMap is called
- Add test for explicit null decision treated as allow
This commit is contained in:
parent
20a57507f5
commit
ceb5b5df8c
3 changed files with 42 additions and 8 deletions
|
|
@ -155,7 +155,13 @@ class PromptGuardGuardrail(CustomGuardrail):
|
|||
except Exception as exc:
|
||||
verbose_proxy_logger.error("PromptGuard API error: %s", str(exc))
|
||||
if self.block_on_error:
|
||||
raise
|
||||
raise GuardrailRaisedException(
|
||||
guardrail_name=self.guardrail_name,
|
||||
message=(
|
||||
f"PromptGuard API unreachable "
|
||||
f"(block_on_error=True): {exc}"
|
||||
),
|
||||
) from exc
|
||||
return inputs
|
||||
|
||||
verbose_proxy_logger.debug(
|
||||
|
|
@ -164,7 +170,7 @@ class PromptGuardGuardrail(CustomGuardrail):
|
|||
result.get("threat_type"),
|
||||
)
|
||||
|
||||
decision = result.get("decision", "allow")
|
||||
decision = result.get("decision") or "allow"
|
||||
|
||||
if decision == "block":
|
||||
threat_type = result.get("threat_type", "unknown")
|
||||
|
|
@ -196,9 +202,16 @@ class PromptGuardGuardrail(CustomGuardrail):
|
|||
|
||||
@staticmethod
|
||||
def _extract_texts_from_messages(messages: list) -> List[str]:
|
||||
"""Extract text content strings from a list of chat messages."""
|
||||
"""Extract text content from user-role messages only.
|
||||
|
||||
Only user messages are extracted to avoid injecting system or
|
||||
assistant content into the ``texts`` list, which should mirror
|
||||
the original user-provided input.
|
||||
"""
|
||||
texts: List[str] = []
|
||||
for message in messages:
|
||||
if message.get("role") != "user":
|
||||
continue
|
||||
content = message.get("content")
|
||||
if isinstance(content, str):
|
||||
texts.append(content)
|
||||
|
|
|
|||
|
|
@ -373,7 +373,7 @@ class TestPromptGuardRedactAction:
|
|||
input_type="request",
|
||||
)
|
||||
assert result["structured_messages"] == redacted
|
||||
assert "My SSN is *********" in result["texts"]
|
||||
assert result["texts"] == ["My SSN is *********"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_redact_structured_only_does_not_create_texts(
|
||||
|
|
@ -629,7 +629,7 @@ class TestPromptGuardErrorHandling:
|
|||
async def test_http_error_propagates_block_on_error(
|
||||
self, promptguard_guardrail, mock_request_data
|
||||
):
|
||||
"""Default block_on_error=True re-raises HTTP errors."""
|
||||
"""Default block_on_error=True wraps HTTP errors in GuardrailRaisedException."""
|
||||
mock_request = httpx.Request("POST", "https://api.test.promptguard.co")
|
||||
mock_resp = httpx.Response(status_code=500, request=mock_request)
|
||||
with patch.object(
|
||||
|
|
@ -641,29 +641,33 @@ class TestPromptGuardErrorHandling:
|
|||
response=mock_resp,
|
||||
),
|
||||
):
|
||||
with pytest.raises(httpx.HTTPStatusError):
|
||||
with pytest.raises(GuardrailRaisedException) as exc_info:
|
||||
await promptguard_guardrail.apply_guardrail(
|
||||
inputs={"texts": ["test"]},
|
||||
request_data=mock_request_data,
|
||||
input_type="request",
|
||||
)
|
||||
assert "block_on_error=True" in str(exc_info.value)
|
||||
assert exc_info.value.__cause__ is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_connection_error_propagates_block_on_error(
|
||||
self, promptguard_guardrail, mock_request_data
|
||||
):
|
||||
"""Default block_on_error=True re-raises connection errors."""
|
||||
"""Default block_on_error=True wraps connection errors in GuardrailRaisedException."""
|
||||
with patch.object(
|
||||
promptguard_guardrail.async_handler,
|
||||
"post",
|
||||
side_effect=httpx.ConnectError("Connection refused"),
|
||||
):
|
||||
with pytest.raises(httpx.ConnectError):
|
||||
with pytest.raises(GuardrailRaisedException) as exc_info:
|
||||
await promptguard_guardrail.apply_guardrail(
|
||||
inputs={"texts": ["test"]},
|
||||
request_data=mock_request_data,
|
||||
input_type="request",
|
||||
)
|
||||
assert "block_on_error=True" in str(exc_info.value)
|
||||
assert exc_info.value.__cause__ is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fail_open_returns_inputs_on_http_error(self, mock_request_data):
|
||||
|
|
@ -747,6 +751,22 @@ class TestPromptGuardErrorHandling:
|
|||
)
|
||||
assert result["texts"] == ["test"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_null_decision_treated_as_allow(
|
||||
self, promptguard_guardrail, mock_request_data
|
||||
):
|
||||
"""Explicit null decision should be treated as allow."""
|
||||
resp = _make_response({"decision": None, "event_id": "evt-null"})
|
||||
with patch.object(
|
||||
promptguard_guardrail.async_handler, "post", return_value=resp
|
||||
):
|
||||
result = await promptguard_guardrail.apply_guardrail(
|
||||
inputs={"texts": ["test"]},
|
||||
request_data=mock_request_data,
|
||||
input_type="request",
|
||||
)
|
||||
assert result["texts"] == ["test"]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Config model
|
||||
|
|
|
|||
|
|
@ -48,6 +48,7 @@ export const guardrail_provider_map: Record<string, string> = {
|
|||
LitellmContentFilter: "litellm_content_filter",
|
||||
ToolPermission: "tool_permission",
|
||||
BlockCodeExecution: "block_code_execution",
|
||||
Promptguard: "promptguard",
|
||||
};
|
||||
|
||||
// Function to populate provider map from API response - updates the original map
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue