|
|
|
|
|
import streamlit as st |
|
|
from datetime import datetime |
|
|
|
|
|
def handle_doc_select(): |
|
|
st.session_state.current_chat = True |
|
|
st.session_state.chat_history = [] |
|
|
|
|
|
def handle_start_chat(): |
|
|
st.session_state.current_chat = True |
|
|
st.session_state.chat_history = [] |
|
|
|
|
|
def display_knowledge_base(conn, backend): |
|
|
st.markdown("### 📚 Knowledge Base") |
|
|
|
|
|
if conn is not None: |
|
|
try: |
|
|
cursor = conn.cursor() |
|
|
cursor.execute("SELECT id, name, upload_date FROM documents ORDER BY upload_date DESC") |
|
|
documents_in_db = cursor.fetchall() |
|
|
|
|
|
if documents_in_db: |
|
|
st.markdown("#### Available Documents") |
|
|
|
|
|
for doc_id, name, upload_date in documents_in_db: |
|
|
col1, col2 = st.columns([3, 1]) |
|
|
|
|
|
with col1: |
|
|
selected = st.checkbox( |
|
|
name, |
|
|
value=doc_id in st.session_state.selected_docs, |
|
|
key=f"doc_{doc_id}", |
|
|
on_change=handle_doc_select |
|
|
) |
|
|
|
|
|
if selected and doc_id not in st.session_state.selected_docs: |
|
|
st.session_state.selected_docs.append(doc_id) |
|
|
elif not selected and doc_id in st.session_state.selected_docs: |
|
|
st.session_state.selected_docs.remove(doc_id) |
|
|
|
|
|
with col2: |
|
|
upload_date = datetime.strptime(upload_date, '%Y-%m-%d %H:%M:%S') |
|
|
st.text(upload_date.strftime('%Y-%m-%d')) |
|
|
|
|
|
initialize_selected_documents(conn, backend) |
|
|
|
|
|
if st.session_state.selected_docs: |
|
|
st.button("🚀 Start New Chat", |
|
|
on_click=handle_start_chat, |
|
|
use_container_width=True) |
|
|
else: |
|
|
st.info("No documents in the knowledge base. Upload some documents to get started!") |
|
|
|
|
|
except Exception as e: |
|
|
st.error(f"Error accessing knowledge base: {e}") |
|
|
|
|
|
def initialize_selected_documents(conn, backend): |
|
|
if st.session_state.selected_docs and not st.session_state.documents_initialized: |
|
|
with st.spinner("Initializing document analysis..."): |
|
|
selected_documents = [] |
|
|
selected_doc_names = [] |
|
|
|
|
|
cursor = conn.cursor() |
|
|
for doc_id in st.session_state.selected_docs: |
|
|
cursor.execute( |
|
|
"SELECT content, name FROM documents WHERE id = ?", |
|
|
(doc_id,) |
|
|
) |
|
|
result = cursor.fetchone() |
|
|
if result: |
|
|
selected_documents.append(result[0]) |
|
|
selected_doc_names.append(result[1]) |
|
|
|
|
|
embeddings = backend.get_embeddings_model() |
|
|
if embeddings: |
|
|
vector_store = backend.initialize_faiss( |
|
|
embeddings, |
|
|
selected_documents, |
|
|
selected_doc_names |
|
|
) |
|
|
if vector_store: |
|
|
st.session_state.qa_system = backend.initialize_qa_system(vector_store) |
|
|
st.session_state.documents_initialized = True |
|
|
|
