@@ -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+
11651243async def test_get_attack_results_by_labels_multiple (sqlite_instance : MemoryInterface ):
11661244 """Test filtering attack results by multiple labels (AND logic)."""
11671245
0 commit comments