MAINT: address PR feedback on model to_dict/from_dict

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

Author: rlundeen2

pyrit/models/attack_result.py

@@ -242,11 +242,8 @@ def to_dict(self) -> dict[str, Any]:
             "outcome_reason": self.outcome_reason,
             "timestamp": self.timestamp.isoformat() if self.timestamp else None,
             "related_conversations": sorted(
-                [
-                    ref.to_dict() if isinstance(ref, ConversationReference) else ref
-                    for ref in self.related_conversations
-                ],
-                key=lambda r: r["conversation_id"] if isinstance(r, dict) else "",
+                [ref.to_dict() for ref in self.related_conversations],
+                key=lambda r: r["conversation_id"],
             ),
             "metadata": self.metadata,
             "labels": self.labels,

pyrit/models/message.py

@@ -287,18 +287,16 @@ def to_dict(self) -> dict[str, object]:
         """
         Convert the message to a dictionary representation including all piece details.
 
-        Serializes each piece individually via MessagePiece.to_dict(). This is the format
-        expected by from_dict().
+        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().
 
         Returns:
-            dict[str, object]: Dictionary with 'role', 'is_simulated', 'conversation_id',
-                'sequence', and 'pieces' (list of MessagePiece.to_dict() dicts).
+            dict[str, object]: Dictionary with a single 'pieces' key containing a list
+                of MessagePiece.to_dict() dicts.
         """
         return {
-            "role": self.api_role,
-            "is_simulated": self.is_simulated,
-            "conversation_id": self.conversation_id,
-            "sequence": self.sequence,
             "pieces": [piece.to_dict() for piece in self.message_pieces],
         }
 

pyrit/models/message_piece.py

@@ -314,7 +314,7 @@ def to_dict(self) -> dict[str, object]:
 
         """
         return {
-            "id": str(self.id),
+            "id": str(self.id) if self.id is not None else None,
             "role": self._role,
             "conversation_id": self.conversation_id,
             "sequence": self.sequence,
@@ -336,7 +336,7 @@ def to_dict(self) -> dict[str, object]:
             "converted_value_sha256": self.converted_value_sha256,
             "response_error": self.response_error,
             "originator": self.originator,
-            "original_prompt_id": str(self.original_prompt_id),
+            "original_prompt_id": str(self.original_prompt_id) if self.original_prompt_id is not None else None,
             "scores": [score.to_dict() for score in self.scores],
         }
 

pyrit/models/score.py

@@ -210,10 +210,10 @@ def from_dict(cls, data: dict[str, Any]) -> Score:
         return cls(
             id=data.get("id"),
             score_value=data["score_value"],
-            score_value_description=data.get("score_value_description", ""),
+            score_value_description=data["score_value_description"],
             score_type=data["score_type"],
             score_category=data.get("score_category"),
-            score_rationale=data.get("score_rationale", ""),
+            score_rationale=data["score_rationale"],
             score_metadata=data.get("score_metadata"),
             scorer_class_identifier=ComponentIdentifier.from_dict(data["scorer_class_identifier"]),
             message_piece_id=data["message_piece_id"],

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 result[0]["role"] == "user"
-    assert "pieces" in result[0]
+    assert set(result[0].keys()) == {"pieces"}
     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,11 +226,9 @@ def test_message_to_dict() -> None:
     message = Message.from_prompt(prompt="Hello world", role="user")
     result = message.to_dict()
 
-    assert result["role"] == "user"
-    assert result["is_simulated"] is False
-    assert "conversation_id" in result
-    assert "sequence" in result
+    assert set(result.keys()) == {"pieces"}
     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"
 

tests/unit/models/test_message_piece.py

@@ -1224,3 +1224,23 @@ def test_to_dict_from_dict_roundtrip():
     )
     roundtripped = MessagePiece.from_dict(original.to_dict())
     assert original.to_dict() == roundtripped.to_dict()
+
+
+def test_to_dict_from_dict_roundtrip_after_set_piece_not_in_database():
+    """Pieces marked not-in-database (id=None) must serialize and deserialize cleanly without ValueError."""
+    piece = MessagePiece(
+        role="user",
+        original_value="Hello world",
+        conversation_id="conv-not-in-db",
+    )
+    piece.set_piece_not_in_database()
+    piece.original_prompt_id = None  # type: ignore[assignment]
+
+    serialized = piece.to_dict()
+    assert serialized["id"] is None
+    assert serialized["original_prompt_id"] is None
+
+    # Must not raise ValueError on the literal string "None" or similar corruption.
+    roundtripped = MessagePiece.from_dict(serialized)
+    assert isinstance(roundtripped.id, uuid.UUID)
+    assert isinstance(roundtripped.original_prompt_id, uuid.UUID)