MAINT: Making Score a Pydantic model (#1891)

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

Author: rlundeen2

pyrit/memory/sqlite_memory.py

@@ -559,7 +559,7 @@ def export_conversations(
             piece_data = piece.model_dump(mode="json")
             # Find associated scores
             piece_scores = [score for score in scores if score.message_piece_id == piece.id]
-            piece_data["scores"] = [score.to_dict() for score in piece_scores]
+            piece_data["scores"] = [score.model_dump(mode="json") for score in piece_scores]
             merged_data.append(piece_data)
 
         if not merged_data:

pyrit/models/attack_result.py

@@ -243,7 +243,7 @@ def to_dict(self) -> dict[str, Any]:
                 self.atomic_attack_identifier.model_dump() if self.atomic_attack_identifier else None
             ),
             "last_response": self.last_response.model_dump(mode="json") if self.last_response else None,
-            "last_score": self.last_score.to_dict() if self.last_score else None,
+            "last_score": self.last_score.model_dump(mode="json") if self.last_score else None,
             "executed_turns": self.executed_turns,
             "execution_time_ms": self.execution_time_ms,
             "outcome": self.outcome.value,
@@ -283,7 +283,7 @@ def from_dict(cls, data: dict[str, Any]) -> AttackResult:
                 else None
             ),
             last_response=(MessagePiece.model_validate(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,
+            last_score=Score.model_validate(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")),

pyrit/models/messages/message_piece.py

@@ -5,28 +5,28 @@
 
 import uuid
 from datetime import datetime, timezone
-from typing import TYPE_CHECKING, Annotated, Any, Literal, Optional
+from typing import TYPE_CHECKING, Any, Literal, Optional
 from uuid import uuid4
 
 from pydantic import (
     AwareDatetime,
     BaseModel,
-    BeforeValidator,
     ConfigDict,
     Field,
-    PlainSerializer,
     model_validator,
 )
 
 from pyrit.common.deprecation import print_deprecation_message
 from pyrit.models.data_type_serializer import data_serializer_factory
-from pyrit.models.identifiers.component_identifier import ComponentIdentifier
 from pyrit.models.literals import (  # noqa: TC001  (runtime-required by Pydantic field annotations)
     ChatMessageRole,
     PromptDataType,
     PromptResponseError,
 )
-from pyrit.models.score import Score
+from pyrit.models.score import (  # noqa: TC001  (runtime-required by Pydantic field annotations)
+    ComponentIdentifierField,
+    Score,
+)
 
 if TYPE_CHECKING:
     from pyrit.models.messages.message import Message
@@ -48,21 +48,9 @@
 )
 
 
-# Annotated alias that round-trips identifier fields through the flat dict
-# storage shape. ``ComponentIdentifier`` is a Pydantic model with a custom
-# flat serializer; ``Score`` is still a plain class needing ``from_dict`` /
-# ``to_dict``. Drop the ``Score`` alias once it becomes a Pydantic model.
-ComponentIdentifierField = Annotated[
-    ComponentIdentifier,
-    BeforeValidator(lambda v: ComponentIdentifier.model_validate(v) if isinstance(v, dict) else v),
-    PlainSerializer(lambda v: v.model_dump() if v is not None else None, return_type=Optional[dict]),
-]
-
-ScoreField = Annotated[
-    Score,
-    BeforeValidator(lambda v: Score.from_dict(v) if isinstance(v, dict) else v),
-    PlainSerializer(lambda v: v.to_dict(), return_type=dict),
-]
+# ``ComponentIdentifierField`` (and ``Score``) are imported from ``pyrit.models.score``
+# above. Both round-trip through the flat dict storage shape via their own Pydantic
+# serializers, so no local annotated aliases are needed here.
 
 
 def __getattr__(name: str) -> Any:
@@ -128,7 +116,7 @@ class MessagePiece(BaseModel):
     prompt_target_identifier: Optional[ComponentIdentifierField] = None
     attack_identifier: Optional[ComponentIdentifierField] = None
     scorer_identifier: Optional[ComponentIdentifierField] = None
-    scores: list[ScoreField] = Field(default_factory=list)
+    scores: list[Score] = Field(default_factory=list)
 
     # When True, the memory layer skips persisting this piece. Used for ephemeral
     # pieces a scorer creates to score arbitrary content; ``exclude=True`` keeps