Update app.py

app.py

@@ -3,8 +3,11 @@ from query import VectaraQuery
 import os
 
 import streamlit as st
+from streamlit_pills import pills
+
 from PIL import Image
 
+max_examples = 4
 
 def isTrue(x) -> bool:
     if isinstance(x, bool):
@@ -19,29 +22,40 @@ def launch_bot():
     def generate_streaming_response(question):
         response = vq.submit_query_streaming(question)
         return response
+    
+    def show_example_questions():        
+        if len(st.session_state.example_messages) > 0 and st.session_state.first_turn:            
+            selected_example = pills("Queries to Try:", st.session_state.example_messages, index=None)
+            if selected_example:
+                st.session_state.ex_prompt = selected_example
+                st.session_state.first_turn = False
+                return True
+        return False
 
     if 'cfg' not in st.session_state:
-        corpus_ids = str(os.environ['corpus_ids']).split(',')
+        corpus_keys = str(os.environ['corpus_keys']).split(',')
         cfg = OmegaConf.create({
-            'customer_id': str(os.environ['customer_id']),
-            'corpus_ids': corpus_ids,
+            'corpus_keys': corpus_keys,
             '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)),
             'prompt_name': os.environ.get('prompt_name', None),
-            'examples': os.environ.get('examples', [])
+            'examples': os.environ.get('examples', None)
         })
         st.session_state.cfg = cfg
-        st.session_state.vq = VectaraQuery(cfg.api_key, cfg.customer_id, cfg.corpus_ids, cfg.prompt_name)
+        st.session_state.ex_prompt = None
+        st.session_state.first_turn = True        
+        example_messages = [example.strip() for example in cfg.examples.split(",")]
+        st.session_state.example_messages = [em for em in example_messages if len(em)>0][:max_examples]
+        
+        st.session_state.vq = VectaraQuery(cfg.api_key, cfg.corpus_keys, cfg.prompt_name)
 
     cfg = st.session_state.cfg
     vq = st.session_state.vq
     st.set_page_config(page_title=cfg.title, layout="wide")
 
-
-    # with streamlit_analytics.track():
     # left side content
     with st.sidebar:
         image = Image.open('Vectara-logo.png')
@@ -61,46 +75,31 @@ def launch_bot():
     st.markdown(f"<center> <h2> Vectara chat demo: {cfg.title} </h2> </center>", unsafe_allow_html=True)
     st.markdown(f"<center> <h4> {cfg.description} </h4> </center>", unsafe_allow_html=True)
 
-    
     if "messages" not in st.session_state.keys():
         st.session_state.messages = [{"role": "assistant", "content": "How may I help you?"}]
 
-
-    # if cfg.examples:
-    #     example_messages = [example.strip() for example in cfg.examples.split(",")]
-    # else:
-    #     example_messages = []
-
-    example_messages = [example.strip() for example in cfg.examples.split(",")]
-    
-    if len(example_messages) > 0:
-        st.markdown("<h6>Queries To Try:</h6>", unsafe_allow_html=True)
-        cols = st.columns(4)
-    
+    example_container = st.empty()
+    with example_container:
+        if show_example_questions():
+            example_container.empty()
+            st.rerun()
+                
     # Display chat messages
     for message in st.session_state.messages:
         with st.chat_message(message["role"]):
             st.write(message["content"])
 
-    # User-provided prompt
-    if prompt := st.chat_input():
+    # select prompt from example question or user provided input
+    if st.session_state.ex_prompt:
+        prompt = st.session_state.ex_prompt
+    else:
+        prompt = st.chat_input()
+    if prompt:
         st.session_state.messages.append({"role": "user", "content": prompt})
         with st.chat_message("user"):
             st.write(prompt)
-
-    # Example prompt
-    for i, example in enumerate(example_messages):
-        button_pressed = False
-        with cols[i]:
-            if st.button(example):
-                prompt = example
-                button_pressed = True
-
-        if button_pressed:
-            st.session_state.messages.append({"role": "user", "content": prompt})
-            with st.chat_message("user"):
-                st.write(prompt)
-    
+        st.session_state.ex_prompt = None
+        
     # Generate a new response if last message is not from assistant
     if st.session_state.messages[-1]["role"] != "assistant":
         with st.chat_message("assistant"):
@@ -113,6 +112,7 @@ def launch_bot():
                     st.write(response)
             message = {"role": "assistant", "content": response}
             st.session_state.messages.append(message)
+            st.rerun()
     
 if __name__ == "__main__":
     launch_bot()