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

[varunj-msft]

Description

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 — which makes the second load of a dataset far more expensive than the first.

Loading the default datasets adds roughly 158,000 seeds, and this lookup dominated the entire load. Measured on main against a real database:

• First load, empty table growing to 158,580 rows: about 210 seconds.
• Second load, table already holding 158,580 rows: over 1140 seconds, and still running when I stopped it.

LoadDefaultDatasets.initialize_async has no already-loaded short-circuit, so it re-resolves the whole selection every time it runs. A session-scoped backend serving many scenarios therefore pays that second, much worse, cost repeatedly.

This prepares the seeds first, then resolves which hashes are already stored using chunked lookups rather than one query per seed, so a full dataset resolves in a handful of queries instead of roughly 158,000. The chunk size of 500 is deliberately well under SQL Server's 2100-parameter ceiling.

Existing behaviour is preserved deliberately, not incidentally:

• The stored keys are snapshotted once, before any comparison, so duplicates within a single call are all still inserted. Only what storage held when the call started counts as already present. This matches what the per-seed loop did, since nothing it inserted was visible to its own subsequent lookups either.
• A seed with no dataset name still matches its hash in any dataset, mirroring the unfiltered lookup it previously performed.

There is no schema change and no migration here. Adding a bounded type and an index on value_sha256 would be the more complete fix, but that needs an Alembic revision, and alembic check runs on every memory init, so it is intentionally left as a follow-up rather than landing during a release window. This change is purely query shape and requires no database changes at all.

Part of the v1.1.0 release wave with #2510, #2511 and #2512.

Tests and Documentation

Four new tests in tests/unit/memory/memory_interface/test_interface_seed_prompts.py:

• test_add_seed_prompts_duplicates_within_one_call_are_all_stored pins the snapshot semantics, which is the subtle way a change here could silently alter behaviour.
• test_add_seed_prompts_without_dataset_name_matches_any_dataset covers the unfiltered lookup path.
• test_add_seed_prompts_dedupes_across_chunk_boundaries uses 1025 seeds, two full chunks plus a partial, then re-adds them with one new seed to confirm only the new one is stored.
• test_add_seed_prompts_queries_are_batched is a regression guard: it asserts get_seeds is called exactly once for 50 seeds, so a revert to per-seed lookups fails loudly instead of quietly getting slow again.

The existing duplicate-handling tests, same dataset and different datasets, are unchanged and still pass, which is the main evidence that this is a pure performance change.

I also diffed old versus new results across seven edge cases, including 2500 hashes in a single call to confirm the chunking respects the parameter ceiling. The retained and skipped sets were identical in every case.

Ran pytest tests/unit/memory/memory_interface/test_interface_seed_prompts.py: 70 passed.

No documentation changes: the public API and its behaviour are unchanged. The add_seeds_to_memory_async docstring now states the within-call duplicate rule explicitly, which was previously only implied. JupyText was not run and is not applicable: no notebooks or code samples are affected.

[hannahwestra25]

Suggested change

	hashes = sorted({prompt.value_sha256 for prompt in seeds if prompt.value_sha256})
	existing_pairs: set[tuple[str, str]] = set()
	existing_hashes: set[str] = set()
	for index in range(0, len(hashes), _SEED_DEDUPE_CHUNK_SIZE):
	chunk = hashes[index : index + _SEED_DEDUPE_CHUNK_SIZE]
	for existing in self.get_seeds(value_sha256=chunk):
	if not existing.value_sha256:
	continue
	existing_hashes.add(existing.value_sha256)
	if existing.dataset_name:
	existing_pairs.add((existing.value_sha256, existing.dataset_name))
	return existing_pairs, existing_hashes
	hashes_by_dataset: dict[str | None, set[str]] = {}
	for seed in seeds:
	if seed.value_sha256:
	# Empty names were previously treated as an unfiltered lookup.
	dataset_name = seed.dataset_name or None
	hashes_by_dataset.setdefault(dataset_name, set()).add(seed.value_sha256)
	existing_pairs: set[tuple[str, str]] = set()
	existing_hashes: set[str] = set()
	for dataset_name, hashes in hashes_by_dataset.items():
	ordered_hashes = sorted(hashes)
	for index in range(0, len(ordered_hashes), _SEED_DEDUPE_CHUNK_SIZE):
	chunk = ordered_hashes[index : index + _SEED_DEDUPE_CHUNK_SIZE]
	matches = self.get_seeds(
	value_sha256=chunk,
	dataset_name=dataset_name,
	)
	for match in matches:
	if not match.value_sha256:
	continue
	if dataset_name is None:
	existing_hashes.add(match.value_sha256)
	else:
	# SQL has already applied the backend's collation rules.
	existing_pairs.add((match.value_sha256, dataset_name))
	return existing_pairs, existing_hashes

[hannahwestra25]

otherwise we're switching from equality in sql to python tuple equality (ie dataset == Dataset in sql)

[hannahwestra25]

dataset_binds = 1 if dataset_name is not None else 0
chunk_size = self._MAX_BIND_VARS - dataset_binds

for index in range(0, len(hashes), chunk_size):
    chunk = hashes[index : index + chunk_size]

bc we can exceed the backend limit by one since we bind _MAX_BIND_VARS hashes with dataset_name,