azure-ai-evaluation 1.0.1__py3-none-any.whl → 1.13.3__py3-none-any.whl

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 ProjectsClient as 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__)
 
@@ -26,9 +28,9 @@ 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 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
 
@@ -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,23 @@
 # noqa: E501
 import asyncio
 import logging
-from typing import Callable, cast
+import random
+from typing import Callable, cast, Union, Optional
 
 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 ProjectsClient as 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__)
 
@@ -28,9 +31,9 @@ class IndirectAttackSimulator(AdversarialSimulator):
     """
     Initializes the XPIA (cross domain prompt injected attack) jailbreak 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
 
@@ -44,30 +47,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):
@@ -90,6 +106,7 @@ class IndirectAttackSimulator(AdversarialSimulator):
         api_call_retry_sleep_sec: int = 1,
         api_call_delay_sec: int = 0,
         concurrent_async_task: int = 3,
+        randomization_seed: Optional[int] = None,
         **kwargs,
     ):
         """
@@ -115,6 +132,9 @@ class IndirectAttackSimulator(AdversarialSimulator):
         :keyword concurrent_async_task: The number of asynchronous tasks to run concurrently during the simulation.
             Defaults to 3.
         :paramtype concurrent_async_task: int
+        :keyword randomization_seed: The seed used to randomize prompt selection. If unset, the system's
+            default seed is used. Defaults to None.
+        :paramtype randomization_seed: Optional[int]
         :return: A list of dictionaries, each representing a simulated conversation. Each dictionary contains:
 
          - 'template_parameters': A dictionary with parameters used in the conversation template,
@@ -175,6 +195,13 @@ class IndirectAttackSimulator(AdversarialSimulator):
             ncols=100,
             unit="simulations",
         )
+
+        # Apply randomization to templates if seed is provided
+        if randomization_seed is not None:
+            # Create a local random instance to avoid polluting global state
+            local_random = random.Random(randomization_seed)
+            local_random.shuffle(templates)
+
         for template in templates:
             for parameter in template.template_parameters:
                 tasks.append(
@@ -189,6 +216,7 @@ class IndirectAttackSimulator(AdversarialSimulator):
                             api_call_delay_sec=api_call_delay_sec,
                             language=language,
                             semaphore=semaphore,
+                            scenario=scenario,
                         )
                     )
                 )

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

@@ -0,0 +1,225 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+
+import logging
+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
+from azure.ai.evaluation._constants import TokenScope
+from azure.ai.evaluation._common.utils import is_onedp_project
+from azure.ai.evaluation._common.onedp import ProjectsClient as 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 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
+        self.logger = logging.getLogger(__name__)
+
+        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,
+                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,
+                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']}/"
+            f"providers/Microsoft.MachineLearningServices/workspaces/{self.azure_ai_project['project_name']}?"
+            f"api-version=2023-08-01-preview",
+            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 "
+                f"correctly, and make sure you have the necessary access permissions. "
+                f"Status code: {response.status_code}."
+            )
+            raise Exception(msg)
+
+        # 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_type: Optional[str] = None,
+        risk_category: Optional[str] = None,
+        application_scenario: str = None,
+        strategy: Optional[str] = None,
+        language: str = "en",
+        scan_session_id: Optional[str] = None,
+        target: Optional[str] = None,
+        client_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 language: Language code for the attack objectives (e.g., "en", "es", "fr")
+        :type language: str
+        :param scan_session_id: Optional unique session ID for the scan
+        :type scan_session_id: Optional[str]
+        :param target: Optional target type (model/agent)
+        :type target: Optional[str]
+        :param client_id: Optional client ID for ACA token authorization
+        :type client_id: Optional[str]
+        :return: The attack objectives
+        :rtype: Dict
+        """
+        try:
+            # Build headers dictionary
+            headers = {}
+            if scan_session_id:
+                headers["x-ms-client-request-id"] = scan_session_id
+            if client_id:
+                from azure.identity import DefaultAzureCredential
+
+                self.logger.info(f"Using client_id: {client_id} to set token in aml-aca-token header ")
+
+                # Get token using the client_id for managed identity
+                managed_identity_credential = DefaultAzureCredential(
+                    managed_identity_client_id=client_id, exclude_interactive_browser_credential=True
+                )
+                token = managed_identity_credential.get_token(TokenScope.DEFAULT_AZURE_MANAGEMENT).token
+                headers["aml-aca-token"] = token
+
+            # Send the request using the autogenerated client
+            response = self._client.get_attack_objectives(
+                risk_types=[risk_type],
+                risk_category=risk_category,
+                lang=language,
+                strategy=strategy,
+                target_type=target,
+                headers=headers,
+            )
+
+            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, 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", headers={"x-ms-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 [""]
+
+    def _fetch_or_reuse_token(self, credential: TokenCredential, token: Optional[str] = None) -> str:
+        """Get token. Fetch a new token if the current token is near expiry
+
+        :param credential: The Azure authentication credential.
+        :type credential:
+        ~azure.core.credentials.TokenCredential
+        :param token: The Azure authentication token. Defaults to None. If none, a new token will be fetched.
+        :type token: str
+        :return: The Azure authentication token.
+        """
+        if token:
+            # Decode the token to get its expiration time
+            try:
+                decoded_token = jwt.decode(token, options={"verify_signature": False})
+            except jwt.PyJWTError:
+                pass
+            else:
+                exp_time = decoded_token["exp"]
+                current_time = time.time()
+
+                # Return current token if not near expiry
+                if (exp_time - current_time) >= 300:
+                    return token
+
+        # Get 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,14 @@ 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 +104,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
@@ -151,6 +152,7 @@ class ManagedIdentityAPITokenManager(APITokenManager):
         ):
             self.last_refresh_time = time.time()
             get_token_method = self.credential.get_token(self.token_scope.value)
+
             if inspect.isawaitable(get_token_method):
                 # If it's awaitable, await it
                 token_response: AccessToken = await get_token_method