MAINT: Add labels to attack results (#1624)

Co-authored-by: Behnam Ousat <behnamousat@microsoft.com>
Co-authored-by: hannahwestra25 <hannahwestra@microsoft.com>

Author: 3 people

pyrit/backend/mappers/attack_mappers.py

@@ -198,8 +198,11 @@ def attack_result_to_summary(
     """
     message_count = stats.message_count
     last_preview = stats.last_message_preview
-    labels = dict(stats.labels) if stats.labels else {}
 
+    # Merge attack-result labels with conversation-level labels.
+    # Conversation labels take precedence on key collision.
+    labels = dict(ar.labels) if ar.labels else {}
+    labels.update(stats.labels or {})
     # Resolution order for created_at: explicit metadata override, then the
     # persisted AttackResult.timestamp, and finally datetime.now() as a
     # last-resort fallback for never-persisted results.

pyrit/backend/services/attack_service.py

@@ -330,6 +330,7 @@ async def create_attack_async(self, *, request: CreateAttackRequest) -> CreateAt
                 "created_at": now.isoformat(),
                 "updated_at": now.isoformat(),
             },
+            labels=labels,
         )
 
         # Store in memory

pyrit/executor/attack/multi_turn/chunked_request.py

@@ -327,6 +327,7 @@ async def _perform_async(self, *, context: ChunkedRequestAttackContext) -> Attac
             outcome_reason=outcome_reason,
             executed_turns=context.executed_turns,
             metadata={"combined_chunks": combined_value, "chunk_count": len(context.chunk_responses)},
+            labels=context.memory_labels,
         )
 
     def _determine_attack_outcome(

pyrit/executor/attack/multi_turn/crescendo.py

@@ -424,6 +424,7 @@ async def _perform_async(self, *, context: CrescendoAttackContext) -> CrescendoA
             last_response=context.last_response.get_piece() if context.last_response else None,
             last_score=context.last_score,
             related_conversations=context.related_conversations,
+            labels=context.memory_labels,
         )
         # setting metadata for backtrack count
         result.backtrack_count = context.backtrack_count

pyrit/executor/attack/multi_turn/multi_prompt_sending.py

@@ -297,6 +297,7 @@ async def _perform_async(self, *, context: MultiTurnAttackContext[Any]) -> Attac
             outcome=outcome,
             outcome_reason=outcome_reason,
             executed_turns=context.executed_turns,
+            labels=context.memory_labels,
         )
 
     def _determine_attack_outcome(

pyrit/executor/attack/multi_turn/red_teaming.py

@@ -338,6 +338,7 @@ async def _perform_async(self, *, context: MultiTurnAttackContext[Any]) -> Attac
             last_response=context.last_response.get_piece() if context.last_response else None,
             last_score=context.last_score,
             related_conversations=context.related_conversations,
+            labels=context.memory_labels,
         )
 
     async def _teardown_async(self, *, context: MultiTurnAttackContext[Any]) -> None:

pyrit/executor/attack/multi_turn/tree_of_attacks.py

@@ -2233,6 +2233,7 @@ def _create_attack_result(
             last_response=last_response,
             last_score=context.best_objective_score,
             related_conversations=context.related_conversations,
+            labels=context.memory_labels,
         )
 
         # Set attack-specific metadata using properties

pyrit/executor/attack/single_turn/prompt_sending.py

@@ -238,6 +238,7 @@ async def _perform_async(self, *, context: SingleTurnAttackContext[Any]) -> Atta
             outcome=outcome,
             outcome_reason=outcome_reason,
             executed_turns=1,
+            labels=context.memory_labels,
         )
 
     def _determine_attack_outcome(

pyrit/executor/attack/single_turn/skeleton_key.py

@@ -181,4 +181,5 @@ def _create_skeleton_key_failure_result(self, *, context: SingleTurnAttackContex
             outcome=AttackOutcome.FAILURE,
             outcome_reason="Skeleton key prompt was filtered or failed",
             executed_turns=1,
+            labels=context.memory_labels,
         )

pyrit/executor/benchmark/fairness_bias.py

@@ -200,6 +200,7 @@ async def _perform_async(self, *, context: FairnessBiasBenchmarkContext) -> Atta
                 atomic_attack_identifier=build_atomic_attack_identifier(
                     attack_identifier=ComponentIdentifier.of(self),
                 ),
+                labels=context.memory_labels,
             )
 
         return last_attack_result