vector_store.py

"""
Singleton FAISS vector store shared by all agents.
"""

from __future__ import annotations

import os
from typing import Optional

from langchain_community.vectorstores import FAISS
from langchain_core.documents import Document
from langchain_community.embeddings import HuggingFaceEmbeddings

_vector_store: Optional[FAISS] = None
_embeddings = None

PERSIST_DIR = os.path.join(os.path.dirname(__file__), ".faiss_store")


def _get_embeddings():
    global _embeddings
    if _embeddings is None:
        _embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
    return _embeddings


def get_vector_store() -> FAISS:
    """Return the singleton FAISS store, loading from disk if available."""
    global _vector_store
    if _vector_store is not None:
        return _vector_store

    emb = _get_embeddings()
    if os.path.exists(PERSIST_DIR):
        _vector_store = FAISS.load_local(
            PERSIST_DIR, emb, allow_dangerous_deserialization=True
        )
    else:
        # Bootstrap with a placeholder so FAISS is initialised
        _vector_store = FAISS.from_documents(
            [Document(page_content="__init__", metadata={"source": "init"})],
            emb,
        )
    return _vector_store


def add_documents(docs: list[Document]) -> None:
    """Add documents to the store and persist to disk."""
    vs = get_vector_store()
    vs.add_documents(docs)
    vs.save_local(PERSIST_DIR)


def reset_store() -> None:
    """Clear the in-memory store (useful between sessions)."""
    global _vector_store
    _vector_store = None
    if os.path.exists(PERSIST_DIR):
        import shutil
        shutil.rmtree(PERSIST_DIR)