Create rag_engine.py

rag_engine.py

@@ -0,0 +1,146 @@
+from __future__ import annotations
+import re
+import textwrap
+from pathlib import Path
+from typing import Any
+import faiss
+import numpy as np
+import requests
+import spacy
+from bs4 import BeautifulSoup
+from huggingface_hub import InferenceClient
+from pypdf import PdfReader
+from sentence_transformers import SentenceTransformer
+
+# ── Config ────────────────────────────────────────────────────────────────────
+EMBED_MODEL   = "sentence-transformers/all-MiniLM-L6-v2"
+LLM_MODEL     = "HuggingFaceH4/zephyr-7b-beta"
+CHUNK_SIZE    = 400
+CHUNK_OVERLAP = 80
+TOP_K         = 4
+
+INTENT_MAP = {
+    "summarise": ["summarise", "summarize", "summary", "overview", "brief", "key points"],
+    "explain": ["explain", "what is", "what are", "define", "describe", "tell me about"],
+    "list": ["list", "enumerate", "give me a list", "what are the types"],
+}
+
+class RAGEngine:
+    def __init__(self):
+        self.embed_model = SentenceTransformer(EMBED_MODEL)
+        self.hf_client = InferenceClient()
+        try:
+            self.nlp = spacy.load("en_core_web_sm")
+        except:
+            # Fallback if model not downloaded
+            import os
+            os.system("python -m spacy download en_core_web_sm")
+            self.nlp = spacy.load("en_core_web_sm")
+        
+        self.reset()
+
+    def reset(self):
+        self.chunks: list[str] = []
+        self.index: faiss.IndexFlatL2 | None = None
+        self.ready = False
+
+    def _process_text_into_chunks(self, text: str):
+        text = re.sub(r'\s+', ' ', text).strip()
+        new_chunks = []
+        for i in range(0, len(text), CHUNK_SIZE - CHUNK_OVERLAP):
+            chunk = text[i : i + CHUNK_SIZE]
+            if len(chunk) > 20:
+                new_chunks.append(chunk)
+        
+        self.chunks = new_chunks
+        embeddings = self.embed_model.encode(self.chunks)
+        self.index = faiss.IndexFlatL2(embeddings.shape[1])
+        self.index.add(np.array(embeddings).astype("float32"))
+        self.ready = True
+
+    def load_pdf(self, path: str) -> str:
+        reader = PdfReader(path)
+        text = "".join([page.extract_text() for page in reader.pages])
+        self._process_text_into_chunks(text)
+        return f"✅ Loaded PDF: {len(self.chunks)} chunks indexed."
+
+    def load_url(self, url: str) -> str:
+        res = requests.get(url, timeout=10)
+        soup = BeautifulSoup(res.text, "html.parser")
+        for script in soup(["script", "style"]):
+            script.decompose()
+        self._process_text_into_chunks(soup.get_text())
+        return f"✅ Loaded URL: {len(self.chunks)} chunks indexed."
+
+    def load_text(self, text: str) -> str:
+        self._process_text_into_chunks(text)
+        return f"✅ Loaded Text: {len(self.chunks)} chunks indexed."
+
+    def detect_intent(self, query: str) -> str:
+        query_lower = query.lower()
+        for intent, keywords in INTENT_MAP.items():
+            if any(k in query_lower for k in keywords):
+                return intent
+        return "general_query"
+
+    def extract_entities(self, text: str) -> dict[str, list[str]]:
+        doc = self.nlp(text)
+        entities = {}
+        for ent in doc.ents:
+            if ent.label_ not in entities:
+                entities[ent.label_] = []
+            if ent.text not in entities[ent.label_]:
+                entities[ent.label_].append(ent.text)
+        return entities
+
+    def _retrieve(self, query: str) -> list[str]:
+        query_vec = self.embed_model.encode([query]).astype("float32")
+        _, indices = self.index.search(query_vec, TOP_K)
+        return [self.chunks[i] for i in indices[0]]
+
+    def _build_prompt(self, query: str, chunks: list[str]) -> str:
+        context = "\n".join([f"- {c}" for c in chunks])
+        return textwrap.dedent(f"""
+            <|system|>
+            You are a helpful assistant. Use the following context to answer the user's question.
+            If the answer is not in the context, say you don't know based on the provided data.
+            </s>
+            <|user|>
+            Context:
+            {context}
+
+            Question: {query}
+            </s>
+            <|assistant|>
+        """).strip()
+
+    def answer(self, query: str) -> str:
+        if not self.ready:
+            return "⚠️ No knowledge source loaded. Please load a PDF, URL, or text first."
+        chunks = self._retrieve(query)
+        prompt = self._build_prompt(query, chunks)
+        try:
+            response = self.hf_client.text_generation(
+                prompt,
+                model=LLM_MODEL,
+                max_new_tokens=512,
+                temperature=0.3,
+                repetition_penalty=1.1,
+                stop_sequences=["</s>", "<|user|>"],
+            )
+            answer_text = response.strip()
+        except Exception as e:
+            answer_text = f"⚠️ LLM API unavailable ({e}).\n\n**Most relevant passage found:**\n\n{chunks[0]}"
+        return answer_text
+
+    def answer_with_nlp(self, query: str) -> tuple[str, dict[str, Any]]:
+        answer_text = self.answer(query)
+        intent = self.detect_intent(query)
+        entities_in_query = self.extract_entities(query)
+        entities_in_answer = self.extract_entities(answer_text)
+        nlp_info = {
+            "detected_intent": intent,
+            "entities_in_query": entities_in_query,
+            "entities_in_answer": entities_in_answer,
+        }
+        return answer_text, nlp_info