From 5f777ecdfbf486e5057d31547bdc53358037dce0 Mon Sep 17 00:00:00 2001
From: Gustaf Rydholm <gustaf.rydholm@gmail.com>
Date: Mon, 8 Apr 2024 00:45:44 +0200
Subject: Update ui

---
 rag/ui.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

(limited to 'rag/ui.py')

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")
-- 
cgit v1.2.3-70-g09d2