|
from omegaconf import OmegaConf |
|
from query import VectaraQuery |
|
import streamlit as st |
|
import os |
|
from PIL import Image |
|
|
|
def launch_bot(): |
|
def generate_response(question): |
|
response = vq.submit_query(question) |
|
return response |
|
|
|
def reset(): |
|
st.session_state.messages = [{"role": "assistant", "content": "Please ask your question about drink names.", "avatar": 'π¦'}] |
|
st.session_state.vq = VectaraQuery(cfg.api_key, cfg.customer_id, cfg.corpus_id, cfg.prompt_name) |
|
|
|
if 'cfg' not in st.session_state: |
|
cfg = OmegaConf.create({ |
|
'customer_id': str(os.environ['VECTARA_CUSTOMER_ID']), |
|
'corpus_id': '46', |
|
'api_key': str(os.environ['VECTARA_API_KEY']), |
|
'prompt_name': 'vectara-experimental-summary-ext-2023-12-11-large', |
|
}) |
|
st.session_state.cfg = cfg |
|
st.session_state.vq = VectaraQuery(cfg.api_key, cfg.customer_id, cfg.corpus_id, cfg.prompt_name) |
|
|
|
cfg = st.session_state.cfg |
|
vq = st.session_state.vq |
|
st.set_page_config(page_title="Drink Name Query Bot", layout="wide") |
|
|
|
|
|
with st.sidebar: |
|
image = Image.open('Vectara-logo.png') |
|
st.image(image, width=250) |
|
st.markdown(f"## Welcome to Drink Name Query Bot.\n\n\n") |
|
|
|
if st.button('Start Over'): |
|
reset() |
|
|
|
st.markdown("---") |
|
st.markdown( |
|
"## How this works?\n" |
|
"This app was built with [Vectara](https://vectara.com).\n\n" |
|
"It demonstrates the use of the Chat functionality along with custom prompts and GPT4-Turbo (as part of our [Scale plan](https://vectara.com/pricing/))" |
|
) |
|
st.markdown("---") |
|
|
|
if "messages" not in st.session_state.keys(): |
|
reset() |
|
|
|
|
|
for message in st.session_state.messages: |
|
with st.chat_message(message["role"], avatar=message["avatar"]): |
|
st.write(message["content"]) |
|
|
|
|
|
if prompt := st.chat_input(): |
|
st.session_state.messages.append({"role": "user", "content": prompt, "avatar": 'π§βπ»'}) |
|
with st.chat_message("user", avatar='π§βπ»'): |
|
st.write(prompt) |
|
|
|
|
|
if st.session_state.messages[-1]["role"] != "assistant": |
|
with st.chat_message("assistant", avatar='π€'): |
|
response = generate_response(prompt) |
|
st.write(response) |
|
message = {"role": "assistant", "content": response, "avatar": 'π€'} |
|
st.session_state.messages.append(message) |
|
|
|
if __name__ == "__main__": |
|
launch_bot() |
|
|