From 51b4886d462a8e5dec3621b13f758187aeb2a019 Mon Sep 17 00:00:00 2001 From: Shivam Sharma <192242476+ShivamSharma43@users.noreply.github.com> Date: Fri, 23 Oct 2026 23:06:39 +0530 Subject: [PATCH] FIX Upload edited images under their own file name in OpenAIImageTarget The edit request used get_data_filename_async() as the multipart file name. That generates a new local storage path with a random name and a .png extension, so the image's real name was never sent, JPEG and WebP inputs were named .png, and the local results directory was sent to the provider. Use Path(image_path).name, as OpenAIVideoTarget already does. Co-Authored-By: Claude Opus 5.5 --- .../openai/openai_image_target.py | 5 +- .../prompt_target/target/test_image_target.py | 46 +++++++++++++++++++ 2 files changed, 50 insertions(+), 1 deletion(-) diff --git a/pyrit/prompt_target/openai/openai_image_target.py b/pyrit/prompt_target/openai/openai_image_target.py index 89dbc703c3..96b6535748 100644 --- a/pyrit/prompt_target/openai/openai_image_target.py +++ b/pyrit/prompt_target/openai/openai_image_target.py @@ -2,6 +2,7 @@ # Licensed under the MIT license. import base64 import logging +from pathlib import Path from typing import Any, Literal from pyrit.common import forward_init_parameters, get_mime_type @@ -232,7 +233,9 @@ async def _send_edit_request_async(self, message: Message) -> Message: category="prompt-memory-entries", value=image_path, data_type="image_path" ) - image_name = str(await img_serializer.get_data_filename_async()) + # Upload the image under its own file name. get_data_filename_async() would generate a new + # local storage path (always .png), not the name of the image being edited. + image_name = Path(image_path).name image_bytes = await img_serializer.read_data_async() image_type = get_mime_type(image_path) diff --git a/tests/unit/prompt_target/target/test_image_target.py b/tests/unit/prompt_target/target/test_image_target.py index 08d8ae5149..4c6072651c 100644 --- a/tests/unit/prompt_target/target/test_image_target.py +++ b/tests/unit/prompt_target/target/test_image_target.py @@ -640,3 +640,49 @@ def test_transparent_background_with_valid_format_succeeds(patch_central_databas ) assert target.background == "transparent" assert target.output_format == valid_format + + +@pytest.mark.parametrize( + "file_names, expected_mime_types", + [ + (["cat_photo.jpg"], ["image/jpeg"]), + (["first.png", "second.jpeg"], ["image/png", "image/jpeg"]), + ], +) +async def test_send_prompt_async_edit_uploads_images_under_their_own_names( + image_target: OpenAIImageTarget, tmp_path, file_names: list[str], expected_mime_types: list[str] +): + image_pieces = [] + for file_name in file_names: + image_path = tmp_path / file_name + image_path.write_bytes(b"image bytes for " + file_name.encode()) + image_pieces.append( + MessagePiece( + role="user", + conversation_id="123", + original_value=str(image_path), + original_value_data_type="image_path", + converted_value_data_type="image_path", + ) + ) + text_piece = MessagePiece(role="user", conversation_id="123", original_value="edit this image") + + mock_response = MagicMock() + mock_image = MagicMock() + mock_image.b64_json = "aGVsbG8=" + mock_response.data = [mock_image] + + with patch.object(image_target._async_client.images, "edit", new_callable=AsyncMock) as mock_edit: + mock_edit.return_value = mock_response + + resp = await image_target.send_prompt_async(message=Message(message_pieces=[text_piece, *image_pieces])) + + sent_images = mock_edit.call_args.kwargs["image"] + if isinstance(sent_images, tuple): + sent_images = [sent_images] + assert sent_images == [ + (file_name, b"image bytes for " + file_name.encode(), mime_type) + for file_name, mime_type in zip(file_names, expected_mime_types, strict=True) + ] + + os.remove(resp[0].message_pieces[0].original_value)