Streamlit chatbot template

ui.py

import uuid

import streamlit as st

#########
# Helpers
#########


def get_completion(question, session_id):
    """TODO: hit RAG API instead."""
    return f"Answer to '{question}'"


def generate_new_session_id():
    return uuid.uuid4()


def get_session_id():
    return st.session_state.session_id


def handler(question, session_id):
    answer = get_completion(question, session_id)
    add_answer_to_session_state(answer)
    display_answer(answer)


def clear_screen():
    st.session_state.session_id = generate_new_session_id()
    st.session_state.messages = [
        {"role": "assistant", "content": "How may I assist you today?"}
    ]


def add_question_to_session_state(question):
    st.session_state.messages.append({"role": "user", "content": question})


def add_answer_to_session_state(answer):
    st.session_state.messages.append({"role": "assistant", "content": answer})


def display_message(message):
    with st.chat_message(message["role"]):
        st.write(message["content"])


def display_question(question):
    with st.chat_message("user"):
        st.write(question)


def display_answer(answer):
    with st.chat_message("assistant"):
        st.markdown(answer)


#########
# UI
#########

st.set_page_config(page_title="RAG POC")

if "session_id" not in st.session_state.keys():
    st.session_state.session_id = generate_new_session_id()

if "messages" not in st.session_state.keys():
    st.session_state.messages = [
        {"role": "assistant", "content": "How may I assist you today?"}
    ]

for message in st.session_state.messages:
    display_message(message)

if question := st.chat_input():
    session_id = get_session_id()
    add_question_to_session_state(question)
    display_question(question)
    handler(question, session_id)

with st.sidebar:
    st.button("Clear Screen", on_click=clear_screen, key="clear_screen")