feat: expand SupportedContentType enum with common media types

doc/code/executor/5_workflow.ipynb

@@ -693,7 +693,7 @@
     "from pyrit.executor.workflow import XPIAWorkflow\n",
     "from pyrit.prompt_normalizer import ConverterConfiguration\n",
     "from pyrit.prompt_target import AzureBlobStorageTarget\n",
-    "from pyrit.prompt_target.azure_blob_storage_target import SupportedContentType\n",
+    "from pyrit.memory.storage.storage import SupportedContentType\n",
     "from pyrit.score import SubStringScorer\n",
     "\n",
     "logging.basicConfig(level=logging.DEBUG)\n",
@@ -912,4 +912,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 5
-}
+}

doc/code/executor/5_workflow.py

@@ -139,7 +139,7 @@ async def processing_callback() -> str:
 from pyrit.executor.workflow import XPIAWorkflow
 from pyrit.prompt_normalizer import ConverterConfiguration
 from pyrit.prompt_target import AzureBlobStorageTarget
-from pyrit.prompt_target.azure_blob_storage_target import SupportedContentType
+from pyrit.memory.storage.storage import SupportedContentType
 from pyrit.score import SubStringScorer
 
 logging.basicConfig(level=logging.DEBUG)

pyrit/memory/storage/serializers.py

@@ -74,20 +74,36 @@ def data_serializer_factory(
             f"The 'category' argument is mandatory and must be one of the following: {get_args(AllowedCategories)}."
         )
     if value is not None:
-        if data_type in ["text", "reasoning", "function_call", "tool_call", "function_call_output"]:
+        if data_type in [
+            "text",
+            "reasoning",
+            "function_call",
+            "tool_call",
+            "function_call_output",
+        ]:
             return TextDataTypeSerializer(prompt_text=value, data_type=data_type)
         if data_type == "image_path":
-            return ImagePathDataTypeSerializer(category=category, prompt_text=value, extension=extension)
+            return ImagePathDataTypeSerializer(
+                category=category, prompt_text=value, extension=extension
+            )
         if data_type == "audio_path":
-            return AudioPathDataTypeSerializer(category=category, prompt_text=value, extension=extension)
+            return AudioPathDataTypeSerializer(
+                category=category, prompt_text=value, extension=extension
+            )
         if data_type == "video_path":
-            return VideoPathDataTypeSerializer(category=category, prompt_text=value, extension=extension)
+            return VideoPathDataTypeSerializer(
+                category=category, prompt_text=value, extension=extension
+            )
         if data_type == "binary_path":
-            return BinaryPathDataTypeSerializer(category=category, prompt_text=value, extension=extension)
+            return BinaryPathDataTypeSerializer(
+                category=category, prompt_text=value, extension=extension
+            )
         if data_type == "error":
             return ErrorDataTypeSerializer(prompt_text=value)
         if data_type == "url":
-            return URLDataTypeSerializer(category=category, prompt_text=value, extension=extension)
+            return URLDataTypeSerializer(
+                category=category, prompt_text=value, extension=extension
+            )
         raise ValueError(f"Data type {data_type} not supported")
     if data_type == "image_path":
         return ImagePathDataTypeSerializer(category=category, extension=extension)
@@ -139,7 +155,9 @@ def _get_storage_io(self) -> StorageIO:
             # Scenarios where a user utilizes an in-memory DuckDB but also needs to interact
             # with an Azure Storage Account, ex., XPIAWorkflow.
             if self._memory.results_storage_io is None:
-                raise RuntimeError("results_storage_io is not configured but Azure storage URL was detected")
+                raise RuntimeError(
+                    "results_storage_io is not configured but Azure storage URL was detected"
+                )
             return self._memory.results_storage_io
         return DiskStorageIO()
 
@@ -153,7 +171,9 @@ def data_on_disk(self) -> bool:
 
         """
 
-    async def save_data_async(self, data: bytes, output_filename: str | None = None) -> None:
+    async def save_data_async(
+        self, data: bytes, output_filename: str | None = None
+    ) -> None:
         """
         Save data to storage.
 
@@ -167,10 +187,16 @@ async def save_data_async(self, data: bytes, output_filename: str | None = None)
         file_path = await self.get_data_filename_async(file_name=output_filename)
         if self._memory.results_storage_io is None:
             raise RuntimeError("Storage IO not initialized")
-        await self._memory.results_storage_io.write_file_async(file_path, data)
+        # Infer content type from file extension
+        content_type = self.get_mime_type(str(file_path))
+        await self._memory.results_storage_io.write_file_async(
+            file_path, data, content_type=content_type
+        )
         self.value = str(file_path)
 
-    async def save_b64_image_async(self, data: str | bytes, output_filename: str | None = None) -> None:
+    async def save_b64_image_async(
+        self, data: str | bytes, output_filename: str | None = None
+    ) -> None:
         """
         Save a base64-encoded image to storage.
 
@@ -185,7 +211,11 @@ async def save_b64_image_async(self, data: str | bytes, output_filename: str | N
         image_bytes = base64.b64decode(data)
         if self._memory.results_storage_io is None:
             raise RuntimeError("Storage IO not initialized")
-        await self._memory.results_storage_io.write_file_async(file_path, image_bytes)
+        # Infer content type from file extension
+        content_type = self.get_mime_type(str(file_path))
+        await self._memory.results_storage_io.write_file_async(
+            file_path, image_bytes, content_type=content_type
+        )
         self.value = str(file_path)
 
     async def save_formatted_audio_async(
@@ -213,7 +243,9 @@ async def save_formatted_audio_async(
 
         # save audio file locally first if in AzureStorageBlob so we can use wave.open to set audio parameters
         if self._is_azure_storage_url(str(file_path)):
-            with tempfile.NamedTemporaryFile(suffix=".wav", dir=DB_DATA_PATH, delete=False) as tmp:
+            with tempfile.NamedTemporaryFile(
+                suffix=".wav", dir=DB_DATA_PATH, delete=False
+            ) as tmp:
                 local_temp_path = Path(tmp.name)
             try:
                 await asyncio.to_thread(
@@ -227,8 +259,14 @@ async def save_formatted_audio_async(
                 async with aiofiles.open(local_temp_path, "rb") as f:
                     audio_data = await f.read()
                 if self._memory.results_storage_io is None:
-                    raise RuntimeError("self._memory.results_storage_io is not initialized")
-                await self._memory.results_storage_io.write_file_async(file_path, audio_data)
+                    raise RuntimeError(
+                        "self._memory.results_storage_io is not initialized"
+                    )
+                # Infer content type from file extension
+                content_type = self.get_mime_type(str(file_path))
+                await self._memory.results_storage_io.write_file_async(
+                    file_path, audio_data, content_type=content_type
+                )
             finally:
                 local_temp_path.unlink(missing_ok=True)
 
@@ -259,7 +297,9 @@ async def read_data_async(self) -> bytes:
 
         """
         if not self.data_on_disk():
-            raise TypeError(f"Data for data Type {self.data_type} is not stored on disk")
+            raise TypeError(
+                f"Data for data Type {self.data_type} is not stored on disk"
+            )
 
         if not self.value:
             raise RuntimeError("Prompt text not set")
@@ -309,7 +349,9 @@ async def get_sha256_async(self) -> str:
             if isinstance(self.value, str):
                 input_bytes = self.value.encode("utf-8")
             else:
-                raise ValueError(f"Invalid data type {self.value}, expected str data type.")
+                raise ValueError(
+                    f"Invalid data type {self.value}, expected str data type."
+                )
 
         hash_object = hashlib.sha256(input_bytes)
         return hash_object.hexdigest()
@@ -349,13 +391,19 @@ async def get_data_filename_async(self, file_name: str | None = None) -> Path |
 
         if self._is_azure_storage_url(results_path):
             full_data_directory_path = results_path + self.data_sub_directory
-            self._file_path = full_data_directory_path + f"/{file_name}.{self.file_extension}"
+            self._file_path = (
+                full_data_directory_path + f"/{file_name}.{self.file_extension}"
+            )
         else:
             full_data_directory_path = results_path + self.data_sub_directory
             if self._memory.results_storage_io is None:
                 raise RuntimeError("self._memory.results_storage_io is not initialized")
-            await self._memory.results_storage_io.create_directory_if_not_exists_async(Path(full_data_directory_path))
-            self._file_path = Path(full_data_directory_path, f"{file_name}.{self.file_extension}")
+            await self._memory.results_storage_io.create_directory_if_not_exists_async(
+                Path(full_data_directory_path)
+            )
+            self._file_path = Path(
+                full_data_directory_path, f"{file_name}.{self.file_extension}"
+            )
 
         return self._file_path
 
@@ -401,7 +449,10 @@ def _is_azure_storage_url(self, path: str) -> bool:
 
         """
         parsed = urlparse(path)
-        return parsed.scheme in ("http", "https") and "blob.core.windows.net" in parsed.netloc
+        return (
+            parsed.scheme in ("http", "https")
+            and "blob.core.windows.net" in parsed.netloc
+        )
 
 
 class TextDataTypeSerializer(DataTypeSerializer):
@@ -458,7 +509,9 @@ def data_on_disk(self) -> bool:
 class URLDataTypeSerializer(DataTypeSerializer):
     """Serializer for URL values and URL-backed local file references."""
 
-    def __init__(self, *, category: str, prompt_text: str, extension: str | None = None) -> None:
+    def __init__(
+        self, *, category: str, prompt_text: str, extension: str | None = None
+    ) -> None:
         """
         Initialize a URL serializer.
 
@@ -488,7 +541,13 @@ def data_on_disk(self) -> bool:
 class ImagePathDataTypeSerializer(DataTypeSerializer):
     """Serializer for image path values stored on disk."""
 
-    def __init__(self, *, category: str, prompt_text: str | None = None, extension: str | None = None) -> None:
+    def __init__(
+        self,
+        *,
+        category: str,
+        prompt_text: str | None = None,
+        extension: str | None = None,
+    ) -> None:
         """
         Initialize an image-path serializer.
 

pyrit/memory/storage/storage.py

@@ -25,8 +25,40 @@ class SupportedContentType(Enum):
     See all options here: https://www.iana.org/assignments/media-types/media-types.xhtml.
     """
 
-    # TODO, add other media supported types
+    # Text types
     PLAIN_TEXT = "text/plain"
+    HTML = "text/html"
+    JSON = "application/json"
+    XML = "application/xml"
+    CSV = "text/csv"
+    MARKDOWN = "text/markdown"
+
+    # Image types
+    PNG = "image/png"
+    JPEG = "image/jpeg"
+    GIF = "image/gif"
+    WEBP = "image/webp"
+    SVG = "image/svg+xml"
+    BMP = "image/bmp"
+
+    # Audio types
+    WAV = "audio/wav"
+    MP3 = "audio/mpeg"
+    OGG = "audio/ogg"
+    FLAC = "audio/flac"
+    M4A = "audio/mp4"
+
+    # Video types
+    MP4 = "video/mp4"
+    WEBM = "video/webm"
+    OGG_VIDEO = "video/ogg"
+    AVI = "video/x-msvideo"
+
+    # Document types
+    PDF = "application/pdf"
+    ZIP = "application/zip"
+    TAR = "application/x-tar"
+    GZIP = "application/gzip"
 
 
 class StorageIO(ABC):
@@ -41,7 +73,9 @@ async def read_file_async(self, path: Path | str) -> bytes:
         """
 
     @abstractmethod
-    async def write_file_async(self, path: Path | str, data: bytes) -> None:
+    async def write_file_async(
+        self, path: Path | str, data: bytes, content_type: str | None = None
+    ) -> None:
         """
         Asynchronously writes data to the given path.
         """
@@ -85,13 +119,16 @@ async def read_file_async(self, path: Path | str) -> bytes:
         async with aiofiles.open(path, "rb") as file:
             return await file.read()
 
-    async def write_file_async(self, path: Path | str, data: bytes) -> None:
+    async def write_file_async(
+        self, path: Path | str, data: bytes, content_type: str | None = None
+    ) -> None:
         """
         Asynchronously writes data to a file on the local disk.
 
         Args:
             path (Path): The path to the file.
             data (bytes): The content to write to the file.
+            content_type (str | None): Optional content type (unused for local disk storage).
 
         """
         path = self._convert_to_path(path)
@@ -211,7 +248,9 @@ async def _create_container_client_async(self) -> AsyncContainerClient:
 
         from azure.identity.aio import DefaultAzureCredential
 
-        logger.info("SAS token not provided. Using DefaultAzureCredential for direct Entra ID authentication.")
+        logger.info(
+            "SAS token not provided. Using DefaultAzureCredential for direct Entra ID authentication."
+        )
         parsed_url = urlparse(self._container_url)
         path_parts = [part for part in parsed_url.path.split("/") if part]
         if not path_parts:
@@ -240,7 +279,9 @@ async def _close_client_async(self) -> None:
             if credential:
                 await credential.close()
 
-    async def _upload_blob_async(self, file_name: str, data: bytes, content_type: str) -> None:
+    async def _upload_blob_async(
+        self, file_name: str, data: bytes, content_type: str
+    ) -> None:
         """
         (Async) Handles uploading blob to given storage container.
 
@@ -368,7 +409,9 @@ async def read_file_async(self, path: Path | str) -> bytes:
         finally:
             await self._close_client_async()
 
-    async def write_file_async(self, path: Path | str, data: bytes) -> None:
+    async def write_file_async(
+        self, path: Path | str, data: bytes, content_type: str | None = None
+    ) -> None:
         """
         Write data to Azure Blob Storage at the specified path.
 
@@ -378,12 +421,25 @@ async def write_file_async(self, path: Path | str, data: bytes) -> None:
         Args:
             path (Path | str): Full blob URL or relative blob path.
             data (bytes): The data to write.
+            content_type (str | None): Optional MIME type for the blob. If not provided,
+                it will be inferred from the file extension, falling back to the default
+                blob_content_type configured in the constructor.
         """
         if not self._client_async:
             self._client_async = await self._create_container_client_async()
         blob_name = self._resolve_blob_name(path)
+
+        # Determine content type: explicit > inferred from extension > default
+        if content_type is None:
+            from mimetypes import guess_type
+
+            inferred_type, _ = guess_type(blob_name)
+            content_type = inferred_type if inferred_type else self._blob_content_type
+
         try:
-            await self._upload_blob_async(file_name=blob_name, data=data, content_type=self._blob_content_type)
+            await self._upload_blob_async(
+                file_name=blob_name, data=data, content_type=content_type
+            )
         except Exception as exc:
             logger.exception(f"Failed to write file at {blob_name}: {exc}")
             raise