Update app.py

app.py

@@ -2,7 +2,7 @@ import os
 import re
 import json
 from pathlib import Path
-from typing import List, Dict, Tuple
+from typing import List, Dict, Tuple, Optional
 
 import numpy as np
 import faiss
@@ -28,6 +28,7 @@ HEADING_RE = re.compile(r"^(#{1,6})\s+(.*)$", re.MULTILINE)
 
 # ----------- Load Markdown -----------
 def read_markdown_files(kb_dir: Path) -> List[Dict]:
+    """Read all markdown files from the knowledge base directory."""
     docs = []
     for md_path in sorted(kb_dir.glob("*.md")):
         text = md_path.read_text(encoding="utf-8", errors="ignore")
@@ -35,10 +36,19 @@ def read_markdown_files(kb_dir: Path) -> List[Dict]:
         m = re.search(r"^#\s+(.*)$", text, flags=re.MULTILINE)
         if m:
             title = m.group(1).strip()
-        docs.append({"filepath": str(md_path), "filename": md_path.name, "title": title, "text": text})
+        docs.append({
+            "filepath": str(md_path), 
+            "filename": md_path.name, 
+            "title": title, 
+            "text": text
+        })
     return docs
 
-def chunk_markdown(doc: Dict, chunk_chars: int = 1200, overlap: int = 150) -> List[Dict]:
+def chunk_markdown(doc: Dict, chunk_chars: int = 800, overlap: int = 200) -> List[Dict]:
+    """
+    Split markdown document into overlapping chunks for better retrieval.
+    Reduced chunk size and increased overlap for more precise matching.
+    """
     text = doc["text"]
     sections = re.split(r"(?=^##\s+|\n##\s+|\n###\s+|^###\s+)", text, flags=re.MULTILINE)
     if len(sections) == 1:
@@ -47,16 +57,19 @@ def chunk_markdown(doc: Dict, chunk_chars: int = 1200, overlap: int = 150) -> Li
     chunks = []
     for sec in sections:
         sec = sec.strip()
-        if not sec:
+        if not sec or len(sec) < 50:  # Skip very short sections
             continue
+            
         heading_match = HEADING_RE.search(sec)
         section_heading = heading_match.group(2).strip() if heading_match else doc["title"]
 
+        # Better chunking logic
         start = 0
         while start < len(sec):
             end = min(start + chunk_chars, len(sec))
             chunk_text = sec[start:end].strip()
-            if chunk_text:
+            
+            if len(chunk_text) > 50:  # Only keep substantial chunks
                 chunks.append({
                     "doc_title": doc["title"],
                     "filename": doc["filename"],
@@ -64,9 +77,11 @@ def chunk_markdown(doc: Dict, chunk_chars: int = 1200, overlap: int = 150) -> Li
                     "section": section_heading,
                     "content": chunk_text
                 })
+            
             if end == len(sec):
                 break
             start = max(0, end - overlap)
+    
     return chunks
 
 # ----------- KB Index -----------
@@ -75,13 +90,20 @@ class KBIndex:
         self.embedder = SentenceTransformer(EMBEDDING_MODEL_NAME)
         self.reader_tokenizer = AutoTokenizer.from_pretrained(READER_MODEL_NAME)
         self.reader_model = AutoModelForQuestionAnswering.from_pretrained(READER_MODEL_NAME)
-        self.reader = pipeline("question-answering", model=self.reader_model, tokenizer=self.reader_tokenizer)
+        self.reader = pipeline(
+            "question-answering", 
+            model=self.reader_model, 
+            tokenizer=self.reader_tokenizer,
+            max_answer_len=200,
+            handle_impossible_answer=True
+        )
 
         self.index = None
         self.embeddings = None
         self.metadata = []
 
     def build(self, kb_dir: Path):
+        """Build the FAISS index from markdown files."""
         docs = read_markdown_files(kb_dir)
         if not docs:
             raise RuntimeError(f"No markdown files found in {kb_dir.resolve()}")
@@ -89,11 +111,17 @@ class KBIndex:
         all_chunks = []
         for d in docs:
             all_chunks.extend(chunk_markdown(d))
-        texts = [c["content"] for c in all_chunks]
-        if not texts:
+        
+        if not all_chunks:
             raise RuntimeError("No content chunks generated from KB.")
 
-        embeddings = self.embedder.encode(texts, batch_size=32, convert_to_numpy=True, show_progress_bar=False)
+        texts = [c["content"] for c in all_chunks]
+        embeddings = self.embedder.encode(
+            texts, 
+            batch_size=32, 
+            convert_to_numpy=True, 
+            show_progress_bar=True
+        )
         faiss.normalize_L2(embeddings)
 
         dim = embeddings.shape[1]
@@ -109,7 +137,8 @@ class KBIndex:
             json.dump(self.metadata, f, ensure_ascii=False, indent=2)
         faiss.write_index(index, str(FAISS_PATH))
 
-    def load(self):
+    def load(self) -> bool:
+        """Load pre-built index from disk."""
         if not (EMBEDDINGS_PATH.exists() and METADATA_PATH.exists() and FAISS_PATH.exists()):
             return False
         self.embeddings = np.load(EMBEDDINGS_PATH)
@@ -118,108 +147,149 @@ class KBIndex:
         self.index = faiss.read_index(str(FAISS_PATH))
         return True
 
-    def retrieve(self, query: str, top_k: int = 4) -> List[Tuple[int, float]]:
+    def retrieve(self, query: str, top_k: int = 6) -> List[Tuple[int, float]]:
+        """Retrieve top-k most similar chunks for a query."""
         q_emb = self.embedder.encode([query], convert_to_numpy=True)
         faiss.normalize_L2(q_emb)
         D, I = self.index.search(q_emb, top_k)
         return list(zip(I[0].tolist(), D[0].tolist()))
 
-    def answer(self, question: str, retrieved: List[Tuple[int, float]]):
-        best = {"text": None, "score": -1e9, "meta": None, "sim": 0.0}
+    def answer(self, question: str, retrieved: List[Tuple[int, float]]) -> Tuple[Optional[str], float, List[Dict], float]:
+        """
+        Extract answer from retrieved chunks using QA model.
+        Returns: (answer_text, qa_score, citations, best_similarity)
+        """
+        candidates = []
+        
         for idx, sim in retrieved:
             meta = self.metadata[idx]
             ctx = meta["content"]
+            
             try:
                 out = self.reader(question=question, context=ctx)
-            except Exception:
+                score = float(out.get("score", 0.0))
+                answer_text = out.get("answer", "").strip()
+                
+                if answer_text and len(answer_text) > 5:
+                    candidates.append({
+                        "text": answer_text,
+                        "score": score,
+                        "meta": meta,
+                        "sim": float(sim)
+                    })
+            except Exception as e:
                 continue
-            score = float(out.get("score", 0.0))
-            if score > best["score"]:
-                best = {"text": out.get("answer", "").strip(), "score": score, "meta": meta, "sim": float(sim)}
-        if not best["text"]:
-            return None, 0.0, []
+        
+        if not candidates:
+            return None, 0.0, [], max([s for _, s in retrieved]) if retrieved else 0.0
+        
+        # Sort by combined score (QA score + similarity)
+        candidates.sort(key=lambda x: x["score"] * 0.7 + x["sim"] * 0.3, reverse=True)
+        best = candidates[0]
+        
+        # Generate citations from top retrieved chunks
         citations = []
         seen = set()
-        for idx, _ in retrieved[:2]:
+        for idx, _ in retrieved[:3]:
             m = self.metadata[idx]
             key = (m["filename"], m["section"])
             if key in seen:
                 continue
             seen.add(key)
-            citations.append({"title": m["doc_title"], "filename": m["filename"], "section": m["section"]})
-        return best["text"], best["score"], citations
-
+            citations.append({
+                "title": m["doc_title"], 
+                "filename": m["filename"], 
+                "section": m["section"]
+            })
+        
+        best_sim = max([s for _, s in retrieved]) if retrieved else 0.0
+        return best["text"], best["score"], citations, best_sim
+
+# Initialize KB
 kb = KBIndex()
 
 def ensure_index():
-    # Build on first run in Space; load if cached
+    """Build index on first run or load from cache."""
     if not kb.load():
-        kb.build(KB_DIR)
+        if KB_DIR.exists():
+            kb.build(KB_DIR)
+        else:
+            print(f"Warning: KB directory {KB_DIR} not found. Please create it and add markdown files.")
+
 ensure_index()
 
 # ----------- Guardrails -----------
-LOW_CONF_THRESHOLD = 0.20
-LOW_SIM_THRESHOLD  = 0.30
-
-HELPFUL_SUGGESTIONS = [
-    ("Connect WhatsApp", "How do I connect my WhatsApp number?"),
-    ("Reset Password", "I can't sign in / forgot my password"),
-    ("First Automation", "How do I create my first automation?"),
-    ("Billing & Invoices", "How do I download invoices for billing?"),
-    ("Fix Instagram Connect", "Why can't I connect Instagram?")
+CONFIDENCE_THRESHOLD = 0.25
+SIMILARITY_THRESHOLD = 0.35
+
+QUICK_ACTIONS = [
+    ("🔗 Connect WhatsApp", "How do I connect my WhatsApp number?"),
+    ("🔑 Reset Password", "I can't sign in / forgot my password"),
+    ("⚡ First Automation", "How do I create my first automation?"),
+    ("💳 Billing & Invoices", "How do I download invoices for billing?"),
+    ("📸 Fix Instagram", "Why can't I connect Instagram?")
 ]
 
 def format_citations(citations: List[Dict]) -> str:
+    """Format citations as markdown list."""
     if not citations:
         return ""
-    return "\n".join([f"• **{c['title']}** — _{c['section']}_  (`{c['filename']}`)" for c in citations])
+    lines = []
+    for c in citations:
+        lines.append(f"• **{c['title']}** — _{c['section']}_")
+    return "\n".join(lines)
 
-def respond(user_msg, history):
+def respond(user_msg: str, history: List) -> str:
+    """Generate response to user query using RAG pipeline."""
     user_msg = (user_msg or "").strip()
+    
     if not user_msg:
-        return "How can I help? Try: **Connect WhatsApp** or **Reset password**."
+        return "👋 How can I help? Ask me anything about the knowledge base, or use a quick action button below."
 
-    retrieved = kb.retrieve(user_msg, top_k=4)
+    # Retrieve relevant chunks
+    retrieved = kb.retrieve(user_msg, top_k=6)
+    
     if not retrieved:
-        return "I couldn’t find anything yet. Try rephrasing or pick a quick action below."
-
-    span, score, citations = kb.answer(user_msg, retrieved)
-    if not span:
-        suggestions = "\n".join([f"- {c['title']} — _{c['section']}_" for c in citations]) or "- Try a different query."
-        return f"I’m not fully sure. Here are the closest matches:\n\n{suggestions}"
+        return "❌ I couldn't find any relevant information. Please try rephrasing your question or contact support."
+
+    # Extract answer using QA model
+    answer, qa_score, citations, best_sim = kb.answer(user_msg, retrieved)
+    
+    if not answer:
+        # Fallback: show closest matches
+        citations_md = format_citations(citations)
+        return (
+            f"🤔 I couldn't extract a specific answer, but here are the most relevant sections:\n\n"
+            f"{citations_md}\n\n"
+            f"💡 Try rephrasing your question or ask me to show more details."
+        )
 
-    best_sim = max([s for _, s in retrieved]) if retrieved else 0.0
-    low_conf = (score < LOW_CONF_THRESHOLD) or (best_sim < LOW_SIM_THRESHOLD)
+    # Check confidence
+    low_confidence = (qa_score < CONFIDENCE_THRESHOLD) or (best_sim < SIMILARITY_THRESHOLD)
     citations_md = format_citations(citations)
-    base_answer = span if len(span) > 3 else "I found a relevant section. Opening the steps in the cited article."
-
-    if low_conf:
+    
+    # Format response based on confidence
+    if low_confidence:
         return (
-            f"{base_answer}\n\n---\n**I may be uncertain.** Here are relevant articles:\n{citations_md}\n\n"
-            f"If this doesn’t solve it, ask me to *escalate to human support*."
+            f"⚠️ **Answer (Low Confidence):**\n{answer}\n\n"
+            f"---\n"
+            f"📚 **Related Sources:**\n{citations_md}\n\n"
+            f"💬 *If this doesn't help, please say \"escalate to support\" for human assistance.*"
+        )
+    else:
+        return (
+            f"✅ **Answer:**\n{answer}\n\n"
+            f"---\n"
+            f"📚 **Sources:**\n{citations_md}\n\n"
+            f"💡 *Say \"show more details\" to see the full context.*"
         )
 
-    return f"{base_answer}\n\n---\n**Sources:**\n{citations_md}\n_Tip: Say **show full steps** for more detail._"
-
-def quick_intent(label):
-    for l, q in HELPFUL_SUGGESTIONS:
-        if l == label:
-            return q
-    return ""
-
-def rebuild_index():
-    kb.build(KB_DIR)
-    return gr.update(value="✅ Index rebuilt from KB.")
-
-# ----------- Gradio UI -----------
-def process_message(user_input, history):
-    """
-    history is a list of dicts: [{"role":"user","content":...}, {"role":"assistant","content":...}, ...]
-    We return updated history and a cleared textbox.
-    """
+def process_message(user_input: str, history: List) -> Tuple[List, Dict]:
+    """Process user message and return updated chat history."""
     user_input = (user_input or "").strip()
     if not user_input:
         return history, gr.update(value="")
+    
     reply = respond(user_input, history or [])
     new_history = (history or []) + [
         {"role": "user", "content": user_input},
@@ -227,52 +297,102 @@ def process_message(user_input, history):
     ]
     return new_history, gr.update(value="")
 
-def process_quick(label, history):
-    # turn a quick button into a user message
-    q = quick_intent(label)
-    return process_message(q, history)
+def process_quick(label: str, history: List) -> Tuple[List, Dict]:
+    """Process quick action button click."""
+    for btn_label, query in QUICK_ACTIONS:
+        if label == btn_label:
+            return process_message(query, history)
+    return history, gr.update(value="")
+
+def rebuild_index_handler():
+    """Rebuild the search index from KB directory."""
+    try:
+        kb.build(KB_DIR)
+        return "✅ Index rebuilt successfully! Ready to answer questions."
+    except Exception as e:
+        return f"❌ Error rebuilding index: {str(e)}"
 
-with gr.Blocks(title="Self-Service KB Assistant", fill_height=True) as demo:
+# ----------- Gradio UI -----------
+with gr.Blocks(
+    title="RAG Knowledge Assistant", 
+    theme=gr.themes.Soft(primary_hue="blue"),
+    css="""
+        .contain { max-width: 1200px; margin: auto; }
+        .quick-btn { min-width: 180px !important; }
+    """
+) as demo:
+    
+    # Header
     gr.Markdown(
         """
-# 🧩 Self-Service Knowledge Assistant
-Ask about setup, automations, billing, or troubleshooting.  
-The assistant answers **only from the Knowledge Base** and cites the articles it used.
-
-**Quick actions:** Use a button below to try a common task.
+        # 🤖 RAG Knowledge Assistant
+        ### AI-powered Q&A with document retrieval and citation
+        Ask questions about your knowledge base, and get answers backed by relevant sources.
         """
     )
-
+    
+    # Main chat interface
     with gr.Row():
-        chat = gr.Chatbot(height=420, show_copy_button=True, type="messages")  # ✅ modern format
+        with gr.Column(scale=1):
+            chat = gr.Chatbot(
+                height=500,
+                show_copy_button=True,
+                type="messages",
+                avatar_images=(None, "https://em-content.zobj.net/source/twitter/376/robot_1f916.png")
+            )
+            
+            with gr.Row():
+                txt = gr.Textbox(
+                    placeholder="💬 Ask a question (e.g., How do I connect WhatsApp?)",
+                    scale=9,
+                    show_label=False,
+                    container=False
+                )
+                send = gr.Button("Send", variant="primary", scale=1)
+    
+    # Quick action buttons
+    gr.Markdown("### ⚡ Quick Actions")
     with gr.Row():
-        txt = gr.Textbox(placeholder="e.g., How do I connect WhatsApp?", scale=10)
-        send = gr.Button("Send", variant="primary")
-
-    with gr.Row():
-        # Quick action buttons
-        btn_whatsapp = gr.Button("Connect WhatsApp")
-        btn_reset = gr.Button("Reset Password")
-        btn_first = gr.Button("First Automation")
-        btn_billing = gr.Button("Billing & Invoices")
-        btn_ig = gr.Button("Fix Instagram Connect")
-
-    with gr.Accordion("Admin", open=False):
-        gr.Markdown("Rebuild the search index after changing files in `/kb`.")
-        rebuild = gr.Button("Rebuild Index")
-        status = gr.Markdown("")
-
-    # Wire up events
+        quick_buttons = []
+        for label, _ in QUICK_ACTIONS:
+            btn = gr.Button(label, elem_classes="quick-btn", size="sm")
+            quick_buttons.append((btn, label))
+    
+    # Admin section
+    with gr.Accordion("🔧 Admin Panel", open=False):
+        gr.Markdown(
+            """
+            **Rebuild Index:** Use this after adding or modifying files in the `/kb` directory.  
+            The system will re-scan all markdown files and update the search index.
+            """
+        )
+        with gr.Row():
+            rebuild_btn = gr.Button("🔄 Rebuild Index", variant="secondary")
+            status_msg = gr.Markdown("")
+    
+    # Event handlers
     send.click(process_message, inputs=[txt, chat], outputs=[chat, txt])
-    txt.submit(process_message, inputs=[txt, chat], outputs=[chat, txt])  # Enter key
-
-    btn_whatsapp.click(process_quick, inputs=[gr.State("Connect WhatsApp"), chat], outputs=[chat, txt])
-    btn_reset.click(process_quick, inputs=[gr.State("Reset Password"), chat], outputs=[chat, txt])
-    btn_first.click(process_quick, inputs=[gr.State("First Automation"), chat], outputs=[chat, txt])
-    btn_billing.click(process_quick, inputs=[gr.State("Billing & Invoices"), chat], outputs=[chat, txt])
-    btn_ig.click(process_quick, inputs=[gr.State("Fix Instagram Connect"), chat], outputs=[chat, txt])
-
-    rebuild.click(lambda: (kb.build(KB_DIR), "✅ Index rebuilt from KB.")[1], outputs=status)
+    txt.submit(process_message, inputs=[txt, chat], outputs=[chat, txt])
+    
+    for btn, label in quick_buttons:
+        btn.click(
+            process_quick, 
+            inputs=[gr.State(label), chat], 
+            outputs=[chat, txt]
+        )
+    
+    rebuild_btn.click(rebuild_index_handler, outputs=status_msg)
+    
+    # Footer
+    gr.Markdown(
+        """
+        ---
+        💡 **Tips:**
+        - Be specific in your questions for better results
+        - Check the cited sources for full context
+        - Use quick actions for common tasks
+        """
+    )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()