diff --git a/pyrit/message_normalizer/json_schema_normalizer.py b/pyrit/message_normalizer/json_schema_normalizer.py index f7acc1cfe2..4937696041 100644 --- a/pyrit/message_normalizer/json_schema_normalizer.py +++ b/pyrit/message_normalizer/json_schema_normalizer.py @@ -4,7 +4,7 @@ import json from pyrit.message_normalizer.message_normalizer import MessageListNormalizer -from pyrit.models import JSON_SCHEMA_METADATA_KEY, JsonSchemaDefinition, Message, MessagePiece +from pyrit.models import JSON_SCHEMA_METADATA_KEY, JsonResponseConfig, JsonSchemaDefinition, Message, MessagePiece # Default text appended to a message piece when the target cannot natively # enforce a JSON schema. The exact phrasing mirrors the legacy in-prompt @@ -35,9 +35,10 @@ class JsonSchemaNormalizer(MessageListNormalizer[Message]): For every message piece carrying a schema in ``JSON_SCHEMA_METADATA_KEY``: - * If the piece is text, the schema is rendered into prompt text appended - to the existing value so the model is still instructed to produce a - conforming JSON object via prompt engineering. + * If the piece is text, the schema (a dict or a JSON string) is rendered + into prompt text appended to the existing value so the model is still + instructed to produce a conforming JSON object via prompt engineering. + An empty string means no schema, so nothing is appended. * Otherwise (non-text pieces — image, audio, video, etc.) the schema metadata is simply removed; non-text modalities have no natural place to embed a textual schema instruction. @@ -132,7 +133,7 @@ def _adapt_piece(self, *, piece: MessagePiece) -> MessagePiece: """ Return a new piece with the schema metadata key removed. - For text pieces, the schema is additionally rendered into the piece's + For text pieces, a non-empty schema is additionally rendered into the piece's ``converted_value`` so the model is still instructed (via prompt text) to produce JSON conforming to the schema. For non-text pieces the schema is simply dropped from metadata. @@ -145,12 +146,13 @@ def _adapt_piece(self, *, piece: MessagePiece) -> MessagePiece: schema rendered into the text value when ``converted_value_data_type`` is ``"text"``. """ - schema = piece.prompt_metadata[JSON_SCHEMA_METADATA_KEY] new_metadata = {key: value for key, value in piece.prompt_metadata.items() if key != JSON_SCHEMA_METADATA_KEY} updates: dict[str, object] = {"prompt_metadata": new_metadata} if piece.converted_value_data_type == "text": - updates["converted_value"] = self._append_schema_instructions(text=piece.converted_value, schema=schema) + schema = JsonResponseConfig.parse_schema(piece.prompt_metadata[JSON_SCHEMA_METADATA_KEY]) + if schema is not None: + updates["converted_value"] = self._append_schema_instructions(text=piece.converted_value, schema=schema) return piece.model_copy(update=updates) diff --git a/pyrit/models/target/json_response_config.py b/pyrit/models/target/json_response_config.py index 763fce93d4..9a78700e23 100644 --- a/pyrit/models/target/json_response_config.py +++ b/pyrit/models/target/json_response_config.py @@ -85,23 +85,36 @@ def from_metadata(cls, *, metadata: dict[str, Any] | None) -> JsonResponseConfig 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, + json_schema=cls.parse_schema(schema_val), schema_name=metadata.get(cls._METADATAKEYS["JSON_SCHEMA_NAME"], "CustomSchema"), strict=metadata.get(cls._METADATAKEYS["JSON_SCHEMA_STRICT"], True), ) return cls(enabled=True) + @staticmethod + def parse_schema(schema_val: dict[str, Any] | str) -> dict[str, Any] | None: + """ + Read a schema stored under ``JSON_SCHEMA_METADATA_KEY``, which may be a dict or a JSON string. + + Args: + schema_val (dict[str, Any] | str): The stored schema. An empty string means no schema. + + Returns: + dict[str, Any] | None: The parsed schema, or ``None`` for an empty string. + + Raises: + ValueError: If a schema string is not valid JSON. + """ + if not isinstance(schema_val, str): + return schema_val + try: + return 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 + def to_metadata(self) -> dict[str, Any]: """ Serialize to the canonical ``response_format`` / ``json_schema`` metadata keys. diff --git a/tests/unit/message_normalizer/test_json_schema_normalizer.py b/tests/unit/message_normalizer/test_json_schema_normalizer.py index eed3ef7679..9f9a83994a 100644 --- a/tests/unit/message_normalizer/test_json_schema_normalizer.py +++ b/tests/unit/message_normalizer/test_json_schema_normalizer.py @@ -58,6 +58,32 @@ async def test_text_piece_gets_schema_appended_to_converted_value(self, normaliz # reason about its structure. assert json.dumps(schema, indent=2) in out_piece.converted_value + async def test_schema_stored_as_json_string_renders_like_dict(self, normalizer: JsonSchemaNormalizer) -> None: + schema = {"type": "object", "properties": {"answer": {"type": "string"}}, "required": ["answer"]} + as_dict = _text_piece(value="Answer.", metadata={JSON_SCHEMA_METADATA_KEY: schema}) + as_string = _text_piece(value="Answer.", metadata={JSON_SCHEMA_METADATA_KEY: json.dumps(schema)}) + + result = await normalizer.normalize_async( + [Message(message_pieces=[as_dict]), Message(message_pieces=[as_string])] + ) + + assert result[1].message_pieces[0].converted_value == result[0].message_pieces[0].converted_value + + async def test_invalid_json_string_schema_raises(self, normalizer: JsonSchemaNormalizer) -> None: + piece = _text_piece(value="Answer.", metadata={JSON_SCHEMA_METADATA_KEY: "{not json"}) + + with pytest.raises(ValueError, match="Invalid JSON schema provided"): + await normalizer.normalize_async([Message(message_pieces=[piece])]) + + async def test_empty_string_schema_is_stripped_without_instructions(self, normalizer: JsonSchemaNormalizer) -> None: + piece = _text_piece(value="Answer.", metadata={JSON_SCHEMA_METADATA_KEY: "", "response_format": "json"}) + + result = await normalizer.normalize_async([Message(message_pieces=[piece])]) + out_piece = result[0].message_pieces[0] + + assert out_piece.converted_value == "Answer." + assert out_piece.prompt_metadata == {"response_format": "json"} + async def test_text_piece_preserves_other_metadata(self, normalizer: JsonSchemaNormalizer) -> None: piece = _text_piece( metadata={