Update app.py

app.py

@@ -1,122 +1,67 @@
 import streamlit as st
-import json
+from database import KodeksProcessor
+from chatbot import Chatbot
 import os
-from sentence_transformers import SentenceTransformer, util
-import torch
-from huggingface_hub import InferenceClient
 
-# Load the Hugging Face token from environment variable
-HF_TOKEN = os.environ.get("HF_TOKEN")
-if not HF_TOKEN:
-    raise ValueError("HF_TOKEN environment variable is not set. Please set it before running the application.")
-
-# Load the processed legal code data
-@st.cache_resource
-def load_data(file_path):
-    with open(file_path, 'r', encoding='utf-8') as f:
-        return json.load(f)
-
-# Initialize the sentence transformer model
-@st.cache_resource
-def load_model():
-    return SentenceTransformer('distiluse-base-multilingual-cased-v1')
-
-def generate_keywords(query):
-    client = InferenceClient(token=HF_TOKEN)
-    
-    prompt = f"Na podstawie poniższego pytania, wygeneruj 3-5 słów kluczowych, które najlepiej opisują główne tematy i koncepcje prawne zawarte w pytaniu. Podaj tylko słowa kluczowe, oddzielone przecinkami.\n\nPytanie: {query}\n\nSłowa kluczowe:"
-    
-    response = client.text_generation(
-        model="Qwen/Qwen2.5-72B-Instruct",
-        prompt=prompt,
-        max_new_tokens=50,
-        temperature=0.3,
-        top_p=0.9
-    )
-    
-    keywords = [keyword.strip() for keyword in response.split(',')]
-    return keywords
+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 search_relevant_chunks(keywords, chunks, model, top_k=3):
-    keyword_embedding = model.encode(keywords, convert_to_tensor=True)
-    chunk_embeddings = model.encode([chunk['text'] for chunk in chunks], convert_to_tensor=True)
+def main():
+    st.title("Asystent Prawny")
     
-    cos_scores = util.pytorch_cos_sim(keyword_embedding, chunk_embeddings)
-    top_results = torch.topk(cos_scores.mean(dim=0), k=top_k)
+    initialize_session_state()
     
-    return [chunks[idx] for idx in top_results.indices]
-
-def generate_ai_response(query, relevant_chunks):
-    client = InferenceClient(token=HF_TOKEN)
-    
-    context = "Kontekst prawny:\n\n"
-    for chunk in relevant_chunks:
-        context += f"{chunk['metadata']['nazwa']} - Artykuł {chunk['metadata']['article']}:\n"
-        context += f"{chunk['text']}\n\n"
-
-    prompt = f"Jesteś asystentem prawniczym. Odpowiedz na poniższe pytanie na podstawie podanego kontekstu prawnego.\n\nKontekst: {context}\n\nPytanie: {query}\n\nOdpowiedź:"
-
-    response = client.text_generation(
-        model="Qwen/Qwen2.5-72B-Instruct",
-        prompt=prompt,
-        max_new_tokens=2048,
-        temperature=0.5,
-        top_p=0.7
-    )
+    # 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
     
-    return response
-
-def main():
-    st.title("Chatbot Prawny z AI")
-
-    # Load data and model
-    data_file = "processed_kodeksy.json"
-    if not os.path.exists(data_file):
-        st.error(f"Plik {data_file} nie istnieje. Najpierw przetwórz dane kodeksów.")
-        return
-
-    chunks = load_data(data_file)
-    model = load_model()
-
-    # Initialize chat history
-    if "messages" not in st.session_state:
+    # Przycisk do czyszczenia historii
+    if st.sidebar.button("Wyczyść historię"):
+        st.session_state.chatbot.clear_history()
         st.session_state.messages = []
+        st.rerun()
 
-    # Display chat history
+    # Wyświetlenie historii czatu
     for message in st.session_state.messages:
         with st.chat_message(message["role"]):
             st.markdown(message["content"])
 
-    # User input
+    # 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)
 
-        # Generate keywords and search for relevant chunks
-        with st.spinner("Analizuję pytanie i szukam odpowiednich informacji..."):
-            keywords = generate_keywords(prompt)
-            relevant_chunks = search_relevant_chunks(keywords, chunks, model)
-
-        # Generate AI response
+        # 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 = generate_ai_response(prompt, relevant_chunks)
+            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})
 
-    # Sidebar for additional options
-    with st.sidebar:
-        st.subheader("Opcje")
-        if st.button("Wyczyść historię czatu"):
-            st.session_state.messages = []
-            st.experimental_rerun()
-
-        st.subheader("Informacje o bazie danych")
-        st.write(f"Liczba chunków: {len(chunks)}")
-        st.write(f"Przykładowy chunk:")
-        st.json(chunks[0] if chunks else {})
-
 if __name__ == "__main__":
     main()