diff --git a/pyrit/memory/alembic/versions/064a1eff19fc_add_is_general_technique_to_seed_entries.py b/pyrit/memory/alembic/versions/064a1eff19fc_add_is_general_technique_to_seed_entries.py new file mode 100644 index 0000000000..9713e38620 --- /dev/null +++ b/pyrit/memory/alembic/versions/064a1eff19fc_add_is_general_technique_to_seed_entries.py @@ -0,0 +1,36 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Persist general technique flag on seed entries.""" + +from collections.abc import Sequence # noqa: TC003 + +import sqlalchemy as sa +from alembic import op + +revision: str = "064a1eff19fc" +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 column and default existing rows.""" + 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())) + seed_prompt_entries = sa.table( + "SeedPromptEntries", + sa.column("is_general_technique", sa.Boolean()), + sa.column("seed_type", sa.String()), + ) + op.execute( + seed_prompt_entries.update() + .where(seed_prompt_entries.c.seed_type == "simulated_conversation") + .values(is_general_technique=True) + ) + + +def downgrade() -> None: + """Drop is_general_technique column from seed entries.""" + 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 72c2c69177..bb2407e90e 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -5117,6 +5117,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=(",", ":"), @@ -5200,7 +5201,10 @@ def _seed_behavior_key(seed: Seed[Any]) -> str: Returns: str: Canonical JSON, preserving condition and template-parameter order. """ - return json.dumps(SeedIdentifier.get_behavioral_parameters(seed), sort_keys=True, separators=(",", ":")) + params = dict(SeedIdentifier.get_behavioral_parameters(seed)) + if getattr(seed, "is_general_technique", False): + params["is_general_technique"] = True + return json.dumps(params, sort_keys=True, separators=(",", ":")) def _get_new_behavior_group_ids( self, diff --git a/pyrit/memory/memory_models.py b/pyrit/memory/memory_models.py index a693c7d71b..96e88b5f23 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(getattr(entry, "is_general_technique", False)) 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=bool(getattr(self, "is_general_technique", False)), ) 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=bool(getattr(self, "is_general_technique", True)), **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=bool(getattr(self, "is_general_technique", False)), ) 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_operations_and_findings_migration.py b/tests/unit/memory/test_operations_and_findings_migration.py index c9b0299cda..651487c469 100644 --- a/tests/unit/memory/test_operations_and_findings_migration.py +++ b/tests/unit/memory/test_operations_and_findings_migration.py @@ -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 } + command.upgrade(config, "head") 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..23fd9072ae --- /dev/null +++ b/tests/unit/memory/test_seed_general_technique.py @@ -0,0 +1,144 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import pytest + +from pyrit.memory import MemoryInterface +from pyrit.memory.memory_models import SeedEntry +from pyrit.models import ( + AttackTechniqueSeedGroup, + SeedIdentifier, + SeedObjective, + SeedPrompt, + SeedRecord, + SeedSimulatedConversation, +) + + +@pytest.mark.usefixtures("patch_central_database") +class TestSeedGeneralTechnique: + """Tests for symmetric persistence of the is_general_technique flag.""" + + async def test_seed_entry_persists_is_general_technique(self) -> None: + general_prompt = SeedPrompt( + value="General technique", data_type="text", is_general_technique=True, added_by="operator" + ) + entry_general = SeedEntry(entry=general_prompt) + assert entry_general.is_general_technique is True + restored_general = entry_general.get_seed() + assert isinstance(restored_general, SeedPrompt) + assert restored_general.is_general_technique is True + + regular_prompt = SeedPrompt( + value="Regular prompt", data_type="text", is_general_technique=False, added_by="operator" + ) + entry_regular = SeedEntry(entry=regular_prompt) + assert entry_regular.is_general_technique is False + restored_regular = entry_regular.get_seed() + assert isinstance(restored_regular, SeedPrompt) + assert restored_regular.is_general_technique is False + + objective = SeedObjective(value="Target objective", added_by="operator") + entry_obj = SeedEntry(entry=objective) + assert entry_obj.is_general_technique is False + restored_obj = entry_obj.get_seed() + assert isinstance(restored_obj, SeedObjective) + assert restored_obj.is_general_technique is False + + conversation = SeedSimulatedConversation( + adversarial_chat_system_prompt=SeedPrompt(value="Adversary system", added_by="operator"), + added_by="operator", + ) + entry_conv = SeedEntry(entry=conversation) + assert entry_conv.is_general_technique is True + restored_conv = entry_conv.get_seed() + assert isinstance(restored_conv, SeedSimulatedConversation) + assert restored_conv.is_general_technique is True + + def test_seed_record_projects_is_general_technique(self) -> None: + general_prompt = SeedPrompt( + value="Technique prompt", data_type="text", is_general_technique=True, added_by="operator" + ) + entry = SeedEntry(entry=general_prompt) + record = entry.get_seed_record() + assert isinstance(record, SeedRecord) + assert record.is_general_technique is True + assert "is_jinja_template" not in record.model_dump() + + regular_prompt = SeedPrompt( + value="Regular prompt", data_type="text", is_general_technique=False, added_by="operator" + ) + record_reg = SeedEntry(entry=regular_prompt).get_seed_record() + assert record_reg.is_general_technique is False + + async def test_attack_technique_seed_group_sqlite_roundtrip_rebuilds_and_matches_identifier( + self, sqlite_instance: MemoryInterface + ) -> None: + technique_seed = SeedPrompt(value="Framing strategy", data_type="text", is_general_technique=True) + group = AttackTechniqueSeedGroup(seeds=[technique_seed]) + + await sqlite_instance.add_seeds_to_memory_async(seeds=group.seeds, added_by="operator") + persisted_seed = group.seeds[0] + original_identifier = SeedIdentifier.from_seed(persisted_seed) + + retrieved_seeds = await sqlite_instance.get_seeds_async() + assert len(retrieved_seeds) == 1 + retrieved_seed = retrieved_seeds[0] + assert isinstance(retrieved_seed, SeedPrompt) + assert retrieved_seed.is_general_technique is True + + retrieved_identifier = SeedIdentifier.from_seed(retrieved_seed) + assert retrieved_identifier.hash == original_identifier.hash + assert retrieved_identifier.is_general_technique is True + + rebuilt_group = AttackTechniqueSeedGroup(seeds=retrieved_seeds) + assert len(rebuilt_group.seeds) == 1 + assert rebuilt_group.seeds[0].is_general_technique is True + + def test_legacy_entry_defaults_to_false_for_prompt(self) -> None: + entry = SeedEntry(entry=SeedPrompt(value="Legacy prompt", added_by="legacy")) + entry.is_general_technique = False + + seed = entry.get_seed() + assert isinstance(seed, SeedPrompt) + assert seed.is_general_technique is False + + record = entry.get_seed_record() + assert record.is_general_technique is False + + async def test_sequential_inserts_with_opposite_flags_both_stored_async( + self, sqlite_instance: MemoryInterface + ) -> None: + regular = SeedPrompt(value="Same value", dataset_name="flags", is_general_technique=False) + technique = SeedPrompt(value="Same value", dataset_name="flags", is_general_technique=True) + + await sqlite_instance.add_seeds_to_memory_async(seeds=[regular], added_by="tester") + await sqlite_instance.add_seeds_to_memory_async(seeds=[technique], added_by="tester") + + stored = await sqlite_instance.get_seeds_async(dataset_name="flags") + assert len(stored) == 2 + flags = {s.is_general_technique for s in stored} + assert flags == {False, True} + + async def test_sequential_inserts_technique_first_opposite_flags_both_stored_async( + self, sqlite_instance: MemoryInterface + ) -> None: + technique = SeedPrompt(value="Same value 2", dataset_name="flags2", is_general_technique=True) + regular = SeedPrompt(value="Same value 2", dataset_name="flags2", is_general_technique=False) + + await sqlite_instance.add_seeds_to_memory_async(seeds=[technique], added_by="tester") + await sqlite_instance.add_seeds_to_memory_async(seeds=[regular], added_by="tester") + + stored = await sqlite_instance.get_seeds_async(dataset_name="flags2") + assert len(stored) == 2 + flags = {s.is_general_technique for s in stored} + assert flags == {False, True} + + def test_seed_behavior_and_group_member_keys_include_general_technique( + self, sqlite_instance: MemoryInterface + ) -> None: + regular = SeedPrompt(value="Key test", is_general_technique=False) + technique = SeedPrompt(value="Key test", is_general_technique=True) + + assert sqlite_instance._seed_behavior_key(regular) != sqlite_instance._seed_behavior_key(technique) + assert sqlite_instance._seed_group_member_key(regular) != sqlite_instance._seed_group_member_key(technique) diff --git a/tests/unit/memory/test_seed_general_technique_migration.py b/tests/unit/memory/test_seed_general_technique_migration.py new file mode 100644 index 0000000000..d94ce9af86 --- /dev/null +++ b/tests/unit/memory/test_seed_general_technique_migration.py @@ -0,0 +1,88 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import importlib +import io +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 +from pyrit.memory.migration import check_schema_migrations, run_schema_migrations + + +def test_general_technique_migration_preserves_existing_rows() -> None: + migration = importlib.import_module( + "pyrit.memory.alembic.versions.064a1eff19fc_add_is_general_technique_to_seed_entries" + ) + 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, seed_type VARCHAR NOT NULL)' + ) + ) + connection.execute( + sa.text( + 'INSERT INTO "SeedPromptEntries" VALUES ' + "(1, 'legacy prompt', 'prompt'), " + "(2, 'legacy objective', 'objective'), " + "(3, 'legacy conversation', 'simulated_conversation')" + ) + ) + operations = Operations(MigrationContext.configure(connection)) + with patch.object(migration, "op", operations): + migration.upgrade() + + 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"] + + rows = connection.execute( + sa.text('SELECT id, is_general_technique FROM "SeedPromptEntries" ORDER BY id') + ).all() + assert rows[0] == (1, 0) + assert rows[1] == (2, 0) + assert rows[2] == (3, 1) + + with patch.object(migration, "op", operations): + migration.downgrade() + + assert "is_general_technique" not in { + column["name"] for column in sa.inspect(connection).get_columns("SeedPromptEntries") + } + preserved = connection.execute( + sa.text('SELECT id, value, seed_type FROM "SeedPromptEntries" ORDER BY id') + ).all() + assert len(preserved) == 3 + assert preserved[0] == (1, "legacy prompt", "prompt") + assert preserved[2] == (3, "legacy conversation", "simulated_conversation") + finally: + engine.dispose() + + +def test_general_technique_migration_sql_server() -> None: + migration = importlib.import_module( + "pyrit.memory.alembic.versions.064a1eff19fc_add_is_general_technique_to_seed_entries" + ) + output = io.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 "is_general_technique BIT NOT NULL" in sql + assert "DROP COLUMN is_general_technique" in sql + assert SeedEntry.__table__.c.is_general_technique.type.python_type is bool + + +def test_schema_matches_models_after_migration() -> None: + engine = sa.create_engine("sqlite:///:memory:") + try: + run_schema_migrations(engine=engine) + check_schema_migrations(engine=engine) + finally: + engine.dispose()