MAINT Breaking: Convert ScenarioResult to Pydantic (#1908)

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

Author: 3 people

pyrit/memory/memory_models.py

@@ -401,14 +401,17 @@ def __init__(self, *, entry: Score) -> None:
         self.score_rationale = entry.score_rationale
         self.score_metadata = entry.score_metadata or {}
         normalized_scorer = entry.scorer_class_identifier
-        # Ensure eval_hash is set before truncation so it survives the DB round-trip
-        if normalized_scorer.eval_hash is None:
-            normalized_scorer = normalized_scorer.with_eval_hash(
-                ScorerEvaluationIdentifier(normalized_scorer).eval_hash
+        if normalized_scorer is None:
+            self.scorer_class_identifier = {}
+        else:
+            # Ensure eval_hash is set before truncation so it survives the DB round-trip
+            if normalized_scorer.eval_hash is None:
+                normalized_scorer = normalized_scorer.with_eval_hash(
+                    ScorerEvaluationIdentifier(normalized_scorer).eval_hash
+                )
+            self.scorer_class_identifier = normalized_scorer.model_dump(
+                context={"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
         # Store in both columns for backward compatibility
@@ -439,7 +442,7 @@ def get_score(self) -> Score:
             score_category=self.score_category,
             score_rationale=self.score_rationale,
             score_metadata=self.score_metadata,
-            scorer_class_identifier=scorer_identifier,  # type: ignore[ty:invalid-argument-type]
+            scorer_class_identifier=scorer_identifier,
             message_piece_id=self.prompt_request_response_id,
             timestamp=_ensure_utc(self.timestamp),
             objective=self.objective,
@@ -975,8 +978,8 @@ class ScenarioResultEntry(Base):
         scenario_init_data (dict): Optional initialization parameters used to configure the scenario.
         objective_target_identifier (dict): Identifier for the target being evaluated in the scenario.
         objective_scorer_identifier (dict): Optional identifier for the scorer used to evaluate results.
-        scenario_run_state (Literal["CREATED", "IN_PROGRESS", "COMPLETED", "FAILED"]): Current execution state
-            of the scenario.
+        scenario_run_state (str): Current execution state of the scenario
+            (one of CREATED, IN_PROGRESS, COMPLETED, FAILED, CANCELLED).
         attack_results_json (str): JSON-serialized dictionary mapping attack names to conversation IDs.
             Format: {"attack_name": ["conversation_id1", "conversation_id2", ...]}.
             The full AttackResult objects are stored in AttackResultEntries and can be queried by conversation_id.
@@ -1003,9 +1006,7 @@ class ScenarioResultEntry(Base):
     scenario_init_data: Mapped[dict[str, Any] | None] = mapped_column(JSON, nullable=True)
     objective_target_identifier: Mapped[dict[str, str]] = mapped_column(JSON, nullable=False)
     objective_scorer_identifier: Mapped[dict[str, str] | None] = mapped_column(JSON, nullable=True)
-    scenario_run_state: Mapped[Literal["CREATED", "IN_PROGRESS", "COMPLETED", "FAILED", "CANCELLED"]] = mapped_column(
-        String, nullable=False, default="CREATED"
-    )
+    scenario_run_state: Mapped[str] = mapped_column(String, nullable=False, default="CREATED")
     attack_results_json: Mapped[str] = mapped_column(Unicode, nullable=False)
     display_group_map_json: Mapped[str | None] = mapped_column(Unicode, nullable=True)
     labels: Mapped[dict[str, str] | None] = mapped_column(JSON, nullable=True)
@@ -1072,7 +1073,7 @@ def __init__(self, *, entry: ScenarioResult) -> None:
         self.attack_results_json = json.dumps(serialized_attack_results)
 
         # Serialize display_group_map if present
-        self.display_group_map_json = json.dumps(entry._display_group_map) if entry._display_group_map else None
+        self.display_group_map_json = json.dumps(entry.display_group_map) if entry.display_group_map else None
 
         self.error_message = entry.error_message
         self.error_type = entry.error_type
@@ -1126,14 +1127,14 @@ def get_scenario_result(self) -> ScenarioResult:
             attack_results=attack_results,
             objective_scorer_identifier=scorer_identifier,
             scenario_run_state=self.scenario_run_state,
-            labels=self.labels,
+            labels=self.labels or {},
             creation_time=self.timestamp,
             number_tries=self.number_tries,
             completion_time=self.completion_time,
-            display_group_map=display_group_map,
+            display_group_map=display_group_map or {},
             error_message=self.error_message,
             error_type=self.error_type,
-            metadata=dict(self.scenario_metadata) if self.scenario_metadata else None,
+            metadata=dict(self.scenario_metadata) if self.scenario_metadata else {},
         )
 
     def get_conversation_ids_by_attack_name(self) -> dict[str, list[str]]:

pyrit/models/identifiers/component_identifier.py

@@ -168,7 +168,7 @@ class ComponentIdentifier(BaseModel):
     #: may have been truncated.
     hash: Optional[str] = None
     #: Version tag for storage. Not included in the content hash.
-    pyrit_version: str = Field(default_factory=lambda: pyrit.__version__)
+    pyrit_version: str = Field(default=pyrit.__version__)
     #: Evaluation hash. Computed by EvaluationIdentifier subclasses and attached
     #: to the identifier so it survives DB round-trips with truncated params.
     eval_hash: Optional[str] = None