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

@@ -10,15 +10,14 @@
 import tempfile
 import time
 import wave
-from mimetypes import guess_type
 from pathlib import Path
 from typing import TYPE_CHECKING, Literal, get_args
 from urllib.parse import urlparse
 
 import aiofiles
 
 from pyrit.common.path import DB_DATA_PATH
-from pyrit.memory.storage.storage import DiskStorageIO, StorageIO
+from pyrit.memory.storage.storage import DiskStorageIO, StorageIO, _EXTENSION_TO_CONTENT_TYPE
 
 if TYPE_CHECKING:
     from pyrit.memory import MemoryInterface
@@ -74,20 +73,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 +154,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 +170,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 +186,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 +210,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(
@@ -209,11 +238,16 @@ async def save_formatted_audio_async(
         Raises:
             RuntimeError: If storage IO is not initialized.
         """
+        # Force .wav extension since we're writing WAV format
+        if output_filename and not output_filename.lower().endswith(".wav"):
+            output_filename = Path(output_filename).stem + ".wav"
         file_path = await self.get_data_filename_async(file_name=output_filename)
 
         # 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 +261,13 @@ 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"
+                    )
+                # Always use audio/wav since we're writing WAV format
+                await self._memory.results_storage_io.write_file_async(
+                    file_path, audio_data, content_type="audio/wav"
+                )
             finally:
                 local_temp_path.unlink(missing_ok=True)
 
@@ -259,7 +298,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 +350,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 +392,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
 
@@ -377,7 +426,7 @@ def get_extension(file_path: str) -> str | None:
     @staticmethod
     def get_mime_type(file_path: str) -> str | None:
         """
-        Get the MIME type of the file path.
+        Get the MIME type of the file path using deterministic mapping.
 
         Args:
             file_path (str): Input file path.
@@ -386,8 +435,9 @@ def get_mime_type(file_path: str) -> str | None:
             str | None: MIME type if detectable; otherwise None.
 
         """
-        mime_type, _ = guess_type(file_path)
-        return mime_type
+        ext = Path(file_path).suffix.lower()
+        content_type = _EXTENSION_TO_CONTENT_TYPE.get(ext)
+        return content_type.value if content_type else None
 
     def _is_azure_storage_url(self, path: str) -> bool:
         """
@@ -401,7 +451,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 +511,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 +543,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.