app/vectorstore/faiss_wrapper.py

"""LangChain-backed FAISS vector store (default backend).

Uses ``langchain_community.vectorstores.FAISS`` — open source, runs locally,
no separate vector database process. Embeddings and persistence use the
configured LangChain embedder and ``settings.faiss_index_path``.

Tenant isolation: FAISS has no server-side metadata filters; we
over-fetch and post-filter by ``tenant_id`` (same pattern as the legacy
custom FAISS wrapper).

Concurrent writes are serialized with a lock so parallel ingest workers do
not corrupt the in-memory index or on-disk files.
"""

import logging
import threading
from pathlib import Path
from typing import Any, cast

from langchain_community.vectorstores.faiss import FAISS
from langchain_core.documents import Document
from langchain_core.embeddings import Embeddings

from app.config import settings
from app.models.schemas import SearchResult
from app.vectorstore.base import VectorStore

logger = logging.getLogger(__name__)

_FETCH_MULTIPLIER = 20  # over-fetch factor to compensate for post-tenant filtering


class FAISSVectorStore(VectorStore):
    """VectorStore backed by LangChain FAISS with post-hoc tenant filtering.

    Args:
        embedding: LangChain-compatible embeddings instance.

    Example::

        vs = FAISSVectorStore(embedding=get_embedding_client())
        vs.add_documents(docs)
        results = vs.search("Victorian terrace", tenant_id="t-abc", k=5)
    """

    def __init__(self, embedding: Embeddings) -> None:
        self._embedding = embedding
        self._index_path = Path(settings.faiss_index_path)
        self._index_path.mkdir(parents=True, exist_ok=True)
        self._store: FAISS | None = None
        self._lock = threading.RLock()
        self._load()

    def _load(self) -> None:
        """Load a persisted LangChain FAISS index from disk if present."""
        index_file = self._index_path / "index.faiss"
        if index_file.exists():
            try:
                self._store = FAISS.load_local(
                    folder_path=str(self._index_path),
                    embeddings=self._embedding,
                    allow_dangerous_deserialization=True,
                )
                logger.info("Loaded LangChain FAISS index from %s", self._index_path)
            except Exception as exc:
                logger.warning("Could not load FAISS index: %s — starting fresh", exc)
                self._store = None
        else:
            logger.info("No FAISS index found at %s — will create on first add", self._index_path)

    def _save(self) -> None:
        """Persist the current FAISS index to disk."""
        if self._store is not None:
            self._store.save_local(str(self._index_path))

    def add_documents(self, documents: list[Document]) -> None:
        """Embed and insert ``documents`` into the FAISS index.

        Args:
            documents: LangChain Documents with full metadata.
        """
        if not documents:
            return

        with self._lock:
            if self._store is None:
                self._store = FAISS.from_documents(documents, self._embedding)
            else:
                self._store.add_documents(documents)

            self._save()
        logger.debug("Added %d documents to LangChain FAISS", len(documents))

    def search(
        self,
        query: str,
        tenant_id: str,
        k: int = 10,
        *,
        hierarchy_level: str | None = None,
        doc_id_in: frozenset[str] | None = None,
    ) -> list[SearchResult]:
        """Search FAISS and post-filter by ``tenant_id`` (and optional metadata filters).

        Over-fetches by ``_FETCH_MULTIPLIER`` to ensure enough tenant-matching
        results are available after filtering.

        Args:
            query: Plain-text search query.
            tenant_id: Only return chunks belonging to this tenant.
            k: Maximum number of results after filtering.
            hierarchy_level: Restrict to this ``hierarchy_level`` when set.
            doc_id_in: Restrict to these ``doc_id`` values when set.

        Returns:
            Filtered list of :class:`~app.models.schemas.SearchResult`.
        """
        fetch_k = max(k * _FETCH_MULTIPLIER, k * 4)
        try:
            if self._store is None:
                return []
            # Embed outside the index lock so concurrent reads do not serialize on embedding.
            query_vector = self._embedding.embed_query(query)
            with self._lock:
                if self._store is None:
                    return []
                # L2 distance from FAISS — not a 0–1 relevance score; map so higher = closer match.
                pairs = self._store.similarity_search_with_score_by_vector(
                    query_vector,
                    k=fetch_k,
                )
        except Exception as exc:
            logger.warning("FAISS search failed: %s", exc)
            return []

        results: list[SearchResult] = []
        for doc, distance in pairs:
            meta = doc.metadata
            if meta.get("tenant_id") != tenant_id:
                continue
            hl = str(meta.get("hierarchy_level") or "paragraph")
            if hierarchy_level is not None and hl != hierarchy_level:
                continue
            did = str(meta.get("doc_id", ""))
            if doc_id_in is not None and did not in doc_id_in:
                # Always allow knowledge-base (KB) documents through — these are
                # the firm's approved templates / training uploads that should be
                # retrievable regardless of which report is being generated.
                if not meta.get("kb"):
                    continue
            d = float(distance)
            sim = 1.0 / (1.0 + d)
            st = meta.get("section_title")
            sid = meta.get("section_id")
            pi = meta.get("paragraph_index")
            src = meta.get("source")
            kb = meta.get("kb")
            kb_path = meta.get("kb_path")
            chunk_role = meta.get("chunk_role")
            pidx: int | None
            if isinstance(pi, int):
                pidx = pi
            else:
                try:
                    pidx = int(pi) if pi is not None else None
                except (TypeError, ValueError):
                    pidx = None
            results.append(
                SearchResult(
                    chunk_id=str(meta.get("chunk_id", "")),
                    doc_id=did,
                    tenant_id=str(meta.get("tenant_id", "")),
                    text=doc.page_content,
                    score=sim,
                    section_type=str(meta.get("section_type", "paragraph")),
                    hierarchy_level=hl,
                    section_title=str(st) if st is not None else None,
                    section_id=str(sid) if sid is not None else None,
                    paragraph_index=pidx,
                    parent_chunk_id=str(meta.get("parent_chunk_id"))
                    if meta.get("parent_chunk_id")
                    else None,
                    source=str(src) if src is not None else None,
                    kb=bool(kb) if kb is not None else None,
                    kb_path=str(kb_path) if kb_path is not None else None,
                    chunk_role=str(chunk_role) if chunk_role is not None else None,
                )
            )
            if len(results) >= k:
                break

        return results

    async def search_async(
        self,
        query: str,
        tenant_id: str,
        k: int = 10,
        *,
        hierarchy_level: str | None = None,
        doc_id_in: frozenset[str] | None = None,
    ) -> list[SearchResult]:
        """Non-blocking search via thread pool (embedding no longer under global lock)."""
        from functools import partial

        from app.async_executor import run_sync_in_executor

        return await run_sync_in_executor(
            partial(
                self.search,
                query,
                tenant_id,
                k=k,
                hierarchy_level=hierarchy_level,
                doc_id_in=doc_id_in,
            ),
        )

    def delete_document(self, doc_id: str) -> None:
        """Remove all chunks for ``doc_id``.

        LangChain FAISS supports deletion by document ID if the index was
        built with ``docstore`` support (the default).

        Args:
            doc_id: Document identifier to remove.
        """
        with self._lock:
            if self._store is None:
                return

            ds = cast(Any, self._store.docstore)._dict
            ids_to_delete = [
                doc_id_key for doc_id_key, doc in ds.items() if doc.metadata.get("doc_id") == doc_id
            ]
            if ids_to_delete:
                self._store.delete(ids_to_delete)
                self._save()
                logger.info("Deleted %d chunks for doc_id=%s from FAISS", len(ids_to_delete), doc_id)

    def count(self, tenant_id: str) -> int:
        """Count live chunks for ``tenant_id``.

        Args:
            tenant_id: Tenant identifier.

        Returns:
            Integer chunk count.
        """
        with self._lock:
            if self._store is None:
                return 0
            ds = cast(Any, self._store.docstore)._dict
            return len([d for d in ds.values() if d.metadata.get("tenant_id") == tenant_id])

    def count_for_doc(self, doc_id: str) -> int:
        """Count live chunks for a specific document.

        Args:
            doc_id: Document identifier.

        Returns:
            Integer chunk count for that document.
        """
        with self._lock:
            if self._store is None:
                return 0
            ds = cast(Any, self._store.docstore)._dict
            return len([d for d in ds.values() if d.metadata.get("doc_id") == doc_id])