Upload rag_app.py

rag_app.py

@@ -0,0 +1,859 @@
+"""
+RAG (Retrieval-Augmented Generation) Application
+==================================================
+A full-featured RAG system with:
+- Document processing (PDF, HTML, DOCX, TXT, MD)
+- Vector database (ChromaDB with persistent storage)
+- Hybrid search (semantic + BM25 keyword search)
+- Conversation memory (last 10 exchanges)
+- Streaming LLM responses with source citations
+- Gradio-based conversational UI
+
+Requirements (install via pip):
+    pip install chromadb sentence-transformers gradio openai pymupdf python-docx \
+                beautifulsoup4 rank_bm25 nltk tiktoken numpy
+
+Usage:
+    1. Place 50+ documents in a ./documents/ folder (PDF, HTML, DOCX, TXT, MD)
+    2. Set your OpenAI API key:  export OPENAI_API_KEY="sk-..."
+    3. Run:  python rag_app.py
+    4. Open the Gradio URL in your browser
+"""
+
+import os
+import re
+import json
+import hashlib
+import logging
+import textwrap
+from pathlib import Path
+from typing import Optional
+from dataclasses import dataclass, field
+from collections import defaultdict
+
+import numpy as np
+
+# -- Document parsing --
+import fitz  # PyMuPDF
+from docx import Document as DocxDocument
+from bs4 import BeautifulSoup
+
+# -- NLP / chunking --
+import nltk
+from nltk.tokenize import sent_tokenize
+
+# -- Embeddings & vector DB --
+from sentence_transformers import SentenceTransformer
+import chromadb
+from chromadb.config import Settings
+
+# -- BM25 keyword search --
+from rank_bm25 import BM25Okapi
+
+# -- LLM --
+import openai
+
+# -- UI --
+import gradio as gr
+
+# ---------------------------------------------------------------------------
+# Configuration
+# ---------------------------------------------------------------------------
+
+@dataclass
+class Config:
+    """Central configuration for the RAG pipeline."""
+
+    # Paths
+    documents_dir: str = "./documents"
+    chroma_persist_dir: str = "./chroma_db"
+
+    # Chunking
+    chunk_size: int = 512          # target tokens per chunk (sentence-based)
+    chunk_overlap: int = 64        # overlap tokens between consecutive chunks
+    min_chunk_length: int = 40     # discard chunks shorter than this (chars)
+
+    # Embedding model (runs locally via sentence-transformers)
+    embedding_model: str = "all-MiniLM-L6-v2"
+    chroma_collection: str = "rag_docs"
+
+    # Retrieval
+    top_k_semantic: int = 20       # initial semantic retrieval
+    top_k_bm25: int = 20           # initial BM25 retrieval
+    top_k_final: int = 5           # after hybrid merge / re-rank
+
+    # Hybrid search weight  (0 = pure BM25, 1 = pure semantic)
+    semantic_weight: float = 0.6
+
+    # LLM
+    openai_model: str = "gpt-4o-mini"
+    temperature: float = 0.2
+    max_context_tokens: int = 6000
+    system_prompt: str = textwrap.dedent("""\
+        You are a knowledgeable assistant. Answer the user's question using ONLY
+        the provided context passages. If the context does not contain enough
+        information, say so honestly.
+
+        Rules:
+        - Cite sources using [Source N] notation after each claim.
+        - Be concise but thorough.
+        - If multiple sources agree, prefer the most specific one.
+        - For follow-up questions, use conversation history for context.
+    """)
+
+    # Conversation memory
+    memory_length: int = 10  # number of past exchanges to keep
+
+    # Server
+    server_port: int = 7860
+    share: bool = True  # set True for public URL via Gradio
+
+
+CFG = Config()
+
+# ---------------------------------------------------------------------------
+# Logging
+# ---------------------------------------------------------------------------
+
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s | %(levelname)-7s | %(message)s",
+    datefmt="%H:%M:%S",
+)
+log = logging.getLogger("rag")
+
+# ---------------------------------------------------------------------------
+# 1. Document Processing
+# ---------------------------------------------------------------------------
+
+@dataclass
+class RawDocument:
+    """A single extracted document before chunking."""
+    text: str
+    metadata: dict = field(default_factory=dict)
+
+
+def extract_pdf(path: str) -> RawDocument:
+    """Extract text and metadata from a PDF using PyMuPDF."""
+    doc = fitz.open(path)
+    pages = []
+    for page in doc:
+        pages.append(page.get_text("text"))
+    meta = doc.metadata or {}
+    return RawDocument(
+        text="\n\n".join(pages),
+        metadata={
+            "source": os.path.basename(path),
+            "path": path,
+            "type": "pdf",
+            "title": meta.get("title", ""),
+            "author": meta.get("author", ""),
+            "pages": len(doc),
+        },
+    )
+
+
+def extract_docx(path: str) -> RawDocument:
+    """Extract text from a DOCX file."""
+    doc = DocxDocument(path)
+    paragraphs = [p.text for p in doc.paragraphs if p.text.strip()]
+    core = doc.core_properties
+    return RawDocument(
+        text="\n\n".join(paragraphs),
+        metadata={
+            "source": os.path.basename(path),
+            "path": path,
+            "type": "docx",
+            "title": core.title or "",
+            "author": core.author or "",
+        },
+    )
+
+
+def extract_html(path: str) -> RawDocument:
+    """Extract text from an HTML file."""
+    with open(path, "r", encoding="utf-8", errors="replace") as f:
+        soup = BeautifulSoup(f.read(), "html.parser")
+    # Remove script and style elements
+    for tag in soup(["script", "style", "nav", "footer", "header"]):
+        tag.decompose()
+    title = soup.title.string if soup.title else ""
+    text = soup.get_text(separator="\n", strip=True)
+    return RawDocument(
+        text=text,
+        metadata={
+            "source": os.path.basename(path),
+            "path": path,
+            "type": "html",
+            "title": title,
+        },
+    )
+
+
+def extract_text(path: str) -> RawDocument:
+    """Extract text from a plain text or markdown file."""
+    with open(path, "r", encoding="utf-8", errors="replace") as f:
+        text = f.read()
+    return RawDocument(
+        text=text,
+        metadata={
+            "source": os.path.basename(path),
+            "path": path,
+            "type": "text",
+        },
+    )
+
+
+EXTRACTORS = {
+    ".pdf": extract_pdf,
+    ".docx": extract_docx,
+    ".html": extract_html,
+    ".htm": extract_html,
+    ".txt": extract_text,
+    ".md": extract_text,
+}
+
+
+def load_documents(directory: str) -> list[RawDocument]:
+    """Recursively load all supported documents from a directory."""
+    docs = []
+    directory = Path(directory)
+    if not directory.exists():
+        log.warning(f"Documents directory not found: {directory}")
+        return docs
+
+    for fpath in sorted(directory.rglob("*")):
+        ext = fpath.suffix.lower()
+        if ext in EXTRACTORS:
+            try:
+                doc = EXTRACTORS[ext](str(fpath))
+                if len(doc.text.strip()) > 50:
+                    docs.append(doc)
+                    log.info(f"  Loaded: {fpath.name} ({len(doc.text):,} chars)")
+            except Exception as e:
+                log.error(f"  Failed: {fpath.name} -> {e}")
+    log.info(f"Total documents loaded: {len(docs)}")
+    return docs
+
+# ---------------------------------------------------------------------------
+# 2. Smart Chunking (Sentence-Based with Overlap)
+# ---------------------------------------------------------------------------
+
+@dataclass
+class Chunk:
+    """A text chunk ready for embedding."""
+    text: str
+    metadata: dict
+    chunk_id: str
+
+
+def _approx_token_count(text: str) -> int:
+    """Rough token count (≈ 4 chars per token for English)."""
+    return len(text) // 4
+
+
+def sentence_chunk(doc: RawDocument, chunk_size: int = 512, overlap: int = 64) -> list[Chunk]:
+    """
+    Sentence-based chunking strategy:
+    - Split text into sentences.
+    - Accumulate sentences until chunk_size tokens is reached.
+    - Overlap by re-including trailing sentences from previous chunk.
+    """
+    try:
+        sentences = sent_tokenize(doc.text)
+    except Exception:
+        nltk.download("punkt_tab", quiet=True)
+        sentences = sent_tokenize(doc.text)
+
+    if not sentences:
+        return []
+
+    chunks: list[Chunk] = []
+    current_sentences: list[str] = []
+    current_tokens = 0
+
+    def _flush(sents: list[str], idx: int):
+        text = " ".join(sents).strip()
+        if len(text) < CFG.min_chunk_length:
+            return
+        chunk_id = hashlib.md5(
+            f"{doc.metadata.get('source', '')}:{idx}:{text[:80]}".encode()
+        ).hexdigest()[:12]
+        chunks.append(Chunk(
+            text=text,
+            metadata={**doc.metadata, "chunk_index": idx},
+            chunk_id=chunk_id,
+        ))
+
+    chunk_idx = 0
+    for sent in sentences:
+        sent_tokens = _approx_token_count(sent)
+        if current_tokens + sent_tokens > chunk_size and current_sentences:
+            _flush(current_sentences, chunk_idx)
+            chunk_idx += 1
+            # Keep overlap sentences from the tail
+            overlap_sents: list[str] = []
+            overlap_tok = 0
+            for s in reversed(current_sentences):
+                t = _approx_token_count(s)
+                if overlap_tok + t > overlap:
+                    break
+                overlap_sents.insert(0, s)
+                overlap_tok += t
+            current_sentences = overlap_sents
+            current_tokens = overlap_tok
+        current_sentences.append(sent)
+        current_tokens += sent_tokens
+
+    if current_sentences:
+        _flush(current_sentences, chunk_idx)
+
+    return chunks
+
+
+def chunk_all_documents(docs: list[RawDocument]) -> list[Chunk]:
+    """Chunk every loaded document."""
+    all_chunks = []
+    for doc in docs:
+        doc_chunks = sentence_chunk(doc, CFG.chunk_size, CFG.chunk_overlap)
+        all_chunks.extend(doc_chunks)
+    log.info(f"Total chunks created: {len(all_chunks)}")
+    return all_chunks
+
+# ---------------------------------------------------------------------------
+# 3. Vector Database (ChromaDB with Persistent Storage)
+# ---------------------------------------------------------------------------
+
+class VectorStore:
+    """Manages ChromaDB collection and embedding model."""
+
+    def __init__(self, config: Config):
+        self.config = config
+        log.info(f"Loading embedding model: {config.embedding_model}")
+        self.embedder = SentenceTransformer(config.embedding_model)
+
+        self.client = chromadb.Client(Settings(
+            persist_directory=config.chroma_persist_dir,
+            anonymized_telemetry=False,
+            is_persistent=True,
+        ))
+        self.collection = self.client.get_or_create_collection(
+            name=config.chroma_collection,
+            metadata={"hnsw:space": "cosine"},
+        )
+        log.info(
+            f"ChromaDB collection '{config.chroma_collection}' "
+            f"has {self.collection.count()} vectors"
+        )
+
+    def embed_text(self, texts: list[str]) -> list[list[float]]:
+        """Generate embeddings for a list of texts."""
+        return self.embedder.encode(texts, show_progress_bar=False).tolist()
+
+    def embed_single(self, text: str) -> list[float]:
+        """Embed a single query string."""
+        return self.embedder.encode(text).tolist()
+
+    def add_chunks(self, chunks: list[Chunk], batch_size: int = 256):
+        """Insert chunks into ChromaDB (skip duplicates by ID)."""
+        existing = set(self.collection.get()["ids"]) if self.collection.count() > 0 else set()
+        new_chunks = [c for c in chunks if c.chunk_id not in existing]
+        if not new_chunks:
+            log.info("No new chunks to add (all already indexed).")
+            return
+
+        for i in range(0, len(new_chunks), batch_size):
+            batch = new_chunks[i : i + batch_size]
+            ids = [c.chunk_id for c in batch]
+            texts = [c.text for c in batch]
+            metas = [c.metadata for c in batch]
+            embeddings = self.embed_text(texts)
+            self.collection.add(
+                ids=ids,
+                documents=texts,
+                metadatas=metas,
+                embeddings=embeddings,
+            )
+            log.info(f"  Indexed batch {i // batch_size + 1} ({len(batch)} chunks)")
+
+        log.info(f"Total vectors in DB: {self.collection.count()}")
+
+    def semantic_search(self, query: str, k: int = 20) -> list[dict]:
+        """Return top-k results by cosine similarity."""
+        count = self.collection.count()
+        if count == 0:
+            return []
+        embedding = self.embed_single(query)
+        results = self.collection.query(
+            query_embeddings=[embedding],
+            n_results=min(k, count),
+            include=["documents", "metadatas", "distances"],
+        )
+        hits = []
+        if results["documents"] and results["documents"][0]:
+            for doc, meta, dist in zip(
+                results["documents"][0],
+                results["metadatas"][0],
+                results["distances"][0],
+            ):
+                hits.append({
+                    "text": doc,
+                    "metadata": meta,
+                    "score": 1 - dist,  # cosine distance -> similarity
+                })
+        return hits
+
+# ---------------------------------------------------------------------------
+# 4. BM25 Keyword Search (for Hybrid Retrieval)
+# ---------------------------------------------------------------------------
+
+class BM25Index:
+    """Maintains a BM25 index over all chunk texts."""
+
+    def __init__(self):
+        self.corpus: list[str] = []
+        self.metadata: list[dict] = []
+        self.bm25: Optional[BM25Okapi] = None
+
+    def build(self, chunks: list[Chunk]):
+        """Build BM25 index from chunks."""
+        self.corpus = [c.text for c in chunks]
+        self.metadata = [c.metadata for c in chunks]
+        tokenized = [self._tokenize(t) for t in self.corpus]
+        self.bm25 = BM25Okapi(tokenized)
+        log.info(f"BM25 index built over {len(self.corpus)} chunks")
+
+    @staticmethod
+    def _tokenize(text: str) -> list[str]:
+        return re.findall(r"\w+", text.lower())
+
+    def search(self, query: str, k: int = 20) -> list[dict]:
+        """Return top-k BM25 results."""
+        if self.bm25 is None:
+            return []
+        tokens = self._tokenize(query)
+        scores = self.bm25.get_scores(tokens)
+        top_idx = np.argsort(scores)[::-1][:k]
+        results = []
+        for idx in top_idx:
+            if scores[idx] > 0:
+                results.append({
+                    "text": self.corpus[idx],
+                    "metadata": self.metadata[idx],
+                    "score": float(scores[idx]),
+                })
+        return results
+
+# ---------------------------------------------------------------------------
+# 5. Hybrid Search: Merge Semantic + BM25 with RRF
+# ---------------------------------------------------------------------------
+
+def reciprocal_rank_fusion(
+    semantic_hits: list[dict],
+    bm25_hits: list[dict],
+    semantic_weight: float = 0.6,
+    k_constant: int = 60,
+    top_k: int = 5,
+) -> list[dict]:
+    """
+    Reciprocal Rank Fusion (RRF) to merge two ranked lists.
+    score(doc) = w_s / (k + rank_semantic) + w_b / (k + rank_bm25)
+    """
+    scores: dict[str, float] = defaultdict(float)
+    doc_map: dict[str, dict] = {}
+
+    bm25_weight = 1.0 - semantic_weight
+
+    for rank, hit in enumerate(semantic_hits, start=1):
+        key = hit["text"][:200]  # use text prefix as dedup key
+        scores[key] += semantic_weight / (k_constant + rank)
+        doc_map[key] = hit
+
+    for rank, hit in enumerate(bm25_hits, start=1):
+        key = hit["text"][:200]
+        scores[key] += bm25_weight / (k_constant + rank)
+        if key not in doc_map:
+            doc_map[key] = hit
+
+    ranked = sorted(scores.items(), key=lambda x: x[1], reverse=True)[:top_k]
+    results = []
+    for key, score in ranked:
+        entry = doc_map[key].copy()
+        entry["hybrid_score"] = score
+        results.append(entry)
+    return results
+
+# ---------------------------------------------------------------------------
+# 6. Conversation Memory
+# ---------------------------------------------------------------------------
+
+class ConversationMemory:
+    """Tracks the last N exchanges for multi-turn support."""
+
+    def __init__(self, max_turns: int = 10):
+        self.max_turns = max_turns
+        self.history: list[dict] = []  # [{"role": "user"/"assistant", "content": ...}]
+
+    def add_user(self, message: str):
+        self.history.append({"role": "user", "content": message})
+        self._trim()
+
+    def add_assistant(self, message: str):
+        self.history.append({"role": "assistant", "content": message})
+        self._trim()
+
+    def _trim(self):
+        # Keep last N *exchanges* (each exchange = 2 messages)
+        max_messages = self.max_turns * 2
+        if len(self.history) > max_messages:
+            self.history = self.history[-max_messages:]
+
+    def get_messages(self) -> list[dict]:
+        return list(self.history)
+
+    def get_context_summary(self) -> str:
+        """Produce a short summary for query rewriting."""
+        if not self.history:
+            return ""
+        recent = self.history[-6:]  # last 3 exchanges
+        lines = []
+        for msg in recent:
+            role = "User" if msg["role"] == "user" else "Assistant"
+            # Truncate long assistant replies
+            content = msg["content"][:300]
+            lines.append(f"{role}: {content}")
+        return "\n".join(lines)
+
+    def clear(self):
+        self.history.clear()
+
+# ---------------------------------------------------------------------------
+# 7. RAG Pipeline (Query → Retrieve → Generate)
+# ---------------------------------------------------------------------------
+
+class RAGPipeline:
+    """Orchestrates the full RAG pipeline."""
+
+    def __init__(self, config: Config):
+        self.config = config
+        self.vector_store = VectorStore(config)
+        self.bm25_index = BM25Index()
+        self.memory = ConversationMemory(max_turns=config.memory_length)
+        self.openai_client = openai.OpenAI(api_key=os.getenv("OPENAI_API_KEY", ""))
+
+    # -- Indexing ----------------------------------------------------------
+
+    def index_documents(self, docs_dir: Optional[str] = None):
+        """Load, chunk, and index all documents."""
+        directory = docs_dir or self.config.documents_dir
+        raw_docs = load_documents(directory)
+        if not raw_docs:
+            log.warning("No documents found. Please add files to the documents folder.")
+            return
+
+        chunks = chunk_all_documents(raw_docs)
+        self.vector_store.add_chunks(chunks)
+        self.bm25_index.build(chunks)
+        log.info("Indexing complete.")
+
+    # -- Query rewriting for follow-ups ------------------------------------
+
+    def _rewrite_query(self, user_query: str) -> str:
+        """Use conversation context to make follow-up queries self-contained."""
+        context = self.memory.get_context_summary()
+        if not context:
+            return user_query
+
+        try:
+            response = self.openai_client.chat.completions.create(
+                model=self.config.openai_model,
+                temperature=0,
+                max_tokens=200,
+                messages=[
+                    {
+                        "role": "system",
+                        "content": (
+                            "Rewrite the user's latest question so it is self-contained, "
+                            "incorporating any necessary context from the conversation. "
+                            "Output ONLY the rewritten question, nothing else."
+                        ),
+                    },
+                    {
+                        "role": "user",
+                        "content": f"Conversation:\n{context}\n\nLatest question: {user_query}",
+                    },
+                ],
+            )
+            rewritten = response.choices[0].message.content.strip()
+            if rewritten:
+                log.info(f"Rewritten query: {rewritten}")
+                return rewritten
+        except Exception as e:
+            log.warning(f"Query rewriting failed: {e}")
+        return user_query
+
+    # -- Retrieval ---------------------------------------------------------
+
+    def retrieve(self, query: str) -> list[dict]:
+        """Hybrid retrieval: semantic + BM25 merged via RRF."""
+        semantic_hits = self.vector_store.semantic_search(
+            query, k=self.config.top_k_semantic
+        )
+        bm25_hits = self.bm25_index.search(query, k=self.config.top_k_bm25)
+
+        merged = reciprocal_rank_fusion(
+            semantic_hits,
+            bm25_hits,
+            semantic_weight=self.config.semantic_weight,
+            top_k=self.config.top_k_final,
+        )
+        return merged
+
+    # -- Generation --------------------------------------------------------
+
+    def generate(self, user_query: str, retrieved_chunks: list[dict]) -> str:
+        """Call the LLM with retrieved context and conversation history."""
+
+        # Build context block with source labels
+        context_parts = []
+        for i, chunk in enumerate(retrieved_chunks, 1):
+            source = chunk["metadata"].get("source", "unknown")
+            title = chunk["metadata"].get("title", "")
+            label = f"[Source {i}: {source}"
+            if title:
+                label += f" — {title}"
+            label += "]"
+            context_parts.append(f"{label}\n{chunk['text']}")
+
+        context_block = "\n\n---\n\n".join(context_parts)
+
+        # Assemble messages
+        messages = [{"role": "system", "content": self.config.system_prompt}]
+
+        # Add conversation history
+        messages.extend(self.memory.get_messages())
+
+        # Add current turn with context
+        user_content = (
+            f"Context passages:\n\n{context_block}\n\n"
+            f"---\n\nQuestion: {user_query}"
+        )
+        messages.append({"role": "user", "content": user_content})
+
+        try:
+            response = self.openai_client.chat.completions.create(
+                model=self.config.openai_model,
+                temperature=self.config.temperature,
+                max_tokens=1500,
+                messages=messages,
+            )
+            return response.choices[0].message.content
+        except Exception as e:
+            return f"LLM generation error: {e}"
+
+    # -- Full pipeline -----------------------------------------------------
+
+    def query(self, user_input: str) -> tuple[str, list[dict]]:
+        """
+        Full RAG pipeline:
+        1. Rewrite query using conversation context
+        2. Hybrid retrieve top-K chunks
+        3. Generate answer with citations
+        4. Update memory
+        Returns (answer_text, retrieved_sources)
+        """
+        # Step 1: Rewrite for follow-ups
+        search_query = self._rewrite_query(user_input)
+
+        # Step 2: Retrieve
+        chunks = self.retrieve(search_query)
+        if not chunks:
+            answer = (
+                "I couldn't find any relevant information in the document collection "
+                "to answer your question. Could you rephrase or ask about a different topic?"
+            )
+            self.memory.add_user(user_input)
+            self.memory.add_assistant(answer)
+            return answer, []
+
+        # Step 3: Generate
+        answer = self.generate(user_input, chunks)
+
+        # Step 4: Update memory
+        self.memory.add_user(user_input)
+        self.memory.add_assistant(answer)
+
+        return answer, chunks
+
+    def reset_conversation(self):
+        """Clear conversation history."""
+        self.memory.clear()
+        return "Conversation history cleared."
+
+# ---------------------------------------------------------------------------
+# 8. Gradio Conversational UI
+# ---------------------------------------------------------------------------
+
+def build_ui(pipeline: RAGPipeline) -> gr.Blocks:
+    """Create the Gradio chat interface with source citations."""
+
+    CUSTOM_CSS = """
+    .gradio-container {
+        max-width: 960px !important;
+        margin: auto !important;
+        font-family: 'Segoe UI', system-ui, sans-serif !important;
+    }
+    .source-card {
+        background: #f8f9fa;
+        border-left: 3px solid #4a90d9;
+        padding: 10px 14px;
+        margin: 6px 0;
+        border-radius: 4px;
+        font-size: 0.88em;
+        line-height: 1.5;
+    }
+    .source-card strong { color: #2c5282; }
+    .status-bar {
+        text-align: center;
+        padding: 6px;
+        font-size: 0.85em;
+        color: #718096;
+    }
+    """
+
+    with gr.Blocks(css=CUSTOM_CSS, title="RAG Assistant", theme=gr.themes.Soft()) as demo:
+        gr.Markdown(
+            "# 📚 RAG Document Assistant\n"
+            "Ask questions about the indexed document collection. "
+            "Sources are cited inline and shown below each answer."
+        )
+
+        chatbot = gr.Chatbot(
+            label="Conversation",
+            height=520,
+            show_copy_button=True,
+            bubble_full_width=False,
+            avatar_images=(None, "https://em-content.zobj.net/source/twitter/376/robot_1f916.png"),
+        )
+
+        sources_display = gr.HTML(
+            value='<div class="status-bar">Sources will appear here after each answer.</div>',
+            label="Retrieved Sources",
+        )
+
+        with gr.Row():
+            msg_input = gr.Textbox(
+                placeholder="Ask a question about your documents...",
+                show_label=False,
+                scale=9,
+                container=False,
+            )
+            send_btn = gr.Button("Send", variant="primary", scale=1)
+
+        with gr.Row():
+            clear_btn = gr.Button("🗑 Clear Chat", size="sm")
+            status = gr.Markdown(
+                f"*{pipeline.vector_store.collection.count()} chunks indexed "
+                f"| Hybrid search (semantic + BM25) "
+                f"| Memory: last {pipeline.config.memory_length} exchanges*"
+            )
+
+        # -- Event handlers ------------------------------------------------
+
+        def respond(user_message: str, chat_history: list):
+            if not user_message.strip():
+                return "", chat_history, ""
+
+            # Guard: check if any documents are indexed
+            if pipeline.vector_store.collection.count() == 0:
+                answer = (
+                    "⚠️ No documents have been indexed yet. Please add at least 50 documents "
+                    "(PDF, DOCX, HTML, TXT, or MD files) to the `./documents/` folder and "
+                    "restart the application."
+                )
+                chat_history = chat_history + [[user_message, answer]]
+                return "", chat_history, '<div class="status-bar">No documents indexed.</div>'
+
+            answer, sources = pipeline.query(user_message)
+            chat_history = chat_history + [[user_message, answer]]
+
+            # Format sources as HTML cards
+            if sources:
+                cards = []
+                for i, s in enumerate(sources, 1):
+                    src = s["metadata"].get("source", "unknown")
+                    title = s["metadata"].get("title", "")
+                    score = s.get("hybrid_score", s.get("score", 0))
+                    preview = s["text"][:250].replace("\n", " ") + "..."
+                    card = (
+                        f'<div class="source-card">'
+                        f"<strong>[Source {i}]</strong> {src}"
+                        f"{f' — <em>{title}</em>' if title else ''}"
+                        f" (score: {score:.4f})<br>"
+                        f"<span style='color:#555'>{preview}</span>"
+                        f"</div>"
+                    )
+                    cards.append(card)
+                sources_html = "".join(cards)
+            else:
+                sources_html = '<div class="status-bar">No relevant sources found.</div>'
+
+            return "", chat_history, sources_html
+
+        def clear_chat():
+            pipeline.reset_conversation()
+            return [], '<div class="status-bar">Conversation cleared. Sources will appear here.</div>'
+
+        # Wire events
+        msg_input.submit(respond, [msg_input, chatbot], [msg_input, chatbot, sources_display])
+        send_btn.click(respond, [msg_input, chatbot], [msg_input, chatbot, sources_display])
+        clear_btn.click(clear_chat, outputs=[chatbot, sources_display])
+
+    return demo
+
+# ---------------------------------------------------------------------------
+# 9. Main Entry Point
+# ---------------------------------------------------------------------------
+
+def main():
+    """Initialize the pipeline, index documents, and launch the UI."""
+    log.info("=" * 60)
+    log.info("RAG Application Starting")
+    log.info("=" * 60)
+
+    # Ensure NLTK data is available
+    try:
+        sent_tokenize("Hello world.")
+    except LookupError:
+        nltk.download("punkt_tab", quiet=True)
+
+    # Validate API key
+    api_key = os.getenv("OPENAI_API_KEY", "")
+    if not api_key:
+        log.warning(
+            "OPENAI_API_KEY not set. LLM generation will fail. "
+            "Set it with: export OPENAI_API_KEY='sk-...'"
+        )
+
+    # Create documents directory if needed
+    os.makedirs(CFG.documents_dir, exist_ok=True)
+
+    # Initialize pipeline
+    pipeline = RAGPipeline(CFG)
+
+    # Index documents (idempotent — skips already-indexed chunks)
+    pipeline.index_documents()
+
+    # Build and launch UI
+    demo = build_ui(pipeline)
+    log.info(f"Launching Gradio on port {CFG.server_port} (share={CFG.share})")
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=CFG.server_port,
+        share=CFG.share,
+        show_error=True,
+    )
+
+
+if __name__ == "__main__":
+    main()