[pf-evals] ContentSafety: Use new way to check RAI availability and perf improve via reusing the token

src/promptflow-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
@@ -20,22 +21,33 @@
 USER_AGENT = "{}/{}".format("promptflow-evals", version)
 
 
-def ensure_service_availability(rai_svc_url: str):
-    svc_liveness_url = rai_svc_url.split("/subscriptions")[0] + "/meta/version"
-    response = requests.get(svc_liveness_url)
+def ensure_service_availability(rai_svc_url: str, token: str, capability: str = None):
+    headers = {
+        "Authorization": f"Bearer {token}",
+        "Content-Type": "application/json",
+        "User-Agent": USER_AGENT,
+    }
+
+    svc_liveness_url = rai_svc_url + "/checkannotation"
+    response = requests.get(svc_liveness_url, headers=headers)
+
     if response.status_code != 200:
-        raise Exception("RAI service is not available in this region")
+        raise Exception(f"RAI service is not available in this region. Status Code: {response.status_code}")
+
+    capabilities = response.json()
+
+    if capability and capability not in capabilities:
+        raise Exception(f"Capability '{capability}' is not available in this region")
 
 
-def submit_request(question: str, answer: str, metric: str, rai_svc_url: str, credential: TokenCredential):
+def submit_request(question: str, answer: str, metric: str, rai_svc_url: str, token: str):
     user_text = f"<Human>{question}</><System>{answer}</>"
     normalized_user_text = user_text.replace("'", '\\"')
     payload = {"UserTextList": [normalized_user_text], "AnnotationTask": Tasks.CONTENT_HARM, "MetricList": [metric]}
 
     url = rai_svc_url + "/submitannotation"
-    bearer_token = credential.get_token("https://management.azure.com/.default").token
     headers = {
-        "Authorization": f"Bearer {bearer_token}",
+        "Authorization": f"Bearer {token}",
         "Content-Type": "application/json",
         "User-Agent": USER_AGENT,
     }
@@ -50,15 +62,14 @@ def submit_request(question: str, answer: str, metric: str, rai_svc_url: str, cr
     return operation_id
 
 
-def fetch_result(operation_id: str, rai_svc_url: str, credential: TokenCredential):
+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
-    bearer_token = credential.get_token("https://management.azure.com/.default").token
-    headers = {"Authorization": f"Bearer {bearer_token}", "Content-Type": "application/json"}
-
     while True:
+        token = fetch_or_reuse_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()
@@ -144,9 +155,8 @@ def parse_response(batch_response: List[dict], metric_name: str) -> List[List[di
     return result
 
 
-def _get_service_discovery_url(azure_ai_project, credential):
-    bearer_token = credential.get_token("https://management.azure.com/.default").token
-    headers = {"Authorization": f"Bearer {bearer_token}", "Content-Type": "application/json"}
+def _get_service_discovery_url(azure_ai_project, token):
+    headers = {"Authorization": f"Bearer {token}", "Content-Type": "application/json"}
     response = requests.get(
         f"https://management.azure.com/subscriptions/{azure_ai_project['subscription_id']}/"
         f"resourceGroups/{azure_ai_project['resource_group_name']}/"
@@ -161,8 +171,8 @@ def _get_service_discovery_url(azure_ai_project, credential):
     return f"{base_url.scheme}://{base_url.netloc}"
 
 
-def get_rai_svc_url(project_scope: dict, credential: TokenCredential):
-    discovery_url = _get_service_discovery_url(azure_ai_project=project_scope, credential=credential)
+def get_rai_svc_url(project_scope: dict, token: str):
+    discovery_url = _get_service_discovery_url(azure_ai_project=project_scope, token=token)
     subscription_id = project_scope["subscription_id"]
     resource_group_name = project_scope["resource_group_name"]
     project_name = project_scope["project_name"]
@@ -176,6 +186,27 @@ def get_rai_svc_url(project_scope: dict, credential: TokenCredential):
     return rai_url
 
 
+def fetch_or_reuse_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
@@ -186,12 +217,13 @@ def evaluate_with_rai_service(
         credential = DefaultAzureCredential()
 
     # Get RAI service URL from discovery service and check service availability
-    rai_svc_url = get_rai_svc_url(project_scope, credential)
-    ensure_service_availability(rai_svc_url)
+    token = fetch_or_reuse_token(credential)
+    rai_svc_url = get_rai_svc_url(project_scope, token)
+    ensure_service_availability(rai_svc_url, token, Tasks.CONTENT_HARM)
 
     # Submit annotation request and fetch result
-    operation_id = submit_request(question, answer, metric_name, rai_svc_url, credential)
-    annotation_response = fetch_result(operation_id, rai_svc_url, credential)
+    operation_id = submit_request(question, answer, metric_name, rai_svc_url, token)
+    annotation_response = fetch_result(operation_id, rai_svc_url, credential, token)
     result = parse_response(annotation_response, metric_name)
 
     return result

src/promptflow-evals/tests/evals/conftest.py

@@ -6,14 +6,14 @@
 
 import jwt
 import pytest
+from azure.ai.ml._ml_client import MLClient
 from pytest_mock import MockerFixture
 
 from promptflow.client import PFClient
 from promptflow.core import AzureOpenAIModelConfiguration
 from promptflow.executor._line_execution_process_pool import _process_wrapper
 from promptflow.executor._process_manager import create_spawned_fork_process_manager
 from promptflow.tracing._integrations._openai_injector import inject_openai_api
-from azure.ai.ml._ml_client import MLClient
 
 try:
     from promptflow.recording.local import recording_array_reset

src/promptflow-evals/tests/evals/e2etests/test_builtin_evaluators.py

@@ -42,6 +42,19 @@ def test_individual_evaluator_service_based(self, project_scope, azure_cred):
         assert score["violence_score"] < 1.0
         assert score["violence_reason"], "violence_reason must not be None or empty."
 
+    @pytest.mark.skip(reason="Not working in ci pipeline. For local run.")
+    def test_content_safety_service_unavailable(self, project_scope, azure_cred):
+        eval_fn = ViolenceEvaluator(project_scope, azure_cred)
+        project_scope["project_name"] = "pf-evals-ws-westus2"
+
+        with pytest.raises(Exception) as exc_info:
+            eval_fn(
+                question="What is the capital of Japan?",
+                answer="The capital of Japan is Tokyo.",
+            )
+
+        assert "RAI service is not available in this region" in exc_info._excinfo[1].inner_exception.args[0]
+
     @pytest.mark.parametrize("parallel", [False, True])
     def test_composite_evaluator_qa(self, model_config, parallel):
         qa_eval = QAEvaluator(model_config, parallel=parallel)

src/promptflow-evals/tests/recordings/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.025'
+      - '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,65 +129594,7 @@ interactions:
       x-content-type-options:
       - nosniff
       x-request-time:
-      - '0.164'
-    status:
-      code: 200
-      message: OK
-- request:
-    body: '[{"ver": 1, "name": "Microsoft.ApplicationInsights.Event", "time": "2024-06-04T18:17:47.066535Z",
-      "sampleRate": 100.0, "iKey": "8b52b368-4c91-4226-b7f7-be52822f0509", "tags":
-      {"ai.device.locale": "en_US", "ai.device.osVersion": "10.0.22631", "ai.device.type":
-      "Other", "ai.internal.sdkVersion": "uwm_py3.10.14:otel1.25.0:ext1.0.0b26", "ai.cloud.role":
-      "***.py", "ai.internal.nodeName": "ninhu-desktop2", "ai.operation.id": "00000000000000000000000000000000",
-      "ai.operation.parentId": "0000000000000000"}, "data": {"baseType": "EventData",
-      "baseData": {"ver": 2, "name": "adversarial.simulator.call.start", "properties":
-      {"level": "INFO", "from_ci": "False", "request_id": "43a91461-447d-48c5-8ef7-63353d0162af",
-      "first_call": "True", "activity_name": "adversarial.simulator.call", "activity_type":
-      "PublicApi", "user_agent": "", "scenario": "AdversarialScenario.ADVERSARIAL_QA",
-      "max_conversation_turns": "1", "max_simulation_results": "1", "python_version":
-      "3.10.14", "installation_id": "ca79f281-c213-5434-91e6-88ee00a05a6a"}}}}, {"ver":
-      1, "name": "Microsoft.ApplicationInsights.Event", "time": "2024-06-04T18:17:47.066535Z",
-      "sampleRate": 100.0, "iKey": "8b52b368-4c91-4226-b7f7-be52822f0509", "tags":
-      {"ai.device.locale": "en_US", "ai.device.osVersion": "10.0.22631", "ai.device.type":
-      "Other", "ai.internal.sdkVersion": "uwm_py3.10.14:otel1.25.0:ext1.0.0b26", "ai.cloud.role":
-      "***.py", "ai.internal.nodeName": "ninhu-desktop2", "ai.operation.id": "00000000000000000000000000000000",
-      "ai.operation.parentId": "0000000000000000"}, "data": {"baseType": "EventData",
-      "baseData": {"ver": 2, "name": "adversarial.simulator.call.complete", "properties":
-      {"level": "INFO", "from_ci": "False", "request_id": "43a91461-447d-48c5-8ef7-63353d0162af",
-      "first_call": "True", "activity_name": "adversarial.simulator.call", "activity_type":
-      "PublicApi", "user_agent": "", "scenario": "AdversarialScenario.ADVERSARIAL_QA",
-      "max_conversation_turns": "1", "max_simulation_results": "1", "completion_status":
-      "Success", "duration_ms": "0.0", "python_version": "3.10.14", "installation_id":
-      "ca79f281-c213-5434-91e6-88ee00a05a6a"}}}}]'
-    headers:
-      Accept:
-      - application/json
-      Accept-Encoding:
-      - gzip, deflate
-      Connection:
-      - keep-alive
-      Content-Length:
-      - '2121'
-      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

src/promptflow-evals/tests/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.025'
+      - '0.033'
     status:
       code: 200
       message: OK

src/promptflow-evals/tests/recordings/azure/test_adv_simulator_TestAdvSimulator_test_adv_summarization_jailbreak_sim_responds_with_responses.yaml

@@ -40,7 +40,7 @@ interactions:
       x-content-type-options:
       - nosniff
       x-request-time:
-      - '0.021'
+      - '0.025'
     status:
       code: 200
       message: OK
@@ -129558,36 +129558,14 @@ interactions:
       x-content-type-options:
       - nosniff
       x-request-time:
-      - '0.078'
+      - '0.020'
     status:
       code: 200
       message: OK
 - request:
-    body: '[{"ver": 1, "name": "Microsoft.ApplicationInsights.Event", "time": "2024-06-04T18:21:13.871575Z",
-      "sampleRate": 100.0, "iKey": "8b52b368-4c91-4226-b7f7-be52822f0509", "tags":
-      {"ai.device.locale": "en_US", "ai.device.osVersion": "10.0.22631", "ai.device.type":
-      "Other", "ai.internal.sdkVersion": "uwm_py3.10.14:otel1.25.0:ext1.0.0b26", "ai.cloud.role":
-      "***.py", "ai.internal.nodeName": "ninhu-desktop2", "ai.operation.id": "00000000000000000000000000000000",
-      "ai.operation.parentId": "0000000000000000"}, "data": {"baseType": "EventData",
-      "baseData": {"ver": 2, "name": "adversarial.simulator.call.start", "properties":
-      {"level": "INFO", "from_ci": "False", "request_id": "92065092-b71d-4ea2-a860-044dbedb93f6",
-      "first_call": "True", "activity_name": "adversarial.simulator.call", "activity_type":
-      "PublicApi", "user_agent": "", "scenario": "AdversarialScenario.ADVERSARIAL_SUMMARIZATION",
-      "max_conversation_turns": "1", "max_simulation_results": "1", "jailbreak": "True",
-      "python_version": "3.10.14", "installation_id": "ca79f281-c213-5434-91e6-88ee00a05a6a"}}}},
-      {"ver": 1, "name": "Microsoft.ApplicationInsights.Event", "time": "2024-06-04T18:21:13.871575Z",
-      "sampleRate": 100.0, "iKey": "8b52b368-4c91-4226-b7f7-be52822f0509", "tags":
-      {"ai.device.locale": "en_US", "ai.device.osVersion": "10.0.22631", "ai.device.type":
-      "Other", "ai.internal.sdkVersion": "uwm_py3.10.14:otel1.25.0:ext1.0.0b26", "ai.cloud.role":
-      "***.py", "ai.internal.nodeName": "ninhu-desktop2", "ai.operation.id": "00000000000000000000000000000000",
-      "ai.operation.parentId": "0000000000000000"}, "data": {"baseType": "EventData",
-      "baseData": {"ver": 2, "name": "adversarial.simulator.call.complete", "properties":
-      {"level": "INFO", "from_ci": "False", "request_id": "92065092-b71d-4ea2-a860-044dbedb93f6",
-      "first_call": "True", "activity_name": "adversarial.simulator.call", "activity_type":
-      "PublicApi", "user_agent": "", "scenario": "AdversarialScenario.ADVERSARIAL_SUMMARIZATION",
-      "max_conversation_turns": "1", "max_simulation_results": "1", "jailbreak": "True",
-      "completion_status": "Success", "duration_ms": "0.0", "python_version": "3.10.14",
-      "installation_id": "ca79f281-c213-5434-91e6-88ee00a05a6a"}}}}]'
+    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
@@ -129596,7 +129574,7 @@ interactions:
       Connection:
       - keep-alive
       Content-Length:
-      - '2185'
+      - '2237'
       Content-Type:
       - application/json
       User-Agent:
@@ -134402,7 +134380,7 @@ interactions:
       x-content-type-options:
       - nosniff
       x-request-time:
-      - '0.025'
+      - '0.019'
     status:
       code: 200
       message: OK