Blame for src/services/rag_service.py

"""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,
        }