feat(ai-chat): Add SQL query field for comparison

In order to compare the (not yet implemented) SQL query generated by
the LLM with an actual query, another text field was added that parses
the query to `pyodbc`, which connects to our database, stores the
resulting rows in a `pandas` dataframe and then visualizes it as a table
in plotly dash.

The SQL functionalities are implemented in the `sql_utils.py` module.

Additionally, some minor updates to the overall behavior and layout of
the app were implemented.
This commit is contained in:
Tobias Quadfasel
2024-09-02 20:43:48 +02:00
parent 923dc3b439
commit 4b9fa0579e
4 changed files with 272 additions and 57 deletions

65
app/sql_utils.py Normal file
View File

@@ -0,0 +1,65 @@
import os
from typing import Union
import pandas as pd
import pyodbc
def test_db_connection() -> bool:
"""Test the connection to Azure SQL Database.
This function attempts to establish a connection to an Azure SQL Database
using the connection string stored in the environment variable
'AZURE_SQL_CONNECTION_STRING'. It makes up to 5 attempts to connect,
with a timeout of 240 seconds for each attempt.
Returns
-------
bool
True if the connection was successful, False otherwise.
"""
connection_string = os.environ.get("AZURE_SQL_CONNECTION_STRING")
for i in range(5):
print(f"Trying to connect to Azure SQL Database... Attempt {i + 1}")
try:
pyodbc.connect(connection_string, timeout=240)
print("Connected to Azure SQL Database successfully!")
connected = True
break
except pyodbc.Error as e:
print(f"Error connecting to Azure SQL Database: {e}")
connected = False
return connected
def execute_query(query: str) -> Union[pd.DataFrame, str]:
"""Execute a SQL query on an Azure SQL Database and return the results.
This function connects to an Azure SQL Database using the connection string
stored in the environment variable 'AZURE_SQL_CONNECTION_STRING', executes
the provided SQL query, and returns the results as a pandas DataFrame.
Parameters
----------
query : str
The SQL query to execute.
Returns
-------
Union[pd.DataFrame, str]
A pandas DataFrame containing the query results if successful,
or a string containing the error message if an exception occurs.
"""
try:
connection_string = os.environ.get("AZURE_SQL_CONNECTION_STRING")
conn = pyodbc.connect(connection_string, timeout=240)
df = pd.read_sql(query, conn)
conn.close()
return df
except Exception as e:
return str(e)
finally:
if conn in locals():
conn.close()