app.py

import streamlit as st
from database import KodeksProcessor
from chatbot import Chatbot
import os
import pandas as pd

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

    # Sidebar do nawigacji
    st.sidebar.title("Nawigacja")
    page = st.sidebar.radio("Wybierz stronę:", ["Chatbot", "Podgląd danych"])

    if page == "Chatbot":
        # Przycisk do czyszczenia historii
        if st.sidebar.button("Wyczyść historię"):
            st.session_state.chatbot.clear_history()
            st.session_state.messages = []
            st.experimental_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})

    elif page == "Podgląd danych":
        st.subheader("Dane w bazie")
        processor = KodeksProcessor()
        
        # Pobierz wszystkie dokumenty z bazy
        all_docs = processor.collection.query(query_texts=[""], n_results=1000)
        
        # Przygotuj dane do wyświetlenia
        data = []
        for doc in all_docs['documents'][0]:
            data.append(doc)

        # Wyświetl dane w tabeli
        if data:
            df = pd.DataFrame(data)
            st.dataframe(df)
        else:
            st.write("Brak danych w bazie.")

if __name__ == "__main__":
    main()