diff --git a/litellm/proxy/guardrails/guardrail_hooks/privaite/engine.py b/litellm/proxy/guardrails/guardrail_hooks/privaite/engine.py index 3cc873f7b91..bcccc53eb32 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/privaite/engine.py +++ b/litellm/proxy/guardrails/guardrail_hooks/privaite/engine.py @@ -46,6 +46,8 @@ class PrivaiteEngine(Protocol): async def process_request(self, messages: Sequence[object]) -> tuple[Sequence[object], RestoreMapping]: ... + async def gate_document(self, value: object) -> None: ... + async def process_response(self, text: str, mapping: RestoreMapping) -> str: ... diff --git a/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py b/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py index 4ec7c3f0748..dcb0003f826 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py +++ b/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py @@ -6,12 +6,17 @@ the reversible map in the request metadata, and the post-call hooks consume it t the real values in the response, streaming included. The request surface scanned matches the PrivAiTe core: chat `messages`, every Responses `input` item shape and `prompt.variables`, the `/v1/completions` `prompt` and `suffix`, and the chat -`prediction.content` and `web_search_options.user_location` fields. Responses API +`prediction.content` and `web_search_options.user_location` fields. Every string leaf +of `extra_body` is scanned too, since LiteLLM merges it over the request it sends. +Other top-level fields are not scanned (`user`, `metadata`, tool definitions, and any +field LiteLLM does not know, which it forwards to the provider as it is). Responses API streaming restore is not implemented yet. -If `block_entities` is configured, a request containing any of those PII types is -rejected with a 400 before anything is forwarded. The error names the types, never the -values. +If `block_entities` is configured, a request with any of those PII types in a scanned +field is rejected with a 400 before anything is forwarded. The error names the types, +never the values. The gate also reads the caller's own prompt (Responses `instructions`, +Anthropic `system`), which is relayed verbatim, and refuses a prompt or input sent as +token IDs, which it could not check. There is no detection or masking logic here, it all lives in the `privaite` package, which is imported lazily so this module stays importable for the guardrail registry @@ -41,11 +46,12 @@ from .engine import ( ) from .payload import JsonObject, as_array, as_object, read_field, read_list, write_field from .request_walk import ( + agent_prompts, aux_field_slots, - has_aux_fields, + carries_request_text, prompt_variable_slots, - prompt_variables, responses_input_slots, + tokenized_fields, ) from .restore import ( StreamRestorer, @@ -102,6 +108,20 @@ def _pop_restore_map(request: JsonObject) -> tuple[tuple[str, str], ...]: return tuple((fake, original) for fake, original in pairs if isinstance(original, str)) +def _refuse_tokenized(data: JsonObject) -> None: + """With block_entities set, a prompt or input sent as token IDs is refused: no + detector reads token IDs, so the gate could not check it. Without block_entities it + passes unscanned, as in the PrivAiTe core.""" + fields: Final = tokenized_fields(data) + if fields: + from fastapi import HTTPException + + raise HTTPException( + status_code=400, + detail={"error": f"block_entities cannot check a tokenized {', '.join(fields)}; send it as text"}, + ) + + def _pseudo_message(text: str) -> JsonObject: """A string leaf wrapped as a message, so it crosses the engine's single choke point like chat content.""" return {"role": "user", "content": text} @@ -194,17 +214,19 @@ class PrivaiteGuardrail(CustomGuardrail): # metadata is caller controlled, so an incoming privaite_map is never trusted: # the post-call hooks may only restore from a map this pre-call pass produced. _drop_restore_map(data) - if ( - not data.get("messages") - and not data.get("input") - and prompt_variables(data) is None - and not has_aux_fields(data) - ): + prompts: Final = agent_prompts(data) if self.block_entities else () + if not prompts and not carries_request_text(data): return data + if self.block_entities: + _refuse_tokenized(data) engine: Final = await self._engine_for(self._languages()) blocked: Final = blocked_error() try: + # The agent prompt is relayed verbatim, but a blocked type in it still + # rejects the request, before the caller's dict is touched. + for prompt in prompts: + await engine.gate_document(prompt) mapping: Final = await self._anonymize_request(data, engine) except blocked as exc: from fastapi import HTTPException diff --git a/litellm/proxy/guardrails/guardrail_hooks/privaite/request_walk.py b/litellm/proxy/guardrails/guardrail_hooks/privaite/request_walk.py index c07b9aa539d..ef387d50cad 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/privaite/request_walk.py +++ b/litellm/proxy/guardrails/guardrail_hooks/privaite/request_walk.py @@ -63,6 +63,10 @@ GENERIC_TEXT_FIELDS: Final = ("output", "arguments", "input", "text", "reason") """Text-bearing fields of an unknown Responses item. Every one is scanned, since an mcp_call carries both `arguments` and `output`.""" +AGENT_PROMPT_FIELDS: Final = ("instructions", "system") +"""The caller's own prompt: Responses `instructions`, Anthropic `system`. Relayed +verbatim, as the PrivAiTe gateway does, but read by the block gate.""" + def leaf(text: object, write: Write) -> Iterator[Slot]: if isinstance(text, str) and text: @@ -208,6 +212,8 @@ def has_aux_fields(data: JsonObject) -> bool: web_search: Final = as_object(data.get("web_search_options")) if web_search is not None and "user_location" in web_search: return True + if data.get("extra_body"): + return True suffix: Final = data.get("suffix") if isinstance(suffix, str) and suffix: return True @@ -217,14 +223,21 @@ def has_aux_fields(data: JsonObject) -> bool: def aux_field_slots(data: JsonObject) -> Iterator[Slot]: """Request text outside messages and input that LiteLLM forwards verbatim: chat - `prediction.content`, `web_search_options.user_location`, and the completions `suffix` - and `prompt`. A tokenized (integer array) prompt passes through unscanned.""" + `prediction.content`, `web_search_options.user_location`, every string leaf of + `extra_body`, and the completions `suffix` and `prompt`. A tokenized (integer array) + prompt passes through unscanned. + + LiteLLM merges `extra_body` over the request it sends, so `extra_body.messages` + replaced the scrubbed messages on the wire. Its keys are provider specific + (`contents` for Gemini, say), so it is walked whole rather than by field name.""" prediction: Final = as_object(data.get("prediction")) if prediction is not None and prediction.get("content") is not None: yield from data_slots(prediction["content"], field_writer(prediction, "content")) web_search: Final = as_object(data.get("web_search_options")) if web_search is not None and web_search.get("user_location") is not None: yield from data_slots(web_search["user_location"], field_writer(web_search, "user_location")) + if data.get("extra_body") is not None: + yield from data_slots(data["extra_body"], top_level_writer(data, "extra_body")) yield from leaf(data.get("suffix"), top_level_writer(data, "suffix")) prompt: Final = data.get("prompt") if isinstance(prompt, str): @@ -235,3 +248,28 @@ def aux_field_slots(data: JsonObject) -> Iterator[Slot]: return for index, part in enumerate(parts): yield from leaf(part, item_writer(parts, index)) + + +def carries_request_text(data: JsonObject) -> bool: + """Whether the request has a field the anonymization pass reads.""" + return bool(data.get("messages") or data.get("input")) or prompt_variables(data) is not None or has_aux_fields(data) + + +def agent_prompts(data: JsonObject) -> tuple[object, ...]: + """The non-empty agent prompt fields, for the block gate.""" + return tuple(data[name] for name in AGENT_PROMPT_FIELDS if data.get(name)) + + +def _is_token_list(value: object) -> bool: + items: Final = as_array(value) or () + return len(items) > 0 and all(type(item) is int for item in items) + + +def tokenized_fields(data: JsonObject) -> tuple[str, ...]: + """The completions `prompt` or embeddings `input` sent as token IDs: a list of + integers, or a batch holding one. No detector reads token IDs.""" + return tuple( + name + for name in ("prompt", "input") + if _is_token_list(data.get(name)) or any(_is_token_list(item) for item in as_array(data.get(name)) or ()) + ) diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py index 5a19ce91fe2..0074cc6646b 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py @@ -103,6 +103,16 @@ class _FakeEngine: raise _FakePIIBlockedError(hit) return anonymized, mapping + async def gate_document(self, value): + """The block gate without the rewrite, like PIIEngine.gate_document.""" + self.gated = [*getattr(self, "gated", []), value] + blocked = set(getattr(self.config, "block_entities", None) or []) + mapping = _FakeMapping() + self._scrub(json.dumps(value), mapping) + detected = {fake.strip("<>").rsplit("_", 1)[0] for fake in mapping.get_all_fakes()} + if detected & blocked: + raise _FakePIIBlockedError(detected & blocked) + async def process_response(self, text, _mapping): for fake, original in _FAKES.items(): text = text.replace(fake, original) @@ -1458,3 +1468,106 @@ async def test_streaming_restores_copies_and_leaves_the_chunks_litellm_reassembl assert [chunk.choices[0].delta.content for chunk in source] == ["Hi "], ( "the chunks LiteLLM keeps to rebuild the logged and cached response must keep the placeholder" ) + + +@pytest.mark.asyncio +async def test_pre_call_scrubs_extra_body_in_place(): + # LiteLLM merges extra_body over the request it sends, so extra_body.messages + # would otherwise replace the scrubbed messages on the wire. + gr = _make_guardrail() + extra_body = {"messages": [{"role": "user", "content": "Email marie@acme.com"}], "top_k": 5} + data = {"messages": [{"role": "user", "content": "hi"}], "extra_body": extra_body} + data["proxy_server_request"] = {"body": dict(data)} + out = await gr.async_pre_call_hook(None, None, data, "completion") + + assert out["extra_body"] is extra_body + assert extra_body == {"messages": [{"role": "user", "content": "Email "}], "top_k": 5} + assert data["proxy_server_request"]["body"]["extra_body"]["messages"][0]["content"] == "Email " + assert out["metadata"]["privaite_map"] == {"": "marie@acme.com"} + + +@pytest.mark.asyncio +async def test_pre_call_block_gate_reads_extra_body(): + from fastapi import HTTPException + + gr = _make_guardrail(block_entities=["EMAIL_ADDRESS"]) + data = { + "messages": [{"role": "user", "content": "hi"}], + "extra_body": {"messages": [{"role": "user", "content": "marie@acme.com"}]}, + } + with pytest.raises(HTTPException) as ei: + await gr.async_pre_call_hook(None, None, data, "completion") + assert ei.value.status_code == 400 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "data", + [ + {"input": "hi", "instructions": "The user is marie@acme.com."}, + {"system": [{"type": "text", "text": "The user is marie@acme.com."}]}, + ], +) +async def test_pre_call_block_gate_reads_the_agent_prompt(data): + from fastapi import HTTPException + + gr = _make_guardrail(block_entities=["EMAIL_ADDRESS"]) + with pytest.raises(HTTPException) as ei: + await gr.async_pre_call_hook(None, None, data, "aresponses") + assert ei.value.status_code == 400 + + +@pytest.mark.asyncio +async def test_pre_call_agent_prompt_is_relayed_verbatim(): + gr = _make_guardrail(block_entities=["EMAIL_ADDRESS"]) + data = {"input": "hi", "instructions": "The user is Marie Dupont."} + out = await gr.async_pre_call_hook(None, None, data, "aresponses") + assert out["instructions"] == "The user is Marie Dupont." + assert gr._engine.gated == ["The user is Marie Dupont."] + + +@pytest.mark.asyncio +async def test_pre_call_agent_prompt_is_not_read_without_block_entities(): + gr = _make_guardrail() + data = {"input": "hi", "instructions": "The user is marie@acme.com."} + out = await gr.async_pre_call_hook(None, None, data, "aresponses") + assert out["instructions"] == "The user is marie@acme.com." + assert getattr(gr._engine, "gated", []) == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("data", "field"), + [ + ({"prompt": [9906, 1917]}, "prompt"), + ({"prompt": [[9906], [1917]]}, "prompt"), + ({"prompt": ["hi", [9906]]}, "prompt"), + ({"input": [9906, 1917]}, "input"), + ], +) +async def test_pre_call_block_gate_refuses_token_ids(data, field): + # No detector reads token IDs, so the gate cannot check them. + from fastapi import HTTPException + + gr = _make_guardrail(block_entities=["EMAIL_ADDRESS"]) + with pytest.raises(HTTPException) as ei: + await gr.async_pre_call_hook(None, None, data, "atext_completion") + assert ei.value.status_code == 400 + assert f"tokenized {field}" in ei.value.detail["error"] + + +@pytest.mark.asyncio +async def test_pre_call_token_ids_pass_without_block_entities(): + gr = _make_guardrail() + out = await gr.async_pre_call_hook(None, None, {"prompt": [9906, 1917]}, "atext_completion") + assert out["prompt"] == [9906, 1917] + + +@pytest.mark.asyncio +async def test_pre_call_extra_body_only_request_is_scanned(): + # No top-level messages at all: the early return must not skip the scan. + gr = _make_guardrail() + extra_body = {"messages": [{"role": "user", "content": "Email marie@acme.com"}]} + out = await gr.async_pre_call_hook(None, None, {"extra_body": extra_body}, "completion") + assert extra_body["messages"][0]["content"] == "Email " + assert out["metadata"]["privaite_map"] == {"": "marie@acme.com"}