MAINT: Minor pyrit.model module reorganization (#1999)

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

Author: rlundeen2

pyrit/backend/routes/scenarios.py

@@ -24,7 +24,7 @@
 )
 from pyrit.backend.services.scenario_run_service import get_scenario_run_service
 from pyrit.backend.services.scenario_service import get_scenario_service
-from pyrit.models.scenario_result import ScenarioResult
+from pyrit.models import ScenarioResult
 
 router = APIRouter(prefix="/scenarios", tags=["scenarios"])
 

pyrit/cli/_output.py

@@ -284,7 +284,7 @@ async def print_scenario_result_async(*, result_dict: dict[str, Any]) -> None:
     Args:
         result_dict: ``ScenarioResult.model_dump(mode="json", by_alias=True)`` 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.model_validate(result_dict)

pyrit/memory/memory_models.py

@@ -56,7 +56,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 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
@@ -85,11 +79,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
@@ -206,6 +207,7 @@
     "TargetCapabilities",
     "TargetIdentifier",
     "TextDataTypeSerializer",
+    "ToolCall",
     "UnvalidatedScore",
     "validate_registry_name",
     "VideoPathDataTypeSerializer",

pyrit/models/attack_result.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/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)
     """