azure-ai-evaluation 1.5.0__py3-none-any.whl → 1.6.0__py3-none-any.whl

azure/ai/evaluation/red_team/_utils/metric_mapping.py

@@ -0,0 +1,23 @@
+"""
+Mapping utilities for metrics and risk categories in Red Team Agent.
+"""
+from .._attack_objective_generator import RiskCategory
+from azure.ai.evaluation._constants import EvaluationMetrics
+
+# Mapping of risk categories to their corresponding evaluation metrics
+RISK_CATEGORY_METRIC_MAP = {
+    RiskCategory.Violence: EvaluationMetrics.VIOLENCE,
+    RiskCategory.HateUnfairness: EvaluationMetrics.HATE_UNFAIRNESS,
+    RiskCategory.Sexual: EvaluationMetrics.SEXUAL,
+    RiskCategory.SelfHarm: EvaluationMetrics.SELF_HARM
+}
+
+def get_metric_from_risk_category(risk_category: RiskCategory) -> str:
+    """Get the appropriate evaluation metric for a given risk category.
+    
+    :param risk_category: The risk category to map to an evaluation metric
+    :type risk_category: RiskCategory
+    :return: The corresponding evaluation metric
+    :rtype: str
+    """
+    return RISK_CATEGORY_METRIC_MAP.get(risk_category, EvaluationMetrics.HATE_UNFAIRNESS)

azure/ai/evaluation/red_team/_utils/strategy_utils.py

@@ -189,4 +189,4 @@ def get_orchestrators_for_attack_strategies(attack_strategies: List[Union[Attack
     # Since we're just returning one orchestrator type for now, simplify the logic
     # This can be expanded later if different orchestrators are needed for different strategies
     return [lambda chat_target, all_prompts, converter, strategy_name, risk_category: 
-            None]  # This will be replaced with the actual orchestrator function in the main class
+            None]  # This will be replaced with the actual orchestrator function in the main class

azure/ai/evaluation/simulator/_adversarial_simulator.py

@@ -12,12 +12,14 @@ import uuid
 from tqdm import tqdm
 
 from azure.ai.evaluation._common._experimental import experimental
-from azure.ai.evaluation._common.utils import validate_azure_ai_project
+from azure.ai.evaluation._common.utils import validate_azure_ai_project, is_onedp_project
+from azure.ai.evaluation._common.onedp._client import AIProjectClient
 from azure.ai.evaluation._exceptions import ErrorBlame, ErrorCategory, ErrorTarget, EvaluationException
 from azure.ai.evaluation._http_utils import get_async_http_client
 from azure.ai.evaluation._model_configurations import AzureAIProject
 from azure.ai.evaluation.simulator import AdversarialScenario, AdversarialScenarioJailbreak
 from azure.ai.evaluation.simulator._adversarial_scenario import _UnstableAdversarialScenario
+from azure.ai.evaluation._constants import TokenScope
 from azure.core.credentials import TokenCredential
 from azure.core.pipeline.policies import AsyncRetryPolicy, RetryMode
 
@@ -35,7 +37,6 @@ from ._model_tools import (
     ManagedIdentityAPITokenManager,
     ProxyChatCompletionsModel,
     RAIClient,
-    TokenScope,
 )
 from ._model_tools._template_handler import AdversarialTemplate, TemplateParameters
 from ._utils import JsonLineList
@@ -48,9 +49,9 @@ class AdversarialSimulator:
     """
     Initializes the adversarial simulator with a project scope.
 
-    :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, AzureAIProject]
     :param credential: The credential for connecting to Azure AI project.
     :type credential: ~azure.core.credentials.TokenCredential
 
@@ -65,26 +66,37 @@ class AdversarialSimulator:
                 2 conversation turns each (4 messages per result).
     """
 
-    def __init__(self, *, azure_ai_project: AzureAIProject, credential: TokenCredential):
+    def __init__(self, *, azure_ai_project: Union[str, AzureAIProject], credential: TokenCredential):
         """Constructor."""
 
-        try:
-            self.azure_ai_project = validate_azure_ai_project(azure_ai_project)
-        except EvaluationException as e:
-            raise EvaluationException(
-                message=e.message,
-                internal_message=e.internal_message,
-                target=ErrorTarget.ADVERSARIAL_SIMULATOR,
-                category=e.category,
-                blame=e.blame,
-            ) from e
-
-        self.token_manager = ManagedIdentityAPITokenManager(
-            token_scope=TokenScope.DEFAULT_AZURE_MANAGEMENT,
-            logger=logging.getLogger("AdversarialSimulator"),
-            credential=cast(TokenCredential, credential),
-        )
-        self.rai_client = RAIClient(azure_ai_project=self.azure_ai_project, token_manager=self.token_manager)
+        if is_onedp_project(azure_ai_project):
+            self.azure_ai_project = azure_ai_project
+            self.credential = cast(TokenCredential, credential)
+            self.token_manager = ManagedIdentityAPITokenManager(
+                token_scope=TokenScope.COGNITIVE_SERVICES_MANAGEMENT,
+                logger=logging.getLogger("AdversarialSimulator"),
+                credential=self.credential,
+            )
+            self.rai_client  = AIProjectClient(endpoint=azure_ai_project, credential=credential)
+        else:
+            try:
+                self.azure_ai_project = validate_azure_ai_project(azure_ai_project)
+            except EvaluationException as e:
+                raise EvaluationException(
+                    message=e.message,
+                    internal_message=e.internal_message,
+                    target=ErrorTarget.ADVERSARIAL_SIMULATOR,
+                    category=e.category,
+                    blame=e.blame,
+                ) from e
+            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
         )
@@ -188,6 +200,12 @@ class AdversarialSimulator:
             )
         self._ensure_service_dependencies()
         templates = await self.adversarial_template_handler._get_content_harm_template_collections(scenario.value)
+        if len(templates) == 0:
+            raise EvaluationException(
+                message="Templates not found. Please check https://aka.ms/azureaiadvsimulator-regionsupport for region support.",
+                internal_message="Please check https://aka.ms/azureaiadvsimulator-regionsupport for region support.",
+                target=ErrorTarget.ADVERSARIAL_SIMULATOR,
+            )
         simulation_id = str(uuid.uuid4())
         logger.warning("Use simulation_id to help debug the issue: %s", str(simulation_id))
         concurrent_async_task = min(concurrent_async_task, 1000)
@@ -206,7 +224,10 @@ class AdversarialSimulator:
         total_tasks = min(total_tasks, max_simulation_results)
         _jailbreak_type = kwargs.get("_jailbreak_type", None)
         if _jailbreak_type:
-            jailbreak_dataset = await self.rai_client.get_jailbreaks_dataset(type=_jailbreak_type)
+            if isinstance(self.rai_client, RAIClient):
+                jailbreak_dataset = await self.rai_client.get_jailbreaks_dataset(type=_jailbreak_type)
+            elif isinstance(self.rai_client, AIProjectClient):
+                jailbreak_dataset = self.rai_client.red_teams.get_jail_break_dataset_with_type(type=_jailbreak_type)
         progress_bar = tqdm(
             total=total_tasks,
             desc="generating jailbreak simulations" if _jailbreak_type else "generating simulations",
@@ -336,23 +357,25 @@ class AdversarialSimulator:
             target=target, role=ConversationRole.ASSISTANT, template=template, parameters=parameters, scenario=scenario
         )
         bots = [user_bot, system_bot]
-        session = get_async_http_client().with_policies(
-            retry_policy=AsyncRetryPolicy(
-                retry_total=api_call_retry_limit,
-                retry_backoff_factor=api_call_retry_sleep_sec,
-                retry_mode=RetryMode.Fixed,
+        
+        if isinstance(self.rai_client, AIProjectClient):
+            session = self.rai_client
+        else:    
+            session = get_async_http_client().with_policies(
+                retry_policy=AsyncRetryPolicy(
+                    retry_total=api_call_retry_limit,
+                    retry_backoff_factor=api_call_retry_sleep_sec,
+                    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,
         )
 
-        async with semaphore, session:
-            _, conversation_history = await simulate_conversation(
-                bots=bots,
-                session=session,
-                turn_limit=max_conversation_turns,
-                api_call_delay_sec=api_call_delay_sec,
-                language=language,
-            )
-
         return self._to_chat_protocol(
             conversation_history=conversation_history,
             template_parameters=cast(Dict[str, Union[str, Dict[str, str]]], parameters),
@@ -361,11 +384,12 @@ 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    
         return ProxyChatCompletionsModel(
             name="raisvc_proxy_model",
             template_key=template_key,
             template_parameters=template_parameters,
-            endpoint_url=self.rai_client.simulation_submit_endpoint,
+            endpoint_url=endpoint_url,
             token_manager=self.token_manager,
             api_version="2023-07-01-preview",
             max_tokens=1200,

azure/ai/evaluation/simulator/_constants.py

@@ -25,3 +25,4 @@ class SupportedLanguages(Enum):
     Portuguese = "pt"
     Japanese = "ja"
     English = "en"
+    Korean = "ko"

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

@@ -8,13 +8,14 @@ import logging
 import time
 from dataclasses import dataclass
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union, cast
-
+import base64
 import re
 import jinja2
 
 from azure.ai.evaluation._exceptions import ErrorBlame, ErrorCategory, ErrorTarget, EvaluationException
 from azure.ai.evaluation._http_utils import AsyncHttpPipeline
 from .._model_tools import LLMBase, OpenAIChatCompletionsModel, RAIClient
+from azure.ai.evaluation._common.onedp._client import AIProjectClient
 from .._model_tools._template_handler import TemplateParameters
 from .constants import ConversationRole
 
@@ -145,7 +146,7 @@ class ConversationBot:
 
     async def generate_response(
         self,
-        session: AsyncHttpPipeline,
+        session: Union[AsyncHttpPipeline, AIProjectClient],
         conversation_history: List[ConversationTurn],
         max_history: int,
         turn_number: int = 0,
@@ -255,7 +256,7 @@ class CallbackConversationBot(ConversationBot):
 
     async def generate_response(
         self,
-        session: AsyncHttpPipeline,
+        session: Union[AsyncHttpPipeline, AIProjectClient],
         conversation_history: List[Any],
         max_history: int,
         turn_number: int = 0,
@@ -329,7 +330,7 @@ class MultiModalConversationBot(ConversationBot):
         callback: Callable,
         user_template: str,
         user_template_parameters: TemplateParameters,
-        rai_client: RAIClient,
+        rai_client: Union[RAIClient, AIProjectClient],
         *args,
         **kwargs,
     ) -> None:
@@ -342,7 +343,7 @@ class MultiModalConversationBot(ConversationBot):
 
     async def generate_response(
         self,
-        session: AsyncHttpPipeline,
+        session: Union[AsyncHttpPipeline, AIProjectClient],
         conversation_history: List[Any],
         max_history: int,
         turn_number: int = 0,
@@ -418,7 +419,13 @@ class MultiModalConversationBot(ConversationBot):
         contents = []
         for msg in messages:
             if msg.startswith("image_understanding/"):
-                encoded_image = await self.rai_client.get_image_data(msg)
+                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/_conversation/_conversation.py

@@ -11,6 +11,7 @@ from azure.ai.evaluation.simulator._constants import SupportedLanguages
 from azure.ai.evaluation.simulator._helpers._language_suffix_mapping import SUPPORTED_LANGUAGES_MAPPING
 from ..._http_utils import AsyncHttpPipeline
 from . import ConversationBot, ConversationTurn
+from azure.ai.evaluation._common.onedp._client import AIProjectClient
 
 
 def is_closing_message(response: Union[Dict, str], recursion_depth: int = 0) -> bool:
@@ -72,7 +73,7 @@ def is_closing_message_helper(response: str) -> bool:
 async def simulate_conversation(
     *,
     bots: List[ConversationBot],
-    session: AsyncHttpPipeline,
+    session: Union[AsyncHttpPipeline, AIProjectClient],
     language: SupportedLanguages,
     stopping_criteria: Callable[[str], bool] = is_closing_message,
     turn_limit: int = 10,

azure/ai/evaluation/simulator/_direct_attack_simulator.py

@@ -5,17 +5,19 @@
 # noqa: E501
 import logging
 from random import randint
-from typing import Callable, Optional, cast
+from typing import Callable, Optional, cast, Union
 
+from azure.ai.evaluation._constants import TokenScope
 from azure.ai.evaluation._common._experimental import experimental
-from azure.ai.evaluation._common.utils import validate_azure_ai_project
+from azure.ai.evaluation._common.utils import validate_azure_ai_project, is_onedp_project
 from azure.ai.evaluation._exceptions import ErrorBlame, ErrorCategory, ErrorTarget, EvaluationException
 from azure.ai.evaluation.simulator import AdversarialScenario
 from azure.ai.evaluation._model_configurations import AzureAIProject
+from azure.ai.evaluation._common.onedp._client import AIProjectClient
 from azure.core.credentials import TokenCredential
 
 from ._adversarial_simulator import AdversarialSimulator
-from ._model_tools import AdversarialTemplateHandler, ManagedIdentityAPITokenManager, RAIClient, TokenScope
+from ._model_tools import AdversarialTemplateHandler, ManagedIdentityAPITokenManager, RAIClient
 
 logger = logging.getLogger(__name__)
 
@@ -42,26 +44,37 @@ class DirectAttackSimulator:
             :caption: Run the DirectAttackSimulator to produce 2 results with 3 conversation turns each (6 messages in each result).
     """
 
-    def __init__(self, *, azure_ai_project: AzureAIProject, credential: TokenCredential):
+    def __init__(self, *, azure_ai_project: Union[str, AzureAIProject], credential: TokenCredential):
         """Constructor."""
-
-        try:
-            self.azure_ai_project = validate_azure_ai_project(azure_ai_project)
-        except EvaluationException as e:
-            raise EvaluationException(
-                message=e.message,
-                internal_message=e.internal_message,
-                target=ErrorTarget.DIRECT_ATTACK_SIMULATOR,
-                category=e.category,
-                blame=e.blame,
-            ) from e
-        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)
+        
+        if is_onedp_project(azure_ai_project):
+            self.azure_ai_project = azure_ai_project
+            self.credential=cast(TokenCredential, credential)
+            self.token_manager = ManagedIdentityAPITokenManager(
+                token_scope=TokenScope.COGNITIVE_SERVICES_MANAGEMENT,
+                logger=logging.getLogger("AdversarialSimulator"),
+                credential=self.credential
+            )
+            self.rai_client  = AIProjectClient(endpoint=azure_ai_project, credential=credential)
+        else:
+            try:
+                self.azure_ai_project = validate_azure_ai_project(azure_ai_project)
+            except EvaluationException as e:
+                raise EvaluationException(
+                    message=e.message,
+                    internal_message=e.internal_message,
+                    target=ErrorTarget.DIRECT_ATTACK_SIMULATOR,
+                    category=e.category,
+                    blame=e.blame,
+                ) from e
+            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/_helpers/_language_suffix_mapping.py

@@ -14,4 +14,5 @@ SUPPORTED_LANGUAGES_MAPPING = {
     SupportedLanguages.SimplifiedChinese: BASE_SUFFIX.replace("__language__", "simplified chinese"),
     SupportedLanguages.Portuguese: BASE_SUFFIX.replace("__language__", "portuguese"),
     SupportedLanguages.Japanese: BASE_SUFFIX.replace("__language__", "japanese"),
+    SupportedLanguages.Korean: BASE_SUFFIX.replace("__language__", "korean"),
 }

azure/ai/evaluation/simulator/_indirect_attack_simulator.py

@@ -5,20 +5,22 @@
 # noqa: E501
 import asyncio
 import logging
-from typing import Callable, cast
+from typing import Callable, cast, Union
 
 from tqdm import tqdm
 
-from azure.ai.evaluation._common.utils import validate_azure_ai_project
+from azure.ai.evaluation._common.utils import validate_azure_ai_project, is_onedp_project
 from azure.ai.evaluation._common._experimental import experimental
 from azure.ai.evaluation._exceptions import ErrorBlame, ErrorCategory, ErrorTarget, EvaluationException
 from azure.ai.evaluation.simulator import AdversarialScenarioJailbreak, SupportedLanguages
 from azure.ai.evaluation._model_configurations import AzureAIProject
+from azure.ai.evaluation._common.onedp._client import AIProjectClient
 from azure.core.credentials import TokenCredential
+from azure.ai.evaluation._constants import TokenScope
 
 from ._adversarial_simulator import AdversarialSimulator, JsonLineList
 
-from ._model_tools import AdversarialTemplateHandler, ManagedIdentityAPITokenManager, RAIClient, TokenScope
+from ._model_tools import AdversarialTemplateHandler, ManagedIdentityAPITokenManager, RAIClient
 
 logger = logging.getLogger(__name__)
 
@@ -44,30 +46,43 @@ class IndirectAttackSimulator(AdversarialSimulator):
             :caption: Run the IndirectAttackSimulator to produce 1 result with 1 conversation turn (2 messages in the result).
     """
 
-    def __init__(self, *, azure_ai_project: AzureAIProject, credential: TokenCredential):
+    def __init__(self, *, azure_ai_project: Union[str, AzureAIProject], credential: TokenCredential):
         """Constructor."""
 
-        try:
-            self.azure_ai_project = validate_azure_ai_project(azure_ai_project)
-        except EvaluationException as e:
-            raise EvaluationException(
-                message=e.message,
-                internal_message=e.internal_message,
-                target=ErrorTarget.DIRECT_ATTACK_SIMULATOR,
-                category=e.category,
-                blame=e.blame,
-            ) from e
-
-        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
-        )
+        if is_onedp_project(azure_ai_project):
+            self.azure_ai_project = azure_ai_project
+            self.credential=cast(TokenCredential, credential)
+            self.token_manager = ManagedIdentityAPITokenManager(
+                token_scope=TokenScope.COGNITIVE_SERVICES_MANAGEMENT,
+                logger=logging.getLogger("AdversarialSimulator"),
+                credential=self.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
+            )
+        else:
+            try:
+                self.azure_ai_project = validate_azure_ai_project(azure_ai_project)
+            except EvaluationException as e:
+                raise EvaluationException(
+                    message=e.message,
+                    internal_message=e.internal_message,
+                    target=ErrorTarget.DIRECT_ATTACK_SIMULATOR,
+                    category=e.category,
+                    blame=e.blame,
+                ) from e
+
+            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
+            )
         super().__init__(azure_ai_project=azure_ai_project, credential=credential)
 
     def _ensure_service_dependencies(self):

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

@@ -4,11 +4,12 @@
 
 """Tooling for model evaluation"""
 
-from ._identity_manager import ManagedIdentityAPITokenManager, PlainTokenManager, TokenScope
+from ._identity_manager import ManagedIdentityAPITokenManager, PlainTokenManager
 from ._proxy_completion_model import ProxyChatCompletionsModel
 from ._rai_client import RAIClient
 from ._template_handler import CONTENT_HARM_TEMPLATES_COLLECTION_KEY, AdversarialTemplateHandler
 from .models import LLMBase, OpenAIChatCompletionsModel
+from ..._constants import TokenScope
 
 __all__ = [
     "ManagedIdentityAPITokenManager",

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

@@ -3,12 +3,15 @@
 # ---------------------------------------------------------
 
 import os
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Union
 
 from azure.core.credentials import TokenCredential
 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
+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
 import jwt
 import time
 import ast
@@ -22,24 +25,27 @@ class GeneratedRAIClient:
     :type token_manager: ~azure.ai.evaluation.simulator._model_tools._identity_manager.APITokenManager
     """
     
-    def __init__(self, azure_ai_project: AzureAIProject, token_manager: ManagedIdentityAPITokenManager):
+    def __init__(self, azure_ai_project: Union[AzureAIProject, str], token_manager: ManagedIdentityAPITokenManager):
         self.azure_ai_project = azure_ai_project
         self.token_manager = token_manager
         
-        # Service URL construction
-        if "RAI_SVC_URL" in os.environ:
-            endpoint = os.environ["RAI_SVC_URL"].rstrip("/")
-        else:
-            endpoint = self._get_service_discovery_url()
+        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,
-            subscription_id=self.azure_ai_project["subscription_id"],
-            resource_group_name=self.azure_ai_project["resource_group_name"],
-            workspace_name=self.azure_ai_project["project_name"],
-            credential=self.token_manager,
-        )
+            # Create the autogenerated client
+            self._client = MachineLearningServicesClient(
+                endpoint=endpoint,
+                subscription_id=self.azure_ai_project["subscription_id"],
+                resource_group_name=self.azure_ai_project["resource_group_name"],
+                workspace_name=self.azure_ai_project["project_name"],
+                credential=self.token_manager,
+            ).rai_svc
+        else:
+            self._client = AIProjectClient(endpoint=azure_ai_project, credential=token_manager).red_teams
         
     def _get_service_discovery_url(self):
         """Get the service discovery URL.
@@ -87,7 +93,7 @@ class GeneratedRAIClient:
         """ 
         try:
             # Send the request using the autogenerated client
-            response = self._client.rai_svc.get_attack_objectives(
+            response = self._client.get_attack_objectives(
                 risk_types=[risk_category],
                 lang="en",
                 strategy=strategy,
@@ -108,7 +114,7 @@ class GeneratedRAIClient:
         """
         try:
             # Send the request using the autogenerated client
-            response = self._client.rai_svc.get_jail_break_dataset_with_type(type="upia")
+            response = self._client.get_jail_break_dataset_with_type(type="upia")
             if isinstance(response, list):
                 return response
             else:
@@ -142,4 +148,4 @@ class GeneratedRAIClient:
                 if (exp_time - current_time) >= 300:
                     return token
 
-        return credential.get_token("https://management.azure.com/.default").token
+        return credential.get_token(TokenScope.DEFAULT_AZURE_MANAGEMENT).token

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

@@ -8,9 +8,9 @@ import logging
 import os
 import time
 from abc import ABC, abstractmethod
-from enum import Enum
-from typing import Optional, Union
+from typing import Optional, Union, Any
 
+from azure.ai.evaluation._constants import TokenScope
 from azure.core.credentials import AccessToken, TokenCredential
 from azure.identity import DefaultAzureCredential, ManagedIdentityCredential
 
@@ -19,12 +19,6 @@ AZURE_TOKEN_REFRESH_INTERVAL = int(
 )  # token refresh interval in seconds
 
 
-class TokenScope(Enum):
-    """Token scopes for Azure endpoints"""
-
-    DEFAULT_AZURE_MANAGEMENT = "https://management.azure.com/.default"
-
-
 class APITokenManager(ABC):
     """Base class for managing API tokens. Subclasses should implement the get_token method.
 
@@ -83,7 +77,8 @@ class APITokenManager(ABC):
         return DefaultAzureCredential()
 
     @abstractmethod
-    def get_token(self) -> str:
+    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:
         """Async method to get the API token. Subclasses should implement this method.
 
         :return: API token
@@ -103,7 +98,7 @@ class ManagedIdentityAPITokenManager(APITokenManager):
     """API Token Manager for Azure Managed Identity
 
     :param token_scope: Token scope for Azure endpoint
-    :type token_scope: ~azure.ai.evaluation.simulator._model_tools.TokenScope
+    :type token_scope: ~azure.ai.evaluation._constants.TokenScope
     :param logger: Logger object
     :type logger: logging.Logger
     :keyword kwargs: Additional keyword arguments