From 9fef84b9a9e6f5fd85acb89389ba71f52c84a8d1 Mon Sep 17 00:00:00 2001 From: ningyixi Date: Thu, 15 Oct 2026 02:12:22 +0800 Subject: [PATCH] fix(guardrails): keep image-generation PII masks after pre_call rebuild The /v1/images/generations endpoint copied prompt into temporary messages for guardrails, then rebuilt prompt from those messages. Unified image handlers wrote the redacted text to data["prompt"] and left messages unchanged, so the rebuild restored the original PII (issue #46473). Prefer a prompt changed by the pre_call hook over the message rebuild, and sync temporary user messages when the image handler redacts the prompt. --- .../guardrail_translation/handler.py | 11 +- litellm/proxy/image_endpoints/endpoints.py | 38 +++++- ...mage_generation_prompt_after_guardrails.py | 121 ++++++++++++++++++ 3 files changed, 165 insertions(+), 5 deletions(-) create mode 100644 tests/unit/llms/openai/image_generation/test_image_generation_prompt_after_guardrails.py diff --git a/litellm/llms/openai/image_generation/guardrail_translation/handler.py b/litellm/llms/openai/image_generation/guardrail_translation/handler.py index b1d64fb1c09..1e9a15534e8 100644 --- a/litellm/llms/openai/image_generation/guardrail_translation/handler.py +++ b/litellm/llms/openai/image_generation/guardrail_translation/handler.py @@ -64,7 +64,16 @@ class OpenAIImageGenerationHandler(BaseTranslation): logging_obj=litellm_logging_obj, ) guardrailed_texts: Final = guardrailed_inputs.get("texts", []) - data["prompt"] = guardrailed_texts[0] if guardrailed_texts else prompt + redacted_prompt: Final = guardrailed_texts[0] if guardrailed_texts else prompt + data["prompt"] = redacted_prompt + # Keep temporary chat messages in sync with the redacted prompt so + # proxy endpoints that rebuild prompt from messages do not restore PII. + messages = data.get("messages") + if isinstance(messages, list): + for message in messages: + if isinstance(message, dict) and message.get("role") == "user": + message["content"] = redacted_prompt + break verbose_proxy_logger.debug( "OpenAI Image Generation: Applied guardrail to prompt. Original length: %d, New length: %d", diff --git a/litellm/proxy/image_endpoints/endpoints.py b/litellm/proxy/image_endpoints/endpoints.py index 28fb9baf2aa..bec7c2d000a 100644 --- a/litellm/proxy/image_endpoints/endpoints.py +++ b/litellm/proxy/image_endpoints/endpoints.py @@ -45,6 +45,37 @@ MASK_ARRAY_FIELD: Final = "mask[]" BRACKETED_FILE_FIELDS: Final = frozenset({IMAGE_ARRAY_FIELD, MASK_ARRAY_FIELD}) +def resolve_image_generation_prompt_after_guardrails( + data: dict, + *, + prompt_before_hook: object, +) -> None: + """ + Restore `data["prompt"]` after `pre_call_hook` and drop temporary messages. + + The image-generation endpoint copies `prompt` into chat-shaped `messages` so + guardrails can run. Unified image handlers rewrite `data["prompt"]` and leave + `messages` unchanged; chat-shaped handlers rewrite `messages` instead. + Prefer a prompt that was changed relative to `prompt_before_hook`; otherwise + rebuild from `messages` (issue #46473). + """ + messages = data.get("messages") + prompt_after_hook = data.get("prompt") + if isinstance(messages, list) and messages: + messages_prompt = get_str_from_messages(messages) + if ( + isinstance(prompt_after_hook, str) + and prompt_before_hook is not None + and prompt_after_hook != prompt_before_hook + ): + data["prompt"] = prompt_after_hook + else: + data["prompt"] = messages_prompt + data.pop("messages", None) + + + + async def uploadfile_to_bytesio(upload: UploadFile) -> io.BytesIO: """ Read a FastAPI UploadFile into a BytesIO and set .name so OpenAI SDK @@ -145,10 +176,9 @@ async def image_generation( user_api_key_dict=user_api_key_dict, data=data, call_type="image_generation" ) - messages: Final = data.get("messages") - if isinstance(messages, list) and messages: - data["prompt"] = get_str_from_messages(messages) - data.pop("messages", None) + resolve_image_generation_prompt_after_guardrails( + data, prompt_before_hook=prompt_value + ) ## ROUTE TO CORRECT ENDPOINT ## llm_call: Final = await route_request( diff --git a/tests/unit/llms/openai/image_generation/test_image_generation_prompt_after_guardrails.py b/tests/unit/llms/openai/image_generation/test_image_generation_prompt_after_guardrails.py new file mode 100644 index 00000000000..01772773159 --- /dev/null +++ b/tests/unit/llms/openai/image_generation/test_image_generation_prompt_after_guardrails.py @@ -0,0 +1,121 @@ +""" +Regression for #46473: pre-call PII masking on /v1/images/generations. + +The endpoint copies prompt -> messages for guardrails, then used to rebuild +prompt from messages. Image handlers write the redacted text to data["prompt"] +and left messages alone, so the rebuild restored the original PII. +""" + +from typing import Any + +import pytest + +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.llms.openai.image_generation.guardrail_translation.handler import ( + OpenAIImageGenerationHandler, +) +from litellm.proxy.image_endpoints.endpoints import ( + resolve_image_generation_prompt_after_guardrails, +) + + +def test_prefers_guardrail_updated_prompt_over_unchanged_messages(): + original = "Draw a portrait of SSN 123-45-6789" + redacted = "Draw a portrait of SSN " + data = { + "prompt": redacted, + "messages": [{"role": "user", "content": original}], + } + + resolve_image_generation_prompt_after_guardrails( + data, prompt_before_hook=original + ) + + assert data["prompt"] == redacted + assert "messages" not in data + + +def test_rebuilds_from_messages_when_prompt_unchanged(): + original = "A red balloon" + redacted_via_messages = "A balloon" + data = { + "prompt": original, + "messages": [{"role": "user", "content": redacted_via_messages}], + } + + resolve_image_generation_prompt_after_guardrails( + data, prompt_before_hook=original + ) + + assert data["prompt"] == redacted_via_messages + assert "messages" not in data + + +def test_leaves_prompt_when_no_messages(): + data = {"prompt": "unchanged"} + resolve_image_generation_prompt_after_guardrails( + data, prompt_before_hook="unchanged" + ) + assert data["prompt"] == "unchanged" + assert "messages" not in data + + +@pytest.mark.asyncio +async def test_handler_syncs_temporary_messages_with_redacted_prompt(): + class AppendTag(CustomGuardrail): + async def apply_guardrail( + self, inputs: dict, request_data: dict, input_type: str, **kwargs: Any + ) -> dict: + return {"texts": [f"{t} [MASKED]" for t in inputs.get("texts", [])]} + + original = "secret@example.com in a landscape" + data = { + "model": "dall-e-3", + "prompt": original, + "messages": [{"role": "user", "content": original}], + } + result = await OpenAIImageGenerationHandler().process_input_messages( + data, AppendTag(guardrail_name="tag") + ) + assert result["prompt"] == f"{original} [MASKED]" + assert result["messages"][0]["content"] == result["prompt"] + + +@pytest.mark.asyncio +async def test_endpoint_rebuild_keeps_handler_redaction(): + """Exact #46473 sequence: handler redacts prompt; rebuild must not restore PII.""" + + class MaskSSN(CustomGuardrail): + async def apply_guardrail( + self, inputs: dict, request_data: dict, input_type: str, **kwargs: Any + ) -> dict: + texts = inputs.get("texts", []) + return { + "texts": [ + t.replace("219-09-9999", "").replace("Jane Doe", "") + for t in texts + ] + } + + original = "Generate an ID badge for Jane Doe, SSN 219-09-9999" + data = { + "model": "dall-e-3", + "prompt": original, + "messages": [{"role": "user", "content": original}], + } + + data = await OpenAIImageGenerationHandler().process_input_messages( + data, MaskSSN(guardrail_name="mask") + ) + assert "219-09-9999" not in data["prompt"] + assert "Jane Doe" not in data["prompt"] + + # Even if messages were somehow still original, prefer updated prompt. + data["messages"] = [{"role": "user", "content": original}] + resolve_image_generation_prompt_after_guardrails( + data, prompt_before_hook=original + ) + assert "219-09-9999" not in data["prompt"] + assert "Jane Doe" not in data["prompt"] + assert "" in data["prompt"] + assert "" in data["prompt"]