diff --git a/pyrit/converter/decomposition_converter.py b/pyrit/converter/decomposition_converter.py index cba8a26531..936e3f8ad6 100644 --- a/pyrit/converter/decomposition_converter.py +++ b/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=[ diff --git a/pyrit/executor/attack/component/adversarial_conversation_manager.py b/pyrit/executor/attack/component/adversarial_conversation_manager.py index ba062fae66..6d8fbd2bd5 100644 --- a/pyrit/executor/attack/component/adversarial_conversation_manager.py +++ b/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: diff --git a/pyrit/executor/promptgen/fuzzer/fuzzer_converter_base.py b/pyrit/executor/promptgen/fuzzer/fuzzer_converter_base.py index 34c4c69f9c..5c730638be 100644 --- a/pyrit/executor/promptgen/fuzzer/fuzzer_converter_base.py +++ b/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( diff --git a/pyrit/executor/promptgen/fuzzer/fuzzer_crossover_converter.py b/pyrit/executor/promptgen/fuzzer/fuzzer_crossover_converter.py index 0217692f42..16534e27e1 100644 --- a/pyrit/executor/promptgen/fuzzer/fuzzer_crossover_converter.py +++ b/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( diff --git a/pyrit/executor/promptgen/fuzzer/fuzzer_expand_converter.py b/pyrit/executor/promptgen/fuzzer/fuzzer_expand_converter.py index 5fd2276b21..0812da760c 100644 --- a/pyrit/executor/promptgen/fuzzer/fuzzer_expand_converter.py +++ b/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( diff --git a/pyrit/models/__init__.py b/pyrit/models/__init__.py index d041e378b5..162c7f7e25 100644 --- a/pyrit/models/__init__.py +++ b/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", diff --git a/pyrit/models/catalog/target.py b/pyrit/models/catalog/target.py index cc9d76f0f8..6fad8fe073 100644 --- a/pyrit/models/catalog/target.py +++ b/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): diff --git a/pyrit/models/seeds/seed_prompt.py b/pyrit/models/seeds/seed_prompt.py index c8c9455f22..e00cf607ba 100644 --- a/pyrit/models/seeds/seed_prompt.py +++ b/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 diff --git a/pyrit/models/target/__init__.py b/pyrit/models/target/__init__.py new file mode 100644 index 0000000000..d43320dda2 --- /dev/null +++ b/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", +] diff --git a/pyrit/models/target/json_response_config.py b/pyrit/models/target/json_response_config.py new file mode 100644 index 0000000000..763fce93d4 --- /dev/null +++ b/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 diff --git a/pyrit/models/json_schema_definition.py b/pyrit/models/target/json_schema_definition.py similarity index 100% rename from pyrit/models/json_schema_definition.py rename to pyrit/models/target/json_schema_definition.py diff --git a/pyrit/models/target_capabilities.py b/pyrit/models/target/target_capabilities.py similarity index 100% rename from pyrit/models/target_capabilities.py rename to pyrit/models/target/target_capabilities.py diff --git a/pyrit/models/token_usage.py b/pyrit/models/target/token_usage.py similarity index 100% rename from pyrit/models/token_usage.py rename to pyrit/models/target/token_usage.py diff --git a/pyrit/prompt_target/common/chat_completions_message_builder.py b/pyrit/prompt_target/common/chat_completions_message_builder.py index 8c30fd6a27..2d9e4896f9 100644 --- a/pyrit/prompt_target/common/chat_completions_message_builder.py +++ b/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: diff --git a/pyrit/prompt_target/common/json_response_config.py b/pyrit/prompt_target/common/json_response_config.py deleted file mode 100644 index 0dc91980c1..0000000000 --- a/pyrit/prompt_target/common/json_response_config.py +++ /dev/null @@ -1,75 +0,0 @@ -# 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 - -from pyrit.models.json_schema_definition import ( - JSON_SCHEMA_METADATA_KEY, # noqa: TC001 (runtime-required by Pydantic field annotations) -) - - -class _JsonResponseConfig(BaseModel): - """ - Configuration for JSON responses (OpenAI request-format vocabulary). - - Parses an OpenAI-shaped ``response_format`` / ``text.format`` block out of - a ``MessagePiece.prompt_metadata`` dict. Lives in the ``pyrit.prompt_target`` - layer because the keys it speaks (``"response_format"``, ``"json_schema_name"``, - ``"json_schema_strict"``) are OpenAI request-shape vocabulary, not - framework-wide concepts. The framework-wide pieces (the - ``JsonSchemaDefinition`` alias and the ``JSON_SCHEMA_METADATA_KEY`` contract) - stay in ``pyrit.models``. - - 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 - """ - - model_config = ConfigDict(extra="forbid") - - # Would prefer StrEnum, but.... Python 3.10 - _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 - - @classmethod - def from_metadata(cls, *, metadata: dict[str, Any] | None) -> _JsonResponseConfig: - 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) diff --git a/pyrit/prompt_target/common/prompt_target.py b/pyrit/prompt_target/common/prompt_target.py index aa5ed47a64..89a4cf1fcc 100644 --- a/pyrit/prompt_target/common/prompt_target.py +++ b/pyrit/prompt_target/common/prompt_target.py @@ -6,8 +6,15 @@ from typing import Any, ClassVar, Literal, final from pyrit.memory import CentralMemory, MemoryInterface -from pyrit.models import ComponentIdentifier, Conversation, Identifiable, Message, MessagePiece, TargetIdentifier -from pyrit.prompt_target.common.json_response_config import _JsonResponseConfig +from pyrit.models import ( + ComponentIdentifier, + Conversation, + Identifiable, + JsonResponseConfig, + Message, + MessagePiece, + TargetIdentifier, +) from pyrit.prompt_target.common.target_capabilities import ( CapabilityName, TargetCapabilities, @@ -505,7 +512,7 @@ def is_response_format_json(self, message_piece: MessagePiece) -> bool: config = self._get_json_response_config(message_piece=message_piece) return config.enabled - def _get_json_response_config(self, *, message_piece: MessagePiece) -> _JsonResponseConfig: + def _get_json_response_config(self, *, message_piece: MessagePiece) -> JsonResponseConfig: """ Get the JSON response configuration from the message piece metadata. @@ -514,12 +521,12 @@ def _get_json_response_config(self, *, message_piece: MessagePiece) -> _JsonResp include JSON response configuration. Returns: - _JsonResponseConfig: The JSON response configuration. + JsonResponseConfig: The JSON response configuration. Raises: ValueError: If JSON response format is requested but unsupported. """ - config = _JsonResponseConfig.from_metadata(metadata=message_piece.prompt_metadata) + config = JsonResponseConfig.from_metadata(metadata=message_piece.prompt_metadata) if config.enabled and not self.capabilities.supports_json_output: target_name = self.get_identifier().class_name diff --git a/pyrit/prompt_target/common/target_capabilities.py b/pyrit/prompt_target/common/target_capabilities.py index 10c1ec0ce2..382feed1c3 100644 --- a/pyrit/prompt_target/common/target_capabilities.py +++ b/pyrit/prompt_target/common/target_capabilities.py @@ -8,7 +8,7 @@ from typing import NoReturn, cast from pyrit.models.literals import PromptDataType -from pyrit.models.target_capabilities import CapabilityName, TargetCapabilities +from pyrit.models.target.target_capabilities import CapabilityName, TargetCapabilities __all__ = [ "CapabilityHandlingPolicy", diff --git a/pyrit/prompt_target/litellm_chat_target.py b/pyrit/prompt_target/litellm_chat_target.py index 870d89b2b2..070ede86c8 100644 --- a/pyrit/prompt_target/litellm_chat_target.py +++ b/pyrit/prompt_target/litellm_chat_target.py @@ -17,6 +17,7 @@ ) from pyrit.models import ( ComponentIdentifier, + JsonResponseConfig, Message, MessagePiece, PromptDataType, @@ -35,7 +36,6 @@ is_content_filter_response, validate_chat_completion_response, ) -from pyrit.prompt_target.common.json_response_config import _JsonResponseConfig from pyrit.prompt_target.common.prompt_target import PromptTarget from pyrit.prompt_target.common.target_capabilities import ( TargetCapabilities, @@ -426,7 +426,7 @@ def _construct_request_body( self, *, messages: list[dict[str, Any]], - json_config: _JsonResponseConfig, + json_config: JsonResponseConfig, api_key: str | None = None, ) -> dict[str, Any]: body: dict[str, Any] = { diff --git a/pyrit/prompt_target/openai/openai_chat_target.py b/pyrit/prompt_target/openai/openai_chat_target.py index 3179046246..4da1ed77c8 100644 --- a/pyrit/prompt_target/openai/openai_chat_target.py +++ b/pyrit/prompt_target/openai/openai_chat_target.py @@ -11,6 +11,7 @@ ) from pyrit.models import ( ComponentIdentifier, + JsonResponseConfig, Message, MessagePiece, ) @@ -30,7 +31,6 @@ save_audio_response_async, validate_chat_completion_response, ) -from pyrit.prompt_target.common.json_response_config import _JsonResponseConfig from pyrit.prompt_target.common.target_capabilities import TargetCapabilities from pyrit.prompt_target.common.target_configuration import TargetConfiguration from pyrit.prompt_target.common.utils import ( @@ -457,7 +457,7 @@ async def _build_chat_messages_for_multi_modal_async( ) async def _construct_request_body_async( - self, *, conversation: MutableSequence[Message], json_config: _JsonResponseConfig + self, *, conversation: MutableSequence[Message], json_config: JsonResponseConfig ) -> dict[str, Any]: messages = await self._build_chat_messages_async(conversation) response_format = self._build_response_format(json_config) @@ -483,5 +483,5 @@ async def _construct_request_body_async( # Filter out None values return {k: v for k, v in body_parameters.items() if v is not None} - def _build_response_format(self, json_config: _JsonResponseConfig) -> dict[str, Any] | None: + def _build_response_format(self, json_config: JsonResponseConfig) -> dict[str, Any] | None: return build_response_format(json_config=json_config) diff --git a/pyrit/prompt_target/openai/openai_response_target.py b/pyrit/prompt_target/openai/openai_response_target.py index 9101fea87b..563dfaeb2d 100644 --- a/pyrit/prompt_target/openai/openai_response_target.py +++ b/pyrit/prompt_target/openai/openai_response_target.py @@ -22,12 +22,12 @@ from pyrit.memory.storage import convert_local_image_to_data_url_async from pyrit.models import ( ComponentIdentifier, + JsonResponseConfig, Message, MessagePiece, PromptDataType, PromptResponseError, ) -from pyrit.prompt_target.common.json_response_config import _JsonResponseConfig from pyrit.prompt_target.common.target_capabilities import TargetCapabilities from pyrit.prompt_target.common.target_configuration import TargetConfiguration from pyrit.prompt_target.common.utils import ( @@ -371,7 +371,7 @@ async def _build_input_for_multi_modal_async(self, conversation: MutableSequence return input_items async def _construct_request_body_async( - self, *, conversation: MutableSequence[Message], json_config: _JsonResponseConfig + self, *, conversation: MutableSequence[Message], json_config: JsonResponseConfig ) -> dict[str, Any]: """ Construct the request body to send to the Responses API. @@ -425,7 +425,7 @@ def _build_reasoning_config(self) -> dict[str, Any] | None: reasoning["summary"] = self._reasoning_summary return reasoning - def _build_text_format(self, json_config: _JsonResponseConfig) -> dict[str, Any] | None: + def _build_text_format(self, json_config: JsonResponseConfig) -> dict[str, Any] | None: if not json_config.enabled: return None diff --git a/pyrit/score/float_scale/insecure_code_scorer.py b/pyrit/score/float_scale/insecure_code_scorer.py index dc2e8c1ff8..e1674af2e3 100644 --- a/pyrit/score/float_scale/insecure_code_scorer.py +++ b/pyrit/score/float_scale/insecure_code_scorer.py @@ -163,7 +163,7 @@ def _build_identifier(self) -> ComponentIdentifier: params={ "system_prompt_template": self._system_prompt, "harm_categories": self._harm_categories, - "response_json_schema": self._response_handler.response_schema, + "response_json_schema": self._response_handler.json_response_config.json_schema, }, prompt_target=self._prompt_target.get_identifier(), ) diff --git a/pyrit/score/float_scale/self_ask_general_float_scale_scorer.py b/pyrit/score/float_scale/self_ask_general_float_scale_scorer.py index f5114ce7c5..fb4de1f084 100644 --- a/pyrit/score/float_scale/self_ask_general_float_scale_scorer.py +++ b/pyrit/score/float_scale/self_ask_general_float_scale_scorer.py @@ -121,7 +121,7 @@ def _build_identifier(self) -> ComponentIdentifier: "system_prompt_template": self._system_prompt_format_string, "user_prompt_template": self._prompt_format_string, "scale": self._scale.model_dump(), - "response_json_schema": self._response_handler.response_schema, + "response_json_schema": self._response_handler.json_response_config.json_schema, }, prompt_target=self._prompt_target.get_identifier(), ) diff --git a/pyrit/score/float_scale/self_ask_likert_scorer.py b/pyrit/score/float_scale/self_ask_likert_scorer.py index e1ae24c027..d57be166ed 100644 --- a/pyrit/score/float_scale/self_ask_likert_scorer.py +++ b/pyrit/score/float_scale/self_ask_likert_scorer.py @@ -11,6 +11,7 @@ from pyrit.exceptions import InvalidJsonException from pyrit.models import ( ComponentIdentifier, + JsonResponseConfig, JsonSchemaDefinition, MessagePiece, Score, @@ -209,14 +210,9 @@ def __init__( self._score_values = frozenset(entry.score_value for entry in likert_scale.entries) @property - def response_format(self) -> str | None: - """The wrapped handler's response-format hint.""" - return self._response_handler.response_format - - @property - def response_schema(self) -> JsonSchemaDefinition | None: - """The wrapped handler's response schema.""" - return self._response_handler.response_schema + def json_response_config(self) -> JsonResponseConfig: + """The wrapped handler's JSON-response request.""" + return self._response_handler.json_response_config def parse( self, @@ -382,7 +378,7 @@ def _build_identifier(self) -> ComponentIdentifier: params={ "system_prompt_template": self._system_prompt, "likert_scale": self._likert_scale.model_dump(), - "response_json_schema": self._response_handler.response_schema, + "response_json_schema": self._response_handler.json_response_config.json_schema, }, prompt_target=self._prompt_target.get_identifier(), ) diff --git a/pyrit/score/float_scale/self_ask_scale_scorer.py b/pyrit/score/float_scale/self_ask_scale_scorer.py index 83df018106..925363b2fb 100644 --- a/pyrit/score/float_scale/self_ask_scale_scorer.py +++ b/pyrit/score/float_scale/self_ask_scale_scorer.py @@ -187,7 +187,7 @@ def _build_identifier(self) -> ComponentIdentifier: "system_prompt_template": self._system_prompt, "user_prompt_template": "objective: {objective}\nresponse: {response}", "scale": self._scale.model_dump(exclude_none=True), - "response_json_schema": self._response_handler.response_schema, + "response_json_schema": self._response_handler.json_response_config.json_schema, }, prompt_target=self._prompt_target.get_identifier(), ) diff --git a/pyrit/score/llm_scoring.py b/pyrit/score/llm_scoring.py index 78183f0a74..45ac0827c5 100644 --- a/pyrit/score/llm_scoring.py +++ b/pyrit/score/llm_scoring.py @@ -4,10 +4,10 @@ from __future__ import annotations import uuid -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING from pyrit.exceptions import EmptyResponseException, ScorerLLMResponseBlockedException, pyrit_json_retry -from pyrit.models import JSON_SCHEMA_METADATA_KEY, Message, MessagePiece +from pyrit.models import Message, MessagePiece if TYPE_CHECKING: from collections.abc import Sequence @@ -39,7 +39,7 @@ async def _run_llm_scoring_async( Perform a single scoring round-trip against an LLM target and delegate parsing. This is the shared LLM evaluation mechanism: it sets the system prompt on the target, sends - the value to be scored (forwarding ``response_handler.response_schema`` so targets that + the value to be scored (forwarding ``response_handler.json_response_config`` so targets that support structured output can enforce it), applies the standard JSON retry behavior, and delegates parsing and validation to ``response_handler``. It is intentionally stateless and independent of any particular ``Scorer`` so that scorers can compose it without inheriting LLM @@ -90,15 +90,10 @@ async def _run_llm_scoring_async( system_prompt=system_prompt, conversation_id=conversation_id, ) - prompt_metadata: dict[str, Any] = {} - response_format = response_handler.response_format - if response_format is not None: - prompt_metadata["response_format"] = response_format - response_schema = response_handler.response_schema - if response_schema is not None: - # Always forward the schema; the target's normalization pipeline omits it - # when the target cannot natively enforce a JSON schema. - prompt_metadata[JSON_SCHEMA_METADATA_KEY] = response_schema + # Forward the JSON-response request (format and any schema together) via the handler's + # canonical config; the target's normalization pipeline omits the schema when it cannot + # natively enforce one. + prompt_metadata = response_handler.json_response_config.to_metadata() # Build message pieces - prepended text context first (if provided), then the main message being scored message_pieces: list[MessagePiece] = [] diff --git a/pyrit/score/response_handler.py b/pyrit/score/response_handler.py index 1318b6f59e..8163b76ddc 100644 --- a/pyrit/score/response_handler.py +++ b/pyrit/score/response_handler.py @@ -10,7 +10,7 @@ from typing import TYPE_CHECKING from pyrit.exceptions import InvalidJsonException, remove_markdown_json -from pyrit.models import UnvalidatedScore +from pyrit.models import JsonResponseConfig, UnvalidatedScore if TYPE_CHECKING: import uuid @@ -96,26 +96,17 @@ class ResponseHandler(abc.ABC): """ @property - def response_format(self) -> str | None: + def json_response_config(self) -> JsonResponseConfig: """ - The ``response_format`` hint forwarded onto the scoring request, or ``None`` for none. + The canonical JSON-response request this handler asks the scoring target for. - Round-trip helpers copy this onto the request metadata. Handlers that require a specific - wire format (e.g. JSON) override it; the default imposes nothing so targets that emit - plain text are not forced into a format they cannot honor. + Format and schema are one coupled unit: the LLM round-trip serializes this onto the + request metadata via ``to_metadata``, and targets that natively support structured output + enforce the schema (others have it omitted by normalization). The default is disabled, + imposing no wire format so targets that emit plain text are not forced into a format they + cannot honor. Handlers that require JSON override this. """ - return None - - @property - def response_schema(self) -> JsonSchemaDefinition | None: - """ - The JSON schema, if any, describing the response the target should return. - - The LLM round-trip forwards this to the scoring target so targets that natively support - structured output can enforce it; targets that cannot have it omitted by normalization. - Handlers that do not constrain the response shape return None (the default). - """ - return None + return JsonResponseConfig(enabled=False) @abstractmethod def parse( @@ -193,14 +184,9 @@ def __init__( self._numeric_value = numeric_value @property - def response_schema(self) -> JsonSchemaDefinition | None: - """The configured JSON schema to forward to the scoring target, if any.""" - return self._response_schema - - @property - def response_format(self) -> str | None: - """The ``"json"`` response format, since this handler parses JSON responses.""" - return "json" + def json_response_config(self) -> JsonResponseConfig: + """The JSON-response request: always JSON, carrying the optional configured schema.""" + return JsonResponseConfig(enabled=True, json_schema=self._response_schema) def parse( self, diff --git a/pyrit/score/true_false/self_ask_category_scorer.py b/pyrit/score/true_false/self_ask_category_scorer.py index 2165acd3e8..8ff148ed7d 100644 --- a/pyrit/score/true_false/self_ask_category_scorer.py +++ b/pyrit/score/true_false/self_ask_category_scorer.py @@ -16,6 +16,7 @@ from pyrit.exceptions import InvalidJsonException from pyrit.models import ( ComponentIdentifier, + JsonResponseConfig, JsonSchemaDefinition, MessagePiece, Score, @@ -144,14 +145,9 @@ def __init__( self._fallback_category = content_classifier.no_category_found @property - def response_format(self) -> str | None: - """The wrapped handler's response-format hint.""" - return self._response_handler.response_format - - @property - def response_schema(self) -> JsonSchemaDefinition | None: - """The wrapped handler's response schema.""" - return self._response_handler.response_schema + def json_response_config(self) -> JsonResponseConfig: + """The wrapped handler's JSON-response request.""" + return self._response_handler.json_response_config def parse( self, @@ -329,7 +325,7 @@ def _build_identifier(self) -> ComponentIdentifier: params={ "system_prompt_template": self._system_prompt, "content_classifier": self._content_classifier.model_dump(), - "response_json_schema": self._response_handler.response_schema, + "response_json_schema": self._response_handler.json_response_config.json_schema, }, score_aggregator=self._score_aggregator.__name__, # type: ignore[ty:unresolved-attribute] prompt_target=self._prompt_target.get_identifier(), diff --git a/pyrit/score/true_false/self_ask_general_true_false_scorer.py b/pyrit/score/true_false/self_ask_general_true_false_scorer.py index 84a3ee48ac..b2ffaa1f2c 100644 --- a/pyrit/score/true_false/self_ask_general_true_false_scorer.py +++ b/pyrit/score/true_false/self_ask_general_true_false_scorer.py @@ -130,7 +130,7 @@ def _build_identifier(self) -> ComponentIdentifier: params={ "system_prompt_template": self._system_prompt_format_string, "user_prompt_template": self._prompt_format_string, - "response_json_schema": self._response_handler.response_schema, + "response_json_schema": self._response_handler.json_response_config.json_schema, }, score_aggregator=self._score_aggregator.__name__, # type: ignore[ty:unresolved-attribute] prompt_target=self._prompt_target.get_identifier(), diff --git a/pyrit/score/true_false/self_ask_refusal_scorer.py b/pyrit/score/true_false/self_ask_refusal_scorer.py index 1f668c6dcf..4b100cde5b 100644 --- a/pyrit/score/true_false/self_ask_refusal_scorer.py +++ b/pyrit/score/true_false/self_ask_refusal_scorer.py @@ -166,7 +166,7 @@ def _build_identifier(self) -> ComponentIdentifier: params={ "system_prompt_template": self._system_prompt, "user_prompt_template": self._prompt_format_string, - "response_json_schema": self._response_handler.response_schema, + "response_json_schema": self._response_handler.json_response_config.json_schema, }, score_aggregator=self._score_aggregator.__name__, # type: ignore[ty:unresolved-attribute] prompt_target=self._prompt_target.get_identifier(), diff --git a/pyrit/score/true_false/self_ask_true_false_scorer.py b/pyrit/score/true_false/self_ask_true_false_scorer.py index 421dd1c385..401a2ed895 100644 --- a/pyrit/score/true_false/self_ask_true_false_scorer.py +++ b/pyrit/score/true_false/self_ask_true_false_scorer.py @@ -276,7 +276,7 @@ def _build_identifier(self) -> ComponentIdentifier: "system_prompt_template": self._system_prompt, "user_prompt_template": "objective: {objective}\nresponse: {response}", "question": self._question.model_dump(), - "response_json_schema": self._response_handler.response_schema, + "response_json_schema": self._response_handler.json_response_config.json_schema, }, score_aggregator=self._score_aggregator.__name__, # type: ignore[ty:unresolved-attribute] prompt_target=self._prompt_target.get_identifier(), diff --git a/tests/unit/executor/attack/component/test_adversarial_conversation_manager.py b/tests/unit/executor/attack/component/test_adversarial_conversation_manager.py index 5e23990d61..1c1f34fa31 100644 --- a/tests/unit/executor/attack/component/test_adversarial_conversation_manager.py +++ b/tests/unit/executor/attack/component/test_adversarial_conversation_manager.py @@ -139,7 +139,7 @@ def test_build_metadata_returns_empty_without_schema(): def test_build_metadata_forwards_schema(): metadata = _build_adversarial_prompt_metadata(response_json_schema=SCHEMA) assert metadata["response_format"] == "json" - assert metadata[JSON_SCHEMA_METADATA_KEY] is SCHEMA + assert metadata[JSON_SCHEMA_METADATA_KEY] == SCHEMA # --- _parse_adversarial_reply ------------------------------------------------ diff --git a/tests/unit/executor/attack/test_attack_parameter_consistency.py b/tests/unit/executor/attack/test_attack_parameter_consistency.py index e392db99a0..a41c02f056 100644 --- a/tests/unit/executor/attack/test_attack_parameter_consistency.py +++ b/tests/unit/executor/attack/test_attack_parameter_consistency.py @@ -36,8 +36,8 @@ PromptDataType, Score, SeedPrompt, + get_common_json_schema, ) -from pyrit.models.json_schema_definition import get_common_json_schema from pyrit.prompt_normalizer import PromptNormalizer from pyrit.prompt_target import PromptTarget from pyrit.score import FloatScaleThresholdScorer, TrueFalseScorer diff --git a/tests/unit/memory/test_memory_models.py b/tests/unit/memory/test_memory_models.py index 63a998778b..bfb72d0379 100644 --- a/tests/unit/memory/test_memory_models.py +++ b/tests/unit/memory/test_memory_models.py @@ -349,7 +349,7 @@ def test_roundtrip_seed_prompt_preserves_inline_response_json_schema(self): def test_roundtrip_seed_prompt_with_named_schema_kwarg(self): """Named-schema construction round-trips as the resolved body.""" - from pyrit.models.json_schema_definition import get_common_json_schema + from pyrit.models import get_common_json_schema expected = get_common_json_schema("true_false_with_rationale") seed = _make_seed_prompt(response_json_schema_name="true_false_with_rationale") 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 0000000000..651a3d547b --- /dev/null +++ b/tests/unit/models/test_json_response_config.py @@ -0,0 +1,173 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import json + +import pytest + +from pyrit.models import JsonResponseConfig + + +def test_with_none(): + config = JsonResponseConfig.from_metadata(metadata=None) + assert config.enabled is False + assert config.json_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.json_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.json_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.json_schema == schema + assert config.schema_name == "CustomSchema" + assert config.strict is True + + +def test_with_empty_json_schema_object(): + metadata = { + "response_format": "json", + "json_schema": {}, + } + config = JsonResponseConfig.from_metadata(metadata=metadata) + assert config.enabled is True + assert config.json_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_improbable_to_have_here", + } + config = JsonResponseConfig.from_metadata(metadata=metadata) + assert config.enabled is False + assert config.json_schema is None + assert config.schema_name == "CustomSchema" + assert config.strict is True + + +def test_schema_without_response_format_is_disabled(): + # A schema is meaningless without the response_format marker, so it must be dropped. + metadata = { + "json_schema": {"type": "object"}, + "json_schema_name": "TestSchema", + } + config = JsonResponseConfig.from_metadata(metadata=metadata) + assert config.enabled is False + assert config.json_schema is None + + +def test_with_empty_json_schema_string(): + metadata = { + "response_format": "json", + "json_schema": "", + } + config = JsonResponseConfig.from_metadata(metadata=metadata) + assert config.enabled is True + assert config.json_schema is None + + +def test_ignores_unrelated_metadata_keys(): + # Real prompt_metadata carries other keys (e.g. token usage) that must be ignored. + metadata = { + "response_format": "json", + "total_tokens": 42, + "some_other_key": "value", + } + config = JsonResponseConfig.from_metadata(metadata=metadata) + assert config.enabled is True + assert config.json_schema is None + + +def test_to_metadata_disabled_is_empty(): + assert JsonResponseConfig(enabled=False).to_metadata() == {} + + +def test_schema_implies_enabled(): + # A schema is meaningless without JSON output, so providing one forces enabled. + config = JsonResponseConfig(json_schema={"type": "object"}) + assert config.enabled is True + # An explicit enabled=False alongside a schema is coerced rather than kept contradictory. + coerced = JsonResponseConfig(enabled=False, json_schema={"type": "object"}) + assert coerced.enabled is True + assert coerced.to_metadata()["response_format"] == "json" + + +def test_to_metadata_enabled_without_schema(): + assert JsonResponseConfig(enabled=True).to_metadata() == {"response_format": "json"} + + +def test_to_metadata_enabled_with_schema(): + schema = {"type": "object", "properties": {"name": {"type": "string"}}} + metadata = JsonResponseConfig( + enabled=True, + json_schema=schema, + schema_name="TestSchema", + strict=False, + ).to_metadata() + assert metadata == { + "response_format": "json", + "json_schema": schema, + "json_schema_name": "TestSchema", + "json_schema_strict": False, + } + + +@pytest.mark.parametrize( + "config", + [ + JsonResponseConfig(enabled=False), + JsonResponseConfig(enabled=True), + JsonResponseConfig(enabled=True, json_schema={"type": "object"}), + JsonResponseConfig( + enabled=True, + json_schema={"type": "object", "properties": {"a": {"type": "integer"}}}, + schema_name="Custom", + strict=False, + ), + ], +) +def test_metadata_round_trips(config): + assert JsonResponseConfig.from_metadata(metadata=config.to_metadata()) == config diff --git a/tests/unit/models/test_json_schema_definition.py b/tests/unit/models/test_json_schema_definition.py index a73d17b7d2..9933025ddd 100644 --- a/tests/unit/models/test_json_schema_definition.py +++ b/tests/unit/models/test_json_schema_definition.py @@ -16,7 +16,7 @@ register_common_json_schema, unregister_common_json_schema, ) -from pyrit.models import json_schema_definition as jsd +from pyrit.models.target import json_schema_definition as jsd @pytest.fixture diff --git a/tests/unit/prompt_target/target/test_chat_completions_helpers.py b/tests/unit/prompt_target/target/test_chat_completions_helpers.py index 5502ae4286..00373fed04 100644 --- a/tests/unit/prompt_target/target/test_chat_completions_helpers.py +++ b/tests/unit/prompt_target/target/test_chat_completions_helpers.py @@ -11,7 +11,7 @@ import pytest from pyrit.exceptions import EmptyResponseException, PyritException -from pyrit.models import Message, MessagePiece +from pyrit.models import JsonResponseConfig, Message, MessagePiece from pyrit.prompt_target.common.chat_completions_message_builder import ( build_multimodal_chat_messages_async, build_response_format, @@ -31,7 +31,6 @@ token_usage_from_chat_completion, validate_chat_completion_response, ) -from pyrit.prompt_target.common.json_response_config import _JsonResponseConfig pytestmark = pytest.mark.usefixtures("patch_central_database") @@ -89,18 +88,18 @@ def test_build_text_content_entry(): def test_build_response_format_disabled_returns_none(): - config = _JsonResponseConfig.from_metadata(metadata={}) + config = JsonResponseConfig.from_metadata(metadata={}) assert build_response_format(json_config=config) is None def test_build_response_format_json_object(): - config = _JsonResponseConfig.from_metadata(metadata={"response_format": "json"}) + config = JsonResponseConfig.from_metadata(metadata={"response_format": "json"}) assert build_response_format(json_config=config) == {"type": "json_object"} def test_build_response_format_json_schema(): schema = {"type": "object", "properties": {"a": {"type": "string"}}} - config = _JsonResponseConfig.from_metadata(metadata={"response_format": "json", "json_schema": json.dumps(schema)}) + config = JsonResponseConfig.from_metadata(metadata={"response_format": "json", "json_schema": json.dumps(schema)}) result = build_response_format(json_config=config) assert result is not None assert result["type"] == "json_schema" diff --git a/tests/unit/prompt_target/target/test_json_response_config.py b/tests/unit/prompt_target/target/test_json_response_config.py deleted file mode 100644 index b2f44bce92..0000000000 --- a/tests/unit/prompt_target/target/test_json_response_config.py +++ /dev/null @@ -1,88 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT license. - -import json - -import pytest - -from pyrit.prompt_target.common.json_response_config import _JsonResponseConfig - - -def test_with_none(): - config = _JsonResponseConfig.from_metadata(metadata=None) - assert config.enabled is False - assert config.json_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.json_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.json_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.json_schema == schema - assert config.schema_name == "CustomSchema" - assert config.strict is True - - -def test_with_empty_json_schema_object(): - metadata = { - "response_format": "json", - "json_schema": {}, - } - config = _JsonResponseConfig.from_metadata(metadata=metadata) - assert config.enabled is True - assert config.json_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.json_schema is None - assert config.schema_name == "CustomSchema" - assert config.strict is True diff --git a/tests/unit/prompt_target/target/test_litellm_chat_target.py b/tests/unit/prompt_target/target/test_litellm_chat_target.py index 857c7b45a9..30cd9245b5 100644 --- a/tests/unit/prompt_target/target/test_litellm_chat_target.py +++ b/tests/unit/prompt_target/target/test_litellm_chat_target.py @@ -15,13 +15,12 @@ RateLimitException, get_retry_max_num_attempts, ) -from pyrit.models import Message, MessagePiece +from pyrit.models import JsonResponseConfig, Message, MessagePiece from pyrit.prompt_target import ( OpenAIChatAudioConfig, TargetCapabilities, TargetConfiguration, ) -from pyrit.prompt_target.common.json_response_config import _JsonResponseConfig from pyrit.prompt_target.litellm_chat_target import LiteLLMChatTarget # --------------------------------------------------------------------------- @@ -135,8 +134,8 @@ def _user_message(text="test prompt", conversation_id="convo"): return piece.to_message() -def _disabled_json_config() -> _JsonResponseConfig: - return _JsonResponseConfig.from_metadata(metadata={}) +def _disabled_json_config() -> JsonResponseConfig: + return JsonResponseConfig.from_metadata(metadata={}) # --------------------------------------------------------------------------- diff --git a/tests/unit/prompt_target/target/test_openai_chat_target.py b/tests/unit/prompt_target/target/test_openai_chat_target.py index 189123ebbd..71a2ca7d63 100644 --- a/tests/unit/prompt_target/target/test_openai_chat_target.py +++ b/tests/unit/prompt_target/target/test_openai_chat_target.py @@ -25,14 +25,13 @@ RateLimitException, ) from pyrit.memory.memory_interface import MemoryInterface -from pyrit.models import Message, MessagePiece, flatten_to_message_pieces +from pyrit.models import JsonResponseConfig, Message, MessagePiece, flatten_to_message_pieces from pyrit.prompt_target import ( OpenAIChatAudioConfig, OpenAIChatTarget, OpenAIResponseTarget, PromptTarget, ) -from pyrit.prompt_target.common.json_response_config import _JsonResponseConfig from pyrit.prompt_target.common.target_capabilities import TargetCapabilities from pyrit.prompt_target.common.target_configuration import TargetConfiguration @@ -157,14 +156,14 @@ async def test_construct_request_body_includes_extra_body_params( request = Message(message_pieces=[dummy_text_message_piece]) - jrc = _JsonResponseConfig.from_metadata(metadata=None) + jrc = JsonResponseConfig.from_metadata(metadata=None) body = await target._construct_request_body_async(conversation=[request], json_config=jrc) assert body["key"] == "value" 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"}) + jrc = JsonResponseConfig.from_metadata(metadata={"response_format": "json"}) body = await target._construct_request_body_async(conversation=[request], json_config=jrc) assert body["response_format"] == {"type": "json_object"} @@ -173,7 +172,7 @@ async def test_construct_request_body_json_object(target: OpenAIChatTarget, dumm 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}) + jrc = JsonResponseConfig.from_metadata(metadata={"response_format": "json", "json_schema": schema_obj}) body = await target._construct_request_body_async(conversation=[request], json_config=jrc) assert body["response_format"] == { @@ -187,7 +186,7 @@ async def test_construct_request_body_json_schema_optional_params( ): schema_obj = {"type": "object", "properties": {"name": {"type": "string"}}} request = Message(message_pieces=[dummy_text_message_piece]) - jrc = _JsonResponseConfig.from_metadata( + jrc = JsonResponseConfig.from_metadata( metadata={ "response_format": "json", "json_schema": schema_obj, @@ -208,7 +207,7 @@ async def test_construct_request_body_removes_empty_values( ): request = Message(message_pieces=[dummy_text_message_piece]) - jrc = _JsonResponseConfig.from_metadata(metadata=None) + jrc = JsonResponseConfig.from_metadata(metadata=None) body = await target._construct_request_body_async(conversation=[request], json_config=jrc) assert "max_completion_tokens" not in body assert "max_tokens" not in body @@ -223,7 +222,7 @@ 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) + jrc = JsonResponseConfig.from_metadata(metadata=None) body = await target._construct_request_body_async(conversation=[request], json_config=jrc) assert body["messages"][0]["content"] == "dummy text", ( @@ -237,7 +236,7 @@ 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) + jrc = JsonResponseConfig.from_metadata(metadata=None) body = await target._construct_request_body_async(conversation=[request], json_config=jrc) messages = body["messages"][0]["content"] @@ -1261,7 +1260,7 @@ async def test_construct_request_body_with_audio_config(patch_central_database, ) request = Message(message_pieces=[dummy_text_message_piece]) - jrc = _JsonResponseConfig.from_metadata(metadata=None) + jrc = JsonResponseConfig.from_metadata(metadata=None) body = await target._construct_request_body_async(conversation=[request], json_config=jrc) diff --git a/tests/unit/prompt_target/target/test_openai_response_target.py b/tests/unit/prompt_target/target/test_openai_response_target.py index 42f0d0266c..5cd0f1f790 100644 --- a/tests/unit/prompt_target/target/test_openai_response_target.py +++ b/tests/unit/prompt_target/target/test_openai_response_target.py @@ -23,9 +23,8 @@ RateLimitException, ) from pyrit.memory.memory_interface import MemoryInterface -from pyrit.models import Message, MessagePiece, flatten_to_message_pieces +from pyrit.models import JsonResponseConfig, Message, MessagePiece, flatten_to_message_pieces from pyrit.prompt_target import OpenAIResponseTarget, PromptTarget -from pyrit.prompt_target.common.json_response_config import _JsonResponseConfig def create_mock_response(response_dict: dict = None) -> MagicMock: @@ -212,13 +211,13 @@ async def test_construct_request_body_includes_extra_body_params( request = Message(message_pieces=[dummy_text_message_piece]) - jrc = _JsonResponseConfig.from_metadata(metadata=None) + jrc = JsonResponseConfig.from_metadata(metadata=None) body = await target._construct_request_body_async(conversation=[request], json_config=jrc) assert body["key"] == "value" async def test_construct_request_body_json_object(target: OpenAIResponseTarget, dummy_text_message_piece: MessagePiece): - json_response_config = _JsonResponseConfig(enabled=True) + json_response_config = JsonResponseConfig(enabled=True) request = Message(message_pieces=[dummy_text_message_piece]) body = await target._construct_request_body_async(conversation=[request], json_config=json_response_config) @@ -227,7 +226,7 @@ async def test_construct_request_body_json_object(target: OpenAIResponseTarget, 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( + json_response_config = JsonResponseConfig.from_metadata( metadata={"response_format": "json", "json_schema": schema_object} ) request = Message(message_pieces=[dummy_text_message_piece]) @@ -248,7 +247,7 @@ async def test_construct_request_body_removes_empty_values( ): request = Message(message_pieces=[dummy_text_message_piece]) - json_response_config = _JsonResponseConfig(enabled=False) + json_response_config = JsonResponseConfig(enabled=False) body = await target._construct_request_body_async(conversation=[request], json_config=json_response_config) assert "max_completion_tokens" not in body assert "max_tokens" not in body @@ -264,7 +263,7 @@ async def test_construct_request_body_serializes_text_message( ): request = Message(message_pieces=[dummy_text_message_piece]) - jrc = _JsonResponseConfig.from_metadata(metadata=None) + jrc = JsonResponseConfig.from_metadata(metadata=None) body = await target._construct_request_body_async(conversation=[request], json_config=jrc) assert body["input"][0]["content"][0]["text"] == "dummy text" @@ -276,7 +275,7 @@ async def test_construct_request_body_serializes_complex_message( 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) + jrc = JsonResponseConfig.from_metadata(metadata=None) body = await target._construct_request_body_async(conversation=[request], json_config=jrc) messages = body["input"][0]["content"] @@ -679,7 +678,7 @@ async def test_construct_request_body_filters_none( target: OpenAIResponseTarget, dummy_text_message_piece: MessagePiece ): req = Message(message_pieces=[dummy_text_message_piece]) - jrc = _JsonResponseConfig.from_metadata(metadata=None) + jrc = JsonResponseConfig.from_metadata(metadata=None) body = await target._construct_request_body_async(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 @@ -1270,7 +1269,7 @@ async def test_construct_request_body_includes_reasoning_effort( reasoning_effort="medium", ) request = Message(message_pieces=[dummy_text_message_piece]) - jrc = _JsonResponseConfig.from_metadata(metadata=None) + jrc = JsonResponseConfig.from_metadata(metadata=None) body = await target._construct_request_body_async(conversation=[request], json_config=jrc) assert body["reasoning"] == {"effort": "medium"} @@ -1285,7 +1284,7 @@ async def test_construct_request_body_includes_reasoning_summary( reasoning_summary="detailed", ) request = Message(message_pieces=[dummy_text_message_piece]) - jrc = _JsonResponseConfig.from_metadata(metadata=None) + jrc = JsonResponseConfig.from_metadata(metadata=None) body = await target._construct_request_body_async(conversation=[request], json_config=jrc) assert body["reasoning"] == {"summary": "detailed"} @@ -1301,7 +1300,7 @@ async def test_construct_request_body_includes_reasoning_effort_and_summary( reasoning_summary="auto", ) request = Message(message_pieces=[dummy_text_message_piece]) - jrc = _JsonResponseConfig.from_metadata(metadata=None) + jrc = JsonResponseConfig.from_metadata(metadata=None) body = await target._construct_request_body_async(conversation=[request], json_config=jrc) assert body["reasoning"] == {"effort": "high", "summary": "auto"} @@ -1310,7 +1309,7 @@ async def test_construct_request_body_omits_reasoning_when_not_set( target: OpenAIResponseTarget, dummy_text_message_piece: MessagePiece ): request = Message(message_pieces=[dummy_text_message_piece]) - jrc = _JsonResponseConfig.from_metadata(metadata=None) + jrc = JsonResponseConfig.from_metadata(metadata=None) body = await target._construct_request_body_async(conversation=[request], json_config=jrc) assert "reasoning" not in body diff --git a/tests/unit/score/test_callable_response_handler.py b/tests/unit/score/test_callable_response_handler.py index c53516cf4b..a3e7c6157d 100644 --- a/tests/unit/score/test_callable_response_handler.py +++ b/tests/unit/score/test_callable_response_handler.py @@ -10,13 +10,13 @@ from pyrit.score import CallableResponseHandler, JsonSchemaResponseHandler -def test_json_schema_handler_response_format_is_json(): - assert JsonSchemaResponseHandler().response_format == "json" +def test_json_schema_handler_requests_json(): + assert JsonSchemaResponseHandler().json_response_config.enabled is True -def test_callable_handler_response_format_is_none(): +def test_callable_handler_requests_no_format(): # The escape hatch imposes no wire format so plain-text classifiers are not forced into JSON. - assert CallableResponseHandler(parser=lambda _text: {}).response_format is None + assert CallableResponseHandler(parser=lambda _text: {}).json_response_config.enabled is False def test_callable_response_handler_parses_dict_from_callable(): diff --git a/tests/unit/score/test_llamaguard_pilot.py b/tests/unit/score/test_llamaguard_pilot.py index 6c6705b099..85bb1db9e6 100644 --- a/tests/unit/score/test_llamaguard_pilot.py +++ b/tests/unit/score/test_llamaguard_pilot.py @@ -18,8 +18,7 @@ from unit.mocks import get_mock_target_identifier from pyrit.exceptions.exception_classes import InvalidJsonException -from pyrit.models import Message, MessagePiece -from pyrit.models.json_schema_definition import JSON_SCHEMA_METADATA_KEY +from pyrit.models import JSON_SCHEMA_METADATA_KEY, Message, MessagePiece from pyrit.score import CallableResponseHandler, SelfAskTrueFalseScorer, TrueFalseQuestion # A compact stand-in for a canonical LlamaGuard system prompt. The real prompt is long and static; diff --git a/tests/unit/score/test_scorer_response_json_schema.py b/tests/unit/score/test_scorer_response_json_schema.py index 0746d9125d..d0c56038ee 100644 --- a/tests/unit/score/test_scorer_response_json_schema.py +++ b/tests/unit/score/test_scorer_response_json_schema.py @@ -14,8 +14,7 @@ import pytest from unit.mocks import get_mock_target_identifier -from pyrit.models import COMMON_JSON_SCHEMAS, Message, MessagePiece -from pyrit.models.json_schema_definition import JSON_SCHEMA_METADATA_KEY +from pyrit.models import COMMON_JSON_SCHEMAS, JSON_SCHEMA_METADATA_KEY, Message, MessagePiece from pyrit.score import ( ContentClassifier, ContentClassifierPaths, @@ -78,12 +77,12 @@ def _loaded_schema(scorer): """Return the response schema regardless of where the scorer stores it. The composition-migrated true/false scorer keeps it on its response handler - (``_response_handler.response_schema``); the other, not-yet-migrated scorers still expose - ``_response_json_schema`` directly. + (``_response_handler.json_response_config.json_schema``); the other, not-yet-migrated scorers + still expose ``_response_json_schema`` directly. """ handler = getattr(scorer, "_response_handler", None) if handler is not None: - return handler.response_schema + return handler.json_response_config.json_schema return scorer._response_json_schema diff --git a/tests/unit/score/test_self_ask_refusal.py b/tests/unit/score/test_self_ask_refusal.py index 3dc8c5b76e..332f4a8511 100644 --- a/tests/unit/score/test_self_ask_refusal.py +++ b/tests/unit/score/test_self_ask_refusal.py @@ -13,8 +13,14 @@ from pyrit.exceptions.exception_classes import InvalidJsonException from pyrit.memory import CentralMemory from pyrit.memory.memory_interface import MemoryInterface -from pyrit.models import COMMON_JSON_SCHEMAS, JSON_SCHEMA_METADATA_KEY, Message, MessagePiece, SeedPrompt -from pyrit.prompt_target.common.json_response_config import _JsonResponseConfig +from pyrit.models import ( + COMMON_JSON_SCHEMAS, + JSON_SCHEMA_METADATA_KEY, + JsonResponseConfig, + Message, + MessagePiece, + SeedPrompt, +) from pyrit.score import JsonSchemaResponseHandler, RefusalScorerPaths, SelfAskRefusalScorer @@ -274,8 +280,8 @@ async def test_refusal_scorer_loads_response_json_schema(scorer_path: RefusalSco system_prompt=SeedPrompt.from_yaml_file(scorer_path.value), ) - assert scorer._response_handler.response_schema is not None - assert scorer._response_handler.response_schema == EXPECTED_REFUSAL_RESPONSE_JSON_SCHEMA + assert scorer._response_handler.json_response_config.json_schema is not None + assert scorer._response_handler.json_response_config.json_schema == EXPECTED_REFUSAL_RESPONSE_JSON_SCHEMA async def test_refusal_scorer_passes_response_json_schema_to_target( @@ -334,7 +340,7 @@ async def test_refusal_scorer_identifier_includes_schema(scorer_path: RefusalSco async def test_refusal_scorer_metadata_round_trips_through_json_response_config( scorer_true_false_response: Message, patch_central_database ): - """The prompt_metadata produced by the scorer must be consumable by _JsonResponseConfig.""" + """The prompt_metadata produced by the scorer must be consumable by JsonResponseConfig.""" chat_target = MagicMock() chat_target.get_identifier.return_value = get_mock_target_identifier("MockChatTarget") chat_target.send_prompt_async = AsyncMock(return_value=[scorer_true_false_response]) @@ -345,7 +351,7 @@ async def test_refusal_scorer_metadata_round_trips_through_json_response_config( _, kwargs = chat_target.send_prompt_async.call_args metadata = kwargs["message"].message_pieces[0].prompt_metadata - config = _JsonResponseConfig.from_metadata(metadata=metadata) + config = JsonResponseConfig.from_metadata(metadata=metadata) assert config.enabled is True assert config.json_schema == EXPECTED_REFUSAL_RESPONSE_JSON_SCHEMA diff --git a/tests/unit/score/test_self_ask_true_false.py b/tests/unit/score/test_self_ask_true_false.py index 45a006e710..001bfd7814 100644 --- a/tests/unit/score/test_self_ask_true_false.py +++ b/tests/unit/score/test_self_ask_true_false.py @@ -231,7 +231,7 @@ def test_init_static_str_system_prompt(patch_central_database): ) assert scorer._system_prompt == "a static classifier prompt" - assert scorer._response_handler.response_schema is None + assert scorer._response_handler.json_response_config.json_schema is None assert scorer._score_category == ["harm"] @@ -253,7 +253,7 @@ def test_init_static_seed_prompt_preserves_schema(patch_central_database): ) assert scorer._system_prompt == "static seed prompt" - assert scorer._response_handler.response_schema == {"type": "object"} + assert scorer._response_handler.json_response_config.json_schema == {"type": "object"} def test_init_default_system_prompt_uses_task_achieved(patch_central_database): @@ -264,7 +264,7 @@ def test_init_default_system_prompt_uses_task_achieved(patch_central_database): scorer = SelfAskTrueFalseScorer(chat_target=chat_target) assert scorer._score_category == ["task_achieved"] - assert scorer._response_handler.response_schema is not None + assert scorer._response_handler.json_response_config.json_schema is not None assert "# Instructions" in scorer._system_prompt @@ -285,7 +285,7 @@ def test_init_templated_seed_prompt_from_separate_files(patch_central_database): assert scorer._score_category == ["grounded"] assert "# Instructions" in scorer._system_prompt # The schema embedded in the template survives model_copy during rendering. - assert scorer._response_handler.response_schema is not None + assert scorer._response_handler.json_response_config.json_schema is not None async def test_init_scores_end_to_end(patch_central_database, scorer_true_false_response: Message):