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

azure/ai/evaluation/simulator/_adversarial_simulator.py

@@ -49,8 +49,8 @@ class AdversarialSimulator:
     """
     Initializes the adversarial simulator with a project scope.
 
-    :param azure_ai_project: The Azure AI project, which can either be a string representing the project endpoint 
-        or an instance of AzureAIProject. It contains subscription id, resource group, and project name. 
+    :param azure_ai_project: The Azure AI project, which can either be a string representing the project endpoint
+        or an instance of AzureAIProject. It contains subscription id, resource group, and project name.
     :type azure_ai_project: Union[str, AzureAIProject]
     :param credential: The credential for connecting to Azure AI project.
     :type credential: ~azure.core.credentials.TokenCredential
@@ -77,7 +77,7 @@ class AdversarialSimulator:
                 logger=logging.getLogger("AdversarialSimulator"),
                 credential=self.credential,
             )
-            self.rai_client  = AIProjectClient(endpoint=azure_ai_project, credential=credential)
+            self.rai_client = AIProjectClient(endpoint=azure_ai_project, credential=credential)
         else:
             try:
                 self.azure_ai_project = validate_azure_ai_project(azure_ai_project)
@@ -96,7 +96,7 @@ class AdversarialSimulator:
                 credential=self.credential,
             )
             self.rai_client = RAIClient(azure_ai_project=self.azure_ai_project, token_manager=self.token_manager)
-                            
+
         self.adversarial_template_handler = AdversarialTemplateHandler(
             azure_ai_project=self.azure_ai_project, rai_client=self.rai_client
         )
@@ -264,7 +264,9 @@ class AdversarialSimulator:
 
         # Limit to max_simulation_results if needed
         if len(template_parameter_pairs) > max_simulation_results:
-            template_parameter_pairs = template_parameter_pairs[:max_simulation_results]        # Create a seeded random instance for jailbreak selection if randomization_seed is provided
+            template_parameter_pairs = template_parameter_pairs[
+                :max_simulation_results
+            ]  # Create a seeded random instance for jailbreak selection if randomization_seed is provided
         jailbreak_random = None
         if _jailbreak_type == "upia" and randomization_seed is not None:
             jailbreak_random = random.Random(randomization_seed)
@@ -364,10 +366,21 @@ class AdversarialSimulator:
             target=target, role=ConversationRole.ASSISTANT, template=template, parameters=parameters, scenario=scenario
         )
         bots = [user_bot, system_bot]
-        
+
+        async def run_simulation(session_obj):
+            async with semaphore:
+                _, conversation_history = await simulate_conversation(
+                    bots=bots,
+                    session=session_obj,
+                    turn_limit=max_conversation_turns,
+                    api_call_delay_sec=api_call_delay_sec,
+                    language=language,
+                )
+            return conversation_history
+
         if isinstance(self.rai_client, AIProjectClient):
             session = self.rai_client
-        else:    
+        else:
             session = get_async_http_client().with_policies(
                 retry_policy=AsyncRetryPolicy(
                     retry_total=api_call_retry_limit,
@@ -375,13 +388,7 @@ class AdversarialSimulator:
                     retry_mode=RetryMode.Fixed,
                 )
             )
-        _, conversation_history = await simulate_conversation(
-            bots=bots,
-            session=session,
-            turn_limit=max_conversation_turns,
-            api_call_delay_sec=api_call_delay_sec,
-            language=language,
-        )
+        conversation_history = await run_simulation(session)
 
         return self._to_chat_protocol(
             conversation_history=conversation_history,
@@ -391,7 +398,11 @@ class AdversarialSimulator:
     def _get_user_proxy_completion_model(
         self, template_key: str, template_parameters: TemplateParameters, simulation_id: str = ""
     ) -> ProxyChatCompletionsModel:
-        endpoint_url = self.rai_client._config.endpoint + "/redTeams/simulation/chat/completions/submit" if isinstance(self.rai_client, AIProjectClient) else self.rai_client.simulation_submit_endpoint    
+        endpoint_url = (
+            self.rai_client._config.endpoint + "/redTeams/simulation/chat/completions/submit"
+            if isinstance(self.rai_client, AIProjectClient)
+            else self.rai_client.simulation_submit_endpoint
+        )
         return ProxyChatCompletionsModel(
             name="raisvc_proxy_model",
             template_key=template_key,

azure/ai/evaluation/simulator/_conversation/__init__.py

@@ -419,13 +419,13 @@ class MultiModalConversationBot(ConversationBot):
         contents = []
         for msg in messages:
             if msg.startswith("image_understanding/"):
-                if(isinstance(self.rai_client, RAIClient)):
+                if isinstance(self.rai_client, RAIClient):
                     encoded_image = await self.rai_client.get_image_data(msg)
                 else:
                     response = self.rai_client.red_teams.get_template_parameters_image(path=msg, stream="true")
                     image_data = b"".join(response)
                     encoded_image = base64.b64encode(image_data).decode("utf-8")
-                    
+
                 contents.append(
                     {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{encoded_image}"}},
                 )

azure/ai/evaluation/simulator/_direct_attack_simulator.py

@@ -28,8 +28,8 @@ class DirectAttackSimulator:
     Initialize a UPIA (user prompt injected attack) jailbreak adversarial simulator with a project scope.
     This simulator converses with your AI system using prompts designed to interrupt normal functionality.
 
-    :param azure_ai_project: The Azure AI project, which can either be a string representing the project endpoint 
-        or an instance of AzureAIProject. It contains subscription id, resource group, and project name. 
+    :param azure_ai_project: The Azure AI project, which can either be a string representing the project endpoint
+        or an instance of AzureAIProject. It contains subscription id, resource group, and project name.
     :type azure_ai_project: Union[str, AzureAIProject]
     :param credential: The credential for connecting to Azure AI project.
     :type credential: ~azure.core.credentials.TokenCredential
@@ -46,16 +46,16 @@ class DirectAttackSimulator:
 
     def __init__(self, *, azure_ai_project: Union[str, AzureAIProject], credential: TokenCredential):
         """Constructor."""
-        
+
         if is_onedp_project(azure_ai_project):
             self.azure_ai_project = azure_ai_project
-            self.credential=cast(TokenCredential, credential)
+            self.credential = cast(TokenCredential, credential)
             self.token_manager = ManagedIdentityAPITokenManager(
                 token_scope=TokenScope.COGNITIVE_SERVICES_MANAGEMENT,
                 logger=logging.getLogger("AdversarialSimulator"),
-                credential=self.credential
+                credential=self.credential,
             )
-            self.rai_client  = AIProjectClient(endpoint=azure_ai_project, credential=credential)
+            self.rai_client = AIProjectClient(endpoint=azure_ai_project, credential=credential)
         else:
             try:
                 self.azure_ai_project = validate_azure_ai_project(azure_ai_project)
@@ -67,14 +67,14 @@ class DirectAttackSimulator:
                     category=e.category,
                     blame=e.blame,
                 ) from e
-            self.credential = cast(TokenCredential, credential)    
+            self.credential = cast(TokenCredential, credential)
             self.token_manager = ManagedIdentityAPITokenManager(
                 token_scope=TokenScope.DEFAULT_AZURE_MANAGEMENT,
                 logger=logging.getLogger("AdversarialSimulator"),
                 credential=self.credential,
             )
             self.rai_client = RAIClient(azure_ai_project=self.azure_ai_project, token_manager=self.token_manager)
-        
+
         self.adversarial_template_handler = AdversarialTemplateHandler(
             azure_ai_project=self.azure_ai_project, rai_client=self.rai_client
         )

azure/ai/evaluation/simulator/_indirect_attack_simulator.py

@@ -30,8 +30,8 @@ class IndirectAttackSimulator(AdversarialSimulator):
     """
     Initializes the XPIA (cross domain prompt injected attack) jailbreak adversarial simulator with a project scope.
 
-    :param azure_ai_project: The Azure AI project, which can either be a string representing the project endpoint 
-        or an instance of AzureAIProject. It contains subscription id, resource group, and project name. 
+    :param azure_ai_project: The Azure AI project, which can either be a string representing the project endpoint
+        or an instance of AzureAIProject. It contains subscription id, resource group, and project name.
     :type azure_ai_project: Union[str, AzureAIProject]
     :param credential: The credential for connecting to Azure AI project.
     :type credential: ~azure.core.credentials.TokenCredential
@@ -51,13 +51,13 @@ class IndirectAttackSimulator(AdversarialSimulator):
 
         if is_onedp_project(azure_ai_project):
             self.azure_ai_project = azure_ai_project
-            self.credential=cast(TokenCredential, credential)
+            self.credential = cast(TokenCredential, credential)
             self.token_manager = ManagedIdentityAPITokenManager(
                 token_scope=TokenScope.COGNITIVE_SERVICES_MANAGEMENT,
                 logger=logging.getLogger("AdversarialSimulator"),
-                credential=self.credential
+                credential=self.credential,
             )
-            self.rai_client  = AIProjectClient(endpoint=azure_ai_project, credential=credential)
+            self.rai_client = AIProjectClient(endpoint=azure_ai_project, credential=credential)
             self.adversarial_template_handler = AdversarialTemplateHandler(
                 azure_ai_project=self.azure_ai_project, rai_client=self.rai_client
             )

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

@@ -6,6 +6,7 @@ import os
 from typing import Dict, List, Optional, Union
 
 from azure.core.credentials import TokenCredential
+from azure.core.pipeline.policies import UserAgentPolicy
 from azure.ai.evaluation._model_configurations import AzureAIProject
 from azure.ai.evaluation.simulator._model_tools import ManagedIdentityAPITokenManager
 from azure.ai.evaluation._common.raiclient import MachineLearningServicesClient
@@ -13,30 +14,35 @@ from azure.ai.evaluation._constants import TokenScope
 from azure.ai.evaluation._common.utils import is_onedp_project
 from azure.ai.evaluation._common.onedp import AIProjectClient
 from azure.ai.evaluation._common import EvaluationServiceOneDPClient
+from azure.ai.evaluation._user_agent import UserAgentSingleton
 import jwt
 import time
 import ast
 
+
 class GeneratedRAIClient:
     """Client for the Responsible AI Service using the auto-generated MachineLearningServicesClient.
-    
-    :param azure_ai_project: The scope of the Azure AI project. It contains subscription id, resource group, and project name.
-    :type azure_ai_project: ~azure.ai.evaluation.AzureAIProject
+
+    :param azure_ai_project: The Azure AI project, which can either be a string representing the project endpoint
+        or an instance of AzureAIProject. It contains subscription id, resource group, and project name.
+    :type azure_ai_project: Union[str, ~azure.ai.evaluation.AzureAIProject]
     :param token_manager: The token manager
     :type token_manager: ~azure.ai.evaluation.simulator._model_tools._identity_manager.APITokenManager
     """
-    
+
     def __init__(self, azure_ai_project: Union[AzureAIProject, str], token_manager: ManagedIdentityAPITokenManager):
         self.azure_ai_project = azure_ai_project
         self.token_manager = token_manager
-        
+
+        user_agent_policy = UserAgentPolicy(base_user_agent=UserAgentSingleton().value)
+
         if not is_onedp_project(azure_ai_project):
             # Service URL construction
             if "RAI_SVC_URL" in os.environ:
                 endpoint = os.environ["RAI_SVC_URL"].rstrip("/")
             else:
                 endpoint = self._get_service_discovery_url()
-            
+
             # Create the autogenerated client
             self._client = MachineLearningServicesClient(
                 endpoint=endpoint,
@@ -46,20 +52,24 @@ class GeneratedRAIClient:
                 credential=self.token_manager,
             ).rai_svc
         else:
-            self._client = AIProjectClient(endpoint=azure_ai_project, credential=token_manager).red_teams
-            self._operations_client = AIProjectClient(endpoint=azure_ai_project, credential=token_manager).evaluations
-            self._evaluation_onedp_client = EvaluationServiceOneDPClient(endpoint=azure_ai_project, credential=token_manager)
-        
+            self._client = AIProjectClient(
+                endpoint=azure_ai_project, credential=token_manager, user_agent_policy=user_agent_policy
+            ).red_teams
+            self._evaluation_onedp_client = EvaluationServiceOneDPClient(
+                endpoint=azure_ai_project, credential=token_manager, user_agent_policy=user_agent_policy
+            )
+
     def _get_service_discovery_url(self):
         """Get the service discovery URL.
-        
+
         :return: The service discovery URL
         :rtype: str
         """
         import requests
+
         bearer_token = self._fetch_or_reuse_token(self.token_manager)
         headers = {"Authorization": f"Bearer {bearer_token}", "Content-Type": "application/json"}
-        
+
         response = requests.get(
             f"https://management.azure.com/subscriptions/{self.azure_ai_project['subscription_id']}/"
             f"resourceGroups/{self.azure_ai_project['resource_group_name']}/"
@@ -68,7 +78,7 @@ class GeneratedRAIClient:
             headers=headers,
             timeout=5,
         )
-        
+
         if response.status_code != 200:
             msg = (
                 f"Failed to connect to your Azure AI project. Please check if the project scope is configured "
@@ -79,51 +89,71 @@ class GeneratedRAIClient:
 
         # Parse the discovery URL
         from urllib.parse import urlparse
+
         base_url = urlparse(response.json()["properties"]["discoveryUrl"])
         return f"{base_url.scheme}://{base_url.netloc}"
-    
-    async def get_attack_objectives(self, risk_category: Optional[str] = None, application_scenario: str = None, strategy: Optional[str] = None) -> Dict:
+
+    async def get_attack_objectives(
+        self,
+        *,
+        risk_type: Optional[str] = None,
+        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.
-        
+
+        :param risk_type: Optional risk type to filter the attack objectives
+        :type risk_type: Optional[str]
         :param risk_category: Optional risk category to filter the attack objectives
         :type risk_category: Optional[str]
         :param application_scenario: Optional description of the application scenario for context
         :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
-        """ 
+        """
         try:
             # Send the request using the autogenerated client
             response = self._client.get_attack_objectives(
-                risk_types=[risk_category],
+                risk_types=[risk_type],
+                risk_category=risk_category,
                 lang="en",
                 strategy=strategy,
+                headers={"client_request_id": scan_session_id},
             )
             return response
-            
+
         except Exception as e:
             # Log the exception for debugging purposes
             import logging
+
             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:
                 self.logger.error("Unexpected response format from get_jail_break_dataset_with_type")
                 raise ValueError("Unexpected response format from get_jail_break_dataset_with_type")
-            
+
         except Exception as e:
             return [""]
 

azure/ai/evaluation/simulator/_model_tools/_identity_manager.py

@@ -78,7 +78,13 @@ class APITokenManager(ABC):
 
     @abstractmethod
     def get_token(
-            self, scopes: Union[str, None] = None, claims: Union[str, None] = None, tenant_id: Union[str, None] = None, enable_cae: bool = False, **kwargs: Any) -> AccessToken:
+        self,
+        scopes: Union[str, None] = None,
+        claims: Union[str, None] = None,
+        tenant_id: Union[str, None] = None,
+        enable_cae: bool = False,
+        **kwargs: Any,
+    ) -> AccessToken:
         """Async method to get the API token. Subclasses should implement this method.
 
         :return: API token

azure/ai/evaluation/simulator/_model_tools/_proxy_completion_model.py

@@ -9,7 +9,7 @@ import uuid
 from typing import Any, Dict, List, Optional, cast, Union
 
 from azure.ai.evaluation._http_utils import AsyncHttpPipeline, get_async_http_client
-from azure.ai.evaluation._user_agent import USER_AGENT
+from azure.ai.evaluation._user_agent import UserAgentSingleton
 from azure.core.exceptions import HttpResponseError
 from azure.core.pipeline.policies import AsyncRetryPolicy, RetryMode
 from azure.ai.evaluation._common.onedp._client import AIProjectClient
@@ -166,7 +166,7 @@ class ProxyChatCompletionsModel(OpenAIChatCompletionsModel):
         proxy_headers = {
             "Authorization": f"Bearer {token}",
             "Content-Type": "application/json",
-            "User-Agent": USER_AGENT,
+            "User-Agent": UserAgentSingleton().value,
         }
 
         headers = {
@@ -193,7 +193,7 @@ class ProxyChatCompletionsModel(OpenAIChatCompletionsModel):
         time_start = time.time()
         full_response = None
 
-        if(isinstance(session, AIProjectClient)):
+        if isinstance(session, AIProjectClient):
             sim_request_dto = SimulationDTO(
                 headers=headers,
                 params=params,
@@ -203,7 +203,7 @@ class ProxyChatCompletionsModel(OpenAIChatCompletionsModel):
             )
             response_data = session.red_teams.submit_simulation(sim_request_dto, headers=headers, params=params)
             operation_id = response_data["location"].split("/")[-1]
-            
+
             request_count = 0
             flag = True
             while flag:
@@ -211,6 +211,7 @@ class ProxyChatCompletionsModel(OpenAIChatCompletionsModel):
                     response = session.evaluations.operation_results(operation_id, headers=headers)
                 except Exception as e:
                     from types import SimpleNamespace  # pylint: disable=forgotten-debug-statement
+
                     response = SimpleNamespace(status_code=202, text=str(e), json=lambda: {"error": str(e)})
                 if isinstance(response, dict):
                     response_data = response
@@ -228,10 +229,11 @@ class ProxyChatCompletionsModel(OpenAIChatCompletionsModel):
             # response.raise_for_status()
             if response.status_code != 202:
                 raise HttpResponseError(
-                    message=f"Received unexpected HTTP status: {response.status_code} {response.text()}", response=response
+                    message=f"Received unexpected HTTP status: {response.status_code} {response.text()}",
+                    response=response,
                 )
             response_data = response.json()
-            
+
             self.result_url = cast(str, response_data["location"])
             retry_policy = AsyncRetryPolicy(  # set up retry configuration
                 retry_on_status_codes=[202],  # on which statuses to retry
@@ -252,14 +254,14 @@ class ProxyChatCompletionsModel(OpenAIChatCompletionsModel):
                 proxy_headers = {
                     "Authorization": f"Bearer {token}",
                     "Content-Type": "application/json",
-                    "User-Agent": USER_AGENT,
+                    "User-Agent": UserAgentSingleton().value,
                 }
                 response = await exp_retry_client.get(  # pylint: disable=too-many-function-args,unexpected-keyword-arg
                     self.result_url, headers=proxy_headers
                 )
             response.raise_for_status()
             response_data = response.json()
-            
+
         self.logger.info("Response: %s", response_data)
 
         # Copy the full response and return it to be saved in jsonl.

azure/ai/evaluation/simulator/_model_tools/_rai_client.py

@@ -10,7 +10,7 @@ import json
 from azure.ai.evaluation._exceptions import ErrorBlame, ErrorCategory, ErrorTarget, EvaluationException
 from azure.ai.evaluation._http_utils import AsyncHttpPipeline, get_async_http_client, get_http_client
 from azure.ai.evaluation._model_configurations import AzureAIProject
-from azure.ai.evaluation._user_agent import USER_AGENT
+from azure.ai.evaluation._user_agent import UserAgentSingleton
 from azure.core.pipeline.policies import AsyncRetryPolicy, RetryMode
 
 from ._identity_manager import APITokenManager
@@ -148,7 +148,7 @@ class RAIClient:  # pylint: disable=client-accepts-api-version-keyword
         headers = {
             "Authorization": f"Bearer {token}",
             "Content-Type": "application/json",
-            "User-Agent": USER_AGENT,
+            "User-Agent": UserAgentSingleton().value,
         }
 
         session = self._create_async_client()
@@ -184,7 +184,7 @@ class RAIClient:  # pylint: disable=client-accepts-api-version-keyword
         headers = {
             "Authorization": f"Bearer {token}",
             "Content-Type": "application/json",
-            "User-Agent": USER_AGENT,
+            "User-Agent": UserAgentSingleton().value,
         }
 
         session = self._create_async_client()
@@ -209,9 +209,11 @@ class RAIClient:  # pylint: disable=client-accepts-api-version-keyword
             blame=ErrorBlame.USER_ERROR,
         )
 
-    async def get_attack_objectives(self, risk_categories: List[str], application_scenario: str = None, strategy: str = None) -> Any:
+    async def get_attack_objectives(
+        self, risk_categories: List[str], application_scenario: str = None, strategy: str = None
+    ) -> Any:
         """Get the attack objectives based on risk categories and application scenario
-        
+
         :param risk_categories: List of risk categories to generate attack objectives for
         :type risk_categories: List[str]
         :param application_scenario: Optional description of the application scenario for context
@@ -224,17 +226,17 @@ class RAIClient:  # pylint: disable=client-accepts-api-version-keyword
         # Create query parameters for the request
         if application_scenario:
             raise NotImplementedError("Application scenario is not supported yet")
-        
+
         params = {
             "api-version": "2022-11-01-preview",
             "riskTypes": ",".join(risk_categories),
-            "lang": "en"  # Default to English
+            "lang": "en",  # Default to English
         }
-        
+
         # Add strategy parameter if provided
         if strategy:
             params["strategy"] = strategy
-            
+
         try:
             # Make the request using the existing get method
             result = await self.get(self.attack_objectives_endpoint)
@@ -254,25 +256,11 @@ class RAIClient:  # pylint: disable=client-accepts-api-version-keyword
         except Exception:
             # If the API fails or isn't implemented yet, return a mock response
             # This is temporary until the API endpoint is fully implemented
-            return [{
-                "metadata": {
-                    "lang": "en",
-                    "target_harms": [
-                        {
-                            "risk-type": "violence",
-                            "risk-subtype": ""
-                        }
-                    ]
-                },
-                "messages": [
-                    {
-                        "role": "user",
-                        "content": "Risky content"
-                    }
-                ],
-                "modality": "text",
-                "source": [
-                    "source"
-                ]
-            }]
- 
+            return [
+                {
+                    "metadata": {"lang": "en", "target_harms": [{"risk-type": "violence", "risk-subtype": ""}]},
+                    "messages": [{"role": "user", "content": "Risky content"}],
+                    "modality": "text",
+                    "source": ["source"],
+                }
+            ]

azure/ai/evaluation/simulator/_model_tools/_template_handler.py

@@ -8,6 +8,7 @@ from typing_extensions import NotRequired
 
 from azure.ai.evaluation._model_configurations import AzureAIProject
 from azure.ai.evaluation._common.onedp._client import AIProjectClient
+from azure.ai.evaluation.simulator._adversarial_scenario import AdversarialScenario
 
 from ._rai_client import RAIClient
 
@@ -148,14 +149,16 @@ class AdversarialTemplateHandler:
     """
     Initialize the AdversarialTemplateHandler.
 
-    :param azure_ai_project: The Azure AI project, which can either be a string representing the project endpoint 
-        or an instance of AzureAIProject. It contains subscription id, resource group, and project name. 
+    :param azure_ai_project: The Azure AI project, which can either be a string representing the project endpoint
+        or an instance of AzureAIProject. It contains subscription id, resource group, and project name.
     :type azure_ai_project: Union[str, AzureAIProject]
     :param rai_client: The RAI client or AI Project client used for fetching parameters.
     :type rai_client: Union[~azure.ai.evaluation.simulator._model_tools.RAIClient, ~azure.ai.evaluation._common.onedp._client.AIProjectClient]
     """
 
-    def __init__(self, azure_ai_project: Union[str, AzureAIProject], rai_client: Union[RAIClient, AIProjectClient]) -> None:
+    def __init__(
+        self, azure_ai_project: Union[str, AzureAIProject], rai_client: Union[RAIClient, AIProjectClient]
+    ) -> None:
         self.azure_ai_project = azure_ai_project
         self.categorized_ch_parameters: Optional[Dict[str, _CategorizedParameter]] = None
         self.rai_client = rai_client
@@ -164,12 +167,11 @@ class AdversarialTemplateHandler:
         if self.categorized_ch_parameters is None:
             categorized_parameters: Dict[str, _CategorizedParameter] = {}
             util = ContentHarmTemplatesUtils
-
             if isinstance(self.rai_client, RAIClient):
                 parameters = await self.rai_client.get_contentharm_parameters()
             elif isinstance(self.rai_client, AIProjectClient):
                 parameters = literal_eval(self.rai_client.red_teams.get_template_parameters())
-            
+
             for k in parameters.keys():
                 template_key = util.get_template_key(k)
                 categorized_parameters[template_key] = {
@@ -181,17 +183,29 @@ class AdversarialTemplateHandler:
 
         template_category = collection_key.split("adv_")[-1]
 
+        # Handle both qa_enterprise and qa_documents mapping to qa
+        if template_category in ["qa_enterprise", "qa_documents"]:
+            template_category = "qa"
+
         plist = self.categorized_ch_parameters
         ch_templates = []
+
         for key, value in plist.items():
+            # Skip enterprise templates for ADVERSARIAL_QA
+            if collection_key == AdversarialScenario.ADVERSARIAL_QA.value and "enterprise" in key:
+                continue
+            # Skip non-enterprise templates for ADVERSARIAL_QA_DOCUMENTS
+            if collection_key == AdversarialScenario.ADVERSARIAL_QA_DOCUMENTS.value and "enterprise" not in key:
+                continue
+
             if value["category"] == template_category:
                 params = value["parameters"]
                 for p in params:
                     p.update({"ch_template_placeholder": "{{ch_template_placeholder}}"})
 
                 template = AdversarialTemplate(template_name=key, text=None, context_key=[], template_parameters=params)
-
                 ch_templates.append(template)
+
         return ch_templates
 
     def get_template(self, template_name: str) -> Optional[AdversarialTemplate]:

azure/ai/evaluation/simulator/_model_tools/models.py

@@ -478,7 +478,7 @@ class OpenAICompletionsModel(LLMBase):
         time_start = time.time()
         full_response = None
 
-        if(isinstance(session, AIProjectClient)):
+        if isinstance(session, AIProjectClient):
             response_data = session.red_teams.submit_simulation(request_data, headers, params)
         else:
             response = await session.post(url=self.endpoint_url, headers=headers, json=request_data, params=params)