[FEAT] Support JSON Schema in Responses

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,
@@ -68,6 +69,7 @@
     "group_message_pieces_into_conversations",
     "Identifier",
     "ImagePathDataTypeSerializer",
+    "JsonResponseConfig",
     "Message",
     "MessagePiece",
     "PromptDataType",

pyrit/models/json_response_config.py

@@ -0,0 +1,44 @@
+# 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:
+    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

@@ -3,7 +3,7 @@
 
 import json
 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 (
@@ -14,6 +14,7 @@
 from pyrit.models import (
     ChatMessage,
     ChatMessageListDictContent,
+    JsonResponseConfig,
     Message,
     MessagePiece,
     construct_response_from_request,
@@ -243,8 +244,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,
@@ -258,7 +262,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:
@@ -274,7 +278,6 @@ def _construct_message_from_openai_json(
         open_ai_str_response: str,
         message_piece: MessagePiece,
     ) -> Message:
-
         try:
             response = json.loads(open_ai_str_response)
         except json.JSONDecodeError as e:
@@ -320,3 +323,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_chat_target_base.py

@@ -15,6 +15,7 @@
 )
 from pyrit.exceptions.exception_classes import RateLimitException
 from pyrit.models import (
+    JsonResponseConfig,
     Message,
     MessagePiece,
 )
@@ -82,9 +83,9 @@ def __init__(
         super().__init__(**kwargs)
 
         if temperature is not None and (temperature < 0 or temperature > 2):
-            raise PyritException("temperature must be between 0 and 2 (inclusive).")
+            raise PyritException(message="temperature must be between 0 and 2 (inclusive).")
         if top_p is not None and (top_p < 0 or top_p > 1):
-            raise PyritException("top_p must be between 0 and 1 (inclusive).")
+            raise PyritException(message="top_p must be between 0 and 1 (inclusive).")
 
         self._temperature = temperature
         self._top_p = top_p
@@ -108,14 +109,14 @@ async def send_prompt_async(self, *, message: Message) -> Message:
 
         message_piece: MessagePiece = message.message_pieces[0]
 
-        is_json_response = self.is_response_format_json(message_piece)
+        json_response_config = self.get_json_response_config(message_piece=message_piece)
 
         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_response_config)
 
         try:
             str_response: httpx.Response = await net_utility.make_request_and_raise_if_error_async(
@@ -157,7 +158,9 @@ async def send_prompt_async(self, *, message: Message) -> Message:
 
         return response
 
-    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:
         raise NotImplementedError
 
     def _construct_message_from_openai_json(

pyrit/prompt_target/openai/openai_response_target.py

@@ -21,6 +21,7 @@
     handle_bad_request_exception,
 )
 from pyrit.models import (
+    JsonResponseConfig,
     Message,
     MessagePiece,
     PromptDataType,
@@ -291,7 +292,9 @@ def _translate_roles(self, conversation: List[Dict[str, Any]]) -> None:
                 request["role"] = "developer"
         return
 
-    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.
 
@@ -300,6 +303,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,
@@ -308,7 +313,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:
@@ -317,6 +322,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 _construct_message_from_openai_json(
         self,
         *,

tests/integration/targets/test_openai_responses_gpt5.py

@@ -2,24 +2,30 @@
 # Licensed under the MIT license.
 
 
+import json
 import os
 import uuid
 
+import jsonschema
 import pytest
 
 from pyrit.models import MessagePiece
 from pyrit.prompt_target import OpenAIResponseTarget
 
 
-@pytest.mark.asyncio
-async def test_openai_responses_gpt5(sqlite_instance):
-    args = {
+@pytest.fixture()
+def gpt5_args():
+    return {
         "endpoint": os.getenv("AZURE_OPENAI_GPT5_RESPONSES_ENDPOINT"),
         "model_name": os.getenv("AZURE_OPENAI_GPT5_MODEL"),
-        "api_key": os.getenv("AZURE_OPENAI_GPT5_KEY"),
+        # "api_key": os.getenv("AZURE_OPENAI_GPT5_KEY"),
+        "use_entra_auth": True,
     }
 
-    target = OpenAIResponseTarget(**args)
+
+@pytest.mark.asyncio
+async def test_openai_responses_gpt5(sqlite_instance, gpt5_args):
+    target = OpenAIResponseTarget(**gpt5_args)
 
     conv_id = str(uuid.uuid4())
 
@@ -46,3 +52,86 @@ async def test_openai_responses_gpt5(sqlite_instance):
     assert result.message_pieces[1].role == "assistant"
     # Hope that the model manages to give the correct answer somewhere (GPT-5 really should)
     assert "Paris" in result.message_pieces[1].converted_value
+
+
+@pytest.mark.asyncio
+async def test_openai_responses_gpt5_json_schema(sqlite_instance, gpt5_args):
+    target = OpenAIResponseTarget(**gpt5_args)
+
+    conv_id = str(uuid.uuid4())
+
+    developer_piece = MessagePiece(
+        role="developer",
+        original_value="You are an expert in the lore of cats.",
+        original_value_data_type="text",
+        conversation_id=conv_id,
+        attack_identifier={"id": str(uuid.uuid4())},
+    )
+    sqlite_instance.add_message_to_memory(request=developer_piece.to_message())
+
+    cat_schema = {
+        "type": "object",
+        "properties": {
+            "name": {"type": "string", "minLength": 12},
+            "age": {"type": "integer", "minimum": 0, "maximum": 20},
+            "colour": {
+                "type": "array",
+                "items": {"type": "integer", "minimum": 0, "maximum": 255},
+                "minItems": 3,
+                "maxItems": 3,
+            },
+        },
+        "required": ["name", "age", "colour"],
+        "additionalProperties": False,
+    }
+
+    prompt = "Create a JSON object that describes a mystical cat "
+    prompt += "with the following properties: name, age, colour."
+
+    user_piece = MessagePiece(
+        role="user",
+        original_value=prompt,
+        original_value_data_type="text",
+        conversation_id=conv_id,
+        prompt_metadata={"response_format": "json", "json_schema": json.dumps(cat_schema)},
+    )
+
+    response = await target.send_prompt_async(message=user_piece.to_message())
+
+    response_content = response.get_value(1)
+    response_json = json.loads(response_content)
+    jsonschema.validate(instance=response_json, schema=cat_schema)
+
+
+@pytest.mark.asyncio
+async def test_openai_responses_gpt5_json_object(sqlite_instance, gpt5_args):
+    target = OpenAIResponseTarget(**gpt5_args)
+
+    conv_id = str(uuid.uuid4())
+
+    developer_piece = MessagePiece(
+        role="developer",
+        original_value="You are an expert in the lore of cats.",
+        original_value_data_type="text",
+        conversation_id=conv_id,
+        attack_identifier={"id": str(uuid.uuid4())},
+    )
+
+    sqlite_instance.add_message_to_memory(request=developer_piece.to_message())
+
+    prompt = "Create a JSON object that describes a mystical cat "
+    prompt += "with the following properties: name, age, colour."
+
+    user_piece = MessagePiece(
+        role="user",
+        original_value=prompt,
+        original_value_data_type="text",
+        conversation_id=conv_id,
+        prompt_metadata={"response_format": "json"},
+    )
+    response = await target.send_prompt_async(message=user_piece.to_message())
+
+    response_content = response.get_value(1)
+    response_json = json.loads(response_content)
+    assert response_json is not None
+    # Can't assert more, since the failure could be due to a bad generation by the model