FIX: query and remove seeds by literal metadata keys

pyrit/memory/azure_sql_memory.py

@@ -374,7 +374,9 @@ def _get_metadata_conditions(self, *, prompt_metadata: dict[str, str | int]) ->
         """
         Generate SQL conditions for filtering by prompt metadata.
 
-        Uses JSON_VALUE() function specific to SQL Azure to query metadata fields in JSON format.
+        Uses OPENJSON() to compare free-form metadata keys literally rather than
+        interpreting them as JSON path expressions. Only strings, numbers, and
+        booleans match, preserving JSON_VALUE() scalar filtering semantics.
 
         Args:
             prompt_metadata (dict[str, str | int]): Dictionary of metadata key-value pairs to filter by.
@@ -383,15 +385,26 @@ def _get_metadata_conditions(self, *, prompt_metadata: dict[str, str | int]) ->
             list: List containing a single SQLAlchemy text condition with bound parameters.
         """
         json_validation = "ISJSON(prompt_metadata) = 1"
-        json_conditions = " AND ".join([f"JSON_VALUE(prompt_metadata, '$.{key}') = :{key}" for key in prompt_metadata])
+        json_conditions, bindparams = [], {}
+        for index, (key, value) in enumerate(prompt_metadata.items()):
+            key_param, value_param = f"md_key_{index}", f"md_value_{index}"
+            alias = f"metadata_{index}"
+            json_conditions.append(
+                f"EXISTS (SELECT 1 FROM OPENJSON(prompt_metadata) AS [{alias}] "
+                f"WHERE [{alias}].[key] = :{key_param} "
+                f"AND DATALENGTH([{alias}].[key]) = DATALENGTH(CAST(:{key_param} AS NVARCHAR(MAX))) "
+                f"AND [{alias}].[type] IN (1, 2, 3) "
+                f"AND [{alias}].[value] = :{value_param})"
+            )
+            bindparams[key_param] = key
+            bindparams[value_param] = str(value)
+
         # Combine both conditions
-        conditions = f"{json_validation} AND {json_conditions}"
+        conditions = f"{json_validation} AND " + " AND ".join(json_conditions)
 
         # Create SQL condition using SQLAlchemy's text() with bindparams
         # for safe parameter passing, preventing SQL injection
-        # Note: JSON_VALUE always returns nvarchar in SQL Server, so we must convert all values to strings
-        # to avoid type conversion errors when comparing mixed types (e.g., int and string)
-        condition = text(conditions).bindparams(**{key: str(value) for key, value in prompt_metadata.items()})
+        condition = text(conditions).bindparams(**bindparams)
         return [condition]
 
     def _get_message_pieces_prompt_metadata_conditions(

pyrit/memory/sqlite_memory.py

@@ -24,7 +24,6 @@
 from sqlalchemy.orm import InstrumentedAttribute, sessionmaker
 from sqlalchemy.orm.session import Session
 from sqlalchemy.pool import StaticPool
-from sqlalchemy.sql.expression import TextClause
 from sqlalchemy.util.concurrency import greenlet_spawn
 
 from pyrit.common.path import DB_DATA_PATH
@@ -36,6 +35,7 @@
     Base,
     PromptMemoryEntry,
     ScenarioResultEntry,
+    SeedEntry,
 )
 from pyrit.memory.memory_session import MemorySession
 from pyrit.memory.storage import DiskStorageIO
@@ -395,36 +395,46 @@ def _get_message_pieces_memory_label_conditions(self, *, memory_labels: dict[str
             )
         ]
 
-    def _get_message_pieces_prompt_metadata_conditions(
-        self, *, prompt_metadata: dict[str, str | int]
-    ) -> list[TextClause]:
+    def _get_message_pieces_prompt_metadata_conditions(self, *, prompt_metadata: dict[str, str | int]) -> list[Any]:
         """
         Generate SQLAlchemy filter conditions for filtering conversation pieces by prompt metadata.
 
+        Match literal object keys through json_each(), without interpreting keys as JSON paths.
+
         Returns:
             list: A list of SQLAlchemy conditions.
         """
-        json_conditions = " AND ".join(
-            [f"JSON_EXTRACT(prompt_metadata, '$.{key}') = :{key}" for key in prompt_metadata]
-        )
-
-        # Create SQL condition using SQLAlchemy's text() with bindparams
-        # Note: We do NOT convert values to string here, to allow integer comparison in JSON
-        condition = text(json_conditions).bindparams(**dict(prompt_metadata.items()))
-        return [condition]
+        conditions = []
+        metadata_entries = func.json_each(PromptMemoryEntry.prompt_metadata).table_valued("key", "value")
+        for key, value in prompt_metadata.items():
+            conditions.append(
+                select(1)
+                .select_from(metadata_entries)
+                .where(metadata_entries.c.key == key, metadata_entries.c.value == value)
+                .correlate(PromptMemoryEntry)
+                .exists()
+            )
+        return conditions
 
     def _get_seed_metadata_conditions(self, *, metadata: dict[str, str | int]) -> Any:
         """
         Generate SQLAlchemy filter conditions for filtering seed prompts by metadata.
 
-        Returns:
-            Any: A SQLAlchemy text condition with bound parameters.
-        """
-        json_conditions = " AND ".join([f"JSON_EXTRACT(prompt_metadata, '$.{key}') = :{key}" for key in metadata])
+        Match literal object keys through json_each(), without interpreting keys as JSON paths.
 
-        # Create SQL condition using SQLAlchemy's text() with bindparams
-        # Note: We do NOT convert values to string here, to allow integer comparison in JSON
-        return text(json_conditions).bindparams(**dict(metadata.items()))
+        Returns:
+            Any: A SQLAlchemy condition matching every requested metadata key and value.
+        """
+        metadata_entries = func.json_each(SeedEntry.prompt_metadata).table_valued("key", "value")
+        conditions = [
+            select(1)
+            .select_from(metadata_entries)
+            .where(metadata_entries.c.key == key, metadata_entries.c.value == value)
+            .correlate(SeedEntry)
+            .exists()
+            for key, value in metadata.items()
+        ]
+        return and_(*conditions)
 
     def _get_condition_json_property_match(
         self,

tests/integration/memory/test_azure_sql_memory_integration.py

@@ -7,6 +7,7 @@
 from uuid import uuid4
 
 import numpy as np
+import pytest
 from sqlalchemy.exc import SQLAlchemyError
 from sqlalchemy.orm import Session
 
@@ -303,6 +304,69 @@ def cleanup(session: Session) -> None:
     assert (await azuresql_instance.get_seeds_async(metadata={"key2": value1}, added_by=test_id)) == []
 
 
+@pytest.mark.run_only_if_all_tests
+@pytest.mark.parametrize(("string_value", "container_value"), [("{}", {}), ("[]", [])])
+async def test_metadata_string_filters_exclude_containers_async(
+    *,
+    azuresql_instance: AzureSQLMemory,
+    string_value: str,
+    container_value: dict[str, object] | list[object],
+) -> None:
+    test_id = str(uuid4())
+    scalar_metadata: dict[str, str | int] = {"payload": string_value}
+    container_metadata = {"payload": container_value}
+    scalar = SeedPrompt(value="scalar", dataset_name=test_id, metadata=scalar_metadata)
+    container = SeedPrompt(value="container", dataset_name=test_id, metadata=container_metadata)
+
+    try:
+        await azuresql_instance.add_seeds_to_memory_async(seeds=[scalar, container], added_by=test_id)
+        async with cleanup_conversation_data_async(azuresql_instance, [test_id]):
+            pieces = [
+                MessagePiece(
+                    conversation_id=test_id, role="user", original_value="scalar", prompt_metadata=scalar_metadata
+                ),
+                MessagePiece(
+                    conversation_id=test_id, role="user", original_value="container", prompt_metadata=container_metadata
+                ),
+            ]
+            await azuresql_instance.add_message_pieces_to_memory_async(message_pieces=pieces)
+            retrieved = await azuresql_instance.get_message_pieces_async(
+                conversation_id=test_id, prompt_metadata=scalar_metadata
+            )
+            assert [piece.id for piece in retrieved] == [pieces[0].id]
+
+        retrieved_seeds = await azuresql_instance.get_seeds_async(dataset_name=test_id, metadata=scalar_metadata)
+        assert [seed.id for seed in retrieved_seeds] == [scalar.id]
+        assert (
+            await azuresql_instance.remove_seeds_from_memory_async(dataset_name=test_id, metadata=scalar_metadata) == 1
+        )
+        assert [seed.id for seed in await azuresql_instance.get_seeds_async(dataset_name=test_id)] == [container.id]
+
+        scalar_group_id, container_group_id = uuid4(), uuid4()
+        groups = [
+            SeedPrompt(
+                value="scalar group", dataset_name=test_id, prompt_group_id=scalar_group_id, metadata=scalar_metadata
+            ),
+            SeedPrompt(value="scalar companion", dataset_name=test_id, prompt_group_id=scalar_group_id),
+            SeedPrompt(
+                value="container group",
+                dataset_name=test_id,
+                prompt_group_id=container_group_id,
+                metadata=container_metadata,
+            ),
+            SeedPrompt(value="container companion", dataset_name=test_id, prompt_group_id=container_group_id),
+        ]
+        await azuresql_instance.add_seeds_to_memory_async(seeds=groups, added_by=test_id)
+        assert (
+            await azuresql_instance.remove_seed_groups_from_memory_async(dataset_name=test_id, metadata=scalar_metadata)
+            == 2
+        )
+        remaining = await azuresql_instance.get_seeds_async(dataset_name=test_id)
+        assert {seed.id for seed in remaining} == {container.id, groups[2].id, groups[3].id}
+    finally:
+        await azuresql_instance.remove_seeds_from_memory_async(dataset_name=test_id)
+
+
 async def test_get_attack_results_by_labels(azuresql_instance: AzureSQLMemory):
     """
     Integration test for SQL Azure JSON filtering on labels.