Update app.py

app.py

@@ -1,7 +1,6 @@
 from omegaconf import OmegaConf
 from query import VectaraQuery
 import os
-import json
 
 import streamlit as st
 from PIL import Image
@@ -20,33 +19,16 @@ def launch_bot():
         response = vq.submit_query_streaming(question)
         return response
 
-    def generate_and_display_response(question):
-        if cfg.streaming:
-            stream = generate_streaming_response(question)
-            response = st.write_stream(stream)
-        else:
-            with st.spinner("Thinking..."):
-                response = generate_response(question)
-                st.write(response)
-        message = {"role": "assistant", "content": response}
-        st.session_state.messages.append(message)
-
-    def submit_question(question):
-        st.session_state.messages.append({"role": "user", "content": question})
-        with st.chat_message("user"):
-            st.write(question)
-        generate_and_display_response(question)
-            
     if 'cfg' not in st.session_state:
+        corpus_ids = str(os.environ['corpus_ids']).split(',')
         cfg = OmegaConf.create({
             'customer_id': str(os.environ['customer_id']),
-            'corpus_ids': [str(cid) for cid in json.loads(os.environ['corpus_ids'])],
+            'corpus_ids': corpus_ids,
             'api_key': str(os.environ['api_key']),
             'title': os.environ['title'],
             'description': os.environ['description'],
             'source_data_desc': os.environ['source_data_desc'],
             'streaming': isTrue(os.environ.get('streaming', False)),
-            'questions': list(eval(os.environ['questions'])),
             'prompt_name': os.environ.get('prompt_name', None)
         })
         st.session_state.cfg = cfg
@@ -77,9 +59,6 @@ def launch_bot():
 
     if "messages" not in st.session_state.keys():
         st.session_state.messages = [{"role": "assistant", "content": "How may I help you?"}]
-        for question in cfg.questions:
-            st.button(question, on_click=lambda q=question: submit_question(q))
-
 
     # Display chat messages
     for message in st.session_state.messages:
@@ -88,11 +67,23 @@ def launch_bot():
 
     # User-provided prompt
     if prompt := st.chat_input():
-        print(f"chat msg = {prompt}")
-        submit_question(prompt)
-
-                
-
+        st.session_state.messages.append({"role": "user", "content": prompt})
+        with st.chat_message("user"):
+            st.write(prompt)
+    
+    # Generate a new response if last message is not from assistant
+    if st.session_state.messages[-1]["role"] != "assistant":
+        with st.chat_message("assistant"):
+            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()
-    
+