@@ -1276,6 +1276,19 @@ async def _multi_turn_orchestrator(
1276
1276
else :
1277
1277
self .logger .debug ("No converters specified" )
1278
1278
1279
+ # Initialize output path for memory labelling
1280
+ base_path = str (uuid .uuid4 ())
1281
+
1282
+ # If scan output directory exists, place the file there
1283
+ if hasattr (self , "scan_output_dir" ) and self .scan_output_dir :
1284
+ # Ensure the directory exists
1285
+ os .makedirs (self .scan_output_dir , exist_ok = True )
1286
+ output_path = os .path .join (self .scan_output_dir , f"{ base_path } { DATA_EXT } " )
1287
+ else :
1288
+ output_path = f"{ base_path } { DATA_EXT } "
1289
+
1290
+ self .red_team_info [strategy_name ][risk_category_name ]["data_file" ] = output_path
1291
+
1279
1292
for prompt_idx , prompt in enumerate (all_prompts ):
1280
1293
prompt_start_time = datetime .now ()
1281
1294
self .logger .debug (f"Processing prompt { prompt_idx + 1 } /{ len (all_prompts )} " )
@@ -1314,17 +1327,6 @@ async def _multi_turn_orchestrator(
1314
1327
# Debug log the first few characters of the current prompt
1315
1328
self .logger .debug (f"Current prompt (truncated): { prompt [:50 ]} ..." )
1316
1329
1317
- # Initialize output path for memory labelling
1318
- base_path = str (uuid .uuid4 ())
1319
-
1320
- # If scan output directory exists, place the file there
1321
- if hasattr (self , "scan_output_dir" ) and self .scan_output_dir :
1322
- output_path = os .path .join (self .scan_output_dir , f"{ base_path } { DATA_EXT } " )
1323
- else :
1324
- output_path = f"{ base_path } { DATA_EXT } "
1325
-
1326
- self .red_team_info [strategy_name ][risk_category_name ]["data_file" ] = output_path
1327
-
1328
1330
try : # Create retry decorator for this specific call with enhanced retry strategy
1329
1331
1330
1332
@retry (** self ._create_retry_config ()["network_retry" ])
@@ -1361,6 +1363,12 @@ async def send_prompt_with_retry():
1361
1363
self .logger .debug (
1362
1364
f"Successfully processed prompt { prompt_idx + 1 } for { strategy_name } /{ risk_category_name } in { prompt_duration :.2f} seconds"
1363
1365
)
1366
+ self ._write_pyrit_outputs_to_file (
1367
+ orchestrator = orchestrator ,
1368
+ strategy_name = strategy_name ,
1369
+ risk_category = risk_category_name ,
1370
+ batch_idx = prompt_idx + 1 ,
1371
+ )
1364
1372
1365
1373
# Print progress to console
1366
1374
if prompt_idx < len (all_prompts ) - 1 : # Don't print for the last prompt
@@ -1385,7 +1393,7 @@ async def send_prompt_with_retry():
1385
1393
orchestrator = orchestrator ,
1386
1394
strategy_name = strategy_name ,
1387
1395
risk_category = risk_category_name ,
1388
- batch_idx = 1 ,
1396
+ batch_idx = prompt_idx + 1 ,
1389
1397
)
1390
1398
# Continue with partial results rather than failing completely
1391
1399
continue
@@ -1404,7 +1412,7 @@ async def send_prompt_with_retry():
1404
1412
orchestrator = orchestrator ,
1405
1413
strategy_name = strategy_name ,
1406
1414
risk_category = risk_category_name ,
1407
- batch_idx = 1 ,
1415
+ batch_idx = prompt_idx + 1 ,
1408
1416
)
1409
1417
# Continue with other batches even if one fails
1410
1418
continue
0 commit comments