LibBee / src /services /rag_service.py
nikeshn's picture
Upload 8 files
a2416df verified
Raw History Blame Contribute Delete
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()
@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,
}