Convert from key to AAD auth

Author: singhnitant

src/copilot_flow/copilot.py

@@ -1,36 +1,29 @@
-# ---------------------------------------------------------
-# Copyright (c) Microsoft Corporation. All rights reserved.
-# ---------------------------------------------------------
 import os
 # set environment variables before importing any other code
 from dotenv import load_dotenv
 load_dotenv()
 
 from pathlib import Path
-
 from typing import TypedDict
-
 from openai import AzureOpenAI
 
-from azure.core.credentials import AzureKeyCredential
+from azure.identity import DefaultAzureCredential, get_bearer_token_provider
 from azure.search.documents import SearchClient
 from azure.search.documents.models import VectorizedQuery
 
-from promptflow.tracing import trace
 from promptflow.core import Prompty, AzureOpenAIModelConfiguration
+from promptflow.tracing import trace
 
 class ChatResponse(TypedDict):
     context: dict
     reply: str
 
 @trace
 def get_chat_response(chat_input: str, chat_history: list = []) -> ChatResponse:
-
     model_config = AzureOpenAIModelConfiguration(
-        azure_deployment=os.environ["AZURE_OPENAI_CHAT_DEPLOYMENT"],
-        api_version=os.environ["AZURE_OPENAI_API_VERSION"],
-        azure_endpoint=os.environ["AZURE_OPENAI_ENDPOINT"],
-        api_key=os.environ["AZURE_OPENAI_API_KEY"]
+        azure_deployment=os.getenv("AZURE_OPENAI_CHAT_DEPLOYMENT"),
+        api_version=os.getenv("AZURE_OPENAI_API_VERSION"),
+        azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT")
     )
 
     searchQuery = chat_input
@@ -48,7 +41,6 @@ def get_chat_response(chat_input: str, chat_history: list = []) -> ChatResponse:
         searchQuery = intentPrompty(query=chat_input, chat_history=chat_history)
 
     # retrieve relevant documents and context given chat_history and current user query (chat_input)
-    print(os.environ["AZURE_OPENAI_API_KEY"])
     documents = get_documents(searchQuery, 3)
 
     # send query + document context to chat completion for a response
@@ -72,23 +64,23 @@ def get_chat_response(chat_input: str, chat_history: list = []) -> ChatResponse:
 @trace
 def get_documents(search_query: str, num_docs=3):
 
-    index_name = os.environ["AZUREAI_SEARCH_INDEX_NAME"]
+    index_name = os.getenv("AZUREAI_SEARCH_INDEX_NAME")
 
-    #  retrieve documents relevant to the user's query from Azure AI Search index
+    #  retrieve documents relevant to the user's question from Cognitive Search
     search_client = SearchClient(
-        endpoint=os.environ["AZURE_SEARCH_ENDPOINT"],
-        credential=AzureKeyCredential(os.environ["AZURE_SEARCH_KEY"]),
+        endpoint=os.getenv("AZURE_SEARCH_ENDPOINT"),
+        credential=DefaultAzureCredential(),
         index_name=index_name)
 
     aoai_client = AzureOpenAI(
-        azure_endpoint=os.environ["AZURE_OPENAI_ENDPOINT"],
-        api_key=os.environ["AZURE_OPENAI_API_KEY"],
-        api_version=os.environ["AZURE_OPENAI_API_VERSION"]
+        azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT"),
+        azure_ad_token_provider=get_bearer_token_provider(DefaultAzureCredential(), "https://cognitiveservices.azure.com/.default"),
+        api_version=os.getenv("AZURE_OPENAI_API_VERSION")
     )
 
     # generate a vector embedding of the user's question
     embedding = aoai_client.embeddings.create(input=search_query,
-                                            model=os.environ["AZURE_OPENAI_EMBEDDING_DEPLOYMENT"])
+                                            model=os.getenv("AZURE_OPENAI_EMBEDDING_DEPLOYMENT"))
     embedding_to_query = embedding.data[0].embedding
 
     context = ""
@@ -102,4 +94,4 @@ def get_documents(search_query: str, num_docs=3):
     for result in results:
         context += f"\n>>> From: {result['id']}\n{result['content']}"
 
-    return context
+    return context

src/custom_evaluators/completeness_cot.prompty

@@ -5,8 +5,8 @@ model:
   api: chat
   configuration:
     type: azure_openai
-    azure_deployment: ${env:AZURE_DEPLOYMENT}
-    api_key: ${env:AZURE_OPENAI_API_KEY}
+    azure_deployment: ${env:AZURE_OPENAI_EVALUATION_DEPLOYMENT}
+    api_version: ${env:AZURE_OPENAI_API_VERSION}
     azure_endpoint: ${env:AZURE_OPENAI_ENDPOINT}
   parameters:
     temperature: 0.0

src/custom_evaluators/friendliness.prompty

@@ -5,8 +5,8 @@ model:
   api: chat
   configuration:
     type: azure_openai
-    azure_deployment: gpt-4
-    api_key: ${env:AZURE_OPENAI_API_KEY}
+    azure_deployment: ${env:AZURE_OPENAI_EVALUATION_DEPLOYMENT}
+    api_version: ${env:AZURE_OPENAI_API_VERSION}
     azure_endpoint: ${env:AZURE_OPENAI_ENDPOINT}
   parameters:
     temperature: 0.1

src/deployment/deploy.py

@@ -1,9 +1,12 @@
-from azure.ai.ml.entities import ManagedOnlineEndpoint, ManagedOnlineDeployment, Model, Environment, BuildContext
-
-import os
+import os, uuid
+# set environment variables before importing any other code
 from dotenv import load_dotenv
 load_dotenv()
 
+from azure.ai.ml.entities import ManagedOnlineEndpoint, ManagedOnlineDeployment, Model, Environment, BuildContext
+from azure.identity import DefaultAzureCredential
+from azure.mgmt.authorization import AuthorizationManagementClient
+from azure.mgmt.authorization.models import RoleAssignmentCreateParameters
 
 from helper_functions import get_client, get_ai_studio_url_for_deploy
 
@@ -20,10 +23,11 @@ def deploy_flow(endpoint_name, deployment_name):
             name=endpoint_name,
             properties={
               "enforce_access_to_default_secret_stores": "enabled" # if you want secret injection support
-            }
+            },
+            auth_mode="aad_token" # using aad auth instead of key-based auth
         )
 
-    deployment = ManagedOnlineDeployment( # defaults to key auth_mode
+    deployment = ManagedOnlineDeployment(
         name=deployment_name,
         endpoint_name=endpoint_name,
         model=Model(
@@ -65,9 +69,7 @@ def deploy_flow(endpoint_name, deployment_name):
             # the following is enabled by secret injection
             # make sure your environment variables here match the environment variables your code depends on
             'AZURE_OPENAI_ENDPOINT': os.getenv('AZURE_OPENAI_ENDPOINT'),
-            'AZURE_OPENAI_API_KEY': os.getenv('AZURE_OPENAI_API_KEY'),
             'AZURE_SEARCH_ENDPOINT':  os.getenv('AZURE_SEARCH_ENDPOINT'),
-            'AZURE_SEARCH_KEY':  os.getenv('AZURE_SEARCH_KEY'),
             'AZURE_OPENAI_API_VERSION': os.getenv('AZURE_OPENAI_API_VERSION'),
             'AZURE_OPENAI_CHAT_DEPLOYMENT': os.getenv('AZURE_OPENAI_CHAT_DEPLOYMENT'),
             'AZURE_OPENAI_EVALUATION_DEPLOYMENT': os.getenv('AZURE_OPENAI_EVALUATION_DEPLOYMENT'),
@@ -77,16 +79,62 @@ def deploy_flow(endpoint_name, deployment_name):
     )
 
     # 1. create endpoint
-    client.begin_create_or_update(endpoint).result() # result() means we wait on this to complete
+    endpoint = client.begin_create_or_update(endpoint).result() # result() means we wait on this to complete
 
     # 2. create deployment
-    client.begin_create_or_update(deployment).result()
+    deployment = client.begin_create_or_update(deployment).result()
 
     # 3. update endpoint traffic for the deployment
     endpoint.traffic = {deployment_name: 100} # 100% of traffic
-    client.begin_create_or_update(endpoint).result()
-  
-    output_deployment_details(client, endpoint_name, deployment_name)
+    endpoint = client.begin_create_or_update(endpoint).result()
+
+    # 4. provide endpoint access to Azure Open AI resource
+    create_role_assignment(
+        scope=f"/subscriptions/{os.getenv("AZURE_SUBSCRIPTION_ID")}/resourceGroups/{os.getenv("AZURE_RESOURCE_GROUP")}/providers/Microsoft.CognitiveServices/accounts/{os.getenv("AZURE_OPENAI_CONNECTION_NAME")}",
+        role_name="Cognitive Services OpenAI User",
+        principal_id=endpoint.identity.principal_id)
+
+    # 5. provide endpoint access to Azure AI Search resource
+    create_role_assignment(
+        scope=f"/subscriptions/{os.getenv("AZURE_SUBSCRIPTION_ID")}/resourceGroups/{os.getenv("AZURE_RESOURCE_GROUP")}/providers/Microsoft.Search/searchServices/{os.getenv("AZURE_SEARCH_CONNECTION_NAME")}",
+        role_name="Search Index Data Contributor",
+        principal_id=endpoint.identity.principal_id)
+    
+    output_deployment_details(
+        client=client,
+        endpoint_name=endpoint_name,
+        deployment_name=deployment_name)
+
+def create_role_assignment(scope, role_name, principal_id):
+    
+    # Get credential
+    credential = DefaultAzureCredential()
+
+    # Instantiate the authorization management client
+    auth_client = AuthorizationManagementClient(
+        credential=credential,
+        subscription_id=os.getenv("AZURE_SUBSCRIPTION_ID"))
+    
+    roles = list(auth_client.role_definitions.list(
+        scope,
+        filter="roleName eq '{}'".format(role_name)))
+    
+    assert len(roles) == 1
+    role = roles[0]
+    
+    # Create role assignment properties
+    parameters = RoleAssignmentCreateParameters(
+        role_definition_id=role.id,
+        principal_id=principal_id,
+        principal_type="ServicePrincipal")
+    
+    # Create role assignment
+    role_assignment = auth_client.role_assignments.create(
+        scope=scope,
+        role_assignment_name=uuid.uuid4(),
+        parameters=parameters)
+    
+    return role_assignment
 
 def output_deployment_details(client, endpoint_name, deployment_name) -> str:
     print("\n ~~~Deployment details~~~")
@@ -107,4 +155,4 @@ def output_deployment_details(client, endpoint_name, deployment_name) -> str:
     endpoint_name = args.endpoint_name if args.endpoint_name else f"rag-copilot-endpoint"
     deployment_name = args.deployment_name if args.deployment_name else f"rag-copilot-deployment"
 
-    deploy_flow(endpoint_name, deployment_name)
+    deploy_flow(endpoint_name, deployment_name)

src/deployment/invoke.py

@@ -1,21 +1,19 @@
+import requests
 from helper_functions import get_client
 
 def invoke_deployment(endpoint_name: str, query: str, stream: bool = False):
     client = get_client()
 
-    import requests
-
     if stream:
         accept_header = "text/event-stream"
     else:
         accept_header = "application/json"
 
     scoring_url = client.online_endpoints.get(endpoint_name).scoring_uri
-    primary_key = client.online_endpoints.get_keys(endpoint_name).primary_key
 
     headers = {
         "Content-Type": "application/json",
-        "Authorization": f"Bearer {primary_key}",
+        "Authorization":  f"Bearer {client._credential.get_token('https://ml.azure.com').token}",
         "Accept": accept_header
     }
 

src/evaluation/evaluate.py

@@ -1,12 +1,10 @@
-
 import json
 import pathlib
-
+import os
 # set environment variables before importing any other code
 from dotenv import load_dotenv
 load_dotenv()
 
-import os
 import pandas as pd
 from pprint import pprint
 
@@ -33,12 +31,12 @@ def copilot_qna(*, chat_input, **kwargs):
     }
     return parsedResult
 
-def run_evaluation(name, dataset_path):
+def run_evaluation(eval_name, dataset_path):
 
     model_config = AzureOpenAIModelConfiguration(
-        azure_endpoint=os.environ.get("AZURE_OPENAI_ENDPOINT"),
-        api_key=os.environ.get("AZURE_OPENAI_API_KEY"),
-        azure_deployment=os.environ.get("AZURE_OPENAI_EVALUATION_DEPLOYMENT"),
+        azure_deployment=os.getenv("AZURE_OPENAI_EVALUATION_DEPLOYMENT"),
+        api_version=os.getenv("AZURE_OPENAI_API_VERSION"),
+        azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT")
     )
 
     # Initializing Evaluators
@@ -59,7 +57,7 @@ def run_evaluation(name, dataset_path):
 
     result = evaluate(
         target=copilot_qna,
-        evaluation_name=name,
+        evaluation_name=eval_name,
         data=path,
         evaluators={
             "groundedness": groundedness_eval,
@@ -68,7 +66,6 @@ def run_evaluation(name, dataset_path):
             "coherence": coherence_eval,
             "friendliness":friendliness_eval,
             #"completeness": completeness_eval
-
         },
         evaluator_config={
             "relevance": {"question": "${data.chat_input}"},
@@ -96,7 +93,7 @@ def run_evaluation(name, dataset_path):
     evaluation_name = args.evaluation_name if args.evaluation_name else "test-sdk-copilot"
     dataset_path = args.dataset_path if args.dataset_path else "./evaluation/evaluation_dataset_small.jsonl"
 
-    result, tabular_result = run_evaluation(name=evaluation_name,
+    result, tabular_result = run_evaluation(eval_name=evaluation_name,
                               dataset_path=dataset_path)
 
     pprint("-----Summarized Metrics-----")

src/evaluation/evaluate_completeness.py

@@ -38,9 +38,9 @@ def copilot_qna(*, chat_input, **kwargs):
 def run_evaluation(name, dataset_path, prompty_filename: str):
 
     model_config = AzureOpenAIModelConfiguration(
-        azure_endpoint=os.environ.get("AZURE_OPENAI_ENDPOINT"),
-        api_key=os.environ.get("AZURE_OPENAI_API_KEY"),
-        azure_deployment=os.environ.get("AZURE_OPENAI_EVALUATION_DEPLOYMENT"),
+        azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT"),
+        api_version=os.getenv("AZURE_OPENAI_API_VERSION"),
+        azure_deployment=os.getenv("AZURE_OPENAI_EVALUATION_DEPLOYMENT"),
     )
 
     # Initializing Evaluators

src/helper_functions.py

@@ -1,13 +1,14 @@
-from azure.ai.ml import MLClient
-from azure.identity import DefaultAzureCredential
-
 import os
+# set environment variables before importing any other code
 from dotenv import load_dotenv
 load_dotenv()
 
+from azure.ai.ml import MLClient
+from azure.identity import DefaultAzureCredential
+
 def get_client() -> MLClient:
   # check if env variables are set and initialize client from those
-  client = MLClient(DefaultAzureCredential(), os.environ["AZURE_SUBSCRIPTION_ID"], os.environ["AZURE_RESOURCE_GROUP"], os.environ["AZUREAI_PROJECT_NAME"])
+  client = MLClient(DefaultAzureCredential(), os.getenv("AZURE_SUBSCRIPTION_ID"), os.getenv("AZURE_RESOURCE_GROUP"), os.getenv("AZUREAI_PROJECT_NAME"))
   if client:
     return client