From fab8680959afe2dda3fe9b96cbb07763f997a131 Mon Sep 17 00:00:00 2001 From: Chen Yufeiyang Date: Sun, 11 Oct 2026 16:33:50 +0800 Subject: [PATCH 1/4] Query prompt and seed metadata whose keys are not SQL identifiers `prompt_metadata` is documented as a free-form dictionary, but the SQLite and Azure SQL backends spelled each key into the JSON path and reused it as the bind-parameter name: JSON_EXTRACT(prompt_metadata, '$.{key}') = :{key} Only `[A-Za-z0-9_]` survives that. Anything else raises before the query runs, so the filter cannot be expressed at all: ArgumentError: This text() construct doesn't define a bound parameter named 'attack.id' `get_message_pieces_async(prompt_metadata={"attack.id": "v"})` against SQLiteMemory raises that today. A dot was also read as a path step, so a key that happened to survive would have looked up a nested object. Both backends now bind the path with the key quoted into it (`'$."{key}"'`) and give each condition positional parameter names, matching how #3063 fixed the same interpolation in the label conditions: JSON_EXTRACT(prompt_metadata, :pm_path_0) = :pm_value_0 Binding the path also keeps the value out of the statement text entirely. Seed metadata filtering goes through the same two code paths (`_get_seed_metadata_conditions`), so it is fixed as well. Regression tests query four keys holding a dot, hyphen, space and colon, plus one query naming three keys at once. Against the parent commit they fail with the ArgumentError above; the plain-key cases the suite already covered still pass. Not addressed: a key containing a double quote cannot be expressed in a SQLite or T-SQL JSON path, so it matches nothing rather than raising. That is the same residual #3063 left in the label path. Validation at this head: tests/unit/memory/memory_interface/, tests/unit/memory/test_sqlite_memory.py, tests/unit/memory/test_azure_sql_memory.py -- 648 passed, 1 skipped; ruff check and format clean. --- pyrit/memory/azure_sql_memory.py | 21 +++++++-- pyrit/memory/sqlite_memory.py | 33 +++++++++---- .../test_interface_prompts.py | 47 +++++++++++++++++++ 3 files changed, 88 insertions(+), 13 deletions(-) diff --git a/pyrit/memory/azure_sql_memory.py b/pyrit/memory/azure_sql_memory.py index 422713b641..89ff0d4d31 100644 --- a/pyrit/memory/azure_sql_memory.py +++ b/pyrit/memory/azure_sql_memory.py @@ -375,6 +375,11 @@ def _get_metadata_conditions(self, *, prompt_metadata: dict[str, str | int]) -> Uses JSON_VALUE() function specific to SQL Azure to query metadata fields in JSON format. + Metadata keys are free-form, so each key is quoted into the JSON path + and the path is bound rather than spelled out: a key holding a dot is a + literal key, not a path step, and a key holding a hyphen is neither a + JSON path nor a legal bind-parameter name. + Args: prompt_metadata (dict[str, str | int]): Dictionary of metadata key-value pairs to filter by. @@ -382,15 +387,21 @@ 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()): + path_param, value_param = f"md_path_{index}", f"md_value_{index}" + json_conditions.append(f"JSON_VALUE(prompt_metadata, :{path_param}) = :{value_param}") + bindparams[path_param] = f'$."{key}"' + # 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) + 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( diff --git a/pyrit/memory/sqlite_memory.py b/pyrit/memory/sqlite_memory.py index 01db1fe859..3be619087b 100644 --- a/pyrit/memory/sqlite_memory.py +++ b/pyrit/memory/sqlite_memory.py @@ -401,30 +401,47 @@ def _get_message_pieces_prompt_metadata_conditions( """ Generate SQLAlchemy filter conditions for filtering conversation pieces by prompt metadata. + Metadata keys are free-form, so each key is quoted into the JSON path + and the path is bound rather than spelled out: a key holding a dot is a + literal key, not a path step, and a key holding a hyphen is neither a + JSON path nor a legal bind-parameter name. + Returns: list: A list of SQLAlchemy conditions. """ - json_conditions = " AND ".join( - [f"JSON_EXTRACT(prompt_metadata, '$.{key}') = :{key}" for key in prompt_metadata] - ) + json_paths, bindparams = [], {} + for index, (key, value) in enumerate(prompt_metadata.items()): + path_param, value_param = f"pm_path_{index}", f"pm_value_{index}" + json_paths.append(f"JSON_EXTRACT(prompt_metadata, :{path_param}) = :{value_param}") + bindparams[path_param] = f'$."{key}"' + # Note: We do NOT convert values to string here, to allow integer comparison in JSON + bindparams[value_param] = value # 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())) + condition = text(" AND ".join(json_paths)).bindparams(**bindparams) return [condition] def _get_seed_metadata_conditions(self, *, metadata: dict[str, str | int]) -> Any: """ Generate SQLAlchemy filter conditions for filtering seed prompts by metadata. + Metadata keys are free-form, so each key is quoted into the JSON path + and the path is bound rather than spelled out, for the same reason as + in _get_message_pieces_prompt_metadata_conditions. + Returns: Any: A SQLAlchemy text condition with bound parameters. """ - json_conditions = " AND ".join([f"JSON_EXTRACT(prompt_metadata, '$.{key}') = :{key}" for key in metadata]) + json_paths, bindparams = [], {} + for index, (key, value) in enumerate(metadata.items()): + path_param, value_param = f"sm_path_{index}", f"sm_value_{index}" + json_paths.append(f"JSON_EXTRACT(prompt_metadata, :{path_param}) = :{value_param}") + bindparams[path_param] = f'$."{key}"' + # Note: We do NOT convert values to string here, to allow integer comparison in JSON + bindparams[value_param] = value # 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())) + return text(" AND ".join(json_paths)).bindparams(**bindparams) def _get_condition_json_property_match( self, diff --git a/tests/unit/memory/memory_interface/test_interface_prompts.py b/tests/unit/memory/memory_interface/test_interface_prompts.py index ed445e114b..d1e26245a5 100644 --- a/tests/unit/memory/memory_interface/test_interface_prompts.py +++ b/tests/unit/memory/memory_interface/test_interface_prompts.py @@ -1043,6 +1043,53 @@ async def test_get_message_pieces_captured_response_metadata(sqlite_instance: Me assert retrieved[0].prompt_metadata[key] == value +@pytest.mark.parametrize("key", ["attack.id", "a-b", "sp ace", "with:colon"]) +async def test_get_message_pieces_prompt_metadata_key_with_non_identifier_character( + sqlite_instance: MemoryInterface, key: str +): + """A free-form metadata key must be queryable whatever characters it holds. + + The key used to be spelled into the JSON path and reused as the bind-parameter + name, so a dot read as a path step and a hyphen was not a legal parameter name: + the query raised ``ArgumentError: This text() construct doesn't define a bound + parameter named ...`` instead of filtering. + """ + matching = MessagePiece( + conversation_id=str(uuid4()), + role="user", + original_value="targeted", + prompt_metadata={key: "value"}, + ) + other = MessagePiece( + conversation_id=str(uuid4()), + role="user", + original_value="unrelated", + prompt_metadata={key: "something_else"}, + ) + sqlite_instance._insert_entries(entries=[PromptMemoryEntry(entry=matching), PromptMemoryEntry(entry=other)]) + + retrieved = await sqlite_instance.get_message_pieces_async(prompt_metadata={key: "value"}) + + assert len(retrieved) == 1 + assert retrieved[0].original_value == "targeted" + + +async def test_get_message_pieces_prompt_metadata_multiple_non_identifier_keys(sqlite_instance: MemoryInterface): + """Several free-form keys in one query each get their own bound path and value.""" + metadata = {"attack.id": "a", "a-b": "b", "plain": "c"} + piece = MessagePiece( + conversation_id=str(uuid4()), + role="user", + original_value="hit", + prompt_metadata=metadata, + ) + sqlite_instance._insert_entries(entries=[PromptMemoryEntry(entry=piece)]) + + assert len(await sqlite_instance.get_message_pieces_async(prompt_metadata=metadata)) == 1 + assert len(await sqlite_instance.get_message_pieces_async(prompt_metadata={"attack.id": "a", "a-b": "b"})) == 1 + assert len(await sqlite_instance.get_message_pieces_async(prompt_metadata={"attack.id": "a", "a-b": "nope"})) == 0 + + async def test_get_message_pieces_id(sqlite_instance: MemoryInterface): entries = [ PromptMemoryEntry( From 2efea6409f434925a7d8717eecfe445eca3e69ca Mon Sep 17 00:00:00 2001 From: Chen Yufeiyang Date: Tue, 13 Oct 2026 07:09:57 +0800 Subject: [PATCH 2/4] Match literal metadata keys when querying and removing seeds --- pyrit/memory/azure_sql_memory.py | 20 ++- pyrit/memory/sqlite_memory.py | 61 ++++----- .../test_interface_prompts.py | 117 ++++++++++++++++++ 3 files changed, 153 insertions(+), 45 deletions(-) diff --git a/pyrit/memory/azure_sql_memory.py b/pyrit/memory/azure_sql_memory.py index 89ff0d4d31..986200cd4d 100644 --- a/pyrit/memory/azure_sql_memory.py +++ b/pyrit/memory/azure_sql_memory.py @@ -373,12 +373,8 @@ 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. - - Metadata keys are free-form, so each key is quoted into the JSON path - and the path is bound rather than spelled out: a key holding a dot is a - literal key, not a path step, and a key holding a hyphen is neither a - JSON path nor a legal bind-parameter name. + Uses OPENJSON() to compare free-form metadata keys literally rather than + interpreting them as JSON path expressions. Args: prompt_metadata (dict[str, str | int]): Dictionary of metadata key-value pairs to filter by. @@ -389,11 +385,13 @@ def _get_metadata_conditions(self, *, prompt_metadata: dict[str, str | int]) -> json_validation = "ISJSON(prompt_metadata) = 1" json_conditions, bindparams = [], {} for index, (key, value) in enumerate(prompt_metadata.items()): - path_param, value_param = f"md_path_{index}", f"md_value_{index}" - json_conditions.append(f"JSON_VALUE(prompt_metadata, :{path_param}) = :{value_param}") - bindparams[path_param] = f'$."{key}"' - # 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) + 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} AND [{alias}].[value] = :{value_param})" + ) + bindparams[key_param] = key bindparams[value_param] = str(value) # Combine both conditions diff --git a/pyrit/memory/sqlite_memory.py b/pyrit/memory/sqlite_memory.py index 3be619087b..4b0438500b 100644 --- a/pyrit/memory/sqlite_memory.py +++ b/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,53 +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. - Metadata keys are free-form, so each key is quoted into the JSON path - and the path is bound rather than spelled out: a key holding a dot is a - literal key, not a path step, and a key holding a hyphen is neither a - JSON path nor a legal bind-parameter name. + Match literal object keys through json_each(), without interpreting keys as JSON paths. Returns: list: A list of SQLAlchemy conditions. """ - json_paths, bindparams = [], {} - for index, (key, value) in enumerate(prompt_metadata.items()): - path_param, value_param = f"pm_path_{index}", f"pm_value_{index}" - json_paths.append(f"JSON_EXTRACT(prompt_metadata, :{path_param}) = :{value_param}") - bindparams[path_param] = f'$."{key}"' - # Note: We do NOT convert values to string here, to allow integer comparison in JSON - bindparams[value_param] = value - - # Create SQL condition using SQLAlchemy's text() with bindparams - condition = text(" AND ".join(json_paths)).bindparams(**bindparams) - 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. - Metadata keys are free-form, so each key is quoted into the JSON path - and the path is bound rather than spelled out, for the same reason as - in _get_message_pieces_prompt_metadata_conditions. + Match literal object keys through json_each(), without interpreting keys as JSON paths. Returns: - Any: A SQLAlchemy text condition with bound parameters. - """ - json_paths, bindparams = [], {} - for index, (key, value) in enumerate(metadata.items()): - path_param, value_param = f"sm_path_{index}", f"sm_value_{index}" - json_paths.append(f"JSON_EXTRACT(prompt_metadata, :{path_param}) = :{value_param}") - bindparams[path_param] = f'$."{key}"' - # Note: We do NOT convert values to string here, to allow integer comparison in JSON - bindparams[value_param] = value - - # Create SQL condition using SQLAlchemy's text() with bindparams - return text(" AND ".join(json_paths)).bindparams(**bindparams) + 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, diff --git a/tests/unit/memory/memory_interface/test_interface_prompts.py b/tests/unit/memory/memory_interface/test_interface_prompts.py index d1e26245a5..968da3d93e 100644 --- a/tests/unit/memory/memory_interface/test_interface_prompts.py +++ b/tests/unit/memory/memory_interface/test_interface_prompts.py @@ -1090,6 +1090,123 @@ async def test_get_message_pieces_prompt_metadata_multiple_non_identifier_keys(s assert len(await sqlite_instance.get_message_pieces_async(prompt_metadata={"attack.id": "a", "a-b": "nope"})) == 0 +@pytest.mark.parametrize( + ("literal_key", "other_metadata"), + [('branch"."leaf', {"branch": {"leaf": "value"}}), (r"\u0061", {"a": "value"})], +) +async def test_get_message_pieces_prompt_metadata_matches_literal_json_keys( + sqlite_instance: MemoryInterface, literal_key: str, other_metadata: dict[str, object] +): + matching = MessagePiece( + conversation_id=str(uuid4()), + role="user", + original_value="literal", + prompt_metadata={literal_key: "value"}, + ) + other = MessagePiece( + conversation_id=str(uuid4()), + role="user", + original_value="other", + prompt_metadata=other_metadata, + ) + sqlite_instance._insert_entries(entries=[PromptMemoryEntry(entry=matching), PromptMemoryEntry(entry=other)]) + + retrieved = await sqlite_instance.get_message_pieces_async(prompt_metadata={literal_key: "value"}) + + assert [piece.original_value for piece in retrieved] == ["literal"] + + +@pytest.mark.parametrize( + ("literal_key", "other_metadata"), + [('branch"."leaf', {"branch": {"leaf": "value"}}), (r"\u0061", {"a": "value"})], +) +async def test_seed_metadata_retrieval_matches_literal_json_keys( + sqlite_instance: MemoryInterface, literal_key: str, other_metadata: dict[str, object] +): + literal_group_id, nested_group_id = uuid4(), uuid4() + seeds = [ + SeedPrompt( + value="literal first", + data_type="text", + metadata={literal_key: "value"}, + prompt_group_id=literal_group_id, + ), + SeedPrompt( + value="literal second", + data_type="text", + metadata={literal_key: "value"}, + prompt_group_id=literal_group_id, + ), + SeedPrompt( + value="nested", + data_type="text", + metadata=other_metadata, + prompt_group_id=nested_group_id, + ), + ] + await sqlite_instance.add_seeds_to_memory_async(seeds=seeds, added_by="tester") + + retrieved = await sqlite_instance.get_seeds_async(metadata={literal_key: "value"}) + assert {seed.value for seed in retrieved} == {"literal first", "literal second"} + + +@pytest.mark.parametrize( + ("literal_key", "other_metadata"), + [('branch"."leaf', {"branch": {"leaf": "value"}}), (r"\u0061", {"a": "value"})], +) +async def test_seed_removal_matches_literal_json_keys( + sqlite_instance: MemoryInterface, literal_key: str, other_metadata: dict[str, object] +): + seeds = [ + SeedPrompt(value="literal", data_type="text", metadata={literal_key: "value"}), + SeedPrompt(value="nested", data_type="text", metadata=other_metadata), + ] + await sqlite_instance.add_seeds_to_memory_async(seeds=seeds, added_by="tester") + + removed = await sqlite_instance.remove_seeds_from_memory_async(metadata={literal_key: "value"}) + + remaining = await sqlite_instance.get_seeds_async() + assert [seed.value for seed in remaining] == ["nested"] + assert removed == 1 + + +@pytest.mark.parametrize( + ("literal_key", "other_metadata"), + [('branch"."leaf', {"branch": {"leaf": "value"}}), (r"\u0061", {"a": "value"})], +) +async def test_seed_group_removal_matches_literal_json_keys( + sqlite_instance: MemoryInterface, literal_key: str, other_metadata: dict[str, object] +): + literal_group_id, nested_group_id = uuid4(), uuid4() + seeds = [ + SeedPrompt( + value="literal first", + data_type="text", + metadata={literal_key: "value"}, + prompt_group_id=literal_group_id, + ), + SeedPrompt( + value="literal second", + data_type="text", + metadata={literal_key: "value"}, + prompt_group_id=literal_group_id, + ), + SeedPrompt( + value="nested", + data_type="text", + metadata=other_metadata, + prompt_group_id=nested_group_id, + ), + ] + await sqlite_instance.add_seeds_to_memory_async(seeds=seeds, added_by="tester") + + removed = await sqlite_instance.remove_seed_groups_from_memory_async(metadata={literal_key: "value"}) + + remaining = await sqlite_instance.get_seeds_async() + assert [seed.value for seed in remaining] == ["nested"] + assert removed == 2 + + async def test_get_message_pieces_id(sqlite_instance: MemoryInterface): entries = [ PromptMemoryEntry( From b9b669f668e974f86a6faf2ec435523acc42dae5 Mon Sep 17 00:00:00 2001 From: Chen Yufeiyang Date: Tue, 13 Oct 2026 14:41:41 +0800 Subject: [PATCH 3/4] FIX: preserve exact Azure metadata key lengths --- pyrit/memory/azure_sql_memory.py | 4 +- .../test_interface_prompts.py | 44 +++++++++++++++++-- tests/unit/memory/test_azure_sql_memory.py | 25 +++++++++++ 3 files changed, 68 insertions(+), 5 deletions(-) diff --git a/pyrit/memory/azure_sql_memory.py b/pyrit/memory/azure_sql_memory.py index 986200cd4d..3ccb74a301 100644 --- a/pyrit/memory/azure_sql_memory.py +++ b/pyrit/memory/azure_sql_memory.py @@ -389,7 +389,9 @@ def _get_metadata_conditions(self, *, prompt_metadata: dict[str, str | int]) -> alias = f"metadata_{index}" json_conditions.append( f"EXISTS (SELECT 1 FROM OPENJSON(prompt_metadata) AS [{alias}] " - f"WHERE [{alias}].[key] = :{key_param} AND [{alias}].[value] = :{value_param})" + f"WHERE [{alias}].[key] = :{key_param} " + f"AND DATALENGTH([{alias}].[key]) = DATALENGTH(CAST(:{key_param} AS NVARCHAR(MAX))) " + f"AND [{alias}].[value] = :{value_param})" ) bindparams[key_param] = key bindparams[value_param] = str(value) diff --git a/tests/unit/memory/memory_interface/test_interface_prompts.py b/tests/unit/memory/memory_interface/test_interface_prompts.py index 968da3d93e..f10f4fbdf3 100644 --- a/tests/unit/memory/memory_interface/test_interface_prompts.py +++ b/tests/unit/memory/memory_interface/test_interface_prompts.py @@ -1092,7 +1092,16 @@ async def test_get_message_pieces_prompt_metadata_multiple_non_identifier_keys(s @pytest.mark.parametrize( ("literal_key", "other_metadata"), - [('branch"."leaf', {"branch": {"leaf": "value"}}), (r"\u0061", {"a": "value"})], + [ + ('branch"."leaf', {"branch": {"leaf": "value"}}), + (r"\u0061", {"a": "value"}), + ("plain", {"plain ": "value"}), + ("plain ", {"plain": "value"}), + ("", {" ": "value"}), + (" ", {"": "value"}), + ("café", {"café ": "value"}), + ("café ", {"café": "value"}), + ], ) async def test_get_message_pieces_prompt_metadata_matches_literal_json_keys( sqlite_instance: MemoryInterface, literal_key: str, other_metadata: dict[str, object] @@ -1118,7 +1127,16 @@ async def test_get_message_pieces_prompt_metadata_matches_literal_json_keys( @pytest.mark.parametrize( ("literal_key", "other_metadata"), - [('branch"."leaf', {"branch": {"leaf": "value"}}), (r"\u0061", {"a": "value"})], + [ + ('branch"."leaf', {"branch": {"leaf": "value"}}), + (r"\u0061", {"a": "value"}), + ("plain", {"plain ": "value"}), + ("plain ", {"plain": "value"}), + ("", {" ": "value"}), + (" ", {"": "value"}), + ("café", {"café ": "value"}), + ("café ", {"café": "value"}), + ], ) async def test_seed_metadata_retrieval_matches_literal_json_keys( sqlite_instance: MemoryInterface, literal_key: str, other_metadata: dict[str, object] @@ -1152,7 +1170,16 @@ async def test_seed_metadata_retrieval_matches_literal_json_keys( @pytest.mark.parametrize( ("literal_key", "other_metadata"), - [('branch"."leaf', {"branch": {"leaf": "value"}}), (r"\u0061", {"a": "value"})], + [ + ('branch"."leaf', {"branch": {"leaf": "value"}}), + (r"\u0061", {"a": "value"}), + ("plain", {"plain ": "value"}), + ("plain ", {"plain": "value"}), + ("", {" ": "value"}), + (" ", {"": "value"}), + ("café", {"café ": "value"}), + ("café ", {"café": "value"}), + ], ) async def test_seed_removal_matches_literal_json_keys( sqlite_instance: MemoryInterface, literal_key: str, other_metadata: dict[str, object] @@ -1172,7 +1199,16 @@ async def test_seed_removal_matches_literal_json_keys( @pytest.mark.parametrize( ("literal_key", "other_metadata"), - [('branch"."leaf', {"branch": {"leaf": "value"}}), (r"\u0061", {"a": "value"})], + [ + ('branch"."leaf', {"branch": {"leaf": "value"}}), + (r"\u0061", {"a": "value"}), + ("plain", {"plain ": "value"}), + ("plain ", {"plain": "value"}), + ("", {" ": "value"}), + (" ", {"": "value"}), + ("café", {"café ": "value"}), + ("café ", {"café": "value"}), + ], ) async def test_seed_group_removal_matches_literal_json_keys( sqlite_instance: MemoryInterface, literal_key: str, other_metadata: dict[str, object] diff --git a/tests/unit/memory/test_azure_sql_memory.py b/tests/unit/memory/test_azure_sql_memory.py index 74ce80ef8d..51f220b9cc 100644 --- a/tests/unit/memory/test_azure_sql_memory.py +++ b/tests/unit/memory/test_azure_sql_memory.py @@ -121,6 +121,31 @@ async def test_insert_embedding_entry_async(memory_interface: AzureSQLMemory): assert persisted_embedding_entry.embedding_type_name == "test_type" +@pytest.mark.parametrize( + "prompt_metadata", + [ + {"plain": "value", "plain ": "padded value"}, + {"": "empty key", " ": "space key"}, + {"café": "unicode key", "café ": "padded unicode key"}, + ], +) +def test_metadata_conditions_guard_literal_key_lengths( + uninitialized_memory_interface: AzureSQLMemory, prompt_metadata: dict[str, str] +): + condition = uninitialized_memory_interface._get_metadata_conditions(prompt_metadata=prompt_metadata)[0] + compiled = condition.compile(dialect=mssql.dialect()) + sql = str(compiled) + + assert sql.count("DATALENGTH(") == 2 * len(prompt_metadata) + for index, (key, value) in enumerate(prompt_metadata.items()): + alias = f"metadata_{index}" + assert f"[{alias}].[key] = :md_key_{index}" in sql + assert f"DATALENGTH([{alias}].[key]) = DATALENGTH(CAST(:md_key_{index} AS NVARCHAR(MAX)))" in sql + assert f"[{alias}].[value] = :md_value_{index}" in sql + assert compiled.params[f"md_key_{index}"] == key + assert compiled.params[f"md_value_{index}"] == value + + def test_disable_embedding(memory_interface: AzureSQLMemory): memory_interface.disable_embedding() From ad72838bd1d8706e85d85226aaa6b0a5085e0d20 Mon Sep 17 00:00:00 2001 From: Roman Lutz Date: Tue, 13 Oct 2026 20:39:35 -0700 Subject: [PATCH 4/4] fix(memory): preserve scalar Azure SQL metadata filtering Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- pyrit/memory/azure_sql_memory.py | 4 +- .../test_azure_sql_memory_integration.py | 64 +++++++++++++++++++ tests/unit/memory/test_azure_sql_memory.py | 60 +++++++++++++++++ 3 files changed, 127 insertions(+), 1 deletion(-) diff --git a/pyrit/memory/azure_sql_memory.py b/pyrit/memory/azure_sql_memory.py index 3ccb74a301..245263a0f4 100644 --- a/pyrit/memory/azure_sql_memory.py +++ b/pyrit/memory/azure_sql_memory.py @@ -374,7 +374,8 @@ def _get_metadata_conditions(self, *, prompt_metadata: dict[str, str | int]) -> Generate SQL conditions for filtering by prompt metadata. Uses OPENJSON() to compare free-form metadata keys literally rather than - interpreting them as JSON path expressions. + 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. @@ -391,6 +392,7 @@ def _get_metadata_conditions(self, *, prompt_metadata: dict[str, str | int]) -> 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 diff --git a/tests/integration/memory/test_azure_sql_memory_integration.py b/tests/integration/memory/test_azure_sql_memory_integration.py index 13389657ea..a652bdc01c 100644 --- a/tests/integration/memory/test_azure_sql_memory_integration.py +++ b/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. diff --git a/tests/unit/memory/test_azure_sql_memory.py b/tests/unit/memory/test_azure_sql_memory.py index 51f220b9cc..117a47b000 100644 --- a/tests/unit/memory/test_azure_sql_memory.py +++ b/tests/unit/memory/test_azure_sql_memory.py @@ -11,6 +11,7 @@ import pytest from sqlalchemy import inspect, or_, select, text from sqlalchemy.dialects import mssql +from sqlalchemy.orm import Query, Session from pyrit.common.singleton import Singleton from pyrit.converter.base64_converter import Base64Converter @@ -146,6 +147,65 @@ def test_metadata_conditions_guard_literal_key_lengths( assert compiled.params[f"md_value_{index}"] == value +@pytest.mark.parametrize( + "metadata", + [ + {"payload": "{}"}, + {"payload": "[]"}, + {"payload": "", "count": 7, "enabled": "true"}, + {'branch"."leaf': "{}", r"\u0061": "[]"}, + ], +) +def test_metadata_conditions_only_match_scalar_json_types( + *, uninitialized_memory_interface: AzureSQLMemory, metadata: dict[str, str | int] +) -> None: + conditions = [ + *uninitialized_memory_interface._get_message_pieces_prompt_metadata_conditions(prompt_metadata=metadata), + uninitialized_memory_interface._get_seed_metadata_conditions(metadata=metadata), + ] + + for condition in conditions: + compiled = condition.compile(dialect=mssql.dialect(paramstyle="qmark")) + sql = str(compiled) + for index, (key, value) in enumerate(metadata.items()): + assert f"[metadata_{index}].[type] IN (1, 2, 3)" in sql + assert compiled.params[f"md_key_{index}"] == key + assert compiled.params[f"md_value_{index}"] == str(value) + assert compiled.positiontup == [ + param + for index in range(len(metadata)) + for param in (f"md_key_{index}", f"md_key_{index}", f"md_value_{index}") + ] + + +@pytest.mark.parametrize("value", ["{}", "[]"]) +@pytest.mark.parametrize("remove_groups", [False, True]) +def test_remove_seeds_metadata_only_matches_scalar_json_types( + *, uninitialized_memory_interface: AzureSQLMemory, value: str, remove_groups: bool +) -> None: + session = MagicMock(spec=Session) + query = MagicMock(spec=Query) + session.query.return_value = query + query.filter.return_value = query + query.delete.return_value = 1 + + with patch.object(uninitialized_memory_interface, "_get_session", return_value=session): + if remove_groups: + removed = uninitialized_memory_interface._execute_remove_seed_groups_from_memory( + metadata={"payload": value} + ) + else: + removed = uninitialized_memory_interface._execute_remove_seeds_from_memory(metadata={"payload": value}) + + assert removed == 1 + condition = query.filter.call_args.args[0] + compiled = condition.compile(dialect=mssql.dialect()) + assert "[metadata_0].[type] IN (1, 2, 3)" in str(compiled) + assert compiled.params["md_key_0"] == "payload" + assert compiled.params["md_value_0"] == value + session.commit.assert_called_once() + + def test_disable_embedding(memory_interface: AzureSQLMemory): memory_interface.disable_embedding()