Skip to content

Commit 19395bf

Browse files
committed
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)
1 parent 648e202 commit 19395bf

File tree

2 files changed

+40
-5
lines changed

2 files changed

+40
-5
lines changed

litellm/proxy/guardrails/guardrail_hooks/promptguard/promptguard.py

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -148,6 +148,7 @@ async def apply_guardrail(
148148
"Content-Type": "application/json",
149149
},
150150
json=payload,
151+
timeout=10.0,
151152
)
152153
response.raise_for_status()
153154
result = response.json()
@@ -184,11 +185,14 @@ async def apply_guardrail(
184185
if redacted:
185186
if structured_messages:
186187
inputs["structured_messages"] = redacted
187-
extracted = self._extract_texts_from_messages(
188-
redacted,
189-
)
190-
if extracted:
191-
inputs["texts"] = extracted
188+
if "texts" in inputs:
189+
extracted = (
190+
self._extract_texts_from_messages(
191+
redacted,
192+
)
193+
)
194+
if extracted:
195+
inputs["texts"] = extracted
192196

193197
return inputs
194198

tests/test_litellm/proxy/guardrails/guardrail_hooks/test_promptguard.py

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -375,6 +375,37 @@ async def test_redact_updates_structured_messages(
375375
assert result["structured_messages"] == redacted
376376
assert "My SSN is *********" in result["texts"]
377377

378+
@pytest.mark.asyncio
379+
async def test_redact_structured_only_does_not_create_texts(
380+
self, promptguard_guardrail, mock_request_data
381+
):
382+
"""When only structured_messages are provided, redact should not inject a texts key."""
383+
original = [
384+
{"role": "user", "content": "My SSN is 123-45-6789"},
385+
]
386+
redacted = [
387+
{"role": "user", "content": "My SSN is *********"},
388+
]
389+
resp = _make_response(
390+
{
391+
"decision": "redact",
392+
"event_id": "evt-009",
393+
"redacted_messages": redacted,
394+
}
395+
)
396+
with patch.object(
397+
promptguard_guardrail.async_handler,
398+
"post",
399+
return_value=resp,
400+
):
401+
result = await promptguard_guardrail.apply_guardrail(
402+
inputs={"structured_messages": original},
403+
request_data=mock_request_data,
404+
input_type="request",
405+
)
406+
assert result["structured_messages"] == redacted
407+
assert "texts" not in result
408+
378409
@pytest.mark.asyncio
379410
async def test_redact_texts_only_without_structured(
380411
self, promptguard_guardrail, mock_request_data

0 commit comments

Comments
 (0)