Download src/services/rag_service.py from nikeshn/LibBee: direct link, hf CLI and curl.
- Browser
- Download file 15.9 kB
-
https://huggingface.co/spaces/nikeshn/LibBee/resolve/main/src/services/rag_service.py
- Command line
-
hf download hf://spaces/nikeshn/LibBee/src/services/rag_service.py
-
curl -L -o rag_service.py https://huggingface.co/spaces/nikeshn/LibBee/resolve/main/src/services/rag_service.py
15.9 kB
| """Local keyword retrieval first, with an optional cached semantic index. | |
| Initialization never needs an API key to serve the local KB. Dense failures | |
| leave keyword search available; callers may run initialize in the background. | |
| """ | |
| import asyncio | |
| import hashlib | |
| import hmac | |
| import json | |
| import logging | |
| import os | |
| import re | |
| import time | |
| from pathlib import Path | |
| from typing import List | |
| from src.config import get_settings | |
| logger = logging.getLogger(__name__) | |
| CHUNK_WORDS = 260 | |
| CHUNK_OVERLAP_WORDS = 40 | |
| EMBEDDING_MODEL = "text-embedding-3-small" | |
| CACHE_SCHEMA = 2 | |
| TOKEN_RE = re.compile(r"[^\W_]+(?:[-_][^\W_]+)*", re.UNICODE) | |
| STOP_WORDS = {"about", "and", "are", "can", "for", "from", "how", "library", "please", "that", "the", "this", "what", "when", "where", "which", "with", "you"} | |
| HEADER_RE = re.compile(r"^(?:SOURCE|TITLE|AUTHORITY|LAST_REVIEWED|LAST_UPDATED):", re.IGNORECASE) | |
| class RAGService: | |
| def __init__(self): | |
| self.vectorstore = None | |
| self.bm25 = None | |
| self.bm25_corpus: List[str] = [] | |
| self.bm25_meta: List[dict] = [] | |
| self._embeddings = None | |
| self._ready = False | |
| self._kb_hash = "" | |
| self._status = "not_started" | |
| self._content_to_idx: dict[str, int] = {} | |
| self._init_lock = asyncio.Lock() | |
| def is_ready(self) -> bool: | |
| return self._ready | |
| def _knowledge_dir(self) -> Path: | |
| return get_settings().kb_dir | |
| def _cache_dir(self) -> Path: | |
| return get_settings().rag_cache_dir | |
| def _hash_knowledge_base(self, txt_files: List[Path]) -> str: | |
| # Deployments commonly reset mtimes. Hash content plus indexing settings | |
| # so unchanged data can reuse persisted embeddings across image rebuilds. | |
| digest = hashlib.sha256() | |
| digest.update(f"{CACHE_SCHEMA}:{EMBEDDING_MODEL}:{CHUNK_WORDS}:{CHUNK_OVERLAP_WORDS}".encode()) | |
| for path in sorted(txt_files): | |
| digest.update(path.name.encode("utf-8")) | |
| digest.update(b"\0") | |
| digest.update(path.read_bytes()) | |
| digest.update(b"\0") | |
| return digest.hexdigest() | |
| def _cache_paths(self) -> tuple[Path, Path, Path]: | |
| cache_dir = self._cache_dir() | |
| return cache_dir / "faiss_index", cache_dir / "bm25_cache.json", cache_dir / "kb_state.json" | |
| def _cache_signature(self, faiss_dir: Path, bm25_path: Path) -> str: | |
| key = get_settings().origin_token.encode("utf-8") | |
| if not key: | |
| return "" | |
| digest = hmac.new(key, digestmod=hashlib.sha256) | |
| for path in sorted([bm25_path, *faiss_dir.glob("*")]): | |
| if path.is_file(): | |
| digest.update(path.name.encode("utf-8")) | |
| digest.update(path.read_bytes()) | |
| return digest.hexdigest() | |
| def _tokens(value: str) -> List[str]: | |
| return TOKEN_RE.findall((value or "").casefold()) | |
| def _build_reverse_index(self) -> None: | |
| self._content_to_idx = {chunk.strip(): i for i, chunk in enumerate(self.bm25_corpus)} | |
| def _build_keyword_index(self): | |
| from rank_bm25 import BM25Okapi | |
| files = sorted(self._knowledge_dir().glob("*.txt")) | |
| kb_hash = self._hash_knowledge_base(files) | |
| chunks, metadata = self._load_and_chunk_all() | |
| if not chunks: | |
| raise ValueError("No usable knowledge-base chunks found") | |
| bm25 = BM25Okapi([self._tokens(chunk) for chunk in chunks]) | |
| return kb_hash, chunks, metadata, bm25 | |
| async def initialize(self, openai_api_key: str = "", *, force_rebuild: bool = False) -> None: | |
| async with self._init_lock: | |
| started = time.monotonic() | |
| self._status = "initializing" | |
| try: | |
| # File reads, tokenization, and heavy imports must not block the | |
| # event loop serving /config, health checks, and chat requests. | |
| kb_hash, chunks, metadata, bm25 = await asyncio.to_thread(self._build_keyword_index) | |
| self.vectorstore = None | |
| self._embeddings = None | |
| self._kb_hash = kb_hash | |
| self.bm25_corpus, self.bm25_meta, self.bm25 = chunks, metadata, bm25 | |
| self._build_reverse_index() | |
| self._ready = True | |
| self._status = "keyword_ready" | |
| logger.info("RAG keyword index ready: %d chunks in %.2fs", len(chunks), time.monotonic() - started) | |
| except Exception: | |
| self._status = ("ready" if self.vectorstore is not None else "keyword_ready") if self._ready else "failed" | |
| logger.exception("RAG keyword initialization failed; retaining any existing index") | |
| return | |
| if not openai_api_key: | |
| return | |
| self._status = "semantic_initializing" | |
| try: | |
| def load_embedding_tools(): | |
| from langchain_openai import OpenAIEmbeddings | |
| from langchain_community.vectorstores import FAISS | |
| return OpenAIEmbeddings, FAISS | |
| OpenAIEmbeddings, FAISS = await asyncio.to_thread(load_embedding_tools) | |
| self._embeddings = OpenAIEmbeddings( | |
| model=EMBEDDING_MODEL, openai_api_key=openai_api_key, | |
| request_timeout=15, max_retries=1, | |
| ) | |
| if not force_rebuild and await asyncio.to_thread(self._try_load_cached_indexes, FAISS): | |
| self._status = "ready" | |
| return | |
| vectorstore = await asyncio.wait_for( | |
| FAISS.afrom_texts(chunks, self._embeddings, metadatas=metadata), timeout=90 | |
| ) | |
| self.vectorstore = vectorstore | |
| self._status = "ready" | |
| try: | |
| await asyncio.to_thread(self._save_cached_indexes) | |
| except Exception: | |
| # An unwritable cache must not disable a working index. | |
| logger.exception("RAG cache save failed; retrieval remains available") | |
| logger.info("RAG semantic index ready in %.2fs", time.monotonic() - started) | |
| except asyncio.CancelledError: | |
| self._status = "keyword_ready" | |
| raise | |
| except Exception: | |
| self._status = "keyword_ready" | |
| logger.exception("RAG semantic index unavailable; using keyword retrieval") | |
| def _try_load_cached_indexes(self, FAISS, _BM25Okapi=None) -> bool: | |
| faiss_dir, bm25_path, state_path = self._cache_paths() | |
| if not (faiss_dir.exists() and bm25_path.exists() and state_path.exists()): | |
| return False | |
| try: | |
| state = json.loads(state_path.read_text(encoding="utf-8")) | |
| if state.get("kb_hash") != self._kb_hash or state.get("schema") != CACHE_SCHEMA: | |
| return False | |
| expected = self._cache_signature(faiss_dir, bm25_path) | |
| if not expected or not hmac.compare_digest(str(state.get("signature", "")), expected): | |
| logger.warning("RAG unsigned or modified cache ignored") | |
| return False | |
| cached = json.loads(bm25_path.read_text(encoding="utf-8")) | |
| if cached.get("corpus") != self.bm25_corpus or cached.get("meta") != self.bm25_meta: | |
| return False | |
| # FAISS uses pickle internally. Only load bytes authenticated with | |
| # the origin secret; never load unsigned/user-supplied index files. | |
| self.vectorstore = FAISS.load_local( | |
| str(faiss_dir), self._embeddings, allow_dangerous_deserialization=True | |
| ) | |
| return True | |
| except Exception: | |
| logger.warning("RAG cache could not be loaded; rebuilding", exc_info=True) | |
| return False | |
| def _save_cached_indexes(self) -> None: | |
| if self.vectorstore is None or not get_settings().origin_token: | |
| return | |
| faiss_dir, bm25_path, state_path = self._cache_paths() | |
| faiss_dir.mkdir(parents=True, exist_ok=True) | |
| self.vectorstore.save_local(str(faiss_dir)) | |
| bm25_tmp = bm25_path.with_suffix(".tmp") | |
| bm25_tmp.write_text(json.dumps({"corpus": self.bm25_corpus, "meta": self.bm25_meta}, ensure_ascii=False), encoding="utf-8") | |
| os.replace(bm25_tmp, bm25_path) | |
| state_tmp = state_path.with_suffix(".tmp") | |
| state_tmp.write_text(json.dumps({ | |
| "schema": CACHE_SCHEMA, "kb_hash": self._kb_hash, | |
| "embedding_model": EMBEDDING_MODEL, "chunk_count": len(self.bm25_corpus), | |
| "signature": self._cache_signature(faiss_dir, bm25_path), | |
| }), encoding="utf-8") | |
| os.replace(state_tmp, state_path) | |
| def _load_and_chunk_all(self) -> tuple[List[str], List[dict]]: | |
| all_chunks, all_meta = [], [] | |
| seen: set[str] = set() | |
| for path in sorted(self._knowledge_dir().glob("*.txt")): | |
| try: | |
| text = path.read_text(encoding="utf-8").strip() | |
| if not text: | |
| continue | |
| header, header_lines = {}, set() | |
| lines = text.splitlines() | |
| for index, line in enumerate(lines[:8]): | |
| if HEADER_RE.match(line.strip()): | |
| key, value = line.split(":", 1) | |
| header[key.strip().upper()] = value.strip() | |
| header_lines.add(index) | |
| body = "\n".join(line for index, line in enumerate(lines) if index not in header_lines) | |
| metadata = { | |
| "source": header.get("SOURCE", ""), | |
| "title": header.get("TITLE", path.stem.replace("_", " ")), | |
| "filename": path.name, | |
| "authority": header.get("AUTHORITY", "supplied_library_content"), | |
| "last_reviewed": header.get("LAST_REVIEWED", ""), | |
| "last_updated": header.get("LAST_UPDATED", ""), | |
| } | |
| for chunk in self._chunk_text(body): | |
| fingerprint = re.sub(r"\s+", " ", chunk).strip() | |
| if fingerprint in seen: | |
| continue | |
| seen.add(fingerprint) | |
| all_chunks.append(chunk) | |
| all_meta.append(dict(metadata)) | |
| except Exception: | |
| logger.warning("Could not read knowledge file %s", path.name, exc_info=True) | |
| return all_chunks, all_meta | |
| def _chunk_text(self, text: str) -> List[str]: | |
| # Source files use RESOURCE 01:, all-caps headings without colons, | |
| # staff names, and === section markers. Keep each entry with its URL. | |
| sections, current = [], [] | |
| for line in text.splitlines(): | |
| stripped = line.strip() | |
| is_heading = bool( | |
| re.match(r"^(?:TOPIC:|RESOURCE\s+\d+:|===)", stripped) | |
| or (4 <= len(stripped) <= 100 and stripped.isupper() | |
| and any(char.isalpha() for char in stripped) | |
| and not stripped.startswith(("http", "SOURCE:", "TITLE:"))) | |
| ) | |
| if is_heading and current: | |
| sections.append("\n".join(current).strip()) | |
| current = [] | |
| current.append(line) | |
| if current: | |
| sections.append("\n".join(current).strip()) | |
| chunks = [] | |
| for section in sections: | |
| if len(section) < 30 or not self._tokens(section): | |
| continue | |
| words = section.split() | |
| if len(words) <= CHUNK_WORDS: | |
| chunks.append(section) | |
| continue | |
| heading = section.splitlines()[0].strip() | |
| for start in range(0, len(words), CHUNK_WORDS - CHUNK_OVERLAP_WORDS): | |
| chunk = " ".join(words[start:start + CHUNK_WORDS]) | |
| if start and len(heading) < 100: | |
| chunk = heading + "\n" + chunk | |
| chunks.append(chunk) | |
| if start + CHUNK_WORDS >= len(words): | |
| break | |
| return chunks | |
| async def hybrid_search(self, query: str, top_k: int = 5, alpha: float = 0.6) -> List[dict]: | |
| if not self._ready or self.bm25 is None or not query.strip() or top_k <= 0: | |
| return [] | |
| # Snapshot before awaiting dense search so an admin rebuild cannot mix | |
| # old vector document ids with a new corpus/metadata list. | |
| corpus, metadata, bm25 = self.bm25_corpus, self.bm25_meta, self.bm25 | |
| vectorstore, reverse = self.vectorstore, self._content_to_idx | |
| tokens = self._tokens(query) | |
| meaningful = {token for token in tokens if token not in STOP_WORDS and len(token) > 1} | |
| if not meaningful: | |
| return [] | |
| fetch_k = min(max(top_k * 3, 15), len(corpus)) | |
| alpha = min(1.0, max(0.0, alpha)) | |
| scores = bm25.get_scores(tokens) | |
| lexical_indices = sorted(range(len(corpus)), key=lambda idx: float(scores[idx]), reverse=True) | |
| lexical_indices = [idx for idx in lexical_indices | |
| if meaningful & set(self._tokens(corpus[idx]))][:fetch_k] | |
| dense_results = [] | |
| if vectorstore is not None: | |
| try: | |
| dense_results = await asyncio.wait_for( | |
| vectorstore.asimilarity_search_with_score(query, k=fetch_k), timeout=4 | |
| ) | |
| except Exception: | |
| logger.warning("Semantic retrieval unavailable; using keywords", exc_info=True) | |
| rrf: dict[int, float] = {} | |
| for rank, (doc, _score) in enumerate(dense_results): | |
| idx = reverse.get(doc.page_content.strip(), -1) | |
| if idx >= 0: | |
| rrf[idx] = rrf.get(idx, 0.0) + alpha / (61 + rank) | |
| lexical_weight = (1 - alpha) if dense_results else 1.0 | |
| for rank, idx in enumerate(lexical_indices): | |
| rrf[idx] = rrf.get(idx, 0.0) + lexical_weight / (61 + rank) | |
| results = [] | |
| # Filter before applying top_k: irrelevant dense hits must not crowd | |
| # relevant keyword hits out of the candidate pool. | |
| for idx in sorted(rrf, key=rrf.get, reverse=True): | |
| overlap = len(meaningful & set(self._tokens(corpus[idx]))) | |
| if not overlap: | |
| continue | |
| results.append({ | |
| "content": corpus[idx], **metadata[idx], "score": round(rrf[idx], 6), | |
| "lexical_overlap": overlap, | |
| "retrieval_mode": "hybrid" if dense_results else "keyword", | |
| }) | |
| if len(results) >= top_k: | |
| break | |
| return results | |
| def _find_corpus_index(self, content: str) -> int: | |
| return self._content_to_idx.get(content.strip(), -1) | |
| def audit_knowledge_base(self) -> dict: | |
| files = sorted(self._knowledge_dir().glob("*.txt")) | |
| titles, duplicates = {}, [] | |
| for path in files: | |
| try: | |
| text = path.read_text(encoding="utf-8")[:500] | |
| title = next((line.split(":", 1)[1].strip() for line in text.splitlines() if line.startswith("TITLE:")), path.stem) | |
| if title in titles: | |
| duplicates.append({"title": title, "files": [titles[title], path.name]}) | |
| else: | |
| titles[title] = path.name | |
| except OSError: | |
| continue | |
| return {"file_count": len(files), "duplicate_titles": duplicates[:20]} | |
| def get_stats(self) -> dict: | |
| return { | |
| "ready": self._ready, "status": self._status, | |
| "retrieval_mode": "hybrid" if self.vectorstore is not None else "keyword" if self._ready else "unavailable", | |
| "chunk_count": len(self.bm25_corpus), "knowledge_dir": str(self._knowledge_dir()), | |
| "vectorstore_loaded": self.vectorstore is not None, "bm25_loaded": self.bm25 is not None, | |
| "kb_hash": self._kb_hash, | |
| } | |