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