MAINT: Evict I/O modules from pyrit.models into pyrit.memory.storage

pyrit/backend/services/attack_service.py

@@ -50,6 +50,7 @@
 from pyrit.backend.models.common import PaginationInfo
 from pyrit.backend.services.converter_service import get_converter_service
 from pyrit.backend.services.target_service import get_target_service
+from pyrit.io import data_serializer_factory
 from pyrit.memory import CentralMemory
 from pyrit.models import (
     AttackOutcome,
@@ -60,7 +61,6 @@
     MessagePiece,
     PromptDataType,
     build_atomic_attack_identifier,
-    data_serializer_factory,
 )
 from pyrit.prompt_normalizer import PromptConverterConfiguration, PromptNormalizer
 

pyrit/backend/services/converter_service.py

@@ -36,8 +36,8 @@
     CreateConverterResponse,
     PreviewStep,
 )
+from pyrit.io import data_serializer_factory
 from pyrit.models import PromptDataType
-from pyrit.models.data_type_serializer import data_serializer_factory
 from pyrit.prompt_converter import PromptConverter
 from pyrit.prompt_target import PromptTarget
 from pyrit.registry.object_registries import ConverterRegistry

pyrit/common/data_url_converter.py

@@ -2,7 +2,7 @@
 # Licensed under the MIT license.
 
 from pyrit.common.deprecation import print_deprecation_message
-from pyrit.models import DataTypeSerializer, data_serializer_factory
+from pyrit.io import DataTypeSerializer, data_serializer_factory
 
 # Supported image formats for Azure OpenAI GPT-4o,
 # https://learn.microsoft.com/en-us/azure/ai-services/openai/concepts/use-your-image-data

pyrit/common/display_response.py

@@ -8,8 +8,9 @@
 
 from pyrit.common.deprecation import print_deprecation_message
 from pyrit.common.notebook_utils import is_in_ipython_session
+from pyrit.io import AzureBlobStorageIO, DiskStorageIO
 from pyrit.memory import CentralMemory
-from pyrit.models import AzureBlobStorageIO, DiskStorageIO, MessagePiece
+from pyrit.models import MessagePiece
 
 logger = logging.getLogger(__name__)
 

pyrit/datasets/seed_datasets/remote/_image_cache.py

@@ -17,7 +17,7 @@
 from typing import Any, Optional
 
 from pyrit.common.net_utility import make_request_and_raise_if_error_async
-from pyrit.models import data_serializer_factory
+from pyrit.io import data_serializer_factory
 
 logger = logging.getLogger(__name__)
 

pyrit/datasets/seed_datasets/remote/msts_dataset.py

@@ -12,7 +12,8 @@
 from pyrit.datasets.seed_datasets.remote.remote_dataset_loader import (
     _RemoteDatasetLoader,
 )
-from pyrit.models import SeedDataset, SeedPrompt, data_serializer_factory
+from pyrit.io import data_serializer_factory
+from pyrit.models import SeedDataset, SeedPrompt
 
 if TYPE_CHECKING:
     from PIL.Image import Image as PILImage

pyrit/io/__init__.py

@@ -0,0 +1,54 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+"""
+I/O layer for PyRIT: storage backends and multi-modal data serializers.
+
+Provides the disk and blob storage adapters (``StorageIO`` and its
+implementations) and the data-type serializers (``data_serializer_factory`` and
+the per-type ``*DataTypeSerializer`` classes) used to read and write prompt
+payloads such as text, images, audio, and video.
+
+Unlike ``pyrit.models``, modules in this package may depend on ``pyrit.memory``
+and ``pyrit.auth`` (resolved lazily to avoid import cycles).
+"""
+
+from pyrit.io.serializers import (
+    AllowedCategories,
+    AudioPathDataTypeSerializer,
+    BinaryPathDataTypeSerializer,
+    DataTypeSerializer,
+    ErrorDataTypeSerializer,
+    ImagePathDataTypeSerializer,
+    TextDataTypeSerializer,
+    URLDataTypeSerializer,
+    VideoPathDataTypeSerializer,
+    data_serializer_factory,
+    set_message_piece_sha256_async,
+    set_seed_sha256_async,
+)
+from pyrit.io.storage import (
+    AzureBlobStorageIO,
+    DiskStorageIO,
+    StorageIO,
+    SupportedContentType,
+)
+
+__all__ = [
+    "AllowedCategories",
+    "AudioPathDataTypeSerializer",
+    "AzureBlobStorageIO",
+    "BinaryPathDataTypeSerializer",
+    "DataTypeSerializer",
+    "data_serializer_factory",
+    "DiskStorageIO",
+    "ErrorDataTypeSerializer",
+    "ImagePathDataTypeSerializer",
+    "set_message_piece_sha256_async",
+    "set_seed_sha256_async",
+    "StorageIO",
+    "SupportedContentType",
+    "TextDataTypeSerializer",
+    "URLDataTypeSerializer",
+    "VideoPathDataTypeSerializer",
+]