mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-30 06:43:27 +00:00
feat(guardrails): match privaite guardrail to core 0.4 PII coverage
Scan the Responses input carriers the batching pass missed (typed action carriers, custom_tool_call_output.output as a list of parts, prompt.variables) plus the auxiliary request fields (completions prompt/suffix, chat prediction.content, web_search_options.user_location), and restore the refusal and audio transcript on both the sync and streaming paths. Coverage is reimplemented inline over the public engine surface (process_request / process_response), with no dependency on privaite internals.
This commit is contained in:
crp4222
parent
b27457862f
commit
ee88ab7183
2 changed files with 556 additions and 82 deletions
|
|
@ -2,14 +2,23 @@
|
|||
PrivAiTe guardrail for the LiteLLM proxy.
|
||||
|
||||
Runs PrivAiTe's engine in-process inside LiteLLM. The pre-call hook anonymizes the
|
||||
request (chat `messages` and Responses API `input`) and stashes the reversible map
|
||||
in the request metadata (consumed and popped by the post-call hook); the post-call
|
||||
hook restores the real values in the response, including chat tool-call arguments,
|
||||
the legacy function_call, and Responses API output_text/function_call output. Chat
|
||||
streaming is restored too: text content, streamed tool-call arguments, and the
|
||||
streamed function_call. (Responses API streaming restore is not yet implemented.)
|
||||
On the failure path the map is dropped from metadata as well, so it cannot reach a
|
||||
failure spend-log.
|
||||
request and stashes the reversible map in the request metadata (consumed and popped
|
||||
by the post-call hook); the post-call hook restores the real values in the response,
|
||||
including chat tool-call arguments, the legacy function_call, and Responses API
|
||||
output_text/function_call output. The request surface scanned matches the PrivAiTe
|
||||
core: chat `messages`; the Responses API `input` (every item type: role message,
|
||||
tool-call arguments, tool output including `custom_tool_call_output.output` as a
|
||||
list of `{type, text}` parts, typed action carriers, bare content parts and bare
|
||||
strings; opaque/binary items relayed whole) and `prompt.variables`; the
|
||||
`/v1/completions` `prompt` (string or batch list) and `suffix`; the chat
|
||||
`prediction.content` and `web_search_options.user_location` auxiliary fields (input
|
||||
side only, nothing to restore; tokenized integer-array inputs pass through
|
||||
unscanned). Restore covers `content`, the reasoning trace, the refusal and the audio
|
||||
transcript besides tool/function arguments. Chat streaming is restored too: text
|
||||
content, the reasoning trace, the refusal, the audio transcript, streamed tool-call
|
||||
arguments and the streamed function_call. (Responses API streaming restore is not
|
||||
yet implemented.) On the failure path the map is dropped from metadata as well, so
|
||||
it cannot reach a failure spend-log.
|
||||
|
||||
If `block_entities` is configured, a request containing any of those PII types is
|
||||
rejected with an HTTP 400 in the pre-call hook, before anything is forwarded to the
|
||||
|
|
@ -25,6 +34,7 @@ See https://github.com/crp4222/PrivAiTe for the package.
|
|||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from collections.abc import Callable
|
||||
from typing import Any, cast
|
||||
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
|
|
@ -40,6 +50,58 @@ _LANG_MODELS = {
|
|||
"nl": "nl_core_news_md",
|
||||
}
|
||||
|
||||
# Responses `input` items relayed byte-for-byte: an opaque/binary payload
|
||||
# (encrypted reasoning, a screenshot, a generated image) or a tool/pointer
|
||||
# definition with no user text. Parity with the PrivAiTe gateway scrubber.
|
||||
_RESPONSES_OPAQUE_TYPES = frozenset(
|
||||
{
|
||||
"reasoning",
|
||||
"compaction",
|
||||
"compaction_trigger",
|
||||
"computer_call_output",
|
||||
"image_generation_call",
|
||||
"item_reference",
|
||||
"mcp_list_tools",
|
||||
"tool_search_call",
|
||||
"tool_search_output",
|
||||
"additional_tools",
|
||||
}
|
||||
)
|
||||
|
||||
# Responses items whose user data sits in a known structured field (a typed
|
||||
# command/action, a patch diff, search queries and results, interpreter code
|
||||
# and logs). Every listed field is walked leaf by leaf.
|
||||
_RESPONSES_DATA_FIELDS: dict[str, tuple[str, ...]] = {
|
||||
"computer_call": ("action", "actions"),
|
||||
"local_shell_call": ("action",),
|
||||
"shell_call": ("action",),
|
||||
"web_search_call": ("action",),
|
||||
"apply_patch_call": ("operation",),
|
||||
"file_search_call": ("queries", "results"),
|
||||
"code_interpreter_call": ("code", "outputs"),
|
||||
"program": ("code",),
|
||||
"program_output": ("result",),
|
||||
}
|
||||
|
||||
# Content/output parts whose payload is binary (base64 or a file id), not text:
|
||||
# scrubbing them would corrupt the payload without removing anything a text
|
||||
# detector could find, so the whole part is relayed whole.
|
||||
_BINARY_PART_TYPES = frozenset(
|
||||
{
|
||||
"input_image",
|
||||
"input_file",
|
||||
"input_audio",
|
||||
"image",
|
||||
"output_image",
|
||||
"computer_screenshot",
|
||||
}
|
||||
)
|
||||
|
||||
# Text-bearing fields on an unknown Responses item shape (e.g. an mcp_call
|
||||
# carries both `arguments` and `output`): every one is scanned, not just the
|
||||
# first, so a second field is never left raw.
|
||||
_GENERIC_TEXT_FIELDS = ("output", "arguments", "input", "text", "reason")
|
||||
|
||||
|
||||
def _obj_get(obj: Any, key: str) -> Any:
|
||||
"""Read a field from a dict-or-object response item."""
|
||||
|
|
@ -175,63 +237,214 @@ class PrivaiteGuardrail(CustomGuardrail):
|
|||
self._engine_key = key
|
||||
return engine
|
||||
|
||||
# Text-bearing fields on non-message Responses input items: a tool output, a
|
||||
# streamed/echoed tool call, and a bare input_text/output_text content part.
|
||||
# Scanned individually so a mixed input list is not left raw.
|
||||
_ITEM_TEXT_FIELDS = ("output", "arguments", "text")
|
||||
|
||||
def _overwrite_snapshot_input(self, data: dict, new_input: Any) -> None:
|
||||
# A plain `data["input"] = ...` rebind leaks for string input: the proxy
|
||||
# snapshots the request body by shallow-copying data before this hook,
|
||||
# so the original string stays in proxy_server_request.body["input"].
|
||||
# Overwrite the snapshot copy too. (List inputs are mutated in place, so
|
||||
# the aliased snapshot already reflects the anonymized values.)
|
||||
def _overwrite_snapshot_field(self, data: dict, field: str, new_value: Any) -> None:
|
||||
# A plain `data[field] = ...` rebind leaks for a top-level string field
|
||||
# (`input`, `suffix`, a string `prompt`): the proxy snapshots the request
|
||||
# body by shallow-copying data before this hook, so the original string
|
||||
# stays in proxy_server_request.body[field]. Overwrite the snapshot copy
|
||||
# too. (Lists and dicts are mutated in place, so the aliased snapshot
|
||||
# already reflects the anonymized values.)
|
||||
psr = data.get("proxy_server_request")
|
||||
body = psr.get("body") if isinstance(psr, dict) else None
|
||||
if isinstance(body, dict) and "input" in body:
|
||||
body["input"] = new_input
|
||||
if isinstance(body, dict) and field in body:
|
||||
body[field] = new_value
|
||||
|
||||
def _responses_item_repr(self, item: Any) -> tuple:
|
||||
"""Map ONE Responses `input` list item to (engine_message, field) so it
|
||||
can be scanned. field=None means the item IS a message (replace it whole
|
||||
with the anonymized copy); a field name means write the scrubbed content
|
||||
back into item[field]; "__str__" means the item is a bare string. Returns
|
||||
(None, None) for a shape with no scannable text."""
|
||||
@staticmethod
|
||||
def _make_setter(container: Any, key: Any) -> Callable[[str], None]:
|
||||
"""A write-back that drops the scrubbed string into container[key]. The
|
||||
container (a list slot or a dict field) is aliased by the proxy's shallow
|
||||
body snapshot, so the in-place write lands in the snapshot too."""
|
||||
|
||||
def _set(value: str) -> None:
|
||||
container[key] = value
|
||||
|
||||
return _set
|
||||
|
||||
def _toplevel_setter(self, data: dict, field: str) -> Callable[[str], None]:
|
||||
"""A write-back for a top-level string field that also overwrites the
|
||||
detached body snapshot copy (see _overwrite_snapshot_field)."""
|
||||
|
||||
def _set(value: str) -> None:
|
||||
data[field] = value
|
||||
self._overwrite_snapshot_field(data, field, value)
|
||||
|
||||
return _set
|
||||
|
||||
def _add_leaf(self, batch: list, setters: list, text: Any, setter: Callable[[str], None]) -> None:
|
||||
"""Register one non-empty string leaf: a pseudo-message carrying it (so it
|
||||
flows through the engine's single choke point exactly like chat content,
|
||||
block gate included) plus the write-back that restores the scrubbed value."""
|
||||
if isinstance(text, str) and text:
|
||||
batch.append({"role": "user", "content": text})
|
||||
setters.append(setter)
|
||||
|
||||
def _collect_content(self, batch: list, setters: list, content: Any, setter: Callable[[str], None]) -> None:
|
||||
"""A role-message `content`: a bare string, or a list of parts whose
|
||||
`text`/`refusal` fields carry user text (a bare string in the list is
|
||||
user text too). Binary parts carry nothing a text detector can find."""
|
||||
if isinstance(content, str):
|
||||
self._add_leaf(batch, setters, content, setter)
|
||||
return
|
||||
if not isinstance(content, list):
|
||||
return
|
||||
for idx, part in enumerate(content):
|
||||
if isinstance(part, str):
|
||||
self._add_leaf(batch, setters, part, self._make_setter(content, idx))
|
||||
elif isinstance(part, dict) and part.get("type") not in _BINARY_PART_TYPES:
|
||||
for field in ("text", "refusal"):
|
||||
self._add_leaf(batch, setters, part.get(field), self._make_setter(part, field))
|
||||
|
||||
def _collect_data_value(self, batch: list, setters: list, value: Any, setter: Callable[[str], None]) -> None:
|
||||
"""Walk a tool payload (a string, a list of parts, or arbitrary nested
|
||||
JSON) registering every string leaf, but relay binary parts whole. Object
|
||||
keys are never scanned, matching the documented boundary."""
|
||||
if isinstance(value, str):
|
||||
self._add_leaf(batch, setters, value, setter)
|
||||
return
|
||||
if isinstance(value, list):
|
||||
for idx, part in enumerate(value):
|
||||
self._collect_data_value(batch, setters, part, self._make_setter(value, idx))
|
||||
return
|
||||
if isinstance(value, dict) and value.get("type") not in _BINARY_PART_TYPES:
|
||||
for key, sub in value.items():
|
||||
self._collect_data_value(batch, setters, sub, self._make_setter(value, key))
|
||||
|
||||
def _collect_responses_item(self, batch: list, setters: list, container: list, idx: int, item: Any) -> None:
|
||||
"""Register the scannable text on ONE Responses `input` list item, by its
|
||||
shape (bare string, role message, typed tool call/output, or an unknown
|
||||
shape scanned field by field). Opaque/binary items are relayed whole."""
|
||||
if isinstance(item, str):
|
||||
return ({"role": "user", "content": item}, "__str__") if item else (None, None)
|
||||
if isinstance(item, dict):
|
||||
if "role" in item:
|
||||
# A message item; the engine scans its content natively.
|
||||
return item, None
|
||||
for field in self._ITEM_TEXT_FIELDS:
|
||||
if isinstance(item.get(field), str) and item[field]:
|
||||
return {"role": "user", "content": item[field]}, field
|
||||
if isinstance(item.get("content"), (str, list)) and item["content"]:
|
||||
return {"role": "user", "content": item["content"]}, "content"
|
||||
return None, None
|
||||
self._add_leaf(batch, setters, item, self._make_setter(container, idx))
|
||||
return
|
||||
if not isinstance(item, dict):
|
||||
return
|
||||
if "role" in item and "content" in item:
|
||||
self._collect_content(batch, setters, item["content"], self._make_setter(item, "content"))
|
||||
return
|
||||
itype = item.get("type") or ""
|
||||
if itype in _RESPONSES_OPAQUE_TYPES:
|
||||
return
|
||||
if not self._collect_typed_item(batch, setters, item, itype):
|
||||
self._collect_generic_item(batch, setters, item)
|
||||
|
||||
async def _anonymize_request(self, data: dict, engine: Any) -> Any:
|
||||
"""Anonymize chat `messages` AND Responses `input` in place using the
|
||||
engine (span-precise), sharing ONE mapping. Every Responses input item is
|
||||
scanned item by item (message, tool output, tool call, bare string), so a
|
||||
mixed input list no longer slips past detection and the block gate.
|
||||
Returns the mapping, or None if there was nothing to anonymize."""
|
||||
messages = data.get("messages")
|
||||
msg_list = messages if isinstance(messages, list) else []
|
||||
batch: list = list(msg_list)
|
||||
def _collect_typed_item(self, batch: list, setters: list, item: dict, itype: str) -> bool:
|
||||
"""Register the text on an item whose user data sits in a type-specific
|
||||
field; returns False to hand an unknown shape to the generic scan. Covers
|
||||
function_call/custom_tool_call, the typed action carriers, and every
|
||||
`*_output` tool output (whose `output` can be a list of {type, text}
|
||||
parts, not just a string)."""
|
||||
if itype == "function_call" and isinstance(item.get("arguments"), str):
|
||||
self._add_leaf(batch, setters, item["arguments"], self._make_setter(item, "arguments"))
|
||||
return True
|
||||
if itype == "custom_tool_call" and isinstance(item.get("input"), str):
|
||||
self._add_leaf(batch, setters, item["input"], self._make_setter(item, "input"))
|
||||
return True
|
||||
fields = _RESPONSES_DATA_FIELDS.get(itype)
|
||||
if fields is not None:
|
||||
for field in fields:
|
||||
if item.get(field) is not None:
|
||||
self._collect_data_value(batch, setters, item[field], self._make_setter(item, field))
|
||||
return True
|
||||
if itype.endswith("_output") and "output" in item:
|
||||
self._collect_data_value(batch, setters, item["output"], self._make_setter(item, "output"))
|
||||
return True
|
||||
return False
|
||||
|
||||
def _collect_generic_item(self, batch: list, setters: list, item: dict) -> None:
|
||||
"""Fallback scan for an unknown item shape: every known text-bearing field
|
||||
plus a role-less `content`, so a second field is never left raw."""
|
||||
for field in _GENERIC_TEXT_FIELDS:
|
||||
self._add_leaf(batch, setters, item.get(field), self._make_setter(item, field))
|
||||
if isinstance(item.get("content"), (str, list)):
|
||||
self._collect_content(batch, setters, item["content"], self._make_setter(item, "content"))
|
||||
|
||||
def _collect_responses_input(self, data: dict, batch: list, setters: list) -> None:
|
||||
input_value = data.get("input")
|
||||
targets: list = [] # (index_or_"str", field) describing each write-back
|
||||
if isinstance(input_value, str) and input_value:
|
||||
batch.append({"role": "user", "content": input_value})
|
||||
targets.append(("str", None))
|
||||
self._add_leaf(batch, setters, input_value, self._toplevel_setter(data, "input"))
|
||||
elif isinstance(input_value, list):
|
||||
for idx, item in enumerate(input_value):
|
||||
rep, field = self._responses_item_repr(item)
|
||||
if rep is None:
|
||||
continue # no scannable text on this item shape
|
||||
batch.append(rep)
|
||||
targets.append((idx, field))
|
||||
self._collect_responses_item(batch, setters, input_value, idx, item)
|
||||
|
||||
@staticmethod
|
||||
def _prompt_variables(data: dict) -> dict | None:
|
||||
"""Responses prompt-template variables carry user data (the template
|
||||
id/version do not). None when there are no scannable variables;
|
||||
/v1/completions sends `prompt` as a string, which is not this surface."""
|
||||
prompt = data.get("prompt")
|
||||
if isinstance(prompt, dict) and isinstance(prompt.get("variables"), dict):
|
||||
return prompt["variables"]
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _has_aux_fields(data: dict) -> bool:
|
||||
"""True when the request carries one of the auxiliary text fields scanned
|
||||
by _collect_aux_fields; keeps the pre-call early return from skipping a
|
||||
request whose only user text sits in those fields."""
|
||||
prediction = data.get("prediction")
|
||||
if isinstance(prediction, dict) and "content" in prediction:
|
||||
return True
|
||||
web_search = data.get("web_search_options")
|
||||
if isinstance(web_search, dict) and "user_location" in web_search:
|
||||
return True
|
||||
suffix = data.get("suffix")
|
||||
if isinstance(suffix, str) and suffix:
|
||||
return True
|
||||
# /v1/completions user text: a string prompt or the batch list shape.
|
||||
# (A dict prompt is the Responses template, handled by _prompt_variables.)
|
||||
prompt = data.get("prompt")
|
||||
return isinstance(prompt, (str, list)) and bool(prompt)
|
||||
|
||||
def _collect_prompt_variables(self, data: dict, batch: list, setters: list) -> None:
|
||||
variables = self._prompt_variables(data)
|
||||
if variables is None:
|
||||
return
|
||||
for key, value in variables.items():
|
||||
self._collect_data_value(batch, setters, value, self._make_setter(variables, key))
|
||||
|
||||
def _collect_aux_fields(self, data: dict, batch: list, setters: list) -> None:
|
||||
"""The request-side text fields outside messages/input that LiteLLM
|
||||
forwards verbatim (core parity): chat `prediction.content`,
|
||||
`web_search_options.user_location`, and the completions `prompt` (string
|
||||
or batch list) and `suffix`. Request inputs only, nothing to restore."""
|
||||
prediction = data.get("prediction")
|
||||
if isinstance(prediction, dict) and prediction.get("content") is not None:
|
||||
self._collect_data_value(batch, setters, prediction["content"], self._make_setter(prediction, "content"))
|
||||
web_search = data.get("web_search_options")
|
||||
if isinstance(web_search, dict) and web_search.get("user_location") is not None:
|
||||
self._collect_data_value(
|
||||
batch, setters, web_search["user_location"], self._make_setter(web_search, "user_location")
|
||||
)
|
||||
suffix = data.get("suffix")
|
||||
if isinstance(suffix, str) and suffix:
|
||||
self._add_leaf(batch, setters, suffix, self._toplevel_setter(data, "suffix"))
|
||||
prompt = data.get("prompt")
|
||||
if isinstance(prompt, str) and prompt:
|
||||
self._add_leaf(batch, setters, prompt, self._toplevel_setter(data, "prompt"))
|
||||
elif isinstance(prompt, list):
|
||||
# The batch shape: string leaves are scrubbed, a tokenized
|
||||
# (integer-array) prompt passes through unscanned as documented.
|
||||
for idx, part in enumerate(prompt):
|
||||
self._add_leaf(batch, setters, part, self._make_setter(prompt, idx))
|
||||
|
||||
async def _anonymize_request(self, data: dict, engine: Any) -> Any:
|
||||
"""Anonymize every scanned request surface in place under ONE shared
|
||||
mapping (chat `messages`, the Responses `input` and `prompt.variables`,
|
||||
and the auxiliary prompt/suffix/prediction/user_location fields), so no
|
||||
source is left untouched when a crafted request carries several. Every
|
||||
scanned string flows through the engine's single choke point, so the block
|
||||
gate and the fail-closed policy apply unchanged. Returns the mapping, or
|
||||
None if there was nothing to anonymize."""
|
||||
messages = data.get("messages")
|
||||
msg_list = messages if isinstance(messages, list) else []
|
||||
# Chat messages are scanned natively (multimodal content, bare-string
|
||||
# lists, per-message block gate); everything else is registered as string
|
||||
# leaves appended after them, sharing this one process_request pass.
|
||||
batch: list = list(msg_list)
|
||||
setters: list = []
|
||||
self._collect_responses_input(data, batch, setters)
|
||||
self._collect_prompt_variables(data, batch, setters)
|
||||
self._collect_aux_fields(data, batch, setters)
|
||||
|
||||
if not batch:
|
||||
return None
|
||||
|
|
@ -242,25 +455,8 @@ class PrivaiteGuardrail(CustomGuardrail):
|
|||
if msg_list:
|
||||
# msg_list is data["messages"] (same object) -> mutate it in place.
|
||||
msg_list[:] = anonymized[:n]
|
||||
|
||||
for (target, field), anon in zip(targets, anonymized[n:]):
|
||||
if target == "str":
|
||||
new_text = anon.get("content", input_value)
|
||||
data["input"] = new_text
|
||||
self._overwrite_snapshot_input(data, new_text)
|
||||
continue
|
||||
# every non-"str" target came from the isinstance(input_value, list)
|
||||
# branch above, so input_value is a list here.
|
||||
if not isinstance(input_value, list):
|
||||
continue
|
||||
if field is None:
|
||||
input_value[target] = anon
|
||||
elif field == "__str__":
|
||||
input_value[target] = anon.get("content", input_value[target])
|
||||
else:
|
||||
new_item = dict(input_value[target])
|
||||
new_item[field] = anon.get("content", new_item[field])
|
||||
input_value[target] = new_item
|
||||
for setter, anon in zip(setters, anonymized[n:]):
|
||||
setter(anon.get("content", ""))
|
||||
return mapping
|
||||
|
||||
async def async_pre_call_hook(self, user_api_key_dict: Any, cache: Any, data: dict, call_type: str) -> dict:
|
||||
|
|
@ -272,7 +468,12 @@ class PrivaiteGuardrail(CustomGuardrail):
|
|||
if isinstance(metadata, dict):
|
||||
metadata.pop("privaite_map", None)
|
||||
|
||||
if not data.get("messages") and not data.get("input"):
|
||||
if (
|
||||
not data.get("messages")
|
||||
and not data.get("input")
|
||||
and self._prompt_variables(data) is None
|
||||
and not self._has_aux_fields(data)
|
||||
):
|
||||
return data
|
||||
|
||||
engine = await self._engine_for(self._languages())
|
||||
|
|
@ -297,16 +498,30 @@ class PrivaiteGuardrail(CustomGuardrail):
|
|||
data.setdefault("metadata", {})["privaite_map"] = dict(mapping.get_all_fakes())
|
||||
return data
|
||||
|
||||
async def _restore_audio_transcript(self, message: Any, engine: Any, mapping: Any) -> None:
|
||||
"""Restore the audio transcript on a response message (restore parity with
|
||||
the core: an audio reply carries its text in audio.transcript, where the
|
||||
model echoes placeholders like anywhere else)."""
|
||||
audio = getattr(message, "audio", None)
|
||||
if audio is None:
|
||||
return
|
||||
transcript = _obj_get(audio, "transcript")
|
||||
if isinstance(transcript, str) and transcript:
|
||||
_obj_set(audio, "transcript", await engine.process_response(transcript, mapping))
|
||||
|
||||
async def _restore_message(self, message: Any, engine: Any, mapping: Any) -> None:
|
||||
"""Restore originals in one response message: content, the reasoning
|
||||
trace, tool-call args and the legacy function_call args."""
|
||||
trace, the refusal, the audio transcript, tool-call args and the legacy
|
||||
function_call args."""
|
||||
content = getattr(message, "content", None)
|
||||
if isinstance(content, str) and content:
|
||||
message.content = await engine.process_response(content, mapping)
|
||||
for field in ("reasoning_content", "reasoning"):
|
||||
# A refusal can quote the request, so it carries placeholders too.
|
||||
for field in ("reasoning_content", "reasoning", "refusal"):
|
||||
value = getattr(message, field, None)
|
||||
if isinstance(value, str) and value:
|
||||
setattr(message, field, await engine.process_response(value, mapping))
|
||||
await self._restore_audio_transcript(message, engine, mapping)
|
||||
for tool_call in getattr(message, "tool_calls", None) or []:
|
||||
fn = getattr(tool_call, "function", None)
|
||||
if fn is None:
|
||||
|
|
@ -356,10 +571,26 @@ class PrivaiteGuardrail(CustomGuardrail):
|
|||
await self._restore_responses_output(response, engine, mapping)
|
||||
return response
|
||||
|
||||
def _restore_delta(self, delta: Any, index: int, finished: bool, restore) -> None:
|
||||
"""Restore one streamed delta in place: text content, the reasoning
|
||||
trace, streamed tool-call argument fragments (per tool_call index) and
|
||||
the legacy function_call."""
|
||||
def _restore_delta_audio(self, delta: Any, index: int, finished: bool, restore: Any) -> None:
|
||||
"""Feed streamed audio transcript fragments through their own restore
|
||||
buffer and flush the held tail on the finish chunk, creating the audio
|
||||
carrier when the finish delta has none so the tail is never dropped."""
|
||||
audio = getattr(delta, "audio", None)
|
||||
fragment = _obj_get(audio, "transcript") if audio is not None else None
|
||||
if not isinstance(fragment, str):
|
||||
fragment = ""
|
||||
if not fragment and not finished:
|
||||
return
|
||||
restored = restore(("audio", index), fragment, finished)
|
||||
if audio is not None:
|
||||
_obj_set(audio, "transcript", restored)
|
||||
elif restored:
|
||||
_obj_set(delta, "audio", {"transcript": restored})
|
||||
|
||||
def _restore_delta(self, delta: Any, index: int, finished: bool, restore: Any) -> None:
|
||||
"""Restore one streamed delta in place: text content, the reasoning trace,
|
||||
the refusal, the audio transcript, streamed tool-call argument fragments
|
||||
(per tool_call index) and the legacy function_call."""
|
||||
content = getattr(delta, "content", None) or ""
|
||||
restored = restore(("content", index), content, finished)
|
||||
# Overwrite whenever there was input text (even if the whole fragment is
|
||||
|
|
@ -368,10 +599,12 @@ class PrivaiteGuardrail(CustomGuardrail):
|
|||
# keeps its None instead of becoming "".
|
||||
if content or restored:
|
||||
delta.content = restored
|
||||
for field in ("reasoning_content", "reasoning"):
|
||||
# A refusal can quote the request, so it carries placeholders too.
|
||||
for field in ("reasoning_content", "reasoning", "refusal"):
|
||||
value = getattr(delta, field, None)
|
||||
if isinstance(value, str) and (value or finished):
|
||||
setattr(delta, field, restore((field, index), value, finished))
|
||||
self._restore_delta_audio(delta, index, finished, restore)
|
||||
for tool_call in getattr(delta, "tool_calls", None) or []:
|
||||
fn = getattr(tool_call, "function", None)
|
||||
if fn is None:
|
||||
|
|
|
|||
|
|
@ -862,3 +862,244 @@ async def test_block_entities_fails_closed_when_privaite_too_old(monkeypatch):
|
|||
data = {"messages": [{"role": "user", "content": "marie@acme.com"}]}
|
||||
with pytest.raises(RuntimeError, match="block_entities"):
|
||||
await gr.async_pre_call_hook(None, None, data, "completion")
|
||||
|
||||
|
||||
# --- Responses API: tool output list-of-parts, typed carriers, prompt variables ---
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pre_call_scans_custom_tool_call_output_list_of_parts():
|
||||
# a shell/custom tool that read a file returns its bytes as a list of
|
||||
# {type, text} parts (not a bare string); every text leaf must be scrubbed
|
||||
# while the binary image part is relayed whole (base64 carries nothing a
|
||||
# text detector can find, and rewriting it would corrupt the payload).
|
||||
gr = _make_guardrail()
|
||||
output = [
|
||||
{"type": "output_text", "text": "the file says reach marie@acme.com"},
|
||||
{"type": "input_image", "image_url": "data:image/png;base64,QUJD"},
|
||||
]
|
||||
data = {"input": [{"type": "custom_tool_call_output", "call_id": "c1", "output": output}]}
|
||||
out = await gr.async_pre_call_hook(None, None, data, "aresponses")
|
||||
|
||||
scrubbed = out["input"][0]["output"]
|
||||
assert scrubbed[0]["text"] == "the file says reach <EMAIL_ADDRESS_1>"
|
||||
assert scrubbed[1]["image_url"] == "data:image/png;base64,QUJD" # binary relayed whole
|
||||
assert out["input"][0]["call_id"] == "c1"
|
||||
assert out["metadata"]["privaite_map"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pre_call_scans_typed_action_carrier():
|
||||
# a file_search_call carries user data in its `queries`/`results` fields, not
|
||||
# in `content`/`output`; the old field scan never looked there.
|
||||
gr = _make_guardrail()
|
||||
data = {
|
||||
"input": [
|
||||
{
|
||||
"type": "file_search_call",
|
||||
"id": "fs1",
|
||||
"queries": ["records for Marie Dupont", "orders by marie@acme.com"],
|
||||
"results": [{"text": "found Marie Dupont", "score": 0.9}],
|
||||
}
|
||||
]
|
||||
}
|
||||
out = await gr.async_pre_call_hook(None, None, data, "aresponses")
|
||||
item = out["input"][0]
|
||||
assert item["queries"] == ["records for <PERSON_1>", "orders by <EMAIL_ADDRESS_1>"]
|
||||
assert item["results"][0]["text"] == "found <PERSON_1>"
|
||||
assert item["results"][0]["score"] == 0.9 # non-text leaf untouched
|
||||
assert item["id"] == "fs1"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pre_call_scans_prompt_variables():
|
||||
# Responses prompt-template variables carry user data; the template id/version
|
||||
# do not. A variable can be a bare string or a typed content part.
|
||||
gr = _make_guardrail()
|
||||
data = {
|
||||
"prompt": {
|
||||
"id": "pmpt_123",
|
||||
"version": "2",
|
||||
"variables": {
|
||||
"customer": "Marie Dupont",
|
||||
"note": {"type": "input_text", "text": "email marie@acme.com"},
|
||||
},
|
||||
}
|
||||
}
|
||||
out = await gr.async_pre_call_hook(None, None, data, "aresponses")
|
||||
variables = out["prompt"]["variables"]
|
||||
assert variables["customer"] == "<PERSON_1>"
|
||||
assert variables["note"]["text"] == "email <EMAIL_ADDRESS_1>"
|
||||
assert out["prompt"]["id"] == "pmpt_123" # template id is not user data
|
||||
assert out["prompt"]["version"] == "2"
|
||||
assert out["metadata"]["privaite_map"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pre_call_block_gate_fires_on_prompt_variables():
|
||||
# the block gate must reach a blocked type sitting only in prompt.variables;
|
||||
# it is only reachable because the variable is scanned in the first place.
|
||||
from fastapi import HTTPException
|
||||
|
||||
gr = _make_guardrail(block_entities=["EMAIL_ADDRESS"])
|
||||
data = {"prompt": {"id": "pmpt_1", "variables": {"to": "marie@acme.com"}}}
|
||||
with pytest.raises(HTTPException) as ei:
|
||||
await gr.async_pre_call_hook(None, None, data, "aresponses")
|
||||
assert ei.value.status_code == 400
|
||||
|
||||
|
||||
# --- Auxiliary request fields: completions prompt/suffix, prediction, user_location ---
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pre_call_scans_completions_prompt_string_and_suffix_and_fixes_snapshot():
|
||||
gr = _make_guardrail()
|
||||
body = {"prompt": "Complete for Marie Dupont", "suffix": "signed marie@acme.com"}
|
||||
data = {
|
||||
"prompt": "Complete for Marie Dupont",
|
||||
"suffix": "signed marie@acme.com",
|
||||
"proxy_server_request": {"body": body},
|
||||
}
|
||||
out = await gr.async_pre_call_hook(None, None, data, "text_completion")
|
||||
|
||||
assert out["prompt"] == "Complete for <PERSON_1>"
|
||||
assert out["suffix"] == "signed <EMAIL_ADDRESS_1>"
|
||||
# the detached snapshot copy must be overwritten too, or raw PII leaks through it
|
||||
assert body["prompt"] == out["prompt"]
|
||||
assert body["suffix"] == out["suffix"]
|
||||
assert "Marie Dupont" not in body["prompt"]
|
||||
assert out["metadata"]["privaite_map"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pre_call_scans_completions_prompt_batch_list_and_skips_tokens():
|
||||
# the /v1/completions batch shape: string leaves are scrubbed; a tokenized
|
||||
# (integer-array) prompt passes through unscanned as documented.
|
||||
gr = _make_guardrail()
|
||||
prompt = ["reach marie@acme.com", [1, 15, 27]]
|
||||
data = {"prompt": prompt}
|
||||
out = await gr.async_pre_call_hook(None, None, data, "text_completion")
|
||||
assert out["prompt"][0] == "reach <EMAIL_ADDRESS_1>"
|
||||
assert out["prompt"][1] == [1, 15, 27] # token array untouched
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pre_call_scans_prediction_content():
|
||||
# chat predicted outputs carry the client's current document verbatim.
|
||||
gr = _make_guardrail()
|
||||
prediction = {"type": "content", "content": "draft addressed to Marie Dupont"}
|
||||
data = {"messages": [{"role": "user", "content": "edit this"}], "prediction": prediction}
|
||||
out = await gr.async_pre_call_hook(None, None, data, "completion")
|
||||
# prediction dict is aliased by the body snapshot, so it is rewritten in place
|
||||
assert prediction["content"] == "draft addressed to <PERSON_1>"
|
||||
assert out["metadata"]["privaite_map"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pre_call_scans_web_search_user_location():
|
||||
gr = _make_guardrail()
|
||||
web_search = {"user_location": {"type": "approximate", "approximate": {"city": "Marie Dupont"}}}
|
||||
data = {"messages": [{"role": "user", "content": "weather?"}], "web_search_options": web_search}
|
||||
await gr.async_pre_call_hook(None, None, data, "completion")
|
||||
assert web_search["user_location"]["approximate"]["city"] == "<PERSON_1>"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pre_call_only_aux_field_is_not_skipped():
|
||||
# a request whose ONLY user text is an auxiliary field (no messages/input)
|
||||
# must not be short-circuited by the pre-call early return.
|
||||
gr = _make_guardrail()
|
||||
data = {"suffix": "signed marie@acme.com"}
|
||||
out = await gr.async_pre_call_hook(None, None, data, "text_completion")
|
||||
assert out["suffix"] == "signed <EMAIL_ADDRESS_1>"
|
||||
assert out["metadata"]["privaite_map"]
|
||||
|
||||
|
||||
# --- Restore parity: refusal and audio transcript (non-streaming + streaming) ---
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_call_restores_refusal_and_audio_transcript():
|
||||
gr = _make_guardrail()
|
||||
data = {"metadata": {"privaite_map": _FAKES}}
|
||||
|
||||
# audio as an object (transcript attribute) and refusal quoting the request.
|
||||
obj_msg = types.SimpleNamespace(
|
||||
content=None,
|
||||
tool_calls=None,
|
||||
function_call=None,
|
||||
refusal="I won't email <PERSON_1>",
|
||||
audio=types.SimpleNamespace(transcript="calling <PERSON_1> now"),
|
||||
)
|
||||
# audio as a dict on a second choice.
|
||||
dict_msg = types.SimpleNamespace(
|
||||
content=None,
|
||||
tool_calls=None,
|
||||
function_call=None,
|
||||
audio={"transcript": "reach <EMAIL_ADDRESS_1>"},
|
||||
)
|
||||
response = types.SimpleNamespace(
|
||||
choices=[types.SimpleNamespace(message=obj_msg), types.SimpleNamespace(message=dict_msg)]
|
||||
)
|
||||
out = await gr.async_post_call_success_hook(data, None, response)
|
||||
|
||||
assert out.choices[0].message.refusal == "I won't email Marie Dupont"
|
||||
assert out.choices[0].message.audio.transcript == "calling Marie Dupont now"
|
||||
assert out.choices[1].message.audio["transcript"] == "reach marie@acme.com"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_restores_refusal():
|
||||
gr = _make_guardrail()
|
||||
request_data = {"metadata": {"privaite_map": _FAKES}}
|
||||
|
||||
def _chunk(refusal, finish=None):
|
||||
return types.SimpleNamespace(
|
||||
choices=[
|
||||
types.SimpleNamespace(
|
||||
index=0,
|
||||
delta=types.SimpleNamespace(content=None, tool_calls=None, function_call=None, refusal=refusal),
|
||||
finish_reason=finish,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
async def _source():
|
||||
yield _chunk("I won't reach <PER") # placeholder split across chunks
|
||||
yield _chunk("SON_1>", finish="stop")
|
||||
|
||||
chunks = await _collect(gr.async_post_call_streaming_iterator_hook(None, _source(), request_data))
|
||||
refusal = "".join(getattr(c.choices[0].delta, "refusal", None) or "" for c in chunks)
|
||||
assert refusal == "I won't reach Marie Dupont"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_restores_audio_transcript():
|
||||
# streamed audio transcript fragments are de-anonymized with a placeholder
|
||||
# split across chunks reassembled in the audio-segment buffer.
|
||||
gr = _make_guardrail()
|
||||
request_data = {"metadata": {"privaite_map": _FAKES}}
|
||||
|
||||
def _chunk(transcript, finish=None):
|
||||
return types.SimpleNamespace(
|
||||
choices=[
|
||||
types.SimpleNamespace(
|
||||
index=0,
|
||||
delta=types.SimpleNamespace(
|
||||
content=None,
|
||||
tool_calls=None,
|
||||
function_call=None,
|
||||
audio=types.SimpleNamespace(transcript=transcript),
|
||||
),
|
||||
finish_reason=finish,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
async def _source():
|
||||
yield _chunk("Hi <PER")
|
||||
yield _chunk("SON_1>", finish="stop")
|
||||
|
||||
chunks = await _collect(gr.async_post_call_streaming_iterator_hook(None, _source(), request_data))
|
||||
transcript = "".join(c.choices[0].delta.audio.transcript or "" for c in chunks if c.choices[0].delta.audio)
|
||||
assert transcript == "Hi Marie Dupont"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue