Update app.py

app.py

@@ -3,6 +3,7 @@ import re
 import json
 from pathlib import Path
 from typing import List, Dict, Tuple, Optional
+import tempfile
 
 import numpy as np
 import faiss
@@ -10,6 +11,8 @@ import gradio as gr
 
 from transformers import pipeline, AutoTokenizer, AutoModelForQuestionAnswering
 from sentence_transformers import SentenceTransformer
+import PyPDF2
+import docx
 
 # ----------- Paths -----------
 KB_DIR = Path("./kb")
@@ -26,7 +29,52 @@ FAISS_PATH = INDEX_DIR / "kb_faiss.index"
 
 HEADING_RE = re.compile(r"^(#{1,6})\s+(.*)$", re.MULTILINE)
 
-# ----------- Load Markdown -----------
+# ----------- Load Documents -----------
+def extract_text_from_pdf(file_path: str) -> str:
+    """Extract text from PDF file."""
+    text = ""
+    try:
+        with open(file_path, 'rb') as file:
+            pdf_reader = PyPDF2.PdfReader(file)
+            for page in pdf_reader.pages:
+                text += page.extract_text() + "\n"
+    except Exception as e:
+        raise RuntimeError(f"Error reading PDF: {str(e)}")
+    return text
+
+def extract_text_from_docx(file_path: str) -> str:
+    """Extract text from DOCX file."""
+    try:
+        doc = docx.Document(file_path)
+        text = "\n".join([paragraph.text for paragraph in doc.paragraphs])
+        return text
+    except Exception as e:
+        raise RuntimeError(f"Error reading DOCX: {str(e)}")
+
+def extract_text_from_txt(file_path: str) -> str:
+    """Extract text from TXT file."""
+    try:
+        with open(file_path, 'r', encoding='utf-8', errors='ignore') as file:
+            return file.read()
+    except Exception as e:
+        raise RuntimeError(f"Error reading TXT: {str(e)}")
+
+def extract_text_from_file(file_path: str) -> Tuple[str, str]:
+    """
+    Extract text from uploaded file based on extension.
+    Returns: (text_content, file_type)
+    """
+    ext = Path(file_path).suffix.lower()
+    
+    if ext == '.pdf':
+        return extract_text_from_pdf(file_path), 'PDF'
+    elif ext == '.docx':
+        return extract_text_from_docx(file_path), 'DOCX'
+    elif ext in ['.txt', '.md']:
+        return extract_text_from_txt(file_path), 'Text'
+    else:
+        raise ValueError(f"Unsupported file type: {ext}. Supported: .pdf, .docx, .txt, .md")
+
 def read_markdown_files(kb_dir: Path) -> List[Dict]:
     """Read all markdown files from the knowledge base directory."""
     docs = []
@@ -101,6 +149,7 @@ class KBIndex:
         self.index = None
         self.embeddings = None
         self.metadata = []
+        self.uploaded_file_active = False  # Track if using uploaded file
 
     def build(self, kb_dir: Path):
         """Build the FAISS index from markdown files."""
@@ -131,12 +180,57 @@ class KBIndex:
         self.index = index
         self.embeddings = embeddings
         self.metadata = all_chunks
+        self.uploaded_file_active = False
 
         np.save(EMBEDDINGS_PATH, embeddings)
         with open(METADATA_PATH, "w", encoding="utf-8") as f:
             json.dump(self.metadata, f, ensure_ascii=False, indent=2)
         faiss.write_index(index, str(FAISS_PATH))
 
+    def build_from_uploaded_file(self, file_path: str, filename: str):
+        """Build temporary index from an uploaded file."""
+        # Extract text from file
+        text_content, file_type = extract_text_from_file(file_path)
+        
+        if not text_content or len(text_content.strip()) < 100:
+            raise RuntimeError("File appears to be empty or too short.")
+        
+        # Create document structure
+        doc = {
+            "filepath": file_path,
+            "filename": filename,
+            "title": Path(filename).stem.replace("_", " ").title(),
+            "text": text_content
+        }
+        
+        # Chunk the document
+        all_chunks = chunk_markdown(doc)
+        
+        if not all_chunks:
+            raise RuntimeError("Could not extract meaningful content from file.")
+        
+        # Build embeddings
+        texts = [c["content"] for c in all_chunks]
+        embeddings = self.embedder.encode(
+            texts, 
+            batch_size=32, 
+            convert_to_numpy=True, 
+            show_progress_bar=False
+        )
+        faiss.normalize_L2(embeddings)
+        
+        # Create new index
+        dim = embeddings.shape[1]
+        index = faiss.IndexFlatIP(dim)
+        index.add(embeddings)
+        
+        self.index = index
+        self.embeddings = embeddings
+        self.metadata = all_chunks
+        self.uploaded_file_active = True
+        
+        return len(all_chunks), file_type
+
     def load(self) -> bool:
         """Load pre-built index from disk."""
         if not (EMBEDDINGS_PATH.exists() and METADATA_PATH.exists() and FAISS_PATH.exists()):
@@ -145,6 +239,7 @@ class KBIndex:
         with open(METADATA_PATH, "r", encoding="utf-8") as f:
             self.metadata = json.load(f)
         self.index = faiss.read_index(str(FAISS_PATH))
+        self.uploaded_file_active = False
         return True
 
     def retrieve(self, query: str, top_k: int = 6) -> List[Tuple[int, float]]:
@@ -170,12 +265,18 @@ class KBIndex:
                 score = float(out.get("score", 0.0))
                 answer_text = out.get("answer", "").strip()
                 
-                if answer_text and len(answer_text) > 5:
+                # Enhanced answer extraction with context
+                if answer_text and len(answer_text) > 3:
+                    # Try to expand the answer with surrounding context
+                    expanded_answer = self._expand_answer(answer_text, ctx)
+                    
                     candidates.append({
-                        "text": answer_text,
+                        "text": expanded_answer,
+                        "original": answer_text,
                         "score": score,
                         "meta": meta,
-                        "sim": float(sim)
+                        "sim": float(sim),
+                        "context": ctx
                     })
             except Exception as e:
                 continue
@@ -204,6 +305,52 @@ class KBIndex:
         
         best_sim = max([s for _, s in retrieved]) if retrieved else 0.0
         return best["text"], best["score"], citations, best_sim
+    
+    def _expand_answer(self, answer: str, context: str, max_chars: int = 300) -> str:
+        """
+        Expand the extracted answer with surrounding context to make it more complete.
+        """
+        # Find the answer in the context
+        answer_pos = context.lower().find(answer.lower())
+        
+        if answer_pos == -1:
+            return answer
+        
+        # Get sentence boundaries around the answer
+        start = answer_pos
+        end = answer_pos + len(answer)
+        
+        # Expand backwards to sentence start
+        while start > 0 and context[start - 1] not in '.!?\n':
+            start -= 1
+            if answer_pos - start > max_chars // 2:
+                break
+        
+        # Expand forwards to sentence end
+        while end < len(context) and context[end] not in '.!?\n':
+            end += 1
+            if end - answer_pos > max_chars // 2:
+                break
+        
+        # Include the punctuation
+        if end < len(context) and context[end] in '.!?':
+            end += 1
+        
+        expanded = context[start:end].strip()
+        
+        # If still too short, try to get the full sentence(s)
+        if len(expanded) < 50:
+            # Look for complete sentences around the answer
+            sentences = context.split('.')
+            for i, sent in enumerate(sentences):
+                if answer.lower() in sent.lower():
+                    # Get this sentence and maybe the next one
+                    result = sent.strip()
+                    if i + 1 < len(sentences) and len(result) < 100:
+                        result += ". " + sentences[i + 1].strip()
+                    return result + ("." if not result.endswith(".") else "")
+        
+        return expanded
 
 # Initialize KB
 kb = KBIndex()
@@ -239,31 +386,47 @@ def format_citations(citations: List[Dict]) -> str:
         lines.append(f"• **{c['title']}** — _{c['section']}_")
     return "\n".join(lines)
 
-def respond(user_msg: str, history: List) -> str:
+def respond(user_msg: str, history: List, uploaded_file_info: str = None) -> str:
     """Generate response to user query using RAG pipeline."""
     user_msg = (user_msg or "").strip()
     
     if not user_msg:
         return "👋 How can I help? Ask me anything about the knowledge base, or use a quick action button below."
 
+    # Check if we have an index
+    if kb.index is None or len(kb.metadata) == 0:
+        return "❌ I don't know the answer to that but if you have any document with details I can learn about it. Please upload a file using the upload section above."
+
+    # Add context about uploaded file
+    source_info = f" in the uploaded file" if kb.uploaded_file_active and uploaded_file_info else " in the knowledge base"
+
     # Retrieve relevant chunks
     retrieved = kb.retrieve(user_msg, top_k=6)
     
-    if not retrieved:
-        return "❌ I couldn't find any relevant information. Please try rephrasing your question or contact support."
+    if not retrieved or (retrieved and max([s for _, s in retrieved]) < 0.20):
+        # Very low similarity - clearly don't know the answer
+        return (
+            f"❌ **I don't know the answer to that** but if you have any document with details I can learn about it.\n\n"
+            f"📤 Upload a relevant document above, and I'll be able to help you find the information you need!"
+        )
 
     # 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)
+    # Stricter threshold for "I don't know" response
+    if not answer or qa_score < 0.15 or best_sim < 0.25:
         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."
+            f"❌ **I don't know the answer to that** but if you have any document with details I can learn about it.\n\n"
+            f"The question seems outside the scope of what I currently know{source_info}. "
+            f"Try uploading a relevant document, or rephrase your question if you think the information might be here."
         )
 
+    # Clean up the answer text
+    answer = answer.strip()
+    # Ensure answer ends with proper punctuation
+    if answer and answer[-1] not in '.!?':
+        answer += "."
+
     # Check confidence
     low_confidence = (qa_score < CONFIDENCE_THRESHOLD) or (best_sim < SIMILARITY_THRESHOLD)
     citations_md = format_citations(citations)
@@ -271,39 +434,72 @@ def respond(user_msg: str, history: List) -> str:
     # Format response based on confidence
     if low_confidence:
         return (
-            f"⚠️ **Answer (Low Confidence):**\n{answer}\n\n"
+            f"⚠️ **Answer (Low Confidence):**\n\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.*"
+            f"💬 *I'm not entirely certain about this answer. If you have a more detailed document about this topic, please upload it for better accuracy.*"
         )
     else:
         return (
-            f"✅ **Answer:**\n{answer}\n\n"
+            f"✅ **Answer:**\n\n{answer}\n\n"
             f"---\n"
             f"📚 **Sources:**\n{citations_md}\n\n"
             f"💡 *Say \"show more details\" to see the full context.*"
         )
 
-def process_message(user_input: str, history: List) -> Tuple[List, Dict]:
+def process_message(user_input: str, history: List, uploaded_file_info: str) -> 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 [])
+    reply = respond(user_input, history or [], uploaded_file_info)
     new_history = (history or []) + [
         {"role": "user", "content": user_input},
         {"role": "assistant", "content": reply}
     ]
     return new_history, gr.update(value="")
 
-def process_quick(label: str, history: List) -> Tuple[List, Dict]:
+def process_quick(label: str, history: List, uploaded_file_info: str) -> 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 process_message(query, history, uploaded_file_info)
     return history, gr.update(value="")
 
+def handle_file_upload(file):
+    """Process uploaded file and build index."""
+    if file is None:
+        return "ℹ️ No file uploaded.", ""
+    
+    try:
+        filename = Path(file.name).name
+        num_chunks, file_type = kb.build_from_uploaded_file(file.name, filename)
+        
+        return (
+            f"✅ **File processed successfully!**\n\n"
+            f"📄 **File:** {filename}\n"
+            f"📋 **Type:** {file_type}\n"
+            f"🔢 **Chunks:** {num_chunks}\n\n"
+            f"You can now ask questions about this document!"
+        ), filename
+    except Exception as e:
+        return f"❌ **Error processing file:** {str(e)}\n\nPlease ensure the file is a valid PDF, DOCX, TXT, or MD file.", ""
+
+def clear_uploaded_file():
+    """Clear uploaded file and reload KB index."""
+    try:
+        if kb.load():
+            return "✅ Switched back to knowledge base.", "", None
+        else:
+            kb.index = None
+            kb.embeddings = None
+            kb.metadata = []
+            kb.uploaded_file_active = False
+            return "ℹ️ No knowledge base found. Please upload a file or build the KB index.", "", None
+    except Exception as e:
+        return f"⚠️ Error: {str(e)}", "", None
+
 def rebuild_index_handler():
     """Rebuild the search index from KB directory."""
     try:
@@ -319,18 +515,35 @@ with gr.Blocks(
     css="""
         .contain { max-width: 1200px; margin: auto; }
         .quick-btn { min-width: 180px !important; }
+        .upload-section { border: 2px dashed #ccc; padding: 20px; border-radius: 8px; }
     """
 ) as demo:
     
+    # State to track uploaded file
+    uploaded_file_state = gr.State("")
+    
     # Header
     gr.Markdown(
         """
         # 🤖 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.
+        Upload a document or use the knowledge base to get answers backed by relevant sources.
         """
     )
     
+    # File upload section
+    with gr.Row():
+        with gr.Column(scale=1):
+            gr.Markdown("### 📤 Upload Document")
+            file_upload = gr.File(
+                label="Upload PDF, DOCX, TXT, or MD file",
+                file_types=[".pdf", ".docx", ".txt", ".md"],
+                type="filepath"
+            )
+            upload_status = gr.Markdown("ℹ️ Upload a file to ask questions about it.")
+            with gr.Row():
+                clear_btn = gr.Button("🔄 Clear & Use KB", variant="secondary", size="sm")
+    
     # Main chat interface
     with gr.Row():
         with gr.Column(scale=1):
@@ -343,20 +556,20 @@ with gr.Blocks(
             
             with gr.Row():
                 txt = gr.Textbox(
-                    placeholder="💬 Ask a question (e.g., How do I connect WhatsApp?)",
+                    placeholder="💬 Ask a question about the document or knowledge base...",
                     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():
-        quick_buttons = []
-        for label, _ in QUICK_ACTIONS:
-            btn = gr.Button(label, elem_classes="quick-btn", size="sm")
-            quick_buttons.append((btn, label))
+    # Quick action buttons (only for KB mode)
+    with gr.Accordion("⚡ Quick Actions (Knowledge Base)", open=False):
+        with gr.Row():
+            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):
@@ -367,17 +580,36 @@ with gr.Blocks(
             """
         )
         with gr.Row():
-            rebuild_btn = gr.Button("🔄 Rebuild Index", variant="secondary")
+            rebuild_btn = gr.Button("🔄 Rebuild KB 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])
+    file_upload.change(
+        handle_file_upload,
+        inputs=[file_upload],
+        outputs=[upload_status, uploaded_file_state]
+    )
+    
+    clear_btn.click(
+        clear_uploaded_file,
+        outputs=[upload_status, uploaded_file_state, file_upload]
+    )
+    
+    send.click(
+        process_message, 
+        inputs=[txt, chat, uploaded_file_state], 
+        outputs=[chat, txt]
+    )
+    txt.submit(
+        process_message, 
+        inputs=[txt, chat, uploaded_file_state], 
+        outputs=[chat, txt]
+    )
     
     for btn, label in quick_buttons:
         btn.click(
             process_quick, 
-            inputs=[gr.State(label), chat], 
+            inputs=[gr.State(label), chat, uploaded_file_state], 
             outputs=[chat, txt]
         )
     
@@ -388,9 +620,10 @@ with gr.Blocks(
         """
         ---
         💡 **Tips:**
+        - Upload a document to ask questions specifically about that file
+        - Use "Clear & Use KB" to switch back to the knowledge base
         - Be specific in your questions for better results
         - Check the cited sources for full context
-        - Use quick actions for common tasks
         """
     )