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:
Abhijoy Sarkar 2026-03-21 11:24:36 +05:30
parent 196039247b
commit 67b8eaa1b2
2 changed files with 40 additions and 5 deletions

View file

@ -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

View file

@ -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