File size: 3,400 Bytes
b22c324 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 | import streamlit as st
from chroma.chatbot_retriever import ask
from chroma.chatbot_chroma_db import create_new_chat, delete_chat, load_chats_from_db
# Initialize session state
if "chats" not in st.session_state:
loaded = load_chats_from_db()
if loaded:
st.session_state.chats = loaded
else:
initial_chat_id = create_new_chat()
st.session_state.chats = {"Chat 1": {"messages": [], "chat_id": initial_chat_id}}
if "active_chat" not in st.session_state:
st.session_state.active_chat = "Chat 1"
# st.write(list(st.session_state.chats.keys()))
# print(st.session_state.active_chat)
# sidebar
with st.sidebar:
st.title("Bank Chatbot")
# New Chat button
if st.button("+ New Chat", use_container_width=True):
new_name = f"Chat {len(st.session_state.chats) + 1}"
new_chat_id = create_new_chat()
st.session_state.chats[new_name] = {"messages": [], "chat_id": new_chat_id}
st.session_state.active_chat = new_name
st.rerun()
st.divider()
# List all chats as buttons
for chat_name in list(st.session_state.chats.keys()): # list() avoids mutation during loop
col1, col2 = st.columns([4, 1])
with col1:
if st.button(chat_name, key=f"select_{chat_name}", use_container_width=True):
st.session_state.active_chat = chat_name
st.rerun()
with col2:
if st.button("🗑", key=f"delete_{chat_name}"):
delete_chat(st.session_state.chats[chat_name]["chat_id"])
del st.session_state.chats[chat_name]
# if deleted chat was active, switch to another
if st.session_state.active_chat == chat_name:
remaining = list(st.session_state.chats.keys())
if remaining:
st.session_state.active_chat = remaining[0]
else:
# no chats left, create a fresh one
new_chat_id = create_new_chat()
st.session_state.chats["Chat 1"] = {
"messages": [],
"chat_id": new_chat_id
}
st.session_state.active_chat = "Chat 1"
st.rerun()
st.header(st.session_state.active_chat)
active = st.session_state.active_chat
# st.write(st.session_state.chats)
# Display history for the active chat
for msg in st.session_state.chats[active]["messages"]:
with st.chat_message(msg["role"]):
st.write(msg["content"])
if prompt := st.chat_input("Ask me anything..."):
# if "chat_id" not in st.session_state:
# st.session_state.chat_id = create_new_chat()
active = st.session_state.active_chat
chat_id = st.session_state.chats[active]["chat_id"]
# saves user message
st.session_state.chats[active]["messages"].append({"role": "user", "content": prompt})
# with st.chat_message("user"):
# st.write(prompt)
# add response from RAG chain, passing session chat history
# with st.chat_message("assistant"):
with st.spinner("Thinking..."):
response = ask(prompt, chat_id = chat_id, chat_name=active)
# st.write(response['answer'])
st.session_state.chats[active]["messages"].append({"role": "assistant", "content": response})
st.rerun()
|