Add context support to red_team attack objectives for enhanced evaluation (#42128)

* Initial plan

* Implement context support in red_team attack objectives and evaluation

Co-authored-by: slister1001 <[email protected]>

* Add documentation and examples for red_team context support

Co-authored-by: slister1001 <[email protected]>

* Apply black formatting to red_team files

Co-authored-by: slister1001 <[email protected]>

* Update sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/red_team/_red_team.py

Co-authored-by: Copilot <[email protected]>

* Update sdk/evaluation/azure-ai-evaluation/samples/red_team_samples.py

Co-authored-by: Copilot <[email protected]>

* Add ISA and XPIA risk categories with similar processing to ProtectedMaterial

Co-authored-by: slister1001 <[email protected]>

* Change RiskCategory.ISA to RiskCategory.UngroundedAttributes

Co-authored-by: slister1001 <[email protected]>

* updates

* updates

* small indentation fix

* run code formatter

* updates for callback

* run code formatter

* code formatter

* updates for multiturn and crescendo with context

* small fix for crescendo

* remove _internalriskcategory typehints

* fix code format

---------

Co-authored-by: copilot-swe-agent[bot] <[email protected]>
Co-authored-by: slister1001 <[email protected]>
Co-authored-by: Copilot <[email protected]>
Co-authored-by: Sydney Lister <[email protected]>

Author: 3 people

sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/_common/onedp/models/_models.py

@@ -1961,19 +1961,24 @@ class Message(_Model):
     :vartype role: str
     :ivar content: The content.
     :vartype content: str
+    :ivar context: The context.
+    :vartype context: str
     """
 
     role: Optional[str] = rest_field(name="Role", visibility=["read", "create", "update", "delete", "query"])
     """The role."""
     content: Optional[str] = rest_field(name="Content", visibility=["read", "create", "update", "delete", "query"])
     """The content."""
+    context: Optional[str] = rest_field(name="Context", visibility=["read", "create", "update", "delete", "query"])
+    """The context."""
 
     @overload
     def __init__(
         self,
         *,
         role: Optional[str] = None,
         content: Optional[str] = None,
+        context: Optional[str] = None,
     ) -> None: ...
 
     @overload

sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/red_team/_attack_objective_generator.py

@@ -20,6 +20,8 @@ class RiskCategory(str, Enum):
     SelfHarm = "self_harm"
     ProtectedMaterial = "protected_material"
     CodeVulnerability = "code_vulnerability"
+    UngroundedAttributes = "ungrounded_attributes"
+    XPIA = "xpia"
 
 
 @experimental

sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/red_team/_callback_chat_target.py

@@ -19,6 +19,7 @@ def __init__(
         *,
         callback: Callable[[List[Dict], bool, Optional[str], Optional[Dict[str, Any]]], Dict],
         stream: bool = False,
+        prompt_to_context: Optional[Dict[str, str]] = None,
     ) -> None:
         """
         Initializes an instance of the _CallbackChatTarget class.
@@ -32,10 +33,12 @@ def __init__(
         Args:
             callback (Callable): The callback function that sends a prompt to a target and receives a response.
             stream (bool, optional): Indicates whether the target supports streaming. Defaults to False.
+            prompt_to_context (Optional[Dict[str, str]], optional): Mapping from prompt content to context. Defaults to None.
         """
         PromptChatTarget.__init__(self)
         self._callback = callback
         self._stream = stream
+        self._prompt_to_context = prompt_to_context or {}
 
     async def send_prompt_async(self, *, prompt_request: PromptRequestResponse) -> PromptRequestResponse:
 
@@ -48,8 +51,18 @@ async def send_prompt_async(self, *, prompt_request: PromptRequestResponse) -> P
 
         logger.info(f"Sending the following prompt to the prompt target: {request}")
 
+        # Get context for the current prompt if available
+        current_prompt_content = request.converted_value
+        context_data = self._prompt_to_context.get(current_prompt_content, "")
+        context_dict = {"context": context_data} if context_data else {}
+
+        # If context is not available via prompt_to_context, it can be fetched from the memory
+        if not context_dict:
+            memory_label_context = request.labels.get("context", None)
+            context_dict = {"context": memory_label_context} if memory_label_context else {}
+
         # response_context contains "messages", "stream", "session_state, "context"
-        response_context = await self._callback(messages=messages, stream=self._stream, session_state=None, context=None)  # type: ignore
+        response_context = await self._callback(messages=messages, stream=self._stream, session_state=None, context=context_dict)  # type: ignore
 
         response_text = response_context["messages"][-1]["content"]
         response_entry = construct_response_from_request(request=request, response_text_pieces=[response_text])

sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/red_team/_evaluation_processor.py

@@ -32,7 +32,11 @@
 from ._attack_strategy import AttackStrategy
 from ._attack_objective_generator import RiskCategory
 from ._utils.constants import RESULTS_EXT, TASK_STATUS
-from ._utils.metric_mapping import get_annotation_task_from_risk_category, get_metric_from_risk_category
+from ._utils.metric_mapping import (
+    get_annotation_task_from_risk_category,
+    get_metric_from_risk_category,
+    get_attack_objective_from_risk_category,
+)
 from ._utils.logging_utils import log_error
 from ._utils.formatting_utils import get_strategy_name
 
@@ -97,12 +101,19 @@ async def evaluate_conversation(
         # Extract all assistant messages for evaluation
         assistant_messages = [msg["content"] for msg in messages if msg.get("role") == "assistant"]
 
+        context = [msg["context"] for msg in messages if msg.get("role") == "user"]
+
         if assistant_messages:
             # Create query-response pair with empty query and all assistant messages
             query_response = {
                 "query": "query",
                 "response": " ".join(assistant_messages),
             }
+
+            # Add context to query_response if found
+            if context[0] is not None:
+                query_response["context"] = context[0]
+
             try:
                 self.logger.debug(f"Evaluating conversation {idx+1} for {risk_category.value}/{strategy_name}")
 

sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/red_team/_orchestrator_manager.py

@@ -38,6 +38,7 @@
 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):
@@ -169,6 +170,7 @@ async def _prompt_sending_orchestrator(
         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.
 
@@ -224,9 +226,6 @@ async def _prompt_sending_orchestrator(
                     task_statuses[task_key] = TASK_STATUS["COMPLETED"]
                 return orchestrator
 
-            # Debug log the first few characters of each prompt
-            self.logger.debug(f"First prompt (truncated): {all_prompts[0][:50]}...")
-
             # Initialize output path for memory labelling
             base_path = str(uuid.uuid4())
 
@@ -313,6 +312,7 @@ async def _multi_turn_orchestrator(
         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).
 
@@ -381,6 +381,7 @@ async def _multi_turn_orchestrator(
         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,
@@ -390,6 +391,7 @@ async def _multi_turn_orchestrator(
                     credential=self.credential,
                     risk_category=risk_category,
                     azure_ai_project=self.azure_ai_project,
+                    context=context,
                 )
 
                 azure_rai_service_target = AzureRAIServiceTarget(
@@ -411,9 +413,6 @@ async def _multi_turn_orchestrator(
                     use_score_as_feedback=False,
                 )
 
-                # Debug log the first few characters of the current prompt
-                self.logger.debug(f"Current prompt (truncated): {prompt[:50]}...")
-
                 try:
                     # Create retry-enabled function using the reusable decorator
                     @network_retry_decorator(
@@ -423,10 +422,7 @@ 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,
-                                },
+                                memory_labels={"risk_strategy_path": output_path, "batch": 1, "context": context},
                             ),
                             timeout=calculated_timeout,
                         )
@@ -438,6 +434,13 @@ async def send_prompt_with_retry():
                         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(
@@ -492,6 +495,7 @@ async def _crescendo_orchestrator(
         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.
 
@@ -542,12 +546,14 @@ async def _crescendo_orchestrator(
         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(
@@ -577,11 +583,9 @@ async def _crescendo_orchestrator(
                     credential=self.credential,
                     risk_category=risk_category,
                     azure_ai_project=self.azure_ai_project,
+                    context=context,
                 )
 
-                # Debug log the first few characters of the current prompt
-                self.logger.debug(f"Current prompt (truncated): {prompt[:50]}...")
-
                 try:
                     # Create retry-enabled function using the reusable decorator
                     @network_retry_decorator(
@@ -594,6 +598,7 @@ async def send_prompt_with_retry():
                                 memory_labels={
                                     "risk_strategy_path": output_path,
                                     "batch": prompt_idx + 1,
+                                    "context": context,
                                 },
                             ),
                             timeout=calculated_timeout,
@@ -606,6 +611,13 @@ async def send_prompt_with_retry():
                         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(