src/rag/KnowledgeBase.py

import hashlib
import logging
import os
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Sequence, Tuple

from langchain_community.vectorstores import Chroma
from langchain_openai import OpenAIEmbeddings

from src.data.chroma_config import COLLECTION_NAME, PERSIST_DIRECTORY, ensure_persist_dir
from src.core.settings import get_settings

logger = logging.getLogger(__name__)


@dataclass(frozen=True)
class RetrievedSource:
    title: str
    category: str
    source_path: str
    excerpt: str


class KnowledgeBase:
    """
    Local, curated knowledge base (KB-first) backed by Chroma.

    Docs live under `src/data/knowledge_base/<category>/*.md` (or .txt).
    """

    def __init__(
        self,
        docs_root: str = "src/data/knowledge_base",
        persist_directory: str = PERSIST_DIRECTORY,
        collection_name: str = COLLECTION_NAME,
    ):
        self.docs_root = docs_root
        self.persist_directory = persist_directory
        self.collection_name = collection_name
        settings = get_settings()
        self.embedding_model = settings.models.embedding_model
        self.min_score = float(settings.kb.min_score)
        self.embeddings = OpenAIEmbeddings(model=self.embedding_model)

        ensure_persist_dir(self.persist_directory)
        self.vector_store = Chroma(
            collection_name=collection_name,
            embedding_function=self.embeddings,
            persist_directory=self.persist_directory,
        )

        logger.info(
            "KnowledgeBase.init docs_root=%s persist_directory=%s collection_name=%s embedding_model=%s",
            self.docs_root,
            self.persist_directory,
            self.collection_name,
            self.embedding_model,
        )

    def _collection_count(self) -> Optional[int]:
        """
        Best-effort collection count for debug logging.
        """
        try:
            col = getattr(self.vector_store, "_collection", None)
            if col is not None and hasattr(col, "count"):
                return int(col.count())
        except Exception:
            return None
        return None

    def describe(self) -> Dict[str, Any]:
        """
        Debug-friendly metadata about the KB instance and backing store.
        """
        return {
            "docs_root": self.docs_root,
            "persist_directory": self.persist_directory,
            "collection_name": self.collection_name,
            "embedding_model": self.embedding_model,
            "min_score": self.min_score,
            "collection_count": self._collection_count(),
        }

    def _read_text(self, path: str) -> str:
        with open(path, "r", encoding="utf-8") as f:
            return f.read()

    def _doc_id(self, path: str, content: str) -> str:
        h = hashlib.sha256()
        h.update(path.encode("utf-8"))
        h.update(b"\x00")
        h.update(content.encode("utf-8"))
        return h.hexdigest()

    def _title_from_text(self, text: str, fallback: str) -> str:
        for line in text.splitlines():
            s = line.strip()
            if s.startswith("#"):
                return s.lstrip("#").strip()[:120] or fallback
        return fallback

    def _iter_docs(self) -> List[Tuple[str, str, str]]:
        """
        Returns list of (path, category, text).
        Category is inferred from the immediate parent folder name.
        """
        if not os.path.isdir(self.docs_root):
            return []

        docs: List[Tuple[str, str, str]] = []
        for root, _, files in os.walk(self.docs_root):
            for name in files:
                if not (name.endswith(".md") or name.endswith(".txt")):
                    continue
                path = os.path.join(root, name)
                category = os.path.basename(os.path.dirname(path)) or "general"
                try:
                    text = self._read_text(path).strip()
                except Exception:
                    logger.exception("Failed reading KB doc: %s", path)
                    continue
                if not text:
                    continue
                docs.append((path, category, text))
        return docs

    def ensure_ingested(self) -> None:
        """
        Idempotent ingestion. Uses deterministic IDs derived from path+content.
        """
        docs = self._iter_docs()
        if not docs:
            return

        texts: List[str] = []
        metadatas: List[Dict[str, Any]] = []
        ids: List[str] = []
        for path, category, text in docs:
            doc_id = self._doc_id(path, text)
            title = self._title_from_text(text, fallback=os.path.basename(path))
            texts.append(text)
            metadatas.append(
                {
                    "namespace": "kb",
                    "source_path": path,
                    "category": category,
                    "title": title,
                }
            )
            ids.append(doc_id)

        # Chroma will error on duplicate IDs; ignore those by checking existing.
        # This is intentionally best-effort; KB ingestion is a startup concern.
        try:
            existing = set(self.vector_store.get(ids=ids).get("ids", []))
        except Exception:
            existing = set()

        to_add = [
            (t, m, i) for t, m, i in zip(texts, metadatas, ids) if i not in existing
        ]
        if not to_add:
            return

        add_texts = [t for t, _, _ in to_add]
        add_metas = [m for _, m, _ in to_add]
        add_ids = [i for _, _, i in to_add]
        try:
            self.vector_store.add_texts(
                texts=add_texts, metadatas=add_metas, ids=add_ids
            )
            logger.info("KB ingested %d docs", len(to_add))
        except Exception:
            # If the underlying store rejects some IDs as already present, keep going.
            logger.exception("KB ingestion failed (possibly duplicate IDs)")

    def retrieve(
        self,
        query: str,
        k: int = 4,
        categories: Optional[Sequence[str]] = None,
    ) -> List[RetrievedSource]:
        self.ensure_ingested()

        query_normalized = (query or "").strip().lower()
        where: Dict[str, Any] = {"namespace": "kb"}
        if categories:
            # Chroma's metadata filtering expects a single top-level operator when using operators.
            # Use $and to combine `namespace` with a category $in constraint.
            where = {
                "$and": [
                    {"namespace": "kb"},
                    {"category": {"$in": list(categories)}},
                ]
            }

        # Use relevance scores so we can avoid returning low-similarity "hits".
        # If we return low-quality KB matches, agents will skip Tavily and answer incorrectly.
        min_score = self.min_score
        docs_with_scores: List[Tuple[Any, float]] = []
        try:
            docs_with_scores = self.vector_store.similarity_search_with_relevance_scores(
                query_normalized, k=k, filter=where
            )
            logger.info(f"KnowledgeBase.retrieve: {docs_with_scores}")
        except TypeError:
            # Older langchain wrappers use `where` not `filter`.
            try:
                docs_with_scores = (
                    self.vector_store.similarity_search_with_relevance_scores(
                        query_normalized, k=k, where=where
                    )
                )  # type: ignore[call-arg]
            except Exception:
                docs_with_scores = []
        except Exception:
            docs_with_scores = []

        # Fallback: no-score search (best-effort).
        docs: List[Any] = []
        if docs_with_scores:
            docs = [d for d, score in docs_with_scores if float(score) >= min_score]
            if not docs:
                # Helpful debug: log the best few candidates even if below threshold.
                top = sorted(docs_with_scores, key=lambda t: float(t[1]), reverse=True)[:3]
                preview = []
                for d, score in top:
                    meta = getattr(d, "metadata", {}) or {}
                    preview.append(
                        {
                            "score": float(score),
                            "title": str(meta.get("title") or "")[:80],
                            "category": str(meta.get("category") or ""),
                        }
                    )
                logger.info(
                    "KnowledgeBase.retrieve no_hits_above_threshold min_score=%.2f categories=%s preview=%s",
                    min_score,
                    list(categories or []),
                    preview,
                )
        else:
            try:
                docs = self.vector_store.similarity_search(query_normalized, k=k, filter=where)
            except TypeError:
                docs = self.vector_store.similarity_search(query_normalized, k=k, where=where)  # type: ignore[call-arg]
            except Exception:
                # Last resort: retrieve without filters and filter client-side.
                try:
                    docs = self.vector_store.similarity_search(query_normalized, k=k)
                except Exception:
                    docs = []

            if categories and docs:
                allowed = {c.strip() for c in categories if str(c).strip()}
                docs = [
                    d
                    for d in docs
                    if str((getattr(d, "metadata", {}) or {}).get("category") or "").strip()
                    in allowed
                ]

        out: List[RetrievedSource] = []
        for d in docs:
            meta = d.metadata or {}
            out.append(
                RetrievedSource(
                    title=str(meta.get("title") or "KB Source"),
                    category=str(meta.get("category") or "general"),
                    source_path=str(meta.get("source_path") or ""),
                    excerpt=(d.page_content or "")[:900],
                )
            )
        return out