src/embeddings.py

"""
embeddings.py — Local embedding generation and FAISS index management.

all-MiniLM-L6-v2 runs on CPU with no API key or cost.
384-dimensional embeddings, ~80MB model, loads in ~3 seconds on first call.

FAISS IndexFlatIP with L2-normalized vectors = exact cosine similarity search.
Correct choice for < 500k chunks (enterprise PDF use case).
"""

import logging
import numpy as np
import faiss
from sentence_transformers import SentenceTransformer
from src.utils import get_env

logger = logging.getLogger("enterprise-rag.embeddings")

_embedding_model = None
EMBEDDING_DIM = 384


def get_embedding_model() -> SentenceTransformer:
    """
    Lazy-load the embedding model once and reuse across all calls.
    Downloading happens on first use; cached in HF Spaces persistent storage.
    """
    global _embedding_model
    if _embedding_model is None:
        model_name = get_env(
            "EMBEDDING_MODEL",
            "sentence-transformers/all-MiniLM-L6-v2"
        )
        logger.info(f"Loading embedding model: {model_name}")
        _embedding_model = SentenceTransformer(model_name)
        logger.info("Embedding model ready")
    return _embedding_model


def embed_texts(texts: list) -> np.ndarray:
    """
    Generate L2-normalized embeddings for a list of texts.

    Normalization converts dot product to cosine similarity,
    which is what IndexFlatIP computes. Standard for semantic search.

    Returns float32 array of shape (len(texts), 384).
    """
    if not texts:
        return np.zeros((0, EMBEDDING_DIM), dtype=np.float32)

    model = get_embedding_model()
    embeddings = model.encode(
        texts,
        batch_size=32,
        normalize_embeddings=True,
        show_progress_bar=False,
        convert_to_numpy=True,
    )
    return embeddings.astype(np.float32)


def build_faiss_index(embeddings: np.ndarray) -> faiss.IndexFlatIP:
    """
    Build a FAISS flat inner-product index.
    With L2-normalized vectors this equals exact cosine similarity search.
    No training required, deterministic results.
    """
    if embeddings.shape[0] == 0:
        raise ValueError("Cannot build FAISS index from empty embeddings.")

    index = faiss.IndexFlatIP(EMBEDDING_DIM)
    index.add(embeddings)
    logger.info(f"FAISS index built: {index.ntotal} vectors")
    return index


def search_index(
    index: faiss.IndexFlatIP,
    query_embedding: np.ndarray,
    top_k: int = 5,
) -> tuple:
    """
    Retrieve top-k most similar chunks.

    Returns:
        scores  — cosine similarity scores, float32 array shape [top_k]
        indices — chunk positions in original list, int64 array shape [top_k]

    Score guide:
        1.0  = identical
        0.8+ = highly relevant
        0.5-0.8 = moderately relevant
        < 0.5 = likely irrelevant
    """
    if index.ntotal == 0:
        return np.array([], dtype=np.float32), np.array([], dtype=np.int64)

    k = min(top_k, index.ntotal)
    query_2d = query_embedding.reshape(1, -1).astype(np.float32)
    scores, indices = index.search(query_2d, k)
    return scores[0], indices[0]