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

azure/ai/evaluation/_azure/__init__.py

@@ -0,0 +1,3 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------

azure/ai/evaluation/_azure/_clients.py

@@ -0,0 +1,204 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+
+from logging import Logger
+from typing import Any, Dict, Final, Optional, Set, Union, cast
+from threading import Lock
+from urllib.parse import quote
+from json.decoder import JSONDecodeError
+
+from azure.core.credentials import TokenCredential, AzureSasCredential
+from azure.core.rest import HttpResponse
+from azure.ai.evaluation._exceptions import ErrorBlame, ErrorCategory, ErrorTarget, EvaluationException
+from azure.ai.evaluation._http_utils import HttpPipeline, get_http_client
+from azure.ai.evaluation._azure._token_manager import AzureMLTokenManager
+from azure.ai.evaluation.simulator._model_tools._identity_manager import TokenScope
+from ._models import BlobStoreInfo, Workspace
+
+
+API_VERSION: Final[str] = "2024-07-01-preview"
+QUERY_KEY_API_VERSION: Final[str] = "api-version"
+PATH_ML_WORKSPACES = ("providers", "Microsoft.MachineLearningServices", "workspaces")
+
+
+class LiteMLClient:
+    """A lightweight Azure ML API client.
+
+    :param subscription_id: Azure subscription ID
+    :type subscription_id: str
+    :param resource_group: Azure resource group name
+    :type resource_group: str
+    :param logger: Logger object
+    :type logger: logging.Logger
+    :keyword credential: Azure credentials
+    :paramtype credential: TokenCredential
+    :keyword kwargs: Additional keyword arguments
+    :paramtype kwargs: Dict
+    :keyword str api_version: The API version. Default is 2024-10-01
+    """
+
+    def __init__(
+        self,
+        subscription_id: str,
+        resource_group: str,
+        logger: Logger,
+        credential: Optional[TokenCredential] = None,
+        **kwargs: Any,
+    ) -> None:
+        subscription_id = quote(subscription_id, safe="")
+        resource_group = quote(resource_group, safe="")
+
+        self._base_url: Final[str] = (
+            f"https://management.azure.com/subscriptions/{subscription_id}/resourceGroups/{resource_group}"
+        )
+        self._logger: Final[Logger] = logger
+        self._api_version: Final[str] = kwargs.get("api_version", API_VERSION)
+        self._http_client: Final[HttpPipeline] = get_http_client(**kwargs)
+        self._lock: Final[Lock] = Lock()
+
+        # things that can change under lock
+        self._token_manager: Optional[AzureMLTokenManager] = None
+        self._credential: Optional[TokenCredential] = credential
+
+    def get_token(self) -> str:
+        return self._get_token_manager().get_token()
+
+    def get_credential(self) -> TokenCredential:
+        # load the token manager to get the credential if needed
+        self._get_token_manager()
+        return cast(TokenCredential, self._credential)
+
+    def workspace_get_default_datastore(
+        self, workspace_name: str, *, include_credentials: bool = False, **kwargs: Any
+    ) -> BlobStoreInfo:
+        # 1. Get the default blob store
+        # REST API documentation:
+        # https://learn.microsoft.com/rest/api/azureml/datastores/list?view=rest-azureml-2024-10-01
+        url = self._generate_path(  # pylint: disable=specify-parameter-names-in-call
+            *PATH_ML_WORKSPACES, workspace_name, "datastores"
+        )
+        headers = self._get_headers()
+
+        stores_response = self._http_client.request(
+            method="GET",
+            url=url,
+            params={QUERY_KEY_API_VERSION: self._api_version, "isDefault": True, "count": 1, "orderByAsc": "false"},
+            headers=headers,
+        )
+        self._throw_on_http_error(stores_response, "list default workspace datastore")
+
+        json = stores_response.json()["value"][0]
+        props_json = json["properties"]
+        name = json["name"]
+        account_name = props_json["accountName"]
+        endpoint = props_json["endpoint"]
+        container_name = props_json["containerName"]
+        credential_type = props_json.get("credentials", {}).get("credentialsType")
+
+        # 2. Get the SAS token to use for accessing the blob store
+        # REST API documentation:
+        # https://learn.microsoft.com/rest/api/azureml/datastores/list-secrets?view=rest-azureml-2024-10-01
+        blob_store_credential: Optional[Union[AzureSasCredential, TokenCredential, str]]
+        if not include_credentials:
+            blob_store_credential = None
+        elif credential_type and credential_type.lower() == "none":
+            # If storage account key access is disabled, and only Microsoft Entra ID authentication is available,
+            # the credentialsType will be "None" and we should not attempt to get the secrets.
+            blob_store_credential = self.get_credential()
+        else:
+            url = self._generate_path(
+                *PATH_ML_WORKSPACES, workspace_name, "datastores", "workspaceblobstore", "listSecrets"
+            )
+            secrets_response = self._http_client.request(
+                method="POST",
+                url=url,
+                json={
+                    "expirableSecret": True,
+                    "expireAfterHours": int(kwargs.get("key_expiration_hours", 1)),
+                },
+                params={
+                    QUERY_KEY_API_VERSION: self._api_version,
+                },
+                headers=headers,
+            )
+            self._throw_on_http_error(secrets_response, "workspace datastore secrets")
+
+            secrets_json = secrets_response.json()
+            secrets_type = secrets_json["secretsType"].lower()
+
+            # As per this website, only SAS tokens, access tokens, or Entra IDs are valid for accessing blob data
+            # stores:
+            # https://learn.microsoft.com/rest/api/storageservices/authorize-requests-to-azure-storage.
+            if secrets_type == "sas":
+                blob_store_credential = AzureSasCredential(secrets_json["sasToken"])
+            elif secrets_type == "accountkey":
+                # To support older versions of azure-storage-blob better, we return a string here instead of
+                # an AzureNamedKeyCredential
+                blob_store_credential = secrets_json["key"]
+            else:
+                raise EvaluationException(
+                    message=f"The '{account_name}' blob store does not use a recognized credential type.",
+                    internal_message=f"The credential type is '{secrets_type}'",
+                    target=ErrorTarget.EVALUATE,
+                    category=ErrorCategory.INVALID_VALUE,
+                    blame=ErrorBlame.SYSTEM_ERROR,
+                )
+
+        return BlobStoreInfo(name, account_name, endpoint, container_name, blob_store_credential)
+
+    def workspace_get_info(self, workspace_name: str) -> Workspace:
+        # https://learn.microsoft.com/rest/api/azureml/workspaces/get?view=rest-azureml-2024-10-01
+        workspace_response = self._http_client.request(
+            "GET",
+            self._generate_path(*PATH_ML_WORKSPACES, workspace_name),
+            params={QUERY_KEY_API_VERSION: self._api_version},
+            headers=self._get_headers(),
+        )
+
+        self._throw_on_http_error(workspace_response, f"get '{workspace_name}' workspace")
+        workspace = Workspace.deserialize(workspace_response)
+        return workspace
+
+    def _get_token_manager(self) -> AzureMLTokenManager:
+        # Lazy init since getting credentials in the constructor can take a long time in some situations
+        if self._token_manager is None:
+            with self._lock:
+                if self._token_manager is None:
+                    self._token_manager = AzureMLTokenManager(
+                        TokenScope.DEFAULT_AZURE_MANAGEMENT.value, self._logger, credential=self._credential
+                    )
+                    self._credential = self._token_manager.credential
+
+        return self._token_manager
+
+    @staticmethod
+    def _throw_on_http_error(response: HttpResponse, description: str, valid_status: Optional[Set[int]] = None) -> None:
+        if valid_status and (response.status_code in valid_status):
+            return
+        if response.status_code >= 200 and response.status_code < 300:
+            # nothing to see here, move along
+            return
+
+        message = f"The {description} request failed with HTTP {response.status_code}"
+        try:
+            error_json = response.json()["error"]
+            additional_info = f"({error_json['code']}) {error_json['message']}"
+            message += f" - {additional_info}"
+        except (JSONDecodeError, ValueError, KeyError):
+            pass
+
+        raise EvaluationException(
+            message=message,
+            target=ErrorTarget.EVALUATE,
+            category=ErrorCategory.FAILED_EXECUTION,
+            blame=ErrorBlame.SYSTEM_ERROR,
+        )
+
+    def _generate_path(self, *paths: str) -> str:
+        sanitized_paths = [quote(path, safe="") for path in paths]
+        url = self._base_url + "/" + str.join("/", sanitized_paths)
+        return url
+
+    def _get_headers(self) -> Dict[str, str]:
+        return {"Authorization": f"Bearer {self.get_token()}", "Content-Type": "application/json"}

azure/ai/evaluation/_azure/_models.py

@@ -0,0 +1,227 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+
+# pylint: disable=too-many-instance-attributes
+# pylint: disable=too-many-locals
+# pylint: disable=line-too-long
+
+from typing import Dict, List, NamedTuple, Optional, Union
+from msrest.serialization import Model
+from azure.core.credentials import AzureSasCredential, TokenCredential
+
+
+class BlobStoreInfo(NamedTuple):
+    name: str
+    account_name: str
+    endpoint: str
+    container_name: str
+    credential: Optional[Union[AzureSasCredential, TokenCredential, str]]
+
+
+class WorkspaceHubConfig(Model):
+    """WorkspaceHub's configuration object."""
+
+    _attribute_map = {
+        "additional_workspace_storage_accounts": {"key": "additionalWorkspaceStorageAccounts", "type": "[str]"},
+        "default_workspace_resource_group": {"key": "defaultWorkspaceResourceGroup", "type": "str"},
+    }
+
+    def __init__(
+        self,
+        *,
+        additional_workspace_storage_accounts: Optional[List[str]] = None,
+        default_workspace_resource_group: Optional[str] = None,
+        **kwargs
+    ):
+        super(WorkspaceHubConfig, self).__init__(**kwargs)
+        self.additional_workspace_storage_accounts = additional_workspace_storage_accounts
+        self.default_workspace_resource_group = default_workspace_resource_group
+
+
+class Workspace(Model):
+    """An object that represents a machine learning workspace.
+
+    Variables are only populated by the server, and will be ignored when sending a request."""
+
+    _validation = {
+        "id": {"readonly": True},
+        "name": {"readonly": True},
+        "type": {"readonly": True},
+        #'system_data': {'readonly': True},
+        "agents_endpoint_uri": {"readonly": True},
+        "ml_flow_tracking_uri": {"readonly": True},
+        #'notebook_info': {'readonly': True},
+        "private_endpoint_connections": {"readonly": True},
+        #'private_link_count': {'readonly': True},
+        "provisioning_state": {"readonly": True},
+        "service_provisioned_resource_group": {"readonly": True},
+        "storage_hns_enabled": {"readonly": True},
+        "tenant_id": {"readonly": True},
+        "workspace_id": {"readonly": True},
+    }
+
+    _attribute_map = {
+        "id": {"key": "id", "type": "str"},
+        "name": {"key": "name", "type": "str"},
+        "type": {"key": "type", "type": "str"},
+        #'system_data': {'key': 'systemData', 'type': 'SystemData'},
+        #'identity': {'key': 'identity', 'type': 'ManagedServiceIdentity'},
+        "kind": {"key": "kind", "type": "str"},
+        "location": {"key": "location", "type": "str"},
+        #'sku': {'key': 'sku', 'type': 'Sku'},
+        "tags": {"key": "tags", "type": "{str}"},
+        "agents_endpoint_uri": {"key": "properties.agentsEndpointUri", "type": "str"},
+        "allow_public_access_when_behind_vnet": {"key": "properties.allowPublicAccessWhenBehindVnet", "type": "bool"},
+        "allow_role_assignment_on_rg": {"key": "properties.allowRoleAssignmentOnRG", "type": "bool"},
+        "application_insights": {"key": "properties.applicationInsights", "type": "str"},
+        "associated_workspaces": {"key": "properties.associatedWorkspaces", "type": "[str]"},
+        "container_registries": {"key": "properties.containerRegistries", "type": "[str]"},
+        "container_registry": {"key": "properties.containerRegistry", "type": "str"},
+        "description": {"key": "properties.description", "type": "str"},
+        "discovery_url": {"key": "properties.discoveryUrl", "type": "str"},
+        "enable_data_isolation": {"key": "properties.enableDataIsolation", "type": "bool"},
+        "enable_service_side_cmk_encryption": {"key": "properties.enableServiceSideCMKEncryption", "type": "bool"},
+        "enable_simplified_cmk": {"key": "properties.enableSimplifiedCmk", "type": "bool"},
+        "enable_software_bill_of_materials": {"key": "properties.enableSoftwareBillOfMaterials", "type": "bool"},
+        #'encryption': {'key': 'properties.encryption', 'type': 'EncryptionProperty'},
+        "existing_workspaces": {"key": "properties.existingWorkspaces", "type": "[str]"},
+        #'feature_store_settings': {'key': 'properties.featureStoreSettings', 'type': 'FeatureStoreSettings'},
+        "friendly_name": {"key": "properties.friendlyName", "type": "str"},
+        "hbi_workspace": {"key": "properties.hbiWorkspace", "type": "bool"},
+        "hub_resource_id": {"key": "properties.hubResourceId", "type": "str"},
+        "image_build_compute": {"key": "properties.imageBuildCompute", "type": "str"},
+        "ip_allowlist": {"key": "properties.ipAllowlist", "type": "[str]"},
+        "key_vault": {"key": "properties.keyVault", "type": "str"},
+        "key_vaults": {"key": "properties.keyVaults", "type": "[str]"},
+        #'managed_network': {'key': 'properties.managedNetwork', 'type': 'ManagedNetworkSettings'},
+        "ml_flow_tracking_uri": {"key": "properties.mlFlowTrackingUri", "type": "str"},
+        #'network_acls': {'key': 'properties.networkAcls', 'type': 'NetworkAcls'},
+        #'notebook_info': {'key': 'properties.notebookInfo', 'type': 'NotebookResourceInfo'},
+        "primary_user_assigned_identity": {"key": "properties.primaryUserAssignedIdentity", "type": "str"},
+        "private_endpoint_connections": {
+            "key": "properties.privateEndpointConnections",
+            "type": "[PrivateEndpointConnection]",
+        },
+        "private_link_count": {"key": "properties.privateLinkCount", "type": "int"},
+        "provision_network_now": {"key": "properties.provisionNetworkNow", "type": "bool"},
+        "provisioning_state": {"key": "properties.provisioningState", "type": "str"},
+        #'public_network_access': {'key': 'properties.publicNetworkAccess', 'type': 'str'},
+        #'serverless_compute_settings': {'key': 'properties.serverlessComputeSettings', 'type': 'ServerlessComputeSettings'},
+        #'service_managed_resources_settings': {'key': 'properties.serviceManagedResourcesSettings', 'type': 'ServiceManagedResourcesSettings'},
+        "service_provisioned_resource_group": {"key": "properties.serviceProvisionedResourceGroup", "type": "str"},
+        #'shared_private_link_resources': {'key': 'properties.sharedPrivateLinkResources', 'type': '[SharedPrivateLinkResource]'},
+        "soft_delete_retention_in_days": {"key": "properties.softDeleteRetentionInDays", "type": "int"},
+        "storage_account": {"key": "properties.storageAccount", "type": "str"},
+        "storage_accounts": {"key": "properties.storageAccounts", "type": "[str]"},
+        "storage_hns_enabled": {"key": "properties.storageHnsEnabled", "type": "bool"},
+        #'system_datastores_auth_mode': {'key': 'properties.systemDatastoresAuthMode', 'type': 'str'},
+        "tenant_id": {"key": "properties.tenantId", "type": "str"},
+        "v1_legacy_mode": {"key": "properties.v1LegacyMode", "type": "bool"},
+        "workspace_hub_config": {"key": "properties.workspaceHubConfig", "type": "WorkspaceHubConfig"},
+        "workspace_id": {"key": "properties.workspaceId", "type": "str"},
+    }
+
+    def __init__(
+        self,
+        *,
+        # system_data: Optional[SystemData] = None,
+        # identity: Optional["ManagedServiceIdentity"] = None,
+        kind: Optional[str] = None,
+        location: Optional[str] = None,
+        # sku: Optional["Sku"] = None,
+        tags: Optional[Dict[str, str]] = None,
+        allow_public_access_when_behind_vnet: Optional[bool] = None,
+        allow_role_assignment_on_rg: Optional[bool] = None,
+        application_insights: Optional[str] = None,
+        associated_workspaces: Optional[List[str]] = None,
+        container_registries: Optional[List[str]] = None,
+        container_registry: Optional[str] = None,
+        description: Optional[str] = None,
+        discovery_url: Optional[str] = None,
+        enable_data_isolation: Optional[bool] = None,
+        enable_service_side_cmk_encryption: Optional[bool] = None,
+        enable_simplified_cmk: Optional[bool] = None,
+        enable_software_bill_of_materials: Optional[bool] = None,
+        # encryption: Optional["EncryptionProperty"] = None,
+        existing_workspaces: Optional[List[str]] = None,
+        # feature_store_settings: Optional["FeatureStoreSettings"] = None,
+        friendly_name: Optional[str] = None,
+        hbi_workspace: Optional[bool] = None,
+        hub_resource_id: Optional[str] = None,
+        image_build_compute: Optional[str] = None,
+        ip_allowlist: Optional[List[str]] = None,
+        key_vault: Optional[str] = None,
+        key_vaults: Optional[List[str]] = None,
+        # managed_network: Optional["ManagedNetworkSettings"] = None,
+        # network_acls: Optional["NetworkAcls"] = None,
+        primary_user_assigned_identity: Optional[str] = None,
+        provision_network_now: Optional[bool] = None,
+        # public_network_access: Optional[Union[str, "PublicNetworkAccessType"]] = None,
+        # serverless_compute_settings: Optional["ServerlessComputeSettings"] = None,
+        # service_managed_resources_settings: Optional["ServiceManagedResourcesSettings"] = None,
+        # shared_private_link_resources: Optional[List["SharedPrivateLinkResource"]] = None,
+        soft_delete_retention_in_days: Optional[int] = None,
+        storage_account: Optional[str] = None,
+        storage_accounts: Optional[List[str]] = None,
+        # system_datastores_auth_mode: Optional[Union[str, "SystemDatastoresAuthMode"]] = None,
+        v1_legacy_mode: Optional[bool] = None,
+        workspace_hub_config: Optional["WorkspaceHubConfig"] = None,
+        **kwargs
+    ):
+        super(Workspace, self).__init__(**kwargs)
+        self.id: Optional[str] = None
+        self.name: Optional[str] = None
+        self.type: Optional[str] = None
+        # self.system_data = system_data
+        # self.identity = identity
+        self.kind = kind
+        self.location = location
+        # self.sku = sku
+        self.tags = tags
+        self.agents_endpoint_uri = None
+        self.allow_public_access_when_behind_vnet = allow_public_access_when_behind_vnet
+        self.allow_role_assignment_on_rg = allow_role_assignment_on_rg
+        self.application_insights = application_insights
+        self.associated_workspaces = associated_workspaces
+        self.container_registries = container_registries
+        self.container_registry = container_registry
+        self.description = description
+        self.discovery_url = discovery_url
+        self.enable_data_isolation = enable_data_isolation
+        self.enable_service_side_cmk_encryption = enable_service_side_cmk_encryption
+        self.enable_simplified_cmk = enable_simplified_cmk
+        self.enable_software_bill_of_materials = enable_software_bill_of_materials
+        # self.encryption = encryption
+        self.existing_workspaces = existing_workspaces
+        # self.feature_store_settings = feature_store_settings
+        self.friendly_name = friendly_name
+        self.hbi_workspace = hbi_workspace
+        self.hub_resource_id = hub_resource_id
+        self.image_build_compute = image_build_compute
+        self.ip_allowlist = ip_allowlist
+        self.key_vault = key_vault
+        self.key_vaults = key_vaults
+        # self.managed_network = managed_network
+        self.ml_flow_tracking_uri = None
+        # self.network_acls = network_acls
+        # self.notebook_info = None
+        self.primary_user_assigned_identity = primary_user_assigned_identity
+        self.private_endpoint_connections = None
+        self.private_link_count = None
+        self.provision_network_now = provision_network_now
+        self.provisioning_state = None
+        # self.public_network_access = public_network_access
+        # self.serverless_compute_settings = serverless_compute_settings
+        # self.service_managed_resources_settings = service_managed_resources_settings
+        self.service_provisioned_resource_group = None
+        # self.shared_private_link_resources = shared_private_link_resources
+        self.soft_delete_retention_in_days = soft_delete_retention_in_days
+        self.storage_account = storage_account
+        self.storage_accounts = storage_accounts
+        self.storage_hns_enabled = None
+        # self.system_datastores_auth_mode = system_datastores_auth_mode
+        self.tenant_id = None
+        self.v1_legacy_mode = v1_legacy_mode
+        self.workspace_hub_config = workspace_hub_config
+        self.workspace_id = None

azure/ai/evaluation/_azure/_token_manager.py

@@ -0,0 +1,118 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+import os
+import logging
+import time
+import inspect
+from typing import cast, Optional, Union
+
+from azure.core.credentials import TokenCredential, AccessToken
+from azure.identity import AzureCliCredential, DefaultAzureCredential, ManagedIdentityCredential
+from azure.ai.evaluation._exceptions import ErrorBlame, ErrorCategory, ErrorTarget, EvaluationException
+from ..simulator._model_tools._identity_manager import APITokenManager, AZURE_TOKEN_REFRESH_INTERVAL
+
+
+class AzureMLTokenManager(APITokenManager):
+    """API Token manager for Azure Management API.
+
+    :param token_scope: Token scopes for Azure endpoint
+    :type token_scope: str
+    :param logger: Logger object
+    :type logger: logging.Logger
+    :keyword kwargs: Additional keyword arguments
+    :paramtype kwargs: Dict
+    """
+
+    def __init__(
+        self,
+        token_scope: str,
+        logger: logging.Logger,
+        credential: Optional[TokenCredential] = None,
+    ):
+        super().__init__(logger, credential=credential)
+        self.token_scope = token_scope
+        self.token_expiry_time: Optional[int] = None
+
+    def get_aad_credential(self) -> Union[DefaultAzureCredential, ManagedIdentityCredential]:
+        """Get the Azure credentials to use for the management APIs.
+
+        :return: Azure credentials
+        :rtype: DefaultAzureCredential or ManagedIdentityCredential
+        """
+        # Adds some of the additional types credentials that the previous Azure AI ML code used
+        # These may or may not be needed but kept here for backwards compatibility
+
+        if os.getenv("AZUREML_OBO_ENABLED"):
+            # using Azure on behalf of credentials requires the use of the azure-ai-ml package
+            try:
+                from azure.ai.ml.identity import AzureMLOnBehalfOfCredential
+
+                self.logger.debug("User identity is configured, use OBO credential.")
+                return AzureMLOnBehalfOfCredential()  # type: ignore
+            except (ModuleNotFoundError, ImportError):
+                raise EvaluationException(  # pylint: disable=raise-missing-from
+                    message=(
+                        "The required packages for OBO credentials are missing.\n"
+                        'To resolve this, please install them by running "pip install azure-ai-ml".'
+                    ),
+                    target=ErrorTarget.EVALUATE,
+                    category=ErrorCategory.MISSING_PACKAGE,
+                    blame=ErrorBlame.USER_ERROR,
+                )
+        elif os.environ.get("PF_USE_AZURE_CLI_CREDENTIAL", "false").lower() == "true":
+            self.logger.debug("Use azure cli credential since specified in environment variable.")
+            return AzureCliCredential()  # type: ignore
+        elif os.environ.get("IS_IN_CI_PIPELINE", "false").lower() == "true":
+            # use managed identity when executing in CI pipeline.
+            self.logger.debug("Use azure cli credential since in CI pipeline.")
+            return AzureCliCredential()  # type: ignore
+        else:
+            # Fall back to using the parent implementation
+            return super().get_aad_credential()
+
+    def get_token(self) -> str:
+        """Get the API token. If the token is not available or has expired, refresh the token.
+
+        :return: API token
+        :rtype: str
+        """
+        if self._token_needs_update():
+            credential = cast(TokenCredential, self.credential)
+            access_token = credential.get_token(self.token_scope)
+            self._update_token(access_token)
+
+        return cast(str, self.token)  # check for none is hidden in the _token_needs_update method
+
+    async def get_token_async(self) -> str:
+        """Get the API token asynchronously. If the token is not available or has expired, refresh it.
+
+        :return: API token
+        :rtype: str
+        """
+        if self._token_needs_update():
+            credential = cast(TokenCredential, self.credential)
+            get_token_method = credential.get_token(self.token_scope)
+            if inspect.isawaitable(get_token_method):
+                access_token = await get_token_method
+            else:
+                access_token = get_token_method
+            self._update_token(access_token)
+
+        return cast(str, self.token)  # check for none is hidden in the _token_needs_update method
+
+    def _token_needs_update(self) -> bool:
+        current_time = time.time()
+        return (
+            self.token is None
+            or self.last_refresh_time is None
+            or self.token_expiry_time is None
+            or self.token_expiry_time - current_time < AZURE_TOKEN_REFRESH_INTERVAL
+            or current_time - self.last_refresh_time > AZURE_TOKEN_REFRESH_INTERVAL
+        )
+
+    def _update_token(self, access_token: AccessToken) -> None:
+        self.token = cast(str, access_token.token)
+        self.token_expiry_time = access_token.expires_on
+        self.last_refresh_time = time.time()
+        self.logger.info("Refreshed Azure management token.")

azure/ai/evaluation/_common/rai_service.py

@@ -72,18 +72,21 @@ def get_formatted_template(data: dict, annotation_task: str) -> str:
     return user_text.replace("'", '\\"')
 
 
-def get_common_headers(token: str) -> Dict:
+def get_common_headers(token: str, evaluator_name: Optional[str] = None) -> Dict:
     """Get common headers for the HTTP request
 
     :param token: The Azure authentication token.
     :type token: str
+    :param evaluator_name: The evaluator name. Default is None.
+    :type evaluator_name: str
     :return: The common headers.
     :rtype: Dict
     """
+    user_agent = f"{USER_AGENT} (type=evaluator; subtype={evaluator_name})" if evaluator_name else USER_AGENT
     return {
         "Authorization": f"Bearer {token}",
         "Content-Type": "application/json",
-        "User-Agent": USER_AGENT,
+        "User-Agent": user_agent,
         # Handle "RuntimeError: Event loop is closed" from httpx AsyncClient
         # https://github.com/encode/httpx/discussions/2959
         "Connection": "close",
@@ -175,7 +178,9 @@ def generate_payload(normalized_user_text: str, metric: str, annotation_task: st
     )
 
 
-async def submit_request(data: dict, metric: str, rai_svc_url: str, token: str, annotation_task: str) -> str:
+async def submit_request(
+    data: dict, metric: str, rai_svc_url: str, token: str, annotation_task: str, evaluator_name: str
+) -> str:
     """Submit request to Responsible AI service for evaluation and return operation ID
 
     :param data: The data to evaluate.
@@ -188,6 +193,8 @@ async def submit_request(data: dict, metric: str, rai_svc_url: str, token: str,
     :type token: str
     :param annotation_task: The annotation task to use.
     :type annotation_task: str
+    :param evaluator_name: The evaluator name.
+    :type evaluator_name: str
     :return: The operation ID.
     :rtype: str
     """
@@ -195,7 +202,7 @@ async def submit_request(data: dict, metric: str, rai_svc_url: str, token: str,
     payload = generate_payload(normalized_user_text, metric, annotation_task=annotation_task)
 
     url = rai_svc_url + "/submitannotation"
-    headers = get_common_headers(token)
+    headers = get_common_headers(token, evaluator_name)
 
     async with get_async_http_client_with_timeout() as client:
         http_response = await client.post(url, json=payload, headers=headers)
@@ -493,24 +500,26 @@ async def evaluate_with_rai_service(
     credential: TokenCredential,
     annotation_task: str = Tasks.CONTENT_HARM,
     metric_display_name=None,
+    evaluator_name=None,
 ) -> Dict[str, Union[str, float]]:
-    """ "Evaluate the content safety of the response using Responsible AI service
+    """Evaluate the content safety of the response using Responsible AI service
 
-       :param data: The data to evaluate.
-       :type data: dict
-       :param metric_name: The evaluation metric to use.
-       :type metric_name: str
-       :param project_scope: The Azure AI project scope details.
-       :type project_scope: Dict
-       :param credential: The Azure authentication credential.
-       :type credential:
-    ~azure.core.credentials.TokenCredential
-       :param annotation_task: The annotation task to use.
-       :type annotation_task: str
-       :param metric_display_name: The display name of metric to use.
-       :type metric_display_name: str
-       :return: The parsed annotation result.
-       :rtype: Dict[str, Union[str, float]]
+    :param data: The data to evaluate.
+    :type data: dict
+    :param metric_name: The evaluation metric to use.
+    :type metric_name: str
+    :param project_scope: The Azure AI project scope details.
+    :type project_scope: Dict
+    :param credential: The Azure authentication credential.
+    :type credential: ~azure.core.credentials.TokenCredential
+    :param annotation_task: The annotation task to use.
+    :type annotation_task: str
+    :param metric_display_name: The display name of metric to use.
+    :type metric_display_name: str
+    :param evaluator_name: The evaluator name to use.
+    :type evaluator_name: str
+    :return: The parsed annotation result.
+    :rtype: Dict[str, Union[str, float]]
     """
 
     # Get RAI service URL from discovery service and check service availability
@@ -519,7 +528,7 @@ async def evaluate_with_rai_service(
     await ensure_service_availability(rai_svc_url, token, annotation_task)
 
     # Submit annotation request and fetch result
-    operation_id = await submit_request(data, metric_name, rai_svc_url, token, annotation_task)
+    operation_id = await submit_request(data, metric_name, rai_svc_url, token, annotation_task, evaluator_name)
     annotation_response = cast(List[Dict], await fetch_result(operation_id, rai_svc_url, credential, token))
     result = parse_response(annotation_response, metric_name, metric_display_name)