address the comment

microsoft · Jun 21, 2024 · 0dfbbb7 · 0dfbbb7
1 parent 8176173
commit 0dfbbb7
Show file tree

Hide file tree

Showing 9 changed files with 292 additions and 1,355 deletions.
diff --git a/...tflow-evals/promptflow/evals/evaluators/_content_safety/flow/evaluate_with_rai_service.py b/...tflow-evals/promptflow/evals/evaluators/_content_safety/flow/evaluate_with_rai_service.py
@@ -4,6 +4,7 @@
 from typing import List
 from urllib.parse import urlparse
 
+import jwt
 import numpy as np
 import requests
 from azure.core.credentials import TokenCredential
@@ -61,14 +62,14 @@ def submit_request(question: str, answer: str, metric: str, rai_svc_url: str, to
     return operation_id
 
 
-def fetch_result(operation_id: str, rai_svc_url: str, token: str):
+def fetch_result(operation_id: str, rai_svc_url: str, credential: TokenCredential, token: str):
     start = time.time()
     request_count = 0
 
     url = rai_svc_url + "/operations/" + operation_id
-    headers = {"Authorization": f"Bearer {token}", "Content-Type": "application/json"}
-
     while True:
+        token = fetch_token(credential, token)
+        headers = {"Authorization": f"Bearer {token}", "Content-Type": "application/json"}
         response = requests.get(url, headers=headers)
         if response.status_code == 200:
             return response.json()
@@ -185,6 +186,27 @@ def get_rai_svc_url(project_scope: dict, token: str):
     return rai_url
 
 
+def fetch_token(credential: TokenCredential, token: str = None):
+    acquire_new_token = True
+    try:
+        if token:
+            # Decode the token to get its expiration time
+            decoded_token = jwt.decode(token, options={"verify_signature": False})
+            exp_time = decoded_token["exp"]
+            current_time = time.time()
+
+            # Check if the token is near expiry
+            if (exp_time - current_time) >= 300:
+                acquire_new_token = False
+    except Exception:
+        pass
+
+    if acquire_new_token:
+        token = credential.get_token("https://management.azure.com/.default").token
+
+    return token
+
+
 @tool
 def evaluate_with_rai_service(
     question: str, answer: str, metric_name: str, project_scope: dict, credential: TokenCredential
@@ -194,15 +216,15 @@ def evaluate_with_rai_service(
     if credential is None or credential == {}:
         credential = DefaultAzureCredential()
 
-    token = credential.get_token("https://management.azure.com/.default").token
+    token = fetch_token(credential)
 
     # Get RAI service URL from discovery service and check service availability
-    rai_svc_url = get_rai_svc_url(project_scope, token)
-    ensure_service_availability(rai_svc_url, token, Tasks.CONTENT_HARM)
+    rai_svc_url = get_rai_svc_url(project_scope, fetch_token(credential, token))
+    ensure_service_availability(rai_svc_url, fetch_token(credential, token), Tasks.CONTENT_HARM)
 
     # Submit annotation request and fetch result
-    operation_id = submit_request(question, answer, metric_name, rai_svc_url, token)
-    annotation_response = fetch_result(operation_id, rai_svc_url, token)
+    operation_id = submit_request(question, answer, metric_name, rai_svc_url, fetch_token(credential, token))
+    annotation_response = fetch_result(operation_id, rai_svc_url, credential, token)
     result = parse_response(annotation_response, metric_name)
 
     return result
diff --git a/...azure/test_adv_simulator_TestAdvSimulator_test_adv_qa_sim_responds_with_one_response.yaml b/...azure/test_adv_simulator_TestAdvSimulator_test_adv_qa_sim_responds_with_one_response.yaml
@@ -40,7 +40,43 @@ interactions:
       x-content-type-options:
       - nosniff
       x-request-time:
-      - '0.031'
+      - '0.024'
+    status:
+      code: 200
+      message: OK
+- request:
+    body: '[{"ver": 1, "name": "Microsoft.ApplicationInsights.Event", "time": "2024-06-06T23:20:59.838896Z",
+      "sampleRate": 100.0, "iKey": "00000000-0000-0000-0000-000000000000", "tags":
+      {"foo": "bar"}}]'
+    headers:
+      Accept:
+      - application/json
+      Accept-Encoding:
+      - gzip, deflate
+      Connection:
+      - keep-alive
+      Content-Length:
+      - '2172'
+      Content-Type:
+      - application/json
+      User-Agent:
+      - azsdk-python-azuremonitorclient/unknown Python/3.10.14 (Windows-10-10.0.22631-SP0)
+    method: POST
+    uri: https://eastus-8.in.applicationinsights.azure.com/v2.1/track
+  response:
+    body:
+      string: '{"itemsReceived": 2, "itemsAccepted": 2, "appId": null, "errors": []}'
+    headers:
+      content-type:
+      - application/json; charset=utf-8
+      server:
+      - Microsoft-HTTPAPI/2.0
+      strict-transport-security:
+      - max-age=31536000
+      transfer-encoding:
+      - chunked
+      x-content-type-options:
+      - nosniff
     status:
       code: 200
       message: OK
@@ -129558,43 +129594,7 @@ interactions:
       x-content-type-options:
       - nosniff
       x-request-time:
-      - '0.394'
-    status:
-      code: 200
-      message: OK
-- request:
-    body: '[{"ver": 1, "name": "Microsoft.ApplicationInsights.Event", "time": "2024-06-06T23:20:59.838896Z",
-      "sampleRate": 100.0, "iKey": "00000000-0000-0000-0000-000000000000", "tags":
-      {"foo": "bar"}}]'
-    headers:
-      Accept:
-      - application/json
-      Accept-Encoding:
-      - gzip, deflate
-      Connection:
-      - keep-alive
-      Content-Length:
-      - '2174'
-      Content-Type:
-      - application/json
-      User-Agent:
-      - azsdk-python-azuremonitorclient/unknown Python/3.10.14 (Windows-10-10.0.22631-SP0)
-    method: POST
-    uri: https://eastus-8.in.applicationinsights.azure.com/v2.1/track
-  response:
-    body:
-      string: '{"itemsReceived": 2, "itemsAccepted": 2, "appId": null, "errors": []}'
-    headers:
-      content-type:
-      - application/json; charset=utf-8
-      server:
-      - Microsoft-HTTPAPI/2.0
-      strict-transport-security:
-      - max-age=31536000
-      transfer-encoding:
-      - chunked
-      x-content-type-options:
-      - nosniff
+      - '0.057'
     status:
       code: 200
       message: OK

diff --git a/...recordings/azure/test_adv_simulator_TestAdvSimulator_test_adv_sim_init_with_prod_url.yaml b/...recordings/azure/test_adv_simulator_TestAdvSimulator_test_adv_sim_init_with_prod_url.yaml
@@ -40,7 +40,7 @@ interactions:
       x-content-type-options:
       - nosniff
       x-request-time:
-      - '0.023'
+      - '0.033'
     status:
       code: 200
       message: OK

diff --git a/...ulator_TestAdvSimulator_test_adv_summarization_jailbreak_sim_responds_with_responses.yaml b/...ulator_TestAdvSimulator_test_adv_summarization_jailbreak_sim_responds_with_responses.yaml
@@ -40,7 +40,7 @@ interactions:
       x-content-type-options:
       - nosniff
       x-request-time:
-      - '0.026'
+      - '0.025'
     status:
       code: 200
       message: OK
@@ -129558,7 +129558,43 @@ interactions:
       x-content-type-options:
       - nosniff
       x-request-time:
-      - '0.018'
+      - '0.020'
+    status:
+      code: 200
+      message: OK
+- request:
+    body: '[{"ver": 1, "name": "Microsoft.ApplicationInsights.Event", "time": "2024-06-06T23:20:59.838896Z",
+      "sampleRate": 100.0, "iKey": "00000000-0000-0000-0000-000000000000", "tags":
+      {"foo": "bar"}}]'
+    headers:
+      Accept:
+      - application/json
+      Accept-Encoding:
+      - gzip, deflate
+      Connection:
+      - keep-alive
+      Content-Length:
+      - '2237'
+      Content-Type:
+      - application/json
+      User-Agent:
+      - azsdk-python-azuremonitorclient/unknown Python/3.10.14 (Windows-10-10.0.22631-SP0)
+    method: POST
+    uri: https://eastus-8.in.applicationinsights.azure.com/v2.1/track
+  response:
+    body:
+      string: '{"itemsReceived": 2, "itemsAccepted": 2, "appId": null, "errors": []}'
+    headers:
+      content-type:
+      - application/json; charset=utf-8
+      server:
+      - Microsoft-HTTPAPI/2.0
+      strict-transport-security:
+      - max-age=31536000
+      transfer-encoding:
+      - chunked
+      x-content-type-options:
+      - nosniff
     status:
       code: 200
       message: OK
@@ -134344,43 +134380,7 @@ interactions:
       x-content-type-options:
       - nosniff
       x-request-time:
-      - '0.021'
-    status:
-      code: 200
-      message: OK
-- request:
-    body: '[{"ver": 1, "name": "Microsoft.ApplicationInsights.Event", "time": "2024-06-06T23:20:59.838896Z",
-      "sampleRate": 100.0, "iKey": "00000000-0000-0000-0000-000000000000", "tags":
-      {"foo": "bar"}}]'
-    headers:
-      Accept:
-      - application/json
-      Accept-Encoding:
-      - gzip, deflate
-      Connection:
-      - keep-alive
-      Content-Length:
-      - '2237'
-      Content-Type:
-      - application/json
-      User-Agent:
-      - azsdk-python-azuremonitorclient/unknown Python/3.10.14 (Windows-10-10.0.22631-SP0)
-    method: POST
-    uri: https://eastus-8.in.applicationinsights.azure.com/v2.1/track
-  response:
-    body:
-      string: '{"itemsReceived": 2, "itemsAccepted": 2, "appId": null, "errors": []}'
-    headers:
-      content-type:
-      - application/json; charset=utf-8
-      server:
-      - Microsoft-HTTPAPI/2.0
-      strict-transport-security:
-      - max-age=31536000
-      transfer-encoding:
-      - chunked
-      x-content-type-options:
-      - nosniff
+      - '0.019'
     status:
       code: 200
       message: OK

diff --git a/...st_adv_simulator_TestAdvSimulator_test_adv_summarization_sim_responds_with_responses.yaml b/...st_adv_simulator_TestAdvSimulator_test_adv_summarization_sim_responds_with_responses.yaml
@@ -40,7 +40,7 @@ interactions:
       x-content-type-options:
       - nosniff
       x-request-time:
-      - '0.032'
+      - '0.027'
     status:
       code: 200
       message: OK
@@ -129558,7 +129558,7 @@ interactions:
       x-content-type-options:
       - nosniff
       x-request-time:
-      - '0.147'
+      - '0.019'
     status:
       code: 200
       message: OK