Prettier red team output via tqdm.write (Azure#41347)

pamelafox · web-flow · commit 1adab86d16af · 2025-05-30T13:00:18.000-07:00
diff --git a/sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/_evaluate/_utils.py b/sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/_evaluate/_utils.py
@@ -13,6 +13,8 @@
 import math
 
 import pandas as pd
+from tqdm import tqdm
+
 from azure.ai.evaluation._legacy._adapters.entities import Run
 
 from azure.ai.evaluation._constants import (
@@ -321,7 +323,8 @@ def _write_output(path: Union[str, os.PathLike], data_dict: Any) -> None:
     with open(p, "w", encoding=DefaultOpenEncoding.WRITE) as f:
         json.dump(data_dict, f, ensure_ascii=False)
 
-    print(f'Evaluation results saved to "{p.resolve()}".\n')
+    # Use tqdm.write to print message without interfering with any current progress bar
+    tqdm.write(f'Evaluation results saved to "{p.resolve()}".\n')
 
 
 def _apply_column_mapping(
diff --git a/sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/red_team/_red_team.py b/sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/red_team/_red_team.py
@@ -974,12 +974,12 @@ async def send_batch_with_retry():
                         
                         # Print progress to console 
                         if batch_idx < len(batches) - 1:  # Don't print for the last batch
-                            print(f"Strategy {strategy_name}, Risk {risk_category_name}: Processed batch {batch_idx+1}/{len(batches)}")
+                            tqdm.write(f"Strategy {strategy_name}, Risk {risk_category_name}: Processed batch {batch_idx+1}/{len(batches)}")
                             
                     except (asyncio.TimeoutError, tenacity.RetryError):
                         self.logger.warning(f"Batch {batch_idx+1} for {strategy_name}/{risk_category_name} timed out after {timeout} seconds, continuing with partial results")
                         self.logger.debug(f"Timeout: Strategy {strategy_name}, Risk {risk_category_name}, Batch {batch_idx+1} after {timeout} seconds.", exc_info=True)
-                        print(f"⚠️ TIMEOUT: Strategy {strategy_name}, Risk {risk_category_name}, Batch {batch_idx+1}")
+                        tqdm.write(f"⚠️ TIMEOUT: Strategy {strategy_name}, Risk {risk_category_name}, Batch {batch_idx+1}")
                         # Set task status to TIMEOUT
                         batch_task_key = f"{strategy_name}_{risk_category_name}_batch_{batch_idx+1}"
                         self.task_statuses[batch_task_key] = TASK_STATUS["TIMEOUT"]
@@ -1020,7 +1020,7 @@ async def send_all_with_retry():
                     self.logger.debug(f"Successfully processed single batch for {strategy_name}/{risk_category_name} in {batch_duration:.2f} seconds")
                 except (asyncio.TimeoutError, tenacity.RetryError):
                     self.logger.warning(f"Prompt processing for {strategy_name}/{risk_category_name} timed out after {timeout} seconds, continuing with partial results")
-                    print(f"⚠️ TIMEOUT: Strategy {strategy_name}, Risk {risk_category_name}")
+                    tqdm.write(f"⚠️ TIMEOUT: Strategy {strategy_name}, Risk {risk_category_name}")
                     # Set task status to TIMEOUT
                     single_batch_task_key = f"{strategy_name}_{risk_category_name}_single_batch"
                     self.task_statuses[single_batch_task_key] = TASK_STATUS["TIMEOUT"]
@@ -2101,7 +2101,7 @@ async def _process_attack(
         
         try:
             start_time = time.time()
-            print(f"▶️ Starting task: {strategy_name} strategy for {risk_category.value} risk category")
+            tqdm.write(f"▶️ Starting task: {strategy_name} strategy for {risk_category.value} risk category")
             log_strategy_start(self.logger, strategy_name, risk_category.value)
             
             converter = self._get_converter_for_strategy(strategy)
@@ -2137,7 +2137,7 @@ async def _process_attack(
                 )
             except Exception as e:
                 log_error(self.logger, f"Error during evaluation for {strategy_name}/{risk_category.value}", e)
-                print(f"⚠️ Evaluation error for {strategy_name}/{risk_category.value}: {str(e)}")
+                tqdm.write(f"⚠️ Evaluation error for {strategy_name}/{risk_category.value}: {str(e)}")
                 self.red_team_info[strategy_name][risk_category.value]["status"] = TASK_STATUS["FAILED"]
                 # Continue processing even if evaluation fails
             
@@ -2156,12 +2156,10 @@ async def _process_attack(
                     
                     # Print task completion message and estimated time on separate lines
                     # This ensures they don't get concatenated with tqdm output
-                    print("")  # Empty line to separate from progress bar
-                    print(f"✅ Completed task {self.completed_tasks}/{self.total_tasks} ({completion_pct:.1f}%) - {strategy_name}/{risk_category.value} in {elapsed_time:.1f}s")
-                    print(f"   Est. remaining: {est_remaining_time/60:.1f} minutes")
+                    tqdm.write(f"✅ Completed task {self.completed_tasks}/{self.total_tasks} ({completion_pct:.1f}%) - {strategy_name}/{risk_category.value} in {elapsed_time:.1f}s")
+                    tqdm.write(f"   Est. remaining: {est_remaining_time/60:.1f} minutes")
                 else:
-                    print("")  # Empty line to separate from progress bar
-                    print(f"✅ Completed task {self.completed_tasks}/{self.total_tasks} ({completion_pct:.1f}%) - {strategy_name}/{risk_category.value} in {elapsed_time:.1f}s")
+                    tqdm.write(f"✅ Completed task {self.completed_tasks}/{self.total_tasks} ({completion_pct:.1f}%) - {strategy_name}/{risk_category.value} in {elapsed_time:.1f}s")
                 
             log_strategy_completion(self.logger, strategy_name, risk_category.value, elapsed_time)
             self.task_statuses[task_key] = TASK_STATUS["COMPLETED"]
@@ -2279,8 +2277,8 @@ def filter(self, record):
         self.logger.debug(f"Timeout: {timeout} seconds")
         
         # Clear, minimal output for start of scan
-        print(f"🚀 STARTING RED TEAM SCAN: {scan_name}")
-        print(f"📂 Output directory: {self.scan_output_dir}")
+        tqdm.write(f"🚀 STARTING RED TEAM SCAN: {scan_name}")
+        tqdm.write(f"📂 Output directory: {self.scan_output_dir}")
         self.logger.info(f"Starting RED TEAM SCAN: {scan_name}")
         self.logger.info(f"Output directory: {self.scan_output_dir}")
         
@@ -2307,7 +2305,7 @@ def filter(self, record):
             
         self.risk_categories = self.attack_objective_generator.risk_categories
         # Show risk categories to user
-        print(f"📊 Risk categories: {[rc.value for rc in self.risk_categories]}")
+        tqdm.write(f"📊 Risk categories: {[rc.value for rc in self.risk_categories]}")
         self.logger.info(f"Risk categories to process: {[rc.value for rc in self.risk_categories]}")
         
         # Prepend AttackStrategy.Baseline to the attack strategy list
@@ -2329,11 +2327,11 @@ def filter(self, record):
                     
                 if strategy == AttackStrategy.Jailbreak:
                     self.logger.warning("Jailbreak strategy with custom attack objectives may not work as expected. The strategy will be run, but results may vary.")
-                    print("⚠️ Warning: Jailbreak strategy with custom attack objectives may not work as expected.")
+                    tqdm.write("⚠️ Warning: Jailbreak strategy with custom attack objectives may not work as expected.")
                     
                 if strategy == AttackStrategy.Tense:
                     self.logger.warning("Tense strategy requires specific formatting in objectives and may not work correctly with custom attack objectives.")
-                    print("⚠️ Warning: Tense strategy requires specific formatting in objectives and may not work correctly with custom attack objectives.")
+                    tqdm.write("⚠️ Warning: Tense strategy requires specific formatting in objectives and may not work correctly with custom attack objectives.")
                 
                 # Check for redundant converters 
                 # TODO: should this be in flattening logic?
@@ -2343,7 +2341,7 @@ def filter(self, record):
                     
                     if converter_type in used_converter_types and strategy != AttackStrategy.Baseline:
                         self.logger.warning(f"Strategy {strategy.name} uses a converter type that has already been used. Skipping redundant strategy.")
-                        print(f"ℹ️ Skipping redundant strategy: {strategy.name} (uses same converter as another strategy)")
+                        tqdm.write(f"ℹ️ Skipping redundant strategy: {strategy.name} (uses same converter as another strategy)")
                         strategies_to_remove.append(strategy)
                     else:
                         used_converter_types.add(converter_type)
@@ -2360,7 +2358,7 @@ def filter(self, record):
             eval_run = self._start_redteam_mlflow_run(self.azure_ai_project, scan_name)
 
             # Show URL for tracking progress
-            print(f"🔗 Track your red team scan in AI Foundry: {self.ai_studio_url}")
+            tqdm.write(f"🔗 Track your red team scan in AI Foundry: {self.ai_studio_url}")
             self.logger.info(f"Started Uploading run: {self.ai_studio_url}")
         
         log_subsection_header(self.logger, "Setting up scan configuration")
@@ -2376,7 +2374,7 @@ def filter(self, record):
         # Calculate total tasks: #risk_categories * #converters
         self.total_tasks = len(self.risk_categories) * len(flattened_attack_strategies) 
         # Show task count for user awareness
-        print(f"📋 Planning {self.total_tasks} total tasks")
+        tqdm.write(f"📋 Planning {self.total_tasks} total tasks")
         self.logger.info(f"Total tasks: {self.total_tasks} ({len(self.risk_categories)} risk categories * {len(flattened_attack_strategies)} strategies)")
         
         # Initialize our tracking dictionary early with empty structures
@@ -2412,10 +2410,10 @@ def filter(self, record):
         # Log the objective source mode
         if using_custom_objectives:
             self.logger.info(f"Using custom attack objectives from {self.attack_objective_generator.custom_attack_seed_prompts}")
-            print(f"📚 Using custom attack objectives from {self.attack_objective_generator.custom_attack_seed_prompts}")
+            tqdm.write(f"📚 Using custom attack objectives from {self.attack_objective_generator.custom_attack_seed_prompts}")
         else:
             self.logger.info("Using attack objectives from Azure RAI service")
-            print("📚 Using attack objectives from Azure RAI service")
+            tqdm.write("📚 Using attack objectives from Azure RAI service")
         
         # Dictionary to store all objectives
         all_objectives = {}
@@ -2434,7 +2432,7 @@ def filter(self, record):
             if "baseline" not in all_objectives:
                 all_objectives["baseline"] = {}
             all_objectives["baseline"][risk_category.value] = baseline_objectives
-            print(f"📝 Fetched baseline objectives for {risk_category.value}: {len(baseline_objectives)} objectives")
+            tqdm.write(f"📝 Fetched baseline objectives for {risk_category.value}: {len(baseline_objectives)} objectives")
         
         # Then fetch objectives for other strategies
         self.logger.info("Fetching objectives for non-baseline strategies")
@@ -2444,7 +2442,7 @@ def filter(self, record):
             if strategy_name == "baseline":
                 continue  # Already fetched
                 
-            print(f"🔄 Fetching objectives for strategy {i+1}/{strategy_count}: {strategy_name}")
+            tqdm.write(f"🔄 Fetching objectives for strategy {i+1}/{strategy_count}: {strategy_name}")
             all_objectives[strategy_name] = {}
             
             for risk_category in self.risk_categories:
@@ -2471,7 +2469,7 @@ def filter(self, record):
             
             if not objectives:
                 self.logger.warning(f"No objectives found for {strategy_name}+{risk_category.value}, skipping")
-                print(f"⚠️ No objectives found for {strategy_name}/{risk_category.value}, skipping")
+                tqdm.write(f"⚠️ No objectives found for {strategy_name}/{risk_category.value}, skipping")
                 self.red_team_info[strategy_name][risk_category.value]["status"] = TASK_STATUS["COMPLETED"]
                 async with progress_bar_lock:
                     progress_bar.update(1)
@@ -2496,7 +2494,7 @@ def filter(self, record):
             
         # Process tasks in parallel with optimized batching
         if parallel_execution and orchestrator_tasks:
-            print(f"⚙️ Processing {len(orchestrator_tasks)} tasks in parallel (max {max_parallel_tasks} at a time)")
+            tqdm.write(f"⚙️ Processing {len(orchestrator_tasks)} tasks in parallel (max {max_parallel_tasks} at a time)")
             self.logger.info(f"Processing {len(orchestrator_tasks)} tasks in parallel (max {max_parallel_tasks} at a time)")
             
             # Create batches for processing
@@ -2514,7 +2512,7 @@ def filter(self, record):
                     )
                 except asyncio.TimeoutError:
                     self.logger.warning(f"Batch {i//max_parallel_tasks+1} timed out after {timeout*2} seconds")
-                    print(f"⚠️ Batch {i//max_parallel_tasks+1} timed out, continuing with next batch")
+                    tqdm.write(f"⚠️ Batch {i//max_parallel_tasks+1} timed out, continuing with next batch")
                     # Set task status to TIMEOUT
                     batch_task_key = f"scan_batch_{i//max_parallel_tasks+1}"
                     self.task_statuses[batch_task_key] = TASK_STATUS["TIMEOUT"]
@@ -2526,7 +2524,7 @@ def filter(self, record):
         else:
             # Sequential execution 
             self.logger.info("Running orchestrator processing sequentially")
-            print("⚙️ Processing tasks sequentially")
+            tqdm.write("⚙️ Processing tasks sequentially")
             for i, task in enumerate(orchestrator_tasks):
                 progress_bar.set_postfix({"current": f"task {i+1}/{len(orchestrator_tasks)}"})
                 self.logger.debug(f"Processing task {i+1}/{len(orchestrator_tasks)}")
@@ -2536,7 +2534,7 @@ def filter(self, record):
                     await asyncio.wait_for(task, timeout=timeout)
                 except asyncio.TimeoutError:
                     self.logger.warning(f"Task {i+1}/{len(orchestrator_tasks)} timed out after {timeout} seconds")
-                    print(f"⚠️ Task {i+1} timed out, continuing with next task")
+                    tqdm.write(f"⚠️ Task {i+1} timed out, continuing with next task")
                     # Set task status to TIMEOUT
                     task_key = f"scan_task_{i+1}"
                     self.task_statuses[task_key] = TASK_STATUS["TIMEOUT"]
@@ -2607,18 +2605,18 @@ def filter(self, record):
             self.scorecard = scorecard
             
             # Print scorecard to console for user visibility (without extra header)
-            print(scorecard)
+            tqdm.write(scorecard)
             
             # Print URL for detailed results (once only)
             studio_url = output.scan_result.get("studio_url", "")
             if studio_url:
-                print(f"\nDetailed results available at:\n{studio_url}")
+                tqdm.write(f"\nDetailed results available at:\n{studio_url}")
             
             # Print the output directory path so the user can find it easily
             if hasattr(self, 'scan_output_dir') and self.scan_output_dir:
-                print(f"\n📂 All scan files saved to: {self.scan_output_dir}")
+                tqdm.write(f"\n📂 All scan files saved to: {self.scan_output_dir}")
         
-        print(f"✅ Scan completed successfully!")
+        tqdm.write(f"✅ Scan completed successfully!")
         self.logger.info("Scan completed successfully")
         for handler in self.logger.handlers:
             if isinstance(handler, logging.FileHandler):
