diff --git a/doc/code/memory/8_seed_database.ipynb b/doc/code/memory/8_seed_database.ipynb index a87a92ef66..9900b0b3ab 100644 --- a/doc/code/memory/8_seed_database.ipynb +++ b/doc/code/memory/8_seed_database.ipynb @@ -228,6 +228,38 @@ "print(\"----------\")\n", "print_group(seed_groups[0])" ] + }, + { + "cell_type": "markdown", + "id": "05f3d1d3", + "metadata": {}, + "source": [ + "## Removing Seeds from the Database\n", + "\n", + "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.\n", + "\n", + "As a safety measure, at least one filter must be provided. Calling it with no filters raises a `ValueError` to prevent accidentally deleting the entire seed database." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "51a51a3b", + "metadata": {}, + "outputs": [], + "source": [ + "# Preview the seeds that will be removed using the same filters\n", + "seeds_to_remove = memory.get_seeds(dataset_name=\"pyrit_example_dataset\")\n", + "print(f\"Seeds matching the filter: {len(seeds_to_remove)}\")\n", + "\n", + "# Remove them and get back the number of seeds deleted\n", + "removed_count = await memory.remove_seeds_from_memory_async(dataset_name=\"pyrit_example_dataset\") # type: ignore\n", + "print(f\"Removed {removed_count} seeds\")\n", + "\n", + "# Confirm they are gone\n", + "seeds = memory.get_seeds(dataset_name=\"pyrit_example_dataset\")\n", + "print(f\"Seeds remaining in dataset: {len(seeds)}\")" + ] } ], "metadata": { diff --git a/doc/code/memory/8_seed_database.py b/doc/code/memory/8_seed_database.py index 8284d16d59..bf76e793f5 100644 --- a/doc/code/memory/8_seed_database.py +++ b/doc/code/memory/8_seed_database.py @@ -113,3 +113,23 @@ def print_group(seed_group): seed_groups = memory.get_seed_groups(data_types=["image_path"], dataset_name="pyrit_example_dataset") print("----------") print_group(seed_groups[0]) + +# %% [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. +# +# As a safety measure, at least one filter must be provided. Calling it with no filters raises a `ValueError` to prevent accidentally deleting the entire seed database. + +# %% +# Preview the seeds that will be removed using the same filters +seeds_to_remove = memory.get_seeds(dataset_name="pyrit_example_dataset") +print(f"Seeds matching the filter: {len(seeds_to_remove)}") + +# Remove them and get back the number of seeds deleted +removed_count = await memory.remove_seeds_from_memory_async(dataset_name="pyrit_example_dataset") # type: ignore +print(f"Removed {removed_count} seeds") + +# Confirm they are gone +seeds = memory.get_seeds(dataset_name="pyrit_example_dataset") +print(f"Seeds remaining in dataset: {len(seeds)}") diff --git a/pyrit/memory/memory_interface.py b/pyrit/memory/memory_interface.py index 602f13b6bf..bb408c10a1 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -2070,7 +2070,7 @@ def cleanup(self) -> None: # Ensure cleanup happens even if the object is garbage collected before process exits weakref.finalize(self, self.dispose_engine) - def get_seeds( + def _build_seed_filter_conditions( self, *, value: str | None = None, @@ -2087,13 +2087,17 @@ def get_seeds( parameters: Sequence[str] | None = None, metadata: dict[str, str | int] | None = None, prompt_group_ids: Sequence[uuid.UUID] | None = None, - ) -> Sequence[Seed]: + ) -> "list[ColumnElement[bool]]": """ - Retrieve a list of seed prompts based on the specified filters. + Build SQLAlchemy filter conditions shared by seed query and removal methods. + + This centralizes the filter-building logic so that get_seeds and + 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. - value_sha256 (str): The SHA256 hash of the value to match. If None, all values are returned. + value_sha256 (Sequence[str] | None): A list of SHA256 hashes of values to match. + If None, all values are returned. dataset_name (str): The dataset name to match exactly. If None, all dataset names are considered. dataset_name_pattern (str): A pattern to match dataset names using SQL LIKE syntax. Supports wildcards: % (any characters) and _ (single character). @@ -2118,9 +2122,9 @@ def get_seeds( 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. + list[ColumnElement[bool]]: A list of SQLAlchemy filter conditions. """ - conditions = [] + conditions: list[ColumnElement[bool]] = [] # Apply filters for non-list fields if value: @@ -2157,6 +2161,77 @@ def get_seeds( if metadata: conditions.append(self._get_seed_metadata_conditions(metadata=metadata)) + return conditions + + def get_seeds( + self, + *, + value: str | None = None, + value_sha256: Sequence[str] | None = None, + dataset_name: str | None = None, + dataset_name_pattern: str | None = None, + data_types: Sequence[str] | None = None, + harm_categories: Sequence[str] | None = None, + added_by: str | None = None, + authors: Sequence[str] | None = None, + groups: Sequence[str] | None = None, + source: str | None = None, + seed_type: SeedType | None = None, + parameters: Sequence[str] | None = None, + metadata: dict[str, str | int] | None = None, + prompt_group_ids: Sequence[uuid.UUID] | None = None, + ) -> Sequence[Seed]: + """ + Retrieve a list of seed prompts based on the specified filters. + + Args: + value (str): The value to match by substring. If None, all values are returned. + value_sha256 (Sequence[str] | None): A list of SHA256 hashes of values to match. + If None, all values are returned. + dataset_name (str): The dataset name to match exactly. If None, all dataset names are considered. + dataset_name_pattern (str): A pattern to match dataset names using SQL LIKE syntax. + Supports wildcards: % (any characters) and _ (single character). + Examples: "harm%" matches names starting with "harm", "%test%" matches names containing "test". + If both dataset_name and dataset_name_pattern are provided, dataset_name takes precedence. + data_types (Sequence[str] | None): List of data types to filter seed prompts by + (e.g., text, image_path). + harm_categories (Sequence[str]): A list of harm categories to filter by. If None, + all harm categories are considered. + Specifying multiple harm categories returns only prompts that are marked with all harm categories. + added_by (str): The user who added the prompts. + authors (Sequence[str]): A list of authors to filter by. + Note that this filters by substring, so a query for "Adam Jones" may not return results if the record + is "A. Jones", "Jones, Adam", etc. If None, all authors are considered. + groups (Sequence[str]): A list of groups to filter by. If None, all groups are considered. + source (str): The source to filter by. If None, all sources are considered. + seed_type (SeedType): The type of seed to filter by ("prompt", "objective", or + "simulated_conversation"). + parameters (Sequence[str]): A list of parameters to filter by. Specifying parameters effectively returns + 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. + + Returns: + Sequence[Seed]: A list of seeds (e.g., SeedPrompt, SeedObjective, SeedSimulatedConversation) + matching the criteria. + """ + conditions = self._build_seed_filter_conditions( + value=value, + value_sha256=value_sha256, + dataset_name=dataset_name, + dataset_name_pattern=dataset_name_pattern, + data_types=data_types, + harm_categories=harm_categories, + added_by=added_by, + authors=authors, + groups=groups, + source=source, + seed_type=seed_type, + parameters=parameters, + metadata=metadata, + prompt_group_ids=prompt_group_ids, + ) + try: memory_entries: Sequence[SeedEntry] = self._query_entries( SeedEntry, @@ -2167,6 +2242,101 @@ def get_seeds( logger.exception(f"Failed to retrieve prompts with dataset name {dataset_name} with error {e}") raise + async def remove_seeds_from_memory_async( + self, + *, + value: str | None = None, + value_sha256: Sequence[str] | None = None, + dataset_name: str | None = None, + dataset_name_pattern: str | None = None, + data_types: Sequence[str] | None = None, + harm_categories: Sequence[str] | None = None, + added_by: str | None = None, + authors: Sequence[str] | None = None, + groups: Sequence[str] | None = None, + source: str | None = None, + seed_type: SeedType | None = None, + parameters: Sequence[str] | None = None, + metadata: dict[str, str | int] | None = None, + prompt_group_ids: Sequence[uuid.UUID] | None = None, + ) -> int: + """ + Remove seeds matching the specified filters from memory. + + Accepts the same filter parameters as get_seeds. It is recommended to call get_seeds with the same + filters first to preview which seeds will be removed. At least one filter must be provided to + prevent accidental deletion of all seeds. The deletion runs in a single transaction, so it works + consistently across SQLite and Azure SQL. + + Args: + value (str): The value to match by substring. If None, all values are considered. + value_sha256 (Sequence[str] | None): A list of SHA256 hashes of values to match. + If None, all values are considered. + dataset_name (str): The dataset name to match exactly. If None, all dataset names are considered. + dataset_name_pattern (str): A pattern to match dataset names using SQL LIKE syntax. + Supports wildcards: % (any characters) and _ (single character). + Examples: "harm%" matches names starting with "harm", "%test%" matches names containing "test". + If both dataset_name and dataset_name_pattern are provided, dataset_name takes precedence. + data_types (Sequence[str] | None): List of data types to filter seed prompts by + (e.g., text, image_path). + harm_categories (Sequence[str]): A list of harm categories to filter by. If None, + all harm categories are considered. + Specifying multiple harm categories matches only prompts that are marked with all harm categories. + added_by (str): The user who added the prompts. + authors (Sequence[str]): A list of authors to filter by. + Note that this filters by substring, so a query for "Adam Jones" may not return results if the record + is "A. Jones", "Jones, Adam", etc. If None, all authors are considered. + groups (Sequence[str]): A list of groups to filter by. If None, all groups are considered. + source (str): The source to filter by. If None, all sources are considered. + seed_type (SeedType): The type of seed to filter by ("prompt", "objective", or + "simulated_conversation"). + 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. + + Returns: + int: The number of seeds removed. + + Raises: + ValueError: If no filters are provided. + """ + conditions = self._build_seed_filter_conditions( + value=value, + value_sha256=value_sha256, + dataset_name=dataset_name, + dataset_name_pattern=dataset_name_pattern, + data_types=data_types, + harm_categories=harm_categories, + added_by=added_by, + authors=authors, + groups=groups, + source=source, + seed_type=seed_type, + parameters=parameters, + metadata=metadata, + prompt_group_ids=prompt_group_ids, + ) + + # Guard against accidental "delete all": require at least one filter. + if not conditions: + raise ValueError( + "At least one filter parameter must be provided to remove_seeds_from_memory_async. " + "Calling without filters would delete all seeds." + ) + + # 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 + def _add_list_conditions( self, field: InstrumentedAttribute[Any], conditions: list[Any], values: Sequence[str] | None = None ) -> None: diff --git a/tests/unit/memory/memory_interface/test_interface_remove_seeds.py b/tests/unit/memory/memory_interface/test_interface_remove_seeds.py new file mode 100644 index 0000000000..348414af3f --- /dev/null +++ b/tests/unit/memory/memory_interface/test_interface_remove_seeds.py @@ -0,0 +1,121 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import pytest + +from pyrit.memory import MemoryInterface +from pyrit.models import SeedPrompt + + +async def test_remove_seeds_by_dataset_name(sqlite_instance: MemoryInterface): + seed_prompts = [ + SeedPrompt(value="prompt1", dataset_name="to_delete", data_type="text"), + SeedPrompt(value="prompt2", dataset_name="to_keep", data_type="text"), + ] + await sqlite_instance.add_seeds_to_memory_async(seeds=seed_prompts, added_by="test") + + removed = await sqlite_instance.remove_seeds_from_memory_async(dataset_name="to_delete") + assert removed == 1 + + remaining = sqlite_instance.get_seeds() + assert len(remaining) == 1 + assert remaining[0].dataset_name == "to_keep" + + +async def test_remove_seeds_by_dataset_name_pattern(sqlite_instance: MemoryInterface): + seed_prompts = [ + SeedPrompt(value="prompt1", dataset_name="harm_category_1", data_type="text"), + SeedPrompt(value="prompt2", dataset_name="harm_category_2", data_type="text"), + SeedPrompt(value="prompt3", dataset_name="safe_dataset", data_type="text"), + ] + await sqlite_instance.add_seeds_to_memory_async(seeds=seed_prompts, added_by="test") + + removed = await sqlite_instance.remove_seeds_from_memory_async(dataset_name_pattern="harm%") + assert removed == 2 + + remaining = sqlite_instance.get_seeds() + assert len(remaining) == 1 + assert remaining[0].dataset_name == "safe_dataset" + + +async def test_remove_seeds_by_added_by(sqlite_instance: MemoryInterface): + seed_prompts = [ + SeedPrompt(value="prompt1", dataset_name="ds1", added_by="user1", data_type="text"), + SeedPrompt(value="prompt2", dataset_name="ds2", added_by="user2", data_type="text"), + ] + await sqlite_instance.add_seeds_to_memory_async(seeds=seed_prompts) + + removed = await sqlite_instance.remove_seeds_from_memory_async(added_by="user1") + assert removed == 1 + + remaining = sqlite_instance.get_seeds() + assert len(remaining) == 1 + assert remaining[0].added_by == "user2" + + +async def test_remove_seeds_multi_filter_narrowing(sqlite_instance: MemoryInterface): + seed_prompts = [ + SeedPrompt(value="prompt1", dataset_name="ds1", added_by="user1", data_type="text"), + SeedPrompt(value="prompt2", dataset_name="ds1", added_by="user2", data_type="text"), + SeedPrompt(value="prompt3", dataset_name="ds2", added_by="user1", data_type="text"), + ] + await sqlite_instance.add_seeds_to_memory_async(seeds=seed_prompts) + + removed = await sqlite_instance.remove_seeds_from_memory_async(dataset_name="ds1", added_by="user1") + assert removed == 1 + + remaining = sqlite_instance.get_seeds() + assert len(remaining) == 2 + + +async def test_remove_seeds_no_filter_raises_value_error(sqlite_instance: MemoryInterface): + seed_prompts = [ + SeedPrompt(value="prompt1", dataset_name="ds1", data_type="text"), + ] + await sqlite_instance.add_seeds_to_memory_async(seeds=seed_prompts, added_by="test") + + with pytest.raises(ValueError, match="At least one filter"): + await sqlite_instance.remove_seeds_from_memory_async() + + +async def test_remove_seeds_returns_zero_when_no_match(sqlite_instance: MemoryInterface): + seed_prompts = [ + SeedPrompt(value="prompt1", dataset_name="ds1", data_type="text"), + ] + await sqlite_instance.add_seeds_to_memory_async(seeds=seed_prompts, added_by="test") + + removed = await sqlite_instance.remove_seeds_from_memory_async(dataset_name="nonexistent") + assert removed == 0 + + remaining = sqlite_instance.get_seeds() + assert len(remaining) == 1 + + +async def test_remove_seeds_by_harm_categories(sqlite_instance: MemoryInterface): + seed_prompts = [ + SeedPrompt(value="prompt1", harm_categories=["violence"], data_type="text"), + SeedPrompt(value="prompt2", harm_categories=["fraud"], data_type="text"), + ] + await sqlite_instance.add_seeds_to_memory_async(seeds=seed_prompts, added_by="test") + + removed = await sqlite_instance.remove_seeds_from_memory_async(harm_categories=["violence"]) + assert removed == 1 + + remaining = sqlite_instance.get_seeds() + assert len(remaining) == 1 + assert remaining[0].harm_categories == ["fraud"] + + +async def test_remove_seeds_by_source(sqlite_instance: MemoryInterface): + seed_prompts = [ + SeedPrompt(value="prompt1", dataset_name="ds1", source="internal", data_type="text"), + SeedPrompt(value="prompt2", dataset_name="ds2", source="external", data_type="text"), + ] + await sqlite_instance.add_seeds_to_memory_async(seeds=seed_prompts, added_by="test") + + removed = await sqlite_instance.remove_seeds_from_memory_async(source="internal") + assert removed == 1 + + remaining = sqlite_instance.get_seeds() + assert len(remaining) == 1 + assert remaining[0].source == "external"