Merge 9fef84b9a9 into f82e4c80ac

litellm/llms/openai/image_generation/guardrail_translation/handler.py

@@ -66,7 +66,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",

litellm/proxy/image_endpoints/endpoints.py

@@ -50,6 +50,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
@@ -163,10 +194,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(

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 <US_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 <COLOR> 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", "<US_SSN>").replace("Jane Doe", "<PERSON>")
+                    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 "<US_SSN>" in data["prompt"]
+    assert "<PERSON>" in data["prompt"]