app/app.py

@@ -5,6 +5,7 @@ from typing import Any, Dict, Tuple
 import dash_auth
 import pandas as pd
 from app_styles import header_style
+from config import check_credentials
 from dash import (
     Dash,
     Input,
@@ -21,6 +22,14 @@ from dash.exceptions import PreventUpdate
 from data_chat import send_message
 from sql_utils import execute_query, test_db_connection
 
+check_credentials()
+
+# first connection to SQL database to mitigate long startup time
+try:
+    test_db_connection()
+except Exception as e:
+    print(f"Error for first connection to Azure SQL Database: {e}")
+
 external_stylesheets = ["https://codepen.io/chriddyp/pen/bWLwgP.css"]
 
 app = Dash(__name__, external_stylesheets=external_stylesheets)
@@ -357,4 +366,5 @@ def run_sql_query(n_clicks: int, value: str) -> str:
 server = app.server
 
 if __name__ == "__main__":
+    print("Hello!")
     app.run(debug=True)

app/config.py

@@ -0,0 +1,78 @@
+"""Global configuration for data preprocessing."""
+
+import os
+
+from azure.identity import (
+    AzureCliCredential,
+    ChainedTokenCredential,
+    ManagedIdentityCredential,
+)
+from azure.keyvault.secrets import SecretClient
+
+
+def check_credentials() -> None:
+    """Check and set up necessary credentials for the application.
+
+    This function verifies the presence of required environment variables.
+    If they are not set, it attempts to retrieve them using Azure-managed identity.
+
+    The function checks for the following environment variables:
+    - OPENAI_API_KEY
+    - AZURE_SQL_CONNECTION_STRING
+    - APP_UNAME
+    - APP_PW
+
+    If AZURE_SQL_CONNECTION_STRING is not set, it constructs the connection string
+    using other environment variables (AZURE_SQL_SERVER, AZURE_SQL_PORT,
+    AZURE_SQL_DATABASE, AZURE_SQL_AUTHENTICATION).
+
+    If any of the required credentials are missing, the function uses Azure Key Vault
+    to retrieve the secrets.
+
+    Raises
+    ------
+    Exception
+        If the required environment variables are not set and cannot be retrieved
+        from Azure Key Vault.
+
+    Notes
+    -----
+    This function modifies the following environment variables:
+    - AZURE_SQL_CONNECTION_STRING (if not already set)
+    - OPENAI_API_KEY (if not already set)
+    - APP_UNAME (if not already set)
+    - APP_PW (if not already set)
+
+    The function uses Azure Managed Identity and Azure CLI credentials to access
+    the Key Vault.
+    """
+    try:
+        assert os.getenv("OPENAI_API_KEY") is not None
+        assert os.getenv("AZURE_SQL_CONNECTION_STRING") is not None
+        assert os.getenv("APP_UNAME") is not None
+        assert os.getenv("APP_PW") is not None
+    except Exception:
+        # Environment variables not set, use azure-managed identity
+        if os.getenv("AZURE_SQL_CONNECTION_STRING") is None:
+            server = os.getenv("AZURE_SQL_SERVER")
+            port = os.getenv("AZURE_SQL_PORT")
+            database = os.getenv("AZURE_SQL_DATABASE")
+            authentication = os.getenv("AZURE_SQL_AUTHENTICATION")
+            os.environ["AZURE_SQL_CONNECTION_STRING"] = (
+                f"Driver={{ODBC Driver 18 for SQL Server}};"
+                f"Server={server},{port};Database={database};"
+                f"Authentication={authentication};Encrypt=yes;"
+            )
+
+        managed_identity = ManagedIdentityCredential()
+        azure_cli = AzureCliCredential()
+        credential_chain = ChainedTokenCredential(managed_identity, azure_cli)
+
+        keyVaultName = os.environ["KEY_VAULT_NAME"]
+        KVUri = f"https://{keyVaultName}.vault.azure.net"
+        client = SecretClient(vault_url=KVUri, credential=credential_chain)
+
+        os.environ["OPENAI_API_KEY"] = client.get_secret("openai-api-key").value
+
+        os.environ["APP_UNAME"] = client.get_secret("app-uname").value
+        os.environ["APP_PW"] = client.get_secret("app-pw").value

app/data_chat.py

@@ -4,22 +4,6 @@ from openai import OpenAI
 
 # from openai import AzureOpenAI
 
-# Set up credentials
-# NOTE: Usually I would use AzureOpenAI, but due to heavy rate
-# limitations on azure trial accounts, I am using OpenAI directly
-# for this project. However, this is how it would look like for
-# AzureOpenAI (credentials must be provided to environment):
-# client = AzureOpenAI(
-#     azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT"),
-#     api_key=os.getenv("AZURE_OPENAI_KEY"),
-#     api_version="2024-02-01",
-# )
-#  MODEL = "sqlai"  # deployment name
-
-# Set up the OpenAI client
-MODEL = "gpt-4o"
-client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
-
 
 def send_message(message: str) -> str:
     """Send a message to the openai chat completion API and return the response.
@@ -103,6 +87,22 @@ def send_message(message: str) -> str:
     Formatiere den Output bestmöglich.
     """
 
+    # Set up credentials
+    # NOTE: Usually I would use AzureOpenAI, but due to heavy rate
+    # limitations on azure trial accounts, I am using OpenAI directly
+    # for this project. However, this is how it would look like for
+    # AzureOpenAI (credentials must be provided to environment):
+    # client = AzureOpenAI(
+    #     azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT"),
+    #     api_key=os.getenv("AZURE_OPENAI_KEY"),
+    #     api_version="2024-02-01",
+    # )
+    #  MODEL = "sqlai"  # deployment name
+
+    # Set up the OpenAI client
+    MODEL = "gpt-4o"
+    client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
+
     response = client.chat.completions.create(
         model=MODEL,
         messages=[