MAINT: Migrating AttackResult to Pydantic (#1899)

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

Author: rlundeen2

.gitignore

@@ -1,6 +1,6 @@
 # PyRIT-specific configs
 submodules/
-results/
+/results/
 dbdata/
 eval/
 default_memory.json.memory

pyrit/auxiliary_attacks/gcg/generator.py

@@ -42,6 +42,7 @@
 
 import numpy as np
 import torch.multiprocessing as mp
+from pydantic import Field
 
 import pyrit.auxiliary_attacks.gcg.attack.gcg.gcg_attack as attack_lib
 from pyrit.auxiliary_attacks.gcg.attack.base.attack_manager import (
@@ -96,7 +97,6 @@ class GCGContext(PromptGeneratorStrategyContext):
     logfile_path: Optional[str] = None
 
 
-@dataclass
 class GCGResult(PromptGeneratorStrategyResult):
     """Result of one GCGGenerator run.
 
@@ -117,10 +117,10 @@ class GCGResult(PromptGeneratorStrategyResult):
     final_suffix: str = ""
     final_loss: float = float("nan")
     step_count: int = 0
-    loss_history: list[float] = field(default_factory=list)
-    control_history: list[str] = field(default_factory=list)
-    log_path: Optional[str] = None
-    memory_labels: dict[str, str] = field(default_factory=dict)
+    loss_history: list[float] = Field(default_factory=list)
+    control_history: list[str] = Field(default_factory=list)
+    log_path: str | None = None
+    memory_labels: dict[str, str] = Field(default_factory=dict)
 
 
 class GCGGenerator(

pyrit/executor/attack/compound/sequential_attack.py

@@ -28,6 +28,8 @@
 from enum import Enum
 from typing import TYPE_CHECKING, Any, Optional
 
+from pydantic import Field
+
 from pyrit.executor.attack.core.attack_executor import AttackExecutor
 from pyrit.executor.attack.core.attack_parameters import AttackParameters
 from pyrit.executor.attack.core.attack_strategy import AttackContext, AttackStrategy
@@ -110,7 +112,6 @@ class SequentialChildAttack:
     memory_labels: Mapping[str, str] = field(default_factory=dict)
 
 
-@dataclass
 class SequentialAttackResult(AttackResult):
     """
     Result of a ``SequentialAttack`` execution.
@@ -138,7 +139,7 @@ class SequentialAttackResult(AttackResult):
             round-trip.
     """
 
-    child_attack_results: list[AttackResult] = field(default_factory=list)
+    child_attack_results: list[AttackResult] = Field(default_factory=list)
     completion_policy: SequenceCompletionPolicy = SequenceCompletionPolicy.FIRST_SUCCESS
 
     @property

pyrit/executor/attack/multi_turn/crescendo.py

@@ -81,7 +81,6 @@ class CrescendoAttackContext(MultiTurnAttackContext[Any]):
     backtrack_count: int = 0
 
 
-@dataclass
 class CrescendoAttackResult(AttackResult):
     """Result of the Crescendo attack strategy execution."""
 
@@ -832,7 +831,9 @@ async def _perform_backtrack_if_refused_async(
 
         # Check for refusal using the scorer (handles blocked/error responses internally)
         refusal_score = await self._check_refusal_async(context, prompt_sent)
-        self._logger.debug(f"Refusal check: {refusal_score.get_value()} - {refusal_score.score_rationale[:100]}...")
+        self._logger.debug(
+            f"Refusal check: {refusal_score.get_value()} - {(refusal_score.score_rationale or '')[:100]}..."
+        )
         is_refusal = bool(refusal_score.get_value())
 
         if not is_refusal:

pyrit/executor/attack/multi_turn/tree_of_attacks.py

@@ -176,7 +176,6 @@ class TAPAttackContext(MultiTurnAttackContext[Any]):
     best_adversarial_conversation_id: Optional[str] = None
 
 
-@dataclass
 class TAPAttackResult(AttackResult):
     """
     Result of the Tree of Attacks with Pruning (TAP) attack strategy execution.
@@ -699,7 +698,8 @@ async def _score_response_async(self, *, response: Message, objective: str) -> N
         # Extract auxiliary scores
         auxiliary_scores = scoring_results["auxiliary_scores"]
         for score in auxiliary_scores:
-            scorer_name = score.scorer_class_identifier.class_name
+            scorer_identifier = score.scorer_class_identifier
+            scorer_name = scorer_identifier.class_name if scorer_identifier else "unknown"
             self.auxiliary_scores[scorer_name] = score
             logger.debug(f"Node {self.node_id}: {scorer_name} score: {score.get_value()}")
 
@@ -904,7 +904,7 @@ async def _generate_red_teaming_prompt_async(self, objective: str) -> str:
             # Generate feedback prompt and get a new response
             feedback_prompt = self._generate_off_topic_feedback_prompt(
                 original_prompt=prompt,
-                off_topic_rationale=on_topic_score.score_rationale,
+                off_topic_rationale=on_topic_score.score_rationale or "",
                 objective=objective,
             )
 

pyrit/executor/promptgen/anecdoctor.py

@@ -57,7 +57,6 @@ class AnecdoctorContext(PromptGeneratorStrategyContext):
     memory_labels: dict[str, str] = field(default_factory=dict)
 
 
-@dataclass
 class AnecdoctorResult(PromptGeneratorStrategyResult):
     """
     Result of Anecdoctor prompt generation.

pyrit/executor/promptgen/core/prompt_generator_strategy.py

@@ -26,8 +26,7 @@ class PromptGeneratorStrategyContext(StrategyContext, ABC):
     """Base class for all prompt generator strategy contexts."""
 
 
-@dataclass
-class PromptGeneratorStrategyResult(StrategyResult, ABC):
+class PromptGeneratorStrategyResult(StrategyResult, ABC):  # noqa: B024
     """Base class for all prompt generator strategy results."""
 
 

pyrit/executor/promptgen/fuzzer/fuzzer.py

@@ -12,6 +12,7 @@
 
 import numpy as np
 from colorama import Fore, Style
+from pydantic import Field
 
 from pyrit.common.utils import combine_dict, get_kwarg_param
 from pyrit.exceptions import MissingPromptPlaceholderException, pyrit_placeholder_retry
@@ -212,7 +213,6 @@ def __post_init__(self) -> None:
             )
 
 
-@dataclass
 class FuzzerResult(PromptGeneratorStrategyResult):
     """
     Result of the Fuzzer prompt generation strategy execution.
@@ -222,8 +222,8 @@ class FuzzerResult(PromptGeneratorStrategyResult):
     """
 
     # Concrete fields instead of metadata storage
-    successful_templates: list[str] = field(default_factory=list)
-    jailbreak_conversation_ids: list[Union[str, uuid.UUID]] = field(default_factory=list)
+    successful_templates: list[str] = Field(default_factory=list)
+    jailbreak_conversation_ids: list[Union[str, uuid.UUID]] = Field(default_factory=list)
     total_queries: int = 0
     templates_explored: int = 0
 

pyrit/executor/workflow/core/workflow_strategy.py

@@ -27,8 +27,7 @@ class WorkflowContext(StrategyContext, ABC):
     """Base class for all workflow contexts."""
 
 
-@dataclass
-class WorkflowResult(StrategyResult, ABC):
+class WorkflowResult(StrategyResult, ABC):  # noqa: B024
     """Base class for all workflow results."""
 
 

pyrit/executor/workflow/xpia.py

@@ -81,7 +81,6 @@ class XPIAContext(WorkflowContext):
     memory_labels: dict[str, str] = field(default_factory=dict)
 
 
-@dataclass
 class XPIAResult(WorkflowResult):
     """
     Result of XPIA workflow execution.