FEAT implement batching for memory interface

[maifeeulasad]

Description

Closes #845

When MemoryInterface query methods (get_scores, get_message_pieces, get_attack_results, get_scenario_results) receive large IN-clause parameter lists (e.g., 1000+ prompt IDs), the generated SQL exceeds SQLite's per-statement bind variable limit of 999, causing the query to fail. This PR adds transparent batching: large IN-clause lists are split into chunks, executed as separate queries, and the results are merged and deduplicated.

What changed

Core batching infrastructure (memory_interface.py):

• _MAX_BIND_VARS class attribute on MemoryInterface (default: 500). Subclasses can override — AzureSQLMemory sets it to 2000 since Azure SQL supports 2100 params.
• _execute_batched_query — Low-level helper that splits a single IN-clause into batches, runs separate queries, and deduplicates results by primary key. Accepts an optional batch_size override for callers that need a smaller batch.
• _query_with_list_params — Higher-level helper used by get_message_pieces and get_attack_results. Classifies list parameters as "small" (fit in one query) or "large" (need batching), adds small params to SQL conditions, batches on the first large param, and filters remaining large params in Python. Computes an effective batch size that accounts for bind variables consumed by small params, preventing cumulative overflow.
• _dedup_attack_entries — Static helper extracted from get_attack_results to deduplicate AttackResultEntry rows by conversation_id (legacy data integrity workaround). Applied consistently across all code paths (batched and non-batched).

Method changes:

• get_scores — Batches score_ids via _execute_batched_query. Adds early return for score_ids=[] for consistency with other methods.
• get_message_pieces — Routes list params (prompt_ids, original_values, converted_values, converted_value_sha256) through _query_with_list_params.
• get_attack_results — Routes attack_result_ids and objective_sha256 through _query_with_list_params. Preserves dedup-by-conversation-id on all paths.
• get_scenario_results — Batches scenario_result_ids and inner all_conversation_ids via _execute_batched_query.

azure_sql_memory.py:

• Overrides _MAX_BIND_VARS = 2000 so Azure SQL uses larger batches (its limit is 2100).

Why this approach

• Multiple separate queries instead of OR-ing IN clauses — An earlier iteration tried combining batches with OR(col.IN(batch1), col.IN(batch2), ...) in a single statement, but this still binds all parameters in one statement and doesn't solve the limit.
• Conservative default of 500 — SQLite's actual limit is 999. Using 500 leaves headroom for bind variables from other conditions in the same query.
• Python-side filtering for multiple large params — When multiple list params exceed the limit simultaneously, only the first is batched in SQL. The rest are filtered in Python after fetching. This avoids nested batching complexity while handling an uncommon edge case correctly.

Tests and Documentation

New test file: tests/unit/memory/memory_interface/test_batching_scale.py (586 lines, 25 tests)

Covers:

• get_message_pieces with many prompt_ids, original_values, sha256 hashes (single and multiple large params)
• get_scores with many score_ids and empty list edge case
• get_attack_results with many IDs, many objective_sha256, and empty list edge case
• _execute_batched_query with small lists (single query), large lists (multiple queries), exact batch boundary, custom batch_size override
• Effective batch size reduction when small + large params are combined
• Filter preservation (role filter works correctly through batched queries)
• Deduplication correctness
• Query count verification via _query_entries spy

All 220 tests/unit/memory/memory_interface/ tests pass. Pre-commit hooks (ruff, mypy strict) clean.

[Copilot]

Pull request overview

This PR aims to address scaling limits in the PyRIT MemoryInterface by introducing “batched” filtering logic for large IN (...) queries, along with unit tests intended to validate behavior when many IDs/values are passed.

Changes:

• Added _SQLITE_MAX_BIND_VARS and _batched_in_condition helper in MemoryInterface and applied it to several query filters.
• Updated multiple retrieval methods (e.g., scores, message pieces, attack results, scenario results) to use the new helper for IN filtering.
• Added a new unit test module to exercise “batching” behavior and large-input queries against the SQLite-backed memory instance.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.

File	Description
pyrit/memory/memory_interface.py	Introduces a helper intended to avoid bind-variable limits and applies it across several query paths.
tests/unit/memory/memory_interface/test_batching_scale.py	Adds tests intended to validate that large ID/value inputs don’t break memory queries.

[Copilot]

Pull request overview

Copilot reviewed 2 out of 2 changed files in this pull request and generated 4 comments.

[maifeeulasad]

@romanlutz eagerly waiting for your feedback!

[romanlutz]

@romanlutz eagerly waiting for your feedback!

I apologize for taking ages. There were A LOT of PRs to work through.

[behnam-o]

Looks Great, Thank you! just a few minor non-blocking comments

[rlundeen2]

Approved with comments

[maifeeulasad]

Love to see this in action! I am quite thankful to all of you. You all are so friendly and helpful.