Add to_dict/from_dict roundtrip serialization to model classes

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

Author: rlundeen2

pyrit/models/attack_result.py

@@ -224,6 +224,90 @@ def __str__(self) -> str:
         """
         return f"AttackResult: {self.conversation_id}: {self.outcome.value}: {self.objective[:50]}..."
 
+    def to_dict(self) -> dict[str, Any]:
+        """
+        Serialize this attack result to a JSON-compatible dictionary.
+
+        Returns:
+            dict[str, Any]: Serialized payload suitable for REST APIs or persistence.
+        """
+        from pyrit.models.conversation_reference import ConversationReference
+
+        return {
+            "conversation_id": self.conversation_id,
+            "objective": self.objective,
+            "attack_result_id": self.attack_result_id,
+            "atomic_attack_identifier": (
+                self.atomic_attack_identifier.to_dict() if self.atomic_attack_identifier else None
+            ),
+            "last_response": self.last_response.to_dict() if self.last_response else None,
+            "last_score": self.last_score.to_dict() if self.last_score else None,
+            "executed_turns": self.executed_turns,
+            "execution_time_ms": self.execution_time_ms,
+            "outcome": self.outcome.value,
+            "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 "",
+            ),
+            "metadata": self.metadata,
+            "labels": self.labels,
+            "error_message": self.error_message,
+            "error_type": self.error_type,
+            "error_traceback": self.error_traceback,
+            "retry_events": [e.to_dict() for e in self.retry_events],
+            "total_retries": self.total_retries,
+        }
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> AttackResult:
+        """
+        Reconstruct an AttackResult from a dictionary.
+
+        Args:
+            data (dict[str, Any]): Dictionary as produced by to_dict().
+
+        Returns:
+            AttackResult: Reconstructed instance.
+        """
+        from pyrit.identifiers.component_identifier import ComponentIdentifier
+        from pyrit.models.conversation_reference import ConversationReference
+        from pyrit.models.message_piece import MessagePiece
+        from pyrit.models.retry_event import RetryEvent
+        from pyrit.models.score import Score
+
+        return cls(
+            conversation_id=data["conversation_id"],
+            objective=data["objective"],
+            attack_result_id=data.get("attack_result_id", str(uuid.uuid4())),
+            atomic_attack_identifier=(
+                ComponentIdentifier.from_dict(data["atomic_attack_identifier"])
+                if data.get("atomic_attack_identifier")
+                else None
+            ),
+            last_response=(MessagePiece.from_dict(data["last_response"]) if data.get("last_response") else None),
+            last_score=Score.from_dict(data["last_score"]) if data.get("last_score") else None,
+            executed_turns=data.get("executed_turns", 0),
+            execution_time_ms=data.get("execution_time_ms", 0),
+            outcome=AttackOutcome(data.get("outcome", "undetermined")),
+            outcome_reason=data.get("outcome_reason"),
+            timestamp=(
+                datetime.fromisoformat(data["timestamp"]) if data.get("timestamp") else datetime.now(timezone.utc)
+            ),
+            related_conversations={ConversationReference.from_dict(r) for r in data.get("related_conversations", [])},
+            metadata=data.get("metadata", {}),
+            labels=data.get("labels", {}),
+            error_message=data.get("error_message"),
+            error_type=data.get("error_type"),
+            error_traceback=data.get("error_traceback"),
+            retry_events=[RetryEvent.from_dict(e) for e in data.get("retry_events", [])],
+            total_retries=data.get("total_retries", 0),
+        )
+
 
 def _add_attack_identifier_compat(cls: type) -> type:
     """

pyrit/models/conversation_reference.py

@@ -36,6 +36,36 @@ def __hash__(self) -> int:
         """
         return hash(self.conversation_id)
 
+    def to_dict(self) -> dict[str, str | None]:
+        """
+        Serialize to a JSON-compatible dictionary.
+
+        Returns:
+            dict[str, str | None]: Dictionary with conversation_id, conversation_type, and description.
+        """
+        return {
+            "conversation_id": self.conversation_id,
+            "conversation_type": self.conversation_type.value,
+            "description": self.description,
+        }
+
+    @classmethod
+    def from_dict(cls, data: dict[str, str | None]) -> ConversationReference:
+        """
+        Reconstruct a ConversationReference from a dictionary.
+
+        Args:
+            data (dict[str, str | None]): Dictionary as produced by to_dict().
+
+        Returns:
+            ConversationReference: Reconstructed instance.
+        """
+        return cls(
+            conversation_id=str(data["conversation_id"]),
+            conversation_type=ConversationType(data["conversation_type"]),
+            description=data.get("description"),
+        )
+
     def __eq__(self, other: object) -> bool:
         """
         Compare two references by conversation ID.

pyrit/models/message.py

@@ -6,7 +6,7 @@
 import copy
 import uuid
 from datetime import datetime, timezone
-from typing import TYPE_CHECKING, Optional, Union
+from typing import TYPE_CHECKING, Any, Optional, Union
 
 from pyrit.common.utils import combine_dict
 from pyrit.models.message_piece import MessagePiece
@@ -285,28 +285,41 @@ def __str__(self) -> str:
 
     def to_dict(self) -> dict[str, object]:
         """
-        Convert the message to a dictionary representation.
+        Convert the message to a dictionary representation including all piece details.
 
-        Returns:
-            dict: A dictionary with 'role', 'converted_value', 'conversation_id', 'sequence',
-                and 'converted_value_data_type' keys.
+        Serializes each piece individually via MessagePiece.to_dict(). 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).
         """
-        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,
+            "is_simulated": self.is_simulated,
             "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],
         }
 
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> Message:
+        """
+        Reconstruct a Message from a dictionary.
+
+        Expects the format produced by to_dict(), which includes a 'pieces' key
+        containing a list of MessagePiece dictionaries.
+
+        Args:
+            data (dict[str, Any]): Dictionary as produced by to_dict().
+
+        Returns:
+            Message: Reconstructed instance.
+        """
+        pieces_data = data.get("pieces", [])
+        message_pieces = [MessagePiece.from_dict(p) for p in pieces_data]
+        return cls(message_pieces, skip_validation=True)
+
     @staticmethod
     def get_all_values(messages: Sequence[Message]) -> list[str]:
         """

pyrit/models/message_piece.py

@@ -5,7 +5,7 @@
 
 import uuid
 from datetime import datetime, timezone
-from typing import TYPE_CHECKING, Literal, Optional, Union, get_args
+from typing import TYPE_CHECKING, Any, Literal, Optional, Union, get_args
 from uuid import uuid4
 
 from pyrit.common.deprecation import print_deprecation_message
@@ -354,6 +354,57 @@ def __str__(self) -> str:
 
     __repr__ = __str__
 
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> MessagePiece:
+        """
+        Reconstruct a MessagePiece from a dictionary.
+
+        Args:
+            data (dict[str, Any]): Dictionary as produced by to_dict().
+
+        Returns:
+            MessagePiece: Reconstructed instance.
+        """
+        from pyrit.identifiers.component_identifier import ComponentIdentifier
+        from pyrit.models.score import Score
+
+        return cls(
+            id=data.get("id"),
+            role=data.get("role", "user"),
+            conversation_id=data.get("conversation_id"),
+            sequence=data.get("sequence", -1),
+            timestamp=(datetime.fromisoformat(str(data["timestamp"])) if data.get("timestamp") else None),
+            labels=data.get("labels"),
+            targeted_harm_categories=data.get("targeted_harm_categories"),
+            prompt_metadata=data.get("prompt_metadata"),
+            converter_identifiers=(
+                [ComponentIdentifier.from_dict(c) for c in data["converter_identifiers"]]
+                if data.get("converter_identifiers")
+                else None
+            ),
+            prompt_target_identifier=(
+                ComponentIdentifier.from_dict(data["prompt_target_identifier"])
+                if data.get("prompt_target_identifier")
+                else None
+            ),
+            attack_identifier=(
+                ComponentIdentifier.from_dict(data["attack_identifier"]) if data.get("attack_identifier") else None
+            ),
+            scorer_identifier=(
+                ComponentIdentifier.from_dict(data["scorer_identifier"]) if data.get("scorer_identifier") else None
+            ),
+            original_value_data_type=data.get("original_value_data_type", "text"),
+            original_value=data.get("original_value", ""),
+            original_value_sha256=data.get("original_value_sha256"),
+            converted_value_data_type=data.get("converted_value_data_type"),
+            converted_value=data.get("converted_value"),
+            converted_value_sha256=data.get("converted_value_sha256"),
+            response_error=data.get("response_error", "none"),
+            originator=data.get("originator", "undefined"),
+            original_prompt_id=(uuid.UUID(str(data["original_prompt_id"])) if data.get("original_prompt_id") else None),
+            scores=([Score.from_dict(s) for s in data["scores"]] if data.get("scores") else None),
+        )
+
     def __eq__(self, other: object) -> bool:
         """
         Compare this message piece with another for semantic equality.