Update app.py

app.py

@@ -4,62 +4,68 @@ from lib import Text2SQLRAG
 from utils import execute_query_and_return_df
 
 
-st.set_page_config(page_title="Text2SQLRAG")
-st.title("Text2SQLRAG")
 
-# Create an instance of Text2SQLRAG
+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()
 
-# Initialize session state for storing chat messages
 if "messages" not in st.session_state:
     st.session_state.messages = []
-    
-    # Display conversation history from session state
+
 for message in st.session_state.messages:
     role = message.get("role", "assistant")
     with st.chat_message(role):
-        if "output" in message:
-            with st.expander("Reasoning", expanded=True):
+        st.markdown(message.get("text", ""))
+        
+        if "reasoning" in message:
+            with st.expander("Reasoning", expanded=False):
                 st.markdown(message["reasoning"])
-        if "sql_query" in message and message["sql_query"]:
-            with st.expander("SQL Query", expanded=True):
-                st.code(message["sql_query"])
+        
+        if "sql_query" in message:
+            with st.expander("SQL Query", expanded=False):
+                st.code(message["sql_query"], language="sql")
 
-# Get user input
 input_text = st.chat_input("Chat with your bot here...")
 
 if input_text:
-    # Display user input
     with st.chat_message("user"):
         st.markdown(input_text)
-
-    # Add user input to chat history
-    st.session_state.chat_history.append({"role": "user", "text": input_text})
-
-    # Get chatbot response
-    response = text2sql.run(input_text)
+    
+    st.session_state.messages.append({"role": "user", "text": input_text})
+    
+    response = text2sql.run(input_text, database=selected_db)
     sql_query = response.query
     reasoning = response.reasoning
     df = execute_query_and_return_df(sql_query)
-
-    if sql_query:
-        with st.expander("SQL Query", expanded=True):
-            st.code(sql_query)
-
-        with st.expander("Reasoning", expanded=True):
-            st.write(reasoning)
-
-    if df is not None:
-        st.dataframe(df)
-    else:
-        st.error("Error executing query")
     
-    # Append assistant response to session state
-    st.session_state.messages.append(
-        {
-            "role": "assistant",
-            "reasoning": reasoning,
-            "sql_query": 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,
+    })