azure-ai-evaluation 1.0.0b2__py3-none-any.whl → 1.13.3__py3-none-any.whl

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

@@ -3,22 +3,20 @@
 # ---------------------------------------------------------
 
 import asyncio
+import inspect
 import logging
 import os
 import time
 from abc import ABC, abstractmethod
-from enum import Enum
-from typing import Dict, 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
 
-AZURE_TOKEN_REFRESH_INTERVAL = 600  # seconds
-
-
-class TokenScope(Enum):
-    """Token scopes for Azure endpoints"""
-
-    DEFAULT_AZURE_MANAGEMENT = "https://management.azure.com/.default"
+AZURE_TOKEN_REFRESH_INTERVAL = int(
+    os.getenv("AZURE_TOKEN_REFRESH_INTERVAL", "600")
+)  # token refresh interval in seconds
 
 
 class APITokenManager(ABC):
@@ -29,24 +27,24 @@ class APITokenManager(ABC):
     :param auth_header: Authorization header prefix. Defaults to "Bearer"
     :type auth_header: str
     :param credential: Azure credential object
-    :type credential: Optional[Union[azure.identity.DefaultAzureCredential, azure.identity.ManagedIdentityCredential]
+    :type credential: Optional[TokenCredential]
     """
 
     def __init__(
         self,
         logger: logging.Logger,
         auth_header: str = "Bearer",
-        credential: Optional[Union[DefaultAzureCredential, ManagedIdentityCredential]] = None,
+        credential: Optional[TokenCredential] = None,
     ) -> None:
         self.logger = logger
         self.auth_header = auth_header
-        self._lock = None
+        self._lock: Optional[asyncio.Lock] = None
         if credential is not None:
             self.credential = credential
         else:
             self.credential = self.get_aad_credential()
-        self.token = None
-        self.last_refresh_time = None
+        self.token: Optional[str] = None
+        self.last_refresh_time: Optional[float] = None
 
     @property
     def lock(self) -> asyncio.Lock:
@@ -73,39 +71,58 @@ class APITokenManager(ABC):
         identity_client_id = os.environ.get("DEFAULT_IDENTITY_CLIENT_ID", None)
         if identity_client_id is not None:
             self.logger.info(f"Using DEFAULT_IDENTITY_CLIENT_ID: {identity_client_id}")
-            credential = ManagedIdentityCredential(client_id=identity_client_id)
-        else:
-            self.logger.info("Environment variable DEFAULT_IDENTITY_CLIENT_ID is not set, using DefaultAzureCredential")
-            credential = DefaultAzureCredential()
-        return credential
+            return ManagedIdentityCredential(client_id=identity_client_id)
+
+        self.logger.info("Environment variable DEFAULT_IDENTITY_CLIENT_ID is not set, using DefaultAzureCredential")
+        return DefaultAzureCredential()
 
     @abstractmethod
-    async 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
+        :rtype: str
+        """
+
+    @abstractmethod
+    async def get_token_async(self) -> str:
         """Async method to get the API token. Subclasses should implement this method.
 
         :return: API token
         :rtype: str
         """
-        pass  # pylint: disable=unnecessary-pass
 
 
 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
     :paramtype kwargs: Dict
     """
 
-    def __init__(self, token_scope: TokenScope, logger: logging.Logger, **kwargs: Dict):
-        super().__init__(logger, **kwargs)
+    def __init__(
+        self,
+        token_scope: TokenScope,
+        logger: logging.Logger,
+        *,
+        auth_header: str = "Bearer",
+        credential: Optional[TokenCredential] = None,
+    ):
+        super().__init__(logger, auth_header=auth_header, credential=credential)
         self.token_scope = token_scope
 
-    # Bug 3353724: This get_token is sync method, but it is defined as async method in the base class
-    def get_token(self) -> str:  # pylint: disable=invalid-overridden-method
+    def get_token(self) -> str:
         """Get the API token. If the token is not available or has expired, refresh the token.
 
         :return: API token
@@ -122,6 +139,32 @@ class ManagedIdentityAPITokenManager(APITokenManager):
 
         return self.token
 
+    async def get_token_async(self) -> str:
+        """Get the API token synchronously. If the token is not available or has expired, refresh it.
+
+        :return: API token
+        :rtype: str
+        """
+        if (
+            self.token is None
+            or self.last_refresh_time is None
+            or time.time() - self.last_refresh_time > AZURE_TOKEN_REFRESH_INTERVAL
+        ):
+            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
+            else:
+                # Otherwise, call it synchronously
+                token_response = get_token_method
+
+            self.token = token_response.token
+            self.logger.info("Refreshed Azure endpoint token.")
+
+        return self.token
+
 
 class PlainTokenManager(APITokenManager):
     """Plain API Token Manager
@@ -134,11 +177,18 @@ class PlainTokenManager(APITokenManager):
     :paramtype kwargs: Dict
     """
 
-    def __init__(self, openapi_key: str, logger: logging.Logger, **kwargs: Dict):
-        super().__init__(logger, **kwargs)
-        self.token = openapi_key
+    def __init__(
+        self,
+        openapi_key: str,
+        logger: logging.Logger,
+        *,
+        auth_header: str = "Bearer",
+        credential: Optional[TokenCredential] = None,
+    ) -> None:
+        super().__init__(logger, auth_header=auth_header, credential=credential)
+        self.token: str = openapi_key
 
-    async def get_token(self) -> str:
+    def get_token(self) -> str:
         """Get the API token
 
         :return: API token

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

@@ -6,13 +6,17 @@ import copy
 import json
 import time
 import uuid
-from typing import Dict, List
+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.core.exceptions import HttpResponseError
+from azure.ai.evaluation._user_agent import UserAgentSingleton
+from azure.core.exceptions import HttpResponseError, ServiceResponseError
 from azure.core.pipeline.policies import AsyncRetryPolicy, RetryMode
+from azure.ai.evaluation._common.onedp._client import ProjectsClient as AIProjectClient
+from azure.ai.evaluation._common.onedp.models import SimulationDTO
+from azure.ai.evaluation._common.constants import RAIService
 
+from .._model_tools._template_handler import TemplateParameters
 from .models import OpenAIChatCompletionsModel
 
 
@@ -33,13 +37,21 @@ class SimulationRequestDTO:
     :type template_parameters: Dict
     """
 
-    def __init__(self, url, headers, payload, params, templatekey, template_parameters):
+    def __init__(
+        self,
+        url: str,
+        headers: Dict[str, str],
+        payload: Dict[str, Any],
+        params: Dict[str, str],
+        templateKey: str,
+        templateParameters: Optional[TemplateParameters],
+    ):
         self.url = url
         self.headers = headers
         self.json = json.dumps(payload)
         self.params = params
-        self.templatekey = templatekey
-        self.templateParameters = template_parameters
+        self.templateKey = templateKey
+        self.templateParameters = templateParameters
 
     def to_dict(self) -> Dict:
         """Convert the DTO to a dictionary.
@@ -47,9 +59,12 @@ class SimulationRequestDTO:
         :return: The DTO as a dictionary.
         :rtype: Dict
         """
-        if self.templateParameters is not None:
-            self.templateParameters = {str(k): str(v) for k, v in self.templateParameters.items()}
-        return self.__dict__
+        toReturn = self.__dict__.copy()
+
+        if toReturn["templateParameters"] is not None:
+            toReturn["templateParameters"] = {str(k): str(v) for k, v in toReturn["templateParameters"].items()}
+
+        return toReturn
 
     def to_json(self):
         """Convert the DTO to a JSON string.
@@ -73,12 +88,13 @@ class ProxyChatCompletionsModel(OpenAIChatCompletionsModel):
     :keyword kwargs: Additional keyword arguments to pass to the parent class.
     """
 
-    def __init__(self, name: str, template_key: str, template_parameters, *args, **kwargs) -> None:
+    def __init__(self, name: str, template_key: str, template_parameters: TemplateParameters, **kwargs) -> None:
         self.tkey = template_key
         self.tparam = template_parameters
-        self.result_url = None
+        self.result_url: Optional[str] = None
+        self.simulation_id: Optional[str] = kwargs.pop("simulation_id", "")
 
-        super().__init__(name=name, *args, **kwargs)
+        super().__init__(name=name, **kwargs)
 
     def format_request_data(self, messages: List[Dict], **request_params) -> Dict:  # type: ignore[override]
         """Format the request data to query the model with.
@@ -98,7 +114,7 @@ class ProxyChatCompletionsModel(OpenAIChatCompletionsModel):
     async def get_conversation_completion(
         self,
         messages: List[Dict],
-        session: AsyncHttpPipeline,
+        session: Union[AsyncHttpPipeline, AIProjectClient],
         role: str = "assistant",  # pylint: disable=unused-argument
         **request_params,
     ) -> dict:
@@ -129,7 +145,7 @@ class ProxyChatCompletionsModel(OpenAIChatCompletionsModel):
 
     async def request_api(
         self,
-        session: AsyncHttpPipeline,
+        session: Union[AsyncHttpPipeline, AIProjectClient],
         request_data: dict,
     ) -> dict:
         """
@@ -150,17 +166,17 @@ class ProxyChatCompletionsModel(OpenAIChatCompletionsModel):
         proxy_headers = {
             "Authorization": f"Bearer {token}",
             "Content-Type": "application/json",
-            "User-Agent": USER_AGENT,
+            "User-Agent": UserAgentSingleton().value,
         }
 
         headers = {
             "Content-Type": "application/json",
             "X-CV": f"{uuid.uuid4()}",
             "X-ModelType": self.model or "",
+            "x-ms-client-request-id": self.simulation_id,
         }
         # add all additional headers
         headers.update(self.additional_headers)  # type: ignore[arg-type]
-
         params = {}
         if self.api_version:
             params["api-version"] = self.api_version
@@ -170,45 +186,101 @@ class ProxyChatCompletionsModel(OpenAIChatCompletionsModel):
             headers=headers,
             payload=request_data,
             params=params,
-            templatekey=self.tkey,
-            template_parameters=self.tparam,
+            templateKey=self.tkey,
+            templateParameters=self.tparam,
         )
 
         time_start = time.time()
         full_response = None
 
-        response = await session.post(url=self.endpoint_url, headers=proxy_headers, json=sim_request_dto.to_dict())
-
-        if response.status_code != 202:
-            raise HttpResponseError(
-                message=f"Received unexpected HTTP status: {response.status_code} {response.text()}", response=response
+        if isinstance(session, AIProjectClient):
+            sim_request_dto = SimulationDTO(
+                headers=headers,
+                params=params,
+                json=json.dumps(request_data),
+                template_key=self.tkey,
+                template_parameters=self.tparam,
+            )
+            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:
+                try:
+                    response = session.red_teams.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
+                    flag = False
+                    break
+                if not isinstance(response, SimpleNamespace) and response.get("object") == "chat.completion":
+                    response_data = response
+                    flag = False
+                    break
+                else:
+                    request_count += 1
+                    sleep_time = RAIService.SLEEP_TIME**request_count
+                    await asyncio.sleep(sleep_time)
+        else:
+            # Retry policy for POST request to RAI service
+            service_call_retry_policy = AsyncRetryPolicy(
+                retry_on_exceptions=[ServiceResponseError],
+                retry_total=7,
+                retry_backoff_factor=10.0,
+                retry_backoff_max=180,
+                retry_mode=RetryMode.Exponential,
             )
 
-        response = response.json()
-        self.result_url = response["location"]
-
-        retry_policy = AsyncRetryPolicy(  # set up retry configuration
-            retry_on_status_codes=[202],  # on which statuses to retry
-            retry_total=7,
-            retry_backoff_factor=10.0,
-            retry_backoff_max=180,
-            retry_mode=RetryMode.Exponential,
-        )
-
-        # initial 15 seconds wait before attempting to fetch result
-        # Need to wait both in this thread and in the async thread for some reason?
-        # Someone not under a crunch and with better async understandings should dig into this more.
-        await asyncio.sleep(15)
-        time.sleep(15)
-
-        async with get_async_http_client().with_policies(retry_policy=retry_policy) as exp_retry_client:
-            response = await exp_retry_client.get(  # pylint: disable=too-many-function-args,unexpected-keyword-arg
-                self.result_url, headers=proxy_headers
+            response = None
+            async with get_async_http_client().with_policies(retry_policy=service_call_retry_policy) as retry_client:
+                try:
+                    response = await retry_client.post(
+                        url=self.endpoint_url, headers=proxy_headers, json=sim_request_dto.to_dict()
+                    )
+                except ServiceResponseError as e:
+                    self.logger.error("ServiceResponseError during POST request to rai svc after retries: %s", str(e))
+                    raise
+
+            # response.raise_for_status()
+            if response.status_code != 202:
+                raise HttpResponseError(
+                    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
+                retry_total=7,
+                retry_backoff_factor=10.0,
+                retry_backoff_max=180,
+                retry_mode=RetryMode.Exponential,
             )
 
-        response.raise_for_status()
+            # initial 15 seconds wait before attempting to fetch result
+            # Need to wait both in this thread and in the async thread for some reason?
+            # Someone not under a crunch and with better async understandings should dig into this more.
+            await asyncio.sleep(15)
+            time.sleep(15)
+
+            async with get_async_http_client().with_policies(retry_policy=retry_policy) as exp_retry_client:
+                token = await self.token_manager.get_token_async()
+                proxy_headers = {
+                    "Authorization": f"Bearer {token}",
+                    "Content-Type": "application/json",
+                    "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()
 
-        response_data = response.json()
         self.logger.info("Response: %s", response_data)
 
         # Copy the full response and return it to be saved in jsonl.