"""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() @staticmethod 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, }