MAINT: Remove deprecated message init in JSON schema normalizer (#2091)

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Roman Lutz <romanlutz13@gmail.com>

Author: 3 people

doc/code/memory/5_advanced_memory.py

@@ -263,7 +263,7 @@
 )
 
 # Wrap each piece in a Message so we can pass it to score_async
-assistant_messages = [Message([piece]) for piece in assistant_pieces]
+assistant_messages = [Message(message_pieces=[piece]) for piece in assistant_pieces]
 
 # Score every response with both scorers — scores are automatically persisted in memory
 for msg in assistant_messages:

doc/code/memory/6_azure_sql_memory.py

@@ -73,9 +73,9 @@
     ),
 ]
 
-memory.add_message_to_memory(request=Message([message_list[0]]))
-memory.add_message_to_memory(request=Message([message_list[1]]))
-memory.add_message_to_memory(request=Message([message_list[2]]))
+memory.add_message_to_memory(request=Message(message_pieces=[message_list[0]]))
+memory.add_message_to_memory(request=Message(message_pieces=[message_list[1]]))
+memory.add_message_to_memory(request=Message(message_pieces=[message_list[2]]))
 
 entries = memory.get_conversation_messages(conversation_id=conversation_id)
 

doc/code/targets/4_openai_video_target.py

@@ -148,7 +148,7 @@
     original_value="Make it a watercolor painting style",
     prompt_metadata={"video_id": video_id},
 )
-remix_result = await video_target.send_prompt_async(message=Message([remix_piece]))  # type: ignore
+remix_result = await video_target.send_prompt_async(message=Message(message_pieces=[remix_piece]))  # type: ignore
 print(f"Remixed video: {remix_result[0].message_pieces[0].converted_value}")
 
 # %% [markdown]
@@ -190,5 +190,5 @@
     converted_value_data_type="image_path",
     conversation_id=conversation_id,
 )
-result = await i2v_target.send_prompt_async(message=Message([text_piece, image_piece]))  # type: ignore
+result = await i2v_target.send_prompt_async(message=Message(message_pieces=[text_piece, image_piece]))  # type: ignore
 print(f"Text+Image-to-video result: {result[0].message_pieces[0].converted_value}")

doc/code/targets/6_1_target_capabilities.py

@@ -304,7 +304,7 @@
 def _ok_response():
     return [
         Message(
-            [
+            message_pieces=[
                 MessagePiece(
                     role="assistant",
                     original_value="ok",

pyrit/message_normalizer/json_schema_normalizer.py

@@ -126,7 +126,7 @@ def _adapt_message(self, *, message: Message) -> Message:
         if not changed:
             return message
 
-        return Message(new_pieces)
+        return Message(message_pieces=new_pieces)
 
     def _adapt_piece(self, *, piece: MessagePiece) -> MessagePiece:
         """

tests/integration/targets/test_entra_auth_targets.py

@@ -332,7 +332,7 @@ async def test_video_target_remix_entra_auth(sqlite_instance):
         original_value="A bird flying over a lake",
         converted_value="A bird flying over a lake",
     )
-    result = await target.send_prompt_async(message=Message([text_piece]))
+    result = await target.send_prompt_async(message=Message(message_pieces=[text_piece]))
     response_piece = result[0].message_pieces[0]
     assert response_piece.response_error == "none"
     video_id = response_piece.prompt_metadata.get("video_id")
@@ -345,7 +345,7 @@ async def test_video_target_remix_entra_auth(sqlite_instance):
         converted_value="Add a sunset",
         prompt_metadata={"video_id": video_id},
     )
-    remix_result = await target.send_prompt_async(message=Message([remix_piece]))
+    remix_result = await target.send_prompt_async(message=Message(message_pieces=[remix_piece]))
     assert remix_result[0].message_pieces[0].response_error == "none"
 
 

tests/integration/targets/test_targets_and_secrets.py

@@ -571,7 +571,7 @@ async def test_video_remix_chain(sqlite_instance):
         original_value="A cat sitting on a windowsill",
         converted_value="A cat sitting on a windowsill",
     )
-    result = await target.send_prompt_async(message=Message([text_piece]))
+    result = await target.send_prompt_async(message=Message(message_pieces=[text_piece]))
     assert len(result) == 1
     response_piece = result[0].message_pieces[0]
     assert response_piece.response_error == "none"
@@ -586,7 +586,7 @@ async def test_video_remix_chain(sqlite_instance):
         converted_value="Make it a watercolor painting style",
         prompt_metadata={"video_id": video_id},
     )
-    remix_result = await target.send_prompt_async(message=Message([remix_piece]))
+    remix_result = await target.send_prompt_async(message=Message(message_pieces=[remix_piece]))
     assert len(remix_result) == 1
     remix_response = remix_result[0].message_pieces[0]
     assert remix_response.response_error == "none"
@@ -636,7 +636,7 @@ async def test_video_image_to_video(sqlite_instance):
         converted_value_data_type="image_path",
         conversation_id=conversation_id,
     )
-    result = await target.send_prompt_async(message=Message([text_piece, image_piece]))
+    result = await target.send_prompt_async(message=Message(message_pieces=[text_piece, image_piece]))
     assert len(result) == 1
     response_piece = result[0].message_pieces[0]
     assert response_piece.response_error == "none", f"Image-to-video failed: {response_piece.converted_value}"

tests/unit/message_normalizer/test_json_schema_normalizer.py

@@ -46,7 +46,7 @@ class TestJsonSchemaNormalizer:
     async def test_text_piece_gets_schema_appended_to_converted_value(self, normalizer: JsonSchemaNormalizer) -> None:
         schema = {"type": "object", "properties": {"answer": {"type": "string"}}}
         piece = _text_piece(value="Answer the question.", metadata={JSON_SCHEMA_METADATA_KEY: schema})
-        message = Message([piece])
+        message = Message(message_pieces=[piece])
 
         result = await normalizer.normalize_async([message])
         out_piece = result[0].message_pieces[0]
@@ -66,7 +66,7 @@ async def test_text_piece_preserves_other_metadata(self, normalizer: JsonSchemaN
                 "other": 7,
             },
         )
-        result = await normalizer.normalize_async([Message([piece])])
+        result = await normalizer.normalize_async([Message(message_pieces=[piece])])
         new_metadata = result[0].message_pieces[0].prompt_metadata
         assert JSON_SCHEMA_METADATA_KEY not in new_metadata
         assert new_metadata == {"response_format": "json", "other": 7}
@@ -76,7 +76,7 @@ async def test_non_text_piece_only_strips_key(self, normalizer: JsonSchemaNormal
         piece = _image_piece(value="fake.jpg", metadata={JSON_SCHEMA_METADATA_KEY: schema, "extra": "stay"})
         original_converted_value = piece.converted_value
 
-        result = await normalizer.normalize_async([Message([piece])])
+        result = await normalizer.normalize_async([Message(message_pieces=[piece])])
         out_piece = result[0].message_pieces[0]
 
         assert JSON_SCHEMA_METADATA_KEY not in out_piece.prompt_metadata
@@ -87,7 +87,7 @@ async def test_non_text_piece_only_strips_key(self, normalizer: JsonSchemaNormal
 
     async def test_no_schema_is_noop(self, normalizer: JsonSchemaNormalizer) -> None:
         piece = _text_piece(value="just say hi", metadata={"unrelated": True})
-        message = Message([piece])
+        message = Message(message_pieces=[piece])
 
         result = await normalizer.normalize_async([message])
 
@@ -98,7 +98,7 @@ async def test_input_pieces_not_mutated(self, normalizer: JsonSchemaNormalizer)
         schema = {"type": "object"}
         piece = _text_piece(value="hi", metadata={JSON_SCHEMA_METADATA_KEY: schema})
 
-        await normalizer.normalize_async([Message([piece])])
+        await normalizer.normalize_async([Message(message_pieces=[piece])])
 
         # The original piece still carries the schema and its unchanged text.
         assert piece.prompt_metadata == {JSON_SCHEMA_METADATA_KEY: schema}
@@ -117,7 +117,7 @@ async def test_mixed_pieces_in_message_each_handled(self, normalizer: JsonSchema
         )
         no_schema_piece = _text_piece(value="z", metadata={"foo": "bar"}, conversation_id=conversation_id)
 
-        result = await normalizer.normalize_async([Message([text_piece, image_piece, no_schema_piece])])
+        result = await normalizer.normalize_async([Message(message_pieces=[text_piece, image_piece, no_schema_piece])])
         out_pieces = result[0].message_pieces
 
         assert JSON_SCHEMA_METADATA_KEY not in out_pieces[0].prompt_metadata
@@ -133,8 +133,8 @@ async def test_mixed_pieces_in_message_each_handled(self, normalizer: JsonSchema
 
     async def test_multiple_messages(self, normalizer: JsonSchemaNormalizer) -> None:
         schema = {"type": "object"}
-        msg_with_schema = Message([_text_piece(value="a", metadata={JSON_SCHEMA_METADATA_KEY: schema})])
-        msg_without_schema = Message([_text_piece(value="b", metadata={})])
+        msg_with_schema = Message(message_pieces=[_text_piece(value="a", metadata={JSON_SCHEMA_METADATA_KEY: schema})])
+        msg_without_schema = Message(message_pieces=[_text_piece(value="b", metadata={})])
 
         result = await normalizer.normalize_async([msg_with_schema, msg_without_schema])
         assert "### Response format" in result[0].message_pieces[0].converted_value
@@ -155,7 +155,7 @@ async def test_appended_text_lists_schema_keys(self, normalizer: JsonSchemaNorma
         }
         piece = _text_piece(value="prompt", metadata={JSON_SCHEMA_METADATA_KEY: schema})
 
-        result = await normalizer.normalize_async([Message([piece])])
+        result = await normalizer.normalize_async([Message(message_pieces=[piece])])
         appended = result[0].message_pieces[0].converted_value
 
         # Sanity-check that the rendered text actually surfaces schema field names.
@@ -170,7 +170,7 @@ async def test_custom_template_is_used(self) -> None:
         schema = {"type": "object"}
         piece = _text_piece(value="hi", metadata={JSON_SCHEMA_METADATA_KEY: schema})
 
-        result = await normalizer.normalize_async([Message([piece])])
+        result = await normalizer.normalize_async([Message(message_pieces=[piece])])
         out_value = result[0].message_pieces[0].converted_value
 
         assert "<<SCHEMA START>>" in out_value

tests/unit/models/test_message.py

@@ -73,7 +73,7 @@ def test_get_pieces_by_type_returns_matching_pieces() -> None:
         converted_value_data_type="image_path",
         conversation_id=conversation_id,
     )
-    msg = Message([text_piece, image_piece])
+    msg = Message(message_pieces=[text_piece, image_piece])
 
     result = msg.get_pieces_by_type(data_type="text")
     assert len(result) == 1
@@ -86,21 +86,21 @@ def test_get_pieces_by_type_returns_matching_pieces() -> None:
 
 def test_get_pieces_by_type_returns_empty_for_no_match() -> None:
     piece = MessagePiece(role="user", original_value="hello", converted_value="hello")
-    msg = Message([piece])
+    msg = Message(message_pieces=[piece])
     assert msg.get_pieces_by_type(data_type="image_path") == []
 
 
 def test_get_piece_by_type_returns_first_match() -> None:
     conversation_id = "test-conv"
     text1 = MessagePiece(role="user", original_value="a", converted_value="a", conversation_id=conversation_id)
     text2 = MessagePiece(role="user", original_value="b", converted_value="b", conversation_id=conversation_id)
-    msg = Message([text1, text2])
+    msg = Message(message_pieces=[text1, text2])
     assert msg.get_piece_by_type(data_type="text") is text1
 
 
 def test_get_piece_by_type_returns_none_for_no_match() -> None:
     piece = MessagePiece(role="user", original_value="hello", converted_value="hello")
-    msg = Message([piece])
+    msg = Message(message_pieces=[piece])
     assert msg.get_piece_by_type(data_type="image_path") is None
 
 

tests/unit/prompt_target/target/test_image_target.py

@@ -80,7 +80,7 @@ async def test_send_prompt_async_generate(
     with patch.object(image_target._async_client.images, "generate", new_callable=AsyncMock) as mock_generate:
         mock_generate.return_value = mock_response
 
-        resp = await image_target.send_prompt_async(message=Message([request]))
+        resp = await image_target.send_prompt_async(message=Message(message_pieces=[request]))
         assert len(resp) == 1
         assert resp
         path = resp[0].message_pieces[0].original_value
@@ -115,7 +115,7 @@ async def test_send_prompt_async_edit(
     with patch.object(image_target._async_client.images, "edit", new_callable=AsyncMock) as mock_edit:
         mock_edit.return_value = mock_response
 
-        resp = await image_target.send_prompt_async(message=Message([text_piece, image_piece]))
+        resp = await image_target.send_prompt_async(message=Message(message_pieces=[text_piece, image_piece]))
         assert len(resp) == 1
         assert resp
         path = resp[0].message_pieces[0].original_value
@@ -151,7 +151,7 @@ async def test_send_prompt_async_edit_single_image_passes_tuple_not_list(
     with patch.object(image_target._async_client.images, "edit", new_callable=AsyncMock) as mock_edit:
         mock_edit.return_value = mock_response
 
-        resp = await image_target.send_prompt_async(message=Message([text_piece, image_piece]))
+        resp = await image_target.send_prompt_async(message=Message(message_pieces=[text_piece, image_piece]))
         assert resp
 
         call_kwargs = mock_edit.call_args[1]
@@ -189,7 +189,9 @@ async def test_send_prompt_async_edit_multiple_images(
     with patch.object(image_target._async_client.images, "edit", new_callable=AsyncMock) as mock_edit:
         mock_edit.return_value = mock_response
 
-        resp = await image_target.send_prompt_async(message=Message([image_piece, text_piece] + image_pieces))
+        resp = await image_target.send_prompt_async(
+            message=Message(message_pieces=[image_piece, text_piece] + image_pieces)
+        )
         assert len(resp) == 1
         assert resp
         path = resp[0].message_pieces[0].original_value
@@ -226,7 +228,7 @@ async def test_send_prompt_async_invalid_image_path(
     )
 
     with pytest.raises(FileNotFoundError):
-        await image_target.send_prompt_async(message=Message([text_piece, image_piece]))
+        await image_target.send_prompt_async(message=Message(message_pieces=[text_piece, image_piece]))
 
 
 async def test_send_prompt_async_empty_response(
@@ -248,7 +250,7 @@ async def test_send_prompt_async_empty_response(
         mock_generate.return_value = mock_response
 
         with pytest.raises(EmptyResponseException):
-            await image_target.send_prompt_async(message=Message([request]))
+            await image_target.send_prompt_async(message=Message(message_pieces=[request]))
 
 
 async def test_send_prompt_async_rate_limit_exception(
@@ -264,7 +266,7 @@ async def test_send_prompt_async_rate_limit_exception(
         mock_generate.side_effect = RateLimitError("Rate Limit Reached", response=MagicMock(), body={})
 
         with pytest.raises(RateLimitException):
-            await image_target.send_prompt_async(message=Message([request]))
+            await image_target.send_prompt_async(message=Message(message_pieces=[request]))
 
 
 async def test_send_prompt_async_bad_request_error(
@@ -290,7 +292,7 @@ async def test_send_prompt_async_bad_request_error(
 
         # Non-content-filter BadRequestError should be re-raised (same as chat target behavior)
         with pytest.raises(Exception):
-            await image_target.send_prompt_async(message=Message([request]))
+            await image_target.send_prompt_async(message=Message(message_pieces=[request]))
 
 
 async def test_send_prompt_async_empty_response_adds_memory(
@@ -316,7 +318,7 @@ async def test_send_prompt_async_empty_response_adds_memory(
         image_target._memory = mock_memory
 
         with pytest.raises(EmptyResponseException):
-            await image_target.send_prompt_async(message=Message([request]))
+            await image_target.send_prompt_async(message=Message(message_pieces=[request]))
 
 
 async def test_send_prompt_async_rate_limit_adds_memory(
@@ -338,7 +340,7 @@ async def test_send_prompt_async_rate_limit_adds_memory(
         image_target._memory = mock_memory
 
         with pytest.raises(RateLimitException):
-            await image_target.send_prompt_async(message=Message([request]))
+            await image_target.send_prompt_async(message=Message(message_pieces=[request]))
 
 
 async def test_send_prompt_async_bad_request_content_filter(
@@ -364,7 +366,7 @@ async def test_send_prompt_async_bad_request_content_filter(
 
     with patch.object(image_target._async_client.images, "generate", new_callable=AsyncMock) as mock_generate:
         mock_generate.side_effect = bad_request_error
-        result = await image_target.send_prompt_async(message=Message([request]))
+        result = await image_target.send_prompt_async(message=Message(message_pieces=[request]))
         assert len(result) == 1
         assert result[0].message_pieces[0].converted_value_data_type == "error"
         assert "content_filter" in result[0].message_pieces[0].converted_value
@@ -393,7 +395,7 @@ async def test_send_prompt_async_bad_request_content_policy_violation(
 
     with patch.object(image_target._async_client.images, "generate", new_callable=AsyncMock) as mock_generate:
         mock_generate.side_effect = bad_request_error
-        result = await image_target.send_prompt_async(message=Message([request]))
+        result = await image_target.send_prompt_async(message=Message(message_pieces=[request]))
         assert len(result) == 1
         assert result[0].message_pieces[0].response_error == "blocked"
         assert result[0].message_pieces[0].converted_value_data_type == "error"
@@ -542,7 +544,7 @@ async def test_generate_request_passes_background(
     with patch.object(image_target._async_client.images, "generate", new_callable=AsyncMock) as mock_generate:
         mock_generate.return_value = mock_response
 
-        resp = await image_target.send_prompt_async(message=Message([request]))
+        resp = await image_target.send_prompt_async(message=Message(message_pieces=[request]))
         assert resp
 
         call_kwargs = mock_generate.call_args[1]
@@ -569,7 +571,7 @@ async def test_generate_request_omits_background_when_none(
     with patch.object(image_target._async_client.images, "generate", new_callable=AsyncMock) as mock_generate:
         mock_generate.return_value = mock_response
 
-        resp = await image_target.send_prompt_async(message=Message([request]))
+        resp = await image_target.send_prompt_async(message=Message(message_pieces=[request]))
         assert resp
 
         call_kwargs = mock_generate.call_args[1]