Bugfix/token expired (Azure#37763)

* Update task_query_response.prompty

remove required keys

* Update task_simulate.prompty

* Update task_query_response.prompty

* Update task_simulate.prompty

* Add a async method to refresh token

* Update changelog

* Update _identity_manager.py

* Fix lint issue

* adding an abstractmethod to help mypy

* black formatting

* Update sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/simulator/_model_tools/_identity_manager.py

Co-authored-by: kdestin <[email protected]>

* Update sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/simulator/_model_tools/_identity_manager.py

Co-authored-by: kdestin <[email protected]>

* Trying a fix for mypy

---------

Co-authored-by: Nagkumar Arkalgud <[email protected]>
Co-authored-by: kdestin <[email protected]>

Author: 3 people

sdk/evaluation/azure-ai-evaluation/CHANGELOG.md

@@ -10,6 +10,7 @@
 - `credential` is now required to be passed in for all content safety evaluators and `ProtectedMaterialsEvaluator`. `DefaultAzureCredential` will no longer be chosen if a credential is not passed. 
 
 ### Bugs Fixed
+- Adversarial Conversation simulations would fail with `Forbidden`. Added logic to re-fetch token in the exponential retry logic to retrive RAI Service response.
 
 ### Other Changes
 

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

@@ -3,14 +3,15 @@
 # ---------------------------------------------------------
 
 import asyncio
+import inspect
 import logging
 import os
 import time
 from abc import ABC, abstractmethod
 from enum import Enum
 from typing import Optional, Union
 
-from azure.core.credentials import TokenCredential
+from azure.core.credentials import TokenCredential, AccessToken
 from azure.identity import DefaultAzureCredential, ManagedIdentityCredential
 
 AZURE_TOKEN_REFRESH_INTERVAL = 600  # seconds
@@ -87,6 +88,14 @@ def get_token(self) -> str:
         :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
+        """
+
 
 class ManagedIdentityAPITokenManager(APITokenManager):
     """API Token Manager for Azure Managed Identity
@@ -127,6 +136,31 @@ def get_token(self) -> str:
 
         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

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

@@ -172,7 +172,6 @@ async def request_api(
         }
         # add all additional headers
         headers.update(self.additional_headers)  # type: ignore[arg-type]
-
         params = {}
         if self.api_version:
             params["api-version"] = self.api_version
@@ -214,6 +213,12 @@ async def request_api(
         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": USER_AGENT,
+            }
             response = await exp_retry_client.get(  # pylint: disable=too-many-function-args,unexpected-keyword-arg
                 self.result_url, headers=proxy_headers
             )