MAINT: preserve Message.to_dict back-compat fields

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: rlundeen2

pyrit/models/message.py

@@ -285,18 +285,30 @@ def __str__(self) -> str:
 
     def to_dict(self) -> dict[str, object]:
         """
-        Convert the message to a dictionary representation including all piece details.
+        Convert the message to a dictionary representation.
 
-        Serializes each piece individually via MessagePiece.to_dict(). All message-level
-        attributes (role, conversation_id, sequence, is_simulated) are derived from the
-        pieces themselves, so only 'pieces' is included. This is the format expected by
-        from_dict().
+        Includes the original top-level fields ('role', 'converted_value', 'conversation_id',
+        'sequence', 'converted_value_data_type') for backward compatibility, plus a 'pieces'
+        list containing each MessagePiece.to_dict() — the latter is the source of truth used
+        by from_dict().
 
         Returns:
-            dict[str, object]: Dictionary with a single 'pieces' key containing a list
-                of MessagePiece.to_dict() dicts.
+            dict[str, object]: Dictionary with 'role', 'converted_value', 'conversation_id',
+                'sequence', 'converted_value_data_type', and 'pieces' keys.
         """
+        if len(self.message_pieces) == 1:
+            converted_value: str | list[str] = self.message_pieces[0].converted_value
+            converted_value_data_type: str | list[str] = self.message_pieces[0].converted_value_data_type
+        else:
+            converted_value = [piece.converted_value for piece in self.message_pieces]
+            converted_value_data_type = [piece.converted_value_data_type for piece in self.message_pieces]
+
         return {
+            "role": self.api_role,
+            "converted_value": converted_value,
+            "conversation_id": self.conversation_id,
+            "sequence": self.sequence,
+            "converted_value_data_type": converted_value_data_type,
             "pieces": [piece.to_dict() for piece in self.message_pieces],
         }
 

tests/unit/message_normalizer/test_generic_system_squash_normalizer.py

@@ -61,9 +61,9 @@ async def test_generic_squash_normalize_to_dicts_async():
     assert isinstance(result, list)
     assert len(result) == 1
     assert isinstance(result[0], dict)
-    assert set(result[0].keys()) == {"pieces"}
+    assert result[0]["role"] == "user"
+    assert "pieces" in result[0]
     assert len(result[0]["pieces"]) == 1
-    assert result[0]["pieces"][0]["role"] == "user"
     assert "### Instructions ###" in result[0]["pieces"][0]["converted_value"]
     assert "System message" in result[0]["pieces"][0]["converted_value"]
     assert "User message" in result[0]["pieces"][0]["converted_value"]

tests/unit/models/test_message.py

@@ -226,9 +226,12 @@ def test_message_to_dict() -> None:
     message = Message.from_prompt(prompt="Hello world", role="user")
     result = message.to_dict()
 
-    assert set(result.keys()) == {"pieces"}
+    assert result["role"] == "user"
+    assert result["converted_value"] == "Hello world"
+    assert result["converted_value_data_type"] == "text"
+    assert "conversation_id" in result
+    assert "sequence" in result
     assert len(result["pieces"]) == 1
-    assert result["pieces"][0]["role"] == "user"
     assert result["pieces"][0]["converted_value"] == "Hello world"
     assert result["pieces"][0]["converted_value_data_type"] == "text"