MAINT: Making Message a Pydantic model

pyrit/backend/mappers/attack_mappers.py

@@ -537,7 +537,7 @@ def request_to_pyrit_message(
         )
         for p in request.pieces
     ]
-    return PyritMessage(pieces)
+    return PyritMessage(message_pieces=pieces)
 
 
 # ============================================================================

pyrit/executor/attack/component/conversation_manager.py

@@ -481,7 +481,7 @@ async def add_prepended_conversation_to_memory_async(
         turn_count = 0
 
         for i, message in enumerate(valid_messages):
-            message_copy = message.duplicate_message()
+            message_copy = message.duplicate()
 
             message_copy.set_simulated_role()
 

pyrit/executor/attack/multi_turn/crescendo.py

@@ -797,7 +797,7 @@ async def _generate_next_prompt_async(self, context: CrescendoAttackContext) ->
         if context.next_message:
             self._logger.debug("Using custom message, bypassing adversarial chat")
             # Duplicate to ensure fresh IDs (avoids conflicts if message was already in memory)
-            message = context.next_message.duplicate_message()
+            message = context.next_message.duplicate()
             context.next_message = None  # Clear for future turns
             return message
 

pyrit/executor/attack/multi_turn/tree_of_attacks.py

@@ -608,7 +608,7 @@ async def _send_initial_prompt_to_target_async(self) -> Message:
             self.objective_target_conversation_id = str(uuid.uuid4())
 
         # Duplicate to ensure fresh IDs (avoids conflicts if message was already in memory)
-        message = self._initial_prompt.duplicate_message()
+        message = self._initial_prompt.duplicate()
         self._initial_prompt = None  # Clear for future turns
 
         # Store the prompt text for reference

pyrit/executor/attack/single_turn/prompt_sending.py

@@ -293,7 +293,7 @@ def _get_message(self, context: SingleTurnAttackContext[Any]) -> Message:
         """
         if context.next_message:
             # Deep copy the message to preserve all fields, then assign new IDs
-            return context.next_message.duplicate_message()
+            return context.next_message.duplicate()
 
         return Message.from_prompt(prompt=context.objective, role="user")
 

pyrit/executor/promptgen/fuzzer/fuzzer_converter_base.py

@@ -90,7 +90,7 @@ async def convert_async(self, *, prompt: str, input_type: PromptDataType = "text
         formatted_prompt = f"===={self.template_label} BEGINS====\n{prompt}\n===={self.template_label} ENDS===="
         prompt_metadata: dict[str, str | int] = {"response_format": "json"}
         request = Message(
-            [
+            message_pieces=[
                 MessagePiece(
                     role="user",
                     original_value=formatted_prompt,

pyrit/executor/promptgen/fuzzer/fuzzer_crossover_converter.py

@@ -92,7 +92,7 @@ async def convert_async(self, *, prompt: str, input_type: PromptDataType = "text
 
         prompt_metadata: dict[str, str | int] = {"response_format": "json"}
         request = Message(
-            [
+            message_pieces=[
                 MessagePiece(
                     role="user",
                     original_value=formatted_prompt,

pyrit/executor/promptgen/fuzzer/fuzzer_expand_converter.py

@@ -66,7 +66,7 @@ async def convert_async(self, *, prompt: str, input_type: PromptDataType = "text
 
         prompt_metadata: dict[str, str | int] = {"response_format": "json"}
         request = Message(
-            [
+            message_pieces=[
                 MessagePiece(
                     role="user",
                     original_value=formatted_prompt,

pyrit/memory/memory_interface.py

@@ -1009,7 +1009,7 @@ def duplicate_messages(self, *, messages: Sequence[Message]) -> tuple[str, Seque
 
         all_pieces: list[MessagePiece] = []
         for message in messages:
-            duplicated_message = message.duplicate_message()
+            duplicated_message = message.duplicate()
 
             for piece in duplicated_message.message_pieces:
                 piece.conversation_id = new_conversation_id

pyrit/message_normalizer/chat_message_normalizer.py

@@ -14,7 +14,7 @@
     apply_system_message_behavior,
 )
 from pyrit.models import ChatMessage, DataTypeSerializer, Message
-from pyrit.models.message_piece import MessagePiece
+from pyrit.models.messages.message_piece import MessagePiece
 
 if TYPE_CHECKING:
     from pyrit.models.literals import ChatMessageRole

pyrit/models/__init__.py

@@ -63,13 +63,16 @@
     validate_registry_name,
 )
 from pyrit.models.literals import ChatMessageRole, Modality, PromptDataType, PromptResponseError, SeedType
-from pyrit.models.message import (
+from pyrit.models.messages import (
     Message,
+    MessagePiece,
     construct_response_from_request,
+    flatten_to_message_pieces,
+    get_all_values,
     group_conversation_message_pieces_by_sequence,
     group_message_pieces_into_conversations,
+    sort_message_pieces,
 )
-from pyrit.models.message_piece import MessagePiece, sort_message_pieces
 from pyrit.models.question_answering import QuestionAnsweringDataset, QuestionAnsweringEntry, QuestionChoice
 from pyrit.models.retry_event import RetryEvent
 from pyrit.models.scenario_result import ScenarioIdentifier, ScenarioResult
@@ -128,7 +131,9 @@
     "EmbeddingUsageInformation",
     "ErrorDataTypeSerializer",
     "EvaluationIdentifier",
+    "flatten_to_message_pieces",
     "get_all_harm_definitions",
+    "get_all_values",
     "group_conversation_message_pieces_by_sequence",
     "group_message_pieces_into_conversations",
     "HarmDefinition",

pyrit/models/attack_result.py

@@ -14,7 +14,7 @@
 from pyrit.models.conversation_reference import ConversationReference, ConversationType
 from pyrit.models.identifiers.atomic_attack_identifier import build_atomic_attack_identifier
 from pyrit.models.identifiers.component_identifier import ComponentIdentifier
-from pyrit.models.message_piece import MessagePiece
+from pyrit.models.messages.message_piece import MessagePiece
 from pyrit.models.retry_event import RetryEvent
 from pyrit.models.score import Score
 from pyrit.models.strategy_result import StrategyResult