FIX: Persist general technique flag on seed entries (#3739)

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)

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,

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)),
         )
 
 

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

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()

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)

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()