Update app.py

app.py

@@ -1,5 +1,3 @@
-import sys
-import toml
 from omegaconf import OmegaConf
 from query import VectaraQuery
 import os
@@ -12,6 +10,10 @@ 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 +24,7 @@ def launch_bot():
             'title': os.environ['title'],
             'description': os.environ['description'],
             'source_data_desc': os.environ['source_data_desc'],
+            'streaming': os.environ.get('streaming', False),
             'prompt_name': os.environ.get('prompt_name', None)
         })
         st.session_state.cfg = cfg
@@ -67,11 +70,16 @@ 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) 
+            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()
+    launch_bot()
+