From 345829a842e26e553c9e55a4b82c376778f8ab5e Mon Sep 17 00:00:00 2001 From: varunj-msft Date: Tue, 1 Sep 2026 02:57:12 +0000 Subject: [PATCH] PERF Batch the seed dedupe lookup when adding seeds to memory add_seeds_to_memory_async ran one SELECT per seed to decide whether that seed was already stored. value_sha256 is not indexed, so each of those scans the table, and the cost grows with the number of rows already present. Loading the default datasets adds roughly 158,000 seeds, where this dominated the whole load, and it was worse still on every later load because the table was already full. Prepare the seeds first, then resolve which hashes are already stored in chunked lookups instead of per seed. get_seeds issues its statement directly rather than through the batching helpers, so the lookup has to bound the IN clause itself. It chunks on _MAX_BIND_VARS, the per-statement bind ceiling each backend already tunes for itself, so SQL Server gets the 2000 AzureSQLMemory sets rather than a separate hardcoded limit. The lookups are grouped by dataset name so the name is still compared by the database. Equality there is a property of the column's collation: Azure SQL's default is case-insensitive and also ignores trailing blanks, so an existing ("Dataset") row matched an incoming ("dataset") one before this change. Comparing the names in Python would impose one fixed rule on every backend and insert a duplicate wherever the stored spelling differed, which no test would catch because CI runs on SQLite, whose default collation is case-sensitive. Since the database decides the match, the pair is keyed by the requested name rather than the stored one. An empty name filters nothing, exactly like None, so the two are normalized together to keep the lookup and the comparison agreeing. The stored keys are snapshotted once, before any comparison, which preserves the existing behaviour that duplicates within a single call are all inserted; only what storage held when the call started counts as already present. Seeds without a dataset name still match their hash in any dataset, matching the unfiltered lookup they previously performed. The dataset name takes one of those binds, so the hashes get the ceiling minus one rather than the whole budget. Binding a full chunk plus the name would have gone one over the limit each backend declares. --- pyrit/memory/memory_interface.py | 75 +++++++- .../test_interface_seed_prompts.py | 181 ++++++++++++++++++ 2 files changed, 251 insertions(+), 5 deletions(-) diff --git a/pyrit/memory/memory_interface.py b/pyrit/memory/memory_interface.py index 758c03bdcd..7603d0017c 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -2976,6 +2976,9 @@ 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. @@ -2983,18 +2986,80 @@ async def add_seeds_to_memory_async(self, *, seeds: Sequence[Seed], added_by: st 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. diff --git a/tests/unit/memory/memory_interface/test_interface_seed_prompts.py b/tests/unit/memory/memory_interface/test_interface_seed_prompts.py index 37d256612d..68f2d24171 100644 --- a/tests/unit/memory/memory_interface/test_interface_seed_prompts.py +++ b/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"),