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

azure/ai/evaluation/__init__.py

@@ -40,6 +40,11 @@ from ._model_configurations import (
     Message,
     OpenAIModelConfiguration,
 )
+from ._aoai.aoai_grader import AzureOpenAIGrader
+from ._aoai.label_grader import AzureOpenAILabelGrader
+from ._aoai.string_check_grader import AzureOpenAIStringCheckGrader
+from ._aoai.text_similarity_grader import AzureOpenAITextSimilarityGrader
+
 
 _patch_all = []
 
@@ -52,22 +57,6 @@ try:
 except ImportError:
     print("[INFO] Could not import AIAgentConverter. Please install the dependency with `pip install azure-ai-projects`.")
 
-# RedTeam requires a dependency on pyrit, but python 3.9 is not supported by pyrit.
-# So we only import it if it's available and the user has pyrit.
-try:
-    from ._red_team._red_team import RedTeam
-    from ._red_team._attack_strategy import AttackStrategy
-    from ._red_team._attack_objective_generator import RiskCategory
-    from ._red_team._red_team_result import RedTeamOutput
-    _patch_all.extend([
-        "RedTeam",
-        "RedTeamOutput",
-        "AttackStrategy",
-        "RiskCategory",
-    ])
-except ImportError:
-    print("[INFO] Could not import RedTeam. Please install the dependency with `pip install azure-ai-evaluation[redteam]`.")
-
 
 __all__ = [
     "evaluate",
@@ -105,6 +94,10 @@ __all__ = [
     "CodeVulnerabilityEvaluator",
     "UngroundedAttributesEvaluator",
     "ToolCallAccuracyEvaluator",
+    "AzureOpenAIGrader",
+    "AzureOpenAILabelGrader",
+    "AzureOpenAIStringCheckGrader",
+    "AzureOpenAITextSimilarityGrader",
 ]
 
 __all__.extend([p for p in _patch_all if p not in __all__])

azure/ai/evaluation/_aoai/__init__.py

@@ -0,0 +1,10 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+
+
+from .aoai_grader import AzureOpenAIGrader
+
+__all__ = [
+    "AzureOpenAIGrader",
+]

azure/ai/evaluation/_aoai/aoai_grader.py

@@ -0,0 +1,89 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+from azure.ai.evaluation._model_configurations import AzureOpenAIModelConfiguration, OpenAIModelConfiguration
+
+from azure.ai.evaluation._constants import DEFAULT_AOAI_API_VERSION
+from azure.ai.evaluation._exceptions import ErrorBlame, ErrorCategory, ErrorTarget, EvaluationException
+from typing import Any, Dict, Union
+from azure.ai.evaluation._common._experimental import experimental
+
+
+@experimental
+class AzureOpenAIGrader():
+    """
+    Base class for Azure OpenAI grader wrappers, recommended only for use by experienced OpenAI API users.
+    Combines a model configuration and any grader configuration
+    into a singular object that can be used in evaluations.
+
+    Supplying an AzureOpenAIGrader to the `evaluate` method will cause an asynchronous request to evaluate
+    the grader via the OpenAI API. The results of the evaluation will then be merged into the standard
+    evaluation results.
+
+    :param model_config: The model configuration to use for the grader.
+    :type model_config: Union[
+        ~azure.ai.evaluation.AzureOpenAIModelConfiguration,
+        ~azure.ai.evaluation.OpenAIModelConfiguration
+    ]
+    :param grader_config: The grader configuration to use for the grader. This is expected
+        to be formatted as a dictionary that matches the specifications of the sub-types of
+        the TestingCriterion alias specified in (OpenAI's SDK)[https://github.com/openai/openai-python/blob/ed53107e10e6c86754866b48f8bd862659134ca8/src/openai/types/eval_create_params.py#L151].
+    :type grader_config: Dict[str, Any]
+    :param kwargs: Additional keyword arguments to pass to the grader.
+    :type kwargs: Any
+
+
+    """
+
+    id = "aoai://general"
+
+    def __init__(self, *, model_config : Union[AzureOpenAIModelConfiguration, OpenAIModelConfiguration], grader_config: Dict[str, Any], **kwargs: Any):
+        self._model_config = model_config
+        self._grader_config = grader_config
+
+        if kwargs.get("validate", True):
+            self._validate_model_config()
+            self._validate_grader_config()
+
+
+
+    def _validate_model_config(self) -> None:
+        """Validate the model configuration that this grader wrapper is using."""
+        if "api_key" not in self._model_config or not self._model_config.get("api_key"):
+            msg = f"{type(self).__name__}: Requires an api_key in the supplied model_config."
+            raise EvaluationException(
+                message=msg,
+                blame=ErrorBlame.USER_ERROR,
+                category=ErrorCategory.INVALID_VALUE,
+                target=ErrorTarget.AOAI_GRADER,
+            )
+    
+    def _validate_grader_config(self) -> None:
+        """Validate the grader configuration that this grader wrapper is using."""
+
+        return
+
+    def get_client(self) -> Any:
+        """Construct an appropriate OpenAI client using this grader's model configuration.
+        Returns a slightly different client depending on whether or not this grader's model
+        configuration is for Azure OpenAI or OpenAI.
+
+        :return: The OpenAI client.
+        :rtype: [~openai.OpenAI, ~openai.AzureOpenAI]
+        """
+        if "azure_endpoint" in self._model_config:
+           from openai import AzureOpenAI
+           # TODO set default values?
+           return AzureOpenAI(
+                azure_endpoint=self._model_config["azure_endpoint"],
+                api_key=self._model_config.get("api_key", None), # Default-style access to appease linters.
+                api_version=self._model_config.get("api_version", DEFAULT_AOAI_API_VERSION),
+                azure_deployment=self._model_config.get("azure_deployment", ""),
+            )
+        from openai import OpenAI
+        # TODO add default values for base_url and organization?
+        return OpenAI(
+            api_key=self._model_config["api_key"],
+            base_url=self._model_config.get("base_url", ""),
+            organization=self._model_config.get("organization", ""),
+        )

azure/ai/evaluation/_aoai/label_grader.py

@@ -0,0 +1,66 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+from typing import Any, Dict, Union, List
+
+from azure.ai.evaluation._model_configurations import AzureOpenAIModelConfiguration, OpenAIModelConfiguration
+from openai.types.eval_create_params import TestingCriterionLabelModel
+from azure.ai.evaluation._common._experimental import experimental
+
+from .aoai_grader import AzureOpenAIGrader
+
+@experimental
+class AzureOpenAILabelGrader(AzureOpenAIGrader):
+    """
+    Wrapper class for OpenAI's label model graders.
+
+    Supplying a LabelGrader to the `evaluate` method will cause an asynchronous request to evaluate
+    the grader via the OpenAI API. The results of the evaluation will then be merged into the standard
+    evaluation results.
+
+    :param model_config: The model configuration to use for the grader.
+    :type model_config: Union[
+        ~azure.ai.evaluation.AzureOpenAIModelConfiguration,
+        ~azure.ai.evaluation.OpenAIModelConfiguration
+    ]
+    :param input: The list of label-based testing criterion for this grader. Individual
+        values of this list are expected to be dictionaries that match the format of any of the valid
+        (TestingCriterionLabelModelInput)[https://github.com/openai/openai-python/blob/ed53107e10e6c86754866b48f8bd862659134ca8/src/openai/types/eval_create_params.py#L125C1-L125C32]
+        subtypes.
+    :type input: List[Dict[str, str]]
+    :param labels: A list of strings representing the classification labels of this grader.
+    :type labels: List[str]
+    :param model: The model to use for the evaluation. Must support structured outputs.
+    :type model: str
+    :param name: The name of the grader.
+    :type name: str
+    :param passing_labels: The labels that indicate a passing result. Must be a subset of labels.
+    :type passing_labels: List[str]
+    :param kwargs: Additional keyword arguments to pass to the grader.
+    :type kwargs: Any
+
+
+    """
+
+    id = "aoai://label_model"
+
+    def __init__(
+        self,
+        *,
+        model_config : Union[AzureOpenAIModelConfiguration, OpenAIModelConfiguration],
+        input: List[Dict[str, str]],
+        labels: List[str],
+        model: str,
+        name: str,
+        passing_labels: List[str],
+        **kwargs: Any
+    ):
+        grader = TestingCriterionLabelModel(
+            input=input,
+            labels=labels,
+            model=model,
+            name=name,
+            passing_labels=passing_labels,
+            type="label_model",
+        )
+        super().__init__(model_config=model_config, grader_config=grader, **kwargs)

azure/ai/evaluation/_aoai/string_check_grader.py

@@ -0,0 +1,65 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+from typing import Any, Dict, Union
+from typing_extensions import Literal
+
+from azure.ai.evaluation._model_configurations import AzureOpenAIModelConfiguration, OpenAIModelConfiguration
+from openai.types.eval_string_check_grader import EvalStringCheckGrader
+from azure.ai.evaluation._common._experimental import experimental
+
+from .aoai_grader import AzureOpenAIGrader
+
+@experimental
+class AzureOpenAIStringCheckGrader(AzureOpenAIGrader):
+    """
+    Wrapper class for OpenAI's string check graders.
+
+    Supplying a StringCheckGrader to the `evaluate` method will cause an asynchronous request to evaluate
+    the grader via the OpenAI API. The results of the evaluation will then be merged into the standard
+    evaluation results.
+
+    :param model_config: The model configuration to use for the grader.
+    :type model_config: Union[
+        ~azure.ai.evaluation.AzureOpenAIModelConfiguration,
+        ~azure.ai.evaluation.OpenAIModelConfiguration
+    ]
+    :param input: The input text. This may include template strings.
+    :type input: str
+    :param name: The name of the grader.
+    :type name: str
+    :param operation: The string check operation to perform. One of `eq`, `ne`, `like`, or `ilike`.
+    :type operation: Literal["eq", "ne", "like", "ilike"]
+    :param reference: The reference text. This may include template strings.
+    :type reference: str
+    :param kwargs: Additional keyword arguments to pass to the grader.
+    :type kwargs: Any
+
+
+    """
+
+    id = "aoai://string_check"
+
+    def __init__(
+        self,
+        *,
+        model_config : Union[AzureOpenAIModelConfiguration, OpenAIModelConfiguration],
+        input: str,
+        name: str,
+        operation: Literal[
+            "eq",
+            "ne",
+            "like",
+            "ilike",
+        ],
+        reference: str,
+        **kwargs: Any
+    ):
+        grader = EvalStringCheckGrader(
+            input=input,
+            name=name,
+            operation=operation,
+            reference=reference,
+            type="string_check",
+        )
+        super().__init__(model_config=model_config, grader_config=grader, **kwargs)

azure/ai/evaluation/_aoai/text_similarity_grader.py

@@ -0,0 +1,88 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+from typing import Any, Dict, Union
+from typing_extensions import Literal
+
+from azure.ai.evaluation._model_configurations import AzureOpenAIModelConfiguration, OpenAIModelConfiguration
+from openai.types.eval_text_similarity_grader import EvalTextSimilarityGrader
+from azure.ai.evaluation._common._experimental import experimental
+
+from .aoai_grader import AzureOpenAIGrader
+
+@experimental
+class AzureOpenAITextSimilarityGrader(AzureOpenAIGrader):
+    """
+    Wrapper class for OpenAI's string check graders.
+
+    Supplying a StringCheckGrader to the `evaluate` method will cause an asynchronous request to evaluate
+    the grader via the OpenAI API. The results of the evaluation will then be merged into the standard
+    evaluation results.
+
+    :param model_config: The model configuration to use for the grader.
+    :type model_config: Union[
+        ~azure.ai.evaluation.AzureOpenAIModelConfiguration,
+        ~azure.ai.evaluation.OpenAIModelConfiguration
+    ]
+    :param evaluation_metric: The evaluation metric to use.
+    :type evaluation_metric: Literal[
+            "fuzzy_match",
+            "bleu",
+            "gleu",
+            "meteor",
+            "rouge_1",
+            "rouge_2",
+            "rouge_3",
+            "rouge_4",
+            "rouge_5",
+            "rouge_l",
+            "cosine",
+        ]
+    :param input: The text being graded.
+    :type input: str
+    :param pass_threshold: A float score where a value greater than or equal indicates a passing grade.
+    :type pass_threshold: float
+    :param reference: The text being graded against.
+    :type reference: str
+    :param name: The name of the grader.
+    :type name: str
+    :param kwargs: Additional keyword arguments to pass to the grader.
+    :type kwargs: Any
+
+
+    """
+
+    id = "aoai://text_similarity"
+
+    def __init__(
+        self,
+        *,
+        model_config : Union[AzureOpenAIModelConfiguration, OpenAIModelConfiguration],
+        evaluation_metric: Literal[
+            "fuzzy_match",
+            "bleu",
+            "gleu",
+            "meteor",
+            "rouge_1",
+            "rouge_2",
+            "rouge_3",
+            "rouge_4",
+            "rouge_5",
+            "rouge_l",
+            "cosine",
+        ],
+        input: str,
+        pass_threshold: float,
+        reference: str,
+        name: str,
+        **kwargs: Any
+    ):
+        grader = EvalTextSimilarityGrader(
+            evaluation_metric=evaluation_metric,
+            input=input,
+            pass_threshold=pass_threshold,
+            name=name,
+            reference=reference,
+            type="text_similarity",
+        )
+        super().__init__(model_config=model_config, grader_config=grader, **kwargs)

azure/ai/evaluation/_azure/_clients.py

@@ -8,12 +8,12 @@ from threading import Lock
 from urllib.parse import quote
 from json.decoder import JSONDecodeError
 
-from azure.core.credentials import TokenCredential, AzureSasCredential
+from azure.core.credentials import TokenCredential, AzureSasCredential, AccessToken
 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 azure.ai.evaluation._constants import TokenScope
 from ._models import BlobStoreInfo, Workspace
 
 
@@ -61,7 +61,7 @@ class LiteMLClient:
         self._token_manager: Optional[AzureMLTokenManager] = None
         self._credential: Optional[TokenCredential] = credential
 
-    def get_token(self) -> str:
+    def get_token(self) -> AccessToken:
         return self._get_token_manager().get_token()
 
     def get_credential(self) -> TokenCredential:
@@ -201,4 +201,4 @@ class LiteMLClient:
         return url
 
     def _get_headers(self) -> Dict[str, str]:
-        return {"Authorization": f"Bearer {self.get_token()}", "Content-Type": "application/json"}
+        return {"Authorization": f"Bearer {self.get_token().token}", "Content-Type": "application/json"}

azure/ai/evaluation/_azure/_envs.py

@@ -0,0 +1,208 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+
+# NOTE:
+# This is a simplified version of the original code from azure-ai-ml:
+# sdk\ml\azure-ai-ml\azure\ai\ml\_azure_environments.py
+
+import asyncio
+import os
+
+from typing import Any, Dict, Final, Mapping, Optional, Sequence, TypedDict
+
+from azure.core import AsyncPipelineClient
+from azure.core.configuration import Configuration
+from azure.core.rest import HttpRequest
+from azure.core.pipeline.policies import ProxyPolicy, AsyncRetryPolicy
+
+
+class AzureEnvironmentMetadata(TypedDict):
+    """Configuration for various Azure environments. All endpoints include a trailing slash."""
+    portal_endpoint: str
+    """The management portal for the Azure environment (e.g. https://portal.azure.com/)"""
+    resource_manager_endpoint: str
+    """The API endpoint for Azure control plan (e.g. https://management.azure.com/)"""
+    active_directory_endpoint: str
+    """The active directory endpoint used for authentication (e.g. https://login.microsoftonline.com/)"""
+    aml_resource_endpoint: str
+    """The endpoint for Azure Machine Learning resources (e.g. https://ml.azure.com/)"""
+    storage_suffix: str
+    """The suffix to use for storage endpoint URLs (e.g. core.windows.net)"""
+    registry_discovery_endpoint: str
+
+
+_ENV_ARM_CLOUD_METADATA_URL: Final[str] = "ARM_CLOUD_METADATA_URL"
+_ENV_DEFAULT_CLOUD_NAME: Final[str] = "AZUREML_CURRENT_CLOUD"
+_ENV_REGISTRY_DISCOVERY_URL: Final[str] = "REGISTRY_DISCOVERY_ENDPOINT_URL"
+_ENV_REGISTRY_DISCOVERY_REGION: Final[str] = "REGISTRY_DISCOVERY_ENDPOINT_REGION"
+_DEFAULT_REGISTRY_DISCOVERY_REGION: Final[str] = "west"
+_DEFAULT_AZURE_ENV_NAME: Final[str] = "AzureCloud"
+
+
+_ASYNC_LOCK = asyncio.Lock()
+_KNOWN_AZURE_ENVIRONMENTS: Dict[str, AzureEnvironmentMetadata] = {
+    _DEFAULT_AZURE_ENV_NAME: {
+        "portal_endpoint": "https://portal.azure.com/",
+        "resource_manager_endpoint": "https://management.azure.com/",
+        "active_directory_endpoint": "https://login.microsoftonline.com/",
+        "aml_resource_endpoint": "https://ml.azure.com/",
+        "storage_suffix": "core.windows.net",
+        "registry_discovery_endpoint": "https://eastus.api.azureml.ms/",
+    },
+    "AzureChinaCloud": {
+        "portal_endpoint": "https://portal.azure.cn/",
+        "resource_manager_endpoint": "https://management.chinacloudapi.cn/",
+        "active_directory_endpoint": "https://login.chinacloudapi.cn/",
+        "aml_resource_endpoint": "https://ml.azure.cn/",
+        "storage_suffix": "core.chinacloudapi.cn",
+        "registry_discovery_endpoint": "https://chinaeast2.api.ml.azure.cn/",
+    },
+    "AzureUSGovernment": {
+        "portal_endpoint": "https://portal.azure.us/",
+        "resource_manager_endpoint": "https://management.usgovcloudapi.net/",
+        "active_directory_endpoint": "https://login.microsoftonline.us/",
+        "aml_resource_endpoint": "https://ml.azure.us/",
+        "storage_suffix": "core.usgovcloudapi.net",
+        "registry_discovery_endpoint": "https://usgovarizona.api.ml.azure.us/",
+    },
+}
+
+
+class AzureEnvironmentClient:
+    DEFAULT_API_VERSION: Final[str] = "2019-05-01"
+    DEFAULT_AZURE_CLOUD_NAME: Final[str] = _DEFAULT_AZURE_ENV_NAME
+
+    def __init__(self, *, base_url: Optional[str] = None, **kwargs: Any) -> None:
+        base_url = base_url if base_url is not None else AzureEnvironmentClient.get_default_metadata_url()
+
+        config: Configuration = kwargs.pop("config", Configuration(**kwargs))
+        if config.retry_policy is None:
+            config.retry_policy = AsyncRetryPolicy(**kwargs)
+        if config.proxy_policy is None and "proxy" in kwargs:
+            config.proxy_policy = ProxyPolicy(proxies={"http": kwargs["proxy"], "https": kwargs["proxy"]})
+
+        self._async_client = AsyncPipelineClient(base_url, config=config, **kwargs)
+
+    async def get_default_cloud_name_async(self, *, update_cached: bool = True) -> str:
+        current_cloud_env = os.getenv(_ENV_DEFAULT_CLOUD_NAME)
+        if current_cloud_env is not None:
+            return current_cloud_env
+
+        arm_metadata_url = os.getenv(_ENV_ARM_CLOUD_METADATA_URL)
+        if arm_metadata_url is None:
+            return _DEFAULT_AZURE_ENV_NAME
+
+        # load clouds from metadata url
+        clouds = await self.get_clouds_async(metadata_url=arm_metadata_url, update_cached=update_cached)
+        matched = next(filter(lambda t: t[1]["resource_manager_endpoint"] in arm_metadata_url, clouds.items()), None)
+        if matched is None:
+            return _DEFAULT_AZURE_ENV_NAME
+
+        os.environ[_ENV_DEFAULT_CLOUD_NAME] = matched[0]
+        return matched[0]
+
+    async def get_cloud_async(self, name: str, *, update_cached: bool = True) -> Optional[AzureEnvironmentMetadata]:
+        default_endpoint: Optional[str]
+
+        def case_insensitive_match(d: Mapping[str, Any], key: str) -> Optional[Any]:
+            key = key.strip().lower()
+            return next((v for k,v in d.items() if k.strip().lower() == key), None)
+
+        async with _ASYNC_LOCK:
+            cloud = _KNOWN_AZURE_ENVIRONMENTS.get(name) or case_insensitive_match(_KNOWN_AZURE_ENVIRONMENTS, name)
+            if cloud:
+                return cloud
+            default_endpoint = (_KNOWN_AZURE_ENVIRONMENTS
+                .get(_DEFAULT_AZURE_ENV_NAME, {})
+                .get("resource_manager_endpoint"))
+
+        metadata_url = self.get_default_metadata_url(default_endpoint)
+        clouds = await self.get_clouds_async(metadata_url=metadata_url, update_cached=update_cached)
+        cloud_metadata = clouds.get(name) or case_insensitive_match(clouds, name)
+
+        return cloud_metadata
+
+    async def get_clouds_async(
+        self,
+        *,
+        metadata_url: Optional[str] = None,
+        update_cached: bool = True
+    ) -> Mapping[str, AzureEnvironmentMetadata]:
+        metadata_url = metadata_url or self.get_default_metadata_url()
+
+        clouds: Mapping[str, AzureEnvironmentMetadata]
+        async with self._async_client.send_request(HttpRequest("GET", metadata_url)) as response:  # type: ignore
+            response.raise_for_status()
+            clouds = await self._parse_cloud_endpoints_async(response.json())
+
+        if update_cached:
+            async with _ASYNC_LOCK:
+                recursive_update(_KNOWN_AZURE_ENVIRONMENTS, clouds)
+        return clouds
+
+    async def close(self) -> None:
+        await self._async_client.close()
+
+    @staticmethod
+    def get_default_metadata_url(default_endpoint: Optional[str] = None) -> str:
+        default_endpoint = default_endpoint or "https://management.azure.com/"
+        metadata_url = os.getenv(
+            _ENV_ARM_CLOUD_METADATA_URL,
+            f"{default_endpoint}metadata/endpoints?api-version={AzureEnvironmentClient.DEFAULT_API_VERSION}")
+        return metadata_url
+
+    @staticmethod
+    async def _get_registry_discovery_url_async(cloud_name: str, cloud_suffix: str) -> str:
+        async with _ASYNC_LOCK:
+            discovery_url = _KNOWN_AZURE_ENVIRONMENTS.get(cloud_name, {}).get("registry_discovery_endpoint")
+            if discovery_url:
+                return discovery_url
+
+        discovery_url = os.getenv(_ENV_REGISTRY_DISCOVERY_URL)
+        if discovery_url is not None:
+            return discovery_url
+
+        region = os.getenv(_ENV_REGISTRY_DISCOVERY_REGION, _DEFAULT_REGISTRY_DISCOVERY_REGION)
+        return f"https://{cloud_name.lower()}{region}.api.ml.azure.{cloud_suffix}/"
+
+    @staticmethod
+    async def _parse_cloud_endpoints_async(data: Any) -> Mapping[str, AzureEnvironmentMetadata]:
+        # If there is only one cloud, you will get a dict, otherwise a list of dicts
+        cloud_data: Sequence[Mapping[str, Any]] = data if not isinstance(data, dict) else [data]
+        clouds: Dict[str, AzureEnvironmentMetadata] = {}
+
+        def append_trailing_slash(url: str) -> str:
+            return url if url.endswith("/") else f"{url}/"
+
+        for cloud in cloud_data:
+            try:
+                name: str = cloud["name"]
+                portal_endpoint: str = cloud["portal"]
+                cloud_suffix = ".".join(portal_endpoint.split(".")[2:]).replace("/", "")
+                discovery_url = await AzureEnvironmentClient._get_registry_discovery_url_async(name, cloud_suffix)
+                clouds[name] = {
+                    "portal_endpoint": append_trailing_slash(portal_endpoint),
+                    "resource_manager_endpoint": append_trailing_slash(cloud["resourceManager"]),
+                    "active_directory_endpoint": append_trailing_slash(cloud["authentication"]["loginEndpoint"]),
+                    "aml_resource_endpoint": append_trailing_slash(f"https://ml.azure.{cloud_suffix}/"),
+                    "storage_suffix": cloud["suffixes"]["storage"],
+                    "registry_discovery_endpoint": append_trailing_slash(discovery_url),
+                }
+            except KeyError:
+                continue
+
+        return clouds
+
+
+def recursive_update(d: Dict, u: Mapping) -> None:
+    """Recursively update a dictionary.
+    
+    :param Dict d: The dictionary to update.
+    :param Mapping u: The mapping to update from.
+    """
+    for k, v in u.items():
+        if isinstance(v, Dict) and k in d:
+            recursive_update(d[k], v)
+        else:
+            d[k] = v

azure/ai/evaluation/_azure/_token_manager.py

@@ -5,11 +5,12 @@ import os
 import logging
 import time
 import inspect
-from typing import cast, Optional, Union
+from typing import cast, Optional, Union, Any
 
 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
 
 
@@ -71,7 +72,8 @@ class AzureMLTokenManager(APITokenManager):
             # Fall back to using the parent implementation
             return super().get_aad_credential()
 
-    def get_token(self) -> str:
+    def get_token(
+            self, scopes = None, claims: Union[str, None] = None, tenant_id: Union[str, None] = None, enable_cae: bool = False, **kwargs: Any) -> AccessToken:
         """Get the API token. If the token is not available or has expired, refresh the token.
 
         :return: API token
@@ -79,12 +81,15 @@ class AzureMLTokenManager(APITokenManager):
         """
         if self._token_needs_update():
             credential = cast(TokenCredential, self.credential)
-            access_token = credential.get_token(self.token_scope)
+            token_scope = self.token_scope
+            if scopes:
+                token_scope = scopes
+            access_token = credential.get_token(token_scope)
             self._update_token(access_token)
 
-        return cast(str, self.token)  # check for none is hidden in the _token_needs_update method
+        return cast(AccessToken, self.token)  # check for none is hidden in the _token_needs_update method
 
-    async def get_token_async(self) -> str:
+    async def get_token_async(self) -> AccessToken:
         """Get the API token asynchronously. If the token is not available or has expired, refresh it.
 
         :return: API token
@@ -99,7 +104,7 @@ class AzureMLTokenManager(APITokenManager):
                 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
+        return cast(AccessToken, self.token)  # check for none is hidden in the _token_needs_update method
 
     def _token_needs_update(self) -> bool:
         current_time = time.time()
@@ -112,7 +117,7 @@ class AzureMLTokenManager(APITokenManager):
         )
 
     def _update_token(self, access_token: AccessToken) -> None:
-        self.token = cast(str, access_token.token)
+        self.token = access_token
         self.token_expiry_time = access_token.expires_on
         self.last_refresh_time = time.time()
         self.logger.info("Refreshed Azure management token.")