ui/chat_page.py

import datetime
import json

import pandas as pd
import streamlit as st
from langchain_core.messages import HumanMessage, FunctionMessage
from streamlit.delta_generator import DeltaGenerator

from backend.chat_bot.json_decoder import CustomJSONDecoder
from backend.constants.streamlit_keys import CHAT_CURRENT_USER_SESSIONS, EL_SESSION_SELECTOR, \
    EL_UPLOAD_FILES_STATUS, USER_PRIVATE_FILES, EL_BUILD_KB_WITH_FILES, \
    EL_PERSONAL_KB_NAME, EL_PERSONAL_KB_DESCRIPTION, \
    USER_PERSONAL_KNOWLEDGE_BASES, AVAILABLE_RETRIEVAL_TOOLS, EL_PERSONAL_KB_NEEDS_REMOVE, \
    CHAT_KNOWLEDGE_TABLE, EL_UPLOAD_FILES, EL_SELECTED_KBS
from backend.constants.variables import DIVIDER_HTML, USER_NAME, RETRIEVER_TOOLS
from backend.construct.build_chat_bot import build_chat_knowledge_table, initialize_session_manager
from backend.chat_bot.chat import refresh_sessions, on_session_change_submit, refresh_agent, \
    create_private_knowledge_base_as_tool, \
    remove_private_knowledge_bases, add_file, clear_files, clear_history, back_to_main, on_chat_submit


def render_session_manager():
    with st.expander("🤖 Session Management"):
        if CHAT_CURRENT_USER_SESSIONS not in st.session_state:
            refresh_sessions()
        st.markdown("Here you can update `session_id` and `system_prompt`")
        st.markdown("- Click empty row to add a new item")
        st.markdown("- If needs to delete an item, just click it and press `DEL` key")
        st.markdown("- Don't forget to submit your change.")

        st.data_editor(
            data=st.session_state[CHAT_CURRENT_USER_SESSIONS],
            num_rows="dynamic",
            key="session_editor",
            use_container_width=True,
        )
        st.button("⏫ Submit", on_click=on_session_change_submit, type="primary")


def render_session_selection():
    with st.expander("✅ Session Selection", expanded=True):
        st.selectbox(
            "Choose a `session` to chat",
            options=st.session_state[CHAT_CURRENT_USER_SESSIONS],
            index=None,
            key=EL_SESSION_SELECTOR,
            format_func=lambda x: x["session_id"],
            on_change=refresh_agent,
        )


def render_files_manager():
    with st.expander("📃 **Upload your personal files**", expanded=False):
        st.markdown("- Files will be parsed by [Unstructured API](https://unstructured.io/api-key).")
        st.markdown("- All files will be converted into vectors and stored in [MyScaleDB](https://myscale.com/).")
        st.file_uploader(label="⏫ **Upload files**", key=EL_UPLOAD_FILES, accept_multiple_files=True)
        # st.markdown("### Uploaded Files")
        st.dataframe(
            data=st.session_state[CHAT_KNOWLEDGE_TABLE].list_files(st.session_state[USER_NAME]),
            use_container_width=True,
        )
        st.session_state[EL_UPLOAD_FILES_STATUS] = st.empty()
        col_1, col_2 = st.columns(2)
        with col_1:
            st.button(label="Upload files", on_click=add_file)
        with col_2:
            st.button(label="Clear all files and tools", on_click=clear_files)


def _render_create_personal_knowledge_bases(div: DeltaGenerator):
    with div:
        st.markdown("- If you haven't upload your personal files, please upload them first.")
        st.markdown("- Select some **files** to build your `personal knowledge base`.")
        st.markdown("- Once the your `personal knowledge base` is built, "
                    "it will answer your questions using information from your personal **files**.")
        st.multiselect(
            label="⚡️Select some files to build a **personal knowledge base**",
            options=st.session_state[USER_PRIVATE_FILES],
            placeholder="You should upload some files first",
            key=EL_BUILD_KB_WITH_FILES,
            format_func=lambda x: x["file_name"],
        )
        st.text_input(
            label="⚡️Personal knowledge base name",
            value="get_relevant_documents",
            key=EL_PERSONAL_KB_NAME
        )
        st.text_input(
            label="⚡️Personal knowledge base description",
            value="Searches from some personal files.",
            key=EL_PERSONAL_KB_DESCRIPTION,
        )
        st.button(
            label="Build 🔧",
            on_click=create_private_knowledge_base_as_tool
        )


def _render_remove_personal_knowledge_bases(div: DeltaGenerator):
    with div:
        st.markdown("> Here is all your personal knowledge bases.")
        if USER_PERSONAL_KNOWLEDGE_BASES in st.session_state and len(st.session_state[USER_PERSONAL_KNOWLEDGE_BASES]) > 0:
            st.dataframe(st.session_state[USER_PERSONAL_KNOWLEDGE_BASES])
        else:
            st.warning("You don't have any personal knowledge bases, please create a new one.")
        st.multiselect(
            label="Choose a personal knowledge base to delete",
            placeholder="Choose a personal knowledge base to delete",
            options=st.session_state[USER_PERSONAL_KNOWLEDGE_BASES],
            format_func=lambda x: x["tool_name"],
            key=EL_PERSONAL_KB_NEEDS_REMOVE,
        )
        st.button("Delete", on_click=remove_private_knowledge_bases, type="primary")


def render_personal_tools_build():
    with st.expander("🔨 **Build your personal knowledge base**", expanded=True):
        create_new_kb, kb_manager = st.tabs(["Create personal knowledge base", "Personal knowledge base management"])
        _render_create_personal_knowledge_bases(create_new_kb)
        _render_remove_personal_knowledge_bases(kb_manager)


def render_knowledge_base_selector():
    with st.expander("🙋 **Select some knowledge bases to query**", expanded=True):
        st.markdown("- Knowledge bases come in two types: `public` and `private`.")
        st.markdown("- All users can access our `public` knowledge bases.")
        st.markdown("- Only you can access your `personal` knowledge bases.")
        options = st.session_state[RETRIEVER_TOOLS].keys()
        if AVAILABLE_RETRIEVAL_TOOLS in st.session_state:
            options = st.session_state[AVAILABLE_RETRIEVAL_TOOLS]
        st.multiselect(
            label="Select some knowledge base tool",
            placeholder="Please select some knowledge bases to query",
            options=options,
            default=["Wikipedia + Self Querying"],
            key=EL_SELECTED_KBS,
            on_change=refresh_agent,
        )


def chat_page():
    # initialize resources
    build_chat_knowledge_table()
    initialize_session_manager()

    # render sidebar
    with st.sidebar:
        left, middle, right = st.columns([1, 1, 2])
        with left:
            st.button(label="↩️ Log Out", help="log out and back to main page", on_click=back_to_main)
        with right:
            st.markdown(f"👤 `{st.session_state[USER_NAME]}`")
        st.markdown(DIVIDER_HTML, unsafe_allow_html=True)
        render_session_manager()
        render_session_selection()
        render_files_manager()
        render_personal_tools_build()
        render_knowledge_base_selector()

    # render chat history
    if "agent" not in st.session_state:
        refresh_agent()
    for msg in st.session_state.agent.memory.chat_memory.messages:
        speaker = "user" if isinstance(msg, HumanMessage) else "assistant"
        if isinstance(msg, FunctionMessage):
            with st.chat_message(name="from knowledge base", avatar="📚"):
                st.write(
                    f"*{datetime.datetime.fromtimestamp(msg.additional_kwargs['timestamp']).isoformat()}*"
                )
                st.write("Retrieved from knowledge base:")
                try:
                    st.dataframe(
                        pd.DataFrame.from_records(
                            json.loads(msg.content, cls=CustomJSONDecoder)
                        ),
                        use_container_width=True,
                    )
                except Exception as e:
                    st.warning(e)
                    st.write(msg.content)
        else:
            if len(msg.content) > 0:
                with st.chat_message(speaker):
                    # print(type(msg), msg.dict())
                    st.write(
                        f"*{datetime.datetime.fromtimestamp(msg.additional_kwargs['timestamp']).isoformat()}*"
                    )
                    st.write(f"{msg.content}")
    st.session_state["next_round"] = st.empty()
    from streamlit import _bottom
    with _bottom:
        col1, col2 = st.columns([1, 16])
        with col1:
            st.button("🗑️", help="Clean chat history", on_click=clear_history, type="secondary")
        with col2:
            st.chat_input("Input Message", on_submit=on_chat_submit, key="chat_input")