Upload rag_engine.py with huggingface_hub

rag_engine.py

@@ -0,0 +1,441 @@
+import os
+import re
+import time
+import math
+from collections import Counter
+import numpy as np
+
+try:
+    from pypdf import PdfReader
+    HAS_PYPDF = True
+except ImportError:
+    HAS_PYPDF = False
+
+try:
+    import docx
+    HAS_DOCX = True
+except ImportError:
+    HAS_DOCX = False
+
+try:
+    import faiss
+    HAS_FAISS = True
+except ImportError:
+    HAS_FAISS = False
+
+try:
+    from sentence_transformers import SentenceTransformer
+    HAS_SENTENCE_TRANSFORMERS = True
+except ImportError:
+    HAS_SENTENCE_TRANSFORMERS = False
+
+try:
+    from groq import Groq
+    HAS_GROQ = True
+except ImportError:
+    HAS_GROQ = False
+
+try:
+    import httpx
+    HAS_HTTPX = True
+except ImportError:
+    HAS_HTTPX = False
+
+try:
+    from dotenv import load_dotenv
+    load_dotenv()
+except ImportError:
+    pass
+
+
+class FastFallbackEmbedder:
+    """Fast TF-IDF / character-word n-gram fallback embedder when SentenceTransformers is not yet loaded."""
+    def __init__(self, dim=384):
+        self.dim = dim
+
+    def encode(self, texts, **kwargs):
+        if isinstance(texts, str):
+            texts = [texts]
+        embeddings = []
+        for text in texts:
+            words = re.findall(r"\w+", text.lower())
+            vec = np.zeros(self.dim, dtype=np.float32)
+            if not words:
+                embeddings.append(vec)
+                continue
+            for i, w in enumerate(words):
+                h = hash(w) % self.dim
+                vec[h] += 1.0 / (1.0 + math.log(1 + i))
+            norm = np.linalg.norm(vec)
+            if norm > 0:
+                vec /= norm
+            embeddings.append(vec)
+        return np.array(embeddings, dtype=np.float32)
+
+
+class RagEngine:
+    def __init__(self):
+        self.chunks = []
+        self.embeddings = None
+        self.faiss_index = None
+        self.embedding_dim = 384
+        self.embed_model = None
+        self.audit_log = []
+
+        # Load SentenceTransformer model
+        if HAS_SENTENCE_TRANSFORMERS:
+            try:
+                print("[RagEngine] Loading embedding model 'all-MiniLM-L6-v2'...")
+                self.embed_model = SentenceTransformer("all-MiniLM-L6-v2")
+                self.embedding_dim = 384
+            except Exception as e:
+                print(f"[RagEngine] Could not load SentenceTransformer: {e}. Using fallback embedder.")
+                self.embed_model = FastFallbackEmbedder(self.embedding_dim)
+        else:
+            print("[RagEngine] sentence-transformers not installed. Using fallback embedder.")
+            self.embed_model = FastFallbackEmbedder(self.embedding_dim)
+
+        self._init_vector_store()
+
+        # Groq client
+        self.groq_api_key = os.environ.get("GROQ_API_KEY", "").strip()
+        self.groq_client = None
+        if self.groq_api_key and HAS_GROQ:
+            try:
+                self.groq_client = Groq(api_key=self.groq_api_key)
+            except Exception as e:
+                print(f"[RagEngine] Groq client init failed: {e}")
+
+    def _init_vector_store(self):
+        if HAS_FAISS:
+            self.faiss_index = faiss.IndexFlatL2(self.embedding_dim)
+        else:
+            self.faiss_index = None
+
+    def _extract_text_and_pages(self, filepath):
+        """Extract text with page numbers from PDF, DOCX, or TXT."""
+        pages = []
+        ext = os.path.splitext(filepath)[1].lower()
+
+        if ext == ".pdf" and HAS_PYPDF:
+            reader = PdfReader(filepath)
+            for i, page in enumerate(reader.pages):
+                text = page.extract_text() or ""
+                if text.strip():
+                    pages.append((i + 1, text))
+        elif ext == ".docx" and HAS_DOCX:
+            doc = docx.Document(filepath)
+            full_text = "\n".join([p.text for p in doc.paragraphs if p.text.strip()])
+            pages.append((1, full_text))
+        elif ext == ".csv":
+            try:
+                import pandas as pd
+                df = pd.read_csv(filepath)
+                cols = list(df.columns)
+                rows_text = []
+                for idx, row in df.iterrows():
+                    row_desc = ", ".join([f"{col}: {row[col]}" for col in cols if pd.notna(row[col])])
+                    rows_text.append(f"Row {idx+1} -> {row_desc}")
+                
+                chunk_size = 15
+                if not rows_text:
+                    pages.append((1, f"CSV File with Columns: {', '.join(map(str, cols))} (No data rows)"))
+                else:
+                    for page_idx, i in enumerate(range(0, len(rows_text), chunk_size), 1):
+                        batch = "\n".join(rows_text[i : i + chunk_size])
+                        pages.append((page_idx, f"CSV File Columns: {', '.join(map(str, cols))}\n\nData:\n{batch}"))
+            except Exception as e:
+                with open(filepath, "r", encoding="utf-8", errors="ignore") as f:
+                    pages.append((1, f.read()))
+
+        elif ext in [".xlsx", ".xls"]:
+            try:
+                import pandas as pd
+                excel_file = pd.ExcelFile(filepath)
+                for sheet_idx, sheet_name in enumerate(excel_file.sheet_names, 1):
+                    df = pd.read_excel(filepath, sheet_name=sheet_name)
+                    cols = list(df.columns)
+                    rows_text = []
+                    for idx, row in df.iterrows():
+                        row_desc = ", ".join([f"{col}: {row[col]}" for col in cols if pd.notna(row[col])])
+                        rows_text.append(f"Row {idx+1} -> {row_desc}")
+                    
+                    chunk_size = 15
+                    if not rows_text:
+                        pages.append((sheet_idx, f"Excel Sheet: '{sheet_name}' | Columns: {', '.join(map(str, cols))} (Empty Sheet)"))
+                    else:
+                        for sub_idx, i in enumerate(range(0, len(rows_text), chunk_size), 1):
+                            batch = "\n".join(rows_text[i : i + chunk_size])
+                            pages.append((sheet_idx, f"Excel Sheet: '{sheet_name}' | Columns: {', '.join(map(str, cols))}\n\nData:\n{batch}"))
+            except Exception as e:
+                pages.append((1, f"Excel file: {os.path.basename(filepath)}"))
+
+        else:
+            encodings = ["utf-8", "latin-1", "cp1252"]
+            content = ""
+            for enc in encodings:
+                try:
+                    with open(filepath, "r", encoding=enc) as f:
+                        content = f.read()
+                    break
+                except Exception:
+                    continue
+            pages.append((1, content))
+
+        return pages
+
+    def _chunk_text(self, text, max_words=120, overlap_words=25):
+        """Split text into overlapping chunks."""
+        words = text.split()
+        if not words:
+            return []
+        chunks = []
+        i = 0
+        while i < len(words):
+            chunk = " ".join(words[i : i + max_words])
+            chunks.append(chunk)
+            i += max_words - overlap_words
+        return chunks
+
+    def ingest_document(self, filepath, doc_role="public"):
+        """Ingest document, chunk text, generate embeddings, and update index."""
+        if not os.path.exists(filepath):
+            raise FileNotFoundError(f"File not found: {filepath}")
+
+        filename = os.path.basename(filepath)
+        pages = self._extract_text_and_pages(filepath)
+        if not pages:
+            raise ValueError(f"Could not extract any text from {filename}")
+
+        new_chunks = []
+        new_texts = []
+        is_stale_file = "stale" in filename.lower() or "archive" in filename.lower()
+        chunk_timestamp = (time.time() - 200 * 86400) if is_stale_file else time.time()
+
+        for page_num, page_text in pages:
+            text_chunks = self._chunk_text(page_text)
+            for t_chunk in text_chunks:
+                chunk_obj = {
+                    "text": t_chunk,
+                    "source": filename,
+                    "page": page_num,
+                    "doc_role": doc_role,
+                    "timestamp": chunk_timestamp,
+                    "images": [],
+                }
+                new_chunks.append(chunk_obj)
+                new_texts.append(t_chunk)
+
+        if not new_texts:
+            raise ValueError(f"No valid text chunks generated from {filename}")
+
+        # Compute embeddings
+        embs = self.embed_model.encode(new_texts)
+        if isinstance(embs, list):
+            embs = np.array(embs, dtype=np.float32)
+        elif not isinstance(embs, np.ndarray):
+            embs = embs.cpu().numpy().astype(np.float32)
+        else:
+            embs = embs.astype(np.float32)
+
+        # Update Vector Index
+        if HAS_FAISS and self.faiss_index is not None:
+            self.faiss_index.add(embs)
+
+        if self.embeddings is None or len(self.embeddings) == 0:
+            self.embeddings = embs
+        else:
+            self.embeddings = np.vstack([self.embeddings, embs])
+
+        self.chunks.extend(new_chunks)
+
+    def get_stats(self):
+        sources = sorted(list(set(c["source"] for c in self.chunks)))
+        return {
+            "total_chunks": len(self.chunks),
+            "unique_sources": len(sources),
+            "sources": sources,
+        }
+
+    def _keyword_search(self, query_tokens, candidate_indices):
+        """Source-aware keyword search returning ranked candidate indices."""
+        scores = {}
+        for idx in candidate_indices:
+            chunk = self.chunks[idx]
+            text = chunk["text"].lower()
+            source = chunk["source"].lower()
+            score = 0
+            for token in query_tokens:
+                if token in source:
+                    score += 3.0
+                if token in text:
+                    score += 1.0
+            if score > 0:
+                scores[idx] = score
+
+        ranked = sorted(scores.items(), key=lambda x: x[1], reverse=True)
+        return [idx for idx, _ in ranked]
+
+    def query(self, message, top_k=8, user_role=None):
+        """Hybrid search combining semantic search & keyword search with Reciprocal Rank Fusion (RRF)."""
+        if not self.chunks:
+            return []
+
+        candidate_indices = list(range(len(self.chunks)))
+
+        # 1. Semantic Search
+        q_emb = self.embed_model.encode([message])
+        if isinstance(q_emb, list):
+            q_emb = np.array(q_emb, dtype=np.float32)
+        elif not isinstance(q_emb, np.ndarray):
+            q_emb = q_emb.cpu().numpy().astype(np.float32)
+        else:
+            q_emb = q_emb.astype(np.float32)
+
+        cand_embs = self.embeddings[candidate_indices]
+        dists = np.linalg.norm(cand_embs - q_emb, axis=1)
+        sem_ranked_order = np.argsort(dists)
+        sem_ranked_indices = [candidate_indices[i] for i in sem_ranked_order]
+
+        # 2. Keyword Search
+        query_tokens = [
+            t.lower()
+            for t in re.findall(r"\w+", message)
+            if len(t) > 2 and t.lower() not in {"what", "how", "why", "when", "where", "the", "and", "for"}
+        ]
+        kw_ranked_indices = self._keyword_search(query_tokens, candidate_indices)
+
+        # 3. Reciprocal Rank Fusion (RRF, k=60)
+        rrf_scores = Counter()
+        for rank, idx in enumerate(sem_ranked_indices):
+            rrf_scores[idx] += 1.0 / (60.0 + rank + 1)
+
+        for rank, idx in enumerate(kw_ranked_indices):
+            rrf_scores[idx] += 1.0 / (60.0 + rank + 1)
+
+        top_candidates = rrf_scores.most_common(top_k)
+
+        results = []
+        retrieved_sources = set()
+        for idx, score in top_candidates:
+            chunk = self.chunks[idx]
+            retrieved_sources.add(chunk["source"])
+            results.append({
+                "source": chunk["source"],
+                "page": chunk.get("page", 1),
+                "score": float(score),
+                "doc_role": chunk.get("doc_role", "public"),
+                "images": chunk.get("images", []),
+                "text": chunk["text"],
+            })
+
+        # Record in audit log
+        self.audit_log.append({
+            "timestamp": time.time(),
+            "question": message,
+            "sources_retrieved": sorted(list(retrieved_sources)),
+        })
+
+        return results
+
+    def generate_answer(self, message, chunks):
+        """Generate direct, beautifully formatted answer with citations using Groq Qwen 3.6 27B."""
+        if not chunks:
+            return "No matching context found."
+
+        # Prepare context blocks with source numbering
+        context_parts = []
+        for i, c in enumerate(chunks, 1):
+            context_parts.append(
+                f"[Source {i}: {c['source']} | Page {c['page']}]\n{c['text']}"
+            )
+        context_str = "\n\n".join(context_parts)
+
+        api_key = os.environ.get("GROQ_API_KEY", self.groq_api_key or "").strip()
+        model_name = os.environ.get("GROQ_MODEL", "qwen/qwen3.8-27b").strip()
+
+        if api_key:
+            candidate_models = [model_name, "qwen/qwen3.8-27b", "openai/gpt-oss-120b", "openai/gpt-oss-20b", "groq/compound"]
+            # remove duplicates while preserving order
+            unique_models = []
+            for m in candidate_models:
+                if m not in unique_models:
+                    unique_models.append(m)
+
+            for target_model in unique_models:
+                try:
+                    system_prompt = (
+                        "You are the Universal RAG AI Knowledge Assistant. "
+                        "Your goal is to provide rich, comprehensive, beautifully structured answers strictly grounded in the provided context.\n\n"
+                        "FORMATTING RULES:\n"
+                        "1. When answering about resumes/candidates: Clearly separate each candidate using headers like '👤 Candidate: [Name]'. "
+                        "Organize their details into '📚 Technical Skills & Core Tools [Source X, Page Y]' (grouped by Languages, Databases, Cloud, Frameworks) "
+                        "and '📁 Main Projects & Associated Technologies' or '🔹 Project X: [Title]' with 'Technologies Used' and 'Key Contributions & Details'.\n"
+                        "2. When answering about tabular CSV/Excel files: Extract specific rows, values, metrics, and comparisons in clean Markdown tables and bullet points.\n"
+                        "3. When answering about research papers or technical reports: Clearly break down methodology, pipeline, dataset, and key benchmarks/metrics.\n"
+                        "4. When referencing information, always cite the source (e.g., '[Source 1, Page 1]' or 'Source: Source 1, Page 1').\n"
+                        "5. Maintain high fidelity to the text (exact libraries, versions, metrics, algorithms).\n"
+                        "6. End with a polite, helpful follow-up note (e.g., 'Let me know if you need a comparison with other candidates or deeper technical breakdowns!')."
+                    )
+                    user_prompt = (
+                        f"CONTEXT INFORMATION:\n{context_str}\n\n"
+                        f"USER QUESTION: {message}\n\n"
+                        f"Provide a structured, complete answer following the formatting rules."
+                    )
+
+                    if HAS_GROQ:
+                        client = self.groq_client or Groq(api_key=api_key)
+                        try:
+                            completion = client.chat.completions.create(
+                                model=target_model,
+                                messages=[
+                                    {"role": "system", "content": system_prompt},
+                                    {"role": "user", "content": user_prompt},
+                                ],
+                                temperature=0.2,
+                                max_tokens=850,
+                                reasoning_format="hidden",
+                            )
+                        except Exception:
+                            completion = client.chat.completions.create(
+                                model=target_model,
+                                messages=[
+                                    {"role": "system", "content": system_prompt},
+                                    {"role": "user", "content": user_prompt},
+                                ],
+                                temperature=0.2,
+                                max_tokens=850,
+                            )
+                        raw_content = completion.choices[0].message.content or ""
+                        clean_content = re.sub(r"<think>.*?</think>", "", raw_content, flags=re.DOTALL).strip()
+                        if clean_content or raw_content:
+                            return clean_content or raw_content
+
+                except Exception as e:
+                    print(f"[RagEngine] Model {target_model} error: {e}")
+                    continue
+
+        # Fallback local synthesis
+        synthesis = "Based on the relevant documents indexed in the knowledge base:\n\n"
+        for i, c in enumerate(chunks[:3], 1):
+            synthesis += f"> **[{c['source']} - Page {c['page']}]**\n> \"{c['text'][:280]}...\"\n\n"
+        return synthesis
+
+    def get_audit_summary(self):
+        """Governance summary: chunks by role/department, staleness, and query audit log."""
+        chunks_by_role = Counter([c.get("doc_role", "public") for c in self.chunks])
+        now = time.time()
+        stale_count = sum(1 for c in self.chunks if (now - c.get("timestamp", now)) > (180 * 86400))
+        return {
+            "chunks_by_role": dict(chunks_by_role),
+            "total_queries": len(self.audit_log),
+            "stale_chunk_count": stale_count,
+            "recent_log": self.audit_log[-15:],
+        }
+
+    def clear_index(self):
+        """Clear all stored chunks and reset vector index."""
+        self.chunks = []
+        self.embeddings = None
+        self._init_vector_store()
+        print("[RagEngine] Vector index and chunks cleared.")