Create app.py

app.py

@@ -0,0 +1,212 @@
+import os
+import glob
+import yaml
+from typing import List, Tuple
+
+import faiss
+import numpy as np
+import gradio as gr
+from sentence_transformers import SentenceTransformer
+from transformers import AutoTokenizer, AutoModelForQuestionAnswering, pipeline
+from PyPDF2 import PdfReader
+import docx
+
+
+# -----------------------------
+# CONFIG
+# -----------------------------
+
+with open("config.yaml", "r", encoding="utf-8") as f:
+    CONFIG = yaml.safe_load(f)
+
+KB_DIR = CONFIG["kb"]["directory"]
+INDEX_DIR = CONFIG["kb"]["index_directory"]
+EMBEDDING_MODEL_NAME = CONFIG["models"]["embedding"]
+QA_MODEL_NAME = CONFIG["models"]["qa"]
+CHUNK_SIZE = CONFIG["chunking"]["chunk_size"]
+CHUNK_OVERLAP = CONFIG["chunking"]["overlap"]
+SIM_THRESHOLD = CONFIG["thresholds"]["similarity"]
+WELCOME_MSG = CONFIG["messages"]["welcome"]
+NO_ANSWER_MSG = CONFIG["messages"]["no_answer"]
+
+
+# -----------------------------
+# UTILITIES
+# -----------------------------
+
+def chunk_text(text: str, chunk_size: int, overlap: int) -> List[str]:
+    if not text:
+        return []
+    chunks = []
+    start = 0
+    while start < len(text):
+        end = min(start + chunk_size, len(text))
+        chunk = text[start:end].strip()
+        if chunk:
+            chunks.append(chunk)
+        start += chunk_size - overlap
+    return chunks
+
+
+def load_file_text(path: str) -> str:
+    ext = os.path.splitext(path)[1].lower()
+    if ext == ".pdf":
+        reader = PdfReader(path)
+        return "\n".join(page.extract_text() or "" for page in reader.pages)
+    elif ext in [".docx", ".doc"]:
+        doc = docx.Document(path)
+        return "\n".join(p.text for p in doc.paragraphs)
+    else:
+        with open(path, "r", encoding="utf-8") as f:
+            return f.read()
+
+
+def load_kb_documents(kb_dir: str) -> List[Tuple[str, str]]:
+    docs = []
+    if os.path.isdir(kb_dir):
+        paths = glob.glob(os.path.join(kb_dir, "*.txt")) \
+              + glob.glob(os.path.join(kb_dir, "*.md")) \
+              + glob.glob(os.path.join(kb_dir, "*.pdf")) \
+              + glob.glob(os.path.join(kb_dir, "*.docx"))
+        for path in paths:
+            try:
+                text = load_file_text(path)
+                if text.strip():
+                    docs.append((os.path.basename(path), text))
+            except Exception as e:
+                print(f"Could not read {path}: {e}")
+    return docs
+
+
+# -----------------------------
+# KB INDEX (FAISS)
+# -----------------------------
+
+class RAGIndex:
+    def __init__(self):
+        print("Loading embedding model...")
+        self.embedder = SentenceTransformer(EMBEDDING_MODEL_NAME)
+        print("Loading QA model...")
+        self.qa_pipeline = pipeline(
+            "question-answering",
+            model=AutoModelForQuestionAnswering.from_pretrained(QA_MODEL_NAME),
+            tokenizer=AutoTokenizer.from_pretrained(QA_MODEL_NAME),
+            handle_impossible_answer=True,
+        )
+        self.chunks: List[str] = []
+        self.chunk_sources: List[str] = []
+        self.index = None
+        self._build_or_load_index()
+
+    def _build_or_load_index(self):
+        os.makedirs(INDEX_DIR, exist_ok=True)
+        idx_path = os.path.join(INDEX_DIR, "kb.index")
+        meta_path = os.path.join(INDEX_DIR, "kb_meta.npy")
+
+        if os.path.exists(idx_path) and os.path.exists(meta_path):
+            print("Loading existing FAISS index...")
+            self.index = faiss.read_index(idx_path)
+            meta = np.load(meta_path, allow_pickle=True).item()
+            self.chunks = meta["chunks"]
+            self.chunk_sources = meta["sources"]
+            print("Index loaded.")
+            return
+
+        print("Building new FAISS index...")
+        docs = load_kb_documents(KB_DIR)
+        all_chunks = []
+        all_sources = []
+        for source, text in docs:
+            for chunk in chunk_text(text, CHUNK_SIZE, CHUNK_OVERLAP):
+                all_chunks.append(chunk)
+                all_sources.append(source)
+
+        if not all_chunks:
+            print("⚠️ No KB documents found, index will stay empty.")
+            self.index = None
+            return
+
+        embeddings = self.embedder.encode(all_chunks, show_progress_bar=True, convert_to_numpy=True)
+        dimension = embeddings.shape[1]
+        index = faiss.IndexFlatIP(dimension)
+
+        # Normalize for cosine similarity
+        faiss.normalize_L2(embeddings)
+        index.add(embeddings)
+
+        faiss.write_index(index, idx_path)
+        np.save(meta_path, {"chunks": np.array(all_chunks, dtype=object), "sources": np.array(all_sources, dtype=object)})
+
+        self.index = index
+        self.chunks = all_chunks
+        self.chunk_sources = all_sources
+        print("FAISS index ready.")
+
+    def retrieve(self, query: str, top_k: int = 5) -> List[Tuple[str, str, float]]:
+        if not query.strip() or self.index is None:
+            return []
+        q_emb = self.embedder.encode([query], convert_to_numpy=True)
+        faiss.normalize_L2(q_emb)
+        scores, idxs = self.index.search(q_emb, top_k)
+        results = []
+        for score, idx in zip(scores[0], idxs[0]):
+            if idx == -1:
+                continue
+            if score < SIM_THRESHOLD:
+                continue
+            results.append((self.chunks[idx], self.chunk_sources[idx], float(score)))
+        return results
+
+    def answer(self, question: str) -> str:
+        contexts = self.retrieve(question, top_k=3)
+        if not contexts:
+            return NO_ANSWER_MSG
+
+        answers = []
+        for ctx, source, score in contexts:
+            qa_input = {"question": question, "context": ctx}
+            try:
+                result = self.qa_pipeline(qa_input)
+                text = result.get("answer", "").strip()
+                if text:
+                    answers.append((text, source, result.get("score", 0.0)))
+            except Exception as e:
+                print(f"QA error: {e}")
+
+        if not answers:
+            return NO_ANSWER_MSG
+
+        # Pick best answer
+        answers.sort(key=lambda x: x[2], reverse=True)
+        best_answer, best_source, best_score = answers[0]
+
+        return (
+            f"**Answer:** {best_answer}\n\n"
+            f"**Source:** {best_source}  (confidence: {best_score:.2f})"
+        )
+
+
+rag_index = RAGIndex()
+
+
+# -----------------------------
+# GRADIO CHAT
+# -----------------------------
+
+def rag_respond(message: str, history):
+    return rag_index.answer(message)
+
+
+description = CONFIG["messages"]["welcome"]
+
+chat = gr.ChatInterface(
+    fn=rag_respond,
+    title=CONFIG["client"]["name"],
+    description=description,
+    type="messages",
+    examples=[qa["query"] for qa in CONFIG.get("quick_actions", [])],
+    cache_examples=False,
+)
+
+if __name__ == "__main__":
+    chat.launch()