diff --git a/build_scripts/check_links.py b/build_scripts/check_links.py index 4cb5a6a14..342d57df1 100644 --- a/build_scripts/check_links.py +++ b/build_scripts/check_links.py @@ -17,6 +17,7 @@ "https://platform.openai.com/docs/api-reference/introduction", # blocks python requests "https://platform.openai.com/docs/api-reference/responses", # blocks python requests "https://platform.openai.com/docs/guides/function-calling", # blocks python requests + "https://platform.openai.com/docs/guides/structured-outputs", # blocks python requests "https://www.anthropic.com/research/many-shot-jailbreaking", # blocks python requests "https://code.visualstudio.com/docs/devcontainers/containers", "https://stackoverflow.com/questions/77134272/pip-install-dev-with-pyproject-toml-not-working", diff --git a/doc/code/targets/1_openai_chat_target.ipynb b/doc/code/targets/1_openai_chat_target.ipynb index 60cda5488..f8a1f09a5 100644 --- a/doc/code/targets/1_openai_chat_target.ipynb +++ b/doc/code/targets/1_openai_chat_target.ipynb @@ -20,6 +20,15 @@ "id": "1", "metadata": {}, "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Found default environment files: ['/home/vscode/.pyrit/.env', '/home/vscode/.pyrit/.env.local']\n", + "Loaded environment file: /home/vscode/.pyrit/.env\n", + "Loaded environment file: /home/vscode/.pyrit/.env.local\n" + ] + }, { "name": "stdout", "output_type": "stream", @@ -88,6 +97,91 @@ "cell_type": "markdown", "id": "2", "metadata": {}, + "source": [ + "## JSON Output\n", + "\n", + "You can also get the output in JSON format for further processing or storage. In this example, we define a simple JSON schema that describes a person with `name` and `age` properties.\n", + "\n", + "For more information about structured outputs with OpenAI, see [the OpenAI documentation](https://platform.openai.com/docs/guides/structured-outputs)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Found default environment files: ['/home/vscode/.pyrit/.env', '/home/vscode/.pyrit/.env.local']\n", + "Loaded environment file: /home/vscode/.pyrit/.env\n", + "Loaded environment file: /home/vscode/.pyrit/.env.local\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "{\n", + " \"name\": \"Bob\",\n", + " \"age\": 32\n", + "}\n" + ] + } + ], + "source": [ + "import json\n", + "\n", + "import jsonschema\n", + "\n", + "from pyrit.models import Message, MessagePiece\n", + "from pyrit.prompt_target import OpenAIChatTarget\n", + "from pyrit.setup import IN_MEMORY, initialize_pyrit_async\n", + "\n", + "await initialize_pyrit_async(memory_db_type=IN_MEMORY) # type: ignore\n", + "\n", + "# Define a simple JSON schema for a person\n", + "person_schema = {\n", + " \"type\": \"object\",\n", + " \"properties\": {\n", + " \"name\": {\"type\": \"string\"},\n", + " \"age\": {\"type\": \"integer\", \"minimum\": 0, \"maximum\": 150},\n", + " },\n", + " \"required\": [\"name\", \"age\"],\n", + " \"additionalProperties\": False,\n", + "}\n", + "\n", + "prompt = \"Create a JSON object describing a person named Bob who is 32 years old.\"\n", + "# Create the message piece and message\n", + "message_piece = MessagePiece(\n", + " role=\"user\",\n", + " original_value=prompt,\n", + " original_value_data_type=\"text\",\n", + " prompt_metadata={\n", + " \"response_format\": \"json\",\n", + " \"json_schema\": json.dumps(person_schema),\n", + " },\n", + ")\n", + "message = Message(message_pieces=[message_piece])\n", + "\n", + "# Create the OpenAI Chat target\n", + "target = OpenAIChatTarget()\n", + "\n", + "# Send the prompt, requesting JSON output\n", + "response = await target.send_prompt_async(message=message) # type: ignore\n", + "\n", + "# Validate and print the response\n", + "response_json = json.loads(response[0].message_pieces[0].converted_value)\n", + "print(json.dumps(response_json, indent=2))\n", + "jsonschema.validate(instance=response_json, schema=person_schema)" + ] + }, + { + "cell_type": "markdown", + "id": "4", + "metadata": {}, "source": [ "## OpenAI Configuration\n", "\n", @@ -125,7 +219,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.11.13" + "version": "3.11.14" } }, "nbformat": 4, diff --git a/doc/code/targets/1_openai_chat_target.py b/doc/code/targets/1_openai_chat_target.py index 9acf4b65e..ea5f0da87 100644 --- a/doc/code/targets/1_openai_chat_target.py +++ b/doc/code/targets/1_openai_chat_target.py @@ -44,6 +44,59 @@ result = await attack.execute_async(objective=jailbreak_prompt) # type: ignore await ConsoleAttackResultPrinter().print_conversation_async(result=result) # type: ignore +# %% [markdown] +# ## JSON Output + +# You can also get the output in JSON format for further processing or storage. In this example, we define a simple JSON schema that describes a person with `name` and `age` properties. +# +# For more information about structured outputs with OpenAI, see [the OpenAI documentation](https://platform.openai.com/docs/guides/structured-outputs). + +# %% +import json + +import jsonschema + +from pyrit.models import Message, MessagePiece +from pyrit.prompt_target import OpenAIChatTarget +from pyrit.setup import IN_MEMORY, initialize_pyrit_async + +await initialize_pyrit_async(memory_db_type=IN_MEMORY) # type: ignore + +# Define a simple JSON schema for a person +person_schema = { + "type": "object", + "properties": { + "name": {"type": "string"}, + "age": {"type": "integer", "minimum": 0, "maximum": 150}, + }, + "required": ["name", "age"], + "additionalProperties": False, +} + +prompt = "Create a JSON object describing a person named Bob who is 32 years old." +# Create the message piece and message +message_piece = MessagePiece( + role="user", + original_value=prompt, + original_value_data_type="text", + prompt_metadata={ + "response_format": "json", + "json_schema": json.dumps(person_schema), + }, +) +message = Message(message_pieces=[message_piece]) + +# Create the OpenAI Chat target +target = OpenAIChatTarget() + +# Send the prompt, requesting JSON output +response = await target.send_prompt_async(message=message) # type: ignore + +# Validate and print the response +response_json = json.loads(response[0].message_pieces[0].converted_value) +print(json.dumps(response_json, indent=2)) +jsonschema.validate(instance=response_json, schema=person_schema) + # %% [markdown] # ## OpenAI Configuration # diff --git a/doc/code/targets/8_openai_responses_target.ipynb b/doc/code/targets/8_openai_responses_target.ipynb index 90b80db95..b2e1c12e4 100644 --- a/doc/code/targets/8_openai_responses_target.ipynb +++ b/doc/code/targets/8_openai_responses_target.ipynb @@ -27,6 +27,15 @@ "id": "1", "metadata": {}, "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Found default environment files: ['/home/vscode/.pyrit/.env', '/home/vscode/.pyrit/.env.local']\n", + "Loaded environment file: /home/vscode/.pyrit/.env\n", + "Loaded environment file: /home/vscode/.pyrit/.env.local\n" + ] + }, { "name": "stdout", "output_type": "stream", @@ -40,8 +49,7 @@ "\u001b[33m────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\n", "\u001b[1m\u001b[33m🔸 ASSISTANT\u001b[0m\n", "\u001b[33m────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\n", - "\u001b[33m Why don’t scientists trust atoms?\u001b[0m\n", - "\u001b[33m Because they make up everything!\u001b[0m\n", + "\u001b[33m Why did the coffee file a police report? It got mugged!\u001b[0m\n", "\n", "\u001b[34m────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\n" ] @@ -74,6 +82,91 @@ "cell_type": "markdown", "id": "2", "metadata": {}, + "source": [ + "## JSON Generation\n", + "\n", + "We can use the OpenAI `Responses API` with a JSON schema to produce structured JSON output. In this example, we define a simple JSON schema that describes a person with `name` and `age` properties.\n", + "\n", + "For more information about structured outputs with OpenAI, see [the OpenAI documentation](https://platform.openai.com/docs/guides/structured-outputs)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Found default environment files: ['/home/vscode/.pyrit/.env', '/home/vscode/.pyrit/.env.local']\n", + "Loaded environment file: /home/vscode/.pyrit/.env\n", + "Loaded environment file: /home/vscode/.pyrit/.env.local\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "{\n", + " \"name\": \"Alice\",\n", + " \"age\": 30\n", + "}\n" + ] + } + ], + "source": [ + "import json\n", + "\n", + "import jsonschema\n", + "\n", + "from pyrit.models import Message, MessagePiece\n", + "from pyrit.prompt_target import OpenAIResponseTarget\n", + "from pyrit.setup import IN_MEMORY, initialize_pyrit_async\n", + "\n", + "await initialize_pyrit_async(memory_db_type=IN_MEMORY) # type: ignore\n", + "\n", + "# Define a simple JSON schema for a person\n", + "person_schema = {\n", + " \"type\": \"object\",\n", + " \"properties\": {\n", + " \"name\": {\"type\": \"string\"},\n", + " \"age\": {\"type\": \"integer\", \"minimum\": 0, \"maximum\": 150},\n", + " },\n", + " \"required\": [\"name\", \"age\"],\n", + " \"additionalProperties\": False,\n", + "}\n", + "\n", + "prompt = \"Create a JSON object describing a person named Alice who is 30 years old.\"\n", + "# Create the message piece and message\n", + "message_piece = MessagePiece(\n", + " role=\"user\",\n", + " original_value=prompt,\n", + " original_value_data_type=\"text\",\n", + " prompt_metadata={\n", + " \"response_format\": \"json\",\n", + " \"json_schema\": json.dumps(person_schema),\n", + " },\n", + ")\n", + "message = Message(message_pieces=[message_piece])\n", + "\n", + "# Create the OpenAI Responses target\n", + "target = OpenAIResponseTarget()\n", + "\n", + "# Send the prompt, requesting JSON output\n", + "response = await target.send_prompt_async(message=message) # type: ignore\n", + "\n", + "# Validate and print the response\n", + "response_json = json.loads(response[0].message_pieces[1].converted_value)\n", + "print(json.dumps(response_json, indent=2))\n", + "jsonschema.validate(instance=response_json, schema=person_schema)" + ] + }, + { + "cell_type": "markdown", + "id": "4", + "metadata": {}, "source": [ "## Tool Use with Custom Functions\n", "\n", @@ -94,17 +187,26 @@ { "cell_type": "code", "execution_count": null, - "id": "3", + "id": "5", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ - "0 | assistant: {\"id\":\"rs_069fc3b631d0d80400693901c520348194ad74ffcc1036e950\",\"type\":\"reasoning\",\"summary\":[],\"content\":null,\"encrypted_content\":null}\n", - "1 | assistant: {\"type\":\"function_call\",\"call_id\":\"call_20XlbZBvvIULm3xn6fVFxtJV\",\"name\":\"get_current_weather\",\"arguments\":\"{\\\"location\\\":\\\"Boston, MA\\\",\\\"unit\\\":\\\"celsius\\\"}\"}\n", - "0 | tool: {\"type\":\"function_call_output\",\"call_id\":\"call_20XlbZBvvIULm3xn6fVFxtJV\",\"output\":\"{\\\"weather\\\":\\\"Sunny\\\",\\\"temp_c\\\":22,\\\"location\\\":\\\"Boston, MA\\\",\\\"unit\\\":\\\"celsius\\\"}\"}\n", - "0 | assistant: The current weather in Boston, MA is Sunny with a temperature of 22°C.\n" + "Found default environment files: ['/home/vscode/.pyrit/.env', '/home/vscode/.pyrit/.env.local']\n", + "Loaded environment file: /home/vscode/.pyrit/.env\n", + "Loaded environment file: /home/vscode/.pyrit/.env.local\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "0 | assistant: {\"id\":\"rs_02cb1830dae3b0d20069541e7716a8819488f35e15b863919d\",\"type\":\"reasoning\",\"summary\":[],\"content\":null,\"encrypted_content\":null}\n", + "1 | assistant: {\"type\":\"function_call\",\"call_id\":\"call_5ooC2LwwJaPlwfFOWkc7uBpm\",\"name\":\"get_current_weather\",\"arguments\":\"{\\\"location\\\":\\\"Boston\\\",\\\"unit\\\":\\\"celsius\\\"}\"}\n", + "0 | tool: {\"type\":\"function_call_output\",\"call_id\":\"call_5ooC2LwwJaPlwfFOWkc7uBpm\",\"output\":\"{\\\"weather\\\":\\\"Sunny\\\",\\\"temp_c\\\":22,\\\"location\\\":\\\"Boston\\\",\\\"unit\\\":\\\"celsius\\\"}\"}\n", + "0 | assistant: The current weather in Boston is Sunny with a temperature of 22°C.\n" ] } ], @@ -168,7 +270,7 @@ }, { "cell_type": "markdown", - "id": "4", + "id": "6", "metadata": {}, "source": [ "## Using the Built-in Web Search Tool\n", @@ -187,15 +289,24 @@ { "cell_type": "code", "execution_count": null, - "id": "5", + "id": "7", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ - "0 | assistant: {\"type\":\"web_search_call\",\"id\":\"ws_0cbecd46b8764f3f00693901ccd09081949a25da86b63d3e54\"}\n", - "1 | assistant: Today, a positive news story comes from India, where conservationists reported a record-breaking 1 million sea turtle nests counted on the country's coast. This number is about 10 times higher than in previous decades and signals strong progress for this endangered species—thanks to dedicated environmental efforts and improved protection measures for nesting sites [Good News Network](https://www.goodnewsnetwork.org/1-million-turtle-nests-counted-on-indias-coast-crazy-high-number-is-10x-more-than-decades-ago/).\n" + "Found default environment files: ['/home/vscode/.pyrit/.env', '/home/vscode/.pyrit/.env.local']\n", + "Loaded environment file: /home/vscode/.pyrit/.env\n", + "Loaded environment file: /home/vscode/.pyrit/.env.local\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "0 | assistant: {\"type\":\"web_search_call\",\"id\":\"ws_0a650bbbf4434ce60069541e7dbc008190b9ead0fb3cae7734\"}\n", + "1 | assistant: One positive news story from today is that the world’s first international treaty to protect the high seas—a critical area for marine biodiversity—has now been ratified by over 79 nations and will become legally binding in 2026. This landmark agreement will provide a new legal framework to preserve marine life in waters outside national boundaries, representing a major victory for global ocean protection efforts and environmental advocates worldwide [Positive News](https://www.positive.news/society/what-went-right-in-2025-the-good-news-that-mattered/).\n" ] } ], @@ -235,7 +346,7 @@ }, { "cell_type": "markdown", - "id": "6", + "id": "8", "metadata": {}, "source": [ "## Grammar-Constrained Generation\n", @@ -250,20 +361,29 @@ { "cell_type": "code", "execution_count": null, - "id": "7", + "id": "9", "metadata": {}, "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Found default environment files: ['/home/vscode/.pyrit/.env', '/home/vscode/.pyrit/.env.local']\n", + "Loaded environment file: /home/vscode/.pyrit/.env\n", + "Loaded environment file: /home/vscode/.pyrit/.env.local\n" + ] + }, { "name": "stdout", "output_type": "stream", "text": [ "Unconstrained Response:\n", - "0 | assistant: {\"id\":\"rs_054a76346e44491900693901d190988193a4245c6c7dc37d40\",\"type\":\"reasoning\",\"summary\":[],\"content\":null,\"encrypted_content\":null}\n", + "0 | assistant: {\"id\":\"rs_0d3f19062cddb30c0069541e82616c8190a7d9dfee276f9c92\",\"type\":\"reasoning\",\"summary\":[],\"content\":null,\"encrypted_content\":null}\n", "1 | assistant: Rome.\n", "\n", "Constrained Response:\n", - "0 | assistant: {\"id\":\"rs_087367e8d137426000693901d587108195a7e06bc9b06bf9d2\",\"type\":\"reasoning\",\"summary\":[],\"content\":null,\"encrypted_content\":null}\n", - "1 | assistant: I think that it is b\n" + "0 | assistant: {\"id\":\"rs_0ec76b80c50f2b190069541e8603e08194aa02993767062a00\",\"type\":\"reasoning\",\"summary\":[],\"content\":null,\"encrypted_content\":null}\n", + "1 | assistant: I think that it is Pisa\n" ] } ], @@ -330,9 +450,6 @@ } ], "metadata": { - "jupytext": { - "main_language": "python" - }, "language_info": { "codemirror_mode": { "name": "ipython", @@ -343,7 +460,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.11.13" + "version": "3.11.14" } }, "nbformat": 4, diff --git a/doc/code/targets/8_openai_responses_target.py b/doc/code/targets/8_openai_responses_target.py index 1955940d2..d08c0cd33 100644 --- a/doc/code/targets/8_openai_responses_target.py +++ b/doc/code/targets/8_openai_responses_target.py @@ -46,6 +46,59 @@ result = await attack.execute_async(objective="Tell me a joke") # type: ignore await ConsoleAttackResultPrinter().print_conversation_async(result=result) # type: ignore +# %% [markdown] +# ## JSON Generation +# +# We can use the OpenAI `Responses API` with a JSON schema to produce structured JSON output. In this example, we define a simple JSON schema that describes a person with `name` and `age` properties. +# +# For more information about structured outputs with OpenAI, see [the OpenAI documentation](https://platform.openai.com/docs/guides/structured-outputs). + +# %% +import json + +import jsonschema + +from pyrit.models import Message, MessagePiece +from pyrit.prompt_target import OpenAIResponseTarget +from pyrit.setup import IN_MEMORY, initialize_pyrit_async + +await initialize_pyrit_async(memory_db_type=IN_MEMORY) # type: ignore + +# Define a simple JSON schema for a person +person_schema = { + "type": "object", + "properties": { + "name": {"type": "string"}, + "age": {"type": "integer", "minimum": 0, "maximum": 150}, + }, + "required": ["name", "age"], + "additionalProperties": False, +} + +prompt = "Create a JSON object describing a person named Alice who is 30 years old." +# Create the message piece and message +message_piece = MessagePiece( + role="user", + original_value=prompt, + original_value_data_type="text", + prompt_metadata={ + "response_format": "json", + "json_schema": json.dumps(person_schema), + }, +) +message = Message(message_pieces=[message_piece]) + +# Create the OpenAI Responses target +target = OpenAIResponseTarget() + +# Send the prompt, requesting JSON output +response = await target.send_prompt_async(message=message) # type: ignore + +# Validate and print the response +response_json = json.loads(response[0].message_pieces[1].converted_value) +print(json.dumps(response_json, indent=2)) +jsonschema.validate(instance=response_json, schema=person_schema) + # %% [markdown] # ## Tool Use with Custom Functions # diff --git a/pyrit/models/json_response_config.py b/pyrit/models/json_response_config.py new file mode 100644 index 000000000..f2ed20032 --- /dev/null +++ b/pyrit/models/json_response_config.py @@ -0,0 +1,61 @@ +# 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 + +# Would prefer StrEnum, but.... Python 3.10 +_METADATAKEYS = { + "RESPONSE_FORMAT": "response_format", + "JSON_SCHEMA": "json_schema", + "JSON_SCHEMA_NAME": "json_schema_name", + "JSON_SCHEMA_STRICT": "json_schema_strict", +} + + +@dataclass +class _JsonResponseConfig: + """ + Configuration for JSON responses (with OpenAI). + + For more details, 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 + """ + + 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(_METADATAKEYS["RESPONSE_FORMAT"]) + if response_format != "json": + return cls(enabled=False) + + schema_val = metadata.get(_METADATAKEYS["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(_METADATAKEYS["JSON_SCHEMA_NAME"], "CustomSchema"), + strict=metadata.get(_METADATAKEYS["JSON_SCHEMA_STRICT"], True), + ) + + return cls(enabled=True) diff --git a/pyrit/prompt_target/common/prompt_chat_target.py b/pyrit/prompt_target/common/prompt_chat_target.py index dcaef747a..823fc2c7b 100644 --- a/pyrit/prompt_target/common/prompt_chat_target.py +++ b/pyrit/prompt_target/common/prompt_chat_target.py @@ -5,6 +5,7 @@ from typing import Optional from pyrit.models import MessagePiece +from pyrit.models.json_response_config import _JsonResponseConfig from pyrit.prompt_target import PromptTarget @@ -95,16 +96,32 @@ 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: + """ + Get the JSON response configuration from the message piece metadata. + + Args: + message_piece: A MessagePiece object with a `prompt_metadata` dictionary that may + include JSON response configuration. + + Returns: + _JsonResponseConfig: The JSON response configuration. + + Raises: + ValueError: If JSON response format is requested but unsupported. + """ + 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 diff --git a/pyrit/prompt_target/openai/openai_chat_target.py b/pyrit/prompt_target/openai/openai_chat_target.py index 501fedd0b..827c62ba3 100644 --- a/pyrit/prompt_target/openai/openai_chat_target.py +++ b/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 ( @@ -17,6 +17,7 @@ MessagePiece, construct_response_from_request, ) +from pyrit.models.json_response_config import _JsonResponseConfig from pyrit.prompt_target import ( OpenAITarget, PromptChatTarget, @@ -186,8 +187,7 @@ 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) @@ -195,7 +195,7 @@ async def send_prompt_async(self, *, message: Message) -> list[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( @@ -393,8 +393,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, @@ -408,7 +411,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: @@ -436,3 +439,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"} diff --git a/pyrit/prompt_target/openai/openai_response_target.py b/pyrit/prompt_target/openai/openai_response_target.py index e8e2f1131..b94d1c6ab 100644 --- a/pyrit/prompt_target/openai/openai_response_target.py +++ b/pyrit/prompt_target/openai/openai_response_target.py @@ -26,6 +26,7 @@ PromptDataType, PromptResponseError, ) +from pyrit.models.json_response_config import _JsonResponseConfig from pyrit.prompt_target import ( OpenAITarget, PromptChatTarget, @@ -316,7 +317,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. @@ -325,13 +328,15 @@ async def _construct_request_body(self, conversation: MutableSequence[Message], Args: conversation: The full conversation history. - is_json_response: Whether the response should be formatted as JSON. + json_config: Specification for JSON formatting. Returns: dict: The request body to send to the Responses API. """ 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, @@ -340,7 +345,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: @@ -349,6 +354,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. @@ -446,7 +468,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( @@ -463,7 +488,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( diff --git a/tests/integration/targets/test_openai_responses_gpt5.py b/tests/integration/targets/test_openai_responses_gpt5.py index 08f0b52a7..82f56bc8c 100644 --- a/tests/integration/targets/test_openai_responses_gpt5.py +++ b/tests/integration/targets/test_openai_responses_gpt5.py @@ -2,24 +2,32 @@ # Licensed under the MIT license. +import json import os import uuid +import jsonschema import pytest +# from pyrit.auth import get_azure_openai_auth from pyrit.models import MessagePiece from pyrit.prompt_target import OpenAIResponseTarget -@pytest.mark.asyncio -async def test_openai_responses_gpt5(sqlite_instance): - args = { - "endpoint": os.getenv("AZURE_OPENAI_GPT5_RESPONSES_ENDPOINT"), +@pytest.fixture() +def gpt5_args(): + endpoint_value = os.environ["AZURE_OPENAI_GPT5_RESPONSES_ENDPOINT"] + return { + "endpoint": endpoint_value, "model_name": os.getenv("AZURE_OPENAI_GPT5_MODEL"), "api_key": os.getenv("AZURE_OPENAI_GPT5_KEY"), + # "api_key": get_azure_openai_auth(endpoint_value), } - 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()) @@ -47,3 +55,91 @@ async def test_openai_responses_gpt5(sqlite_instance): assert result[0].message_pieces[1].role == "assistant" # Hope that the model manages to give the correct answer somewhere (GPT-5 really should) assert "Paris" in result[0].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}, + "fur_rgb": { + "type": "array", + "items": {"type": "integer", "minimum": 0, "maximum": 255}, + "minItems": 3, + "maxItems": 3, + }, + }, + "required": ["name", "age", "fur_rgb"], + "additionalProperties": False, + } + + prompt = "Create a JSON object that describes a mystical cat " + prompt += "with the following properties: name, age, fur_rgb." + + 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()) + + assert len(response) == 1 + assert len(response[0].message_pieces) == 2 + response_piece = response[0].message_pieces[1] + assert response_piece.role == "assistant" + response_json = json.loads(response_piece.converted_value) + 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, fur_rgb." + + 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()) + + assert len(response) == 1 + assert len(response[0].message_pieces) == 2 + response_piece = response[0].message_pieces[1] + assert response_piece.role == "assistant" + _ = json.loads(response_piece.converted_value) + # Can't assert more, since the failure could be due to a bad generation by the model diff --git a/tests/unit/models/test_json_response_config.py b/tests/unit/models/test_json_response_config.py new file mode 100644 index 000000000..d91c4bd54 --- /dev/null +++ b/tests/unit/models/test_json_response_config.py @@ -0,0 +1,76 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import json + +import pytest + +from pyrit.models.json_response_config import _JsonResponseConfig + + +def test_with_none(): + config = _JsonResponseConfig.from_metadata(metadata=None) + assert config.enabled is False + assert config.schema is None + assert config.schema_name == "CustomSchema" + assert config.strict is True + + +def test_with_json_object(): + metadata = { + "response_format": "json", + } + config = _JsonResponseConfig.from_metadata(metadata=metadata) + assert config.enabled is True + assert config.schema is None + assert config.schema_name == "CustomSchema" + assert config.strict is True + + +def test_with_json_string_schema(): + schema = {"type": "object", "properties": {"name": {"type": "string"}}} + metadata = { + "response_format": "json", + "json_schema": json.dumps(schema), + "json_schema_name": "TestSchema", + "json_schema_strict": False, + } + config = _JsonResponseConfig.from_metadata(metadata=metadata) + assert config.enabled is True + assert config.schema == schema + assert config.schema_name == "TestSchema" + assert config.strict is False + + +def test_with_json_schema_object(): + schema = {"type": "object", "properties": {"age": {"type": "integer"}}} + metadata = { + "response_format": "json", + "json_schema": schema, + } + config = _JsonResponseConfig.from_metadata(metadata=metadata) + assert config.enabled is True + assert config.schema == schema + assert config.schema_name == "CustomSchema" + assert config.strict is True + + +def test_with_invalid_json_schema_string(): + metadata = { + "response_format": "json", + "json_schema": "{invalid_json: true}", + } + with pytest.raises(ValueError) as e: + _JsonResponseConfig.from_metadata(metadata=metadata) + assert "Invalid JSON schema provided" in str(e.value) + + +def test_other_response_format(): + metadata = { + "response_format": "something_really_improbably_to_have_here", + } + config = _JsonResponseConfig.from_metadata(metadata=metadata) + assert config.enabled is False + assert config.schema is None + assert config.schema_name == "CustomSchema" + assert config.strict is True diff --git a/tests/unit/target/test_openai_chat_target.py b/tests/unit/target/test_openai_chat_target.py index 875417dfb..fb0af3aa7 100644 --- a/tests/unit/target/test_openai_chat_target.py +++ b/tests/unit/target/test_openai_chat_target.py @@ -24,6 +24,7 @@ ) from pyrit.memory.memory_interface import MemoryInterface from pyrit.models import Message, MessagePiece +from pyrit.models.json_response_config import _JsonResponseConfig from pyrit.prompt_target import OpenAIChatTarget, PromptChatTarget @@ -130,7 +131,6 @@ def test_init_is_json_supported_can_be_set_to_true(patch_central_database): @pytest.mark.asyncio() async def test_build_chat_messages_for_multi_modal(target: OpenAIChatTarget): - image_request = get_image_message_piece() entries = [ Message( @@ -183,23 +183,53 @@ async def test_construct_request_body_includes_extra_body_params( request = Message(message_pieces=[dummy_text_message_piece]) - body = await target._construct_request_body(conversation=[request], is_json_response=False) + jrc = _JsonResponseConfig.from_metadata(metadata=None) + body = await target._construct_request_body(conversation=[request], json_config=jrc) assert body["key"] == "value" @pytest.mark.asyncio -@pytest.mark.parametrize("is_json", [True, False]) -async def test_construct_request_body_includes_json( - is_json, target: OpenAIChatTarget, dummy_text_message_piece: MessagePiece -): +async def test_construct_request_body_json_object(target: OpenAIChatTarget, dummy_text_message_piece: MessagePiece): + request = Message(message_pieces=[dummy_text_message_piece]) + jrc = _JsonResponseConfig.from_metadata(metadata={"response_format": "json"}) + + body = await target._construct_request_body(conversation=[request], json_config=jrc) + assert body["response_format"] == {"type": "json_object"} + + +@pytest.mark.asyncio +async def test_construct_request_body_json_schema(target: OpenAIChatTarget, dummy_text_message_piece: MessagePiece): + schema_obj = {"type": "object", "properties": {"name": {"type": "string"}}} + request = Message(message_pieces=[dummy_text_message_piece]) + jrc = _JsonResponseConfig.from_metadata(metadata={"response_format": "json", "json_schema": schema_obj}) + body = await target._construct_request_body(conversation=[request], json_config=jrc) + assert body["response_format"] == { + "type": "json_schema", + "json_schema": {"name": "CustomSchema", "schema": schema_obj, "strict": True}, + } + + +@pytest.mark.asyncio +async def test_construct_request_body_json_schema_optional_params( + target: OpenAIChatTarget, dummy_text_message_piece: MessagePiece +): + schema_obj = {"type": "object", "properties": {"name": {"type": "string"}}} request = Message(message_pieces=[dummy_text_message_piece]) + jrc = _JsonResponseConfig.from_metadata( + metadata={ + "response_format": "json", + "json_schema": schema_obj, + "json_schema_name": "MySchema", + "json_schema_strict": False, + } + ) - body = await target._construct_request_body(conversation=[request], is_json_response=is_json) - if is_json: - assert body["response_format"] == {"type": "json_object"} - else: - assert "response_format" not in body + body = await target._construct_request_body(conversation=[request], json_config=jrc) + assert body["response_format"] == { + "type": "json_schema", + "json_schema": {"name": "MySchema", "schema": schema_obj, "strict": False}, + } @pytest.mark.asyncio @@ -208,13 +238,15 @@ async def test_construct_request_body_removes_empty_values( ): request = Message(message_pieces=[dummy_text_message_piece]) - body = await target._construct_request_body(conversation=[request], is_json_response=False) + jrc = _JsonResponseConfig.from_metadata(metadata=None) + body = await target._construct_request_body(conversation=[request], json_config=jrc) assert "max_completion_tokens" not in body assert "max_tokens" not in body assert "temperature" not in body assert "top_p" not in body assert "frequency_penalty" not in body assert "presence_penalty" not in body + assert "response_format" not in body @pytest.mark.asyncio @@ -222,8 +254,9 @@ async def test_construct_request_body_serializes_text_message( target: OpenAIChatTarget, dummy_text_message_piece: MessagePiece ): request = Message(message_pieces=[dummy_text_message_piece]) + jrc = _JsonResponseConfig.from_metadata(metadata=None) - body = await target._construct_request_body(conversation=[request], is_json_response=False) + body = await target._construct_request_body(conversation=[request], json_config=jrc) assert ( body["messages"][0]["content"] == "dummy text" ), "Text messages are serialized in a simple way that's more broadly supported" @@ -236,8 +269,9 @@ async def test_construct_request_body_serializes_complex_message( image_piece = get_image_message_piece() image_piece.conversation_id = dummy_text_message_piece.conversation_id # Match conversation IDs request = Message(message_pieces=[dummy_text_message_piece, image_piece]) + jrc = _JsonResponseConfig.from_metadata(metadata=None) - body = await target._construct_request_body(conversation=[request], is_json_response=False) + body = await target._construct_request_body(conversation=[request], json_config=jrc) messages = body["messages"][0]["content"] assert len(messages) == 2, "Complex messages are serialized as a list" assert messages[0]["type"] == "text", "Text messages are serialized properly when multi-modal" @@ -441,7 +475,6 @@ async def test_send_prompt_async_empty_response_retries(openai_response_json: di @pytest.mark.asyncio async def test_send_prompt_async_rate_limit_exception_retries(target: OpenAIChatTarget): - message = Message(message_pieces=[MessagePiece(role="user", conversation_id="12345", original_value="Hello")]) # Create proper mock request and response for RateLimitError @@ -458,7 +491,6 @@ async def test_send_prompt_async_rate_limit_exception_retries(target: OpenAIChat @pytest.mark.asyncio async def test_send_prompt_async_bad_request_error(target: OpenAIChatTarget): - # Create proper mock request and response for BadRequestError (without content_filter) mock_request = httpx.Request("POST", "https://api.openai.com/v1/chat/completions") mock_response = httpx.Response(400, text="Bad Request Error", request=mock_request) @@ -476,7 +508,6 @@ async def test_send_prompt_async_bad_request_error(target: OpenAIChatTarget): @pytest.mark.asyncio async def test_send_prompt_async_content_filter_200(target: OpenAIChatTarget): - message = Message( message_pieces=[ MessagePiece( @@ -503,7 +534,6 @@ async def test_send_prompt_async_content_filter_200(target: OpenAIChatTarget): def test_validate_request_unsupported_data_types(target: OpenAIChatTarget): - image_piece = get_image_message_piece() image_piece.converted_value_data_type = "new_unknown_type" # type: ignore message = Message( @@ -548,7 +578,6 @@ def test_inheritance_from_prompt_chat_target_base(): def test_is_response_format_json_supported(target: OpenAIChatTarget): - message_piece = MessagePiece( role="user", original_value="original prompt text", @@ -559,10 +588,28 @@ def test_is_response_format_json_supported(target: OpenAIChatTarget): ) result = target.is_response_format_json(message_piece) - + assert isinstance(result, bool) assert result is True +def test_is_response_format_json_schema_supported(target: OpenAIChatTarget): + schema = {"type": "object", "properties": {"name": {"type": "string"}}} + message_piece = MessagePiece( + role="user", + original_value="original prompt text", + converted_value="Hello, how are you?", + conversation_id="conversation_1", + sequence=0, + prompt_metadata={ + "response_format": "json", + "json_schema": json.dumps(schema), + }, + ) + + result = target.is_response_format_json(message_piece) + assert result + + def test_is_response_format_json_no_metadata(target: OpenAIChatTarget): message_piece = MessagePiece( role="user", @@ -589,7 +636,6 @@ async def test_send_prompt_async_content_filter_400(target: OpenAIChatTarget): patch.object(target, "_validate_request"), patch.object(target, "_construct_request_body", new_callable=AsyncMock) as mock_construct, ): - mock_construct.return_value = {"model": "gpt-4", "messages": [], "stream": False} # Create proper mock request and response for BadRequestError with content_filter diff --git a/tests/unit/target/test_openai_response_target.py b/tests/unit/target/test_openai_response_target.py index 4f905a718..a20f6be11 100644 --- a/tests/unit/target/test_openai_response_target.py +++ b/tests/unit/target/test_openai_response_target.py @@ -23,6 +23,7 @@ ) from pyrit.memory.memory_interface import MemoryInterface from pyrit.models import Message, MessagePiece +from pyrit.models.json_response_config import _JsonResponseConfig from pyrit.prompt_target import OpenAIResponseTarget, PromptChatTarget @@ -138,7 +139,6 @@ def test_init_with_no_additional_request_headers_var_raises(): @pytest.mark.asyncio() async def test_build_input_for_multi_modal(target: OpenAIResponseTarget): - image_request = get_image_message_piece() conversation_id = image_request.conversation_id entries = [ @@ -217,23 +217,37 @@ async def test_construct_request_body_includes_extra_body_params( request = Message(message_pieces=[dummy_text_message_piece]) - body = await target._construct_request_body(conversation=[request], is_json_response=False) + jrc = _JsonResponseConfig.from_metadata(metadata=None) + body = await target._construct_request_body(conversation=[request], json_config=jrc) assert body["key"] == "value" @pytest.mark.asyncio -@pytest.mark.parametrize("is_json", [True, False]) -async def test_construct_request_body_includes_json( - is_json, target: OpenAIResponseTarget, dummy_text_message_piece: MessagePiece -): +async def test_construct_request_body_json_object(target: OpenAIResponseTarget, dummy_text_message_piece: MessagePiece): + json_response_config = _JsonResponseConfig(enabled=True) + request = Message(message_pieces=[dummy_text_message_piece]) + + body = await target._construct_request_body(conversation=[request], json_config=json_response_config) + assert body["text"] == {"format": {"type": "json_object"}} + +@pytest.mark.asyncio +async def test_construct_request_body_json_schema(target: OpenAIResponseTarget, dummy_text_message_piece: MessagePiece): + schema_object = {"type": "object", "properties": {"name": {"type": "string"}}} + json_response_config = _JsonResponseConfig.from_metadata( + metadata={"response_format": "json", "json_schema": schema_object} + ) request = Message(message_pieces=[dummy_text_message_piece]) - body = await target._construct_request_body(conversation=[request], is_json_response=is_json) - if is_json: - assert body["response_format"] == {"type": "json_object"} - else: - assert "response_format" not in body + body = await target._construct_request_body(conversation=[request], json_config=json_response_config) + assert body["text"] == { + "format": { + "type": "json_schema", + "schema": schema_object, + "name": "CustomSchema", + "strict": True, + } + } @pytest.mark.asyncio @@ -242,13 +256,15 @@ async def test_construct_request_body_removes_empty_values( ): request = Message(message_pieces=[dummy_text_message_piece]) - body = await target._construct_request_body(conversation=[request], is_json_response=False) + json_response_config = _JsonResponseConfig(enabled=False) + body = await target._construct_request_body(conversation=[request], json_config=json_response_config) assert "max_completion_tokens" not in body assert "max_tokens" not in body assert "temperature" not in body assert "top_p" not in body assert "frequency_penalty" not in body assert "presence_penalty" not in body + assert "text" not in body @pytest.mark.asyncio @@ -257,7 +273,8 @@ async def test_construct_request_body_serializes_text_message( ): request = Message(message_pieces=[dummy_text_message_piece]) - body = await target._construct_request_body(conversation=[request], is_json_response=False) + jrc = _JsonResponseConfig.from_metadata(metadata=None) + body = await target._construct_request_body(conversation=[request], json_config=jrc) assert body["input"][0]["content"][0]["text"] == "dummy text" @@ -265,13 +282,13 @@ async def test_construct_request_body_serializes_text_message( async def test_construct_request_body_serializes_complex_message( target: OpenAIResponseTarget, dummy_text_message_piece: MessagePiece ): - image_piece = get_image_message_piece() dummy_text_message_piece.conversation_id = image_piece.conversation_id request = Message(message_pieces=[dummy_text_message_piece, image_piece]) + jrc = _JsonResponseConfig.from_metadata(metadata=None) - body = await target._construct_request_body(conversation=[request], is_json_response=False) + body = await target._construct_request_body(conversation=[request], json_config=jrc) messages = body["input"][0]["content"] assert len(messages) == 2 assert messages[0]["type"] == "input_text" @@ -479,7 +496,6 @@ async def test_send_prompt_async_empty_response_retries(openai_response_json: di @pytest.mark.asyncio async def test_send_prompt_async_rate_limit_exception_retries(target: OpenAIResponseTarget): - message = Message(message_pieces=[MessagePiece(role="user", conversation_id="12345", original_value="Hello")]) # Mock SDK to raise RateLimitError @@ -556,7 +572,6 @@ async def test_send_prompt_async_content_filter(target: OpenAIResponseTarget): def test_validate_request_unsupported_data_types(target: OpenAIResponseTarget): - image_piece = get_image_message_piece() image_piece.converted_value_data_type = "new_unknown_type" # type: ignore message = Message( @@ -589,7 +604,6 @@ def test_inheritance_from_prompt_chat_target(target: OpenAIResponseTarget): def test_is_response_format_json_supported(target: OpenAIResponseTarget): - message_piece = MessagePiece( role="user", original_value="original prompt text", @@ -601,9 +615,28 @@ def test_is_response_format_json_supported(target: OpenAIResponseTarget): result = target.is_response_format_json(message_piece) + assert isinstance(result, bool) assert result is True +def test_is_response_format_json_schema_supported(target: OpenAIResponseTarget): + schema = {"type": "object", "properties": {"name": {"type": "string"}}} + message_piece = MessagePiece( + role="user", + original_value="original prompt text", + converted_value="Hello, how are you?", + conversation_id="conversation_1", + sequence=0, + prompt_metadata={ + "response_format": "json", + "json_schema": json.dumps(schema), + }, + ) + + result = target.is_response_format_json(message_piece) + assert result + + def test_is_response_format_json_no_metadata(target: OpenAIResponseTarget): message_piece = MessagePiece( role="user", @@ -683,7 +716,8 @@ async def test_construct_request_body_filters_none( target: OpenAIResponseTarget, dummy_text_message_piece: MessagePiece ): req = Message(message_pieces=[dummy_text_message_piece]) - body = await target._construct_request_body([req], is_json_response=False) + jrc = _JsonResponseConfig.from_metadata(metadata=None) + body = await target._construct_request_body(conversation=[req], json_config=jrc) assert "max_output_tokens" not in body or body["max_output_tokens"] is None assert "temperature" not in body or body["temperature"] is None assert "top_p" not in body or body["top_p"] is None