diff --git a/pyrit/models/score/scorable.py b/pyrit/models/score/scorable.py index e3de7f4474..10ac8848e7 100644 --- a/pyrit/models/score/scorable.py +++ b/pyrit/models/score/scorable.py @@ -5,7 +5,7 @@ import uuid # noqa: TC003 (runtime-required by Pydantic field annotations) from abc import ABC -from typing import TYPE_CHECKING, Annotated, Any, Literal, get_args +from typing import TYPE_CHECKING, Annotated, Any, Literal, Self, get_args from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, model_validator @@ -66,7 +66,7 @@ def _validate_message_piece_ids(self) -> MessageScorable: def from_message( cls, message: Message, - ) -> MessageScorable: + ) -> Self: """ Name the pieces of a persisted message. @@ -74,7 +74,7 @@ def from_message( message (Message): The message whose pieces to name. Returns: - MessageScorable: A scorable naming the message's pieces. + Self: A scorable naming the message's pieces. """ return cls(message_piece_ids=tuple(piece.id for piece in message.message_pieces)) @@ -93,7 +93,7 @@ class ContentScorable(Scorable): data_type: PromptDataType = "text" @classmethod - def from_message(cls, message: Message) -> ContentScorable: + def from_message(cls, message: Message) -> Self: """ Describe the converted content of a single-piece ephemeral message. @@ -107,7 +107,7 @@ def from_message(cls, message: Message) -> ContentScorable: message (Message): The ephemeral message whose converted content to take. Returns: - ContentScorable: A scorable holding the converted message content. + Self: A scorable holding the converted message content. """ piece = message.get_piece() return cls(value=piece.converted_value, data_type=piece.converted_value_data_type) diff --git a/pyrit/models/seeds/seed_dataset.py b/pyrit/models/seeds/seed_dataset.py index f57ff98dcb..f471ef44e3 100644 --- a/pyrit/models/seeds/seed_dataset.py +++ b/pyrit/models/seeds/seed_dataset.py @@ -12,7 +12,7 @@ import uuid from collections import defaultdict from datetime import UTC, datetime -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Self from pydantic import BaseModel, ConfigDict, Field, model_validator @@ -233,7 +233,7 @@ def get_random_values(self, *, number: PositiveInt, harm_categories: Sequence[st return random.sample(prompts, min(len(prompts), number)) @classmethod - def from_dict(cls, data: dict[str, Any]) -> SeedDataset: + def from_dict(cls, data: dict[str, Any]) -> Self: """ Build a SeedDataset, assigning per-seed ``prompt_group_id`` by alias. @@ -246,7 +246,7 @@ def from_dict(cls, data: dict[str, Any]) -> SeedDataset: data (dict[str, Any]): Dataset payload with top-level defaults and seed entries. Returns: - SeedDataset: Constructed dataset. + Self: Constructed dataset. Raises: ValueError: If any seed entry includes a pre-set ``prompt_group_id``. diff --git a/tests/unit/models/test_scorable.py b/tests/unit/models/test_scorable.py index 79f9cd5798..f5b6684f7d 100644 --- a/tests/unit/models/test_scorable.py +++ b/tests/unit/models/test_scorable.py @@ -2,6 +2,7 @@ # Licensed under the MIT license. import uuid +from typing import assert_type import pytest from pydantic import ValidationError @@ -80,6 +81,18 @@ def test_message_scorable_from_message_names_pieces(): assert not hasattr(scorable, "message") +def test_message_scorable_from_message_preserves_subclass_type(): + class CustomMessageScorable(MessageScorable): + pass + + message = _message() + scorable = CustomMessageScorable.from_message(message) + + assert_type(scorable, CustomMessageScorable) + assert type(scorable) is CustomMessageScorable + assert scorable.message_piece_ids == (message.get_piece().id,) + + def test_message_scorable_rejects_empty_ids(): with pytest.raises(ValueError, match="at least one message piece"): MessageScorable(message_piece_ids=()) @@ -118,6 +131,26 @@ def test_content_scorable_from_message_uses_converted_view(): assert scorable.data_type == "text" +def test_content_scorable_from_message_preserves_subclass_type(): + class CustomContentScorable(ContentScorable): + pass + + message = MessagePiece( + role="user", + original_value="original", + converted_value="converted", + original_value_data_type="text", + converted_value_data_type="text", + ).to_message() + + scorable = CustomContentScorable.from_message(message) + + assert_type(scorable, CustomContentScorable) + assert type(scorable) is CustomContentScorable + assert scorable.value == "converted" + assert scorable.data_type == "text" + + @pytest.mark.parametrize( "scorable", [ diff --git a/tests/unit/models/test_seed_dataset.py b/tests/unit/models/test_seed_dataset.py index 9b9e96e748..8988f4892e 100644 --- a/tests/unit/models/test_seed_dataset.py +++ b/tests/unit/models/test_seed_dataset.py @@ -2,6 +2,7 @@ # Licensed under the MIT license. import uuid +from typing import assert_type import pytest @@ -195,3 +196,14 @@ def test_seed_dataset_group_without_group_id(): def test_seed_dataset_group_rejects_unsupported_seed_type(): with pytest.raises(ValueError, match="Unsupported seed type: Seed"): SeedDataset.group_seed_prompts_by_prompt_group_id([Seed(value="unsupported")]) + + +def test_seed_dataset_from_dict_preserves_subclass_type(): + class CustomSeedDataset(SeedDataset): + pass + + ds = CustomSeedDataset.from_dict({"seeds": [{"value": "hello"}]}) + + assert_type(ds, CustomSeedDataset) + assert type(ds) is CustomSeedDataset + assert isinstance(ds.seeds[0], SeedPrompt)