""" main.py ------- Streamlit UI for the RAG chatbot. Step 3 Enhancements: - Score Threshold slider (0.0–1.0) in sidebar - Streaming responses via st.write_stream() — no more spinner waiting - Reset Knowledge Base button — deletes FAISS index, clears cache, resets state All previous features retained: - Upload PDF / TXT / DOCX / MD - Top-K slider - Clear Chat History button - Conversation memory (last 6 messages) - Source citations with similarity scores - Cross-platform temp file handling """ import logging import shutil import sys import tempfile from pathlib import Path import streamlit as st sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) from app.chatbot import Chatbot from app.config import APP_DESCRIPTION, APP_TITLE, VECTOR_DB_PATH from components.document_loader import load_document from components.embedder import HuggingFaceEmbedder from components.text_splitter import split_documents from components.vector_store import VectorStore logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) # ── Page config ─────────────────────────────────────────────────────────────── st.set_page_config( page_title=APP_TITLE, page_icon="🤖", layout="wide", ) # ── Session state ───────────────────────────────────────────────────────────── if "messages" not in st.session_state: st.session_state.messages = [] if "store_ready" not in st.session_state: st.session_state.store_ready = False if "top_k" not in st.session_state: st.session_state.top_k = 4 if "score_threshold" not in st.session_state: st.session_state.score_threshold = 0.0 # ── Helper: initialise chatbot (cached) ─────────────────────────────────────── @st.cache_resource(show_spinner="Loading models …") def get_chatbot() -> tuple: embedder = HuggingFaceEmbedder() store = VectorStore(embedder=embedder, index_path=VECTOR_DB_PATH) loaded = store.load() chatbot = Chatbot(vector_store=store) return chatbot, loaded # ── Sidebar ─────────────────────────────────────────────────────────────────── with st.sidebar: st.header("📂 Knowledge Base") st.caption("Upload documents and ingest them into the vector store.") uploaded_files = st.file_uploader( "Upload PDF / TXT / DOCX / MD", type=["pdf", "txt", "docx", "md"], accept_multiple_files=True, ) if st.button("⚙️ Ingest Documents", use_container_width=True): if not uploaded_files: st.warning("Please upload at least one document first.") else: with st.spinner("Ingesting documents …"): chatbot_obj, _ = get_chatbot() all_chunks = [] for uf in uploaded_files: # Cross-platform temp dir (works on Windows + Linux/Mac) tmp_path = Path(tempfile.gettempdir()) / uf.name with open(tmp_path, "wb") as f: f.write(uf.getbuffer()) try: docs = load_document(tmp_path) for d in docs: d.metadata["source"] = uf.name chunks = split_documents(docs) all_chunks.extend(chunks) st.success(f"✅ {uf.name} — {len(chunks)} chunks") except Exception as exc: st.error(f"❌ {uf.name}: {exc}") if all_chunks: chatbot_obj.vector_store.build(all_chunks) st.session_state.store_ready = True st.success(f"Vector store built with {len(all_chunks)} total chunks.") st.divider() st.caption("Or pre-load documents by placing files in `data/raw/` and running `scripts/ingest.py`.") # ── Status badge ───────────────────────────────────────────────────────── chatbot_obj, preloaded = get_chatbot() ready = preloaded or st.session_state.store_ready if ready: st.success("📚 Knowledge base ready") else: st.warning("⚠️ No knowledge base loaded") # ── Top-K slider ────────────────────────────────────────────────────────── st.divider() top_k = st.slider( "Retrieved chunks (Top-K)", min_value=1, max_value=10, value=st.session_state.top_k, help="How many document chunks the retriever fetches per query. Higher = more context but slower.", ) st.session_state.top_k = top_k # ── Score Threshold slider (Part A) ─────────────────────────────────────── score_threshold = st.slider( "Score Threshold", min_value=0.0, max_value=1.0, value=st.session_state.score_threshold, step=0.05, help="Minimum similarity score a chunk must have to be used as context. " "0.0 = include everything. 0.5 = only confident matches.", ) st.session_state.score_threshold = score_threshold # ── Clear Chat button ───────────────────────────────────────────────────── st.divider() if st.button("🗑️ Clear Chat History", use_container_width=True): st.session_state.messages = [] st.rerun() # ── Reset Knowledge Base button (Part C) ────────────────────────────────── st.divider() if st.button("🔄 Reset Knowledge Base", use_container_width=True): index_path = Path(VECTOR_DB_PATH) if index_path.exists(): shutil.rmtree(index_path) # delete FAISS index folder + contents st.session_state.store_ready = False st.cache_resource.clear() # force get_chatbot() to run fresh st.success("Knowledge base cleared. Upload new documents to start fresh.") st.rerun() # ── Main chat UI ────────────────────────────────────────────────────────────── st.title(f"🤖 {APP_TITLE}") st.caption(APP_DESCRIPTION) st.divider() # Render chat history for msg in st.session_state.messages: with st.chat_message(msg["role"]): st.markdown(msg["content"]) if msg.get("sources"): with st.expander("📄 Sources"): for src in msg["sources"]: st.markdown(f"- `{src}`") # Chat input if prompt := st.chat_input("Ask a question about your documents …"): # Display user message st.session_state.messages.append({"role": "user", "content": prompt}) with st.chat_message("user"): st.markdown(prompt) # ── Streaming response (Part B) ─────────────────────────────────────────── with st.chat_message("assistant"): top_k = st.session_state.get("top_k", 4) score_threshold = st.session_state.get("score_threshold", 0.0) history = st.session_state.messages[-6:] if st.session_state.messages else [] chatbot_obj, _ = get_chatbot() token_stream, sources, _ = chatbot_obj.chat_stream( prompt, top_k=top_k, history=history, score_threshold=score_threshold, ) if token_stream is None: # Early exit — store not ready or no relevant chunks found if not chatbot_obj.vector_store.is_ready: answer = "No documents have been ingested yet. Please upload documents first." else: answer = "I couldn't find any relevant information to answer your question." st.markdown(answer) else: # Stream tokens into UI — st.write_stream returns full text when done answer = st.write_stream(token_stream) if sources: with st.expander("📄 Sources"): for src in sources: st.markdown(f"- `{src}`") # Save full answer to session state for chat history st.session_state.messages.append({ "role": "assistant", "content": answer, "sources": sources, })