rag/vectorstore.py

from __future__ import annotations

import logging
from typing import Optional, Iterable, List
from langchain_core.embeddings import Embeddings
from langchain_community.vectorstores import FAISS
from .config import settings
from .splitter import documents

logger = logging.getLogger(__name__)

# Directory where FAISS index is persisted
_VECTORSTORE_DIR = settings.vectorstore_dir
_INDEX_NAME = settings.vectorstore_index_name
_VECTORSTORE_DIR.mkdir(parents=True, exist_ok=True)


class _LazyHuggingFaceEmbeddings(Embeddings):

    def __init__(self, model_name: str, encode_kwargs: Optional[dict] = None) -> None:
        self.model_name = model_name
        self.encode_kwargs = encode_kwargs or {}
        self._real = None

    def _init(self) -> None:
        if self._real is None:
            from langchain_huggingface import HuggingFaceEmbeddings

            self._real = HuggingFaceEmbeddings(
                model_name=self.model_name, encode_kwargs=self.encode_kwargs
            )

    async def aembed_documents(self, texts: Iterable[str]) -> List[List[float]]:
        self._init()
        return await self._real.aembed_documents(texts)  # type: ignore[attr-defined]

    def embed_documents(self, texts: Iterable[str]) -> List[List[float]]:
        self._init()
        return self._real.embed_documents(texts)  # type: ignore[attr-defined]

    def embed_query(self, text: str) -> List[float]:
        self._init()
        return self._real.embed_query(text)  # type: ignore[attr-defined]

    async def aembed_query(self, text: str) -> List[float]:
        self._init()
        return await self._real.aembed_query(text)  # type: ignore[attr-defined]

_EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
_embeddings = _LazyHuggingFaceEmbeddings(_EMBEDDING_MODEL, encode_kwargs={"normalize_embeddings": True})

_vectorstore: Optional[FAISS] = None


def get_vectorstore() -> FAISS:
    """Return a persisted FAISS vectorstore, building it on first run only.

    This attempts to load an existing index from disk. If the index is not
    available it computes embeddings for the documents and saves the index.
    """
    global _vectorstore

    if _vectorstore is not None:
        return _vectorstore

    try:
        _vectorstore = FAISS.load_local(
            str(_VECTORSTORE_DIR),
            _embeddings,
            index_name=_INDEX_NAME,
            allow_dangerous_deserialization=True,
        )
        return _vectorstore
    except (FileNotFoundError, OSError, ValueError, RuntimeError) as error:
        logger.info(
            "FAISS index missing or unreadable; rebuilding index. %s",
            error,
        )
        _vectorstore = FAISS.from_documents(documents, _embeddings)  # will call embed_documents
        _vectorstore.save_local(str(_VECTORSTORE_DIR), index_name=_INDEX_NAME)
        return _vectorstore