From 250e8115708eb1b319ff5c38ddc795be21f748e0 Mon Sep 17 00:00:00 2001 From: Hrushikesh Ramilla Date: Tue, 3 Nov 2026 03:21:50 +0530 Subject: [PATCH] FIX: Persist general technique flag on seed entries (#3739) Persist is_general_technique on SeedEntry and project it onto SeedRecord, reconstructing the execution flag on retrieved seeds and preventing identity drift or reconstruction errors when restoring technique groups. Also adds Alembic schema migration and unit regression tests. Signed-off-by: Hrushikesh Ramilla --- ...dd_is_general_technique_to_seed_entries.py | 36 +++++ pyrit/memory/memory_interface.py | 6 +- pyrit/memory/memory_models.py | 5 + pyrit/models/seeds/seed_record.py | 1 + .../test_operations_and_findings_migration.py | 1 + .../memory/test_seed_general_technique.py | 144 ++++++++++++++++++ .../test_seed_general_technique_migration.py | 88 +++++++++++ 7 files changed, 280 insertions(+), 1 deletion(-) create mode 100644 pyrit/memory/alembic/versions/064a1eff19fc_add_is_general_technique_to_seed_entries.py create mode 100644 tests/unit/memory/test_seed_general_technique.py create mode 100644 tests/unit/memory/test_seed_general_technique_migration.py 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()