File size: 2,354 Bytes
cdd85c7 |
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 |
import streamlit as st
from database import KodeksProcessor
from chatbot import Chatbot
import os
def initialize_session_state():
if 'chatbot' not in st.session_state:
st.session_state.chatbot = Chatbot()
if 'messages' not in st.session_state:
st.session_state.messages = []
def main():
st.title("Asystent Prawny")
initialize_session_state()
# Inicjalizacja bazy danych (jeśli potrzebna)
if 'db_initialized' not in st.session_state:
with st.spinner("Inicjalizacja bazy danych..."):
processor = KodeksProcessor()
if not os.path.exists("chroma_db"):
processor.process_all_files("data/kodeksy")
st.session_state.db_initialized = True
# Przycisk do czyszczenia historii
if st.sidebar.button("Wyczyść historię"):
st.session_state.chatbot.clear_history()
st.session_state.messages = []
st.rerun()
# Wyświetlenie historii czatu
for message in st.session_state.messages:
with st.chat_message(message["role"]):
st.markdown(message["content"])
# Input użytkownika
if prompt := st.chat_input("Zadaj pytanie dotyczące prawa..."):
# Dodaj pytanie użytkownika do historii
st.session_state.messages.append({"role": "user", "content": prompt})
with st.chat_message("user"):
st.markdown(prompt)
# Wyszukaj odpowiednie fragmenty w bazie
processor = KodeksProcessor()
relevant_chunks = processor.search(prompt)
# Wygeneruj odpowiedź
with st.chat_message("assistant"):
message_placeholder = st.empty()
full_response = ""
context = st.session_state.chatbot.generate_context(
[{"text": doc} for doc in relevant_chunks['documents'][0]]
)
for response_chunk in st.session_state.chatbot.get_response(prompt, context):
full_response += response_chunk
message_placeholder.markdown(full_response + "▌")
message_placeholder.markdown(full_response)
# Dodaj odpowiedź asystenta do historii
st.session_state.messages.append({"role": "assistant", "content": full_response})
if __name__ == "__main__":
main()
|