diff --git a/frontend/src/types/index.ts b/frontend/src/types/index.ts index 5b156682c0..df61ef7497 100644 --- a/frontend/src/types/index.ts +++ b/frontend/src/types/index.ts @@ -1347,6 +1347,7 @@ export interface SeedRecord { parameters: string[] | null conditions: SeedCondition[] | null response_json_schema: Record | null + is_general_technique?: boolean } /** Wire shape of one logical example from `GET /api/datasets/seeds`. */ diff --git a/pyrit/memory/alembic/versions/e7a1b3c5d9f0_add_seed_is_general_technique.py b/pyrit/memory/alembic/versions/e7a1b3c5d9f0_add_seed_is_general_technique.py new file mode 100644 index 0000000000..d3b7c36c4b --- /dev/null +++ b/pyrit/memory/alembic/versions/e7a1b3c5d9f0_add_seed_is_general_technique.py @@ -0,0 +1,34 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +""" +Add is_general_technique flag to SeedPromptEntries. + +Existing rows default to False. Downgrading removes the column. + +Revision ID: e7a1b3c5d9f0 +Revises: c8d3e5f7a901 +Create Date: 2026-10-28 00:00:00.000000 +""" + +from collections.abc import Sequence # noqa: TC003 + +import sqlalchemy as sa +from alembic import op + +revision: str = "e7a1b3c5d9f0" +down_revision: str | None = "c8d3e5f7a901" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +def upgrade() -> None: + """Add is_general_technique storage and default existing rows to False.""" + with op.batch_alter_table("SeedPromptEntries") as batch_op: + batch_op.add_column(sa.Column("is_general_technique", sa.Boolean(), nullable=False, server_default=sa.false())) + + +def downgrade() -> None: + """Remove is_general_technique column.""" + with op.batch_alter_table("SeedPromptEntries") as batch_op: + batch_op.drop_column("is_general_technique", mssql_drop_default=True) diff --git a/pyrit/memory/memory_interface.py b/pyrit/memory/memory_interface.py index e230165ae0..de9ee276e3 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -4397,6 +4397,7 @@ def _build_seed_filter_conditions( source: str | None = None, seed_type: SeedType | None = None, origin: SeedOrigin | None = None, + is_general_technique: bool | None = None, parameters: Sequence[str] | None = None, metadata: dict[str, str | int] | None = None, prompt_group_ids: Sequence[uuid.UUID] | None = None, @@ -4409,6 +4410,7 @@ def _build_seed_filter_conditions( Args: origin (SeedOrigin | None): Match the recorded ingestion origin. + is_general_technique (bool | None): Match whether the seed represents a general attack technique. value (str): The value to match. By default this matches by substring; pass exact=True to require full-string equality instead. If None, all values are returned. exact (bool): When True, ``value`` is matched by full-string equality rather than substring, @@ -4472,6 +4474,9 @@ def _build_seed_filter_conditions( elif seed_type is not None: conditions.append(SeedEntry.seed_type == seed_type) + if is_general_technique is not None: + conditions.append(SeedEntry.is_general_technique == is_general_technique) + self._add_list_conditions( field=SeedEntry.harm_categories, values=harm_categories, conditions=conditions, exact=exact ) @@ -4564,6 +4569,7 @@ def _execute_get_seeds( source: str | None = None, seed_type: SeedType | None = None, origin: SeedOrigin | None = None, + is_general_technique: bool | None = None, parameters: Sequence[str] | None = None, metadata: dict[str, str | int] | None = None, prompt_group_ids: Sequence[uuid.UUID] | None = None, @@ -4573,6 +4579,7 @@ def _execute_get_seeds( Args: origin (SeedOrigin | None): Match the recorded ingestion origin. + is_general_technique (bool | None): Match whether the seed represents a general attack technique. value (str): The value to match by substring. If None, all values are returned. value_sha256 (Sequence[str] | None): A list of SHA256 hashes of values to match. If None, all values are returned. @@ -4616,6 +4623,7 @@ def _execute_get_seeds( source=source, seed_type=seed_type, origin=origin, + is_general_technique=is_general_technique, parameters=parameters, metadata=metadata, prompt_group_ids=prompt_group_ids, @@ -4647,6 +4655,7 @@ def _execute_remove_seeds_from_memory( source: str | None = None, seed_type: SeedType | None = None, origin: SeedOrigin | None = None, + is_general_technique: bool | None = None, parameters: Sequence[str] | None = None, metadata: dict[str, str | int] | None = None, prompt_group_ids: Sequence[uuid.UUID] | None = None, @@ -4664,6 +4673,7 @@ def _execute_remove_seeds_from_memory( Args: origin (SeedOrigin | None): Match the recorded ingestion origin. + is_general_technique (bool | None): Match whether the seed represents a general attack technique. value (str): The value to match. For the remove methods this defaults to full-string equality (exact=True) so a short or common value does not delete far more seeds than intended; pass exact=False to match by substring instead. If None, all values are considered. @@ -4718,6 +4728,7 @@ def _execute_remove_seeds_from_memory( source=source, seed_type=seed_type, origin=origin, + is_general_technique=is_general_technique, parameters=parameters, metadata=metadata, prompt_group_ids=prompt_group_ids, @@ -4758,6 +4769,7 @@ def _execute_remove_seed_groups_from_memory( source: str | None = None, seed_type: SeedType | None = None, origin: SeedOrigin | None = None, + is_general_technique: bool | None = None, parameters: Sequence[str] | None = None, metadata: dict[str, str | int] | None = None, prompt_group_ids: Sequence[uuid.UUID] | None = None, @@ -4783,6 +4795,8 @@ def _execute_remove_seed_groups_from_memory( Args: origin (SeedOrigin | None): Match origin before expanding to complete groups. + is_general_technique (bool | None): Match whether the seed represents a general attack technique + before expanding to complete groups. value (str): The value to match. For the remove methods this defaults to full-string equality (exact=True) so a short or common value does not delete far more seeds than intended; pass exact=False to match by substring instead. If None, all values are considered. @@ -4837,6 +4851,7 @@ def _execute_remove_seed_groups_from_memory( source=source, seed_type=seed_type, origin=origin, + is_general_technique=is_general_technique, parameters=parameters, metadata=metadata, prompt_group_ids=prompt_group_ids, @@ -5115,6 +5130,7 @@ def _seed_group_member_key(seed: Seed[Any]) -> str: "authors": entry.authors or [], "groups": entry.groups or [], "source": entry.source, + "is_general_technique": entry.is_general_technique, }, sort_keys=True, separators=(",", ":"), @@ -5778,6 +5794,7 @@ def _execute_get_seed_groups( source: str | None = None, seed_type: SeedType | None = None, origin: SeedOrigin | None = None, + is_general_technique: bool | None = None, parameters: Sequence[str] | None = None, metadata: dict[str, str | int] | None = None, prompt_group_ids: Sequence[uuid.UUID] | None = None, @@ -5788,6 +5805,8 @@ def _execute_get_seed_groups( Args: origin (SeedOrigin | None): Match origin before expanding to complete groups. + is_general_technique (bool | None, Optional): Match whether the seed represents a general attack technique + before expanding to complete groups. value (str | None, Optional): The value to match by substring. value_sha256 (Sequence[str] | None, Optional): SHA256 hash of value to filter seed groups by. dataset_name (str | None, Optional): Name of the dataset to match exactly. @@ -5828,6 +5847,7 @@ def _execute_get_seed_groups( source=source, seed_type=seed_type, origin=origin, + is_general_technique=is_general_technique, parameters=parameters, metadata=metadata, prompt_group_ids=prompt_group_ids, @@ -10302,6 +10322,7 @@ def get_seeds( source: str | None = None, seed_type: SeedType | None = None, origin: SeedOrigin | None = None, + is_general_technique: bool | None = None, parameters: Sequence[str] | None = None, metadata: dict[str, str | int] | None = None, prompt_group_ids: Sequence[uuid.UUID] | None = None, @@ -10330,6 +10351,7 @@ def get_seeds( source=source, seed_type=seed_type, origin=origin, + is_general_technique=is_general_technique, parameters=parameters, metadata=metadata, prompt_group_ids=prompt_group_ids, @@ -10351,6 +10373,7 @@ async def get_seeds_async( source: str | None = None, seed_type: SeedType | None = None, origin: SeedOrigin | None = None, + is_general_technique: bool | None = None, parameters: Sequence[str] | None = None, metadata: dict[str, str | int] | None = None, prompt_group_ids: Sequence[uuid.UUID] | None = None, @@ -10363,6 +10386,7 @@ async def get_seeds_async( Args: origin (SeedOrigin | None): Match the recorded ingestion origin. + is_general_technique (bool | None): Match whether the seed represents a general attack technique. value (str): The value to match by substring. If None, all values are returned. value_sha256 (Sequence[str] | None): A list of SHA256 hashes of values to match. If None, all values are returned. @@ -10407,6 +10431,7 @@ async def get_seeds_async( source=source, seed_type=seed_type, origin=origin, + is_general_technique=is_general_technique, parameters=parameters, metadata=metadata, prompt_group_ids=prompt_group_ids, @@ -10428,6 +10453,7 @@ def remove_seeds_from_memory( source: str | None = None, seed_type: SeedType | None = None, origin: SeedOrigin | None = None, + is_general_technique: bool | None = None, parameters: Sequence[str] | None = None, metadata: dict[str, str | int] | None = None, prompt_group_ids: Sequence[uuid.UUID] | None = None, @@ -10459,6 +10485,7 @@ def remove_seeds_from_memory( source=source, seed_type=seed_type, origin=origin, + is_general_technique=is_general_technique, parameters=parameters, metadata=metadata, prompt_group_ids=prompt_group_ids, @@ -10481,6 +10508,7 @@ async def remove_seeds_from_memory_async( source: str | None = None, seed_type: SeedType | None = None, origin: SeedOrigin | None = None, + is_general_technique: bool | None = None, parameters: Sequence[str] | None = None, metadata: dict[str, str | int] | None = None, prompt_group_ids: Sequence[uuid.UUID] | None = None, @@ -10498,6 +10526,7 @@ async def remove_seeds_from_memory_async( Args: origin (SeedOrigin | None): Match the recorded ingestion origin. + is_general_technique (bool | None): Match whether the seed represents a general attack technique. value (str): The value to match. For the remove methods this defaults to full-string equality (exact=True) so a short or common value does not delete far more seeds than intended; pass exact=False to match by substring instead. If None, all values are considered. @@ -10552,6 +10581,7 @@ async def remove_seeds_from_memory_async( source=source, seed_type=seed_type, origin=origin, + is_general_technique=is_general_technique, parameters=parameters, metadata=metadata, prompt_group_ids=prompt_group_ids, @@ -10573,6 +10603,7 @@ def remove_seed_groups_from_memory( source: str | None = None, seed_type: SeedType | None = None, origin: SeedOrigin | None = None, + is_general_technique: bool | None = None, parameters: Sequence[str] | None = None, metadata: dict[str, str | int] | None = None, prompt_group_ids: Sequence[uuid.UUID] | None = None, @@ -10604,6 +10635,7 @@ def remove_seed_groups_from_memory( source=source, seed_type=seed_type, origin=origin, + is_general_technique=is_general_technique, parameters=parameters, metadata=metadata, prompt_group_ids=prompt_group_ids, @@ -10626,6 +10658,7 @@ async def remove_seed_groups_from_memory_async( source: str | None = None, seed_type: SeedType | None = None, origin: SeedOrigin | None = None, + is_general_technique: bool | None = None, parameters: Sequence[str] | None = None, metadata: dict[str, str | int] | None = None, prompt_group_ids: Sequence[uuid.UUID] | None = None, @@ -10651,6 +10684,8 @@ async def remove_seed_groups_from_memory_async( Args: origin (SeedOrigin | None): Match origin before expanding to complete groups. + is_general_technique (bool | None): Match whether the seed represents a general attack technique + before expanding to complete groups. value (str): The value to match. For the remove methods this defaults to full-string equality (exact=True) so a short or common value does not delete far more seeds than intended; pass exact=False to match by substring instead. If None, all values are considered. @@ -10705,6 +10740,7 @@ async def remove_seed_groups_from_memory_async( source=source, seed_type=seed_type, origin=origin, + is_general_technique=is_general_technique, parameters=parameters, metadata=metadata, prompt_group_ids=prompt_group_ids, @@ -10852,6 +10888,7 @@ def get_seed_groups( source: str | None = None, seed_type: SeedType | None = None, origin: SeedOrigin | None = None, + is_general_technique: bool | None = None, parameters: Sequence[str] | None = None, metadata: dict[str, str | int] | None = None, prompt_group_ids: Sequence[uuid.UUID] | None = None, @@ -10883,6 +10920,7 @@ def get_seed_groups( source=source, seed_type=seed_type, origin=origin, + is_general_technique=is_general_technique, parameters=parameters, metadata=metadata, prompt_group_ids=prompt_group_ids, @@ -10905,6 +10943,7 @@ async def get_seed_groups_async( source: str | None = None, seed_type: SeedType | None = None, origin: SeedOrigin | None = None, + is_general_technique: bool | None = None, parameters: Sequence[str] | None = None, metadata: dict[str, str | int] | None = None, prompt_group_ids: Sequence[uuid.UUID] | None = None, @@ -10915,6 +10954,8 @@ async def get_seed_groups_async( Args: origin (SeedOrigin | None): Match origin before expanding to complete groups. + is_general_technique (bool | None, Optional): Match whether the seed represents a general attack technique + before expanding to complete groups. value (str | None, Optional): The value to match by substring. value_sha256 (Sequence[str] | None, Optional): SHA256 hash of value to filter seed groups by. dataset_name (str | None, Optional): Name of the dataset to match exactly. @@ -10954,6 +10995,7 @@ async def get_seed_groups_async( source=source, seed_type=seed_type, origin=origin, + is_general_technique=is_general_technique, parameters=parameters, metadata=metadata, prompt_group_ids=prompt_group_ids, diff --git a/pyrit/memory/memory_models.py b/pyrit/memory/memory_models.py index a693c7d71b..36c83bfad8 100644 --- a/pyrit/memory/memory_models.py +++ b/pyrit/memory/memory_models.py @@ -1697,6 +1697,7 @@ class SeedEntry(Base): String(16), nullable=False, default=SeedOrigin.UNKNOWN.value, server_default="unknown", index=True ) conditions: Mapped[list[dict[str, Any]] | None] = mapped_column(JSON, nullable=True) + is_general_technique: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False, server_default="0") def __init__(self, *, entry: Seed[Any]) -> None: """ @@ -1730,6 +1731,7 @@ def __init__(self, *, entry: Seed[Any]) -> None: self.prompt_group_id = entry.prompt_group_id self.seed_type = seed_type self.origin = entry.origin.value + self.is_general_technique = bool(entry.is_general_technique) self.conditions = ( entry.model_dump(mode="json", include={"conditions"})["conditions"] or None if isinstance(entry, SeedObjective) @@ -1856,6 +1858,7 @@ def get_seed_record(self) -> SeedRecord: parameters=self.parameters, conditions=self.conditions, response_json_schema=response_json_schema, + is_general_technique=self.is_general_technique, ) def get_seed(self) -> SeedPrompt | SeedObjective | SeedSimulatedConversation: @@ -1930,6 +1933,7 @@ def get_seed(self) -> SeedPrompt | SeedObjective | SeedSimulatedConversation: num_turns=config.get("num_turns", 3), sequence=config.get("sequence", 0), pyrit_version=config.get("pyrit_version"), + is_general_technique=self.is_general_technique, **prompt_config, ) except (OSError, ValueError) as exc: @@ -1961,6 +1965,7 @@ def get_seed(self) -> SeedPrompt | SeedObjective | SeedSimulatedConversation: prompt_group_id=self.prompt_group_id, sequence=self.sequence or 0, role=self.role, + is_general_technique=self.is_general_technique, ) diff --git a/pyrit/models/seeds/seed_record.py b/pyrit/models/seeds/seed_record.py index d081f20c69..585a3e57cd 100644 --- a/pyrit/models/seeds/seed_record.py +++ b/pyrit/models/seeds/seed_record.py @@ -40,3 +40,4 @@ class SeedRecord(BaseModel): parameters: list[str] | None conditions: ConditionTuple | None response_json_schema: JsonSchemaDefinition | None + is_general_technique: bool = False diff --git a/tests/unit/memory/test_memory_models.py b/tests/unit/memory/test_memory_models.py index a94b5b176a..ccd9c141f1 100644 --- a/tests/unit/memory/test_memory_models.py +++ b/tests/unit/memory/test_memory_models.py @@ -598,8 +598,21 @@ def test_get_seed_record_preserves_stored_fields_without_generated_defaults(self assert record.response_json_schema == {"type": "object"} assert record.date_added == seed.date_added assert record.added_by == seed.added_by + assert record.is_general_technique is False assert "is_jinja_template" not in record.model_dump() + def test_seed_prompt_roundtrip_preserves_is_general_technique(self) -> None: + seed = _make_seed_prompt(is_general_technique=True) + entry = SeedEntry(entry=seed) + assert entry.is_general_technique is True + + record = entry.get_seed_record() + assert record.is_general_technique is True + + recovered = entry.get_seed() + assert isinstance(recovered, SeedPrompt) + assert recovered.is_general_technique is True + # ---- response_json_schema persistence --------------------------------- def test_roundtrip_seed_prompt_preserves_inline_response_json_schema(self): diff --git a/tests/unit/memory/test_operations_and_findings_migration.py b/tests/unit/memory/test_operations_and_findings_migration.py index c9b0299cda..d92c5b0616 100644 --- a/tests/unit/memory/test_operations_and_findings_migration.py +++ b/tests/unit/memory/test_operations_and_findings_migration.py @@ -24,7 +24,7 @@ OperationEntry, ScenarioResultEntry, ) -from pyrit.memory.migration import check_schema_migrations +from pyrit.memory.migration import check_schema_migrations, run_schema_migrations PREVIOUS_REVISION = "3f6a9d2c8b1e" REVISION = "c8d3e5f7a901" @@ -65,6 +65,7 @@ def test_migration_adds_tables_without_changing_existing_ones() -> None: assert {name: columns for name, columns in after.items() if name not in ACTIVITY_TABLES} == { name: columns for name, columns in existing.items() if name not in ACTIVITY_TABLES } + run_schema_migrations(engine=engine) check_schema_migrations(engine=engine) finally: engine.dispose() diff --git a/tests/unit/memory/test_seed_general_technique.py b/tests/unit/memory/test_seed_general_technique.py new file mode 100644 index 0000000000..6e8c5b3252 --- /dev/null +++ b/tests/unit/memory/test_seed_general_technique.py @@ -0,0 +1,194 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +from uuid import uuid4 + +import pytest + +from pyrit.memory import MemoryInterface +from pyrit.models import ( + AttackTechniqueSeedGroup, + SeedIdentifier, + SeedPrompt, + SeedRecord, + SeedSimulatedConversation, +) + + +@pytest.mark.usefixtures("patch_central_database") +class TestSeedGeneralTechnique: + async def test_attack_technique_seed_group_roundtrip_async(self, sqlite_instance: MemoryInterface) -> None: + group = AttackTechniqueSeedGroup.from_system_prompt("System attack prompt.") + seed = group.seeds[0] + assert seed.is_general_technique is True + + # Prepare and verify identity + await sqlite_instance.add_seeds_to_memory_async(seeds=group.seeds, added_by="operator") + identifier_before = SeedIdentifier.from_seed(seed) + + [retrieved_seed] = await sqlite_instance.get_seeds_async(prompt_group_ids=[seed.prompt_group_id]) + assert isinstance(retrieved_seed, SeedPrompt) + assert retrieved_seed.is_general_technique is True + identifier_after = SeedIdentifier.from_seed(retrieved_seed) + + assert identifier_before.hash == identifier_after.hash + + # Reconstructing group from retrieved seeds must succeed without validation error + rebuilt_group = AttackTechniqueSeedGroup(seeds=[retrieved_seed]) + assert len(rebuilt_group.seeds) == 1 + assert rebuilt_group.seeds[0].is_general_technique is True + + async def test_seed_simulated_conversation_roundtrip_async(self, sqlite_instance: MemoryInterface) -> None: + seed = SeedSimulatedConversation( + adversarial_chat_system_prompt=SeedPrompt(value="Generate a conversation."), + dataset_name="techniques", + ) + assert seed.is_general_technique is True + + await sqlite_instance.add_seeds_to_memory_async(seeds=[seed], added_by="operator") + identifier_before = SeedIdentifier.from_seed(seed) + + [retrieved] = await sqlite_instance.get_seeds_async(dataset_name="techniques") + assert isinstance(retrieved, SeedSimulatedConversation) + assert retrieved.is_general_technique is True + + identifier_after = SeedIdentifier.from_seed(retrieved) + assert identifier_before.hash == identifier_after.hash + + async def test_filter_by_is_general_technique_async(self, sqlite_instance: MemoryInterface) -> None: + general_seed = SeedPrompt(value="General technique seed", is_general_technique=True, dataset_name="mixed") + specific_seed = SeedPrompt(value="Specific objective prompt", is_general_technique=False, dataset_name="mixed") + + await sqlite_instance.add_seeds_to_memory_async(seeds=[general_seed, specific_seed], added_by="operator") + + all_seeds = await sqlite_instance.get_seeds_async(dataset_name="mixed") + assert len(all_seeds) == 2 + + general_only = await sqlite_instance.get_seeds_async(dataset_name="mixed", is_general_technique=True) + assert len(general_only) == 1 + assert general_only[0].value == "General technique seed" + assert general_only[0].is_general_technique is True + + specific_only = await sqlite_instance.get_seeds_async(dataset_name="mixed", is_general_technique=False) + assert len(specific_only) == 1 + assert specific_only[0].value == "Specific objective prompt" + assert specific_only[0].is_general_technique is False + + async def test_remove_by_is_general_technique_async(self, sqlite_instance: MemoryInterface) -> None: + general_seed = SeedPrompt(value="General prompt", is_general_technique=True, dataset_name="removal") + specific_seed = SeedPrompt(value="Specific prompt", is_general_technique=False, dataset_name="removal") + + await sqlite_instance.add_seeds_to_memory_async(seeds=[general_seed, specific_seed], added_by="operator") + + removed_count = await sqlite_instance.remove_seeds_from_memory_async( + dataset_name="removal", is_general_technique=True + ) + assert removed_count == 1 + + remaining = await sqlite_instance.get_seeds_async(dataset_name="removal") + assert len(remaining) == 1 + assert remaining[0].value == "Specific prompt" + assert remaining[0].is_general_technique is False + + async def test_seed_groups_filter_by_is_general_technique_async(self, sqlite_instance: MemoryInterface) -> None: + group_id_1 = uuid4() + group_id_2 = uuid4() + + group1_seeds = [ + SeedPrompt( + value="General turn 1", + is_general_technique=True, + prompt_group_id=group_id_1, + sequence=0, + role="user", + ), + SeedPrompt( + value="General turn 2", + is_general_technique=True, + prompt_group_id=group_id_1, + sequence=1, + role="user", + ), + ] + group2_seeds = [ + SeedPrompt( + value="Regular turn 1", + is_general_technique=False, + prompt_group_id=group_id_2, + sequence=0, + role="user", + ), + SeedPrompt( + value="Regular turn 2", + is_general_technique=False, + prompt_group_id=group_id_2, + sequence=1, + role="user", + ), + ] + + await sqlite_instance.add_seeds_to_memory_async(seeds=group1_seeds + group2_seeds, added_by="operator") + + general_groups = await sqlite_instance.get_seed_groups_async(is_general_technique=True) + assert len(general_groups) == 1 + assert general_groups[0].seeds[0].prompt_group_id == group_id_1 + + non_general_groups = await sqlite_instance.get_seed_groups_async(is_general_technique=False) + assert len(non_general_groups) == 1 + assert non_general_groups[0].seeds[0].prompt_group_id == group_id_2 + + async def test_remove_seed_groups_by_is_general_technique_async(self, sqlite_instance: MemoryInterface) -> None: + group_id_1 = uuid4() + group_id_2 = uuid4() + + group1_seeds = [ + SeedPrompt( + value="Delete me general", + is_general_technique=True, + prompt_group_id=group_id_1, + sequence=0, + role="user", + ), + SeedPrompt( + value="Delete me companion", + is_general_technique=True, + prompt_group_id=group_id_1, + sequence=1, + role="user", + ), + ] + group2_seeds = [ + SeedPrompt( + value="Keep me regular", + is_general_technique=False, + prompt_group_id=group_id_2, + sequence=0, + role="user", + ), + SeedPrompt( + value="Keep me companion", + is_general_technique=False, + prompt_group_id=group_id_2, + sequence=1, + role="user", + ), + ] + + await sqlite_instance.add_seeds_to_memory_async(seeds=group1_seeds + group2_seeds, added_by="operator") + + removed = await sqlite_instance.remove_seed_groups_from_memory_async(is_general_technique=True) + assert removed == 2 + + remaining_groups = await sqlite_instance.get_seed_groups_async() + assert len(remaining_groups) == 1 + assert remaining_groups[0].seeds[0].prompt_group_id == group_id_2 + + async def test_seed_record_projection(self, sqlite_instance: MemoryInterface) -> None: + seed = SeedPrompt(value="Record projection test", is_general_technique=True, dataset_name="projection") + await sqlite_instance.add_seeds_to_memory_async(seeds=[seed], added_by="operator") + + examples = await sqlite_instance.get_seed_example_async(dataset_name="projection", example_id=seed.id) + assert len(examples) == 1 + record = examples[0] + assert isinstance(record, SeedRecord) + assert record.is_general_technique is True diff --git a/tests/unit/memory/test_seed_is_general_technique_migration.py b/tests/unit/memory/test_seed_is_general_technique_migration.py new file mode 100644 index 0000000000..cb8a8e88d6 --- /dev/null +++ b/tests/unit/memory/test_seed_is_general_technique_migration.py @@ -0,0 +1,63 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import importlib +from io import StringIO +from unittest.mock import patch + +import sqlalchemy as sa +from alembic.migration import MigrationContext +from alembic.operations import Operations + +from pyrit.memory.memory_models import SeedEntry + + +def test_seed_is_general_technique_migration_preserves_legacy_rows() -> None: + migration = importlib.import_module("pyrit.memory.alembic.versions.e7a1b3c5d9f0_add_seed_is_general_technique") + engine = sa.create_engine("sqlite:///:memory:") + try: + with engine.begin() as connection: + connection.execute(sa.text('CREATE TABLE "SeedPromptEntries" (id INTEGER PRIMARY KEY, value TEXT)')) + connection.execute(sa.text("INSERT INTO \"SeedPromptEntries\" VALUES (1, 'legacy')")) + operations = Operations(MigrationContext.configure(connection)) + with patch.object(migration, "op", operations): + migration.upgrade() + row = connection.execute(sa.text('SELECT id, value, is_general_technique FROM "SeedPromptEntries"')).one() + assert tuple(row) == (1, "legacy", 0) + + connection.execute( + sa.text("INSERT INTO \"SeedPromptEntries\" (id, value, is_general_technique) VALUES (2, 'general', 1)") + ) + row2 = connection.execute( + sa.text('SELECT id, value, is_general_technique FROM "SeedPromptEntries" WHERE id = 2') + ).one() + assert tuple(row2) == (2, "general", 1) + + columns = {column["name"]: column for column in sa.inspect(connection).get_columns("SeedPromptEntries")} + assert "is_general_technique" in columns + assert not columns["is_general_technique"]["nullable"] + + with patch.object(migration, "op", operations): + migration.downgrade() + assert connection.execute(sa.text('SELECT id, value FROM "SeedPromptEntries"')).fetchall() == [ + (1, "legacy"), + (2, "general"), + ] + assert "is_general_technique" not in { + column["name"] for column in sa.inspect(connection).get_columns("SeedPromptEntries") + } + finally: + engine.dispose() + + +def test_seed_is_general_technique_migration_uses_sql_server_bit() -> None: + migration = importlib.import_module("pyrit.memory.alembic.versions.e7a1b3c5d9f0_add_seed_is_general_technique") + output = StringIO() + context = MigrationContext.configure(dialect_name="mssql", opts={"as_sql": True, "output_buffer": output}) + with patch.object(migration, "op", Operations(context)): + migration.upgrade() + migration.downgrade() + sql = output.getvalue() + assert "ADD is_general_technique BIT NOT NULL DEFAULT 0" in sql + assert "DROP COLUMN is_general_technique" in sql + assert SeedEntry.__table__.c.is_general_technique.type.python_type is bool