Update app.py

app.py

@@ -1,76 +1,141 @@
-import sys
-import toml
+Hugging Face's logo
+Hugging Face
+Search models, datasets, users...
+Models
+Datasets
+Spaces
+Posts
+Docs
+Solutions
+Pricing
+
+
+
+Spaces:
+
+vectara
+/
+IRS-chat
+
+
+like
+3
+
+Logs
+App
+Files
+Community
+Settings
+IRS-chat
+/
+app.py
+
+ofermend's picture
+ofermend
+Update app.py
+81cb44a
+VERIFIED
+1 day ago
+raw
+history
+blame
+edit
+delete
+No virus
+3.56 kB
 from omegaconf import OmegaConf
 from query import VectaraQuery
 import os
 
 import streamlit as st
 from PIL import Image
-from functools import partial
 
-def set_query(q: str):
-    st.session_state['query'] = q
+def isTrue(x) -> bool:
+    if isinstance(x, bool):
+        return s
+    return x.strip().lower() == 'true'
 
 def launch_bot():
-    def get_answer(question):
+    def generate_response(question):
         response = vq.submit_query(question)
         return response
+    
+    def generate_streaming_response(question):
+        response = vq.submit_query_streaming(question)
+        return response
 
-    corpus_ids = list(eval(os.environ['corpus_ids']))
-    questions = list(eval(os.environ['examples']))
-    cfg = OmegaConf.create({
-        'customer_id': os.environ['customer_id'],
-        'corpus_ids': corpus_ids,
-        'api_key': os.environ['api_key'],
-        'title': os.environ['title'],
-        'description': os.environ['description'],
-        'examples': questions,
-        'source_data_desc': os.environ['source_data_desc']
-    })
-    vq = VectaraQuery(cfg.api_key, cfg.customer_id, cfg.corpus_ids)
+    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': 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': os.environ['questions'],
+            'prompt_name': os.environ.get('prompt_name', None)
+        })
+        st.session_state.cfg = cfg
+        st.session_state.vq = VectaraQuery(cfg.api_key, cfg.customer_id, cfg.corpus_ids, cfg.prompt_name)
+
+    cfg = st.session_state.cfg
+    vq = st.session_state.vq
     st.set_page_config(page_title=cfg.title, layout="wide")
 
     # left side content
     with st.sidebar:
         image = Image.open('Vectara-logo.png')
         st.markdown(f"## Welcome to {cfg.title}\n\n"
-                    f"With this demo uses Retieval Augmented Generation to ask questions about {cfg.source_data_desc}\n\n")
+                    f"This demo uses Retrieval Augmented Generation to ask questions about {cfg.source_data_desc}\n\n")
 
         st.markdown("---")
         st.markdown(
             "## How this works?\n"
             "This app was built with [Vectara](https://vectara.com).\n"
             "Vectara's [Indexing API](https://docs.vectara.com/docs/api-reference/indexing-apis/indexing) was used to ingest the data into a Vectara corpus (or index).\n\n"
-            "This app uses Vectara API to query the corpus and present the results to you, answering your question.\n\n"
+            "This app uses Vectara [Chat API](https://docs.vectara.com/docs/console-ui/vectara-chat-overview) to query the corpus and present the results to you, answering your question.\n\n"
         )
         st.markdown("---")
         st.image(image, width=250)
 
-    st.markdown(f"<center> <h2> Vectara demo app: {cfg.title} </h2> </center>", unsafe_allow_html=True)
+    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)
 
-    # Setup a split column layout
-    main_col, questions_col = st.columns([4, 2], gap="medium")
-    with main_col:        
-        cols = st.columns([1, 8], gap="small")
-        cols[0].markdown("""<h5>Search</h5>""", unsafe_allow_html=True)             
-        cols[1].text_input(label="search", key='query', max_chars=256, label_visibility='collapsed', help="Enter your question here")
-
-        st.markdown("<h5>Response</h5>", unsafe_allow_html=True)
-        response_text = st.empty()
-        response_text.text_area(f" ", placeholder="The answer will appear here.", disabled=True,
-                                key="response", height=1, label_visibility='collapsed')
-    with questions_col:
-        st.markdown("<h5 style='text-align:center; color: red'> Sample questions </h5>", unsafe_allow_html=True)
-        for q in list(cfg.examples):
-            st.button(q, on_click=partial(set_query, q), use_container_width=True)
-
-
-    # run the main flow
-    if st.session_state.get('query'):
-        query = st.session_state['query']
-        response = get_answer(query)
-        response_text.markdown(response)
+    if "messages" not in st.session_state.keys():
+        st.session_state.messages = [{"role": "assistant", "content": "How may I help you?"}]
+        for question in questions:
+            st.button(question, on_click=lambda q=question: submit_question(q))
+
+
+    # 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():
+        submit_question(prompt)
+
+    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)
+        
+    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)
+                
 
 if __name__ == "__main__":
     launch_bot()
+