MAINT: Consolidating Target Metadata

pyrit/converter/decomposition_converter.py

@@ -14,8 +14,8 @@
 from pyrit.converter.converter import Converter, ConverterResult
 from pyrit.exceptions import InvalidJsonException, pyrit_json_retry, remove_markdown_json
 from pyrit.models import (
-    JSON_SCHEMA_METADATA_KEY,
     ComponentIdentifier,
+    JsonResponseConfig,
     Message,
     MessagePiece,
     PromptDataType,
@@ -236,9 +236,10 @@ async def _decompose_async(self, *, objective: str) -> tuple[list[str], list[str
             conversation_id=conversation_id,
         )
 
-        prompt_metadata: dict[str, Any] = {"response_format": "json"}
-        if self._decomposition_prompt.response_json_schema is not None:
-            prompt_metadata[JSON_SCHEMA_METADATA_KEY] = self._decomposition_prompt.response_json_schema
+        prompt_metadata = JsonResponseConfig(
+            enabled=True,
+            json_schema=self._decomposition_prompt.response_json_schema,
+        ).to_metadata()
 
         request = Message(
             message_pieces=[

pyrit/executor/attack/component/adversarial_conversation_manager.py

@@ -27,7 +27,7 @@
     resolve_adversarial_system_prompt,
 )
 from pyrit.models import (
-    JSON_SCHEMA_METADATA_KEY,
+    JsonResponseConfig,
     JsonSchemaDefinition,
     Message,
     Score,
@@ -208,7 +208,7 @@ def _build_adversarial_prompt_metadata(*, response_json_schema: JsonSchemaDefini
     """
     if response_json_schema is None:
         return {}
-    return {"response_format": "json", JSON_SCHEMA_METADATA_KEY: response_json_schema}
+    return JsonResponseConfig(enabled=True, json_schema=response_json_schema).to_metadata()
 
 
 def _parse_adversarial_reply(response_text: str, *, schema: JsonSchemaDefinition | None = None) -> AdversarialReply:

pyrit/executor/promptgen/fuzzer/fuzzer_converter_base.py

@@ -14,6 +14,7 @@
     remove_markdown_json,
 )
 from pyrit.models import (
+    JsonResponseConfig,
     Message,
     MessagePiece,
     PromptDataType,
@@ -87,7 +88,7 @@ async def convert_async(self, *, prompt: str, input_type: PromptDataType = "text
         )
 
         formatted_prompt = f"===={self.template_label} BEGINS====\n{prompt}\n===={self.template_label} ENDS===="
-        prompt_metadata: dict[str, str | int] = {"response_format": "json"}
+        prompt_metadata = JsonResponseConfig(enabled=True).to_metadata()
         request = Message(
             message_pieces=[
                 MessagePiece(

pyrit/executor/promptgen/fuzzer/fuzzer_crossover_converter.py

@@ -12,7 +12,7 @@
 from pyrit.executor.promptgen.fuzzer.fuzzer_converter_base import (
     FuzzerConverter,
 )
-from pyrit.models import Message, MessagePiece, PromptDataType, SeedPrompt
+from pyrit.models import JsonResponseConfig, Message, MessagePiece, PromptDataType, SeedPrompt
 from pyrit.prompt_target import PromptTarget
 
 
@@ -89,7 +89,7 @@ async def convert_async(self, *, prompt: str, input_type: PromptDataType = "text
             f"\n====TEMPLATE 2 BEGINS====\n{random.choice(self.prompt_templates)}\n====TEMPLATE 2 ENDS====\n"
         )
 
-        prompt_metadata: dict[str, str | int] = {"response_format": "json"}
+        prompt_metadata = JsonResponseConfig(enabled=True).to_metadata()
         request = Message(
             message_pieces=[
                 MessagePiece(

pyrit/executor/promptgen/fuzzer/fuzzer_expand_converter.py

@@ -10,7 +10,7 @@
 from pyrit.executor.promptgen.fuzzer.fuzzer_converter_base import (
     FuzzerConverter,
 )
-from pyrit.models import Message, MessagePiece, PromptDataType, SeedPrompt
+from pyrit.models import JsonResponseConfig, Message, MessagePiece, PromptDataType, SeedPrompt
 from pyrit.prompt_target import PromptTarget
 
 
@@ -62,7 +62,7 @@ async def convert_async(self, *, prompt: str, input_type: PromptDataType = "text
 
         formatted_prompt = f"===={self.template_label} BEGINS====\n{prompt}\n===={self.template_label} ENDS===="
 
-        prompt_metadata: dict[str, str | int] = {"response_format": "json"}
+        prompt_metadata = JsonResponseConfig(enabled=True).to_metadata()
         request = Message(
             message_pieces=[
                 MessagePiece(

pyrit/models/__init__.py

@@ -49,15 +49,6 @@
     snake_case_to_class_name,
     validate_registry_name,
 )
-from pyrit.models.json_schema_definition import (
-    COMMON_JSON_SCHEMAS,
-    JSON_SCHEMA_METADATA_KEY,
-    SEED_RESPONSE_JSON_SCHEMA_METADATA_KEY,
-    JsonSchemaDefinition,
-    get_common_json_schema,
-    register_common_json_schema,
-    unregister_common_json_schema,
-)
 from pyrit.models.literals import (
     MEDIA_PATH_DATA_TYPES,
     ChatMessageRole,
@@ -114,8 +105,19 @@
     SimulatedTargetSystemPromptPaths,
     group_seeds_into_attack_groups,
 )
-from pyrit.models.target_capabilities import CapabilityName, TargetCapabilities
-from pyrit.models.token_usage import TokenUsage
+from pyrit.models.target import (
+    COMMON_JSON_SCHEMAS,
+    JSON_SCHEMA_METADATA_KEY,
+    SEED_RESPONSE_JSON_SCHEMA_METADATA_KEY,
+    CapabilityName,
+    JsonResponseConfig,
+    JsonSchemaDefinition,
+    TargetCapabilities,
+    TokenUsage,
+    get_common_json_schema,
+    register_common_json_schema,
+    unregister_common_json_schema,
+)
 
 __all__ = [
     "ALLOWED_CHAT_MESSAGE_ROLES",
@@ -161,6 +163,7 @@
     "IdentifierType",
     "JSONValue",
     "COMMON_JSON_SCHEMAS",
+    "JsonResponseConfig",
     "get_common_json_schema",
     "register_common_json_schema",
     "unregister_common_json_schema",

pyrit/models/catalog/target.py

@@ -26,7 +26,7 @@
 
 from pyrit.models.identifiers.component_identifier import JSONValue
 from pyrit.models.identifiers.target_identifier import TargetIdentifier
-from pyrit.models.target_capabilities import TargetCapabilities
+from pyrit.models.target.target_capabilities import TargetCapabilities
 
 
 class TargetInstance(BaseModel):

pyrit/models/seeds/seed_prompt.py

@@ -15,15 +15,15 @@
 from tinytag import TinyTag
 
 from pyrit.common.path import PATHS_DICT
-from pyrit.models.json_schema_definition import (  # noqa: TC001  (runtime-required by Pydantic field annotations)
-    JsonSchemaDefinition,
-    get_common_json_schema,
-)
 from pyrit.models.literals import (  # noqa: TC001  (runtime-required by Pydantic field annotations)
     ChatMessageRole,
     PromptDataType,
 )
 from pyrit.models.seeds.seed import Seed
+from pyrit.models.target.json_schema_definition import (  # noqa: TC001  (runtime-required by Pydantic field annotations)
+    JsonSchemaDefinition,
+    get_common_json_schema,
+)
 
 if TYPE_CHECKING:
     import uuid

pyrit/models/target/__init__.py

@@ -0,0 +1,45 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+"""
+Canonical data models for how PyRIT interacts with targets.
+
+This sub-package groups the value objects that describe a target interaction and
+own their own ``MessagePiece.prompt_metadata`` (de)serialization:
+
+* ``TokenUsage`` — provider-agnostic token accounting for a model call.
+* ``JsonResponseConfig`` — PyRIT's canonical JSON-response request config.
+* ``TargetCapabilities`` / ``CapabilityName`` — what a target natively supports.
+* ``JsonSchemaDefinition`` and the shared JSON-schema registry / metadata keys.
+
+Everything here is re-exported from the top-level ``pyrit.models`` package, so
+callers should keep importing from ``pyrit.models`` (e.g.
+``from pyrit.models import TokenUsage``).
+"""
+
+from pyrit.models.target.json_response_config import JsonResponseConfig
+from pyrit.models.target.json_schema_definition import (
+    COMMON_JSON_SCHEMAS,
+    JSON_SCHEMA_METADATA_KEY,
+    SEED_RESPONSE_JSON_SCHEMA_METADATA_KEY,
+    JsonSchemaDefinition,
+    get_common_json_schema,
+    register_common_json_schema,
+    unregister_common_json_schema,
+)
+from pyrit.models.target.target_capabilities import CapabilityName, TargetCapabilities
+from pyrit.models.target.token_usage import TokenUsage
+
+__all__ = [
+    "COMMON_JSON_SCHEMAS",
+    "CapabilityName",
+    "JSON_SCHEMA_METADATA_KEY",
+    "JsonResponseConfig",
+    "JsonSchemaDefinition",
+    "SEED_RESPONSE_JSON_SCHEMA_METADATA_KEY",
+    "TargetCapabilities",
+    "TokenUsage",
+    "get_common_json_schema",
+    "register_common_json_schema",
+    "unregister_common_json_schema",
+]

pyrit/models/target/json_response_config.py

@@ -0,0 +1,126 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+from __future__ import annotations
+
+import json
+from typing import Any, ClassVar
+
+from pydantic import BaseModel, ConfigDict, model_validator
+
+from pyrit.models.target.json_schema_definition import (
+    JSON_SCHEMA_METADATA_KEY,  # noqa: TC001  (runtime-required by Pydantic field annotations)
+)
+
+
+class JsonResponseConfig(BaseModel):
+    """
+    Canonical PyRIT configuration for requesting a JSON response from a target.
+
+    A value object that owns PyRIT's canonical ``MessagePiece.prompt_metadata`` keys for
+    JSON responses (``"response_format"``, ``JSON_SCHEMA_METADATA_KEY``, ``"json_schema_name"``,
+    ``"json_schema_strict"``) and (de)serializes them via ``from_metadata`` / ``to_metadata``.
+    Producers (scorers, attacks, converters) build one and call ``to_metadata`` to attach it to
+    a piece; targets read it back with ``from_metadata``.
+
+    Providing a ``json_schema`` implies ``enabled`` (a schema is meaningless without JSON output),
+    so JSON-object mode is ``enabled=True`` with no schema and JSON-schema mode is just
+    ``json_schema=...``.
+
+    These are PyRIT keys, not a provider's wire format. Translating this config into a specific
+    provider's request block (e.g. the OpenAI chat ``response_format`` or Responses ``text.format``
+    shape) is the target's job and lives in ``pyrit.prompt_target`` (``build_response_format`` /
+    ``_build_text_format``).
+
+    For the provider shapes those translators emit, see:
+    https://platform.openai.com/docs/api-reference/chat/create#chat_create-response_format-json_schema
+    and
+    https://platform.openai.com/docs/api-reference/responses/create#responses_create-text
+    """
+
+    model_config = ConfigDict(extra="forbid")
+
+    _METADATAKEYS: ClassVar[dict[str, str]] = {
+        "RESPONSE_FORMAT": "response_format",
+        "JSON_SCHEMA": JSON_SCHEMA_METADATA_KEY,
+        "JSON_SCHEMA_NAME": "json_schema_name",
+        "JSON_SCHEMA_STRICT": "json_schema_strict",
+    }
+
+    enabled: bool = False
+    json_schema: dict[str, Any] | None = None
+    schema_name: str = "CustomSchema"
+    strict: bool = True
+
+    @model_validator(mode="after")
+    def _schema_implies_enabled(self) -> JsonResponseConfig:
+        if self.json_schema is not None:
+            self.enabled = True
+        return self
+
+    @classmethod
+    def from_metadata(cls, *, metadata: dict[str, Any] | None) -> JsonResponseConfig:
+        """
+        Reconstruct a config from a ``MessagePiece``'s ``prompt_metadata``.
+
+        Reads the canonical ``response_format`` / ``json_schema`` keys written by ``to_metadata``.
+        Returns a disabled config when JSON output was not requested (``response_format`` absent or
+        not ``"json"``). A schema stored as a JSON string is parsed back into a dict.
+
+        Args:
+            metadata (dict[str, Any] | None): The prompt metadata to read.
+
+        Returns:
+            JsonResponseConfig: The reconstructed config (``enabled=False`` when none requested).
+
+        Raises:
+            ValueError: If a schema string is present but is not valid JSON.
+        """
+        if not metadata:
+            return cls(enabled=False)
+
+        response_format = metadata.get(cls._METADATAKEYS["RESPONSE_FORMAT"])
+        if response_format != "json":
+            return cls(enabled=False)
+
+        schema_val = metadata.get(cls._METADATAKEYS["JSON_SCHEMA"])
+        if schema_val is not None:
+            if isinstance(schema_val, str):
+                try:
+                    schema = json.loads(schema_val) if schema_val else None
+                except json.JSONDecodeError as e:
+                    raise ValueError(f"Invalid JSON schema provided: {schema_val}") from e
+            else:
+                schema = schema_val
+
+            return cls(
+                enabled=True,
+                json_schema=schema,
+                schema_name=metadata.get(cls._METADATAKEYS["JSON_SCHEMA_NAME"], "CustomSchema"),
+                strict=metadata.get(cls._METADATAKEYS["JSON_SCHEMA_STRICT"], True),
+            )
+
+        return cls(enabled=True)
+
+    def to_metadata(self) -> dict[str, Any]:
+        """
+        Serialize to the canonical ``response_format`` / ``json_schema`` metadata keys.
+
+        Symmetric with ``from_metadata``: a disabled config produces an empty dict
+        (nothing to request), an enabled config without a schema produces just the
+        ``response_format`` marker, and an enabled config with a schema also writes the
+        schema body plus its name and strict flag. The result is meant to be merged into
+        a ``MessagePiece.prompt_metadata`` dict.
+
+        Returns:
+            dict[str, Any]: The metadata fragment to attach to a piece.
+        """
+        if not self.enabled:
+            return {}
+
+        metadata: dict[str, Any] = {self._METADATAKEYS["RESPONSE_FORMAT"]: "json"}
+        if self.json_schema is not None:
+            metadata[self._METADATAKEYS["JSON_SCHEMA"]] = self.json_schema
+            metadata[self._METADATAKEYS["JSON_SCHEMA_NAME"]] = self.schema_name
+            metadata[self._METADATAKEYS["JSON_SCHEMA_STRICT"]] = self.strict
+        return metadata

pyrit/prompt_target/common/chat_completions_message_builder.py

@@ -18,10 +18,10 @@
 from pyrit.memory.storage import convert_local_image_to_data_url_async
 from pyrit.models import (
     ChatMessage,
+    JsonResponseConfig,
     Message,
     MessagePiece,
 )
-from pyrit.prompt_target.common.json_response_config import _JsonResponseConfig
 
 # Data types that render as a plain text content part.
 _TEXT_DATA_TYPES = ("text", "error")
@@ -238,12 +238,12 @@ async def build_multimodal_chat_messages_async(
     return chat_messages
 
 
-def build_response_format(*, json_config: _JsonResponseConfig) -> dict[str, Any] | None:
+def build_response_format(*, json_config: JsonResponseConfig) -> dict[str, Any] | None:
     """
     Build the ``response_format`` request parameter from a JSON response config.
 
     Args:
-        json_config (_JsonResponseConfig): The JSON response configuration derived from the
+        json_config (JsonResponseConfig): The JSON response configuration derived from the
             request metadata.
 
     Returns: