FIX Preserve user message structure in generic system squash

pyrit/message_normalizer/generic_system_squash.py

@@ -1,10 +1,9 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT license.
 
-
 from pyrit.message_normalizer._helpers import build_squashed_user_message
 from pyrit.message_normalizer.message_normalizer import MessageListNormalizer
-from pyrit.models import Message
+from pyrit.models import Message, MessagePiece
 
 
 class GenericSystemSquashNormalizer(MessageListNormalizer[Message]):
@@ -48,15 +47,56 @@ async def normalize_async(self, messages: list[Message]) -> list[Message]:
                 )
             ]
 
-        # Combine system with first user message
+        user_message_index = next(
+            (i for i, message in enumerate(messages[1:], start=1) if message.api_role == "user"),
+            -1,
+        )
+        if user_message_index == -1:
+            # Preserve the instruction content without rewriting non-user messages.
+            return [
+                build_squashed_user_message(
+                    new_message_content=first_piece.converted_value, source_messages=messages[:1]
+                )
+            ] + list(messages[1:])
+
+        # Combine system with the first user message, preserving non-text pieces (e.g. images) and their order.
         system_content = first_piece.converted_value
-        user_piece = messages[1].get_piece()
-        user_content = user_piece.converted_value
+        user_message = messages[user_message_index]
+        # Propagate prompt_metadata from the user message's first piece so downstream normalizers
+        # (e.g. JsonSchemaNormalizer) still see request-level metadata after squashing.
+        propagated_metadata = dict(user_message.message_pieces[0].prompt_metadata)
+        text_piece_index = next(
+            (i for i, piece in enumerate(user_message.message_pieces) if piece.converted_value_data_type == "text"),
+            -1,
+        )
 
-        combined_content = f"### Instructions ###\n\n{system_content}\n\n######\n\n{user_content}"
+        if text_piece_index == -1:
+            # No text piece to merge into; prepend an instruction-only text piece so non-text pieces are preserved.
+            template_piece = user_message.get_piece()
+            instruction_piece = MessagePiece(
+                role="user",
+                original_value=f"### Instructions ###\n\n{system_content}\n\n######",
+                conversation_id=template_piece.conversation_id,
+                sequence=template_piece.sequence,
+                prompt_metadata=propagated_metadata,
+            )
+            squashed_pieces = [instruction_piece] + list(user_message.message_pieces)
+        else:
+            text_piece = user_message.message_pieces[text_piece_index]
+            combined_piece = MessagePiece(
+                role="user",
+                original_value=f"### Instructions ###\n\n{system_content}\n\n######\n\n{text_piece.converted_value}",
+                conversation_id=text_piece.conversation_id,
+                sequence=text_piece.sequence,
+                prompt_metadata=propagated_metadata,
+            )
+            squashed_pieces = (
+                list(user_message.message_pieces[:text_piece_index])
+                + [combined_piece]
+                + list(user_message.message_pieces[text_piece_index + 1 :])
+            )
 
-        squashed_message = build_squashed_user_message(
-            new_message_content=combined_content, source_messages=messages[:2]
-        )
-        # Return the squashed message followed by remaining messages (skip first two)
-        return [squashed_message] + list(messages[2:])
+        squashed_message = Message(message_pieces=squashed_pieces)
+
+        # Remove system (index 0), replace the first user message with the squashed version, preserve all others
+        return list(messages[1:user_message_index]) + [squashed_message] + list(messages[user_message_index + 1 :])

tests/unit/message_normalizer/test_generic_system_squash_normalizer.py

@@ -69,6 +69,140 @@ async def test_generic_squash_normalize_to_dicts_async():
     assert "User message" in result[0]["pieces"][0]["converted_value"]
 
 
+async def test_generic_squash_preserves_multipart_user_message():
+    """Test that squashing keeps non-text user pieces instead of collapsing to plain text."""
+    conversation_id = "conv-1"
+    messages = [
+        _make_message("system", "System message"),
+        Message(
+            message_pieces=[
+                MessagePiece(
+                    role="user",
+                    original_value="User message",
+                    conversation_id=conversation_id,
+                    sequence=0,
+                ),
+                MessagePiece(
+                    role="user",
+                    original_value="/tmp/example.png",
+                    original_value_data_type="image_path",
+                    conversation_id=conversation_id,
+                    sequence=0,
+                ),
+            ]
+        ),
+    ]
+
+    result = await GenericSystemSquashNormalizer().normalize_async(messages)
+
+    assert len(result) == 1
+    assert result[0].api_role == "user"
+    assert len(result[0].message_pieces) == 2
+    assert result[0].get_value() == "### Instructions ###\n\nSystem message\n\n######\n\nUser message"
+    assert result[0].message_pieces[1].converted_value == "/tmp/example.png"
+    assert result[0].message_pieces[1].converted_value_data_type == "image_path"
+
+
+async def test_generic_squash_uses_first_user_message_instead_of_rewriting_assistant():
+    """Test that squash targets the first user message even if assistant messages appear first."""
+    messages = [
+        _make_message("system", "System message"),
+        _make_message("assistant", "Assistant message"),
+        _make_message("user", "User message"),
+    ]
+
+    result = await GenericSystemSquashNormalizer().normalize_async(messages)
+
+    assert len(result) == 2
+    assert result[0].api_role == "assistant"
+    assert result[0].get_value() == "Assistant message"
+    assert result[1].api_role == "user"
+    assert result[1].get_value() == "### Instructions ###\n\nSystem message\n\n######\n\nUser message"
+
+
+async def test_generic_squash_no_user_message_converts_system_to_user():
+    """Test that system is converted to user when no user messages exist."""
+    messages = [
+        _make_message("system", "System message"),
+        _make_message("assistant", "Assistant message"),
+    ]
+
+    result = await GenericSystemSquashNormalizer().normalize_async(messages)
+
+    assert len(result) == 2
+    assert result[0].api_role == "user"
+    assert result[0].get_value() == "System message"
+    assert result[1].api_role == "assistant"
+    assert result[1].get_value() == "Assistant message"
+
+
+async def test_generic_squash_preserves_image_first_multipart_user_message():
+    """Test that squashing merges into the first text piece when an image piece comes first."""
+    conversation_id = "conv-image-first"
+    messages = [
+        _make_message("system", "System message"),
+        Message(
+            message_pieces=[
+                MessagePiece(
+                    role="user",
+                    original_value="/tmp/example.png",
+                    original_value_data_type="image_path",
+                    conversation_id=conversation_id,
+                    sequence=0,
+                ),
+                MessagePiece(
+                    role="user",
+                    original_value="Describe this image",
+                    conversation_id=conversation_id,
+                    sequence=0,
+                ),
+            ]
+        ),
+    ]
+
+    result = await GenericSystemSquashNormalizer().normalize_async(messages)
+
+    assert len(result) == 1
+    assert result[0].api_role == "user"
+    assert len(result[0].message_pieces) == 2
+    assert result[0].message_pieces[0].converted_value == "/tmp/example.png"
+    assert result[0].message_pieces[0].converted_value_data_type == "image_path"
+    assert result[0].message_pieces[1].converted_value_data_type == "text"
+    assert (
+        result[0].message_pieces[1].converted_value
+        == "### Instructions ###\n\nSystem message\n\n######\n\nDescribe this image"
+    )
+
+
+async def test_generic_squash_user_message_without_text_pieces_prepends_instructions():
+    """Test that an instruction-only text piece is prepended when no text piece exists to merge into."""
+    conversation_id = "conv-no-text"
+    messages = [
+        _make_message("system", "System message"),
+        Message(
+            message_pieces=[
+                MessagePiece(
+                    role="user",
+                    original_value="/tmp/example.png",
+                    original_value_data_type="image_path",
+                    conversation_id=conversation_id,
+                    sequence=0,
+                ),
+            ]
+        ),
+    ]
+
+    result = await GenericSystemSquashNormalizer().normalize_async(messages)
+
+    assert len(result) == 1
+    assert result[0].api_role == "user"
+    assert len(result[0].message_pieces) == 2
+    assert result[0].message_pieces[0].converted_value_data_type == "text"
+    assert result[0].message_pieces[0].converted_value == "### Instructions ###\n\nSystem message\n\n######"
+    assert result[0].message_pieces[1].converted_value == "/tmp/example.png"
+    assert result[0].message_pieces[1].converted_value_data_type == "image_path"
+
+
 async def test_generic_squash_propagates_user_piece_metadata():
     """
     Regression: when squashing system + user, the squashed piece must carry the