Skip to content

Commit dc8510a

Browse files
bIackr0seromanlutzCopilot
authored
FIX: match label keys containing dots or hyphens in attack result and message filters (#3063)
Co-authored-by: Roman Lutz <romanlutz13@gmail.com> Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
1 parent 620091f commit dc8510a

4 files changed

Lines changed: 130 additions & 19 deletions

File tree

‎pyrit/memory/azure_sql_memory.py‎

Lines changed: 10 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -348,9 +348,11 @@ def _get_message_pieces_memory_label_conditions(self, *, memory_labels: dict[str
348348
are_label_parts: list[str] = []
349349
are_bindparams: dict[str, str] = {}
350350

351-
for key, value in memory_labels.items():
352-
are_param = f"are_ml_{key}"
353-
are_label_parts.append(f"JSON_VALUE(\"AttackResultEntries\".labels, '$.{key}') = :{are_param}")
351+
for key_index, (key, value) in enumerate(memory_labels.items()):
352+
path_param = f"are_ml_path_{key_index}"
353+
are_param = f"are_ml_{key_index}"
354+
are_label_parts.append(f'JSON_VALUE("AttackResultEntries".labels, :{path_param}) = :{are_param}')
355+
are_bindparams[path_param] = f'$."{key}"'
354356
are_bindparams[are_param] = str(value)
355357

356358
combined_are = " AND ".join(are_label_parts)
@@ -545,17 +547,19 @@ def _get_attack_result_label_condition(self, *, labels: dict[str, str | Sequence
545547
are_label_conditions: list[str] = []
546548
are_bindparams: dict[str, str] = {}
547549

548-
for key, raw_value in labels.items():
550+
for key_index, (key, raw_value) in enumerate(labels.items()):
549551
values = [raw_value] if isinstance(raw_value, str) else list(raw_value)
550552
if not values:
551553
continue
554+
path_param = f"are_label_path_{key_index}"
555+
are_bindparams[path_param] = f'$."{key}"'
552556
are_placeholders = []
553557
for idx, v in enumerate(values):
554-
are_param = f"are_label_{key}_{idx}"
558+
are_param = f"are_label_{key_index}_{idx}"
555559
are_placeholders.append(f":{are_param}")
556560
are_bindparams[are_param] = str(v)
557561
are_in = ", ".join(are_placeholders)
558-
are_label_conditions.append(f"JSON_VALUE(\"AttackResultEntries\".labels, '$.{key}') IN ({are_in})")
562+
are_label_conditions.append(f'JSON_VALUE("AttackResultEntries".labels, :{path_param}) IN ({are_in})')
559563

560564
are_parts: list[Any] = [AttackResultEntry.labels.isnot(None)]
561565
if are_label_conditions:

‎pyrit/memory/sqlite_memory.py‎

Lines changed: 11 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -368,18 +368,23 @@ def _create_engine(self, *, has_echo: bool) -> Engine:
368368
def _get_message_pieces_memory_label_conditions(self, *, memory_labels: dict[str, str]) -> list[Any]:
369369
"""
370370
Generate SQLAlchemy filter conditions for filtering conversation pieces by memory labels.
371-
For SQLite, we use JSON_EXTRACT function to handle JSON fields.
371+
Match literal object keys through json_each(), without interpreting keys as JSON paths.
372372
373-
Matches if labels are on the PromptMemoryEntry itself OR on any
374-
AttackResultEntry that shares the same conversation_id.
373+
All labels must match one AttackResultEntry that shares the same conversation_id.
375374
376375
Returns:
377376
list: A list of SQLAlchemy conditions.
378377
"""
379378
per_key_are_conditions = []
379+
label_entries = func.json_each(AttackResultEntry.labels).table_valued("key", "value")
380380
for key, value in memory_labels.items():
381-
are_col = func.json_extract(AttackResultEntry.labels, f"$.{key}")
382-
per_key_are_conditions.append(are_col == str(value))
381+
per_key_are_conditions.append(
382+
select(1)
383+
.select_from(label_entries)
384+
.where(label_entries.c.key == key, label_entries.c.value == str(value))
385+
.correlate(AttackResultEntry)
386+
.exists()
387+
)
383388
return [
384389
exists().where(
385390
and_(
@@ -626,7 +631,7 @@ def _get_attack_result_label_condition(self, *, labels: dict[str, str | Sequence
626631
values = [raw_value] if isinstance(raw_value, str) else list(raw_value)
627632
if not values:
628633
continue
629-
are_col = func.json_extract(AttackResultEntry.labels, f"$.{key}")
634+
are_col = func.json_extract(AttackResultEntry.labels, f'$."{key}"')
630635
per_key_are_conditions.append(are_col.in_(values))
631636

632637
return and_(

‎tests/unit/memory/memory_interface/test_interface_attack_results.py‎

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1162,6 +1162,84 @@ async def test_get_attack_results_rejects_invalid_label_keys(sqlite_instance: Me
11621162
(await sqlite_instance.get_attack_results_async(labels={bad_key: "value"}))
11631163

11641164

1165+
@pytest.mark.parametrize("key", ["team.name", "run-id"])
1166+
async def test_get_attack_results_by_labels_key_with_dot_or_hyphen(sqlite_instance: MemoryInterface, key: str):
1167+
"""A label key the allowlist accepts is matched as one key, not as a nested JSON path."""
1168+
await sqlite_instance.add_attack_results_to_memory_async(
1169+
attack_results=[
1170+
create_attack_result("conv_1", 1, labels={key: "safety"}),
1171+
create_attack_result("conv_2", 2, labels={"team": "safety"}),
1172+
]
1173+
)
1174+
1175+
results = await sqlite_instance.get_attack_results_async(labels={key: "safety"})
1176+
assert [r.conversation_id for r in results] == ["conv_1"]
1177+
1178+
results = await sqlite_instance.get_attack_results_async(labels={key: ["other", "safety"]})
1179+
assert [r.conversation_id for r in results] == ["conv_1"]
1180+
1181+
1182+
@pytest.mark.parametrize(
1183+
"key, value",
1184+
[
1185+
("team.name", "safety"),
1186+
("run-id", "safety"),
1187+
('team"name', "safety"),
1188+
(r"team\name", "safety"),
1189+
('team\\"name', "safety"),
1190+
("team[0]", "safety"),
1191+
("team name", "safety"),
1192+
("caf\u00e9", "safety"),
1193+
("", "safety"),
1194+
("team.name", ""),
1195+
('team"name', 'quoted"\\value'),
1196+
(r"team\name", "caf\u00e9"),
1197+
],
1198+
)
1199+
async def test_get_message_pieces_by_literal_label_key_async(
1200+
sqlite_instance: MemoryInterface, key: str, value: str
1201+
) -> None:
1202+
"""Message labels match complete string keys and values without JSON-path interpretation."""
1203+
for conversation_id in ("conv_1", "conv_2", "conv_3"):
1204+
await sqlite_instance.add_message_pieces_to_memory_async(
1205+
message_pieces=[MessagePiece(role="user", original_value="hello", conversation_id=conversation_id)]
1206+
)
1207+
await sqlite_instance.add_attack_results_to_memory_async(
1208+
attack_results=[
1209+
create_attack_result("conv_1", 1, labels={key: value}),
1210+
create_attack_result("conv_2", 2, labels={"team": value}),
1211+
create_attack_result("conv_3", 3, labels={key: f"{value}-other"}),
1212+
]
1213+
)
1214+
1215+
pieces = await sqlite_instance.get_message_pieces_async(labels={key: value})
1216+
assert [p.conversation_id for p in pieces] == ["conv_1"]
1217+
1218+
1219+
async def test_get_message_pieces_literal_labels_match_one_attack_without_duplicates_async(
1220+
sqlite_instance: MemoryInterface,
1221+
) -> None:
1222+
"""All labels must match one attack result, without multiplying its conversation's pieces."""
1223+
labels = {'team"name': "safety", r"run\id": "r1"}
1224+
for conversation_id in ("conv_1", "conv_2", "conv_3"):
1225+
await sqlite_instance.add_message_pieces_to_memory_async(
1226+
message_pieces=[MessagePiece(role="user", original_value="hello", conversation_id=conversation_id)]
1227+
)
1228+
await sqlite_instance.add_attack_results_to_memory_async(
1229+
attack_results=[
1230+
create_attack_result("conv_1", 1, labels=labels),
1231+
create_attack_result("conv_1", 2, labels=labels),
1232+
create_attack_result("conv_2", 3, labels={'team"name': "safety"}),
1233+
create_attack_result("conv_2", 4, labels={r"run\id": "r1"}),
1234+
create_attack_result("conv_3", 5, labels={**labels, r"run\id": "other"}),
1235+
]
1236+
)
1237+
1238+
pieces = await sqlite_instance.get_message_pieces_async(labels=labels, role="user")
1239+
assert [p.conversation_id for p in pieces] == ["conv_1"]
1240+
assert await sqlite_instance.get_message_pieces_async(labels=labels, conversation_id="conv_2") == []
1241+
1242+
11651243
async def test_get_attack_results_by_labels_multiple(sqlite_instance: MemoryInterface):
11661244
"""Test filtering attack results by multiple labels (AND logic)."""
11671245

‎tests/unit/memory/test_azure_sql_memory.py‎

Lines changed: 31 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -308,7 +308,31 @@ def test_get_message_pieces_memory_label_conditions_bind_params(uninitialized_me
308308
memory_labels={"operation": "test_op"}
309309
)
310310
params = conditions[0].compile().params
311-
assert params == {"are_ml_operation": "test_op"}
311+
assert params == {"are_ml_path_0": '$."operation"', "are_ml_0": "test_op"}
312+
313+
314+
def test_label_conditions_bind_whole_keys_with_dot_or_hyphen(uninitialized_memory_interface: AzureSQLMemory):
315+
"""Allowlisted keys with ``.`` or ``-`` build a query that looks up each whole key."""
316+
attack_condition = uninitialized_memory_interface._get_attack_result_label_condition(
317+
labels={"team.name": ["red", "blue"], "run-id": "r1"}
318+
)
319+
assert attack_condition.compile().params == {
320+
"are_label_path_0": '$."team.name"',
321+
"are_label_0_0": "red",
322+
"are_label_0_1": "blue",
323+
"are_label_path_1": '$."run-id"',
324+
"are_label_1_0": "r1",
325+
}
326+
327+
piece_condition = uninitialized_memory_interface._get_message_pieces_memory_label_conditions(
328+
memory_labels={"team.name": "red", "run-id": "r1"}
329+
)[0]
330+
assert piece_condition.compile().params == {
331+
"are_ml_path_0": '$."team.name"',
332+
"are_ml_0": "red",
333+
"are_ml_path_1": '$."run-id"',
334+
"are_ml_1": "r1",
335+
}
312336

313337

314338
async def test_update_entries_async(memory_interface: AzureSQLMemory):
@@ -421,17 +445,18 @@ def test_get_attack_result_label_condition_with_string_value(memory_interface: A
421445
"""String values produce a single-placeholder IN clause with the stringified value."""
422446
condition = memory_interface._get_attack_result_label_condition(labels={"operator": "roakey"})
423447
params = condition.compile().params
424-
assert params == {"are_label_operator_0": "roakey"}
448+
assert params == {"are_label_path_0": '$."operator"', "are_label_0_0": "roakey"}
425449

426450

427451
def test_get_attack_result_label_condition_with_sequence_value(memory_interface: AzureSQLMemory):
428452
"""Sequence values produce one placeholder per element."""
429453
condition = memory_interface._get_attack_result_label_condition(labels={"operation": ["op_a", "op_b", "op_c"]})
430454
params = condition.compile().params
431455
assert params == {
432-
"are_label_operation_0": "op_a",
433-
"are_label_operation_1": "op_b",
434-
"are_label_operation_2": "op_c",
456+
"are_label_path_0": '$."operation"',
457+
"are_label_0_0": "op_a",
458+
"are_label_0_1": "op_b",
459+
"are_label_0_2": "op_c",
435460
}
436461

437462

@@ -440,8 +465,7 @@ def test_get_attack_result_label_condition_skips_empty_sequence(memory_interface
440465
condition = memory_interface._get_attack_result_label_condition(labels={"operator": "roakey", "operation": []})
441466
params = condition.compile().params
442467
# operator gets bind params; operation (empty) does not.
443-
assert params == {"are_label_operator_0": "roakey"}
444-
assert not any("label_operation_" in k for k in params)
468+
assert params == {"are_label_path_0": '$."operator"', "are_label_0_0": "roakey"}
445469

446470

447471
def test_get_attack_result_label_condition_empty_labels_dict(memory_interface: AzureSQLMemory):

0 commit comments

Comments
 (0)