MAINT: Batch the seed dedupe lookup when adding seeds to memory

pyrit/memory/memory_interface.py

@@ -2976,25 +2976,90 @@ async def add_seeds_to_memory_async(self, *, seeds: Sequence[Seed], added_by: st
         """
         Insert a list of seeds into the memory storage.
 
+        Seeds already present in storage are skipped. Duplicates *within* ``seeds`` are all
+        inserted, because the check looks at what storage held when the call started.
+
         Args:
             seeds (Sequence[Seed]): A list of seeds to insert.
             added_by (str): The user who added the seeds.
 
         Raises:
             ValueError: If the 'added_by' attribute is not set for each prompt.
         """
-        entries: MutableSequence[SeedEntry] = []
         current_time = datetime.now(tz=timezone.utc)
         for prompt in seeds:
             await self._prepare_seed_for_storage_async(prompt=prompt, added_by=added_by, current_time=current_time)
 
-            if prompt.value_sha256 and not self.get_seeds(
-                value_sha256=[prompt.value_sha256], dataset_name=prompt.dataset_name
-            ):
-                entries.append(SeedEntry(entry=prompt))
+        existing_pairs, existing_hashes = self._get_existing_seed_keys(seeds=seeds)
+
+        entries: MutableSequence[SeedEntry] = []
+        for prompt in seeds:
+            if not prompt.value_sha256:
+                continue
+            # A seed without a dataset name matches the hash in any dataset, mirroring the
+            # filter that is applied when dataset_name is not supplied.
+            if prompt.dataset_name:
+                if (prompt.value_sha256, prompt.dataset_name) in existing_pairs:
+                    continue
+            elif prompt.value_sha256 in existing_hashes:
+                continue
+            entries.append(SeedEntry(entry=prompt))
 
         self._insert_entries(entries=entries)
 
+    def _get_existing_seed_keys(self, *, seeds: Sequence[Seed]) -> tuple[set[tuple[str, str]], set[str]]:
+        """
+        Look up which of these seeds' hashes are already stored.
+
+        Queries in chunks rather than once per seed, which otherwise dominates the cost of
+        loading a large dataset. ``get_seeds`` issues the statement directly instead of going
+        through the batching helpers, so the bound has to be applied here.
+
+        The seeds are grouped by dataset name so the name is still compared by the database
+        rather than in Python. Equality here is a property of the column's collation: Azure
+        SQL's default is case-insensitive, T-SQL also ignores trailing blanks, and an
+        accent-insensitive collation folds further still. Comparing the names in Python would
+        silently impose one fixed rule on every backend and insert a duplicate wherever the
+        stored spelling differs from the incoming one.
+
+        Args:
+            seeds (Sequence[Seed]): The seeds whose hashes should be looked up.
+
+        Returns:
+            tuple[set[tuple[str, str]], set[str]]: The stored (value_sha256, dataset_name)
+                pairs keyed by the requested dataset name, and the stored hashes irrespective
+                of dataset.
+        """
+        hashes_by_dataset: dict[str | None, set[str]] = {}
+        for prompt in seeds:
+            if prompt.value_sha256:
+                # An empty name filters nothing, exactly like None, and the caller below
+                # treats both as "match the hash in any dataset". Normalize so the two
+                # cannot disagree.
+                hashes_by_dataset.setdefault(prompt.dataset_name or None, set()).add(prompt.value_sha256)
+
+        existing_pairs: set[tuple[str, str]] = set()
+        existing_hashes: set[str] = set()
+        for dataset_name, dataset_hashes in hashes_by_dataset.items():
+            hashes = sorted(dataset_hashes)
+            # _MAX_BIND_VARS is the whole statement's budget, and the name takes one of those
+            # binds, so the hashes get what is left rather than the full ceiling.
+            chunk_size = self._MAX_BIND_VARS - (1 if dataset_name is not None else 0)
+            for index in range(0, len(hashes), chunk_size):
+                chunk = hashes[index : index + chunk_size]
+                for existing in self.get_seeds(value_sha256=chunk, dataset_name=dataset_name):
+                    if not existing.value_sha256:
+                        continue
+                    if dataset_name:
+                        # The database decided the name matched, so record the name that was
+                        # asked for; the stored spelling can differ under a case-insensitive
+                        # collation.
+                        existing_pairs.add((existing.value_sha256, dataset_name))
+                    else:
+                        existing_hashes.add(existing.value_sha256)
+
+        return existing_pairs, existing_hashes
+
     async def add_seed_datasets_to_memory_async(self, *, datasets: Sequence[SeedDataset], added_by: str) -> None:
         """
         Insert a list of seed datasets into the memory storage.

tests/unit/memory/memory_interface/test_interface_seed_prompts.py

@@ -8,9 +8,11 @@
 from uuid import uuid4
 
 import pytest
+from sqlalchemy import String
 from sqlalchemy.exc import SQLAlchemyError
 
 from pyrit.memory import MemoryInterface
+from pyrit.memory.memory_models import SeedEntry
 from pyrit.models import MessagePiece, SeedDataset, SeedGroup, SeedObjective, SeedPrompt
 
 
@@ -429,6 +431,185 @@ async def test_add_seed_prompts_duplicate_entries_same_dataset(sqlite_instance:
     assert len(stored_prompts) == 3
 
 
+async def test_add_seed_prompts_duplicates_within_one_call_are_all_stored(sqlite_instance: MemoryInterface):
+    """Existing behaviour: the dedupe check looks at storage, not at the batch being added."""
+    prompts: Sequence[SeedPrompt] = [
+        SeedPrompt(value="prompt1", dataset_name="test_dataset", data_type="text"),
+        SeedPrompt(value="prompt1", dataset_name="test_dataset", data_type="text"),
+        SeedPrompt(value="prompt1", dataset_name="test_dataset", data_type="text"),
+    ]
+    await sqlite_instance.add_seeds_to_memory_async(seeds=prompts, added_by="tester")
+
+    assert len(sqlite_instance.get_seeds(dataset_name="test_dataset")) == 3
+
+
+async def test_add_seed_prompts_without_dataset_name_matches_any_dataset(sqlite_instance: MemoryInterface):
+    """Without a dataset name the lookup is unfiltered, so the hash matches in any dataset."""
+    await sqlite_instance.add_seeds_to_memory_async(
+        seeds=[SeedPrompt(value="prompt1", dataset_name="test_dataset", data_type="text")],
+        added_by="tester",
+    )
+
+    await sqlite_instance.add_seeds_to_memory_async(
+        seeds=[SeedPrompt(value="prompt1", data_type="text")],
+        added_by="tester",
+    )
+
+    assert len(sqlite_instance.get_seeds()) == 1
+
+
+async def test_add_seed_prompts_dedupes_across_chunk_boundaries(sqlite_instance: MemoryInterface):
+    """The lookup is chunked, so duplicates have to be caught across chunk boundaries."""
+    count = sqlite_instance._MAX_BIND_VARS * 2 + 25
+    first: Sequence[SeedPrompt] = [
+        SeedPrompt(value=f"prompt{index}", dataset_name="test_dataset", data_type="text") for index in range(count)
+    ]
+    await sqlite_instance.add_seeds_to_memory_async(seeds=first, added_by="tester")
+    assert len(sqlite_instance.get_seeds(dataset_name="test_dataset")) == count
+
+    # Re-adding the same seeds plus one new one must only store the new one.
+    second: Sequence[SeedPrompt] = [
+        SeedPrompt(value=f"prompt{index}", dataset_name="test_dataset", data_type="text") for index in range(count)
+    ] + [SeedPrompt(value="brand_new", dataset_name="test_dataset", data_type="text")]
+    await sqlite_instance.add_seeds_to_memory_async(seeds=second, added_by="tester")
+
+    assert len(sqlite_instance.get_seeds(dataset_name="test_dataset")) == count + 1
+
+
+async def test_add_seed_prompts_queries_are_batched(sqlite_instance: MemoryInterface):
+    """Regression guard: the dedupe lookup must not run one query per seed."""
+    prompts: Sequence[SeedPrompt] = [
+        SeedPrompt(value=f"prompt{index}", dataset_name="test_dataset", data_type="text") for index in range(50)
+    ]
+
+    with patch.object(sqlite_instance, "get_seeds", wraps=sqlite_instance.get_seeds) as spied:
+        await sqlite_instance.add_seeds_to_memory_async(seeds=prompts, added_by="tester")
+
+    assert spied.call_count == 1
+
+
+async def test_add_seed_prompts_dedupe_honors_backend_bind_var_limit(sqlite_instance: MemoryInterface):
+    """The whole statement has to fit the backend's ceiling, and the dataset name takes a bind."""
+    prompts: Sequence[SeedPrompt] = [
+        SeedPrompt(value=f"prompt{index}", dataset_name="test_dataset", data_type="text") for index in range(120)
+    ]
+
+    with patch.object(type(sqlite_instance), "_MAX_BIND_VARS", 50):
+        with patch.object(sqlite_instance, "get_seeds", wraps=sqlite_instance.get_seeds) as spied:
+            await sqlite_instance.add_seeds_to_memory_async(seeds=prompts, added_by="tester")
+
+    assert spied.call_count == 3
+    # 49 hashes plus the name is 50 binds; a full 50 hashes would exceed the ceiling by one.
+    assert [len(call.kwargs["value_sha256"]) for call in spied.call_args_list] == [49, 49, 22]
+    assert all(len(call.kwargs["value_sha256"]) + 1 <= 50 for call in spied.call_args_list)
+    assert len(sqlite_instance.get_seeds(dataset_name="test_dataset")) == 120
+
+
+async def test_add_seed_prompts_dedupe_uses_the_full_budget_without_a_dataset_name(
+    sqlite_instance: MemoryInterface,
+):
+    """With no name to bind there is nothing to reserve, so the whole ceiling goes to hashes."""
+    prompts: Sequence[SeedPrompt] = [SeedPrompt(value=f"prompt{index}", data_type="text") for index in range(120)]
+
+    with patch.object(type(sqlite_instance), "_MAX_BIND_VARS", 50):
+        with patch.object(sqlite_instance, "get_seeds", wraps=sqlite_instance.get_seeds) as spied:
+            await sqlite_instance.add_seeds_to_memory_async(seeds=prompts, added_by="tester")
+
+    assert [len(call.kwargs["value_sha256"]) for call in spied.call_args_list] == [50, 50, 20]
+
+
+async def test_add_seed_prompts_dedupe_delegates_dataset_name_to_the_database(sqlite_instance: MemoryInterface):
+    """Case sensitivity belongs to the column's collation, so the name must be compared in SQL."""
+    prompts: Sequence[SeedPrompt] = [
+        SeedPrompt(value=f"prompt{index}", dataset_name="test_dataset", data_type="text") for index in range(3)
+    ]
+
+    with patch.object(sqlite_instance, "get_seeds", wraps=sqlite_instance.get_seeds) as spied:
+        await sqlite_instance.add_seeds_to_memory_async(seeds=prompts, added_by="tester")
+
+    assert all(call.kwargs["dataset_name"] == "test_dataset" for call in spied.call_args_list)
+
+
+async def test_add_seed_prompts_dedupe_follows_a_case_insensitive_collation(sqlite_instance: MemoryInterface):
+    """Azure SQL's default collation is case-insensitive, so a differently cased name is a duplicate.
+
+    Rebuilds the real column with COLLATE NOCASE so the comparison is made by SQL rather than a
+    stubbed lookup; a Python-side comparison would pass a stub but still insert a duplicate here.
+    """
+    table = SeedEntry.__table__
+    original_type = table.c.dataset_name.type
+    table.drop(sqlite_instance.engine)
+    table.c.dataset_name.type = String(collation="NOCASE")
+    try:
+        table.create(sqlite_instance.engine)
+
+        await sqlite_instance.add_seeds_to_memory_async(
+            seeds=[SeedPrompt(value="prompt1", dataset_name="Dataset", data_type="text")], added_by="tester"
+        )
+        await sqlite_instance.add_seeds_to_memory_async(
+            seeds=[SeedPrompt(value="prompt1", dataset_name="dataset", data_type="text")], added_by="tester"
+        )
+
+        assert len(sqlite_instance.get_seeds()) == 1
+    finally:
+        table.c.dataset_name.type = original_type
+
+
+async def test_add_seed_prompts_dedupe_follows_a_trailing_blank_insensitive_collation(
+    sqlite_instance: MemoryInterface,
+):
+    """T-SQL ignores trailing blanks too, so normalizing case in Python would not have been enough."""
+    await sqlite_instance.add_seeds_to_memory_async(
+        seeds=[SeedPrompt(value="prompt1", dataset_name="alpha", data_type="text")], added_by="tester"
+    )
+
+    real_get_seeds = sqlite_instance.get_seeds
+
+    def blank_insensitive_get_seeds(*, dataset_name=None, **kwargs):
+        rows = real_get_seeds(**kwargs)
+        if not dataset_name:
+            return rows
+        return [row for row in rows if (row.dataset_name or "").rstrip() == dataset_name.rstrip()]
+
+    with patch.object(sqlite_instance, "get_seeds", side_effect=blank_insensitive_get_seeds):
+        await sqlite_instance.add_seeds_to_memory_async(
+            seeds=[SeedPrompt(value="prompt1", dataset_name="alpha   ", data_type="text")], added_by="tester"
+        )
+
+    assert len(sqlite_instance.get_seeds()) == 1
+
+
+async def test_add_seed_prompts_dedupe_groups_each_dataset_name_separately(sqlite_instance: MemoryInterface):
+    """Seeds arriving for several datasets at once still get one query per dataset name."""
+    prompts: Sequence[SeedPrompt] = [
+        SeedPrompt(value="shared", dataset_name="alpha", data_type="text"),
+        SeedPrompt(value="shared", dataset_name="beta", data_type="text"),
+        SeedPrompt(value="loose", data_type="text"),
+    ]
+
+    with patch.object(sqlite_instance, "get_seeds", wraps=sqlite_instance.get_seeds) as spied:
+        await sqlite_instance.add_seeds_to_memory_async(seeds=prompts, added_by="tester")
+
+    queried = [call.kwargs["dataset_name"] for call in spied.call_args_list]
+    assert sorted(queried, key=lambda name: (name is None, name or "")) == ["alpha", "beta", None]
+    assert len(sqlite_instance.get_seeds()) == 3
+
+
+async def test_add_seed_prompts_dedupe_treats_an_empty_dataset_name_as_unfiltered(
+    sqlite_instance: MemoryInterface,
+):
+    """An empty name filters nothing in SQL, so it must match a hash stored under any dataset."""
+    await sqlite_instance.add_seeds_to_memory_async(
+        seeds=[SeedPrompt(value="prompt1", dataset_name="alpha", data_type="text")], added_by="tester"
+    )
+
+    await sqlite_instance.add_seeds_to_memory_async(
+        seeds=[SeedPrompt(value="prompt1", dataset_name="", data_type="text")], added_by="tester"
+    )
+
+    assert len(sqlite_instance.get_seeds()) == 1
+
+
 async def test_add_seed_prompts_duplicate_entries_different_datasets(sqlite_instance: MemoryInterface):
     prompts: Sequence[SeedPrompt] = [
         SeedPrompt(value="prompt1", dataset_name="test_dataset", data_type="text"),