mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-11 03:38:38 +00:00
Address remaining Greptile feedback: timeout, redact guard
- Add explicit 10s timeout to async_handler.post() to prevent indefinite hangs when PromptGuard API is unresponsive - Guard redact path: only update inputs["texts"] when the key was originally present, avoiding phantom key injection - Add test: redact with structured_messages only does not create texts key (41 tests total)
This commit is contained in:
parent
196039247b
commit
67b8eaa1b2
2 changed files with 40 additions and 5 deletions
|
|
@ -148,6 +148,7 @@ class PromptGuardGuardrail(CustomGuardrail):
|
|||
"Content-Type": "application/json",
|
||||
},
|
||||
json=payload,
|
||||
timeout=10.0,
|
||||
)
|
||||
response.raise_for_status()
|
||||
result = response.json()
|
||||
|
|
@ -184,11 +185,14 @@ class PromptGuardGuardrail(CustomGuardrail):
|
|||
if redacted:
|
||||
if structured_messages:
|
||||
inputs["structured_messages"] = redacted
|
||||
extracted = self._extract_texts_from_messages(
|
||||
redacted,
|
||||
)
|
||||
if extracted:
|
||||
inputs["texts"] = extracted
|
||||
if "texts" in inputs:
|
||||
extracted = (
|
||||
self._extract_texts_from_messages(
|
||||
redacted,
|
||||
)
|
||||
)
|
||||
if extracted:
|
||||
inputs["texts"] = extracted
|
||||
|
||||
return inputs
|
||||
|
||||
|
|
|
|||
|
|
@ -375,6 +375,37 @@ class TestPromptGuardRedactAction:
|
|||
assert result["structured_messages"] == redacted
|
||||
assert "My SSN is *********" in result["texts"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_redact_structured_only_does_not_create_texts(
|
||||
self, promptguard_guardrail, mock_request_data
|
||||
):
|
||||
"""When only structured_messages are provided, redact should not inject a texts key."""
|
||||
original = [
|
||||
{"role": "user", "content": "My SSN is 123-45-6789"},
|
||||
]
|
||||
redacted = [
|
||||
{"role": "user", "content": "My SSN is *********"},
|
||||
]
|
||||
resp = _make_response(
|
||||
{
|
||||
"decision": "redact",
|
||||
"event_id": "evt-009",
|
||||
"redacted_messages": redacted,
|
||||
}
|
||||
)
|
||||
with patch.object(
|
||||
promptguard_guardrail.async_handler,
|
||||
"post",
|
||||
return_value=resp,
|
||||
):
|
||||
result = await promptguard_guardrail.apply_guardrail(
|
||||
inputs={"structured_messages": original},
|
||||
request_data=mock_request_data,
|
||||
input_type="request",
|
||||
)
|
||||
assert result["structured_messages"] == redacted
|
||||
assert "texts" not in result
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_redact_texts_only_without_structured(
|
||||
self, promptguard_guardrail, mock_request_data
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue