Update app.py

app.py

@@ -1,5 +1,3 @@
-import sys
-import toml
 from omegaconf import OmegaConf
 from query import VectaraQuery
 import os
@@ -7,11 +5,19 @@ import os
 import streamlit as st
 from PIL import Image
 
+def isTrue(x) -> bool:
+    if isinstance(x, bool):
+        return s
+    return x.strip().lower() == 'true'
 
 def launch_bot():
     def generate_response(question):
         response = vq.submit_query(question)
         return response
+    
+    def generate_streaming_response(question):
+        response = vq.submit_query_streaming(question)
+        return response
 
     if 'cfg' not in st.session_state:
         corpus_ids = str(os.environ['corpus_ids']).split(',')
@@ -22,6 +28,7 @@ def launch_bot():
             'title': os.environ['title'],
             'description': os.environ['description'],
             'source_data_desc': os.environ['source_data_desc'],
+            'streaming': isTrue(os.environ.get('streaming', False)),
             'prompt_name': os.environ.get('prompt_name', None)
         })
         st.session_state.cfg = cfg
@@ -67,11 +74,15 @@ def launch_bot():
     # Generate a new response if last message is not from assistant
     if st.session_state.messages[-1]["role"] != "assistant":
         with st.chat_message("assistant"):
-            with st.spinner("Thinking..."):
-                response = generate_response(prompt) 
-                st.write(response) 
-        message = {"role": "assistant", "content": response}
-        st.session_state.messages.append(message)
+            if cfg.streaming:
+                stream = generate_streaming_response(prompt) 
+                response = st.write_stream(stream) 
+            else:
+                with st.spinner("Thinking..."):
+                    response = generate_response(prompt)
+                    st.write(response)
+            message = {"role": "assistant", "content": response}
+            st.session_state.messages.append(message)
     
 if __name__ == "__main__":
     launch_bot()