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")