| 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_dotenv() |
|
|
| |
| 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) |
|
|
| |
| 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: |
| 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") |
|
|
| |
| if "messages" not in st.session_state: |
| st.session_state.messages = [] |
|
|
| |
| for message in st.session_state.messages: |
| with st.chat_message(message["role"]): |
| st.markdown(message["content"]) |
|
|
| |
| if prompt := st.chat_input("Ask about your data"): |
| |
| st.session_state.messages.append({"role": "user", "content": prompt}) |
| |
| |
| with st.chat_message("user"): |
| st.markdown(prompt) |
| |
| try: |
| |
| sql_query = vn.generate_sql(prompt) |
| |
| |
| df = vn.run_sql(sql_query) |
| |
| |
| 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) |
| |
| |
| 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)}") |
|
|
| |
| 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." |
| ) |
|
|
