azure-ai-evaluation 1.9.0__py3-none-any.whl → 1.11.0__py3-none-any.whl

azure/ai/evaluation/red_team/_orchestrator_manager.py

@@ -0,0 +1,661 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+"""
+Orchestrator management module for Red Team Agent.
+
+This module handles PyRIT orchestrator initialization, execution, and management
+for different attack strategies including single-turn, multi-turn, and crescendo attacks.
+"""
+
+import asyncio
+import math
+import os
+import uuid
+from datetime import datetime
+from typing import Dict, List, Optional, Union, Callable
+from tqdm import tqdm
+
+# PyRIT imports
+from pyrit.orchestrator.single_turn.prompt_sending_orchestrator import PromptSendingOrchestrator
+from pyrit.orchestrator.multi_turn.red_teaming_orchestrator import RedTeamingOrchestrator
+from pyrit.orchestrator.multi_turn.crescendo_orchestrator import CrescendoOrchestrator
+from pyrit.orchestrator import Orchestrator
+from pyrit.prompt_converter import PromptConverter
+from pyrit.prompt_target import PromptChatTarget
+
+# Retry imports
+import httpx
+import httpcore
+import tenacity
+from tenacity import retry
+
+# Local imports
+from ._attack_strategy import AttackStrategy
+from ._attack_objective_generator import RiskCategory
+from ._utils._rai_service_target import AzureRAIServiceTarget
+from ._utils._rai_service_true_false_scorer import AzureRAIServiceTrueFalseScorer
+from ._utils._rai_service_eval_chat_target import RAIServiceEvalChatTarget
+from ._utils.constants import DATA_EXT, TASK_STATUS
+from ._utils.logging_utils import log_strategy_start, log_error
+from ._utils.formatting_utils import write_pyrit_outputs_to_file
+
+
+def network_retry_decorator(retry_config, logger, strategy_name, risk_category_name, prompt_idx=None):
+    """Create a reusable retry decorator for network operations.
+
+    :param retry_config: Retry configuration dictionary
+    :param logger: Logger instance for logging warnings
+    :param strategy_name: Name of the attack strategy
+    :param risk_category_name: Name of the risk category
+    :param prompt_idx: Optional prompt index for detailed logging
+    :return: Configured retry decorator
+    """
+
+    def decorator(func):
+        @retry(**retry_config["network_retry"])
+        async def wrapper(*args, **kwargs):
+            try:
+                return await func(*args, **kwargs)
+            except (
+                httpx.ConnectTimeout,
+                httpx.ReadTimeout,
+                httpx.ConnectError,
+                httpx.HTTPError,
+                ConnectionError,
+                TimeoutError,
+                OSError,
+                asyncio.TimeoutError,
+                httpcore.ReadTimeout,
+                httpx.HTTPStatusError,
+            ) as e:
+                prompt_detail = f" for prompt {prompt_idx}" if prompt_idx is not None else ""
+                logger.warning(
+                    f"Network error{prompt_detail} for {strategy_name}/{risk_category_name}: {type(e).__name__}: {str(e)}"
+                )
+                await asyncio.sleep(2)
+                raise
+
+        return wrapper
+
+    return decorator
+
+
+class OrchestratorManager:
+    """Manages PyRIT orchestrators for different attack strategies."""
+
+    def __init__(
+        self,
+        logger,
+        generated_rai_client,
+        credential,
+        azure_ai_project,
+        one_dp_project,
+        retry_config,
+        scan_output_dir=None,
+    ):
+        """Initialize the orchestrator manager.
+
+        :param logger: Logger instance for logging
+        :param generated_rai_client: RAI client for service interactions
+        :param credential: Authentication credential
+        :param azure_ai_project: Azure AI project configuration
+        :param one_dp_project: Whether this is a OneDP project
+        :param retry_config: Retry configuration for network errors
+        :param scan_output_dir: Directory for scan outputs
+        """
+        self.logger = logger
+        self.generated_rai_client = generated_rai_client
+        self.credential = credential
+        self.azure_ai_project = azure_ai_project
+        self._one_dp_project = one_dp_project
+        self.retry_config = retry_config
+        self.scan_output_dir = scan_output_dir
+
+    def _calculate_timeout(self, base_timeout: int, orchestrator_type: str) -> int:
+        """Calculate appropriate timeout based on orchestrator type.
+
+        Multi-turn and crescendo orchestrators need more generous timeouts due to their
+        iterative nature and multiple API calls per prompt.
+
+        :param base_timeout: Base timeout value in seconds
+        :param orchestrator_type: Type of orchestrator ('single', 'multi_turn', 'crescendo')
+        :return: Calculated timeout in seconds
+        """
+        timeout_multipliers = {
+            "single": 1.0,  # Standard timeout for single-turn
+            "multi_turn": 3.0,  # 3x timeout for multi-turn interactions
+            "crescendo": 4.0,  # 4x timeout for crescendo with backtracks
+        }
+
+        multiplier = timeout_multipliers.get(orchestrator_type, 1.0)
+        calculated_timeout = int(base_timeout * multiplier)
+
+        self.logger.debug(
+            f"Calculated timeout for {orchestrator_type} orchestrator: {calculated_timeout}s "
+            f"(base: {base_timeout}s, multiplier: {multiplier}x)"
+        )
+
+        return calculated_timeout
+
+    def get_orchestrator_for_attack_strategy(
+        self, attack_strategy: Union[AttackStrategy, List[AttackStrategy]]
+    ) -> Callable:
+        """Get appropriate orchestrator function for the specified attack strategy.
+
+        :param attack_strategy: Attack strategy to get orchestrator for
+        :type attack_strategy: Union[AttackStrategy, List[AttackStrategy]]
+        :return: Callable orchestrator function
+        :rtype: Callable
+        """
+        if isinstance(attack_strategy, list):
+            if AttackStrategy.MultiTurn in attack_strategy or AttackStrategy.Crescendo in attack_strategy:
+                self.logger.error("MultiTurn and Crescendo strategies are not supported in composed attacks.")
+                raise ValueError("MultiTurn and Crescendo strategies are not supported in composed attacks.")
+        elif AttackStrategy.MultiTurn == attack_strategy:
+            return self._multi_turn_orchestrator
+        elif AttackStrategy.Crescendo == attack_strategy:
+            return self._crescendo_orchestrator
+        return self._prompt_sending_orchestrator
+
+    async def _prompt_sending_orchestrator(
+        self,
+        chat_target: PromptChatTarget,
+        all_prompts: List[str],
+        converter: Union[PromptConverter, List[PromptConverter]],
+        *,
+        strategy_name: str = "unknown",
+        risk_category_name: str = "unknown",
+        risk_category: Optional[RiskCategory] = None,
+        timeout: int = 120,
+        red_team_info: Dict = None,
+        task_statuses: Dict = None,
+        prompt_to_context: Dict[str, str] = None,
+    ) -> Orchestrator:
+        """Send prompts via the PromptSendingOrchestrator.
+
+        :param chat_target: The target to send prompts to
+        :type chat_target: PromptChatTarget
+        :param all_prompts: List of prompts to process and send
+        :type all_prompts: List[str]
+        :param converter: Prompt converter or list of converters to transform prompts
+        :type converter: Union[PromptConverter, List[PromptConverter]]
+        :param strategy_name: Name of the attack strategy being used
+        :type strategy_name: str
+        :param risk_category_name: Name of the risk category being evaluated
+        :type risk_category_name: str
+        :param risk_category: Risk category being evaluated
+        :type risk_category: Optional[RiskCategory]
+        :param timeout: Timeout in seconds for each prompt
+        :type timeout: int
+        :param red_team_info: Dictionary to store file paths and results
+        :type red_team_info: Dict
+        :param task_statuses: Dictionary to track task statuses
+        :type task_statuses: Dict
+        :return: Configured and initialized orchestrator
+        :rtype: Orchestrator
+        """
+        task_key = f"{strategy_name}_{risk_category_name}_orchestrator"
+        if task_statuses:
+            task_statuses[task_key] = TASK_STATUS["RUNNING"]
+
+        log_strategy_start(self.logger, strategy_name, risk_category_name)
+
+        # Create converter list from single converter or list of converters
+        converter_list = (
+            [converter] if converter and isinstance(converter, PromptConverter) else converter if converter else []
+        )
+
+        # Log which converter is being used
+        if converter_list:
+            if isinstance(converter_list, list) and len(converter_list) > 0:
+                converter_names = [c.__class__.__name__ for c in converter_list if c is not None]
+                self.logger.debug(f"Using converters: {', '.join(converter_names)}")
+            elif converter is not None:
+                self.logger.debug(f"Using converter: {converter.__class__.__name__}")
+        else:
+            self.logger.debug("No converters specified")
+
+        # Initialize orchestrator
+        try:
+            orchestrator = PromptSendingOrchestrator(objective_target=chat_target, prompt_converters=converter_list)
+
+            if not all_prompts:
+                self.logger.warning(f"No prompts provided to orchestrator for {strategy_name}/{risk_category_name}")
+                if task_statuses:
+                    task_statuses[task_key] = TASK_STATUS["COMPLETED"]
+                return orchestrator
+
+            # Initialize output path for memory labelling
+            base_path = str(uuid.uuid4())
+
+            # If scan output directory exists, place the file there
+            if self.scan_output_dir:
+                output_path = os.path.join(self.scan_output_dir, f"{base_path}{DATA_EXT}")
+            else:
+                output_path = f"{base_path}{DATA_EXT}"
+
+            if red_team_info:
+                red_team_info[strategy_name][risk_category_name]["data_file"] = output_path
+
+            # Process all prompts at once
+            self.logger.debug(f"Processing {len(all_prompts)} prompts for {strategy_name}/{risk_category_name}")
+            start_time = datetime.now()
+
+            # Calculate appropriate timeout for single-turn orchestrator
+            calculated_timeout = self._calculate_timeout(timeout, "single")
+
+            try:
+                # Create retry-enabled function using the reusable decorator
+                @network_retry_decorator(self.retry_config, self.logger, strategy_name, risk_category_name)
+                async def send_all_with_retry():
+                    return await asyncio.wait_for(
+                        orchestrator.send_prompts_async(
+                            prompt_list=all_prompts,
+                            memory_labels={
+                                "risk_strategy_path": output_path,
+                                "batch": 1,
+                            },
+                        ),
+                        timeout=calculated_timeout,
+                    )
+
+                # Execute the retry-enabled function
+                await send_all_with_retry()
+                duration = (datetime.now() - start_time).total_seconds()
+                self.logger.debug(
+                    f"Successfully processed all prompts for {strategy_name}/{risk_category_name} in {duration:.2f} seconds"
+                )
+            except (asyncio.TimeoutError, tenacity.RetryError):
+                self.logger.warning(
+                    f"Prompt processing for {strategy_name}/{risk_category_name} timed out after {calculated_timeout} seconds, continuing with partial results"
+                )
+                print(f"⚠️ TIMEOUT: Strategy {strategy_name}, Risk {risk_category_name}")
+                if task_statuses:
+                    task_statuses[task_key] = TASK_STATUS["TIMEOUT"]
+                if red_team_info:
+                    red_team_info[strategy_name][risk_category_name]["status"] = TASK_STATUS["INCOMPLETE"]
+            except Exception as e:
+                log_error(
+                    self.logger,
+                    "Error processing prompts",
+                    e,
+                    f"{strategy_name}/{risk_category_name}",
+                )
+                if red_team_info:
+                    red_team_info[strategy_name][risk_category_name]["status"] = TASK_STATUS["INCOMPLETE"]
+
+            if task_statuses:
+                task_statuses[task_key] = TASK_STATUS["COMPLETED"]
+            return orchestrator
+
+        except Exception as e:
+            log_error(
+                self.logger,
+                "Failed to initialize orchestrator",
+                e,
+                f"{strategy_name}/{risk_category_name}",
+            )
+            if task_statuses:
+                task_statuses[task_key] = TASK_STATUS["FAILED"]
+            raise
+
+    async def _multi_turn_orchestrator(
+        self,
+        chat_target: PromptChatTarget,
+        all_prompts: List[str],
+        converter: Union[PromptConverter, List[PromptConverter]],
+        *,
+        strategy_name: str = "unknown",
+        risk_category_name: str = "unknown",
+        risk_category: Optional[RiskCategory] = None,
+        timeout: int = 120,
+        red_team_info: Dict = None,
+        task_statuses: Dict = None,
+        prompt_to_context: Dict[str, str] = None,
+    ) -> Orchestrator:
+        """Send prompts via the RedTeamingOrchestrator (multi-turn orchestrator).
+
+        :param chat_target: The target to send prompts to
+        :type chat_target: PromptChatTarget
+        :param all_prompts: List of prompts to process and send
+        :type all_prompts: List[str]
+        :param converter: Prompt converter or list of converters to transform prompts
+        :type converter: Union[PromptConverter, List[PromptConverter]]
+        :param strategy_name: Name of the attack strategy being used
+        :type strategy_name: str
+        :param risk_category_name: Name of the risk category being evaluated
+        :type risk_category_name: str
+        :param risk_category: Risk category being evaluated
+        :type risk_category: Optional[RiskCategory]
+        :param timeout: Timeout in seconds for each prompt
+        :type timeout: int
+        :param red_team_info: Dictionary to store file paths and results
+        :type red_team_info: Dict
+        :param task_statuses: Dictionary to track task statuses
+        :type task_statuses: Dict
+        :return: Configured and initialized orchestrator
+        :rtype: Orchestrator
+        """
+        max_turns = 5  # Set a default max turns value
+        task_key = f"{strategy_name}_{risk_category_name}_orchestrator"
+        if task_statuses:
+            task_statuses[task_key] = TASK_STATUS["RUNNING"]
+
+        log_strategy_start(self.logger, strategy_name, risk_category_name)
+        converter_list = []
+        # Create converter list from single converter or list of converters
+        if converter and isinstance(converter, PromptConverter):
+            converter_list = [converter]
+        elif converter and isinstance(converter, list):
+            # Filter out None values from the converter list
+            converter_list = [c for c in converter if c is not None]
+
+        # Log which converter is being used
+        if converter_list:
+            if isinstance(converter_list, list) and len(converter_list) > 0:
+                converter_names = [c.__class__.__name__ for c in converter_list if c is not None]
+                self.logger.debug(f"Using converters: {', '.join(converter_names)}")
+            elif converter is not None:
+                self.logger.debug(f"Using converter: {converter.__class__.__name__}")
+        else:
+            self.logger.debug("No converters specified")
+
+        # Initialize output path for memory labelling
+        base_path = str(uuid.uuid4())
+
+        # If scan output directory exists, place the file there
+        if self.scan_output_dir:
+            # Ensure the directory exists
+            os.makedirs(self.scan_output_dir, exist_ok=True)
+            output_path = os.path.join(self.scan_output_dir, f"{base_path}{DATA_EXT}")
+        else:
+            output_path = f"{base_path}{DATA_EXT}"
+
+        if red_team_info:
+            red_team_info[strategy_name][risk_category_name]["data_file"] = output_path
+
+        # Calculate appropriate timeout for multi-turn orchestrator
+        calculated_timeout = self._calculate_timeout(timeout, "multi_turn")
+
+        for prompt_idx, prompt in enumerate(all_prompts):
+            prompt_start_time = datetime.now()
+            self.logger.debug(f"Processing prompt {prompt_idx+1}/{len(all_prompts)}")
+            context = prompt_to_context.get(prompt, None) if prompt_to_context else None
+            try:
+                azure_rai_service_scorer = AzureRAIServiceTrueFalseScorer(
+                    client=self.generated_rai_client,
+                    api_version=None,
+                    model="gpt-4",
+                    logger=self.logger,
+                    credential=self.credential,
+                    risk_category=risk_category,
+                    azure_ai_project=self.azure_ai_project,
+                    context=context,
+                )
+
+                azure_rai_service_target = AzureRAIServiceTarget(
+                    client=self.generated_rai_client,
+                    api_version=None,
+                    model="gpt-4",
+                    prompt_template_key="orchestrators/red_teaming/text_generation.yaml",
+                    objective=prompt,
+                    logger=self.logger,
+                    is_one_dp_project=self._one_dp_project,
+                )
+
+                orchestrator = RedTeamingOrchestrator(
+                    objective_target=chat_target,
+                    adversarial_chat=azure_rai_service_target,
+                    max_turns=max_turns,
+                    prompt_converters=converter_list,
+                    objective_scorer=azure_rai_service_scorer,
+                    use_score_as_feedback=False,
+                )
+
+                try:
+                    # Create retry-enabled function using the reusable decorator
+                    @network_retry_decorator(
+                        self.retry_config, self.logger, strategy_name, risk_category_name, prompt_idx + 1
+                    )
+                    async def send_prompt_with_retry():
+                        return await asyncio.wait_for(
+                            orchestrator.run_attack_async(
+                                objective=prompt,
+                                memory_labels={"risk_strategy_path": output_path, "batch": 1, "context": context},
+                            ),
+                            timeout=calculated_timeout,
+                        )
+
+                    # Execute the retry-enabled function
+                    await send_prompt_with_retry()
+                    prompt_duration = (datetime.now() - prompt_start_time).total_seconds()
+                    self.logger.debug(
+                        f"Successfully processed prompt {prompt_idx+1} for {strategy_name}/{risk_category_name} in {prompt_duration:.2f} seconds"
+                    )
+
+                    # Write outputs to file after each prompt is processed
+                    write_pyrit_outputs_to_file(
+                        output_path=output_path,
+                        logger=self.logger,
+                        prompt_to_context=prompt_to_context,
+                    )
+
+                    # Print progress to console
+                    if prompt_idx < len(all_prompts) - 1:  # Don't print for the last prompt
+                        print(
+                            f"Strategy {strategy_name}, Risk {risk_category_name}: Processed prompt {prompt_idx+1}/{len(all_prompts)}"
+                        )
+
+                except (asyncio.TimeoutError, tenacity.RetryError):
+                    self.logger.warning(
+                        f"Batch {prompt_idx+1} for {strategy_name}/{risk_category_name} timed out after {calculated_timeout} seconds, continuing with partial results"
+                    )
+                    print(f"⚠️ TIMEOUT: Strategy {strategy_name}, Risk {risk_category_name}, Batch {prompt_idx+1}")
+                    # Set task status to TIMEOUT
+                    batch_task_key = f"{strategy_name}_{risk_category_name}_prompt_{prompt_idx+1}"
+                    if task_statuses:
+                        task_statuses[batch_task_key] = TASK_STATUS["TIMEOUT"]
+                    if red_team_info:
+                        red_team_info[strategy_name][risk_category_name]["status"] = TASK_STATUS["INCOMPLETE"]
+                    continue
+                except Exception as e:
+                    log_error(
+                        self.logger,
+                        f"Error processing prompt {prompt_idx+1}",
+                        e,
+                        f"{strategy_name}/{risk_category_name}",
+                    )
+                    if red_team_info:
+                        red_team_info[strategy_name][risk_category_name]["status"] = TASK_STATUS["INCOMPLETE"]
+                    continue
+            except Exception as e:
+                log_error(
+                    self.logger,
+                    "Failed to initialize orchestrator",
+                    e,
+                    f"{strategy_name}/{risk_category_name}",
+                )
+                if task_statuses:
+                    task_statuses[task_key] = TASK_STATUS["FAILED"]
+                raise
+        if task_statuses:
+            task_statuses[task_key] = TASK_STATUS["COMPLETED"]
+        return orchestrator
+
+    async def _crescendo_orchestrator(
+        self,
+        chat_target: PromptChatTarget,
+        all_prompts: List[str],
+        converter: Union[PromptConverter, List[PromptConverter]],
+        *,
+        strategy_name: str = "unknown",
+        risk_category_name: str = "unknown",
+        risk_category: Optional[RiskCategory] = None,
+        timeout: int = 120,
+        red_team_info: Dict = None,
+        task_statuses: Dict = None,
+        prompt_to_context: Dict[str, str] = None,
+    ) -> Orchestrator:
+        """Send prompts via the CrescendoOrchestrator with optimized performance.
+
+        :param chat_target: The target to send prompts to
+        :type chat_target: PromptChatTarget
+        :param all_prompts: List of prompts to process and send
+        :type all_prompts: List[str]
+        :param converter: Prompt converter or list of converters to transform prompts
+        :type converter: Union[PromptConverter, List[PromptConverter]]
+        :param strategy_name: Name of the attack strategy being used
+        :type strategy_name: str
+        :param risk_category_name: Name of the risk category being evaluated
+        :type risk_category_name: str
+        :param risk_category: Risk category being evaluated
+        :type risk_category: Optional[RiskCategory]
+        :param timeout: Timeout in seconds for each prompt
+        :type timeout: int
+        :param red_team_info: Dictionary to store file paths and results
+        :type red_team_info: Dict
+        :param task_statuses: Dictionary to track task statuses
+        :type task_statuses: Dict
+        :return: Configured and initialized orchestrator
+        :rtype: Orchestrator
+        """
+        max_turns = 10  # Set a default max turns value
+        max_backtracks = 5
+        task_key = f"{strategy_name}_{risk_category_name}_orchestrator"
+        if task_statuses:
+            task_statuses[task_key] = TASK_STATUS["RUNNING"]
+
+        log_strategy_start(self.logger, strategy_name, risk_category_name)
+
+        # Initialize output path for memory labelling
+        base_path = str(uuid.uuid4())
+
+        # If scan output directory exists, place the file there
+        if self.scan_output_dir:
+            output_path = os.path.join(self.scan_output_dir, f"{base_path}{DATA_EXT}")
+        else:
+            output_path = f"{base_path}{DATA_EXT}"
+
+        if red_team_info:
+            red_team_info[strategy_name][risk_category_name]["data_file"] = output_path
+
+        # Calculate appropriate timeout for crescendo orchestrator
+        calculated_timeout = self._calculate_timeout(timeout, "crescendo")
+
+        for prompt_idx, prompt in enumerate(all_prompts):
+            prompt_start_time = datetime.now()
+            self.logger.debug(f"Processing prompt {prompt_idx+1}/{len(all_prompts)}")
+            context = prompt_to_context.get(prompt, None) if prompt_to_context else None
+            try:
+                red_llm_scoring_target = RAIServiceEvalChatTarget(
+                    logger=self.logger,
+                    credential=self.credential,
+                    risk_category=risk_category,
+                    azure_ai_project=self.azure_ai_project,
+                    context=context,
+                )
+
+                azure_rai_service_target = AzureRAIServiceTarget(
+                    client=self.generated_rai_client,
+                    api_version=None,
+                    model="gpt-4",
+                    prompt_template_key="orchestrators/crescendo/crescendo_variant_1.yaml",
+                    objective=prompt,
+                    logger=self.logger,
+                    crescendo_format=True,
+                    is_one_dp_project=self._one_dp_project,
+                )
+
+                orchestrator = CrescendoOrchestrator(
+                    objective_target=chat_target,
+                    adversarial_chat=azure_rai_service_target,
+                    max_turns=max_turns,
+                    scoring_target=red_llm_scoring_target,
+                    max_backtracks=max_backtracks,
+                )
+
+                orchestrator._objective_scorer = AzureRAIServiceTrueFalseScorer(
+                    client=self.generated_rai_client,
+                    api_version=None,
+                    model="gpt-4",
+                    logger=self.logger,
+                    credential=self.credential,
+                    risk_category=risk_category,
+                    azure_ai_project=self.azure_ai_project,
+                    context=context,
+                )
+
+                try:
+                    # Create retry-enabled function using the reusable decorator
+                    @network_retry_decorator(
+                        self.retry_config, self.logger, strategy_name, risk_category_name, prompt_idx + 1
+                    )
+                    async def send_prompt_with_retry():
+                        return await asyncio.wait_for(
+                            orchestrator.run_attack_async(
+                                objective=prompt,
+                                memory_labels={
+                                    "risk_strategy_path": output_path,
+                                    "batch": prompt_idx + 1,
+                                    "context": context,
+                                },
+                            ),
+                            timeout=calculated_timeout,
+                        )
+
+                    # Execute the retry-enabled function
+                    await send_prompt_with_retry()
+                    prompt_duration = (datetime.now() - prompt_start_time).total_seconds()
+                    self.logger.debug(
+                        f"Successfully processed prompt {prompt_idx+1} for {strategy_name}/{risk_category_name} in {prompt_duration:.2f} seconds"
+                    )
+
+                    # Write outputs to file after each prompt is processed
+                    write_pyrit_outputs_to_file(
+                        output_path=output_path,
+                        logger=self.logger,
+                        prompt_to_context=prompt_to_context,
+                    )
+
+                    # Print progress to console
+                    if prompt_idx < len(all_prompts) - 1:  # Don't print for the last prompt
+                        print(
+                            f"Strategy {strategy_name}, Risk {risk_category_name}: Processed prompt {prompt_idx+1}/{len(all_prompts)}"
+                        )
+
+                except (asyncio.TimeoutError, tenacity.RetryError):
+                    self.logger.warning(
+                        f"Batch {prompt_idx+1} for {strategy_name}/{risk_category_name} timed out after {calculated_timeout} seconds, continuing with partial results"
+                    )
+                    print(f"⚠️ TIMEOUT: Strategy {strategy_name}, Risk {risk_category_name}, Batch {prompt_idx+1}")
+                    # Set task status to TIMEOUT
+                    batch_task_key = f"{strategy_name}_{risk_category_name}_prompt_{prompt_idx+1}"
+                    if task_statuses:
+                        task_statuses[batch_task_key] = TASK_STATUS["TIMEOUT"]
+                    if red_team_info:
+                        red_team_info[strategy_name][risk_category_name]["status"] = TASK_STATUS["INCOMPLETE"]
+                    continue
+                except Exception as e:
+                    log_error(
+                        self.logger,
+                        f"Error processing prompt {prompt_idx+1}",
+                        e,
+                        f"{strategy_name}/{risk_category_name}",
+                    )
+                    if red_team_info:
+                        red_team_info[strategy_name][risk_category_name]["status"] = TASK_STATUS["INCOMPLETE"]
+                    continue
+            except Exception as e:
+                log_error(
+                    self.logger,
+                    "Failed to initialize orchestrator",
+                    e,
+                    f"{strategy_name}/{risk_category_name}",
+                )
+                if task_statuses:
+                    task_statuses[task_key] = TASK_STATUS["FAILED"]
+                raise
+        if task_statuses:
+            task_statuses[task_key] = TASK_STATUS["COMPLETED"]
+        return orchestrator