mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-19 00:01:29 +00:00
fix(guardrails): write per-message guardrail rewrites back onto Responses input items
A guardrail that answers one rewritten text per message it saw no longer matches the texts the Responses handler extracted once the request carries instructions or tool items, so the rewrite was rejected with a 500. Spread such an answer over the structured messages' text slots and write it back through the structured path, have Prompt Security modify return structured_messages directly, and give the chat completions pairing the same named rejection instead of a silent misalignment when the counts differ.
This commit is contained in:
mateo-berri
parent
e240997529
commit
a635d7be6a
7 changed files with 373 additions and 9 deletions
|
|
@ -1,8 +1,10 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from collections.abc import Callable, Iterator, Sequence
|
||||
from typing import Final, TypeVar
|
||||
from collections.abc import Callable, Iterator, Mapping, Sequence
|
||||
from itertools import accumulate
|
||||
from types import MappingProxyType
|
||||
from typing import Final, TypeVar, cast # noqa: TID251 # a rebuilt chat row has no typed constructor across roles
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
|
|
@ -364,3 +366,74 @@ def merge_guardrailed_scoped_messages(
|
|||
yield from appended
|
||||
|
||||
return list(_merged())
|
||||
|
||||
|
||||
def _content_part_text(part: object) -> str | None:
|
||||
if not isinstance(part, Mapping):
|
||||
return None
|
||||
text: Final = part.get("text")
|
||||
return text if isinstance(text, str) else None
|
||||
|
||||
|
||||
def message_text_slot_count(message: AllMessageValues) -> int:
|
||||
content: Final = message.get("content")
|
||||
if isinstance(content, str):
|
||||
return 1
|
||||
if isinstance(content, list):
|
||||
return sum(1 for part in content if _content_part_text(part) is not None)
|
||||
return 0
|
||||
|
||||
|
||||
def _part_with_text(part: object, text: str) -> object:
|
||||
if not isinstance(part, Mapping):
|
||||
return part
|
||||
return {**part, "text": text} # mutable-ok: content parts stay JSON-plain dicts
|
||||
|
||||
|
||||
def _content_with_slot_texts(content: Sequence[object], texts: Sequence[str]) -> list[object]:
|
||||
text_part_indices: Final = tuple(
|
||||
index for index, part in enumerate(content) if _content_part_text(part) is not None
|
||||
)
|
||||
replacement_by_index: Final = MappingProxyType(dict(zip(text_part_indices, texts)))
|
||||
return [ # mutable-ok: message content stays a JSON list
|
||||
_part_with_text(part, replacement_by_index[index]) if index in replacement_by_index else part
|
||||
for index, part in enumerate(content)
|
||||
]
|
||||
|
||||
|
||||
def _message_with_slot_texts(message: AllMessageValues, texts: Sequence[str]) -> AllMessageValues:
|
||||
content: Final = message.get("content")
|
||||
if not isinstance(content, (str, list)):
|
||||
return message
|
||||
rewritten_content: Final = texts[0] if isinstance(content, str) else _content_with_slot_texts(content, texts)
|
||||
rewritten: Final = {**message, "content": rewritten_content} # mutable-ok: chat rows stay JSON-plain dicts
|
||||
return cast("AllMessageValues", rewritten) # cast-ok: the same row with only its text slots swapped
|
||||
|
||||
|
||||
def message_with_slot_texts(message: AllMessageValues, texts: Sequence[str]) -> AllMessageValues | None:
|
||||
if message_text_slot_count(message) != len(texts):
|
||||
return None
|
||||
return _message_with_slot_texts(message, texts)
|
||||
|
||||
|
||||
def messages_with_slot_texts(
|
||||
messages: Sequence[AllMessageValues],
|
||||
texts: Sequence[str],
|
||||
) -> list[AllMessageValues] | None:
|
||||
"""Spread one flat list of rewritten texts over the messages' text slots, in order.
|
||||
|
||||
A slot is a string ``content`` or one list part carrying a string ``text``;
|
||||
images and other parts ride along untouched. A guardrail that answers one
|
||||
text per message it saw produces exactly this shape, which stops matching
|
||||
the endpoint's own per-text extraction as soon as the request carries
|
||||
instructions or tool items. Returns None unless the counts line up exactly,
|
||||
so a rewrite never lands on the wrong slot.
|
||||
"""
|
||||
slot_counts: Final = tuple(message_text_slot_count(message) for message in messages)
|
||||
if sum(slot_counts) != len(texts):
|
||||
return None
|
||||
offsets: Final = tuple(accumulate(slot_counts, initial=0))
|
||||
return [ # mutable-ok: guardrail rows travel as a list
|
||||
_message_with_slot_texts(message, texts[start:end])
|
||||
for message, start, end in zip(messages, offsets, offsets[1:])
|
||||
]
|
||||
|
|
|
|||
|
|
@ -196,6 +196,10 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
else:
|
||||
# Step 3: Map guardrail responses back to original message structure
|
||||
if guardrailed_texts and texts_to_check:
|
||||
if len(guardrailed_texts) != len(text_task_mappings):
|
||||
from litellm.proxy.policy_engine.pipeline_executor import UnappliableRequestRewrite
|
||||
|
||||
raise UnappliableRequestRewrite(guardrail_to_apply.guardrail_name or "unknown")
|
||||
await self._apply_guardrail_responses_to_input_texts(
|
||||
messages=messages,
|
||||
responses=guardrailed_texts,
|
||||
|
|
|
|||
|
|
@ -53,6 +53,7 @@ from litellm.llms.base_llm.guardrail_translation.base_translation import (
|
|||
)
|
||||
from litellm.llms.base_llm.guardrail_translation.utils import (
|
||||
blocked_responses_stream_usage,
|
||||
messages_with_slot_texts,
|
||||
stream_item_field,
|
||||
stream_item_fingerprint,
|
||||
stream_item_items,
|
||||
|
|
@ -395,6 +396,20 @@ def _patched_request_fields(
|
|||
)
|
||||
|
||||
|
||||
def _guardrailed_structured_messages(
|
||||
structured_messages: Sequence[AllMessageValues] | None,
|
||||
sent_text_count: int,
|
||||
guardrailed_inputs: GenericGuardrailAPIInputs,
|
||||
) -> Sequence[AllMessageValues] | None:
|
||||
returned: Final = guardrailed_inputs.get("structured_messages")
|
||||
if returned is not None and returned is not structured_messages:
|
||||
return returned
|
||||
rewritten_texts: Final = guardrailed_inputs.get("texts")
|
||||
if not structured_messages or rewritten_texts is None or len(rewritten_texts) == sent_text_count:
|
||||
return None
|
||||
return messages_with_slot_texts(structured_messages, rewritten_texts)
|
||||
|
||||
|
||||
def _patch_or_convert_request_fields(
|
||||
raw_input: object,
|
||||
instructions: object,
|
||||
|
|
@ -473,7 +488,8 @@ class OpenAIResponsesHandler(BaseTranslation):
|
|||
form.chat_tools for form in LiteLLMCompletionResponsesConfig.responses_tools_to_chat_forms(original_tools)
|
||||
)
|
||||
extracted: Final = self._extract_guardrail_inputs(data, input_data, flattened_tool_groups)
|
||||
if not extracted.inputs.get("texts"):
|
||||
sent_texts: Final = extracted.inputs.get("texts")
|
||||
if not sent_texts:
|
||||
return data
|
||||
if structured_messages:
|
||||
extracted.inputs["structured_messages"] = structured_messages
|
||||
|
|
@ -486,7 +502,9 @@ class OpenAIResponsesHandler(BaseTranslation):
|
|||
self._apply_guardrailed_tools_to_data(
|
||||
data, original_tools, flattened_tool_groups, guardrailed_inputs.get("tools")
|
||||
)
|
||||
written_back: Final = self._written_back_request_fields(data, structured_messages, guardrailed_inputs)
|
||||
written_back: Final = self._written_back_request_fields(
|
||||
data, structured_messages, len(sent_texts), guardrailed_inputs
|
||||
)
|
||||
if written_back is not None:
|
||||
data["input"] = list(written_back.input) # mutable-ok: JSON body
|
||||
if written_back.instructions is None:
|
||||
|
|
@ -553,10 +571,11 @@ class OpenAIResponsesHandler(BaseTranslation):
|
|||
def _written_back_request_fields(
|
||||
data: Mapping[str, object],
|
||||
structured_messages: Sequence[AllMessageValues] | None,
|
||||
sent_text_count: int,
|
||||
guardrailed_inputs: GenericGuardrailAPIInputs,
|
||||
) -> _RequestFields | None:
|
||||
guardrailed: Final = guardrailed_inputs.get("structured_messages")
|
||||
if guardrailed is None or guardrailed is structured_messages:
|
||||
guardrailed: Final = _guardrailed_structured_messages(structured_messages, sent_text_count, guardrailed_inputs)
|
||||
if guardrailed is None:
|
||||
return None
|
||||
return _patch_or_convert_request_fields(
|
||||
data.get("input"),
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@ import asyncio
|
|||
import base64
|
||||
import os
|
||||
from collections.abc import Mapping, Sequence
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final, Literal, Optional
|
||||
|
||||
import httpx
|
||||
|
|
@ -14,11 +15,13 @@ from litellm.integrations.custom_guardrail import (
|
|||
CustomGuardrail,
|
||||
log_guardrail_information,
|
||||
)
|
||||
from litellm.llms.base_llm.guardrail_translation.utils import message_with_slot_texts
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
get_async_httpx_client,
|
||||
httpxSpecialProvider,
|
||||
)
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.types.utils import GenericGuardrailAPIInputs
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -27,6 +30,7 @@ if TYPE_CHECKING:
|
|||
|
||||
|
||||
_SANITIZE_FILE_FAIL_OPEN_TIMEOUT_SECONDS: Final = 30.0
|
||||
_PROTECT_ROLES: Final = frozenset({"system", "user", "assistant"})
|
||||
|
||||
|
||||
class PromptSecurityGuardrailMissingSecrets(Exception):
|
||||
|
|
@ -275,14 +279,44 @@ class PromptSecurityGuardrail(CustomGuardrail):
|
|||
detail="Blocked by Prompt Security, Violations: " + ", ".join(violations),
|
||||
)
|
||||
elif action == "modify":
|
||||
# Extract modified texts from modified_messages
|
||||
modified_messages: Final = result.get("modified_messages", [])
|
||||
modified_texts: Final = self._extract_texts_from_messages(modified_messages)
|
||||
if modified_texts:
|
||||
inputs["texts"] = modified_texts
|
||||
rewritten_messages: Final = self._structured_messages_with_modifications(
|
||||
structured_messages, modified_messages
|
||||
)
|
||||
if rewritten_messages is not None:
|
||||
inputs["structured_messages"] = rewritten_messages
|
||||
|
||||
return inputs
|
||||
|
||||
def _is_sent_to_protect(self, message: Mapping[str, object]) -> bool:
|
||||
return self.check_tool_results or message.get("role") in _PROTECT_ROLES
|
||||
|
||||
def _structured_messages_with_modifications(
|
||||
self,
|
||||
structured_messages: Sequence[AllMessageValues],
|
||||
modified_messages: Sequence[Mapping[str, object]],
|
||||
) -> list[AllMessageValues] | None:
|
||||
sent_indices: Final = tuple(
|
||||
index for index, message in enumerate(structured_messages) if self._is_sent_to_protect(message)
|
||||
)
|
||||
if not sent_indices or len(sent_indices) != len(modified_messages):
|
||||
return None
|
||||
rewritten: Final = tuple(
|
||||
message_with_slot_texts(structured_messages[index], self._extract_texts_from_messages((modified,)))
|
||||
for index, modified in zip(sent_indices, modified_messages)
|
||||
)
|
||||
replacements: Final = MappingProxyType(
|
||||
{index: message for index, message in zip(sent_indices, rewritten) if message is not None}
|
||||
)
|
||||
if len(replacements) != len(sent_indices):
|
||||
return None
|
||||
return [ # mutable-ok: guardrail inputs take a list
|
||||
replacements.get(index, message) for index, message in enumerate(structured_messages)
|
||||
]
|
||||
|
||||
async def _apply_guardrail_on_response(
|
||||
self,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
|
|
@ -678,14 +712,13 @@ class PromptSecurityGuardrail(CustomGuardrail):
|
|||
|
||||
This allows checking tool results for indirect prompt injection when enabled.
|
||||
"""
|
||||
supported_roles: Final = ["system", "user", "assistant"]
|
||||
filtered_messages: Final = []
|
||||
transformed_count = 0
|
||||
filtered_count = 0
|
||||
|
||||
for message in messages:
|
||||
role = message.get("role", "")
|
||||
if role in supported_roles:
|
||||
if role in _PROTECT_ROLES:
|
||||
filtered_messages.append(message)
|
||||
else:
|
||||
if self.check_tool_results:
|
||||
|
|
|
|||
|
|
@ -1893,6 +1893,49 @@ class TestScanOnlyToolResults:
|
|||
assert data["messages"][4]["content"] == "and then?"
|
||||
|
||||
|
||||
class ToolDroppingTextGuardrail(CustomGuardrail):
|
||||
"""Answers one text per non-tool message it saw, the way a guardrail that
|
||||
filters tool rows out before scanning does, and hands back only texts."""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__(guardrail_name="tool-dropping-redactor")
|
||||
|
||||
async def apply_guardrail(
|
||||
self,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
request_data: dict,
|
||||
input_type: Literal["request", "response"],
|
||||
logging_obj: Optional[Any] = None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
kept = [m for m in inputs.get("structured_messages") or [] if m.get("role") != "tool"]
|
||||
return {**inputs, "texts": [str(m.get("content")).replace("POISON", "[BLOCKED]") for m in kept]}
|
||||
|
||||
|
||||
class TestPerMessageTextWriteBack:
|
||||
"""Texts that no longer pair one-to-one with what the handler extracted must be
|
||||
rejected by name instead of sliding onto the wrong messages."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fewer_texts_than_extracted_over_a_tool_message_is_rejected(self):
|
||||
from litellm.proxy.policy_engine.pipeline_executor import UnappliableRequestRewrite
|
||||
|
||||
handler = OpenAIChatCompletionsHandler()
|
||||
original_messages = [
|
||||
{"role": "system", "content": "SYSTEM-PROMPT"},
|
||||
{"role": "user", "content": "fetch the page"},
|
||||
{"role": "assistant", "content": "fetching"},
|
||||
{"role": "tool", "tool_call_id": "call_1", "content": "page says POISON here"},
|
||||
{"role": "user", "content": "and then?"},
|
||||
]
|
||||
data = {"messages": json.loads(json.dumps(original_messages))}
|
||||
|
||||
with pytest.raises(UnappliableRequestRewrite) as excinfo:
|
||||
await handler.process_input_messages(data=data, guardrail_to_apply=ToolDroppingTextGuardrail())
|
||||
|
||||
assert excinfo.value.guardrail_name == "tool-dropping-redactor"
|
||||
assert data["messages"] == original_messages, "a rejected rewrite must leave the request untouched"
|
||||
|
||||
|
||||
class TestBuildBlockSseChunks:
|
||||
"""build_block_sse_chunks turns a streaming ModifyResponseException into 200 SSE chunks"""
|
||||
|
||||
|
|
|
|||
|
|
@ -2338,6 +2338,109 @@ def _parallel_tool_call_input() -> list:
|
|||
]
|
||||
|
||||
|
||||
SSN = "123-45-6789"
|
||||
REDACTED_SSN = "<US_SSN>"
|
||||
|
||||
|
||||
def _slot_texts(message: dict) -> list[str]:
|
||||
content = message.get("content")
|
||||
if isinstance(content, str):
|
||||
return [content]
|
||||
if isinstance(content, list):
|
||||
return [part["text"] for part in content if isinstance(part, dict) and isinstance(part.get("text"), str)]
|
||||
return []
|
||||
|
||||
|
||||
class PerMessageRedactionGuardrail(CustomGuardrail):
|
||||
"""Guardrail that answers one redacted text per message it was shown and hands
|
||||
back only texts, the way Prompt Security in modify mode and a generic guardrail
|
||||
API server that scans per message do."""
|
||||
|
||||
def __init__(self, extra_texts: int = 0):
|
||||
super().__init__(guardrail_name="per-message-redactor")
|
||||
self.extra_texts = extra_texts
|
||||
|
||||
async def apply_guardrail(
|
||||
self,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
request_data: dict,
|
||||
input_type: Literal["request", "response"],
|
||||
logging_obj: Optional[Any] = None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
messages = inputs.get("structured_messages") or []
|
||||
texts = [text.replace(SSN, REDACTED_SSN) for message in messages for text in _slot_texts(message)]
|
||||
return {**inputs, "texts": texts + ["junk"] * self.extra_texts}
|
||||
|
||||
|
||||
class TestPerMessageTextWriteBack:
|
||||
"""A guardrail that rewrites one text per message it saw must land on the
|
||||
instructions and the input items those messages came from, not be rejected."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_instructions_plus_tool_replay_gets_each_rewrite_in_place(self):
|
||||
handler = OpenAIResponsesHandler()
|
||||
function_call_item = {
|
||||
"type": "function_call",
|
||||
"call_id": "call_1",
|
||||
"name": "lookup_customer",
|
||||
"arguments": '{"query": "' + SSN + '"}',
|
||||
}
|
||||
data = {
|
||||
"model": "gpt-5.6",
|
||||
"instructions": "Never repeat the SSN " + SSN + " back.",
|
||||
"input": [
|
||||
{"role": "user", "content": "Look up " + SSN + " for me."},
|
||||
function_call_item,
|
||||
{"type": "function_call_output", "call_id": "call_1", "output": '{"ssn": "' + SSN + '"}'},
|
||||
],
|
||||
}
|
||||
|
||||
result = await handler.process_input_messages(data, PerMessageRedactionGuardrail())
|
||||
|
||||
assert result["instructions"] == "Never repeat the SSN " + REDACTED_SSN + " back."
|
||||
assert [item.get("type", item.get("role")) for item in result["input"]] == [
|
||||
"user",
|
||||
"function_call",
|
||||
"function_call_output",
|
||||
]
|
||||
assert _slot_texts(result["input"][0]) == ["Look up " + REDACTED_SSN + " for me."]
|
||||
assert result["input"][1] == function_call_item
|
||||
assert result["input"][2]["output"] == '{"ssn": "' + REDACTED_SSN + '"}'
|
||||
assert result["input"][2]["call_id"] == "call_1"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_string_input_with_instructions_keeps_the_two_apart(self):
|
||||
handler = OpenAIResponsesHandler()
|
||||
data = {
|
||||
"model": "gpt-5.6",
|
||||
"instructions": "Redact " + SSN + " everywhere.",
|
||||
"input": "My SSN is " + SSN + ".",
|
||||
}
|
||||
|
||||
result = await handler.process_input_messages(data, PerMessageRedactionGuardrail())
|
||||
|
||||
assert result["instructions"] == "Redact " + REDACTED_SSN + " everywhere."
|
||||
assert [_slot_texts(item) for item in result["input"]] == [["My SSN is " + REDACTED_SSN + "."]]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_count_matching_neither_texts_nor_messages_is_still_rejected(self):
|
||||
from litellm.proxy.policy_engine.pipeline_executor import UnappliableRequestRewrite
|
||||
|
||||
handler = OpenAIResponsesHandler()
|
||||
original_input = [
|
||||
{"role": "user", "content": "Look up " + SSN + " for me."},
|
||||
{"type": "function_call_output", "call_id": "call_1", "output": '{"ssn": "' + SSN + '"}'},
|
||||
]
|
||||
data = {"model": "gpt-5.6", "instructions": "Be terse.", "input": copy.deepcopy(original_input)}
|
||||
|
||||
with pytest.raises(UnappliableRequestRewrite) as excinfo:
|
||||
await handler.process_input_messages(data, PerMessageRedactionGuardrail(extra_texts=1))
|
||||
|
||||
assert excinfo.value.guardrail_name == "per-message-redactor"
|
||||
assert data["input"] == original_input
|
||||
assert data["instructions"] == "Be terse."
|
||||
|
||||
|
||||
class TestProvenancePatching:
|
||||
"""The O(n) provenance pass must keep patching rewritten rows in place for the
|
||||
shapes real agent loops produce, and fall back safely everywhere else."""
|
||||
|
|
|
|||
|
|
@ -174,6 +174,95 @@ async def test_apply_guardrail_modify_request(monkeypatch: pytest.MonkeyPatch):
|
|||
assert result["texts"] == ["User prompt with PII: SSN [REDACTED]"]
|
||||
|
||||
|
||||
def _modify_response(modified_messages: list) -> Response:
|
||||
mock_response = Response(
|
||||
json={"result": {"prompt": {"action": "modify", "modified_messages": modified_messages}}},
|
||||
status_code=200,
|
||||
request=Request(method="POST", url="https://test.prompt.security/api/protect"),
|
||||
)
|
||||
mock_response.raise_for_status = lambda: None
|
||||
return mock_response
|
||||
|
||||
|
||||
def _tool_replay_messages() -> list:
|
||||
return [
|
||||
{"role": "system", "content": "Never echo an SSN like 123-45-6789."},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Look up 123-45-6789"},
|
||||
{"type": "image_url", "image_url": {"url": "https://example.com/id-card.png"}},
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [
|
||||
{"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": "{}"}}
|
||||
],
|
||||
},
|
||||
{"role": "tool", "tool_call_id": "call_1", "content": '{"ssn": "123-45-6789"}'},
|
||||
{"role": "user", "content": "Summarize what you found."},
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_modify_returns_structured_messages_with_tool_rows_kept(monkeypatch: pytest.MonkeyPatch):
|
||||
"""A per-message modify verdict comes back as structured_messages so the
|
||||
endpoint handler can write it back by message, with the rows Prompt Security
|
||||
never saw (tool results) and the non-text parts (images) left in place."""
|
||||
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)
|
||||
messages = _tool_replay_messages()
|
||||
inputs = {"texts": ["Look up 123-45-6789", "Summarize what you found."], "structured_messages": messages}
|
||||
modified_messages = [
|
||||
{"role": "system", "content": "Never echo an SSN like [REDACTED]."},
|
||||
{"role": "user", "content": [{"type": "text", "text": "Look up [REDACTED]"}]},
|
||||
{"role": "assistant", "content": None},
|
||||
{"role": "user", "content": "Summarize what you found."},
|
||||
]
|
||||
|
||||
with patch.object(guardrail.async_handler, "post", return_value=_modify_response(modified_messages)):
|
||||
result = await guardrail.apply_guardrail(inputs=inputs, request_data={"messages": messages}, input_type="request")
|
||||
|
||||
assert result["structured_messages"] == [
|
||||
{"role": "system", "content": "Never echo an SSN like [REDACTED]."},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Look up [REDACTED]"},
|
||||
{"type": "image_url", "image_url": {"url": "https://example.com/id-card.png"}},
|
||||
],
|
||||
},
|
||||
messages[2],
|
||||
messages[3],
|
||||
{"role": "user", "content": "Summarize what you found."},
|
||||
]
|
||||
assert result["structured_messages"] is not messages
|
||||
assert result["texts"] == [
|
||||
"Never echo an SSN like [REDACTED].",
|
||||
"Look up [REDACTED]",
|
||||
"Summarize what you found.",
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_modify_with_unexpected_message_count_keeps_texts_only(monkeypatch: pytest.MonkeyPatch):
|
||||
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)
|
||||
messages = _tool_replay_messages()
|
||||
inputs = {"texts": ["Look up 123-45-6789", "Summarize what you found."], "structured_messages": messages}
|
||||
modified_messages = [{"role": "user", "content": "Look up [REDACTED]"}]
|
||||
|
||||
with patch.object(guardrail.async_handler, "post", return_value=_modify_response(modified_messages)):
|
||||
result = await guardrail.apply_guardrail(inputs=inputs, request_data={"messages": messages}, input_type="request")
|
||||
|
||||
assert result["structured_messages"] is messages
|
||||
assert result["texts"] == ["Look up [REDACTED]"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_allow_request(monkeypatch: pytest.MonkeyPatch):
|
||||
"""Test that apply_guardrail allows safe prompts"""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue