ingestion/embedder.py

from __future__ import annotations

import numpy as np
from sentence_transformers import SentenceTransformer
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from ingestion.chunker import Chunk

MODEL_NAME = "all-MiniLM-L6-v2"


class Embedder:
    def __init__(self, model_name: str = MODEL_NAME) -> None:
        self._model = SentenceTransformer(model_name)
        # Cache dimension by embedding an empty probe rather than hard-coding it.
        probe = self._model.encode(["probe"], normalize_embeddings=False)
        self._dimension = int(probe.shape[1])

    @property
    def dimension(self) -> int:
        """Output vector size (384 for all-MiniLM-L6-v2)."""
        return self._dimension

    def embed(self, texts: list[str], batch_size: int = 64) -> np.ndarray:
        """Embed a list of texts and return an L2-normalised float32 array.

        Args:
            texts:      Non-empty list of strings to embed.
            batch_size: How many texts to encode per forward pass.

        Returns:
            np.ndarray of shape (len(texts), dimension), dtype float32,
            each row unit-normalised so dot product == cosine similarity.

        Raises:
            ValueError: If texts is empty.
        """
        if not texts:
            raise ValueError("texts must be a non-empty list")

        vectors = self._model.encode(
            texts,
            batch_size=batch_size,
            normalize_embeddings=True,
            show_progress_bar=False,
        )
        return vectors.astype(np.float32)

    def embed_chunks(self, chunks: list[Chunk], batch_size: int = 64) -> np.ndarray:
        """Build context-enriched text for each chunk, embed it, and return the vectors.

        The enriched text prepends source metadata so the embedding captures
        document context alongside chunk content (Contextual Retrieval pattern):

            "<source> | Page <n> [| <section_header>]\\n<chunk text>"

        chunk.text is never modified — it stays as the original text for
        display and citation.  chunk.embedded_text is set to the enriched
        string so callers can inspect or log exactly what was embedded.

        Args:
            chunks:     List of Chunk objects to embed.
            batch_size: How many texts to encode per forward pass.

        Returns:
            np.ndarray of shape (len(chunks), dimension), dtype float32,
            L2-normalised — one row per chunk in input order.
        """
        enriched: list[str] = []
        for chunk in chunks:
            prefix = f"{chunk.source} | Page {chunk.page_num}"
            if chunk.section_header:
                prefix += f" | {chunk.section_header}"
            text = f"{prefix}\n{chunk.text}"
            chunk.embedded_text = text
            enriched.append(text)
        return self.embed(enriched, batch_size=batch_size)

    def embed_query(self, query: str) -> np.ndarray:
        """Embed a single query string.

        Returns:
            np.ndarray of shape (dimension,), dtype float32, L2-normalised.

        Raises:
            ValueError: If query is empty or whitespace-only.
        """
        if not query or not query.strip():
            raise ValueError("query must be a non-empty string")

        vector = self._model.encode(
            query,
            normalize_embeddings=True,
            show_progress_bar=False,
        )
        return vector.astype(np.float32)