src/services/rag_service.py

"""
rag_service.py — LibBee v3.1

Fixes applied:
1. _find_corpus_index replaced with O(1) reverse-lookup dict (_content_to_idx).
   Old implementation was O(n * fetch_k) per query — a major bottleneck.
2. fetch_k corrected: max(top_k * 3, 15) — old expression max(top_k*3, top_k)
   was always top_k*3 (max was a no-op). Now enforces a floor of 15.
3. _content_to_idx rebuilt whenever corpus is loaded (init, cache load, rebuild).
4. audit_knowledge_base removed from get_stats() — now only called explicitly
   via a dedicated audit() method to avoid disk reads on every /rag-status poll.
5. _write / cache saves use os.replace for atomic writes.
"""
import hashlib
import json
import logging
import os
import re
import time
from pathlib import Path
from typing import List, Optional

import numpy as np

from src.config import get_settings

logger = logging.getLogger(__name__)
CHUNK_MAX_CHARS = 1600
CHUNK_OVERLAP_CHARS = 200


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 = ""
        # O(1) reverse lookup: chunk_text -> corpus index
        self._content_to_idx: dict[str, int] = {}

    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:
        digest = hashlib.sha256()
        for path in txt_files:
            stat = path.stat()
            digest.update(path.name.encode("utf-8"))
            digest.update(str(stat.st_mtime_ns).encode("utf-8"))
            digest.update(str(stat.st_size).encode("utf-8"))
        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 _build_reverse_index(self) -> None:
        """Build O(1) content -> index lookup. Call after any corpus change."""
        self._content_to_idx = {
            chunk.strip(): i for i, chunk in enumerate(self.bm25_corpus)
        }

    async def initialize(self, openai_api_key: str) -> None:
        t0 = time.time()
        logger.info("RAGService: starting initialization")
        settings = get_settings()
        knowledge_dir = settings.kb_dir
        txt_files = sorted(knowledge_dir.glob("*.txt")) if knowledge_dir.exists() else []
        self._kb_hash = self._hash_knowledge_base(txt_files) if txt_files else ""

        try:
            from langchain_openai import OpenAIEmbeddings
            from langchain_community.vectorstores import FAISS
            from rank_bm25 import BM25Okapi

            self._embeddings = OpenAIEmbeddings(
                model="text-embedding-3-small", openai_api_key=openai_api_key
            )
            if self._try_load_cached_indexes(FAISS, BM25Okapi):
                self._build_reverse_index()
                self._ready = True
                logger.info(
                    "RAGService: loaded cached indexes — %d chunks", len(self.bm25_corpus)
                )
                return

            chunks, metadatas = self._load_and_chunk_all()
            if not chunks:
                logger.error("RAGService: no chunks loaded — check knowledge_dir: %s", knowledge_dir)
                return

            self.vectorstore = await FAISS.afrom_texts(
                chunks, self._embeddings, metadatas=metadatas
            )
            self.bm25_corpus = chunks
            self.bm25_meta = metadatas
            tokenized = [c.lower().split() for c in chunks]
            self.bm25 = BM25Okapi(tokenized)
            self._build_reverse_index()
            self._save_cached_indexes()
            self._ready = True
            logger.info(
                "RAGService: ready — %d chunks indexed in %.1fs", len(chunks), time.time() - t0
            )
        except Exception as exc:
            logger.error("RAGService initialization failed: %s", exc, exc_info=True)

    def _try_load_cached_indexes(self, FAISS, BM25Okapi) -> 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:
                logger.info("RAGService: KB changed (hash mismatch) — rebuilding indexes")
                return False
            self.vectorstore = FAISS.load_local(
                str(faiss_dir), self._embeddings, allow_dangerous_deserialization=True
            )
            bm25_cache = json.loads(bm25_path.read_text(encoding="utf-8"))
            self.bm25_corpus = bm25_cache["corpus"]
            self.bm25_meta = bm25_cache["meta"]
            tokenized = [c.lower().split() for c in self.bm25_corpus]
            self.bm25 = BM25Okapi(tokenized)
            return True
        except Exception as exc:
            logger.warning("Failed loading cached RAG index: %s — will rebuild", exc)
            return False

    def _save_cached_indexes(self) -> None:
        if not self.vectorstore:
            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))

        # Atomic write for bm25 cache
        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)

        # Atomic write for state
        state_tmp = state_path.with_suffix(".tmp")
        state_tmp.write_text(
            json.dumps({"kb_hash": self._kb_hash, "chunk_count": len(self.bm25_corpus)}),
            encoding="utf-8",
        )
        os.replace(state_tmp, state_path)

    def _load_and_chunk_all(self) -> tuple[List[str], List[dict]]:
        all_chunks: List[str] = []
        all_meta: List[dict] = []
        knowledge_dir = self._knowledge_dir()
        if not knowledge_dir.exists():
            logger.error("Knowledge directory not found: %s", knowledge_dir)
            return [], []
        txt_files = sorted(knowledge_dir.glob("*.txt"))
        for fpath in txt_files:
            try:
                text = fpath.read_text(encoding="utf-8").strip()
                if not text:
                    continue
                source_url = ""
                title = fpath.stem.replace("_", " ").replace("-", " ")
                for line in text.splitlines()[:8]:
                    line = line.strip()
                    if line.startswith("SOURCE:"):
                        source_url = line.replace("SOURCE:", "").strip()
                    elif line.startswith("TITLE:"):
                        title = line.replace("TITLE:", "").strip()
                for chunk in self._chunk_text(text):
                    all_chunks.append(chunk)
                    all_meta.append(
                        {"source": source_url, "title": title, "filename": fpath.name}
                    )
            except Exception as exc:
                logger.warning("Failed to load %s: %s", fpath.name, exc)
        logger.info("RAGService: loaded %d chunks from %d files", len(all_chunks), len(txt_files))
        return all_chunks, all_meta

    def _chunk_text(self, text: str) -> List[str]:
        chunks: List[str] = []
        sections = re.split(r"\n(?=(?:TOPIC:|[A-Z][A-Z\s,/&()\-]+:)\s)", text)
        for section in sections:
            section = section.strip()
            if not section or len(section) < 30:
                continue
            if len(section) <= CHUNK_MAX_CHARS:
                chunks.append(section)
            else:
                start = 0
                while start < len(section):
                    end = start + CHUNK_MAX_CHARS
                    chunk = section[start:end].strip()
                    if chunk:
                        chunks.append(chunk)
                    if end >= len(section):
                        break
                    start += CHUNK_MAX_CHARS - CHUNK_OVERLAP_CHARS
        return chunks

    async def hybrid_search(self, query: str, top_k: int = 5, alpha: float = 0.6) -> List[dict]:
        if not self._ready:
            logger.warning("hybrid_search called before ready — returning empty results")
            return []
        try:
            # fetch_k: retrieve 3x candidates, minimum floor of 15
            fetch_k = max(top_k * 3, 15)

            dense_results = await self.vectorstore.asimilarity_search_with_score(
                query, k=fetch_k
            )
            tokenized_query = query.lower().split()
            bm25_scores = self.bm25.get_scores(tokenized_query)
            top_bm25_idx = np.argsort(bm25_scores)[::-1][:fetch_k].tolist()

            # Weighted RRF fusion
            # alpha controls dense list weight; (1-alpha) controls BM25 list weight
            K = 60
            rrf: dict[int, float] = {}
            for rank, (doc, _score) in enumerate(dense_results):
                idx = self._find_corpus_index(doc.page_content)
                if idx >= 0:
                    rrf[idx] = rrf.get(idx, 0.0) + alpha * (1.0 / (K + rank + 1))
            for rank, idx in enumerate(top_bm25_idx):
                rrf[idx] = rrf.get(idx, 0.0) + (1 - alpha) * (1.0 / (K + rank + 1))

            sorted_idx = sorted(rrf, key=lambda i: rrf[i], reverse=True)[:top_k]
            results = []
            for idx in sorted_idx:
                if idx < len(self.bm25_corpus):
                    meta = self.bm25_meta[idx]
                    results.append(
                        {
                            "content": self.bm25_corpus[idx],
                            "source": meta.get("source", ""),
                            "title": meta.get("title", ""),
                            "filename": meta.get("filename", ""),
                            "score": round(rrf[idx], 6),
                        }
                    )
            return results
        except Exception as exc:
            logger.error("hybrid_search error: %s", exc, exc_info=True)
            return []

    def _find_corpus_index(self, content: str) -> int:
        """O(1) reverse lookup using pre-built dict. Falls back to -1 if not found."""
        return self._content_to_idx.get(content.strip(), -1)

    def audit_knowledge_base(self) -> dict:
        """Scan KB files for duplicate titles. Call explicitly — not on every stats poll."""
        files = sorted(self._knowledge_dir().glob("*.txt"))
        titles: dict[str, str] = {}
        duplicates = []
        for path in files:
            title = path.stem
            try:
                text = path.read_text(encoding="utf-8", errors="ignore")[:500]
                first_title = next(
                    (
                        line.replace("TITLE:", "").strip()
                        for line in text.splitlines()
                        if line.startswith("TITLE:")
                    ),
                    title,
                )
            except Exception:
                first_title = title
            if first_title in titles:
                duplicates.append(
                    {"title": first_title, "files": [titles[first_title], path.name]}
                )
            else:
                titles[first_title] = path.name
        return {"file_count": len(files), "duplicate_titles": duplicates[:20]}

    def get_stats(self) -> dict:
        """Lightweight stats — no disk reads beyond what's already in memory."""
        return {
            "ready": self._ready,
            "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,
        }