File size: 2,660 Bytes
3351f47 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 | import streamlit as st
import os
from vanna import VannaBase
import pandas as pd
from dotenv import load_dotenv
from sqlalchemy import create_engine
from sqlalchemy.exc import SQLAlchemyError
# Load environment variables
load_dotenv()
# Initialize Vanna AI
vanna_api_key = os.getenv("VANNA_API_KEY")
if not vanna_api_key:
st.error("VANNA_API_KEY is not set in the environment variables. Please set it and restart the application.")
st.stop()
vn = VannaBase(api_key=vanna_api_key)
# Check if DATABASE_URL is set
database_url = os.getenv("DATABASE_URL")
if not database_url:
st.error("DATABASE_URL is not set in the environment variables. Please set it and restart the application.")
st.stop()
# Try to connect to the database
try:
engine = create_engine(database_url)
with engine.connect() as connection:
st.success("Successfully connected to the database.")
vn.connect_to_postgres(database_url)
except SQLAlchemyError as e:
st.error(f"Failed to connect to the database: {str(e)}")
st.stop()
st.title("Data Explorer")
# Initialize chat history
if "messages" not in st.session_state:
st.session_state.messages = []
# Display chat messages
for message in st.session_state.messages:
with st.chat_message(message["role"]):
st.markdown(message["content"])
# Chat input
if prompt := st.chat_input("Ask about your data"):
# Add user message to chat history
st.session_state.messages.append({"role": "user", "content": prompt})
# Display user message
with st.chat_message("user"):
st.markdown(prompt)
try:
# Generate SQL query
sql_query = vn.generate_sql(prompt)
# Execute SQL query and get results
df = vn.run_sql(sql_query)
# Display assistant response
with st.chat_message("assistant"):
st.markdown(f"Here's the SQL query I generated:\n```sql\n{sql_query}\n```")
st.markdown("And here are the results:")
st.dataframe(df)
# Add assistant message to chat history
st.session_state.messages.append({
"role": "assistant",
"content": f"Here's the SQL query I generated:\n```sql\n{sql_query}\n```\n\nAnd here are the results:\n{df.to_markdown()}"
})
except Exception as e:
st.error(f"An error occurred: {str(e)}")
# Sidebar with additional information
st.sidebar.header("About")
st.sidebar.info(
"This is a data exploration tool using Streamlit and Vanna AI. "
"Ask questions about your data in natural language, and the app will "
"generate SQL queries and display the results."
)
|