feat: Enable AVM WAF implementation to align with updated AVM standards (#1898)

Signed-off-by: dependabot[bot] <[email protected]>
Co-authored-by: Roopan-Microsoft <[email protected]>
Co-authored-by: Ajit Padhi <[email protected]>
Co-authored-by: Roopan P M <[email protected]>
Co-authored-by: Ross Smith <[email protected]>
Co-authored-by: gpickett <[email protected]>
Co-authored-by: Francia Riesco <[email protected]>
Co-authored-by: Francia Riesco <[email protected]>
Co-authored-by: Prajwal D C <[email protected]>
Co-authored-by: Harmanpreet-Microsoft <[email protected]>
Co-authored-by: UtkarshMishra-Microsoft <[email protected]>
Co-authored-by: Priyanka-Microsoft <[email protected]>
Co-authored-by: Prasanjeet-Microsoft <[email protected]>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
Co-authored-by: Kiran-Siluveru-Microsoft <[email protected]>
Co-authored-by: Prashant-Microsoft <[email protected]>
Co-authored-by: Rohini-Microsoft <[email protected]>
Co-authored-by: Avijit-Microsoft <[email protected]>
Co-authored-by: RaviKiran-Microsoft <[email protected]>
Co-authored-by: Somesh Joshi <[email protected]>
Co-authored-by: Himanshi Agrawal <[email protected]>
Co-authored-by: pradeepjha-microsoft <[email protected]>
Co-authored-by: Harmanpreet Kaur <[email protected]>
Co-authored-by: Bangarraju-Microsoft <[email protected]>
Co-authored-by: Harsh-Microsoft <[email protected]>
Co-authored-by: Kanchan-Microsoft <[email protected]>
Co-authored-by: Cristopher Coronado <[email protected]>
Co-authored-by: Cristopher Coronado Moreira <[email protected]>
Co-authored-by: Vamshi-Microsoft <[email protected]>
Co-authored-by: Thanusree-Microsoft <[email protected]>
Co-authored-by: Niraj Chaudhari (Persistent Systems Inc) <[email protected]>
Co-authored-by: Rohini-Microsoft <[email protected]>

Author: 29 people

azure.yaml

@@ -1,8 +1,8 @@
 # yaml-language-server: $schema=https://raw.githubusercontent.com/Azure/azure-dev/main/schemas/v1.0/azure.yaml.json
 
 name: chat-with-your-data-solution-accelerator
-metadata:
-  template: [email protected]
+# metadata:
+#   template: [email protected]
 hooks:
   postprovision:
     # run: ./infra/prompt-flow/create-prompt-flow.sh

code/backend/batch/utilities/chat_history/database_factory.py

@@ -25,7 +25,7 @@ def get_conversation_client():
                 f"https://{env_helper.AZURE_COSMOSDB_ACCOUNT}.documents.azure.com:443/"
             )
             credential = (
-                get_azure_credential()
+                get_azure_credential(env_helper.MANAGED_IDENTITY_CLIENT_ID)
                 if not env_helper.AZURE_COSMOSDB_ACCOUNT_KEY
                 else env_helper.AZURE_COSMOSDB_ACCOUNT_KEY
             )

code/backend/batch/utilities/chat_history/postgresdbservice.py

@@ -2,6 +2,7 @@
 import asyncpg
 from datetime import datetime, timezone
 from ..helpers.azure_credential_utils import get_azure_credential
+from ..helpers.env_helper import EnvHelper
 
 from .database_client_base import DatabaseClientBase
 
@@ -13,6 +14,7 @@ class PostgresConversationClient(DatabaseClientBase):
     def __init__(
         self, user: str, host: str, database: str, enable_message_feedback: bool = False
     ):
+        self.env_helper = EnvHelper()
         self.user = user
         self.host = host
         self.database = database
@@ -21,7 +23,7 @@ def __init__(
 
     async def connect(self):
         try:
-            credential = get_azure_credential()
+            credential = get_azure_credential(self.env_helper.MANAGED_IDENTITY_CLIENT_ID)
             token = credential.get_token(
                 "https://ossrdbms-aad.database.windows.net/.default"
             ).token

code/backend/batch/utilities/helpers/azure_blob_storage_client.py

@@ -25,7 +25,7 @@ def create_queue_client():
         return QueueClient(
             account_url=f"https://{env_helper.AZURE_BLOB_ACCOUNT_NAME}.queue.core.windows.net/",
             queue_name=env_helper.DOCUMENT_PROCESSING_QUEUE_NAME,
-            credential=get_azure_credential(),
+            credential=get_azure_credential(env_helper.MANAGED_IDENTITY_CLIENT_ID),
             message_encode_policy=BinaryBase64EncodePolicy(),
         )
 
@@ -56,7 +56,7 @@ def __init__(
         if self.auth_type == "rbac":
             self.account_key = None
             self.blob_service_client = BlobServiceClient(
-                account_url=self.endpoint, credential=get_azure_credential()
+                account_url=self.endpoint, credential=get_azure_credential(env_helper.MANAGED_IDENTITY_CLIENT_ID)
             )
             self.user_delegation_key = self.request_user_delegation_key(
                 blob_service_client=self.blob_service_client

code/backend/batch/utilities/helpers/azure_computer_vision_client.py

@@ -27,6 +27,7 @@ def __init__(self, env_helper: EnvHelper) -> None:
         self.model_version = (
             env_helper.AZURE_COMPUTER_VISION_VECTORIZE_IMAGE_MODEL_VERSION
         )
+        self.managed_identity_client_id = env_helper.MANAGED_IDENTITY_CLIENT_ID
 
     def vectorize_image(self, image_url: str) -> list[float]:
         logger.info(f"Making call to computer vision to vectorize image: {image_url}")
@@ -57,7 +58,7 @@ def __make_request(self, path: str, body) -> Response:
                 headers["Ocp-Apim-Subscription-Key"] = self.key
             else:
                 token_provider = get_bearer_token_provider(
-                    get_azure_credential(), self.__TOKEN_SCOPE
+                    get_azure_credential(self.managed_identity_client_id), self.__TOKEN_SCOPE
                 )
                 headers["Authorization"] = "Bearer " + token_provider()
 

code/backend/batch/utilities/helpers/azure_form_recognizer_helper.py

@@ -19,7 +19,7 @@ def __init__(self) -> None:
         if env_helper.AZURE_AUTH_TYPE == "rbac":
             self.document_analysis_client = DocumentAnalysisClient(
                 endpoint=self.AZURE_FORM_RECOGNIZER_ENDPOINT,
-                credential=get_azure_credential(),
+                credential=get_azure_credential(env_helper.MANAGED_IDENTITY_CLIENT_ID),
                 headers={
                     "x-ms-useragent": "chat-with-your-data-solution-accelerator/1.0.0"
                 },

code/backend/batch/utilities/helpers/azure_postgres_helper.py

@@ -24,7 +24,7 @@ def _create_search_client(self):
             dbname = self.env_helper.POSTGRESQL_DATABASE
 
             # Acquire the access token
-            credential = get_azure_credential()
+            credential = get_azure_credential(self.env_helper.MANAGED_IDENTITY_CLIENT_ID)
             access_token = credential.get_token(
                 "https://ossrdbms-aad.database.windows.net/.default"
             )

code/backend/batch/utilities/helpers/azure_search_helper.py

@@ -49,7 +49,7 @@ def _search_credential(self):
         if self.env_helper.is_auth_type_keys():
             return AzureKeyCredential(self.env_helper.AZURE_SEARCH_KEY)
         else:
-            return get_azure_credential()
+            return get_azure_credential(self.env_helper.MANAGED_IDENTITY_CLIENT_ID)
 
     def _create_search_client(
         self, search_credential: Union[AzureKeyCredential, get_azure_credential]
@@ -285,7 +285,7 @@ def get_conversation_logger(self):
         ]
 
         if self.env_helper.AZURE_AUTH_TYPE == "rbac":
-            credential = get_azure_credential()
+            credential = get_azure_credential(self.env_helper.MANAGED_IDENTITY_CLIENT_ID)
             return AzureSearch(
                 azure_search_endpoint=self.env_helper.AZURE_SEARCH_SERVICE,
                 azure_search_key=None,  # Remove API key

code/backend/batch/utilities/helpers/env_helper.py

@@ -39,6 +39,7 @@ def __load_config(self, **kwargs) -> None:
         # Azure
         self.AZURE_SUBSCRIPTION_ID = os.getenv("AZURE_SUBSCRIPTION_ID", "")
         self.AZURE_RESOURCE_GROUP = os.getenv("AZURE_RESOURCE_GROUP", "")
+        self.MANAGED_IDENTITY_CLIENT_ID = os.getenv("MANAGED_IDENTITY_CLIENT_ID", "")
 
         # Azure Search
         self.AZURE_SEARCH_SERVICE = os.getenv("AZURE_SEARCH_SERVICE", "")
@@ -217,7 +218,7 @@ def __load_config(self, **kwargs) -> None:
         )
 
         self.AZURE_TOKEN_PROVIDER = get_bearer_token_provider(
-            get_azure_credential(), "https://cognitiveservices.azure.com/.default"
+            get_azure_credential(self.MANAGED_IDENTITY_CLIENT_ID), "https://cognitiveservices.azure.com/.default"
         )
         self.ADVANCED_IMAGE_PROCESSING_MAX_IMAGES = self.get_env_var_int(
             "ADVANCED_IMAGE_PROCESSING_MAX_IMAGES", 1
@@ -234,7 +235,7 @@ def __load_config(self, **kwargs) -> None:
         self.AZURE_COMPUTER_VISION_VECTORIZE_IMAGE_MODEL_VERSION = os.getenv(
             "AZURE_COMPUTER_VISION_VECTORIZE_IMAGE_MODEL_VERSION", "2023-04-15"
         )
-        self.FUNCTION_KEY = os.getenv("FUNCTION_KEY", "")
+        self.FUNCTION_KEY = self.secretHelper.get_secret("FUNCTION_KEY")
 
         # Initialize Azure keys based on authentication type and environment settings.
         # When AZURE_AUTH_TYPE is "rbac", azure keys are None or an empty string.
@@ -243,7 +244,6 @@ def __load_config(self, **kwargs) -> None:
             self.AZURE_OPENAI_API_KEY = ""
             self.AZURE_SPEECH_KEY = None
             self.AZURE_COMPUTER_VISION_KEY = None
-            self.FUNCTION_KEY = self.secretHelper.get_secret("FUNCTION_KEY")
         else:
             self.AZURE_SEARCH_KEY = self.secretHelper.get_secret("AZURE_SEARCH_KEY")
             self.AZURE_OPENAI_API_KEY = self.secretHelper.get_secret(
@@ -429,8 +429,11 @@ def __init__(self) -> None:
         self.USE_KEY_VAULT = os.getenv("USE_KEY_VAULT", "").lower() == "true"
         self.secret_client = None
         if self.USE_KEY_VAULT:
+            vault_endpoint = os.environ.get("AZURE_KEY_VAULT_ENDPOINT")
+            if not vault_endpoint:
+                raise ValueError("AZURE_KEY_VAULT_ENDPOINT environment variable is required when USE_KEY_VAULT is true")
             self.secret_client = SecretClient(
-                os.environ.get("AZURE_KEY_VAULT_ENDPOINT"), get_azure_credential()
+                vault_endpoint, get_azure_credential(client_id=os.getenv("MANAGED_IDENTITY_CLIENT_ID", None))
             )
 
     def get_secret(self, secret_name: str) -> str:

code/backend/batch/utilities/helpers/llm_helper.py

@@ -166,7 +166,7 @@ def get_sk_service_settings(self, service: AzureChatCompletion):
     def get_ml_client(self):
         if not hasattr(self, "_ml_client"):
             self._ml_client = MLClient(
-                get_azure_credential(),
+                get_azure_credential(self.env_helper.MANAGED_IDENTITY_CLIENT_ID),
                 self.env_helper.AZURE_SUBSCRIPTION_ID,
                 self.env_helper.AZURE_RESOURCE_GROUP,
                 self.env_helper.AZURE_ML_WORKSPACE_NAME,