Spaces:
Running
Running
| import os | |
| import io | |
| import json | |
| import re | |
| import torch | |
| import httpx | |
| from bs4 import BeautifulSoup | |
| from typing import List, Dict, Any, Optional | |
| from sentence_transformers import SentenceTransformer | |
| from ddgs import DDGS | |
| import pypdf | |
| KNOWLEDGE_DIR = os.path.join(os.path.dirname(__file__), "knowledge_base") | |
| os.makedirs(KNOWLEDGE_DIR, exist_ok=True) | |
| INDEX_PATH = os.path.join(KNOWLEDGE_DIR, "index.pt") | |
| METADATA_PATH = os.path.join(KNOWLEDGE_DIR, "metadata.json") | |
| class RAGEngine: | |
| def __init__(self, model_name: str = "all-MiniLM-L6-v2"): | |
| self.device = "mps" if torch.backends.mps.is_available() else "cuda" if torch.cuda.is_available() else "cpu" | |
| print(f"⚡ Initializing RAG Embedding Engine on {self.device}...") | |
| self.encoder = SentenceTransformer(model_name, device=self.device) | |
| self.chunks: List[Dict[str, Any]] = [] # [{ "text": str, "source": str, "chunk_id": int }] | |
| self.embeddings: Optional[torch.Tensor] = None # Tensor [N, 384] | |
| self.load_index() | |
| def load_index(self): | |
| if os.path.exists(INDEX_PATH) and os.path.exists(METADATA_PATH): | |
| try: | |
| self.embeddings = torch.load(INDEX_PATH, map_location=self.device, weights_only=True) | |
| with open(METADATA_PATH, "r", encoding="utf-8") as f: | |
| self.chunks = json.load(f) | |
| print(f"📚 Loaded {len(self.chunks)} document chunks from disk cache.") | |
| except Exception as e: | |
| print(f"⚠️ Failed to load RAG cache: {e}") | |
| self.chunks = [] | |
| self.embeddings = None | |
| def save_index(self): | |
| try: | |
| if self.embeddings is not None and len(self.chunks) > 0: | |
| torch.save(self.embeddings, INDEX_PATH) | |
| with open(METADATA_PATH, "w", encoding="utf-8") as f: | |
| json.dump(self.chunks, f, indent=2) | |
| print(f"💾 Saved {len(self.chunks)} chunks to knowledge base.") | |
| else: | |
| if os.path.exists(INDEX_PATH): os.remove(INDEX_PATH) | |
| if os.path.exists(METADATA_PATH): os.remove(METADATA_PATH) | |
| except Exception as e: | |
| print(f"⚠️ Failed to save RAG cache: {e}") | |
| def chunk_text(self, text: str, source: str, chunk_size: int = 500, overlap: int = 80) -> List[Dict[str, Any]]: | |
| cleaned = re.sub(r'\s+', ' ', text).strip() | |
| if not cleaned: | |
| return [] | |
| chunks = [] | |
| start = 0 | |
| chunk_id = 0 | |
| text_len = len(cleaned) | |
| while start < text_len: | |
| end = min(start + chunk_size, text_len) | |
| # Try to break at nearest space or punctuation if not at the end | |
| if end < text_len: | |
| last_space = cleaned.rfind(' ', start, end) | |
| if last_space > start + (chunk_size // 2): | |
| end = last_space | |
| chunk_slice = cleaned[start:end].strip() | |
| if len(chunk_slice) > 20: | |
| chunks.append({ | |
| "text": chunk_slice, | |
| "source": source, | |
| "chunk_id": chunk_id | |
| }) | |
| chunk_id += 1 | |
| start = end - overlap if end < text_len else text_len | |
| return chunks | |
| def ingest_file(self, filename: str, file_bytes: bytes) -> Dict[str, Any]: | |
| extracted_text = "" | |
| ext = os.path.splitext(filename)[1].lower() | |
| if ext == ".pdf": | |
| try: | |
| reader = pypdf.PdfReader(io.BytesIO(file_bytes)) | |
| pages_text = [] | |
| for i, page in enumerate(reader.pages): | |
| t = page.extract_text() | |
| if t: | |
| pages_text.append(f"[Page {i+1}] {t}") | |
| extracted_text = "\n".join(pages_text) | |
| except Exception as e: | |
| raise ValueError(f"Failed to extract text from PDF '{filename}': {e}") | |
| else: | |
| # Plaintext, Markdown, Python, Code, JSON, CSV | |
| try: | |
| extracted_text = file_bytes.decode("utf-8") | |
| except UnicodeDecodeError: | |
| extracted_text = file_bytes.decode("latin-1", errors="ignore") | |
| if not extracted_text.strip(): | |
| raise ValueError(f"No extractable text found in '{filename}'") | |
| # Remove existing chunks for this file if re-uploading | |
| self.delete_document(filename, auto_save=False) | |
| new_chunks = self.chunk_text(extracted_text, source=filename) | |
| if not new_chunks: | |
| raise ValueError("Document was too short to create meaningful chunks.") | |
| texts = [c["text"] for c in new_chunks] | |
| new_embeddings = self.encoder.encode(texts, batch_size=32, convert_to_tensor=True, show_progress_bar=False).to(self.device) | |
| # Normalize for cosine similarity via dot product | |
| new_embeddings = torch.nn.functional.normalize(new_embeddings, p=2, dim=1) | |
| if self.embeddings is None or len(self.chunks) == 0: | |
| self.embeddings = new_embeddings | |
| self.chunks = new_chunks | |
| else: | |
| self.embeddings = torch.cat([self.embeddings, new_embeddings], dim=0) | |
| self.chunks.extend(new_chunks) | |
| self.save_index() | |
| return { | |
| "filename": filename, | |
| "chunks_added": len(new_chunks), | |
| "total_chunks": len(self.chunks), | |
| "total_documents": len(self.list_documents()) | |
| } | |
| def search(self, query: str, top_k: int = 3, min_similarity: float = 0.20) -> List[Dict[str, Any]]: | |
| if self.embeddings is None or len(self.chunks) == 0: | |
| return [] | |
| query_emb = self.encoder.encode([query], convert_to_tensor=True, show_progress_bar=False).to(self.device) | |
| query_emb = torch.nn.functional.normalize(query_emb, p=2, dim=1) | |
| # Dot product of normalized vectors = Cosine Similarity | |
| scores = torch.mm(query_emb, self.embeddings.T).squeeze(0) # [N] | |
| top_k_val = min(top_k, len(self.chunks)) | |
| top_scores, top_indices = torch.topk(scores, k=top_k_val) | |
| results = [] | |
| for score, idx in zip(top_scores.tolist(), top_indices.tolist()): | |
| if score >= min_similarity: | |
| chunk = dict(self.chunks[idx]) | |
| chunk["score"] = round(float(score), 4) | |
| results.append(chunk) | |
| return results | |
| def list_documents(self) -> List[Dict[str, Any]]: | |
| docs: Dict[str, int] = {} | |
| for c in self.chunks: | |
| src = c.get("source", "Unknown") | |
| docs[src] = docs.get(src, 0) + 1 | |
| return [{"filename": name, "chunks": count} for name, count in docs.items()] | |
| def delete_document(self, filename: str, auto_save: bool = True) -> int: | |
| if not self.chunks: | |
| return 0 | |
| keep_indices = [i for i, c in enumerate(self.chunks) if c.get("source") != filename] | |
| removed_count = len(self.chunks) - len(keep_indices) | |
| if removed_count > 0: | |
| self.chunks = [self.chunks[i] for i in keep_indices] | |
| if self.chunks and self.embeddings is not None: | |
| self.embeddings = self.embeddings[keep_indices] | |
| else: | |
| self.embeddings = None | |
| if auto_save: | |
| self.save_index() | |
| return removed_count | |
| def clear(self): | |
| self.chunks = [] | |
| self.embeddings = None | |
| self.save_index() | |
| class WebSearchEngine: | |
| def __init__(self, rag_encoder: Optional[SentenceTransformer] = None): | |
| self.rag_encoder = rag_encoder | |
| self.client = httpx.Client( | |
| headers={ | |
| "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/124.0.0.0 Safari/537.36" | |
| }, | |
| follow_redirects=True, | |
| timeout=4.0 | |
| ) | |
| def scrape_url_paragraphs(self, url: str) -> List[str]: | |
| try: | |
| resp = self.client.get(url) | |
| if resp.status_code != 200: | |
| return [] | |
| soup = BeautifulSoup(resp.text, "html.parser") | |
| # Strip noise | |
| for tag in soup(["script", "style", "nav", "footer", "header", "noscript", "aside", "form"]): | |
| tag.decompose() | |
| paragraphs = [] | |
| for p in soup.find_all(["p", "h1", "h2", "h3", "li"]): | |
| txt = re.sub(r'\s+', ' ', p.get_text()).strip() | |
| if len(txt) > 40 and not any(skip in txt.lower() for skip in ["cookie", "privacy policy", "all rights reserved", "subscribe"]): | |
| paragraphs.append(txt) | |
| return paragraphs[:25] # top 25 candidate paragraphs | |
| except Exception: | |
| return [] | |
| def search_wikipedia_fallback(self, query: str, max_results: int = 2) -> List[Dict[str, Any]]: | |
| try: | |
| import urllib.parse | |
| encoded_query = urllib.parse.quote_plus(query) | |
| search_url = f"https://en.wikipedia.org/w/api.php?action=query&list=search&srsearch={encoded_query}&utf8=&format=json" | |
| resp = self.client.get(search_url, timeout=3.0) | |
| if resp.status_code == 200: | |
| data = resp.json() | |
| items = data.get("query", {}).get("search", []) | |
| results = [] | |
| for item in items[:max_results]: | |
| title = item.get("title", "") | |
| clean_snippet = BeautifulSoup(item.get("snippet", ""), "html.parser").get_text() | |
| results.append({ | |
| "title": f"Wikipedia: {title}", | |
| "url": f"https://en.wikipedia.org/wiki/{title.replace(' ', '_')}", | |
| "snippet": clean_snippet, | |
| "body": clean_snippet | |
| }) | |
| return results | |
| except Exception: | |
| pass | |
| return [] | |
| def search_and_deep_scrape(self, query: str, max_results: int = 3) -> Dict[str, Any]: | |
| raw_results = [] | |
| try: | |
| ddgs_res = DDGS().text(query, max_results=max_results) | |
| if ddgs_res: | |
| for r in ddgs_res: | |
| raw_results.append({ | |
| "title": r.get("title", "Web Page"), | |
| "url": r.get("href", ""), | |
| "snippet": r.get("body", ""), | |
| "body": r.get("body", "") | |
| }) | |
| except Exception as e: | |
| print(f"⚠️ DDGS primary search notice: {e}") | |
| # If primary search produced nothing, invoke Wikipedia fallback | |
| if not raw_results: | |
| raw_results = self.search_wikipedia_fallback(query, max_results=max_results) | |
| if not raw_results: | |
| return {"context": "", "sources": []} | |
| sources = [] | |
| detailed_context_blocks = [] | |
| for i, item in enumerate(raw_results[:max_results]): | |
| title = item["title"] | |
| url = item["url"] | |
| snippet = item["snippet"] | |
| # Record source metadata for frontend citations | |
| sources.append({ | |
| "id": i + 1, | |
| "title": title, | |
| "url": url, | |
| "snippet": snippet[:180] + ("..." if len(snippet) > 180 else "") | |
| }) | |
| # Deep scrape top 2 web pages if valid URL | |
| page_text = "" | |
| if url and url.startswith("http") and i < 2: | |
| paragraphs = self.scrape_url_paragraphs(url) | |
| if paragraphs: | |
| # Select top most informative paragraphs | |
| page_text = " ".join(paragraphs[:3]) | |
| block_content = page_text if page_text and len(page_text) > len(snippet) else snippet | |
| clean_slice = re.sub(r'\s+', ' ', block_content).strip()[:350] | |
| if clean_slice: | |
| detailed_context_blocks.append(f"- {clean_slice}") | |
| combined_context = "\n".join(detailed_context_blocks) | |
| return { | |
| "context": combined_context, | |
| "sources": sources | |
| } | |