Spaces:
Running on Zero
Running on Zero
| 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.") | |