Code Editor for app.py

import streamlit as st
from streamlit_ace import st_ace

# Set page config
st.set_page_config(page_title="Welcome to Vanna.AI", layout="wide")

# Custom CSS for dark theme and styling
st.markdown("""
<style>
.stApp {
    background-color: #0f1117;
    color: #ffffff;
}
.css-1d391kg {
    padding-top: 3rem;
}
.output-container {
    background-color: #1a1f2d;
    border-radius: 5px;
    padding: 10px;
    margin: 10px 0;
}
</style>
""", unsafe_allow_html=True)

# Title
st.title("Welcome to Vanna.AI")
st.subheader("Your AI-powered copilot for SQL queries.")

# Chat input
if "messages" not in st.session_state:
    st.session_state.messages = []
    
# Display chat history
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.write(message["content"])
        if "sql" in message:
            st.code(message["sql"], language="sql")
            # Show results in a container
            with st.container():
                st.markdown("**Results:**")
                st.dataframe({"TOTAL_ALARMS": [0], "EARLIEST_ALARM_TIME": ["null"]})

# Chat input
prompt = st.chat_input("Ask me a question about your data that I can turn into SQL")

if prompt:
    # Add user message to chat history
    st.session_state.messages.append({"role": "user", "content": prompt})
    
    # Add mock assistant response
    sample_sql = """SELECT COUNT(*) AS total_alarms, MIN(production_time) as earliest_alarm_time 
FROM fa_jdm.fa_vector_dev1_2023_alarm_top"""
    
    st.session_state.messages.append({
        "role": "assistant", 
        "content": "Here's the SQL query to answer your question:",
        "sql": sample_sql
    })
    
    st.rerun()