MAINT: Minor pyrit.model module reorganization

pyrit/cli/_output.py

@@ -284,7 +284,7 @@ async def print_scenario_result_async(*, result_dict: dict[str, Any]) -> None:
     Args:
         result_dict: ``ScenarioResult.to_dict()`` payload from the REST API.
     """
-    from pyrit.models.scenario_result import ScenarioResult
+    from pyrit.models import ScenarioResult
     from pyrit.output.scenario_result.pretty import PrettyScenarioResultMemoryPrinter
 
     scenario_result = ScenarioResult.from_dict(result_dict)

pyrit/memory/memory_models.py

@@ -55,7 +55,6 @@
     SeedSimulatedConversation,
     SeedType,
 )
-from pyrit.models.scenario_result import ScenarioRunState
 
 logger = logging.getLogger(__name__)
 

pyrit/models/__init__.py

@@ -21,12 +21,6 @@
 from typing import TYPE_CHECKING, Any
 
 from pyrit.common.deprecation import print_deprecation_message
-from pyrit.models.chat_message import (
-    ALLOWED_CHAT_MESSAGE_ROLES,
-    ChatMessage,
-    ChatMessagesDataset,
-)
-from pyrit.models.conversation_reference import ConversationReference, ConversationType
 from pyrit.models.conversation_stats import ConversationStats
 from pyrit.models.embeddings import EmbeddingData, EmbeddingResponse, EmbeddingSupport, EmbeddingUsageInformation
 from pyrit.models.harm_definition import HarmDefinition, ScaleDescription, get_all_harm_definitions
@@ -79,11 +73,18 @@
     group_message_pieces_into_conversations,
     sort_message_pieces,
 )
+from pyrit.models.messages.chat_message import (
+    ALLOWED_CHAT_MESSAGE_ROLES,
+    ChatMessage,
+    ChatMessagesDataset,
+    ToolCall,
+)
+from pyrit.models.messages.conversation_reference import ConversationReference, ConversationType
 from pyrit.models.question_answering import QuestionAnsweringDataset, QuestionAnsweringEntry, QuestionChoice
 from pyrit.models.results.attack_result import AttackOutcome, AttackResult, AttackResultT
+from pyrit.models.results.scenario_result import ScenarioIdentifier, ScenarioResult, ScenarioRunState
 from pyrit.models.results.strategy_result import StrategyResult, StrategyResultT
 from pyrit.models.retry_event import RetryEvent
-from pyrit.models.scenario_result import ScenarioIdentifier, ScenarioResult, ScenarioRunState
 from pyrit.models.score import Score, ScoreType, UnvalidatedScore
 
 # Seeds - import from new seeds submodule for forward compatibility
@@ -192,6 +193,7 @@
     "TARGET_EVAL_PARAM_FALLBACKS",
     "TARGET_EVAL_PARAMS",
     "TextDataTypeSerializer",
+    "ToolCall",
     "UnvalidatedScore",
     "validate_registry_name",
     "VideoPathDataTypeSerializer",

pyrit/models/messages/__init__.py

@@ -6,11 +6,21 @@
 
 - MessagePiece: A single piece of a message exchanged with a target.
 - Message: One request/response to a target, made up of one or more pieces.
+- ChatMessage: OpenAI-style wire shape consumed/emitted by prompt targets.
+- Conversation: Conversation-scoped metadata shared by every piece.
+- ConversationReference: Immutable reference to a conversation in an attack.
 - conversations: Free functions that operate on collections of messages/pieces.
 """
 
-from pyrit.models.messages.conversation import Conversation
+from pyrit.models.messages.chat_message import (
+    ALLOWED_CHAT_MESSAGE_ROLES,
+    ChatMessage,
+    ChatMessagesDataset,
+    ToolCall,
+)
+from pyrit.models.messages.conversation_reference import ConversationReference, ConversationType
 from pyrit.models.messages.conversations import (
+    Conversation,
     construct_response_from_request,
     flatten_to_message_pieces,
     get_all_values,
@@ -21,9 +31,15 @@
 from pyrit.models.messages.message_piece import MessagePiece, sort_message_pieces
 
 __all__ = [
+    "ALLOWED_CHAT_MESSAGE_ROLES",
+    "ChatMessage",
+    "ChatMessagesDataset",
     "Conversation",
+    "ConversationReference",
+    "ConversationType",
     "Message",
     "MessagePiece",
+    "ToolCall",
     "construct_response_from_request",
     "flatten_to_message_pieces",
     "get_all_values",

pyrit/models/chat_message.py → pyrit/models/messages/chat_message.py

@@ -1,6 +1,20 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT license.
 
+"""
+OpenAI-format chat message types.
+
+``ChatMessage`` is the OpenAI Chat Completions wire shape — a ``role`` plus a
+string-or-multipart ``content``, with the OpenAI ``name`` / ``tool_calls`` /
+``tool_call_id`` fields. Prompt targets that speak the OpenAI API (and the many
+providers that mirror it) consume and emit these objects directly.
+
+It is intentionally distinct from the PyRIT domain ``Message`` / ``MessagePiece``
+types in this same package: those model a persisted request/response exchange,
+whereas ``ChatMessage`` is the lightweight OpenAI-shaped transport representation
+handed to a model API.
+"""
+
 from typing import Any
 
 from pydantic import BaseModel, ConfigDict
@@ -21,9 +35,9 @@ class ToolCall(BaseModel):
 
 class ChatMessage(BaseModel):
     """
-    Represents a chat message for API consumption.
+    Represents a single OpenAI Chat Completions message.
 
-    The content field can be:
+    Mirrors the OpenAI message schema. The content field can be:
     - A simple string for single-part text messages
     - A list of dicts for multipart messages (e.g., text + images)
     """

pyrit/models/messages/conversations.py

@@ -1,21 +1,47 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT license.
 
-"""Helpers that operate on collections of ``Message`` / ``MessagePiece``."""
+"""``Conversation`` model plus helpers that operate on collections of ``Message`` / ``MessagePiece``."""
 
 from __future__ import annotations
 
 from typing import TYPE_CHECKING
 
+from pydantic import BaseModel, ConfigDict
+
 from pyrit.models.messages.message import Message
 from pyrit.models.messages.message_piece import MessagePiece
+from pyrit.models.score import (  # noqa: TC001  (runtime-required by Pydantic field annotations)
+    ComponentIdentifierField,
+)
 
 if TYPE_CHECKING:
     from collections.abc import MutableSequence, Sequence
 
     from pyrit.models.literals import PromptDataType, PromptResponseError
 
 
+class Conversation(BaseModel):
+    """
+    Conversation-scoped metadata shared by every piece in a conversation.
+
+    A ``Conversation`` records identifiers that belong to the conversation as a
+    whole rather than to any individual ``MessagePiece`` -- most importantly the
+    target the conversation is held with. Persisting these once per conversation
+    (instead of stamping them onto every piece/row) is what keeps ``MessagePiece``
+    small.
+    """
+
+    model_config = ConfigDict(
+        arbitrary_types_allowed=True,
+        extra="forbid",
+        validate_assignment=False,
+    )
+
+    conversation_id: str
+    target_identifier: ComponentIdentifierField | None = None
+
+
 def get_all_values(messages: Sequence[Message]) -> list[str]:
     """
     Return all converted values across the provided messages.

pyrit/models/results/__init__.py

@@ -2,20 +2,31 @@
 # Licensed under the MIT license.
 
 """
-Results module - strategy and attack result types for PyRIT.
+Results module - strategy, attack, and scenario result types for PyRIT.
 
 - StrategyResult: Base class for all strategy results.
 - AttackResult: Result of an attack execution, with conversation/scoring evidence.
 - AttackOutcome: Enum of possible attack outcomes.
+- ScenarioResult: Aggregate result of a scenario run.
+- ScenarioIdentifier: Identifier describing the executed scenario.
+- ScenarioRunState: Lifecycle state of a scenario run.
 """
 
 from pyrit.models.results.attack_result import AttackOutcome, AttackResult, AttackResultT
+from pyrit.models.results.scenario_result import (
+    ScenarioIdentifier,
+    ScenarioResult,
+    ScenarioRunState,
+)
 from pyrit.models.results.strategy_result import StrategyResult, StrategyResultT
 
 __all__ = [
     "AttackOutcome",
     "AttackResult",
     "AttackResultT",
+    "ScenarioIdentifier",
+    "ScenarioResult",
+    "ScenarioRunState",
     "StrategyResult",
     "StrategyResultT",
 ]

pyrit/models/results/attack_result.py

@@ -11,8 +11,8 @@
 from pydantic import AwareDatetime, Field
 
 from pyrit.common.deprecation import print_deprecation_message
-from pyrit.models.conversation_reference import ConversationReference, ConversationType
 from pyrit.models.identifiers.component_identifier import ComponentIdentifier
+from pyrit.models.messages.conversation_reference import ConversationReference, ConversationType
 from pyrit.models.messages.message_piece import MessagePiece
 from pyrit.models.results.strategy_result import StrategyResult
 from pyrit.models.retry_event import RetryEvent

pyrit/output/helpers.py

@@ -11,8 +11,7 @@
 
 import os
 
-from pyrit.models import AttackResult, ComponentIdentifier, Message, Score
-from pyrit.models.scenario_result import ScenarioResult
+from pyrit.models import AttackResult, ComponentIdentifier, Message, ScenarioResult, Score
 from pyrit.output.attack_result.markdown import MarkdownAttackResultMemoryPrinter
 from pyrit.output.attack_result.pretty import PrettyAttackResultMemoryPrinter
 from pyrit.output.conversation.pretty import PrettyConversationMemoryPrinter

pyrit/output/scenario_result/base.py

@@ -4,7 +4,7 @@
 from abc import abstractmethod
 
 from pyrit.common.deprecation import print_deprecation_message
-from pyrit.models.scenario_result import ScenarioResult
+from pyrit.models import ScenarioResult
 from pyrit.output.base import PrinterBase
 
 

pyrit/output/scenario_result/pretty.py

@@ -6,8 +6,7 @@
 from colorama import Fore, Style
 
 from pyrit.common.deprecation import print_deprecation_message
-from pyrit.models import AttackOutcome
-from pyrit.models.scenario_result import ScenarioResult
+from pyrit.models import AttackOutcome, ScenarioResult
 from pyrit.output.scenario_result.base import ScenarioResultPrinterBase
 from pyrit.output.scorer.base import ScorerPrinterBase
 from pyrit.output.sink import Sink

pyrit/scenario/__init__.py

@@ -19,7 +19,7 @@
 from types import ModuleType
 
 from pyrit.common.parameter import Parameter
-from pyrit.models.scenario_result import ScenarioIdentifier, ScenarioResult
+from pyrit.models import ScenarioIdentifier, ScenarioResult
 from pyrit.scenario.core import (
     AtomicAttack,
     AttackTechnique,