[FEAT] Support JSON Schema in Responses

doc/api.rst

@@ -331,6 +331,7 @@ API Reference
     AttackOutcome
     AttackResult
     DecomposedSeedGroup
+    JsonResponseConfig
     Message
     MessagePiece
     PromptDataType

pyrit/models/__init__.py

@@ -25,6 +25,7 @@
 )
 from pyrit.models.embeddings import EmbeddingData, EmbeddingResponse, EmbeddingSupport, EmbeddingUsageInformation
 from pyrit.models.identifiers import Identifier
+from pyrit.models.json_response_config import JsonResponseConfig
 from pyrit.models.literals import ChatMessageRole, PromptDataType, PromptResponseError
 from pyrit.models.message import (
     Message,
@@ -69,6 +70,7 @@
     "group_message_pieces_into_conversations",
     "Identifier",
     "ImagePathDataTypeSerializer",
+    "JsonResponseConfig",
     "Message",
     "MessagePiece",
     "PromptDataType",

pyrit/models/json_response_config.py

@@ -0,0 +1,48 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from typing import Any, Dict, Optional
+
+
+@dataclass
+class JsonResponseConfig:
+    """
+    Configuration for JSON responses (with OpenAI).
+    """
+
+    enabled: bool = False
+    schema: Optional[Dict[str, Any]] = None
+    schema_name: str = "CustomSchema"
+    strict: bool = True
+
+    @classmethod
+    def from_metadata(cls, *, metadata: Optional[Dict[str, Any]]) -> "JsonResponseConfig":
+        if not metadata:
+            return cls(enabled=False)
+
+        response_format = metadata.get("response_format")
+        if response_format != "json":
+            return cls(enabled=False)
+
+        schema_val = metadata.get("json_schema")
+        if schema_val:
+            if isinstance(schema_val, str):
+                try:
+                    schema = json.loads(schema_val) if schema_val else None
+                except json.JSONDecodeError:
+                    raise ValueError(f"Invalid JSON schema provided: {schema_val}")
+            else:
+                schema = schema_val
+
+            return cls(
+                enabled=True,
+                schema=schema,
+                schema_name=metadata.get("schema_name", "CustomSchema"),
+                strict=metadata.get("strict", True),
+            )
+
+        return cls(enabled=True)

pyrit/prompt_target/common/prompt_chat_target.py

@@ -4,7 +4,7 @@
 import abc
 from typing import Optional
 
-from pyrit.models import MessagePiece
+from pyrit.models import JsonResponseConfig, MessagePiece
 from pyrit.prompt_target import PromptTarget
 
 
@@ -75,16 +75,19 @@ def is_response_format_json(self, message_piece: MessagePiece) -> bool:
                 include a "response_format" key.
 
         Returns:
-            bool: True if the response format is JSON and supported, False otherwise.
+                bool: True if the response format is JSON, False otherwise.
 
         Raises:
             ValueError: If "json" response format is requested but unsupported.
         """
-        if message_piece.prompt_metadata:
-            response_format = message_piece.prompt_metadata.get("response_format")
-            if response_format == "json":
-                if not self.is_json_response_supported():
-                    target_name = self.get_identifier()["__type__"]
-                    raise ValueError(f"This target {target_name} does not support JSON response format.")
-                return True
-        return False
+        config = self.get_json_response_config(message_piece=message_piece)
+        return config.enabled
+
+    def get_json_response_config(self, *, message_piece: MessagePiece) -> JsonResponseConfig:
+        config = JsonResponseConfig.from_metadata(metadata=message_piece.prompt_metadata)
+
+        if config.enabled and not self.is_json_response_supported():
+            target_name = self.get_identifier()["__type__"]
+            raise ValueError(f"This target {target_name} does not support JSON response format.")
+
+        return config

pyrit/prompt_target/openai/openai_chat_target.py

@@ -2,7 +2,7 @@
 # Licensed under the MIT license.
 
 import logging
-from typing import Any, MutableSequence, Optional
+from typing import Any, Dict, MutableSequence, Optional
 
 from pyrit.common import convert_local_image_to_data_url
 from pyrit.exceptions import (
@@ -13,6 +13,7 @@
 from pyrit.models import (
     ChatMessage,
     ChatMessageListDictContent,
+    JsonResponseConfig,
     Message,
     MessagePiece,
     construct_response_from_request,
@@ -182,16 +183,15 @@ async def send_prompt_async(self, *, message: Message) -> list[Message]:
         self._validate_request(message=message)
 
         message_piece: MessagePiece = message.message_pieces[0]
-
-        is_json_response = self.is_response_format_json(message_piece)
+        json_config = self.get_json_response_config(message_piece=message_piece)
 
         # Get conversation from memory and append the current message
         conversation = self._memory.get_conversation(conversation_id=message_piece.conversation_id)
         conversation.append(message)
 
         logger.info(f"Sending the following prompt to the prompt target: {message}")
 
-        body = await self._construct_request_body(conversation=conversation, is_json_response=is_json_response)
+        body = await self._construct_request_body(conversation=conversation, json_config=json_config)
 
         # Use unified error handling - automatically detects ChatCompletion and validates
         response = await self._handle_openai_request(
@@ -376,8 +376,11 @@ async def _build_chat_messages_for_multi_modal_async(self, conversation: Mutable
             chat_messages.append(chat_message.model_dump(exclude_none=True))
         return chat_messages
 
-    async def _construct_request_body(self, conversation: MutableSequence[Message], is_json_response: bool) -> dict:
+    async def _construct_request_body(
+        self, *, conversation: MutableSequence[Message], json_config: JsonResponseConfig
+    ) -> dict:
         messages = await self._build_chat_messages_async(conversation)
+        response_format = self._build_response_format(json_config)
 
         body_parameters = {
             "model": self._model_name,
@@ -391,7 +394,7 @@ async def _construct_request_body(self, conversation: MutableSequence[Message],
             "seed": self._seed,
             "n": self._n,
             "messages": messages,
-            "response_format": {"type": "json_object"} if is_json_response else None,
+            "response_format": response_format,
         }
 
         if self._extra_body_parameters:
@@ -419,3 +422,19 @@ def _validate_request(self, *, message: Message) -> None:
         for prompt_data_type in converted_prompt_data_types:
             if prompt_data_type not in ["text", "image_path"]:
                 raise ValueError(f"This target only supports text and image_path. Received: {prompt_data_type}.")
+
+    def _build_response_format(self, json_config: JsonResponseConfig) -> Optional[Dict[str, Any]]:
+        if not json_config.enabled:
+            return None
+
+        if json_config.schema:
+            return {
+                "type": "json_schema",
+                "json_schema": {
+                    "name": json_config.schema_name,
+                    "schema": json_config.schema,
+                    "strict": json_config.strict,
+                },
+            }
+
+        return {"type": "json_object"}

pyrit/prompt_target/openai/openai_response_target.py

@@ -21,6 +21,7 @@
     pyrit_target_retry,
 )
 from pyrit.models import (
+    JsonResponseConfig,
     Message,
     MessagePiece,
     PromptDataType,
@@ -313,7 +314,9 @@ async def _build_input_for_multi_modal_async(self, conversation: MutableSequence
 
         return input_items
 
-    async def _construct_request_body(self, conversation: MutableSequence[Message], is_json_response: bool) -> dict:
+    async def _construct_request_body(
+        self, *, conversation: MutableSequence[Message], json_config: JsonResponseConfig
+    ) -> dict:
         """
         Construct the request body to send to the Responses API.
 
@@ -322,6 +325,8 @@ async def _construct_request_body(self, conversation: MutableSequence[Message],
         """
         input_items = await self._build_input_for_multi_modal_async(conversation)
 
+        text_format = self._build_text_format(json_config=json_config)
+
         body_parameters = {
             "model": self._model_name,
             "max_output_tokens": self._max_output_tokens,
@@ -330,7 +335,7 @@ async def _construct_request_body(self, conversation: MutableSequence[Message],
             "stream": False,
             "input": input_items,
             # Correct JSON response format per Responses API
-            "response_format": {"type": "json_object"} if is_json_response else None,
+            "text": text_format,
         }
 
         if self._extra_body_parameters:
@@ -339,6 +344,23 @@ async def _construct_request_body(self, conversation: MutableSequence[Message],
         # Filter out None values
         return {k: v for k, v in body_parameters.items() if v is not None}
 
+    def _build_text_format(self, json_config: JsonResponseConfig) -> Optional[Dict[str, Any]]:
+        if not json_config.enabled:
+            return None
+
+        if json_config.schema:
+            return {
+                "format": {
+                    "type": "json_schema",
+                    "name": json_config.schema_name,
+                    "schema": json_config.schema,
+                    "strict": json_config.strict,
+                }
+            }
+
+        logger.info("Using json_object format without schema - consider providing a schema for better results")
+        return {"format": {"type": "json_object"}}
+
     def _check_content_filter(self, response: Any) -> bool:
         """
         Check if a Response API response has a content filter error.
@@ -436,7 +458,10 @@ async def send_prompt_async(self, *, message: Message) -> list[Message]:
         self._validate_request(message=message)
 
         message_piece: MessagePiece = message.message_pieces[0]
-        is_json_response = self.is_response_format_json(message_piece)
+        json_config = JsonResponseConfig(enabled=False)
+        if message.message_pieces:
+            last_piece = message.message_pieces[-1]
+            json_config = self.get_json_response_config(message_piece=last_piece)
 
         # Get full conversation history from memory and append the current message
         conversation: MutableSequence[Message] = self._memory.get_conversation(
@@ -453,7 +478,7 @@ async def send_prompt_async(self, *, message: Message) -> list[Message]:
         while True:
             logger.info(f"Sending conversation with {len(conversation)} messages to the prompt target")
 
-            body = await self._construct_request_body(conversation=conversation, is_json_response=is_json_response)
+            body = await self._construct_request_body(conversation=conversation, json_config=json_config)
 
             # Use unified error handling - automatically detects Response and validates
             result = await self._handle_openai_request(