diff --git a/doc/code/memory/8_seed_database.ipynb b/doc/code/memory/8_seed_database.ipynb index a87a92ef66..8e5b97681f 100644 --- a/doc/code/memory/8_seed_database.ipynb +++ b/doc/code/memory/8_seed_database.ipynb @@ -228,6 +228,59 @@ "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`. It accepts the same filtering parameters as `get_seeds` (plus an `exact` flag), 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 = memory.remove_seeds_from_memory(dataset_name=\"pyrit_example_dataset\")\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)}\")" + ] + }, + { + "cell_type": "markdown", + "id": "7b2c9e14", + "metadata": {}, + "source": [ + "### Removing entire groups\n", + "\n", + "`remove_seeds_from_memory` deletes only the individual seeds that match your filters. Because a seed group (for example a multimodal prompt made of text plus an image, or a multi-turn conversation) is stored as several seeds sharing a `prompt_group_id`, filtering by a single modality or attribute can leave a **partial group** behind. Some consequences to be aware of:\n", + "\n", + "- Deleting the sole objective while leaving its prompts produces an invalid `AttackSeedGroup`, and scenario initialization will raise a `ValueError`.\n", + "- Deleting one modality (e.g. the image) leaves a group that is still valid but now sends only text.\n", + "- Deleting one turn of a multi-turn conversation leaves the group with an incomplete context.\n", + "- Deleting the only role-bearing prompt in a sequence can cause a surviving multi-sequence group to fail role validation.\n", + "\n", + "For the most part these are user errors, but when you want to remove whole groups rather than individual seeds, use `remove_seed_groups_from_memory`. It applies the same filters, but removes every seed that shares a `prompt_group_id` with any match, so groups are never left partial. Note that it only affects seeds that belong to a group: a matching seed added individually (with no `prompt_group_id`) is skipped, so use `remove_seeds_from_memory` for those.\n", + "\n", + "> **Note on deleting by `value`.** For the remove methods, the `value` filter defaults to full-string equality (`exact=True`), so `remove_seeds_from_memory(value=\"the\")` deletes only seeds whose value is exactly `\"the\"` — not everything containing it. This differs from `get_seeds`, which always matches `value` by substring. Pass `exact=False` to opt into substring deletion when you really want it. As a general rule, preview with the same filters via `get_seeds(...)` first and prefer a specific filter (such as `dataset_name` or `value_sha256`) for deletion.\n", + "\n", + "> **Note on file-backed seeds.** For `image_path`, `audio_path`, and `video_path` seeds, removal deletes only the database record; the serialized file on disk is left in place. Delete those files separately if they are no longer needed." + ] } ], "metadata": { diff --git a/doc/code/memory/8_seed_database.py b/doc/code/memory/8_seed_database.py index 8284d16d59..c231215a7d 100644 --- a/doc/code/memory/8_seed_database.py +++ b/doc/code/memory/8_seed_database.py @@ -113,3 +113,39 @@ 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`. It accepts the same filtering parameters as `get_seeds` (plus an `exact` flag), 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 = memory.remove_seeds_from_memory(dataset_name="pyrit_example_dataset") +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)}") + +# %% [markdown] +# ### Removing entire groups +# +# `remove_seeds_from_memory` deletes only the individual seeds that match your filters. Because a seed group (for example a multimodal prompt made of text plus an image, or a multi-turn conversation) is stored as several seeds sharing a `prompt_group_id`, filtering by a single modality or attribute can leave a **partial group** behind. Some consequences to be aware of: +# +# - Deleting the sole objective while leaving its prompts produces an invalid `AttackSeedGroup`, and scenario initialization will raise a `ValueError`. +# - Deleting one modality (e.g. the image) leaves a group that is still valid but now sends only text. +# - Deleting one turn of a multi-turn conversation leaves the group with an incomplete context. +# - Deleting the only role-bearing prompt in a sequence can cause a surviving multi-sequence group to fail role validation. +# +# For the most part these are user errors, but when you want to remove whole groups rather than individual seeds, use `remove_seed_groups_from_memory`. It applies the same filters, but removes every seed that shares a `prompt_group_id` with any match, so groups are never left partial. Note that it only affects seeds that belong to a group: a matching seed added individually (with no `prompt_group_id`) is skipped, so use `remove_seeds_from_memory` for those. +# +# > **Note on deleting by `value`.** For the remove methods, the `value` filter defaults to full-string equality (`exact=True`), so `remove_seeds_from_memory(value="the")` deletes only seeds whose value is exactly `"the"` — not everything containing it. This differs from `get_seeds`, which always matches `value` by substring. Pass `exact=False` to opt into substring deletion when you really want it. As a general rule, preview with the same filters via `get_seeds(...)` first and prefer a specific filter (such as `dataset_name` or `value_sha256`) for deletion. +# +# > **Note on file-backed seeds.** For `image_path`, `audio_path`, and `video_path` seeds, removal deletes only the database record; the serialized file on disk is left in place. Delete those files separately if they are no longer needed. diff --git a/pyrit/memory/memory_interface.py b/pyrit/memory/memory_interface.py index e8f258da88..69c0209e1f 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -2190,10 +2190,11 @@ 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, + exact: bool = False, value_sha256: Sequence[str] | None = None, dataset_name: str | None = None, dataset_name_pattern: str | None = None, @@ -2207,13 +2208,20 @@ 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 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 (str): The value to match. By default this matches by substring; pass exact=True to + require full-string equality instead. If None, all values are returned. + exact (bool): When True, ``value`` is matched by full-string equality rather than substring. + Has no effect unless ``value`` is provided. Defaults to False (substring matching). + 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). @@ -2238,13 +2246,13 @@ 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: - conditions.append(SeedEntry.value.contains(value)) + conditions.append(SeedEntry.value == value if exact else SeedEntry.value.contains(value)) if value_sha256: conditions.append(SeedEntry.value_sha256.in_(value_sha256)) if dataset_name: @@ -2277,6 +2285,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, @@ -2287,8 +2366,245 @@ def get_seeds( logger.exception(f"Failed to retrieve prompts with dataset name {dataset_name} with error {e}") raise + def remove_seeds_from_memory( + self, + *, + value: str | None = None, + exact: bool = True, + 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: + """ + Delete a list of seed prompts based on the specified filters. + + Accepts the same filtering parameters as get_seeds (plus an exact flag). 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. + + Only database records are removed. For file-backed seeds (image_path, audio_path, video_path) the + serialized file on disk is left in place; delete those files separately if they are no longer needed. + + Args: + value (str): The value to match. For the remove methods this defaults to full-string equality + (exact=True) so a short or common value does not delete far more seeds than intended; pass + exact=False to match by substring instead. If None, all values are considered. + exact (bool): When True, ``value`` is matched by full-string equality rather than substring. + Has no effect unless ``value`` is provided. Defaults to True for the remove methods (the + safer choice for deletion). Note this differs from get_seeds, which always matches ``value`` + by substring. + 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. + SQLAlchemyError: If the database deletion fails (the transaction is rolled back). + """ + conditions = self._build_seed_filter_conditions( + value=value, + exact=exact, + 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. " + "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 remove_seed_groups_from_memory( + self, + *, + value: str | None = None, + exact: bool = True, + 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: + """ + Delete groups of seed prompts based on the provided filtering criteria. + + Unlike remove_seeds_from_memory, which deletes only the individual seeds that match, this + removes every seed that shares a prompt_group_id with any matching seed. This preserves group + integrity: filtering by a single modality (e.g. data_types=["image_path"]) or attribute removes + the whole group rather than leaving a partial group behind. It accepts the same filtering + parameters as get_seeds (plus an exact flag). It is recommended to call get_seed_groups with the + same filters first to preview which groups 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. + + Only seeds that belong to a group are affected: a matching seed with no prompt_group_id (for example, + one added individually via add_seeds_to_memory_async rather than as part of a group) is skipped, since + it has no group to expand. Use remove_seeds_from_memory to delete ungrouped seeds. + + Only database records are removed. For file-backed seeds (image_path, audio_path, video_path) the + serialized file on disk is left in place; delete those files separately if they are no longer needed. + + Args: + value (str): The value to match. For the remove methods this defaults to full-string equality + (exact=True) so a short or common value does not delete far more seeds than intended; pass + exact=False to match by substring instead. If None, all values are considered. + exact (bool): When True, ``value`` is matched by full-string equality rather than substring. + Has no effect unless ``value`` is provided. Defaults to True for the remove methods (the + safer choice for deletion). Note this differs from get_seeds, which always matches ``value`` + by substring. + 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 across all affected groups. + + Raises: + ValueError: If no filters are provided. + SQLAlchemyError: If the database deletion fails (the transaction is rolled back). + """ + conditions = self._build_seed_filter_conditions( + value=value, + exact=exact, + 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_seed_groups_from_memory. " + "Calling without filters would delete all seeds." + ) + + # Expand matches to whole groups and delete them in a single transaction. The matching group IDs + # are selected with a server-side subquery rather than materialized into Python and sent back as an + # IN (...) list, so a broad filter cannot exceed the backend's bound-parameter limit (e.g. Azure SQL). + # Seeds with a NULL prompt_group_id are excluded, so ungrouped matches are skipped. + with closing(self.get_session()) as session: + try: + group_id_subquery = ( + select(SeedEntry.prompt_group_id) + .where(and_(*conditions)) + .where(SeedEntry.prompt_group_id.isnot(None)) + .distinct() + ) + count = ( + session.query(SeedEntry) + .filter(SeedEntry.prompt_group_id.in_(group_id_subquery)) + .delete(synchronize_session=False) + ) + session.commit() + return count + except SQLAlchemyError as e: + session.rollback() + logger.exception(f"Failed to remove seed groups from memory: {e}") + raise + def _add_list_conditions( - self, field: InstrumentedAttribute[Any], conditions: list[Any], values: Sequence[str] | None = None + self, + field: InstrumentedAttribute[Any], + conditions: "list[ColumnElement[bool]]", + values: Sequence[str] | None = None, ) -> None: if values: conditions.extend(field.contains(value) for value in values) 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..8f78e9eac7 --- /dev/null +++ b/tests/unit/memory/memory_interface/test_interface_remove_seeds.py @@ -0,0 +1,398 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +from unittest.mock import patch + +import pytest +from sqlalchemy.exc import SQLAlchemyError + +from pyrit.memory import MemoryInterface +from pyrit.models import SeedGroup, SeedObjective, SeedPrompt + +# ========================================================================= +# remove_seeds_from_memory +# ========================================================================= + + +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 = sqlite_instance.remove_seeds_from_memory(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 = sqlite_instance.remove_seeds_from_memory(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 = sqlite_instance.remove_seeds_from_memory(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_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 = sqlite_instance.remove_seeds_from_memory(source="internal") + assert removed == 1 + + remaining = sqlite_instance.get_seeds() + assert len(remaining) == 1 + assert remaining[0].source == "external" + + +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 = sqlite_instance.remove_seeds_from_memory(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_authors(sqlite_instance: MemoryInterface): + seed_prompts = [ + SeedPrompt(value="prompt1", authors=["author1"], data_type="text"), + SeedPrompt(value="prompt2", authors=["author2"], data_type="text"), + ] + await sqlite_instance.add_seeds_to_memory_async(seeds=seed_prompts, added_by="test") + + removed = sqlite_instance.remove_seeds_from_memory(authors=["author1"]) + assert removed == 1 + + remaining = sqlite_instance.get_seeds() + assert len(remaining) == 1 + assert remaining[0].authors == ["author2"] + + +async def test_remove_seeds_by_groups(sqlite_instance: MemoryInterface): + seed_prompts = [ + SeedPrompt(value="prompt1", groups=["group1"], data_type="text"), + SeedPrompt(value="prompt2", groups=["group2"], data_type="text"), + ] + await sqlite_instance.add_seeds_to_memory_async(seeds=seed_prompts, added_by="test") + + removed = sqlite_instance.remove_seeds_from_memory(groups=["group1"]) + assert removed == 1 + + remaining = sqlite_instance.get_seeds() + assert len(remaining) == 1 + assert remaining[0].groups == ["group2"] + + +async def test_remove_seeds_by_parameters(sqlite_instance: MemoryInterface): + seed_prompts = [ + SeedPrompt(value="prompt1", parameters=["param1"], data_type="text"), + SeedPrompt(value="prompt2", parameters=["param2"], data_type="text"), + ] + await sqlite_instance.add_seeds_to_memory_async(seeds=seed_prompts, added_by="test") + + removed = sqlite_instance.remove_seeds_from_memory(parameters=["param1"]) + assert removed == 1 + assert len(sqlite_instance.get_seeds(parameters=["param1"])) == 0 + assert len(sqlite_instance.get_seeds(parameters=["param2"])) == 1 + + +async def test_remove_seeds_by_metadata(sqlite_instance: MemoryInterface): + seed_prompts = [ + SeedPrompt(value="prompt1", data_type="text", metadata={"key1": "value1"}), + SeedPrompt(value="prompt2", data_type="text", metadata={"key1": "value2"}), + ] + await sqlite_instance.add_seeds_to_memory_async(seeds=seed_prompts, added_by="test") + + removed = sqlite_instance.remove_seeds_from_memory(metadata={"key1": "value1"}) + assert removed == 1 + + remaining = sqlite_instance.get_seeds() + assert len(remaining) == 1 + assert remaining[0].value == "prompt2" + + +async def test_remove_seeds_by_data_type(sqlite_instance: MemoryInterface): + seed_prompts = [ + SeedPrompt(value="prompt1", dataset_name="ds1", data_type="text"), + SeedPrompt(value="prompt2", dataset_name="ds2", data_type="text"), + ] + await sqlite_instance.add_seeds_to_memory_async(seeds=seed_prompts, added_by="test") + + assert sqlite_instance.remove_seeds_from_memory(data_types=["image_path"]) == 0 + assert sqlite_instance.remove_seeds_from_memory(data_types=["text"]) == 2 + assert len(sqlite_instance.get_seeds()) == 0 + + +async def test_remove_seeds_by_seed_type(sqlite_instance: MemoryInterface): + prompt = SeedPrompt(value="a prompt", dataset_name="ds", data_type="text") + objective = SeedObjective(value="an objective") + await sqlite_instance.add_seeds_to_memory_async(seeds=[prompt, objective], added_by="test") + + removed = sqlite_instance.remove_seeds_from_memory(seed_type="objective") + assert removed == 1 + + remaining = sqlite_instance.get_seeds() + assert len(remaining) == 1 + assert remaining[0].seed_type == "prompt" + + +async def test_remove_seeds_by_value_sha256(sqlite_instance: MemoryInterface): + seed_prompts = [ + SeedPrompt(value="prompt1", dataset_name="ds", data_type="text"), + SeedPrompt(value="prompt2", dataset_name="ds", data_type="text"), + ] + await sqlite_instance.add_seeds_to_memory_async(seeds=seed_prompts, added_by="test") + target = next(seed for seed in sqlite_instance.get_seeds() if seed.value == "prompt1") + + removed = sqlite_instance.remove_seeds_from_memory(value_sha256=[target.value_sha256]) + assert removed == 1 + + remaining = sqlite_instance.get_seeds() + assert len(remaining) == 1 + assert remaining[0].value == "prompt2" + + +async def test_remove_seeds_by_prompt_group_ids(sqlite_instance: MemoryInterface): + group = SeedGroup(seeds=[SeedPrompt(value="grouped", dataset_name="ds", data_type="text", sequence=0)]) + await sqlite_instance.add_seed_groups_to_memory_async(prompt_groups=[group], added_by="test") + group_id = sqlite_instance.get_seeds()[0].prompt_group_id + + removed = sqlite_instance.remove_seeds_from_memory(prompt_group_ids=[group_id]) + assert removed == 1 + assert len(sqlite_instance.get_seeds()) == 0 + + +async def test_remove_seeds_by_value_substring_is_broad(sqlite_instance: MemoryInterface): + seed_prompts = [ + SeedPrompt(value="the quick fox", dataset_name="ds", data_type="text"), + SeedPrompt(value="the lazy dog", dataset_name="ds", data_type="text"), + SeedPrompt(value="a cat", dataset_name="ds", data_type="text"), + ] + await sqlite_instance.add_seeds_to_memory_async(seeds=seed_prompts, added_by="test") + + # exact=False opts into substring matching, which is broad: "the" matches both values that contain it. + removed = sqlite_instance.remove_seeds_from_memory(value="the", exact=False) + assert removed == 2 + + remaining = sqlite_instance.get_seeds() + assert len(remaining) == 1 + assert remaining[0].value == "a cat" + + +async def test_remove_seeds_by_value_defaults_to_exact(sqlite_instance: MemoryInterface): + seed_prompts = [ + SeedPrompt(value="the quick fox", dataset_name="ds", data_type="text"), + SeedPrompt(value="the lazy dog", dataset_name="ds", data_type="text"), + ] + await sqlite_instance.add_seeds_to_memory_async(seeds=seed_prompts, added_by="test") + + # The remove methods default to exact matching, so a substring like "the" removes nothing. + assert sqlite_instance.remove_seeds_from_memory(value="the") == 0 + # The full value matches exactly one seed. + assert sqlite_instance.remove_seeds_from_memory(value="the quick fox") == 1 + + remaining = sqlite_instance.get_seeds() + assert len(remaining) == 1 + assert remaining[0].value == "the lazy dog" + + +async def test_remove_seeds_by_value_exact_is_narrow(sqlite_instance: MemoryInterface): + seed_prompts = [ + SeedPrompt(value="the quick fox", dataset_name="ds", data_type="text"), + SeedPrompt(value="the lazy dog", dataset_name="ds", data_type="text"), + ] + await sqlite_instance.add_seeds_to_memory_async(seeds=seed_prompts, added_by="test") + + # exact=True only removes the fully-matching value, not substring siblings. + removed = sqlite_instance.remove_seeds_from_memory(value="the quick fox", exact=True) + assert removed == 1 + + remaining = sqlite_instance.get_seeds() + assert len(remaining) == 1 + assert remaining[0].value == "the lazy dog" + + +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 = sqlite_instance.remove_seeds_from_memory(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"): + sqlite_instance.remove_seeds_from_memory() + + +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 = sqlite_instance.remove_seeds_from_memory(dataset_name="nonexistent") + assert removed == 0 + + remaining = sqlite_instance.get_seeds() + assert len(remaining) == 1 + + +async def test_remove_seeds_rolls_back_on_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 patch("sqlalchemy.orm.session.Session.commit", side_effect=SQLAlchemyError("boom")): + with pytest.raises(SQLAlchemyError): + sqlite_instance.remove_seeds_from_memory(dataset_name="ds1") + + # The failed deletion is rolled back, so the seed is still present. + assert len(sqlite_instance.get_seeds()) == 1 + + +# ========================================================================= +# remove_seed_groups_from_memory +# ========================================================================= + + +async def test_remove_seed_groups_removes_entire_group(sqlite_instance: MemoryInterface): + group = SeedGroup( + seeds=[ + SeedPrompt(value="match_me", dataset_name="grouped", data_type="text", sequence=0, role="user"), + SeedPrompt(value="sibling", dataset_name="grouped", data_type="text", sequence=1, role="user"), + ] + ) + standalone = SeedGroup(seeds=[SeedPrompt(value="keep", dataset_name="other", data_type="text", sequence=0)]) + await sqlite_instance.add_seed_groups_to_memory_async(prompt_groups=[group, standalone], added_by="test") + + # Filter matches only one seed in the group, but the whole group is removed. + removed = sqlite_instance.remove_seed_groups_from_memory(value="match_me", exact=True) + assert removed == 2 + + remaining = sqlite_instance.get_seeds() + assert len(remaining) == 1 + assert remaining[0].value == "keep" + + +async def test_remove_seed_groups_spanning_multiple_datasets(sqlite_instance: MemoryInterface): + group = SeedGroup( + seeds=[ + SeedPrompt(value="seed_a", dataset_name="ds_a", data_type="text", sequence=0, role="user"), + SeedPrompt(value="seed_b", dataset_name="ds_b", data_type="text", sequence=1, role="user"), + ] + ) + await sqlite_instance.add_seed_groups_to_memory_async(prompt_groups=[group], added_by="test") + + # Matching a seed in one dataset removes the whole group, including the seed in the other dataset. + removed = sqlite_instance.remove_seed_groups_from_memory(dataset_name="ds_a") + assert removed == 2 + assert len(sqlite_instance.get_seeds()) == 0 + + +async def test_remove_seed_groups_preserves_non_matching_groups(sqlite_instance: MemoryInterface): + group_to_remove = SeedGroup(seeds=[SeedPrompt(value="g1", dataset_name="remove", data_type="text", sequence=0)]) + group_to_keep = SeedGroup(seeds=[SeedPrompt(value="g2", dataset_name="keep", data_type="text", sequence=0)]) + await sqlite_instance.add_seed_groups_to_memory_async( + prompt_groups=[group_to_remove, group_to_keep], added_by="test" + ) + + removed = sqlite_instance.remove_seed_groups_from_memory(dataset_name="remove") + assert removed == 1 + + remaining = sqlite_instance.get_seeds() + assert len(remaining) == 1 + assert remaining[0].value == "g2" + + +async def test_remove_seed_groups_skips_ungrouped_seeds(sqlite_instance: MemoryInterface): + # Seeds added individually keep prompt_group_id=None, so they belong to no group. + seed_prompts = [ + SeedPrompt(value="ungrouped", dataset_name="ds1", data_type="text"), + ] + await sqlite_instance.add_seeds_to_memory_async(seeds=seed_prompts, added_by="test") + + removed = sqlite_instance.remove_seed_groups_from_memory(dataset_name="ds1") + assert removed == 0 + assert len(sqlite_instance.get_seeds()) == 1 + + +async def test_remove_seed_groups_no_filter_raises_value_error(sqlite_instance: MemoryInterface): + group = SeedGroup(seeds=[SeedPrompt(value="prompt1", dataset_name="ds1", data_type="text", sequence=0)]) + await sqlite_instance.add_seed_groups_to_memory_async(prompt_groups=[group], added_by="test") + + with pytest.raises(ValueError, match="At least one filter"): + sqlite_instance.remove_seed_groups_from_memory() + + +async def test_remove_seed_groups_returns_zero_when_no_match(sqlite_instance: MemoryInterface): + group = SeedGroup(seeds=[SeedPrompt(value="prompt1", dataset_name="ds1", data_type="text", sequence=0)]) + await sqlite_instance.add_seed_groups_to_memory_async(prompt_groups=[group], added_by="test") + + removed = sqlite_instance.remove_seed_groups_from_memory(dataset_name="nonexistent") + assert removed == 0 + + remaining = sqlite_instance.get_seeds() + assert len(remaining) == 1 + + +async def test_remove_seed_groups_rolls_back_on_error(sqlite_instance: MemoryInterface): + group = SeedGroup(seeds=[SeedPrompt(value="prompt1", dataset_name="ds1", data_type="text", sequence=0)]) + await sqlite_instance.add_seed_groups_to_memory_async(prompt_groups=[group], added_by="test") + + with patch("sqlalchemy.orm.session.Session.commit", side_effect=SQLAlchemyError("boom")): + with pytest.raises(SQLAlchemyError): + sqlite_instance.remove_seed_groups_from_memory(dataset_name="ds1") + + # The failed deletion is rolled back, so the group is still present. + assert len(sqlite_instance.get_seeds()) == 1