FEAT: Add remove_seeds_from_memory_async to memory#2243
Conversation
There was a problem hiding this comment.
Pull request overview
This PR adds first-class support for deleting seed prompts from PyRIT memory by introducing remove_seeds_from_memory_async, mirroring the existing get_seeds filter surface so users no longer need backend-specific SQL for seed cleanup.
Changes:
- Extracts seed filter construction into a shared private helper (
_build_seed_filter_conditions) used by both seed querying and deletion. - Adds
remove_seeds_from_memory_asyncto delete matchingSeedEntryrows transactionally and return the number removed (with a no-filter safety guard). - Adds unit tests and documentation demonstrating a “preview via
get_seeds, then delete viaremove_seeds_from_memory_async” workflow.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 13 comments.
| File | Description |
|---|---|
pyrit/memory/memory_interface.py |
Adds shared seed-filter builder and new async seed removal API implemented via SQLAlchemy. |
tests/unit/memory/memory_interface/test_interface_remove_seeds.py |
Adds unit coverage for common filter cases and safety behavior. |
doc/code/memory/8_seed_database.py |
Documents seed removal usage and recommended preview-then-delete workflow. |
doc/code/memory/8_seed_database.ipynb |
Syncs the same documentation updates into the notebook format. |
| @pytest.mark.asyncio | ||
| async def test_remove_seeds_by_dataset_name(sqlite_instance: MemoryInterface): |
| @pytest.mark.asyncio | ||
| async def test_remove_seeds_by_dataset_name_pattern(sqlite_instance: MemoryInterface): |
| @pytest.mark.asyncio | ||
| async def test_remove_seeds_by_added_by(sqlite_instance: MemoryInterface): |
| @pytest.mark.asyncio | ||
| async def test_remove_seeds_multi_filter_narrowing(sqlite_instance: MemoryInterface): |
| @pytest.mark.asyncio | ||
| async def test_remove_seeds_no_filter_raises_value_error(sqlite_instance: MemoryInterface): |
| remove_seeds_from_memory_async stay in sync and cannot drift. | ||
|
|
||
| Args: | ||
| value (str): The value to match by substring. If None, all values are returned. |
|
|
||
| Args: | ||
| value (str): The value to match by substring. If None, all values are returned. | ||
| value_sha256 (str): The SHA256 hash of the value to match. If None, all values are returned. |
| prompt_group_ids (Sequence[uuid.UUID]): A list of prompt group IDs to filter by. | ||
|
|
||
| Returns: | ||
| Sequence[SeedPrompt]: A list of prompts matching the criteria. |
|
|
||
| Args: | ||
| value (str): The value to match by substring. If None, all values are considered. | ||
| value_sha256 (str): The SHA256 hash of the value to match. If None, all values are considered. |
| query = session.query(SeedEntry).filter(and_(*conditions)) | ||
| count = query.delete(synchronize_session="fetch") | ||
| session.commit() |
Adds an API to remove seed prompts from memory using the same filter parameters as get_seeds (dataset_name, dataset_name_pattern, added_by, harm_categories, authors, groups, source, value_sha256, etc.). Implementation: - Extract filter-building logic from get_seeds into shared _build_seed_filter_conditions helper to prevent drift - Add remove_seeds_from_memory_async that uses the shared helper - Require at least one filter (raises ValueError if none provided) - Single transaction with rollback on failure - Works across DuckDB, SQLite, Azure SQL (uses SQLAlchemy ORM) Unit tests cover: exact name filter, LIKE pattern, multi-filter narrowing, no-filter ValueError, zero-match return, harm_categories, and source filter. Docs: add a 'Removing Seeds from the Database' section to the seed database notebook (8_seed_database.py/.ipynb) demonstrating the preview-then-remove workflow. Resolves AB#5688 Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: f9e3cc10-5c77-4384-92d9-d27dad9baea2
2d45b57 to
0e3ae02
Compare
|
Thanks for the review! Addressed all points in the latest push:
|
| # Delete matching rows in a single transaction so it is atomic across backends. | ||
| with closing(self.get_session()) as session: | ||
| try: | ||
| query = session.query(SeedEntry).filter(and_(*conditions)) | ||
| count = query.delete(synchronize_session=False) | ||
| session.commit() | ||
| return count | ||
| except SQLAlchemyError as e: | ||
| session.rollback() | ||
| logger.exception(f"Failed to remove seeds from memory: {e}") | ||
| raise |
|
@microsoft-github-policy-service agree company="Microsoft" |
| prevent accidental deletion of all seeds. The deletion runs in a single transaction, so it works | ||
| consistently across DuckDB, SQLite, and Azure SQL. |
There was a problem hiding this comment.
DuckDB is no longer a supported memory backend. This should just say SQLite and Azure SQL.
| logger.exception(f"Failed to retrieve prompts with dataset name {dataset_name} with error {e}") | ||
| raise | ||
|
|
||
| async def remove_seeds_from_memory_async( |
| async def remove_seeds_from_memory_async( | ||
| self, | ||
| *, | ||
| value: str | None = None, |
There was a problem hiding this comment.
hmmm i'm kinda concerned about using value for deletion. It's much safer when we're just getting a seed but if a user inputs "the" for example or a commonly used word the substring matching would delete a lot of seeds.
| metadata: dict[str, str | int] | None = None, | ||
| prompt_group_ids: Sequence[uuid.UUID] | None = None, | ||
| ) -> Sequence[Seed]: | ||
| ) -> list[Any]: |
There was a problem hiding this comment.
Can we make this type stricter than list[Any]? The file already uses list[ColumnElement[bool]] for condition builders. Using Any here removes useful checking from both query and deletion paths.
| with closing(self.get_session()) as session: | ||
| try: | ||
| query = session.query(SeedEntry).filter(and_(*conditions)) | ||
| count = query.delete(synchronize_session=False) |
There was a problem hiding this comment.
What are we doing to file-backed image, audio, and video seeds? This removes the database row but leaves the serialized file behind. We should either clean up owned artifacts or clearly document that this only removes database records.
| parameters (Sequence[str]): A list of parameters to filter by. Specifying parameters effectively targets | ||
| prompt templates instead of prompts. | ||
| metadata (dict[str, str | int]): A free-form dictionary for tagging prompts with custom metadata. | ||
| prompt_group_ids (Sequence[uuid.UUID]): A list of prompt group IDs to filter by. |
There was a problem hiding this comment.
so prompt_group_ids would remove entire groups ? I'm wondering if we should have a separate remove_seed_groups_from_memory (and similarly for datasets)
| # %% [markdown] | ||
| # ## Removing Seeds from the Database | ||
| # | ||
| # Just as you can add and query seeds, you can remove them using `remove_seeds_from_memory_async`. It accepts the same filter parameters as `get_seeds`, so the recommended workflow is to preview the matching seeds with `get_seeds(...)` first, then remove them with the same filters. The method returns the number of seeds removed. |
There was a problem hiding this comment.
just noting that there are some possible cases where deleting will cause unintended consequences:
- Delete the sole objective, leave prompts: invalid AttackSeedGroup; scenario initialization raises ValueError.
- Delete one modality: group remains valid but sends only text instead of text plus image/audio.
- Delete one conversation turn: group may remain valid but its intended multi-turn context is incomplete.
- Delete the only role-bearing prompt in a sequence: surviving multi-sequence group may fail role validation.
these are just some examples and for the most part can be classified as user errors but might be something good to document
Description
PyRIT memory can add seed prompt datasets (via
add_seed_datasets_to_memory_async/add_seeds_to_memory_async) and query them withget_seeds, but there is no counterpart for removing seeds. Today users have to drop into raw SQL, which is error-prone and differs across the supported backends (DuckDB, SQLite, Azure SQL).This PR adds
remove_seeds_from_memory_async, which accepts the same filter parameters asget_seeds(dataset_name,dataset_name_pattern,added_by,harm_categories,authors,groups,source,value_sha256,data_types,seed_type,parameters,metadata,prompt_group_ids) and returns the number of seeds removed.Recommended workflow — preview, then delete with the same filters:
Implementation
get_seedsinto a shared private helper (_build_seed_filter_conditions) so both query and removal build conditions from a single source and cannot drift.remove_seeds_from_memory_asyncdeletes matching rows in a single transaction using the SQLAlchemy ORM (rollback on error), so it behaves consistently across DuckDB, SQLite, and Azure SQL.Safety
ValueErrorto prevent accidentally wiping the entire seed database.Tests and Documentation
Tests —
tests/unit/memory/memory_interface/test_interface_remove_seeds.pycovers:dataset_namematchdataset_name_pattern(SQLLIKE)dataset_name+added_by)ValueError0and deletes nothingharm_categories(list field) andsourcefiltersDocumentation — added a "Removing Seeds from the Database" section to
doc/code/memory/8_seed_database.py(and the synced.ipynb) demonstrating the preview-then-remove workflow.