MAINT: tighten model from_dict edge cases (labels, pyrit_version, completion_time)

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

Author: rlundeen2

pyrit/models/attack_result.py

@@ -240,7 +240,7 @@ def to_dict(self) -> dict[str, Any]:
             "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,
+            "timestamp": self.timestamp.isoformat(),
             "related_conversations": sorted(
                 [ref.to_dict() for ref in self.related_conversations],
                 key=lambda r: r["conversation_id"],

pyrit/models/message_piece.py

@@ -370,7 +370,7 @@ def from_dict(cls, data: dict[str, Any]) -> MessagePiece:
             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"),
+            labels=data.get("labels") or None,
             targeted_harm_categories=data.get("targeted_harm_categories"),
             prompt_metadata=data.get("prompt_metadata"),
             converter_identifiers=(

pyrit/models/scenario_result.py

@@ -79,7 +79,7 @@ def from_dict(cls, data: dict[str, Any]) -> ScenarioIdentifier:
             description=data.get("description", ""),
             scenario_version=data.get("scenario_version", 1),
             init_data=data.get("init_data"),
-            pyrit_version=data.get("pyrit_version"),
+            pyrit_version=data.get("pyrit_version") or "unknown",
         )
 
 
@@ -339,7 +339,7 @@ def from_dict(cls, data: dict[str, Any]) -> ScenarioResult:
         """
         from pyrit.identifiers.component_identifier import ComponentIdentifier
 
-        return cls(
+        result = cls(
             id=uuid.UUID(data["id"]) if data.get("id") else None,
             scenario_identifier=ScenarioIdentifier.from_dict(data["scenario_identifier"]),
             objective_target_identifier=(
@@ -366,3 +366,8 @@ def from_dict(cls, data: dict[str, Any]) -> ScenarioResult:
             error_message=data.get("error_message"),
             error_type=data.get("error_type"),
         )
+        # Preserve missing completion_time: __init__ defaults it to now(), but a
+        # still-running scenario shouldn't be marked as completed-at-load-time.
+        if not data.get("completion_time"):
+            result.completion_time = None  # type: ignore[ty:invalid-assignment]
+        return result

tests/unit/models/test_message_piece.py

@@ -1244,3 +1244,13 @@ def test_to_dict_from_dict_roundtrip_after_set_piece_not_in_database():
     roundtripped = MessagePiece.from_dict(serialized)
     assert isinstance(roundtripped.id, uuid.UUID)
     assert isinstance(roundtripped.original_prompt_id, uuid.UUID)
+
+
+def test_from_dict_does_not_emit_deprecation_warning_for_default_labels():
+    """from_dict on a piece with no labels should not trigger the labels deprecation warning."""
+    piece = MessagePiece(role="user", original_value="hello", conversation_id="conv-1")
+    serialized = piece.to_dict()
+
+    with warnings.catch_warnings():
+        warnings.simplefilter("error", DeprecationWarning)
+        MessagePiece.from_dict(serialized)

tests/unit/models/test_scenario_result.py

@@ -268,3 +268,35 @@ def test_scenario_result_to_dict_from_dict_roundtrip():
     )
     roundtripped = ScenarioResult.from_dict(original.to_dict())
     assert original.to_dict() == roundtripped.to_dict()
+
+
+def test_scenario_identifier_from_dict_missing_pyrit_version_yields_unknown():
+    """from_dict should preserve 'unknown' rather than fabricating the current version when missing."""
+    data = {
+        "name": "Legacy",
+        "description": "loaded from older payload",
+        "scenario_version": 1,
+        "init_data": None,
+        # pyrit_version intentionally absent
+    }
+    identifier = ScenarioIdentifier.from_dict(data)
+    assert identifier.pyrit_version == "unknown"
+
+
+def test_scenario_result_from_dict_preserves_missing_completion_time():
+    """An in-progress scenario serialized without completion_time should round-trip with completion_time=None."""
+    scenario_id = ScenarioIdentifier(name="Test", scenario_version=1, pyrit_version="0.14.0")
+    target_id = ComponentIdentifier(class_name="OpenAIChatTarget", class_module="pyrit.prompt_target")
+
+    original = ScenarioResult(
+        scenario_identifier=scenario_id,
+        objective_target_identifier=target_id,
+        objective_scorer_identifier=None,
+        attack_results={},
+        scenario_run_state="IN_PROGRESS",
+    )
+    original.completion_time = None  # type: ignore[ty:invalid-assignment]
+
+    roundtripped = ScenarioResult.from_dict(original.to_dict())
+    assert roundtripped.completion_time is None
+    assert roundtripped.scenario_run_state == "IN_PROGRESS"