diff --git a/app/app.py b/app/app.py index 752ed67..92bb713 100644 --- a/app/app.py +++ b/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) diff --git a/app/config.py b/app/config.py new file mode 100644 index 0000000..a1a38fb --- /dev/null +++ b/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 diff --git a/app/data_chat.py b/app/data_chat.py index 4cc3106..9d3c267 100644 --- a/app/data_chat.py +++ b/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=[