Update ui

1 files changed, 6 insertions, 2 deletions

diff --git a/rag/ui.py b/rag/ui.py
index 277c084..37c50dd 100644
--- a/rag/ui.py
+++ b/rag/ui.py
@@ -22,6 +22,7 @@ def upload_pdfs():
 
 
 if __name__ == "__main__":
+    ss = st.session_state
     st.header("RAG-UI")
 
     upload_pdfs()
@@ -35,10 +36,13 @@ if __name__ == "__main__":
         disabled=False,
     )
 
+    (b,) = st.columns(1)
     (result_column, context_column) = st.columns(2)
 
-    if query:
-        response = rag.retrive(query)
+    if b.button("Generate", disabled=False, type="primary", use_container_width=True):
+        query = ss.get("query", "")
+        with st.spinner("Generating answer..."):
+            response = rag.retrieve(query)
 
         with result_column:
             st.markdown("### Answer")