Adding scan_Session_id

Author: Nagkumar Arkalgud

sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/_common/rai_service.py

@@ -259,7 +259,7 @@ async def submit_request(
 
 
 async def submit_request_onedp(
-    client: AIProjectClient, data: dict, metric: str, token: str, annotation_task: str, evaluator_name: str
+    client: AIProjectClient, data: dict, metric: str, token: str, annotation_task: str, evaluator_name: str, scan_session_id: Optional[str] = None
 ) -> str:
     """Submit request to Responsible AI service for evaluation and return operation ID
 
@@ -275,12 +275,16 @@ async def submit_request_onedp(
     :type annotation_task: str
     :param evaluator_name: The evaluator name.
     :type evaluator_name: str
+    :param scan_session_id: The scan session ID to use for the evaluation.
+    :type scan_session_id: Optional[str]
     :return: The operation ID.
     :rtype: str
     """
     normalized_user_text = get_formatted_template(data, annotation_task)
     payload = generate_payload(normalized_user_text, metric, annotation_task=annotation_task)
     headers = get_common_headers(token, evaluator_name)
+    if scan_session_id:
+        headers["client_request_id"] = scan_session_id
     response = client.evaluations.submit_annotation(payload, headers=headers)
     result = json.loads(response)
     operation_id = result["location"].split("/")[-1]
@@ -631,6 +635,7 @@ async def evaluate_with_rai_service(
     annotation_task: str = Tasks.CONTENT_HARM,
     metric_display_name=None,
     evaluator_name=None,
+    scan_session_id: Optional[str]=None
 ) -> Dict[str, Union[str, float]]:
     """Evaluate the content safety of the response using Responsible AI service
 
@@ -649,6 +654,8 @@ async def evaluate_with_rai_service(
     :type metric_display_name: str
     :param evaluator_name: The evaluator name to use.
     :type evaluator_name: str
+    :param scan_session_id: The scan session ID to use for the evaluation.
+    :type scan_session_id: Optional[str]
     :return: The parsed annotation result.
     :rtype: Dict[str, Union[str, float]]
     """
@@ -661,7 +668,7 @@ async def evaluate_with_rai_service(
         )
         token = await fetch_or_reuse_token(credential=credential, workspace=COG_SRV_WORKSPACE)
         await ensure_service_availability_onedp(client, token, annotation_task)
-        operation_id = await submit_request_onedp(client, data, metric_name, token, annotation_task, evaluator_name)
+        operation_id = await submit_request_onedp(client, data, metric_name, token, annotation_task, evaluator_name, scan_session_id)
         annotation_response = cast(List[Dict], await fetch_result_onedp(client, operation_id, token))
         result = parse_response(annotation_response, metric_name, metric_display_name)
         return result

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

@@ -292,6 +292,7 @@ def __init__(
         self.failed_tasks = 0
         self.start_time = None
         self.scan_id = None
+        self.scan_session_id = None
         self.scan_output_dir = None
 
         self.generated_rai_client = GeneratedRAIClient(azure_ai_project=self.azure_ai_project, token_manager=self.token_manager.credential)  # type: ignore
@@ -759,13 +760,15 @@ async def get_jailbreak_prefixes_with_retry():
                         risk_category=other_risk,
                         application_scenario=application_scenario or "",
                         strategy="tense",
+                        scan_session_id=self.scan_session_id
                     )
                 else:
                     objectives_response = await self.generated_rai_client.get_attack_objectives(
                         risk_type=content_harm_risk,
                         risk_category=other_risk,
                         application_scenario=application_scenario or "",
                         strategy=None,
+                        scan_session_id=self.scan_session_id
                     )
                 if isinstance(objectives_response, list):
                     self.logger.debug(f"API returned {len(objectives_response)} objectives")
@@ -775,7 +778,7 @@ async def get_jailbreak_prefixes_with_retry():
                 # Handle jailbreak strategy - need to apply jailbreak prefixes to messages
                 if strategy == "jailbreak":
                     self.logger.debug("Applying jailbreak prefixes to objectives")
-                    jailbreak_prefixes = await self.generated_rai_client.get_jailbreak_prefixes()
+                    jailbreak_prefixes = await self.generated_rai_client.get_jailbreak_prefixes(scan_session_id=self.scan_session_id)
                     for objective in objectives_response:
                         if "messages" in objective and len(objective["messages"]) > 0:
                             message = objective["messages"][0]
@@ -2350,6 +2353,7 @@ async def evaluate_with_rai_service_with_retry():
                             project_scope=self.azure_ai_project,
                             credential=self.credential,
                             annotation_task=annotation_task,
+                            scan_session_id=self.scan_session_id,
                         )
                     except (
                         httpx.ConnectTimeout,
@@ -2737,6 +2741,8 @@ async def scan(
         )
         self.scan_id = self.scan_id.replace(" ", "_")
 
+        self.scan_session_id = str(uuid.uuid4())  # Unique session ID for this scan
+
         # Create output directory for this scan
         # If DEBUG environment variable is set, use a regular folder name; otherwise, use a hidden folder
         is_debug = os.environ.get("DEBUG", "").lower() in ("true", "1", "yes", "y")

sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/simulator/_model_tools/_generated_rai_client.py

@@ -102,6 +102,7 @@ async def get_attack_objectives(
         risk_category: Optional[str] = None,
         application_scenario: str = None,
         strategy: Optional[str] = None,
+        scan_session_id: Optional[str] = None,
     ) -> Dict:
         """Get attack objectives using the auto-generated operations.
 
@@ -113,6 +114,8 @@ async def get_attack_objectives(
         :type application_scenario: str
         :param strategy: Optional strategy to filter the attack objectives
         :type strategy: Optional[str]
+        :param scan_session_id: Optional unique session ID for the scan
+        :type scan_session_id: Optional[str]
         :return: The attack objectives
         :rtype: Dict
         """
@@ -123,6 +126,7 @@ async def get_attack_objectives(
                 risk_categories=[risk_category],
                 lang="en",
                 strategy=strategy,
+                headers={"client_request_id": scan_session_id},
             )
             return response
 
@@ -133,15 +137,19 @@ async def get_attack_objectives(
             logging.error(f"Error in get_attack_objectives: {str(e)}")
             raise
 
-    async def get_jailbreak_prefixes(self) -> List[str]:
+    async def get_jailbreak_prefixes(self, scan_session_id: Optional[str] = None) -> List[str]:
         """Get jailbreak prefixes using the auto-generated operations.
 
+        :param scan_session_id: Optional unique session ID for the scan
+        :type scan_session_id: Optional[str]
         :return: The jailbreak prefixes
         :rtype: List[str]
         """
         try:
             # Send the request using the autogenerated client
-            response = self._client.get_jail_break_dataset_with_type(type="upia")
+            response = self._client.get_jail_break_dataset_with_type(
+                type="upia", headers={"client_request_id": scan_session_id}
+            )
             if isinstance(response, list):
                 return response
             else: