File size: 2,296 Bytes
5eda983
4180c11
af733da
 
4180c11
 
af733da
37c6c29
 
 
 
 
 
 
 
 
 
 
 
 
 
af733da
5eda983
af733da
 
 
37c6c29
af733da
 
 
37c6c29
 
 
 
af733da
37c6c29
 
 
 
7c33e91
 
af733da
 
 
 
 
 
37c6c29
 
 
9fc65c0
af733da
 
 
 
37c6c29
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
7c33e91
37c6c29
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
import chromadb
import streamlit as st
from lib import Text2SQLRAG
from utils import execute_query_and_return_df



st.set_page_config(page_title="Text2SQLRAG", layout="wide")
st.title("Text2SQLRAG: Conversational SQL Generator")

with st.sidebar:
    
    st.markdown("""
    # Example Question
    - Calculate the average invoice amount
    - Find all artists whose names start with 'A'
    - Find all playlists containing tracks composed by Bach
    - Count how many different artists are in each playlist
    - Find the 5 tracks that appear on the most invoices
    """)

text2sql = Text2SQLRAG()
chromadb.api.client.SharedSystemClient.clear_system_cache()

if "messages" not in st.session_state:
    st.session_state.messages = []

for message in st.session_state.messages:
    role = message.get("role", "assistant")
    with st.chat_message(role):
        st.markdown(message.get("text", ""))
        
        if "reasoning" in message:
            with st.expander("Reasoning", expanded=False):
                st.markdown(message["reasoning"])
        
        if "sql_query" in message:
            with st.expander("SQL Query", expanded=False):
                st.code(message["sql_query"], language="sql")
        if "dataframe" in message:
            st.dataframe(message["dataframe"])

input_text = st.chat_input("Chat with your bot here...")

if input_text:
    with st.chat_message("user"):
        st.markdown(input_text)
    
    st.session_state.messages.append({"role": "user", "text": input_text})
    
    response = text2sql.run(input_text)
    sql_query = response.query
    reasoning = response.reasoning
    df = execute_query_and_return_df(sql_query)
    
    with st.chat_message("assistant"):
        if sql_query:
            with st.expander("SQL Query", expanded=True):
                st.code(sql_query, language="sql")
            
            with st.expander("Reasoning", expanded=True):
                st.markdown(reasoning)
        
        if df is not None and not df.empty:
            st.dataframe(df)
        else:
            st.error("Error executing query or no data returned.")
    
    st.session_state.messages.append({
        "role": "assistant",
        "reasoning": reasoning,
        "sql_query": sql_query,
        "dataframe": df,
    })