Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion doc/code/memory/5_advanced_memory.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
6 changes: 3 additions & 3 deletions doc/code/memory/6_azure_sql_memory.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down
4 changes: 2 additions & 2 deletions doc/code/targets/4_openai_video_target.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down Expand Up @@ -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}")
2 changes: 1 addition & 1 deletion doc/code/targets/6_1_target_capabilities.py
Original file line number Diff line number Diff line change
Expand Up @@ -304,7 +304,7 @@
def _ok_response():
return [
Message(
[
message_pieces=[
MessagePiece(
role="assistant",
original_value="ok",
Expand Down
2 changes: 1 addition & 1 deletion pyrit/message_normalizer/json_schema_normalizer.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
"""
Expand Down
4 changes: 2 additions & 2 deletions tests/integration/targets/test_entra_auth_targets.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand All @@ -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"


Expand Down
6 changes: 3 additions & 3 deletions tests/integration/targets/test_targets_and_secrets.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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"
Expand Down Expand Up @@ -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}"
Expand Down
20 changes: 10 additions & 10 deletions tests/unit/message_normalizer/test_json_schema_normalizer.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand All @@ -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}
Expand All @@ -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
Expand All @@ -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])

Expand All @@ -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}
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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.
Expand All @@ -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
Expand Down
8 changes: 4 additions & 4 deletions tests/unit/models/test_message.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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


Expand Down
30 changes: 16 additions & 14 deletions tests/unit/prompt_target/target/test_image_target.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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]
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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(
Expand All @@ -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(
Expand All @@ -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(
Expand All @@ -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(
Expand All @@ -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(
Expand All @@ -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(
Expand All @@ -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
Expand Down Expand Up @@ -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"
Expand Down Expand Up @@ -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]
Expand All @@ -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]
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -72,7 +72,7 @@ async def test_prompt_shield_reject_non_text(
promptshield_target: PromptShieldTarget, audio_message_piece: MessagePiece
):
with pytest.raises(ValueError):
await promptshield_target.send_prompt_async(message=Message([audio_message_piece]))
await promptshield_target.send_prompt_async(message=Message(message_pieces=[audio_message_piece]))


async def test_prompt_shield_document_parsing(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -132,7 +132,7 @@ async def test_send_prompt_async(

message_piece = sample_entries[0]
message_piece.converted_value = "Test content"
request = Message([message_piece])
request = Message(message_pieces=[message_piece])

response = await azure_blob_storage_target.send_prompt_async(message=request)

Expand Down
Loading
Loading