File size: 3,471 Bytes
04126a0 |
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 |
# src/components/knowledge_base.py
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
|