fix: added retry logic for SQL DB connection, updated agent and thread management

infra/deploy_app_service.bicep

@@ -184,6 +184,10 @@ resource Website 'Microsoft.Web/sites@2020-06-01' = {
     serverFarmId: HostingPlanName
     siteConfig: {
       appSettings: [
+        {
+          name: 'APP_ENV'
+          value: 'Prod'
+        }
         {
           name: 'APPINSIGHTS_INSTRUMENTATIONKEY'
           value: reference(applicationInsightsId, '2015-05-01').InstrumentationKey

infra/main.bicep

@@ -74,8 +74,7 @@ param aiDeploymentsLocation string
 param AZURE_LOCATION string = ''
 var solutionLocation = empty(AZURE_LOCATION) ? resourceGroup().location : AZURE_LOCATION
 
-var uniqueId = toLower(uniqueString(environmentName, subscription().id, solutionLocation))
-
+var uniqueId = toLower(uniqueString(environmentName, subscription().id, solutionLocation, resourceGroup().name))
 var solutionPrefix = 'ca${padLeft(take(uniqueId, 12), 12, '0')}'
 
 // Load the abbrevations file required to name the azure resources.

infra/main.json

@@ -5,7 +5,7 @@
     "_generator": {
       "name": "bicep",
       "version": "0.36.177.2456",
-      "templateHash": "2238194529646818649"
+      "templateHash": "3253509031453285119"
     }
   },
   "parameters": {
@@ -359,7 +359,7 @@
       }
     },
     "solutionLocation": "[if(empty(parameters('AZURE_LOCATION')), resourceGroup().location, parameters('AZURE_LOCATION'))]",
-    "uniqueId": "[toLower(uniqueString(parameters('environmentName'), subscription().id, variables('solutionLocation')))]",
+    "uniqueId": "[toLower(uniqueString(parameters('environmentName'), subscription().id, variables('solutionLocation'), resourceGroup().name))]",
     "solutionPrefix": "[format('ca{0}', padLeft(take(variables('uniqueId'), 12), 12, '0'))]",
     "abbrs": "[variables('$fxv#0')]",
     "functionAppSqlPrompt": "Generate a valid T-SQL query to find {query} for tables and columns provided below:\r\n   1. Table: Clients\r\n   Columns: ClientId, Client, Email, Occupation, MaritalStatus, Dependents\r\n   2. Table: InvestmentGoals\r\n   Columns: ClientId, InvestmentGoal\r\n   3. Table: Assets\r\n   Columns: ClientId, AssetDate, Investment, ROI, Revenue, AssetType\r\n   4. Table: ClientSummaries\r\n   Columns: ClientId, ClientSummary\r\n   5. Table: InvestmentGoalsDetails\r\n   Columns: ClientId, InvestmentGoal, TargetAmount, Contribution\r\n   6. Table: Retirement\r\n   Columns: ClientId, StatusDate, RetirementGoalProgress, EducationGoalProgress\r\n   7. Table: ClientMeetings\r\n   Columns: ClientId, ConversationId, Title, StartTime, EndTime, Advisor, ClientEmail\r\n   Always use the Investment column from the Assets table as the value.\r\n   Assets table has snapshots of values by date. Do not add numbers across different dates for total values.\r\n   Do not use client name in filters.\r\n   Do not include assets values unless asked for.\r\n   ALWAYS use ClientId = {clientid} in the query filter.\r\n   ALWAYS select Client Name (Column: Client) in the query.\r\n   Query filters are IMPORTANT. Add filters like AssetType, AssetDate, etc. if needed.\r\n   When answering scheduling or time-based meeting questions, always use the StartTime column from ClientMeetings table. Use correct logic to return the most recent past meeting (last/previous) or the nearest future meeting (next/upcoming), and ensure only StartTime column is used for meeting timing comparisons.\r\n   Only return the generated SQL query. Do not return anything else.",
@@ -744,7 +744,7 @@
             "_generator": {
               "name": "bicep",
               "version": "0.36.177.2456",
-              "templateHash": "1124249040831466979"
+              "templateHash": "1343961496887433815"
             }
           },
           "parameters": {
@@ -1458,7 +1458,7 @@
                     "_generator": {
                       "name": "bicep",
                       "version": "0.36.177.2456",
-                      "templateHash": "11899270249637077405"
+                      "templateHash": "10199364008784095733"
                     }
                   },
                   "parameters": {
@@ -2244,7 +2244,7 @@
             "_generator": {
               "name": "bicep",
               "version": "0.36.177.2456",
-              "templateHash": "10507186896960913919"
+              "templateHash": "7899619253922538038"
             }
           },
           "parameters": {
@@ -2615,6 +2615,10 @@
                 "serverFarmId": "[parameters('HostingPlanName')]",
                 "siteConfig": {
                   "appSettings": [
+                    {
+                      "name": "APP_ENV",
+                      "value": "Prod"
+                    },
                     {
                       "name": "APPINSIGHTS_INSTRUMENTATIONKEY",
                       "value": "[reference(parameters('applicationInsightsId'), '2015-05-01').InstrumentationKey]"
@@ -2894,7 +2898,7 @@
                     "_generator": {
                       "name": "bicep",
                       "version": "0.36.177.2456",
-                      "templateHash": "11899270249637077405"
+                      "templateHash": "10199364008784095733"
                     }
                   },
                   "parameters": {

infra/scripts/fabric_scripts/create_fabric_items.py

@@ -1,4 +1,3 @@
-from azure.identity import DefaultAzureCredential
 import base64
 import json
 import requests
@@ -8,7 +7,6 @@
 import time
 
 
-# credential = DefaultAzureCredential()
 from azure.identity import AzureCliCredential
 credential = AzureCliCredential()
 

infra/scripts/index_scripts/create_search_index.py

@@ -5,7 +5,8 @@
 import time
 
 import pandas as pd
-from azure.identity import DefaultAzureCredential, get_bearer_token_provider
+from azure.identity import get_bearer_token_provider
+from azure.identity import AzureCliCredential
 from azure.keyvault.secrets import SecretClient
 from azure.search.documents import SearchClient
 from azure.search.documents.indexes import SearchIndexClient
@@ -44,9 +45,7 @@
     "clienttranscripts/meeting_transcripts_metadata/transcripts_metadata.csv"
 )
 
-credential = DefaultAzureCredential(
-    managed_identity_client_id=managed_identity_client_id
-)
+credential = AzureCliCredential()
 token_provider = get_bearer_token_provider(
     credential,
     "https://cognitiveservices.azure.com/.default"

infra/scripts/index_scripts/create_sql_tables.py

@@ -7,15 +7,13 @@
 
 import pandas as pd
 import pyodbc
-from azure.identity import DefaultAzureCredential
+from azure.identity import AzureCliCredential
 from azure.keyvault.secrets import SecretClient
 
 
 def get_secrets_from_kv(kv_name, secret_name):
     key_vault_name = kv_name  # Set the name of the Azure Key Vault
-    credential = DefaultAzureCredential(
-        managed_identity_client_id=managed_identity_client_id
-    )
+    credential = AzureCliCredential()  # Use Azure CLI Credential
     secret_client = SecretClient(
         vault_url=f"https://{key_vault_name}.vault.azure.net/", credential=credential
     )  # Create a secret client object using the credential and Key Vault name
@@ -27,9 +25,7 @@ def get_secrets_from_kv(kv_name, secret_name):
 driver = "{ODBC Driver 18 for SQL Server}"
 
 
-credential = DefaultAzureCredential(
-    managed_identity_client_id=managed_identity_client_id
-)
+credential = AzureCliCredential()  # Use Azure CLI Credential
 
 token_bytes = credential.get_token(
     "https://database.windows.net/.default"
@@ -49,9 +45,7 @@ def get_secrets_from_kv(kv_name, secret_name):
 from azure.storage.filedatalake import DataLakeServiceClient
 
 account_name = get_secrets_from_kv(key_vault_name, "ADLS-ACCOUNT-NAME")
-credential = DefaultAzureCredential(
-    managed_identity_client_id=managed_identity_client_id
-)
+credential = AzureCliCredential()  # Use Azure CLI Credential
 
 account_url = f"https://{account_name}.dfs.core.windows.net"
 

infra/scripts/index_scripts/create_update_sql_dates.py

@@ -6,15 +6,13 @@
 
 import pandas as pd
 import pymssql
-from azure.identity import DefaultAzureCredential
+from azure.identity import AzureCliCredential
 from azure.keyvault.secrets import SecretClient
 
 
 def get_secrets_from_kv(kv_name, secret_name):
     key_vault_name = kv_name  # Set the name of the Azure Key Vault
-    credential = DefaultAzureCredential(
-        managed_identity_client_id=managed_identity_client_id
-    )
+    credential = AzureCliCredential()  # Use Azure CLI Credential
     secret_client = SecretClient(
         vault_url=f"https://{key_vault_name}.vault.azure.net/", credential=credential
     )  # Create a secret client object using the credential and Key Vault name
@@ -32,9 +30,7 @@ def get_secrets_from_kv(kv_name, secret_name):
 from azure.storage.filedatalake import DataLakeServiceClient
 
 account_name = get_secrets_from_kv(key_vault_name, "ADLS-ACCOUNT-NAME")
-credential = DefaultAzureCredential(
-    managed_identity_client_id=managed_identity_client_id
-)
+credential = AzureCliCredential()  # Use Azure CLI Credential
 
 account_url = f"https://{account_name}.dfs.core.windows.net"
 

src/App/.env.sample

@@ -67,4 +67,5 @@ AZURE_SQL_SYSTEM_PROMPT="Generate a valid T-SQL query to find {query} for tables
 # Misc
 APPINSIGHTS_INSTRUMENTATIONKEY=
 AUTH_ENABLED="false"
-USE_INTERNAL_STREAM="True"
+USE_INTERNAL_STREAM="True"
+APP_ENV="dev"

src/App/WebApp.Dockerfile

@@ -36,4 +36,4 @@ COPY --from=frontend /home/node/app/static  /usr/src/app/static/
 WORKDIR /usr/src/app  
 EXPOSE 80  
 
-CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "80", "--workers", "4", "--log-level", "info", "--access-log"]
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "80", "--workers", "1", "--log-level", "info", "--access-log"]

src/App/app.py

@@ -6,7 +6,8 @@
 import uuid
 from types import SimpleNamespace
 
-from azure.identity import DefaultAzureCredential, get_bearer_token_provider
+from azure.identity import get_bearer_token_provider
+from backend.helpers.azure_credential_utils import get_azure_credential
 from azure.monitor.opentelemetry import configure_azure_monitor
 
 # from quart.sessions import SecureCookieSessionInterface
@@ -90,6 +91,7 @@ async def shutdown():
         await AgentFactory.delete_all_agent_instance()
         app.wealth_advisor_agent = None
         app.search_agent = None
+        app.sql_agent = None
         logging.info("Agents cleaned up during application shutdown")
 
     # app.secret_key = secrets.token_hex(16)
@@ -185,7 +187,7 @@ def init_openai_client(use_data=SHOULD_USE_DATA):
         if not aoai_api_key:
             logging.debug("No AZURE_OPENAI_KEY found, using Azure AD auth")
             ad_token_provider = get_bearer_token_provider(
-                DefaultAzureCredential(), "https://cognitiveservices.azure.com/.default"
+                get_azure_credential(), "https://cognitiveservices.azure.com/.default"
             )
 
         # Deployment
@@ -233,7 +235,7 @@ def init_cosmosdb_client():
             )
 
             if not config.AZURE_COSMOSDB_ACCOUNT_KEY:
-                credential = DefaultAzureCredential()
+                credential = get_azure_credential()
             else:
                 credential = config.AZURE_COSMOSDB_ACCOUNT_KEY
 

src/App/backend/agents/agent_factory.py

@@ -10,8 +10,8 @@
 from typing import Optional
 
 from azure.ai.projects import AIProjectClient
-from azure.identity import DefaultAzureCredential as DefaultAzureCredentialSync
-from azure.identity.aio import DefaultAzureCredential
+from backend.helpers.azure_credential_utils import get_azure_credential
+from backend.helpers.azure_credential_utils import get_azure_credential_async
 from semantic_kernel.agents import AzureAIAgent, AzureAIAgentSettings
 
 from backend.common.config import config
@@ -35,7 +35,7 @@ async def get_wealth_advisor_agent(cls):
         async with cls._lock:
             if cls._wealth_advisor_agent is None:
                 ai_agent_settings = AzureAIAgentSettings()
-                creds = DefaultAzureCredential()
+                creds = await get_azure_credential_async()
                 client = AzureAIAgent.create_client(
                     credential=creds, endpoint=ai_agent_settings.endpoint
                 )
@@ -76,7 +76,7 @@ async def get_search_agent(cls):
 
                 project_client = AIProjectClient(
                     endpoint=config.AI_PROJECT_ENDPOINT,
-                    credential=DefaultAzureCredentialSync(),
+                    credential=get_azure_credential(),
                     api_version="2025-05-01",
                 )
 
@@ -137,7 +137,7 @@ async def get_sql_agent(cls) -> dict:
 
                 project_client = AIProjectClient(
                     endpoint=config.AI_PROJECT_ENDPOINT,
-                    credential=DefaultAzureCredentialSync(),
+                    credential=get_azure_credential(),
                     api_version="2025-05-01",
                 )
 

src/App/backend/helpers/azure_credential_utils.py

@@ -0,0 +1,41 @@
+import os
+from azure.identity import ManagedIdentityCredential, DefaultAzureCredential
+from azure.identity.aio import ManagedIdentityCredential as AioManagedIdentityCredential, DefaultAzureCredential as AioDefaultAzureCredential
+
+
+async def get_azure_credential_async(client_id=None):
+    """
+    Returns an Azure credential asynchronously based on the application environment.
+
+    If the environment is 'dev', it uses AioDefaultAzureCredential.
+    Otherwise, it uses AioManagedIdentityCredential.
+
+    Args:
+        client_id (str, optional): The client ID for the Managed Identity Credential.
+
+    Returns:
+        Credential object: Either AioDefaultAzureCredential or AioManagedIdentityCredential.
+    """
+    if os.getenv("APP_ENV", "prod").lower() == 'dev':
+        return AioDefaultAzureCredential()  # CodeQL [SM05139] Okay use of DefaultAzureCredential as it is only used in development
+    else:
+        return AioManagedIdentityCredential(client_id=client_id)
+
+
+def get_azure_credential(client_id=None):
+    """
+    Returns an Azure credential based on the application environment.
+
+    If the environment is 'dev', it uses DefaultAzureCredential.
+    Otherwise, it uses ManagedIdentityCredential.
+
+    Args:
+        client_id (str, optional): The client ID for the Managed Identity Credential.
+
+    Returns:
+        Credential object: Either DefaultAzureCredential or ManagedIdentityCredential.
+    """
+    if os.getenv("APP_ENV", "prod").lower() == 'dev':
+        return DefaultAzureCredential()  # CodeQL [SM05139] Okay use of DefaultAzureCredential as it is only used in development
+    else:
+        return ManagedIdentityCredential(client_id=client_id)

src/App/backend/plugins/chat_with_data_plugin.py

@@ -9,7 +9,8 @@
     MessageRole,
 )
 from azure.ai.projects import AIProjectClient
-from azure.identity import DefaultAzureCredential, get_bearer_token_provider
+from azure.identity import get_bearer_token_provider
+from backend.helpers.azure_credential_utils import get_azure_credential
 from semantic_kernel.functions.kernel_function_decorator import kernel_function
 
 from backend.common.config import config
@@ -202,7 +203,7 @@ async def get_answers_from_calltranscripts(
 
     def get_openai_client(self):
         token_provider = get_bearer_token_provider(
-            DefaultAzureCredential(), "https://cognitiveservices.azure.com/.default"
+            get_azure_credential(), "https://cognitiveservices.azure.com/.default"
         )
         openai_client = openai.AzureOpenAI(
             azure_endpoint=config.AZURE_OPENAI_ENDPOINT,
@@ -213,7 +214,7 @@ def get_openai_client(self):
 
     def get_project_openai_client(self):
         project = AIProjectClient(
-            endpoint=config.AI_PROJECT_ENDPOINT, credential=DefaultAzureCredential()
+            endpoint=config.AI_PROJECT_ENDPOINT, credential=get_azure_credential()
         )
         openai_client = project.inference.get_azure_openai_client(
             api_version=config.AZURE_OPENAI_PREVIEW_API_VERSION