MAINT: Refactoring Identifiers to be Pydantic classes (#1881)

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

Author: rlundeen2

pyrit/backend/services/attack_service.py

@@ -756,7 +756,7 @@ async def _update_attack_after_message_async(
                     children=new_children,
                 )
                 if ar.atomic_attack_identifier:
-                    atomic = ComponentIdentifier.from_dict(ar.atomic_attack_identifier.to_dict())
+                    atomic = ComponentIdentifier.model_validate(ar.atomic_attack_identifier.model_dump())
                     atomic_children = dict(atomic.children)
                     # Navigate into attack_technique child to update the nested attack child.
                     technique = atomic_children.get("attack_technique")
@@ -778,7 +778,7 @@ async def _update_attack_after_message_async(
                         params=dict(atomic.params),
                         children=atomic_children,
                     )
-                    update_fields["atomic_attack_identifier"] = new_atomic.to_dict()
+                    update_fields["atomic_attack_identifier"] = new_atomic.model_dump()
 
         self._memory.update_attack_result_by_id(
             attack_result_id=attack_result_id,

pyrit/memory/memory_models.py

@@ -57,7 +57,7 @@
 # Default pyrit_version for database records created before version tracking was added
 LEGACY_PYRIT_VERSION = "<0.10.0"
 
-# Maximum length for string values in ComponentIdentifier.to_dict() when storing to the database.
+# Maximum length for string values in ComponentIdentifier.model_dump() when storing to the database.
 # Longer values are truncated with a "..." suffix.
 MAX_IDENTIFIER_VALUE_LENGTH: int = 80
 
@@ -231,16 +231,17 @@ def __init__(self, *, entry: MessagePiece) -> None:
         self.prompt_metadata = entry.prompt_metadata
         self.targeted_harm_categories = entry.targeted_harm_categories
         self.converter_identifiers = [
-            conv.to_dict(max_value_length=MAX_IDENTIFIER_VALUE_LENGTH) for conv in entry.converter_identifiers
+            conv.model_dump(context={"max_value_length": MAX_IDENTIFIER_VALUE_LENGTH})
+            for conv in entry.converter_identifiers
         ]
         # Normalize prompt_target_identifier and convert to dict for JSON serialization
         self.prompt_target_identifier = (
-            entry.prompt_target_identifier.to_dict(max_value_length=MAX_IDENTIFIER_VALUE_LENGTH)
+            entry.prompt_target_identifier.model_dump(context={"max_value_length": MAX_IDENTIFIER_VALUE_LENGTH})
             if entry.prompt_target_identifier
             else {}
         )
         self.attack_identifier = (
-            entry.attack_identifier.to_dict(max_value_length=MAX_IDENTIFIER_VALUE_LENGTH)
+            entry.attack_identifier.model_dump(context={"max_value_length": MAX_IDENTIFIER_VALUE_LENGTH})
             if entry.attack_identifier
             else {}
         )
@@ -270,21 +271,21 @@ def get_message_piece(self) -> MessagePiece:
         stored_version = self.pyrit_version or LEGACY_PYRIT_VERSION
         if self.converter_identifiers:
             converter_ids = [
-                ComponentIdentifier.from_dict({**c, "pyrit_version": stored_version})
+                ComponentIdentifier.model_validate({**c, "pyrit_version": stored_version})
                 for c in self.converter_identifiers
             ]
 
         # Reconstruct ComponentIdentifier with the stored pyrit_version
         target_id: ComponentIdentifier | None = None
         if self.prompt_target_identifier:
-            target_id = ComponentIdentifier.from_dict(
+            target_id = ComponentIdentifier.model_validate(
                 {**self.prompt_target_identifier, "pyrit_version": stored_version}
             )
 
         # Reconstruct ComponentIdentifier with the stored pyrit_version
         attack_id: ComponentIdentifier | None = None
         if self.attack_identifier:
-            attack_id = ComponentIdentifier.from_dict({**self.attack_identifier, "pyrit_version": stored_version})
+            attack_id = ComponentIdentifier.model_validate({**self.attack_identifier, "pyrit_version": stored_version})
 
         message_piece = MessagePiece(
             role=self.role,
@@ -405,8 +406,8 @@ def __init__(self, *, entry: Score) -> None:
             normalized_scorer = normalized_scorer.with_eval_hash(
                 ScorerEvaluationIdentifier(normalized_scorer).eval_hash
             )
-        self.scorer_class_identifier = normalized_scorer.to_dict(
-            max_value_length=MAX_IDENTIFIER_VALUE_LENGTH,
+        self.scorer_class_identifier = normalized_scorer.model_dump(
+            context={"max_value_length": MAX_IDENTIFIER_VALUE_LENGTH},
         )
         self.prompt_request_response_id = entry.message_piece_id if entry.message_piece_id else None
         self.timestamp = entry.timestamp
@@ -427,7 +428,7 @@ def get_score(self) -> Score:
         scorer_identifier = None
         stored_version = self.pyrit_version or LEGACY_PYRIT_VERSION
         if self.scorer_class_identifier:
-            scorer_identifier = ComponentIdentifier.from_dict(
+            scorer_identifier = ComponentIdentifier.model_validate(
                 {**self.scorer_class_identifier, "pyrit_version": stored_version}
             )
         return Score(
@@ -779,16 +780,18 @@ def __init__(self, *, entry: AttackResult) -> None:
         # Will be removed in 0.15.0.
         _attack_strategy_id = entry.get_attack_strategy_identifier()
         self.attack_identifier = (
-            _attack_strategy_id.to_dict(max_value_length=MAX_IDENTIFIER_VALUE_LENGTH) if _attack_strategy_id else {}
+            _attack_strategy_id.model_dump(context={"max_value_length": MAX_IDENTIFIER_VALUE_LENGTH})
+            if _attack_strategy_id
+            else {}
         )
         # Ensure eval_hash is set before truncation so it survives the DB round-trip
         if entry.atomic_attack_identifier and entry.atomic_attack_identifier.eval_hash is None:
             entry.atomic_attack_identifier = entry.atomic_attack_identifier.with_eval_hash(
                 AtomicAttackEvaluationIdentifier(entry.atomic_attack_identifier).eval_hash
             )
         self.atomic_attack_identifier = (
-            entry.atomic_attack_identifier.to_dict(
-                max_value_length=MAX_IDENTIFIER_VALUE_LENGTH,
+            entry.atomic_attack_identifier.model_dump(
+                context={"max_value_length": MAX_IDENTIFIER_VALUE_LENGTH},
             )
             if entry.atomic_attack_identifier
             else None
@@ -911,13 +914,13 @@ def get_attack_result(self) -> AttackResult:
         # Reconstruct atomic_attack_identifier, with backward compatibility for
         # legacy rows that only have the attack_identifier column.
         atomic_id = (
-            ComponentIdentifier.from_dict(self.atomic_attack_identifier) if self.atomic_attack_identifier else None
+            ComponentIdentifier.model_validate(self.atomic_attack_identifier) if self.atomic_attack_identifier else None
         )
         if atomic_id is None and self.attack_identifier:
             from pyrit.models import build_atomic_attack_identifier
 
             atomic_id = build_atomic_attack_identifier(
-                attack_identifier=ComponentIdentifier.from_dict(self.attack_identifier),
+                attack_identifier=ComponentIdentifier.model_validate(self.attack_identifier),
             )
 
         # Deserialize retry events from JSON
@@ -1037,8 +1040,8 @@ def __init__(self, *, entry: ScenarioResult) -> None:
         self.scenario_init_data = entry.scenario_identifier.init_data
         # Convert ComponentIdentifier to dict for JSON storage
         self.objective_target_identifier = (
-            entry.objective_target_identifier.to_dict(
-                max_value_length=MAX_IDENTIFIER_VALUE_LENGTH,
+            entry.objective_target_identifier.model_dump(
+                context={"max_value_length": MAX_IDENTIFIER_VALUE_LENGTH},
             )
             if entry.objective_target_identifier
             else None
@@ -1050,8 +1053,8 @@ def __init__(self, *, entry: ScenarioResult) -> None:
             )
 
         self.objective_scorer_identifier = (
-            entry.objective_scorer_identifier.to_dict(
-                max_value_length=MAX_IDENTIFIER_VALUE_LENGTH,
+            entry.objective_scorer_identifier.model_dump(
+                context={"max_value_length": MAX_IDENTIFIER_VALUE_LENGTH},
             )
             if entry.objective_scorer_identifier
             else None
@@ -1104,12 +1107,12 @@ def get_scenario_result(self) -> ScenarioResult:
         # Convert dict back to ComponentIdentifier with the stored pyrit_version
         scorer_identifier = None
         if self.objective_scorer_identifier:
-            scorer_identifier = ComponentIdentifier.from_dict(
+            scorer_identifier = ComponentIdentifier.model_validate(
                 {**self.objective_scorer_identifier, "pyrit_version": stored_version}
             )
 
         # Convert dict back to ComponentIdentifier for reconstruction
-        target_identifier = ComponentIdentifier.from_dict(self.objective_target_identifier)
+        target_identifier = ComponentIdentifier.model_validate(self.objective_target_identifier)
 
         # Deserialize display_group_map if stored
         display_group_map: dict[str, str] | None = None

pyrit/models/attack_result.py

@@ -240,7 +240,7 @@ def to_dict(self) -> dict[str, Any]:
             "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
+                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,
@@ -278,7 +278,7 @@ def from_dict(cls, data: dict[str, Any]) -> AttackResult:
             objective=data["objective"],
             attack_result_id=data.get("attack_result_id", str(uuid.uuid4())),
             atomic_attack_identifier=(
-                ComponentIdentifier.from_dict(data["atomic_attack_identifier"])
+                ComponentIdentifier.model_validate(data["atomic_attack_identifier"])
                 if data.get("atomic_attack_identifier")
                 else None
             ),