[FEAT] Support JSON Schema in Responses (#1177)

Co-authored-by: =Richard Edgar <=>
Co-authored-by: Roman Lutz <romanlutz13@gmail.com>

Author: riedgar-ms

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",

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,

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
 #