feat: optimise embeds

.env-template

@@ -11,8 +11,20 @@ INTERNAL_KEY=<internal key for worker-to-backend authentication>
 # NOVITA_API_KEY=<your-novita-api-key>
 # OPEN_ROUTER_API_KEY=<your-openrouter-api-key>
 
-# Remote Embeddings (Optional - for using a remote embeddings API instead of local SentenceTransformer)
-# When set, the app will use the remote API and won't load SentenceTransformer (saves RAM)
+# Embedding model. New installs should use granite: multilingual, 32k context,
+# and the same 768 dimensions as the legacy model.
+#
+# UPGRADING AN EXISTING INSTALL? Leave this commented out. Vectors already in
+# your index were produced by the old model, and changing this without
+# re-embedding gives you retrieval against a different vector space -- which
+# fails silently, because both models are 768-dimensional. To switch, set this
+# and then run:
+#     python -m application.scripts.reembed
+EMBEDDINGS_NAME=ibm-granite/granite-embedding-311m-multilingual-r2
+
+# Remote Embeddings (Optional - for using a remote embeddings API instead of
+# running the model in-process). When set, the app calls the remote API and
+# never loads a local model, which keeps the API and worker containers small.
 EMBEDDINGS_BASE_URL=
 EMBEDDINGS_KEY=
 

AGENTS.md

@@ -96,6 +96,20 @@ ruff check .
 python -m pytest
 ```
 
+On **macOS**, run the suite with `KMP_DUPLICATE_LIB_OK=TRUE`:
+
+```bash
+KMP_DUPLICATE_LIB_OK=TRUE python -m pytest
+```
+
+`faiss-cpu` and `torch` each ship their own LLVM OpenMP runtime, and loading
+both into one process makes `libomp.dylib` abort the interpreter
+(`OMP: Error #15`). It is a macOS-only packaging clash, not a code fault: Linux
+resolves both to `libgomp`, which tolerates duplicates, so CI (`ubuntu-latest`)
+and the Docker images are unaffected. Without the variable, whether the run
+aborts depends on which tests happen to load faiss and torch in the same
+process, so a green run on one selection and an abort on another is expected.
+
 ### Frontend changes
 
 ```bash

application/Dockerfile

@@ -17,11 +17,6 @@ RUN if [ -f /usr/bin/python3.12 ]; then \
         echo "Python 3.12 not found"; exit 1; \
     fi
 
-# Download and unzip the model
-RUN wget https://d3dg1063dc54p9.cloudfront.net/models/embeddings/mpnet-base-v2.zip && \
-    unzip mpnet-base-v2.zip -d models && \
-    rm mpnet-base-v2.zip
-
 # Install Rust
 RUN wget -q -O - https://sh.rustup.rs | sh -s -- -y
 
@@ -42,6 +37,7 @@ RUN pip install --no-cache-dir --upgrade pip && \
     pip install --no-cache-dir tiktoken && \
     pip install --no-cache-dir -r requirements.txt
 
+
 # Final Stage
 FROM ubuntu:24.04 as final
 
@@ -67,11 +63,17 @@ RUN groupadd -r appuser && \
 # Copy the virtual environment and model from the builder stage
 COPY --from=builder /venv /venv
 
-COPY --from=builder /models /app/models
-
 # Copy your application code
 COPY . /app/application
 
+# Pre-fetch the embedding models into FastEmbed's cache so a fresh container
+# does not download on first ingest and an air-gapped install works at all.
+# Both defaults are baked: an upgraded deployment keeps using mpnet until it
+# runs the re-embed script, while a new one starts on granite.
+ENV EMBEDDINGS_CACHE_DIR=/app/models
+ARG EMBEDDINGS_PREFETCH=""
+RUN PYTHONPATH=/app /venv/bin/python -m application.scripts.prefetch_models ${EMBEDDINGS_PREFETCH}
+
 # Change the ownership of the /app directory to the appuser
 
 RUN mkdir -p /app/application/inputs/local

application/core/settings.py

@@ -45,11 +45,23 @@ class Settings(BaseSettings):
 
     LLM_PROVIDER: str = "docsgpt"
     LLM_NAME: Optional[str] = None  # if LLM_PROVIDER is openai, LLM_NAME can be gpt-4 or gpt-3.5-turbo
+    # Deliberately the legacy model, not the current recommendation. An
+    # existing deployment that never pinned EMBEDDINGS_NAME falls through to
+    # this default, and its stored vectors were produced by this model --
+    # changing the default here would silently retrieve against a different
+    # vector space (both are 768-dimensional, so no dimension check fires).
+    # New installs get granite from .env-template / setup.sh; existing ones
+    # switch by setting this and running application.scripts.reembed.
     EMBEDDINGS_NAME: str = "huggingface_sentence-transformers/all-mpnet-base-v2"
     EMBEDDINGS_BASE_URL: Optional[str] = None  # Remote embeddings API URL (OpenAI-compatible)
     EMBEDDINGS_KEY: Optional[str] = None  # api key for embeddings (if using openai, just copy API_KEY)
     EMBEDDINGS_MAX_INPUT_TOKENS: Optional[int] = None  # truncate each remote embed input to N tokens (overflow lost)
     EMBEDDINGS_BATCH_SIZE: int = 32  # chunks per embed request during ingest (1 = legacy per-chunk behaviour)
+    # Intra-op threads for the local ONNX runner. None = every core. ONNX Runtime
+    # scales sub-linearly across threads, so several single-threaded worker
+    # processes beat one many-threaded process on the same cores.
+    EMBEDDINGS_THREADS: Optional[int] = None
+    EMBEDDINGS_CACHE_DIR: Optional[str] = None  # where FastEmbed caches model artifacts
     GITHUB_INGEST_MAX_FILE_BYTES: int = 1048576  # skip repo blobs larger than this (0 = no cap)
     GITHUB_INGEST_MAX_WORKERS: int = 8  # parallel file fetches per GitHub repo ingest
     # Optional directory of operator-supplied model YAMLs, loaded after the

application/parser/chunking.py

@@ -3,7 +3,7 @@ from typing import List, Tuple
 import logging
 from application.parser.chunking_creator import ChunkerCreator
 from application.parser.schema.base import Document
-from application.utils import get_encoding
+from application.parser.tokenization import get_token_counter
 
 logger = logging.getLogger(__name__)
 
@@ -27,7 +27,10 @@ class Chunker:
         self.max_tokens = max_tokens
         self.min_tokens = min_tokens
         self.duplicate_headers = duplicate_headers
-        self.encoding = get_encoding()
+        # Counted in the embedding model's tokenizer, not cl100k: ``max_tokens``
+        # is compared against a limit the embedding server enforces in its own
+        # units, so counting in any other unit is a guess.
+        self.counter = get_token_counter()
 
     def separate_header_and_body(self, text: str) -> Tuple[str, str]:
         header_pattern = r"^(.*?\n){3}"
@@ -42,28 +45,40 @@ class Chunker:
 
 
     def split_document(self, doc: Document) -> List[Document]:
-        split_docs = []
-        header, body = self.separate_header_and_body(doc.text)
-        header_tokens = self.encoding.encode_ordinary(header) if header else []
-        body_tokens = self.encoding.encode_ordinary(body)
+        """Split one oversized document into ``max_tokens``-sized chunks.
 
-        current_position = 0
-        part_index = 0
-        while current_position < len(body_tokens):
-            end_position = current_position + self.max_tokens - len(header_tokens)
-            chunk_tokens = (header_tokens + body_tokens[current_position:end_position]
-                            if self.duplicate_headers or part_index == 0 else body_tokens[current_position:end_position])
-            chunk_text = self.encoding.decode(chunk_tokens)
-            new_doc = Document(
-                text=chunk_text,
-                doc_id=f"{doc.doc_id}-{part_index}",
-                embedding=doc.embedding,
-                extra_info={**(doc.extra_info or {}), "token_count": len(chunk_tokens)}
+        Pieces are sliced out of the original text rather than decoded back
+        from token ids. WordPiece tokenizers normalise as they decode --
+        all-mpnet-base-v2 lowercases -- so a decode round-trip would rewrite
+        every stored document.
+        """
+        header, body = self.separate_header_and_body(doc.text)
+        header_tokens = self.counter.count(header) if header else 0
+
+        # A chunk carrying the header has that much less room for body text.
+        with_header_budget = max(1, self.max_tokens - header_tokens)
+        if self.duplicate_headers:
+            body_pieces = self.counter.split(body, with_header_budget)
+        else:
+            body_pieces = self.counter.split(
+                body, self.max_tokens, first_max_tokens=with_header_budget
+            )
+
+        split_docs = []
+        for part_index, piece in enumerate(body_pieces):
+            include_header = bool(header) and (self.duplicate_headers or part_index == 0)
+            chunk_text = f"{header}{piece}" if include_header else piece
+            split_docs.append(
+                Document(
+                    text=chunk_text,
+                    doc_id=f"{doc.doc_id}-{part_index}",
+                    embedding=doc.embedding,
+                    extra_info={
+                        **(doc.extra_info or {}),
+                        "token_count": self.counter.count(chunk_text),
+                    },
+                )
             )
-            split_docs.append(new_doc)
-            current_position = end_position
-            part_index += 1
-            header_tokens = []
         return split_docs
 
     def classic_chunk(self, documents: List[Document]) -> List[Document]:
@@ -71,8 +86,7 @@ class Chunker:
         i = 0
         while i < len(documents):
             doc = documents[i]
-            tokens = self.encoding.encode_ordinary(doc.text)
-            token_count = len(tokens)
+            token_count = self.counter.count(doc.text)
 
             if self.min_tokens <= token_count <= self.max_tokens:
                 doc.extra_info = doc.extra_info or {}

application/parser/chunking_strategies.py

@@ -16,7 +16,7 @@ from typing import List
 from application.parser.chunking import Chunker
 from application.parser.chunking_creator import ChunkerCreator
 from application.parser.schema.base import Document
-from application.utils import get_encoding
+from application.parser.tokenization import get_token_counter
 
 logger = logging.getLogger(__name__)
 
@@ -41,19 +41,16 @@ class _BaseStrategyChunker:
         self.max_tokens = max(1, int(max_tokens))
         self.min_tokens = max(0, int(min_tokens))
         self.duplicate_headers = duplicate_headers
-        self.encoding = get_encoding()
+        # Same unit as the embedding server counts in; see
+        # ``application.parser.tokenization``.
+        self.counter = get_token_counter()
 
     def _token_count(self, text: str) -> int:
-        return len(self.encoding.encode_ordinary(text))
+        return self.counter.count(text)
 
     def _split_by_tokens(self, text: str) -> List[str]:
         """Split ``text`` into pieces no larger than ``max_tokens`` tokens."""
-        tokens = self.encoding.encode_ordinary(text)
-        pieces = []
-        for start in range(0, len(tokens), self.max_tokens):
-            chunk_tokens = tokens[start:start + self.max_tokens]
-            pieces.append(self.encoding.decode(chunk_tokens))
-        return pieces
+        return self.counter.split(text, self.max_tokens)
 
     def _emit(self, base: Document, part_index: int, text: str) -> Document:
         """Build a child Document carrying token_count and inherited info."""
@@ -188,16 +185,11 @@ class ParentChildChunker(_BaseStrategyChunker):
         processed: List[Document] = []
         child_size = self._child_size()
         for doc in documents:
-            tokens = self.encoding.encode_ordinary(doc.text)
             part_index = 0
-            for p_start in range(0, len(tokens), self.max_tokens):
-                parent_tokens = tokens[p_start:p_start + self.max_tokens]
-                parent_text = self.encoding.decode(parent_tokens)
+            for parent_text in self.counter.split(doc.text, self.max_tokens):
                 if not parent_text.strip():
                     continue
-                for c_start in range(0, len(parent_tokens), child_size):
-                    child_tokens = parent_tokens[c_start:c_start + child_size]
-                    child_text = self.encoding.decode(child_tokens)
+                for child_text in self.counter.split(parent_text, child_size):
                     if not child_text.strip():
                         continue
                     child = Document(
@@ -208,7 +200,7 @@ class ParentChildChunker(_BaseStrategyChunker):
                         embedding=doc.embedding,
                         extra_info={
                             **(doc.extra_info or {}),
-                            "token_count": len(child_tokens),
+                            "token_count": self.counter.count(child_text),
                             "parent_text": parent_text,
                         },
                     )

application/parser/tokenization.py

@@ -0,0 +1,217 @@
+"""Counting and splitting text in the embedding model's own tokenizer.
+
+Chunk sizes only mean something in the units the embedding server counts. The
+chunker used to count cl100k (tiktoken) while the server counted whatever its
+model used, so ``max_tokens`` was a value in one unit compared against a limit
+in another. Every recalibration of that number was really an attempt to guess
+the conversion factor.
+
+This module removes the conversion. :func:`get_token_counter` returns a counter
+backed by the configured embedding model's tokenizer, falling back to cl100k
+when that tokenizer cannot be loaded -- an offline install, or a model the
+registry does not describe.
+
+Splitting never round-trips through ``decode``. Byte-level BPE decodes
+losslessly, but WordPiece tokenizers do not: all-mpnet-base-v2 lowercases, so
+decoding ``"Hello World"`` yields ``"hello world"`` and would silently rewrite
+every stored document. Counters therefore cut the *original* string at
+character offsets the tokenizer reports.
+"""
+
+from __future__ import annotations
+
+import logging
+import threading
+from typing import Iterator, List, Optional, Tuple
+
+from application.core.settings import settings
+from application.utils import get_encoding
+from application.vectorstore.model_registry import resolve
+
+logger = logging.getLogger(__name__)
+
+_cache: dict = {}
+_cache_lock = threading.Lock()
+
+
+def _windows(total: int, first: int, rest: int) -> Iterator[Tuple[int, int]]:
+    """Yield ``(start, end)`` token windows, the first sized independently.
+
+    A header consumes part of the first chunk's budget but none of the rest,
+    so the first window is often smaller than those that follow.
+    """
+    start = 0
+    budget = max(1, first)
+    while start < total:
+        end = min(start + budget, total)
+        yield start, end
+        start = end
+        budget = max(1, rest)
+
+
+class TokenCounter:
+    """Counts and splits text in one tokenizer's units.
+
+    Attributes:
+        name: Human-readable identifier of the underlying tokenizer, for logs.
+    """
+
+    name = "unknown"
+
+    def count(self, text: str) -> int:
+        """Number of tokens ``text`` occupies."""
+        raise NotImplementedError
+
+    def split(
+        self, text: str, max_tokens: int, first_max_tokens: Optional[int] = None
+    ) -> List[str]:
+        """Cut ``text`` into consecutive pieces of at most ``max_tokens``.
+
+        The concatenation of the returned pieces equals ``text`` exactly.
+
+        Args:
+            text: Source text.
+            max_tokens: Token budget per piece; values below 1 are treated as 1.
+            first_max_tokens: Budget for the first piece only, when it must
+                leave room for a header. Defaults to ``max_tokens``.
+
+        Returns:
+            The pieces, in order. An empty ``text`` yields an empty list.
+        """
+        raise NotImplementedError
+
+
+class TiktokenCounter(TokenCounter):
+    """cl100k counter -- the historical behaviour, and the fallback."""
+
+    name = "cl100k_base"
+
+    def __init__(self) -> None:
+        self._encoding = get_encoding()
+
+    def count(self, text: str) -> int:
+        if not text:
+            return 0
+        return len(self._encoding.encode_ordinary(text))
+
+    def split(
+        self, text: str, max_tokens: int, first_max_tokens: Optional[int] = None
+    ) -> List[str]:
+        if not text:
+            return []
+        rest = max(1, max_tokens)
+        first = max(1, first_max_tokens if first_max_tokens is not None else rest)
+        tokens = self._encoding.encode_ordinary(text)
+        if len(tokens) <= first:
+            return [text]
+        return [
+            self._encoding.decode(tokens[start:end])
+            for start, end in _windows(len(tokens), first, rest)
+        ]
+
+
+class HuggingFaceCounter(TokenCounter):
+    """Counts in a Hugging Face tokenizer, slicing by character offsets."""
+
+    def __init__(self, tokenizer, name: str) -> None:
+        self._tokenizer = tokenizer
+        self.name = name
+
+    def _encode(self, text: str):
+        return self._tokenizer.encode(text, add_special_tokens=False)
+
+    def count(self, text: str) -> int:
+        if not text:
+            return 0
+        return len(self._encode(text).ids)
+
+    def split(
+        self, text: str, max_tokens: int, first_max_tokens: Optional[int] = None
+    ) -> List[str]:
+        if not text:
+            return []
+        rest = max(1, max_tokens)
+        first = max(1, first_max_tokens if first_max_tokens is not None else rest)
+        offsets = self._encode(text).offsets
+        if len(offsets) <= first:
+            return [text]
+
+        pieces: List[str] = []
+        cursor = 0
+        for start_token, end_token in _windows(len(offsets), first, rest):
+            window = offsets[start_token:end_token]
+            # Some tokenizers emit (0, 0) for specials or normalised-away
+            # characters; those carry no span to cut on.
+            spans = [end for _, end in window if end > cursor]
+            if not spans:
+                continue
+            end_char = max(spans)
+            pieces.append(text[cursor:end_char])
+            cursor = end_char
+        if cursor < len(text):
+            # Trailing characters the tokenizer dropped (e.g. whitespace) belong
+            # to the final piece, so no input is lost.
+            if pieces:
+                pieces[-1] = pieces[-1] + text[cursor:]
+            else:
+                pieces.append(text[cursor:])
+        return pieces
+
+
+def _load_hf_counter(repo: str) -> Optional[HuggingFaceCounter]:
+    """Load ``repo``'s tokenizer, or ``None`` if it is not reachable."""
+    try:
+        from tokenizers import Tokenizer
+
+        tokenizer = Tokenizer.from_pretrained(repo)
+        # Repos ship padding and truncation defaults meant for inference
+        # batches. Left on, every count returns the padded width (128 for
+        # mpnet) and the offsets carry (0, 0) entries for the padding, so both
+        # counting and slicing are wrong.
+        tokenizer.no_padding()
+        tokenizer.no_truncation()
+        return HuggingFaceCounter(tokenizer, repo)
+    except Exception as exc:  # noqa: BLE001 -- chunking must never hard-fail here
+        logger.warning(
+            "Could not load the tokenizer for %s (%s); counting chunk sizes in "
+            "cl100k instead. Chunk sizes will be approximate for this model.",
+            repo,
+            exc,
+        )
+        return None
+
+
+def get_token_counter(embeddings_name: Optional[str] = None) -> TokenCounter:
+    """Return the counter for ``embeddings_name``, cached per process.
+
+    Args:
+        embeddings_name: Model name; defaults to ``settings.EMBEDDINGS_NAME``.
+
+    Returns:
+        A :class:`HuggingFaceCounter` for a model whose tokenizer could be
+        loaded, else a :class:`TiktokenCounter`.
+    """
+    name = embeddings_name or getattr(settings, "EMBEDDINGS_NAME", None)
+    key = name or "__default__"
+    with _cache_lock:
+        if key in _cache:
+            return _cache[key]
+
+    spec = resolve(name)
+    repo = spec.repo if spec else name
+    counter: TokenCounter
+    if repo and (spec is None or spec.provider == "fastembed"):
+        counter = _load_hf_counter(repo) or TiktokenCounter()
+    else:
+        counter = TiktokenCounter()
+
+    with _cache_lock:
+        _cache[key] = counter
+    logger.info("Chunking will count tokens with %s", counter.name)
+    return counter
+
+
+def reset_cache() -> None:
+    """Drop cached counters. For tests and for a settings change at runtime."""
+    with _cache_lock:
+        _cache.clear()

application/requirements.txt

@@ -20,6 +20,7 @@ ddgs>=8.0.0
 fast-ebook
 elevenlabs==2.62.0
 Flask==3.1.3
+fastembed==0.8.0
 faiss-cpu==1.15.0
 fastmcp==3.4.6
 flask-restx==1.3.2
@@ -90,13 +91,14 @@ referencing>=0.28.0,<0.38.0
 regex==2026.7.19
 requests==2.34.2
 retry==0.9.2
-sentence-transformers==5.7.0
 sqlalchemy>=2.0,<3
 starlette>=1.0,<2
 tiktoken==0.13.0
 tokenizers==0.22.2
 torch==2.11.0
 tqdm==4.67.3
+# torch and transformers are here for docling's model stack only -- embeddings
+# moved to fastembed/onnxruntime. Both can go when docling does.
 # Capped <5.9.0 to match docling-core's own darwin pin: 5.9+ breaks docling's
 # PDF layout model on Apple Silicon ("Cannot convert a MPS Tensor to float64").
 transformers==5.8.1

application/scripts/prefetch_models.py

@@ -0,0 +1,89 @@
+"""Download embedding model artifacts into FastEmbed's cache.
+
+Run at image build time so a fresh container does not download a model on its
+first ingest, and an air-gapped install works at all. Both the legacy and the
+current default are baked: an upgraded deployment keeps using mpnet until it
+runs ``reembed``, while a new one starts on granite.
+
+Usage::
+
+    python -m application.scripts.prefetch_models                 # the defaults
+    python -m application.scripts.prefetch_models granite-311m    # a subset
+"""
+
+from __future__ import annotations
+
+import logging
+import sys
+from typing import List, Optional, Sequence
+
+from application.vectorstore.model_registry import (
+    DEFAULT_LEGACY,
+    DEFAULT_NEW_INSTALL,
+    known_names,
+    resolve,
+)
+
+logger = logging.getLogger("prefetch_models")
+
+#: Fetched when no names are given.
+DEFAULT_MODELS = (DEFAULT_LEGACY, DEFAULT_NEW_INSTALL)
+
+
+def prefetch(names: Sequence[str], cache_dir: Optional[str] = None) -> List[str]:
+    """Fetch each named model's artifacts.
+
+    Args:
+        names: Registry names or aliases.
+        cache_dir: FastEmbed cache directory; its default when omitted.
+
+    Returns:
+        The repositories actually fetched.
+
+    Raises:
+        SystemExit: If a name is not in the registry, since a silent skip at
+            build time becomes a download at run time on an offline host.
+    """
+    from fastembed import TextEmbedding
+    from fastembed.common.model_description import ModelSource, PoolingType
+
+    pooling_types = {"cls": PoolingType.CLS, "mean": PoolingType.MEAN}
+    fetched: List[str] = []
+    for name in names:
+        spec = resolve(name)
+        if spec is None:
+            raise SystemExit(
+                f"Unknown embedding model {name!r}. Known: {', '.join(known_names())}"
+            )
+        if spec.provider != "fastembed":
+            logger.info("Skipping %s: served remotely, nothing to cache.", spec.name)
+            continue
+        logger.info("Fetching %s", spec.repo)
+        TextEmbedding.add_custom_model(
+            model=spec.repo,
+            pooling=pooling_types[spec.pooling],
+            normalization=spec.normalize,
+            sources=ModelSource(hf=spec.repo),
+            dim=spec.dimension,
+            model_file=spec.onnx_file,
+        )
+        kwargs = {"model_name": spec.repo}
+        if cache_dir:
+            kwargs["cache_dir"] = cache_dir
+        TextEmbedding(**kwargs)
+        fetched.append(spec.repo)
+    return fetched
+
+
+def main(argv: Optional[Sequence[str]] = None) -> int:
+    logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s")
+    import os
+
+    names = list(argv) if argv else list(DEFAULT_MODELS)
+    fetched = prefetch(names, os.environ.get("EMBEDDINGS_CACHE_DIR"))
+    logger.info("Cached %d model(s): %s", len(fetched), ", ".join(fetched))
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main(sys.argv[1:]))

application/scripts/reembed.py

@@ -0,0 +1,322 @@
+"""Re-embed an existing index with the currently configured embedding model.
+
+Changing ``EMBEDDINGS_NAME`` does not change vectors already stored. Because
+several models share a width -- granite-311m and all-mpnet-base-v2 are both
+768-dimensional -- a swapped model does not fail any dimension check. It simply
+retrieves against vectors that mean something else, which shows up as bad
+answers rather than as an error.
+
+This script rebuilds the vectors from the chunk text already held in the store,
+so it never re-downloads, re-parses or re-chunks anything: no source files, no
+crawl, no docling. It works against pgvector and FAISS.
+
+Usage::
+
+    python -m application.scripts.reembed --dry-run     # report, change nothing
+    python -m application.scripts.reembed               # re-embed everything
+    python -m application.scripts.reembed --sources a,b # only these sources
+"""
+
+from __future__ import annotations
+
+import argparse
+import logging
+import sys
+import time
+from dataclasses import dataclass
+from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple
+
+from application.core.settings import settings
+from application.vectorstore.model_registry import resolve
+from application.vectorstore.vector_creator import VectorCreator
+
+logger = logging.getLogger("reembed")
+
+#: Stores this script knows how to rewrite in place.
+SUPPORTED_STORES = ("pgvector", "faiss")
+
+#: Chunks embedded (and, for pgvector, written) per transaction.
+DEFAULT_BATCH_SIZE = 64
+
+
+@dataclass
+class _RebuildDoc:
+    """Minimal document shape ``FaissStore(docs_init=...)`` accepts."""
+
+    page_content: str
+    metadata: Dict[str, Any]
+
+
+class ReembedError(RuntimeError):
+    """Raised for a misconfiguration the user needs to fix before running."""
+
+
+def _log_setup(verbose: bool) -> None:
+    logging.basicConfig(
+        level=logging.DEBUG if verbose else logging.INFO,
+        format="%(asctime)s %(levelname)s %(message)s",
+    )
+
+
+def _batched(items: Sequence[Any], size: int) -> Iterable[Sequence[Any]]:
+    for start in range(0, len(items), size):
+        yield items[start : start + size]
+
+
+def list_source_ids(store_type: str) -> List[str]:
+    """Every source id that has vectors in the configured store.
+
+    Args:
+        store_type: ``"pgvector"`` or ``"faiss"``.
+
+    Returns:
+        Source ids, sorted for a stable and resumable ordering.
+    """
+    if store_type == "pgvector":
+        return _pgvector_source_ids()
+    return _faiss_source_ids()
+
+
+def _pgvector_source_ids() -> List[str]:
+    store = VectorCreator.create_vectorstore("pgvector", source_id="")
+    # ``_get_connection`` hands back a pooled connection, not a context
+    # manager: closing it via ``with`` would cost the pool a slot.
+    conn = store._get_connection()
+    cursor = conn.cursor()
+    try:
+        cursor.execute(
+            f"SELECT DISTINCT source_id FROM {store._table_name} "
+            f"WHERE source_id IS NOT NULL ORDER BY source_id"
+        )
+        return [row[0] for row in cursor.fetchall()]
+    finally:
+        cursor.close()
+        store.close()
+
+
+def _faiss_source_ids() -> List[str]:
+    """Directories under the FAISS root that hold an index."""
+    from application.storage.storage_creator import StorageCreator
+
+    storage = StorageCreator.get_storage()
+    root = "indexes"
+    ids: List[str] = []
+    try:
+        for path in storage.list_files(root):
+            # ``indexes/<source_id>/index.faiss``
+            parts = str(path).split("/")
+            if parts and parts[-1].endswith(".faiss") and len(parts) >= 2:
+                ids.append(parts[-2])
+    except Exception as exc:  # noqa: BLE001 -- surfaced to the user below
+        raise ReembedError(
+            f"Could not list FAISS indexes under {root!r}: {exc}"
+        ) from exc
+    return sorted(set(ids))
+
+
+def reembed_pgvector(source_id: str, batch_size: int, dry_run: bool) -> Tuple[int, int]:
+    """Rewrite one source's vectors in place.
+
+    Rows are updated, never deleted and re-inserted, so an interrupted run
+    leaves every chunk present and simply re-does its batch next time.
+
+    Args:
+        source_id: Source whose chunks to re-embed.
+        batch_size: Chunks per embed call and per transaction.
+        dry_run: When true, count the work and change nothing.
+
+    Returns:
+        ``(chunks_seen, chunks_written)``.
+    """
+    store = VectorCreator.create_vectorstore("pgvector", source_id=source_id)
+    table, vector_column = store._table_name, store._vector_column
+    conn = store._get_connection()
+
+    cursor = conn.cursor()
+    try:
+        cursor.execute(
+            f"SELECT id, text FROM {table} WHERE source_id = %s ORDER BY id",
+            (source_id,),
+        )
+        rows = cursor.fetchall()
+    finally:
+        cursor.close()
+
+    if dry_run or not rows:
+        store.close()
+        return len(rows), 0
+
+    written = 0
+    try:
+        for batch in _batched(rows, batch_size):
+            vectors = store._embedding.embed_documents([row[1] or "" for row in batch])
+            cursor = conn.cursor()
+            try:
+                cursor.executemany(
+                    f"UPDATE {table} SET {vector_column} = %s::vector WHERE id = %s",
+                    [
+                        (str(list(vector)), row[0])
+                        for vector, row in zip(vectors, batch)
+                    ],
+                )
+                conn.commit()
+            except Exception:
+                conn.rollback()
+                raise
+            finally:
+                cursor.close()
+            written += len(batch)
+    finally:
+        store.close()
+    return len(rows), written
+
+
+def reembed_faiss(source_id: str, batch_size: int, dry_run: bool) -> Tuple[int, int]:
+    """Rebuild one FAISS index from the chunk text in its sidecar.
+
+    A FAISS index is a flat array, so this rebuilds rather than updates. Every
+    new vector is computed *before* the old index is touched, so an interrupt
+    during embedding leaves the existing index intact.
+
+    Args:
+        source_id: Source whose index to rebuild.
+        batch_size: Chunks per embed call.
+        dry_run: When true, count the work and change nothing.
+
+    Returns:
+        ``(chunks_seen, chunks_written)``.
+    """
+    store = VectorCreator.create_vectorstore(
+        "faiss", source_id=source_id, embeddings_key="embeddings"
+    )
+    chunks: List[Dict[str, Any]] = store.get_chunks() or []
+    if dry_run or not chunks:
+        return len(chunks), 0
+
+    docs = [
+        _RebuildDoc(chunk.get("text") or "", chunk.get("metadata") or {})
+        for chunk in chunks
+    ]
+
+    # Constructing with ``docs_init`` embeds everything into a fresh in-memory
+    # index and touches nothing on disk. The existing index is only replaced by
+    # ``save_local`` below, so a failure while embedding leaves it intact.
+    rebuilt = VectorCreator.create_vectorstore(
+        "faiss", source_id=source_id, embeddings_key="embeddings", docs_init=docs
+    )
+    rebuilt.save_local()
+    return len(chunks), len(docs)
+
+
+def run(
+    store_type: str,
+    source_ids: Optional[Sequence[str]],
+    batch_size: int,
+    dry_run: bool,
+) -> int:
+    """Re-embed the requested sources. Returns a process exit code."""
+    model = settings.EMBEDDINGS_NAME
+    spec = resolve(model)
+    logger.info(
+        "Re-embedding %s with %s (%s)",
+        store_type,
+        model,
+        f"{spec.dimension} dims" if spec else "width determined at runtime",
+    )
+    if dry_run:
+        logger.info("Dry run: nothing will be written.")
+
+    ids = list(source_ids) if source_ids else list_source_ids(store_type)
+    if not ids:
+        logger.warning("No sources found; nothing to do.")
+        return 0
+    logger.info("%d source(s) to process", len(ids))
+
+    handler = reembed_pgvector if store_type == "pgvector" else reembed_faiss
+    total_seen = total_written = 0
+    failures: List[str] = []
+    started = time.monotonic()
+
+    for position, source_id in enumerate(ids, start=1):
+        try:
+            seen, written = handler(source_id, batch_size, dry_run)
+            total_seen += seen
+            total_written += written
+            logger.info(
+                "[%d/%d] %s: %d chunk(s)%s",
+                position,
+                len(ids),
+                source_id,
+                seen,
+                "" if dry_run else f", {written} re-embedded",
+            )
+        except Exception as exc:  # noqa: BLE001 -- one bad source must not stop the run
+            failures.append(source_id)
+            logger.error("[%d/%d] %s FAILED: %s", position, len(ids), source_id, exc)
+
+    elapsed = time.monotonic() - started
+    logger.info(
+        "Done in %.1fs: %d chunk(s) seen, %d re-embedded, %d source(s) failed",
+        elapsed,
+        total_seen,
+        total_written,
+        len(failures),
+    )
+    if failures:
+        logger.error("Failed sources: %s", ", ".join(failures))
+        logger.error("Re-run with --sources %s to retry just those.", ",".join(failures))
+        return 1
+    if dry_run:
+        logger.info("Re-run without --dry-run to apply.")
+    return 0
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        prog="python -m application.scripts.reembed",
+        description=__doc__.split("\n\n")[0],
+    )
+    parser.add_argument(
+        "--dry-run", action="store_true", help="report what would change, write nothing"
+    )
+    parser.add_argument(
+        "--sources", help="comma-separated source ids; default is every source"
+    )
+    parser.add_argument(
+        "--batch-size",
+        type=int,
+        default=DEFAULT_BATCH_SIZE,
+        help=f"chunks per embed call (default {DEFAULT_BATCH_SIZE})",
+    )
+    parser.add_argument("-v", "--verbose", action="store_true", help="debug logging")
+    return parser
+
+
+def main(argv: Optional[Sequence[str]] = None) -> int:
+    args = build_parser().parse_args(argv)
+    _log_setup(args.verbose)
+
+    store_type = (settings.VECTOR_STORE or "").lower()
+    if store_type not in SUPPORTED_STORES:
+        logger.error(
+            "VECTOR_STORE is %r; this script supports %s. For other stores, "
+            "delete and re-ingest the affected sources.",
+            settings.VECTOR_STORE,
+            " and ".join(SUPPORTED_STORES),
+        )
+        return 2
+
+    sources = (
+        [s.strip() for s in args.sources.split(",") if s.strip()]
+        if args.sources
+        else None
+    )
+    try:
+        return run(store_type, sources, max(1, args.batch_size), args.dry_run)
+    except ReembedError as exc:
+        logger.error("%s", exc)
+        return 2
+
+
+if __name__ == "__main__":
+    sys.exit(main())

application/vectorstore/base.py

@@ -1,5 +1,4 @@
 import logging
-import os
 from abc import ABC, abstractmethod
 from typing import Optional
 
@@ -7,7 +6,11 @@ import requests
 
 from application.core.settings import settings
 from application.vectorstore.embeddings_openai import OpenAIEmbeddings
-from application.utils import get_encoding
+from application.vectorstore.model_registry import (
+    dimension_for,
+    max_input_tokens_for,
+    resolve,
+)
 
 
 class RemoteEmbeddings:
@@ -23,45 +26,75 @@ class RemoteEmbeddings:
         self.headers = {"Content-Type": "application/json"}
         if api_key:
             self.headers["Authorization"] = f"Bearer {api_key}"
-        self.dimension = 768
+        # Width comes from the registry. This used to be a hardcoded 768 that
+        # ``embed_query`` claimed to correct on first use -- but the correction
+        # was guarded by ``if self.dimension is None``, which the hardcode made
+        # unreachable, so a remote model of any other width silently produced a
+        # ``vector(768)`` column. ``None`` here means "unknown", and the probe
+        # below now genuinely runs.
+        self.dimension = dimension_for(model_name)
+
+    def _token_counter(self):
+        """Counter matching the remote model's tokenizer, cached per process."""
+        from application.parser.tokenization import get_token_counter
+
+        return get_token_counter(self.model_name)
+
+    def _resolve_input_limit(self):
+        """Token ceiling for a single embed input, or ``None`` for no limit.
+
+        ``EMBEDDINGS_MAX_INPUT_TOKENS`` wins when set. Otherwise a registered
+        model contributes its own context window, so a request that the server
+        would reject -- or silently truncate -- is clipped here instead of
+        being sent and paid for.
+        """
+        configured = settings.EMBEDDINGS_MAX_INPUT_TOKENS
+        if configured and configured > 0:
+            return configured
+        model_limit = max_input_tokens_for(self.model_name)
+        return model_limit if model_limit and model_limit > 0 else None
 
     def _truncate_inputs(self, inputs):
-        """Clip each input to ``EMBEDDINGS_MAX_INPUT_TOKENS`` tokens.
+        """Clip each input to the resolved token limit.
 
         The remote server (e.g. llama.cpp) hard-rejects any single input
-        larger than its physical batch size with a 500. When the setting is
-        configured, each input is truncated to that many tokens before the
-        request and the overflow is dropped (lossy by design). Token counts
-        use the shared tiktoken encoding, which differs from the server's
-        tokenizer, so set the limit with headroom under the server's true
-        limit to absorb tokenizer skew.
+        larger than its physical batch size with a 500, so oversized inputs are
+        truncated before the request and the overflow is dropped (lossy by
+        design).
+
+        Counting uses the embedding model's own tokenizer where it is known, so
+        the limit and the count are in the same unit. When it is not -- an
+        unregistered model, or no tokenizer available -- this falls back to
+        tiktoken, and the limit should then carry headroom to absorb the skew
+        between the two tokenizers.
 
         Args:
             inputs: A single string or a list of strings to embed.
 
         Returns:
             The inputs with each string clipped to the token limit, or the
-            inputs unchanged when the limit is unset or non-positive.
+            inputs unchanged when no limit applies.
         """
-        limit = settings.EMBEDDINGS_MAX_INPUT_TOKENS
-        if not limit or limit <= 0:
+        limit = self._resolve_input_limit()
+        if not limit:
             return inputs
 
-        encoding = get_encoding()
+        counter = self._token_counter()
 
         def clip(text):
             if not isinstance(text, str):
                 return text
-            tokens = encoding.encode_ordinary(text)
-            if len(tokens) <= limit:
+            count = counter.count(text)
+            if count <= limit:
                 return text
             logging.warning(
                 "Truncating remote embeddings input from %d to %d tokens (%d dropped)",
-                len(tokens),
+                count,
                 limit,
-                len(tokens) - limit,
+                count - limit,
             )
-            return encoding.decode(tokens[:limit])
+            pieces = counter.split(text, limit)
+            return pieces[0] if pieces else text
 
         if isinstance(inputs, list):
             return [clip(text) for text in inputs]
@@ -129,7 +162,7 @@ class RemoteEmbeddings:
 
 
 def _get_embeddings_wrapper():
-    """Lazy import of EmbeddingsWrapper to avoid loading SentenceTransformer when using remote embeddings."""
+    """Lazy import of EmbeddingsWrapper, so a remote setup never loads ONNX."""
     from application.vectorstore.embeddings_local import EmbeddingsWrapper
 
     return EmbeddingsWrapper
@@ -178,36 +211,28 @@ class EmbeddingsSingleton:
 
     @staticmethod
     def _create_instance(embeddings_name, *args, **kwargs):
-        if embeddings_name == "openai_text-embedding-ada-002":
+        """Build the runner for ``embeddings_name``, per the model registry.
+
+        The registry replaced a hand-maintained factory dict whose entries
+        existed only to rewrite a configured name into a repository id. That
+        rewrite is now a registry field, so an unknown name needs no entry
+        here: it is passed through as a Hugging Face repository.
+        """
+        spec = resolve(embeddings_name)
+        if spec is not None and spec.provider == "openai":
             return OpenAIEmbeddings(*args, **kwargs)
 
-        # Lazy import EmbeddingsWrapper only when needed (avoids loading SentenceTransformer)
         EmbeddingsWrapper = _get_embeddings_wrapper()
-
-        embeddings_factory = {
-            "huggingface_sentence-transformers/all-mpnet-base-v2": lambda: EmbeddingsWrapper(
-                "sentence-transformers/all-mpnet-base-v2"
-            ),
-            "huggingface_sentence-transformers-all-mpnet-base-v2": lambda: EmbeddingsWrapper(
-                "sentence-transformers/all-mpnet-base-v2"
-            ),
-            "huggingface_hkunlp/instructor-large": lambda: EmbeddingsWrapper(
-                "hkunlp/instructor-large"
-            ),
-        }
-
-        if embeddings_name in embeddings_factory:
-            if args or kwargs:
-                logging.debug(
-                    "Dropping %d positional and %d keyword argument(s) for pinned "
-                    "embeddings model %s: its factory takes none.",
-                    len(args),
-                    len(kwargs),
-                    embeddings_name,
-                )
-            return embeddings_factory[embeddings_name]()
-        else:
-            return EmbeddingsWrapper(embeddings_name, *args, **kwargs)
+        if spec is not None and (args or kwargs):
+            logging.debug(
+                "Dropping %d positional and %d keyword argument(s) for registered "
+                "embeddings model %s: the registry supplies its configuration.",
+                len(args),
+                len(kwargs),
+                embeddings_name,
+            )
+            return EmbeddingsWrapper(embeddings_name)
+        return EmbeddingsWrapper(embeddings_name, *args, **kwargs)
 
 
 def _azure_configured() -> bool:
@@ -225,11 +250,8 @@ def get_embeddings(
     """Resolve the configured embeddings instance. The single entry point.
 
     Callers that reach for :meth:`EmbeddingsSingleton.get_instance` directly
-    reproduce neither the bundled local-model path (the Docker image ships
-    ``/app/models/all-mpnet-base-v2``) nor the OpenAI/Azure key handling: they
-    download a second copy of the model from the hub, and passing the key
-    positionally to the pinned HuggingFace names raises ``TypeError`` because
-    those factories take no arguments. Route every caller through here.
+    skip the remote dispatch and the OpenAI/Azure key handling. Route every
+    caller through here.
 
     Args:
         embeddings_name: Model name; defaults to ``settings.EMBEDDINGS_NAME``.
@@ -259,31 +281,10 @@ def get_embeddings(
             embedding_instance = EmbeddingsSingleton.get_instance(
                 embeddings_name, openai_api_key=embeddings_key
             )
-    elif embeddings_name == "huggingface_sentence-transformers/all-mpnet-base-v2":
-        possible_paths = [
-            "/app/models/all-mpnet-base-v2",  # Docker absolute path
-            "./models/all-mpnet-base-v2",  # Relative path
-        ]
-        local_model_path = None
-        for path in possible_paths:
-            if os.path.exists(path):
-                local_model_path = path
-                logging.info(f"Found local model at path: {path}")
-                break
-            else:
-                logging.info(f"Path does not exist: {path}")
-        if local_model_path:
-            embedding_instance = EmbeddingsSingleton.get_instance(
-                local_model_path,
-            )
-        else:
-            logging.warning(
-                f"Local model not found in any of the paths: {possible_paths}. Falling back to HuggingFace download."
-            )
-            embedding_instance = EmbeddingsSingleton.get_instance(
-                embeddings_name,
-            )
     else:
+        # No per-model branching: the registry resolves names and FastEmbed
+        # caches artifacts under EMBEDDINGS_CACHE_DIR, which is where the
+        # image warms them at build time.
         embedding_instance = EmbeddingsSingleton.get_instance(embeddings_name)
     return embedding_instance
 

application/vectorstore/embeddings_local.py

@@ -1,54 +1,189 @@
-"""
-Local embeddings using SentenceTransformer.
-This module is only imported when EMBEDDINGS_BASE_URL is not set,
-to avoid loading SentenceTransformer into memory when using remote embeddings.
+"""Local embeddings via FastEmbed (ONNX Runtime).
+
+Replaces the previous SentenceTransformer implementation. Both run the same
+weights; FastEmbed reaches them through ONNX Runtime instead of torch, which
+removes torch, transformers and sentence-transformers from the dependency set
+and measurably reduces both resident memory and import cost.
+
+The swap is numerically transparent for existing indexes: embedding the same
+text with SentenceTransformer and with FastEmbed's ONNX graph of the same
+model returns vectors at cosine 1.0, so an mpnet index built before this
+change keeps working unchanged.
+
+Models are described in :mod:`application.vectorstore.model_registry`. A name
+the registry does not know is treated as a Hugging Face repository, which is
+what someone configuring an arbitrary model expects.
 """
 
 import logging
+import threading
+from typing import Any, List, Optional
 
-from sentence_transformers import SentenceTransformer
+from application.core.settings import settings
+from application.vectorstore.model_registry import EmbeddingModel, resolve, known_names
+
+logger = logging.getLogger(__name__)
+
+# ``add_custom_model`` mutates a process-global registry inside FastEmbed, so
+# repeated registration of the same name is both wasteful and racy under the
+# thread pool the API serves requests from.
+_registered: set = set()
+_register_lock = threading.Lock()
+
+# Assumed layout for a model the registry does not describe.
+_FALLBACK_ONNX_FILE = "onnx/model.onnx"
+_FALLBACK_POOLING = "mean"
+
+
+def _pooling_type(pooling: str):
+    """Map our ``"cls"``/``"mean"`` spelling onto FastEmbed's enum."""
+    from fastembed.common.model_description import PoolingType
+
+    return {"cls": PoolingType.CLS, "mean": PoolingType.MEAN}[pooling]
+
+
+def _register(model: EmbeddingModel) -> None:
+    """Teach FastEmbed about a model, exactly once per process."""
+    from fastembed import TextEmbedding
+    from fastembed.common.model_description import ModelSource
+
+    with _register_lock:
+        if model.repo in _registered:
+            return
+        TextEmbedding.add_custom_model(
+            model=model.repo,
+            pooling=_pooling_type(model.pooling),
+            normalization=model.normalize,
+            sources=ModelSource(hf=model.repo),
+            dim=model.dimension,
+            model_file=model.onnx_file,
+        )
+        _registered.add(model.repo)
+
+
+def _spec_for(model_name: str) -> EmbeddingModel:
+    """Registry entry for ``model_name``, or a best-effort one for a raw repo.
+
+    Args:
+        model_name: Configured ``EMBEDDINGS_NAME``.
+
+    Returns:
+        A registry entry. For an unregistered name the returned entry assumes
+        the standard Hugging Face ONNX layout and mean pooling, and carries
+        ``dimension = 0`` so the caller knows to probe for the real width.
+    """
+    spec = resolve(model_name)
+    if spec is not None:
+        return spec
+    logger.info(
+        "Embedding model %r is not in the registry (known: %s); treating it as a "
+        "Hugging Face repository with %s and %s pooling.",
+        model_name,
+        ", ".join(known_names()),
+        _FALLBACK_ONNX_FILE,
+        _FALLBACK_POOLING,
+    )
+    return EmbeddingModel(
+        name=model_name,
+        dimension=0,
+        max_input_tokens=512,
+        pooling=_FALLBACK_POOLING,
+        normalize=True,
+        repo=model_name,
+        onnx_file=_FALLBACK_ONNX_FILE,
+    )
 
 
 class EmbeddingsWrapper:
-    def __init__(self, model_name, *args, **kwargs):
-        logging.info(f"Initializing EmbeddingsWrapper with model: {model_name}")
+    """Runs an embedding model locally through FastEmbed.
+
+    Exposes the ``embed_query``/``embed_documents``/``dimension`` interface the
+    vector stores rely on, matching ``RemoteEmbeddings`` and ``OpenAIEmbeddings``.
+    """
+
+    def __init__(self, model_name: str, *args: Any, **kwargs: Any) -> None:
+        """Load ``model_name`` locally.
+
+        Args:
+            model_name: Registry name, alias, or a Hugging Face repository id.
+
+        Raises:
+            RuntimeError: If the model cannot be loaded, with the configured
+                name and the registry's known names in the message.
+        """
+        from fastembed import TextEmbedding
+
+        self.spec = _spec_for(model_name)
+        logger.info("Loading embeddings model %s via FastEmbed", self.spec.repo)
         try:
-            kwargs.setdefault("trust_remote_code", True)
-            self.model = SentenceTransformer(
-                model_name,
-                config_kwargs={"allow_dangerous_deserialization": True},
-                *args,
-                **kwargs,
-            )
-            if self.model is None or self.model._first_module() is None:
-                raise ValueError(
-                    f"SentenceTransformer model failed to load properly for: {model_name}"
-                )
-            # Renamed in sentence-transformers 5.4; keep the old name as a fallback.
-            get_dimension = getattr(
-                self.model,
-                "get_embedding_dimension",
-                getattr(self.model, "get_sentence_embedding_dimension", None),
-            )
-            self.dimension = get_dimension()
-            logging.info(f"Successfully loaded model with dimension: {self.dimension}")
-        except Exception as e:
-            logging.error(
-                f"Failed to initialize SentenceTransformer with model {model_name}: {str(e)}",
-                exc_info=True,
-            )
-            raise
+            _register(self.spec)
+            init_kwargs = {"model_name": self.spec.repo}
+            threads = getattr(settings, "EMBEDDINGS_THREADS", None)
+            if isinstance(threads, int) and threads > 0:
+                init_kwargs["threads"] = threads
+            cache_dir = getattr(settings, "EMBEDDINGS_CACHE_DIR", None)
+            if cache_dir:
+                init_kwargs["cache_dir"] = cache_dir
+            self.model = TextEmbedding(**init_kwargs)
+        except Exception as exc:
+            raise RuntimeError(
+                f"Could not load embeddings model {model_name!r} via FastEmbed: "
+                f"{exc}. Known models: {', '.join(known_names())}."
+            ) from exc
 
-    def embed_query(self, query: str):
-        return self.model.encode(query).tolist()
+        self.dimension = self.spec.dimension or self._probe_dimension()
+        logger.info("Embeddings model ready (dimension=%d)", self.dimension)
 
-    def embed_documents(self, documents: list):
-        return self.model.encode(documents).tolist()
+    def _probe_dimension(self) -> int:
+        """Determine the vector width of a model the registry does not describe."""
+        return len(self.embed_query("dimension probe"))
+
+    @property
+    def tokenizer(self):
+        """The model's own tokenizer, so chunking can count in its units."""
+        return self.model.model.tokenizer
+
+    def embed_query(self, query: str) -> List[float]:
+        """Embed a single query string."""
+        return self.embed_documents([query])[0]
+
+    def embed_documents(self, documents: List[str]) -> List[List[float]]:
+        """Embed a list of documents, preserving input order.
+
+        Inputs are grouped by length before batching. ONNX needs a rectangular
+        tensor, so every input in a batch is padded up to the longest one in
+        it; with mixed lengths that padding is most of the work. Grouping
+        similar lengths together measured 19% faster and 45% lower peak memory
+        on production-sized chunks, and at full context an unsorted batch of 4
+        was slower than no batching at all.
+
+        The original order is restored before returning, so callers zipping
+        these against their texts are unaffected.
+        """
+        if not documents:
+            return []
+        batch_size: Optional[int] = None
+        raw = getattr(settings, "EMBEDDINGS_BATCH_SIZE", None)
+        if isinstance(raw, int) and not isinstance(raw, bool) and raw > 0:
+            batch_size = raw
+
+        documents = list(documents)
+        if batch_size is None or len(documents) <= batch_size:
+            # One batch either way: sorting would only add work.
+            return [v.tolist() for v in self.model.embed(documents, batch_size=batch_size)]
+
+        order = sorted(range(len(documents)), key=lambda i: len(documents[i]))
+        grouped = [documents[i] for i in order]
+        vectors = [v.tolist() for v in self.model.embed(grouped, batch_size=batch_size)]
+
+        restored: List[Optional[List[float]]] = [None] * len(documents)
+        for position, original_index in enumerate(order):
+            restored[original_index] = vectors[position]
+        return restored
 
     def __call__(self, text):
         if isinstance(text, str):
             return self.embed_query(text)
         elif isinstance(text, list):
             return self.embed_documents(text)
-        else:
-            raise ValueError("Input must be a string or a list of strings")
+        raise ValueError("Input must be a string or a list of strings")

application/vectorstore/faiss.py

@@ -301,24 +301,29 @@ class FaissStore(BaseVectorStore):
     # -- Introspection ---------------------------------------------------
 
     def assert_embedding_dimensions(self, embeddings) -> None:
-        """Check the index width matches the embedding model's width."""
-        if (
-            settings.EMBEDDINGS_NAME
-            == "huggingface_sentence-transformers/all-mpnet-base-v2"
-        ):
-            word_embedding_dimension = getattr(embeddings, "dimension", None)
-            if word_embedding_dimension is None:
-                raise AttributeError(
-                    "'dimension' attribute not found in embeddings instance."
-                )
-            if self.index is None:
-                return
-            if word_embedding_dimension != self.index.d:
-                raise ValueError(
-                    f"Embedding dimension mismatch: embeddings.dimension "
-                    f"({word_embedding_dimension}) != docsearch index dimension "
-                    f"({self.index.d})"
-                )
+        """Check the index width matches the embedding model's width.
+
+        This used to run only when ``EMBEDDINGS_NAME`` was mpnet, so every
+        other model skipped the check entirely -- exactly the models most
+        likely to differ from an index built earlier. It now runs for any
+        model that reports a width.
+        """
+        word_embedding_dimension = getattr(embeddings, "dimension", None)
+        if word_embedding_dimension is None:
+            # A remote model of unknown width reports None until its first
+            # call; there is nothing to compare yet.
+            return
+        if self.index is None:
+            return
+        if word_embedding_dimension != self.index.d:
+            raise ValueError(
+                f"Embedding dimension mismatch: {settings.EMBEDDINGS_NAME} produces "
+                f"{word_embedding_dimension}-dim vectors but this FAISS index is "
+                f"{self.index.d}-dim. The index was built with a different "
+                f"embedding model; re-embed it with "
+                f"`python -m application.scripts.reembed` or point "
+                f"EMBEDDINGS_NAME back at the original model."
+            )
 
     def get_chunks(self) -> List[Dict[str, Any]]:
         """Return every chunk held in the index."""

application/vectorstore/model_registry.py

@@ -0,0 +1,170 @@
+"""Canonical description of every embedding model DocsGPT knows how to run.
+
+``EMBEDDINGS_NAME`` used to be a free-form string interpreted in half a dozen
+places: a factory dict here, a bundled-model path probe there, a dimension
+assertion gated on one model's name, and a hardcoded ``dimension = 768`` on the
+remote client. Each of those encoded a different subset of the same facts, and
+they drifted.
+
+This module is the single place those facts live. A model is described once and
+every consumer -- the local runner, the remote client, the schema bootstrap, the
+chunker -- reads the same entry.
+
+Unknown names are not an error: :func:`resolve` returns ``None`` and callers
+fall back to treating the name as a Hugging Face repository, which is what a
+user configuring an arbitrary model expects.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Dict, Optional, Tuple
+
+
+@dataclass(frozen=True)
+class EmbeddingModel:
+    """Everything the application needs to know about one embedding model.
+
+    Attributes:
+        name: Canonical ``EMBEDDINGS_NAME`` value.
+        dimension: Width of the vectors it produces.
+        max_input_tokens: The model's own context window, counted in *its*
+            tokenizer. Used to bound what we send it, never to silently
+            reshape chunks.
+        pooling: ``"cls"`` or ``"mean"`` -- how token vectors become one vector.
+        normalize: Whether outputs are L2-normalised to unit length.
+        provider: Which runner handles it (``"fastembed"`` or ``"openai"``).
+        repo: Hugging Face repository holding weights and tokenizer.
+        onnx_file: Path within ``repo`` to the ONNX graph to run.
+        aliases: Other spellings of ``name`` accepted from configuration.
+    """
+
+    name: str
+    dimension: int
+    max_input_tokens: int
+    pooling: str = "mean"
+    normalize: bool = True
+    provider: str = "fastembed"
+    repo: Optional[str] = None
+    onnx_file: Optional[str] = None
+    aliases: Tuple[str, ...] = field(default_factory=tuple)
+
+
+#: The model DocsGPT installed before the granite migration. Kept as the
+#: default so an existing deployment that upgrades keeps its index working;
+#: new installs are pointed at granite by the setup script and env template.
+MPNET = EmbeddingModel(
+    name="huggingface_sentence-transformers/all-mpnet-base-v2",
+    dimension=768,
+    max_input_tokens=384,
+    pooling="mean",
+    normalize=True,
+    repo="sentence-transformers/all-mpnet-base-v2",
+    onnx_file="onnx/model.onnx",
+    aliases=(
+        "huggingface_sentence-transformers-all-mpnet-base-v2",
+        "sentence-transformers/all-mpnet-base-v2",
+        "all-mpnet-base-v2",
+    ),
+)
+
+#: Default for new installs: same 768 dimensions as mpnet, an 8x wider
+#: effective input, and multilingual retrieval.
+GRANITE_311M = EmbeddingModel(
+    name="ibm-granite/granite-embedding-311m-multilingual-r2",
+    dimension=768,
+    max_input_tokens=32768,
+    pooling="cls",
+    normalize=True,
+    repo="ibm-granite/granite-embedding-311m-multilingual-r2",
+    onnx_file="onnx/model_quint8_avx2.onnx",
+    aliases=("granite-embedding-311m-multilingual-r2", "granite-311m"),
+)
+
+#: Smaller granite. Half the vector width, roughly three times the speed.
+GRANITE_97M = EmbeddingModel(
+    name="ibm-granite/granite-embedding-97m-multilingual-r2",
+    dimension=384,
+    max_input_tokens=32768,
+    pooling="cls",
+    normalize=True,
+    repo="ibm-granite/granite-embedding-97m-multilingual-r2",
+    onnx_file="onnx/model_quint8_avx2.onnx",
+    aliases=("granite-embedding-97m-multilingual-r2", "granite-97m"),
+)
+
+OPENAI_ADA_002 = EmbeddingModel(
+    name="openai_text-embedding-ada-002",
+    dimension=1536,
+    max_input_tokens=8191,
+    pooling="mean",
+    normalize=True,
+    provider="openai",
+    repo=None,
+    onnx_file=None,
+    aliases=("text-embedding-ada-002",),
+)
+
+MODELS: Tuple[EmbeddingModel, ...] = (
+    MPNET,
+    GRANITE_311M,
+    GRANITE_97M,
+    OPENAI_ADA_002,
+)
+
+#: Fallback vector width when the configured model is unknown -- the width
+#: every DocsGPT install has used to date, so an unrecognised model does not
+#: silently reshape an existing table.
+DEFAULT_EMBEDDING_DIMENSION = 768
+
+#: Name that a fresh install should be configured with.
+DEFAULT_NEW_INSTALL = GRANITE_311M.name
+
+#: Name that ``settings.EMBEDDINGS_NAME`` defaults to, i.e. what an existing
+#: deployment falls back to when it never pinned one.
+DEFAULT_LEGACY = MPNET.name
+
+
+def _index() -> Dict[str, EmbeddingModel]:
+    """Build the lookup table of every accepted spelling."""
+    table: Dict[str, EmbeddingModel] = {}
+    for model in MODELS:
+        for key in (model.name, *model.aliases):
+            table[key.lower()] = model
+    return table
+
+
+_LOOKUP = _index()
+
+
+def resolve(name: Optional[str]) -> Optional[EmbeddingModel]:
+    """Return the registry entry for ``name``, or ``None`` when unknown.
+
+    Args:
+        name: A configured ``EMBEDDINGS_NAME`` value, in any accepted spelling.
+
+    Returns:
+        The matching :class:`EmbeddingModel`, or ``None`` for a name the
+        registry does not describe -- which callers treat as a Hugging Face
+        repository rather than an error.
+    """
+    if not name:
+        return None
+    return _LOOKUP.get(name.strip().lower())
+
+
+def dimension_for(name: Optional[str]) -> Optional[int]:
+    """Vector width for ``name``, or ``None`` when unknown."""
+    model = resolve(name)
+    return model.dimension if model else None
+
+
+def max_input_tokens_for(name: Optional[str]) -> Optional[int]:
+    """Context window for ``name``, or ``None`` when unknown."""
+    model = resolve(name)
+    return model.max_input_tokens if model else None
+
+
+def known_names() -> Tuple[str, ...]:
+    """Canonical names of every registered model, for error messages."""
+    return tuple(model.name for model in MODELS)

docs/content/upgrading.mdx

@@ -11,6 +11,45 @@ import { Callout } from 'nextra/components'
   **Upgrading from 0.16.x?** User data moved from MongoDB to Postgres in 0.17.0. Follow the [Postgres Migration guide](/Deploying/Postgres-Migration) before running `docker compose pull` or `git pull` — existing deployments will not start cleanly without it.
 </Callout>
 
+## Embedding models
+
+DocsGPT now runs embeddings through [FastEmbed](https://github.com/qdrant/fastembed) (ONNX Runtime) instead of SentenceTransformer. The models are the same and the vectors are identical, so **an existing deployment needs no action** — `all-mpnet-base-v2` keeps working exactly as before.
+
+New installs default to `ibm-granite/granite-embedding-311m-multilingual-r2`: multilingual, a 32k-token context, and the same 768 dimensions.
+
+### Switching an existing deployment to granite
+
+<Callout type="error">
+  Changing `EMBEDDINGS_NAME` on an index that already has vectors **breaks retrieval silently**. Both models are 768-dimensional, so nothing raises an error — queries are simply compared against vectors that mean something else, and answers quietly get worse. Always re-embed.
+</Callout>
+
+Set the model, then rebuild the vectors:
+
+```bash
+# 1. In your .env
+EMBEDDINGS_NAME=ibm-granite/granite-embedding-311m-multilingual-r2
+
+# 2. Rebuild the vectors from the chunk text already in your index
+docker compose exec backend python -m application.scripts.reembed --dry-run
+docker compose exec backend python -m application.scripts.reembed
+```
+
+Re-embedding reads the chunk text already stored in your index. It does not re-download, re-parse or re-chunk your documents, so no source files are needed and the run is proportional to index size, not corpus size. Both `pgvector` and `faiss` are supported.
+
+Useful flags:
+
+| Flag | Effect |
+| --- | --- |
+| `--dry-run` | Report how many chunks would change, write nothing |
+| `--sources a,b` | Only these source ids — also how you retry a failed source |
+| `--batch-size N` | Chunks per embed call (default 64) |
+
+The script processes sources independently: one failing source is reported and skipped rather than aborting the run, and the exit code is non-zero if any failed. For `pgvector` it updates rows in place, so an interrupted run simply re-does its last batch. For `faiss` it builds the replacement index in memory and only writes once the whole rebuild succeeds, leaving the existing index untouched on failure.
+
+<Callout type="info">
+  Staying on `all-mpnet-base-v2` is a supported choice — it remains in the model registry and in `setup.sh`. You only need this section if you want to move to granite.
+</Callout>
+
 ## Check your version
 
 ```bash

setup.sh

@@ -250,17 +250,18 @@ configure_vector_store() {
 configure_embeddings() {
     echo -e "\n${DEFAULT_FG}${BOLD}Embeddings Configuration${NC}"
     echo -e "${DEFAULT_FG}Choose your embeddings provider:${NC}"
-    echo -e "${YELLOW}1) HuggingFace (default, local)${NC}"
+    echo -e "${YELLOW}1) Granite multilingual (default, local)${NC}"
     echo -e "${YELLOW}2) OpenAI Embeddings${NC}"
     echo -e "${YELLOW}3) Custom Remote Embeddings (OpenAI-compatible API)${NC}"
+    echo -e "${YELLOW}4) all-mpnet-base-v2 (legacy local, English-only)${NC}"
     echo -e "${YELLOW}b) Back${NC}"
     echo
-    read -p "$(echo -e "${DEFAULT_FG}Choose option (1-3, or b): ${NC}")" emb_choice
+    read -p "$(echo -e "${DEFAULT_FG}Choose option (1-4, or b): ${NC}")" emb_choice
 
     case "$emb_choice" in
         1)
-            echo "EMBEDDINGS_NAME=huggingface_sentence-transformers/all-mpnet-base-v2" >> "$ENV_FILE"
-            echo -e "${GREEN}Embeddings set to HuggingFace (local).${NC}"
+            echo "EMBEDDINGS_NAME=ibm-granite/granite-embedding-311m-multilingual-r2" >> "$ENV_FILE"
+            echo -e "${GREEN}Embeddings set to granite-311m-multilingual-r2 (local).${NC}"
             ;;
         2)
             echo "EMBEDDINGS_NAME=openai_text-embedding-ada-002" >> "$ENV_FILE"
@@ -277,6 +278,10 @@ configure_embeddings() {
             [ -n "$emb_key" ] && echo "EMBEDDINGS_KEY=$emb_key" >> "$ENV_FILE"
             echo -e "${GREEN}Custom remote embeddings configured.${NC}"
             ;;
+        4)
+            echo "EMBEDDINGS_NAME=huggingface_sentence-transformers/all-mpnet-base-v2" >> "$ENV_FILE"
+            echo -e "${GREEN}Embeddings set to all-mpnet-base-v2 (legacy local).${NC}"
+            ;;
         b|B) return ;;
         *) echo -e "\n${RED}Invalid choice.${NC}" ; sleep 1 ;;
     esac
@@ -515,7 +520,7 @@ serve_local_ollama() {
     echo "LLM_NAME=$model_name" >> "$ENV_FILE"
     echo "VITE_API_STREAMING=true" >> "$ENV_FILE"
     echo "OPENAI_BASE_URL=http://ollama:11434/v1" >> "$ENV_FILE"
-    echo "EMBEDDINGS_NAME=huggingface_sentence-transformers/all-mpnet-base-v2" >> "$ENV_FILE"
+    echo "EMBEDDINGS_NAME=ibm-granite/granite-embedding-311m-multilingual-r2" >> "$ENV_FILE"
     echo -e "${GREEN}.env file configured for Ollama ($(echo "$docker_compose_file_suffix" | tr '[:lower:]' '[:upper:]')${NC}${GREEN}).${NC}"
 
     prompt_advanced_settings
@@ -628,7 +633,7 @@ connect_local_inference_engine() {
     echo "LLM_NAME=$model_name" >> "$ENV_FILE"
     echo "VITE_API_STREAMING=true" >> "$ENV_FILE"
     echo "OPENAI_BASE_URL=$openai_base_url" >> "$ENV_FILE"
-    echo "EMBEDDINGS_NAME=huggingface_sentence-transformers/all-mpnet-base-v2" >> "$ENV_FILE"
+    echo "EMBEDDINGS_NAME=ibm-granite/granite-embedding-311m-multilingual-r2" >> "$ENV_FILE"
     echo -e "${GREEN}.env file configured for ${BOLD}${engine_name}${NC}${GREEN} with OpenAI API format.${NC}"
     echo -e "${YELLOW}Note: MODEL_NAME is set to '${BOLD}$model_name${NC}${YELLOW}'. You can change it later in the .env file.${NC}"
 

tests/parser/conftest.py

@@ -0,0 +1,23 @@
+import pytest
+
+
+@pytest.fixture(scope="session")
+def _mpnet_tokenizer():
+    """Real WordPiece tokenizer; skipped when the hub is unreachable."""
+    pytest.importorskip("tokenizers")
+    from tokenizers import Tokenizer
+
+    try:
+        tokenizer = Tokenizer.from_pretrained("sentence-transformers/all-mpnet-base-v2")
+    except Exception as exc:  # offline CI
+        pytest.skip(f"tokenizer unavailable: {exc}")
+    tokenizer.no_padding()
+    tokenizer.no_truncation()
+    return tokenizer
+
+
+@pytest.fixture
+def hf_counter(_mpnet_tokenizer):
+    from application.parser.tokenization import HuggingFaceCounter
+
+    return HuggingFaceCounter(_mpnet_tokenizer, "sentence-transformers/all-mpnet-base-v2")

tests/parser/test_chunking.py

@@ -108,18 +108,24 @@ class TestSplitDocument:
 
         result = chunker.split_document(doc)
         assert len(result) > 1
-        # First chunk should contain header
         assert "h1" in result[0].text
+        # Only the first: this is what duplicate_headers=False means.
+        assert all("h1" not in chunk.text for chunk in result[1:])
 
     def test_split_duplicates_header(self):
+        """Every chunk carries the header when the flag is set.
+
+        Asserting only the first chunk passed even while the flag did nothing:
+        the implementation cleared the header after the first iteration, so
+        duplicate_headers was unreachable.
+        """
         chunker = Chunker(max_tokens=50, min_tokens=5, duplicate_headers=True)
         text = "h1\nh2\nh3\n" + "word " * 200
         doc = Document(text=text, doc_id="doc1")
 
         result = chunker.split_document(doc)
         assert len(result) > 1
-        # First chunk should contain header
-        assert "h1" in result[0].text
+        assert all("h1" in chunk.text for chunk in result)
 
     def test_split_preserves_embedding(self):
         chunker = Chunker(max_tokens=50, min_tokens=5)

tests/parser/test_chunking_strategies.py

@@ -15,11 +15,13 @@ from application.parser.chunking_strategies import (
     SemanticChunker,
 )
 from application.parser.schema.base import Document
-from application.utils import get_encoding
+from application.parser.tokenization import get_token_counter
 
 
 def _tok(text: str) -> int:
-    return len(get_encoding().encode(text))
+    # Chunk budgets are enforced in the embedding model's tokenizer, so a
+    # test that measures them in cl100k measures the wrong thing.
+    return get_token_counter().count(text)
 
 
 @pytest.mark.unit

tests/parser/test_tokenization.py

@@ -0,0 +1,107 @@
+"""Chunk sizes must be counted in the embedding model's units, and splitting
+must never rewrite the text it splits."""
+
+import pytest
+
+from application.parser import tokenization
+from application.parser.tokenization import (
+    HuggingFaceCounter,
+    TiktokenCounter,
+    get_token_counter,
+)
+
+SAMPLES = [
+    "Hello World: DocsGPT ANSWERS Questions.",
+    "The quick brown fox jumps over the lazy dog. " * 40,
+    "Comment configurer l'authentification avec une clé API ?",
+    "def embed(text: str) -> list[float]:\n    return model.encode(text)\n",
+    "Ünïcödé — em-dashes, curly “quotes”, and 日本語 text.",
+    "a,b,c\n1,2,3\n4,5,6\n" * 30,
+]
+
+
+class _StubEncoding:
+    """Whitespace tokenizer standing in for tiktoken."""
+
+    def encode_ordinary(self, text):
+        return [ord(c) for c in text]
+
+    def decode(self, ids):
+        return "".join(chr(i) for i in ids)
+
+
+@pytest.fixture(autouse=True)
+def _clear_cache():
+    tokenization.reset_cache()
+    yield
+    tokenization.reset_cache()
+
+
+class TestSplittingPreservesText:
+    """The property that protects every stored document."""
+
+    @pytest.mark.parametrize("text", SAMPLES)
+    def test_tiktoken_split_reassembles_exactly(self, text, monkeypatch):
+        monkeypatch.setattr(tokenization, "get_encoding", lambda: _StubEncoding())
+        counter = TiktokenCounter()
+        pieces = counter.split(text, 7)
+        assert "".join(pieces) == text
+
+    @pytest.mark.parametrize("text", SAMPLES)
+    def test_hf_split_reassembles_exactly(self, text, hf_counter):
+        """WordPiece lowercases on decode, so splitting must slice, not decode."""
+        pieces = hf_counter.split(text, 7)
+        assert "".join(pieces) == text
+
+    def test_hf_split_does_not_lowercase(self, hf_counter):
+        text = "Hello World: DocsGPT ANSWERS Questions."
+        assert "".join(hf_counter.split(text, 3)) == text
+        assert "DocsGPT" in "".join(hf_counter.split(text, 3))
+
+    @pytest.mark.parametrize("text", SAMPLES)
+    def test_every_piece_is_within_budget(self, text, hf_counter):
+        budget = 10
+        for piece in hf_counter.split(text, budget):
+            # The final piece can absorb trailing characters the tokenizer
+            # dropped, so allow a small overshoot there only.
+            assert hf_counter.count(piece) <= budget + 2
+
+    def test_short_text_is_returned_whole(self, hf_counter):
+        assert hf_counter.split("short", 100) == ["short"]
+
+    def test_empty_text_yields_no_pieces(self, hf_counter):
+        assert hf_counter.split("", 10) == []
+
+    def test_zero_budget_is_clamped_not_infinite_loop(self, hf_counter):
+        pieces = hf_counter.split("some words here to split", 0)
+        assert "".join(pieces) == "some words here to split"
+
+
+class TestCounting:
+    def test_counts_differ_between_tokenizers(self, hf_counter, monkeypatch):
+        """The whole point: mpnet and cl100k disagree, so units matter."""
+        monkeypatch.setattr(tokenization, "get_encoding", lambda: _StubEncoding())
+        text = "internationalisation tokenization"
+        assert hf_counter.count(text) != TiktokenCounter().count(text)
+
+    def test_empty_text_counts_zero(self, hf_counter):
+        assert hf_counter.count("") == 0
+
+
+class TestSelection:
+    def test_registered_model_uses_its_own_tokenizer(self):
+        counter = get_token_counter("huggingface_sentence-transformers/all-mpnet-base-v2")
+        assert isinstance(counter, HuggingFaceCounter)
+        assert counter.name == "sentence-transformers/all-mpnet-base-v2"
+
+    def test_openai_model_falls_back_to_cl100k(self):
+        """OpenAI models are served remotely and genuinely count cl100k."""
+        assert isinstance(get_token_counter("openai_text-embedding-ada-002"), TiktokenCounter)
+
+    def test_unreachable_tokenizer_falls_back_rather_than_raising(self, monkeypatch):
+        monkeypatch.setattr(tokenization, "_load_hf_counter", lambda repo: None)
+        assert isinstance(get_token_counter("granite-311m"), TiktokenCounter)
+
+    def test_counter_is_cached_per_model(self):
+        first = get_token_counter("granite-311m")
+        assert get_token_counter("granite-311m") is first

tests/scripts/test_reembed.py

@@ -0,0 +1,124 @@
+"""Re-embed script: CLI contract, orchestration, and the FAISS rebuild path."""
+
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from application.scripts import reembed
+
+
+class TestCLI:
+    def test_defaults(self):
+        args = reembed.build_parser().parse_args([])
+        assert args.dry_run is False
+        assert args.sources is None
+        assert args.batch_size == reembed.DEFAULT_BATCH_SIZE
+
+    def test_unsupported_store_exits_without_touching_anything(self):
+        with patch.object(reembed.settings, "VECTOR_STORE", "qdrant", create=True):
+            with patch.object(reembed, "run") as run:
+                assert reembed.main([]) == 2
+        run.assert_not_called()
+
+    def test_supported_stores_are_pgvector_and_faiss(self):
+        assert set(reembed.SUPPORTED_STORES) == {"pgvector", "faiss"}
+
+    def test_sources_are_split_and_trimmed(self):
+        with patch.object(reembed.settings, "VECTOR_STORE", "faiss", create=True):
+            with patch.object(reembed, "run", return_value=0) as run:
+                reembed.main(["--sources", " a , b ,, c "])
+        assert run.call_args.args[1] == ["a", "b", "c"]
+
+    def test_batch_size_is_clamped_to_at_least_one(self):
+        with patch.object(reembed.settings, "VECTOR_STORE", "faiss", create=True):
+            with patch.object(reembed, "run", return_value=0) as run:
+                reembed.main(["--batch-size", "0"])
+        assert run.call_args.args[2] == 1
+
+
+class TestRun:
+    def test_no_sources_is_a_clean_exit(self):
+        with patch.object(reembed, "list_source_ids", return_value=[]):
+            assert reembed.run("faiss", None, 8, False) == 0
+
+    def test_processes_every_discovered_source(self):
+        with patch.object(reembed, "list_source_ids", return_value=["a", "b"]):
+            with patch.object(reembed, "reembed_faiss", return_value=(3, 3)) as handler:
+                assert reembed.run("faiss", None, 8, False) == 0
+        assert [call.args[0] for call in handler.call_args_list] == ["a", "b"]
+
+    def test_explicit_sources_skip_discovery(self):
+        with patch.object(reembed, "list_source_ids") as discover:
+            with patch.object(reembed, "reembed_faiss", return_value=(1, 1)):
+                reembed.run("faiss", ["only-this"], 8, False)
+        discover.assert_not_called()
+
+    def test_one_failing_source_does_not_stop_the_others(self):
+        def handler(source_id, batch_size, dry_run):
+            if source_id == "bad":
+                raise RuntimeError("boom")
+            return (1, 1)
+
+        with patch.object(reembed, "reembed_faiss", side_effect=handler) as spy:
+            code = reembed.run("faiss", ["good", "bad", "also-good"], 8, False)
+        assert code == 1, "a failure must be reported in the exit code"
+        assert spy.call_count == 3, "later sources must still be attempted"
+
+    def test_dry_run_is_reported_as_success(self):
+        with patch.object(reembed, "reembed_faiss", return_value=(5, 0)):
+            assert reembed.run("faiss", ["a"], 8, True) == 0
+
+    def test_pgvector_uses_the_pgvector_handler(self):
+        with patch.object(reembed, "reembed_pgvector", return_value=(1, 1)) as handler:
+            reembed.run("pgvector", ["a"], 8, False)
+        handler.assert_called_once()
+
+
+class TestFaissRebuild:
+    @pytest.fixture
+    def stores(self):
+        """An existing store to read from and the rebuilt one written back."""
+        existing = MagicMock()
+        existing.get_chunks.return_value = [
+            {"doc_id": "1", "text": "alpha", "metadata": {"i": 0}},
+            {"doc_id": "2", "text": "beta", "metadata": {"i": 1}},
+        ]
+        rebuilt = MagicMock()
+        with patch.object(
+            reembed.VectorCreator, "create_vectorstore", side_effect=[existing, rebuilt]
+        ) as factory:
+            yield existing, rebuilt, factory
+
+    def test_rebuilds_from_stored_text_and_saves(self, stores):
+        existing, rebuilt, factory = stores
+        seen, written = reembed.reembed_faiss("s1", batch_size=8, dry_run=False)
+
+        assert (seen, written) == (2, 2)
+        rebuilt.save_local.assert_called_once()
+        docs = factory.call_args_list[1].kwargs["docs_init"]
+        assert [d.page_content for d in docs] == ["alpha", "beta"]
+        assert [d.metadata for d in docs] == [{"i": 0}, {"i": 1}]
+
+    def test_existing_index_is_not_deleted(self, stores):
+        """The rebuild must not destroy the old index before the new one exists."""
+        existing, rebuilt, _ = stores
+        reembed.reembed_faiss("s1", batch_size=8, dry_run=False)
+        existing.delete_index.assert_not_called()
+
+    def test_dry_run_reads_but_never_rebuilds(self):
+        existing = MagicMock()
+        existing.get_chunks.return_value = [{"doc_id": "1", "text": "a", "metadata": {}}]
+        with patch.object(
+            reembed.VectorCreator, "create_vectorstore", return_value=existing
+        ) as factory:
+            seen, written = reembed.reembed_faiss("s1", batch_size=8, dry_run=True)
+        assert (seen, written) == (1, 0)
+        assert factory.call_count == 1, "no rebuild store may be constructed"
+
+    def test_empty_index_is_a_no_op(self):
+        existing = MagicMock()
+        existing.get_chunks.return_value = []
+        with patch.object(
+            reembed.VectorCreator, "create_vectorstore", return_value=existing
+        ):
+            assert reembed.reembed_faiss("s1", batch_size=8, dry_run=False) == (0, 0)

tests/scripts/test_reembed_pgvector_live.py

@@ -0,0 +1,187 @@
+"""Live pgvector run of the re-embed script.
+
+Uses the ephemeral pytest-postgresql cluster with a stub embeddings model, so
+a real ``UPDATE ... ::vector`` round trip is exercised without downloading a
+model. Skips when the cluster has no pgvector build.
+"""
+
+from __future__ import annotations
+
+from unittest.mock import patch
+
+import pytest
+
+from application.scripts import reembed
+from application.vectorstore import pgvector as pgvector_module
+from application.vectorstore.pgvector import PGVectorStore
+
+pytestmark = pytest.mark.integration
+
+DIM = 8
+
+
+class _Embeddings:
+    """Returns a distinct constant per generation, so a rewrite is visible."""
+
+    dimension = DIM
+
+    def __init__(self, seed: float):
+        self.seed = seed
+        self.calls = 0
+
+    def embed_documents(self, texts):
+        self.calls += 1
+        return [[self.seed] + [0.0] * (DIM - 1) for _ in texts]
+
+    def embed_query(self, query):
+        return [self.seed] + [0.0] * (DIM - 1)
+
+
+def _dsn(info) -> str:
+    password = f":{info.password}" if info.password else ""
+    return f"postgresql://{info.user}{password}@{info.host}:{info.port}/{info.dbname}"
+
+
+@pytest.fixture(autouse=True)
+def _close_pools():
+    yield
+    for dsn, pool in list(pgvector_module._POOLS.items()):
+        try:
+            pool.close()
+        except Exception:
+            pass
+        pgvector_module._POOLS.pop(dsn, None)
+
+
+@pytest.fixture
+def live_dsn(postgresql, monkeypatch):
+    try:
+        with postgresql.cursor() as cursor:
+            cursor.execute("CREATE EXTENSION vector;")
+        postgresql.rollback()
+    except Exception as exc:
+        postgresql.rollback()
+        pytest.skip(f"pgvector extension unavailable: {exc}")
+
+    dsn = _dsn(postgresql.info)
+    from application.core import settings as settings_module
+
+    settings = settings_module.settings
+    monkeypatch.setattr(settings, "VECTOR_STORE", "pgvector", raising=False)
+    monkeypatch.setattr(settings, "PGVECTOR_CONNECTION_STRING", dsn, raising=False)
+    monkeypatch.setattr(settings, "GRAPHRAG_ENABLED", False, raising=False)
+    monkeypatch.setattr(settings, "PGVECTOR_IVFFLAT_PROBES", None, raising=False)
+    monkeypatch.setattr(settings, "PGVECTOR_POOL_MAX_SIZE", 4, raising=False)
+    monkeypatch.setattr(settings, "EMBEDDINGS_NAME", "granite-311m", raising=False)
+    return dsn
+
+
+def _seed(dsn, source_id, texts, embeddings):
+    """Create the schema and insert ``texts`` embedded by ``embeddings``."""
+    with patch(
+        "application.vectorstore.base.BaseVectorStore._get_embeddings",
+        return_value=embeddings,
+    ):
+        store = PGVectorStore(source_id=source_id, connection_string=dsn)
+        conn = store._get_connection()
+        PGVectorStore.create_schema(conn, dimension=DIM)
+        conn.commit()
+        store.add_texts(list(texts), metadatas=[{"i": i} for i in range(len(texts))])
+        store.close()
+
+
+def _as_list(vector):
+    """Normalise a stored vector to a list of floats.
+
+    Depending on whether pgvector's adapter is registered on the reading
+    connection, the value comes back as a ``Vector`` or as its text form
+    ``'[1,0,...]'``. Both are valid; the assertions should not care.
+    """
+    if vector is None:
+        return []
+    if hasattr(vector, "to_list"):
+        return list(vector.to_list())
+    if isinstance(vector, str):
+        return [float(part) for part in vector.strip("[]").split(",") if part]
+    return list(vector)
+
+
+def _vectors(dsn, source_id):
+    store = PGVectorStore(source_id=source_id, connection_string=dsn)
+    conn = store._get_connection()
+    cursor = conn.cursor()
+    try:
+        cursor.execute(
+            "SELECT text, embedding FROM documents WHERE source_id = %s ORDER BY id",
+            (source_id,),
+        )
+        # pgvector hands back a ``Vector``; normalise to a plain list so the
+        # assertions read the same whichever adapter is registered.
+        return [(text, _as_list(vector)) for text, vector in cursor.fetchall()]
+    finally:
+        cursor.close()
+        store.close()
+
+
+TEXTS = ["alpha document", "beta document", "gamma document"]
+
+
+class TestReembedPgvectorLive:
+    def test_rewrites_vectors_and_preserves_text(self, live_dsn):
+        _seed(live_dsn, "src-a", TEXTS, _Embeddings(1.0))
+        before = _vectors(live_dsn, "src-a")
+        assert [row[0] for row in before] == TEXTS
+        assert all(row[1][0] == pytest.approx(1.0) for row in before)
+
+        new_model = _Embeddings(9.0)
+        with patch(
+            "application.vectorstore.base.BaseVectorStore._get_embeddings",
+            return_value=new_model,
+        ):
+            seen, written = reembed.reembed_pgvector("src-a", batch_size=2, dry_run=False)
+
+        assert (seen, written) == (3, 3)
+        after = _vectors(live_dsn, "src-a")
+        # Text is untouched; only the vectors moved.
+        assert [row[0] for row in after] == TEXTS
+        assert all(row[1][0] == pytest.approx(9.0) for row in after)
+
+    def test_dry_run_counts_without_writing(self, live_dsn):
+        _seed(live_dsn, "src-b", TEXTS, _Embeddings(1.0))
+        new_model = _Embeddings(9.0)
+        with patch(
+            "application.vectorstore.base.BaseVectorStore._get_embeddings",
+            return_value=new_model,
+        ):
+            seen, written = reembed.reembed_pgvector("src-b", batch_size=2, dry_run=True)
+
+        assert (seen, written) == (3, 0)
+        assert new_model.calls == 0, "dry run must not embed"
+        assert all(row[1][0] == pytest.approx(1.0) for row in _vectors(live_dsn, "src-b"))
+
+    def test_batches_are_respected(self, live_dsn):
+        _seed(live_dsn, "src-c", TEXTS, _Embeddings(1.0))
+        new_model = _Embeddings(9.0)
+        with patch(
+            "application.vectorstore.base.BaseVectorStore._get_embeddings",
+            return_value=new_model,
+        ):
+            reembed.reembed_pgvector("src-c", batch_size=2, dry_run=False)
+        assert new_model.calls == 2, "3 chunks at batch 2 is two embed calls"
+
+    def test_only_the_named_source_is_touched(self, live_dsn):
+        _seed(live_dsn, "src-d", TEXTS, _Embeddings(1.0))
+        _seed(live_dsn, "src-e", TEXTS, _Embeddings(1.0))
+        with patch(
+            "application.vectorstore.base.BaseVectorStore._get_embeddings",
+            return_value=_Embeddings(9.0),
+        ):
+            reembed.reembed_pgvector("src-d", batch_size=64, dry_run=False)
+
+        assert all(row[1][0] == pytest.approx(9.0) for row in _vectors(live_dsn, "src-d"))
+        assert all(row[1][0] == pytest.approx(1.0) for row in _vectors(live_dsn, "src-e"))
+
+    def test_source_discovery_lists_every_source(self, live_dsn):
+        _seed(live_dsn, "src-f", TEXTS, _Embeddings(1.0))
+        _seed(live_dsn, "src-g", TEXTS, _Embeddings(1.0))
+        assert reembed.list_source_ids("pgvector") == ["src-f", "src-g"]

tests/vectorstore/test_base.py

@@ -263,11 +263,11 @@ class TestEmbeddingsSingleton:
     def test_get_instance_hf_ignores_positional_key(
         self, mock_get_wrapper, mock_settings
     ):
-        """A stray key must not reach the zero-arg HuggingFace factory.
+        """A stray key must not reach the wrapper for a registered model.
 
-        The factories are ``lambda: EmbeddingsWrapper(...)``, so a caller that
-        passed ``settings.EMBEDDINGS_KEY`` positionally used to blow up with
-        ``TypeError: <lambda>() takes 0 positional arguments``.
+        Registered models take their whole configuration from the registry, so
+        a caller that passes ``settings.EMBEDDINGS_KEY`` positionally (as the
+        vector stores do) must have it dropped rather than forwarded.
         """
         mock_settings.EMBEDDINGS_BASE_URL = None
         mock_wrapper_cls = Mock()
@@ -278,9 +278,8 @@ class TestEmbeddingsSingleton:
         result = EmbeddingsSingleton.get_instance(HF_MPNET, None)
 
         assert result is mock_instance
-        mock_wrapper_cls.assert_called_once_with(
-            "sentence-transformers/all-mpnet-base-v2"
-        )
+        # The configured name is passed through; the registry maps it to a repo.
+        mock_wrapper_cls.assert_called_once_with(HF_MPNET)
 
     @patch("application.vectorstore.base.settings")
     @patch("application.vectorstore.base._get_embeddings_wrapper")
@@ -293,9 +292,7 @@ class TestEmbeddingsSingleton:
 
         EmbeddingsSingleton.get_instance(HF_MPNET, openai_api_key="sk-nope")
 
-        mock_wrapper_cls.assert_called_once_with(
-            "sentence-transformers/all-mpnet-base-v2"
-        )
+        mock_wrapper_cls.assert_called_once_with(HF_MPNET)
 
 
 # --- BaseVectorStore ---
@@ -401,12 +398,15 @@ class TestBaseVectorStore:
 
     @patch("application.vectorstore.base.settings")
     @patch("application.vectorstore.base.EmbeddingsSingleton.get_instance")
-    @patch("os.path.exists")
-    def test_get_embeddings_huggingface_local_model(
-        self, mock_exists, mock_get_instance, mock_settings
+    def test_get_embeddings_registered_model_passes_configured_name(
+        self, mock_get_instance, mock_settings
     ):
+        """No bundled-path branch any more: the name goes straight through.
+
+        FastEmbed resolves artifacts through its own cache (warmed in the
+        image), so the old ``/app/models/...`` probe has no job to do.
+        """
         mock_settings.EMBEDDINGS_BASE_URL = None
-        mock_exists.side_effect = lambda p: p == "/app/models/all-mpnet-base-v2"
         mock_emb = Mock()
         mock_get_instance.return_value = mock_emb
 
@@ -415,7 +415,9 @@ class TestBaseVectorStore:
             "huggingface_sentence-transformers/all-mpnet-base-v2"
         )
         assert result is mock_emb
-        mock_get_instance.assert_called_with("/app/models/all-mpnet-base-v2")
+        mock_get_instance.assert_called_with(
+            "huggingface_sentence-transformers/all-mpnet-base-v2"
+        )
 
     @patch("application.vectorstore.base.settings")
     @patch("application.vectorstore.base.EmbeddingsSingleton.get_instance")
@@ -518,18 +520,18 @@ class TestGetEmbeddingsResolver:
 
     @patch("application.vectorstore.base.settings")
     @patch("application.vectorstore.base._get_embeddings_wrapper")
-    @patch("os.path.exists")
-    def test_uses_local_model_path_and_caches_it(
-        self, mock_exists, mock_get_wrapper, mock_settings
+    def test_repeated_resolution_loads_one_model(
+        self, mock_get_wrapper, mock_settings
     ):
-        """With the bundled model present, the instance is keyed by its path.
+        """A second call must not load a second copy of the model.
 
-        A second call must not load a second copy of the model.
+        The instance is keyed by the configured name. It used to be keyed by a
+        bundled filesystem path when one happened to exist, which meant the
+        same model could be cached twice under two keys.
         """
         mock_settings.EMBEDDINGS_BASE_URL = None
         mock_settings.EMBEDDINGS_NAME = HF_MPNET
         mock_settings.EMBEDDINGS_KEY = None
-        mock_exists.side_effect = lambda path: path == LOCAL_MPNET
         mock_wrapper_cls = Mock()
         mock_wrapper_cls.return_value = Mock()
         mock_get_wrapper.return_value = mock_wrapper_cls
@@ -538,8 +540,8 @@ class TestGetEmbeddingsResolver:
         second = get_embeddings()
 
         assert first is second
-        assert set(EmbeddingsSingleton._instances) == {LOCAL_MPNET}
-        mock_wrapper_cls.assert_called_once_with(LOCAL_MPNET)
+        assert set(EmbeddingsSingleton._instances) == {HF_MPNET}
+        mock_wrapper_cls.assert_called_once_with(HF_MPNET)
 
     @patch("application.vectorstore.base.settings")
     def test_remote_when_base_url_configured(self, mock_settings):

tests/vectorstore/test_embeddings_local.py

@@ -1,156 +1,167 @@
-from unittest.mock import MagicMock, Mock, patch
+"""Local embeddings run through FastEmbed, configured from the model registry."""
 
+from unittest.mock import MagicMock, patch
+
+import numpy as np
 import pytest
 
+from application.vectorstore import embeddings_local
+from application.vectorstore.embeddings_local import EmbeddingsWrapper
+from application.vectorstore.model_registry import GRANITE_97M, MPNET
 
-@pytest.mark.unit
-class TestEmbeddingsWrapper:
-    @patch("application.vectorstore.embeddings_local.SentenceTransformer")
-    def test_init_success(self, mock_st_cls):
-        mock_model = MagicMock()
-        mock_model._first_module.return_value = MagicMock()
-        mock_model.get_embedding_dimension.return_value = 768
-        mock_st_cls.return_value = mock_model
 
-        from application.vectorstore.embeddings_local import EmbeddingsWrapper
+@pytest.fixture(autouse=True)
+def _clear_registration():
+    """``add_custom_model`` writes to a FastEmbed global; keep tests isolated."""
+    embeddings_local._registered.clear()
+    yield
+    embeddings_local._registered.clear()
 
-        wrapper = EmbeddingsWrapper("test-model")
 
-        mock_st_cls.assert_called_once()
-        assert wrapper.dimension == 768
+@pytest.fixture
+def fake_fastembed():
+    """Patch FastEmbed so no model is downloaded or run."""
+    text_embedding = MagicMock()
+    instance = MagicMock()
+    instance.embed.return_value = iter([np.array([0.1, 0.2, 0.3])])
+    text_embedding.return_value = instance
+    with patch("fastembed.TextEmbedding", text_embedding):
+        yield text_embedding, instance
 
-    @patch("application.vectorstore.embeddings_local.SentenceTransformer")
-    def test_init_falls_back_to_legacy_dimension_api(self, mock_st_cls):
-        """sentence-transformers < 5.4 only exposes get_sentence_embedding_dimension."""
-        mock_model = MagicMock()
-        mock_model._first_module.return_value = MagicMock()
-        del mock_model.get_embedding_dimension
-        mock_model.get_sentence_embedding_dimension.return_value = 768
-        mock_st_cls.return_value = mock_model
 
-        from application.vectorstore.embeddings_local import EmbeddingsWrapper
+class TestRegistryDrivenLoading:
+    def test_registered_model_loads_by_repo_not_by_configured_name(self, fake_fastembed):
+        text_embedding, _ = fake_fastembed
+        wrapper = EmbeddingsWrapper(MPNET.name)
+        assert text_embedding.call_args.kwargs["model_name"] == MPNET.repo
+        assert wrapper.dimension == MPNET.dimension
 
-        wrapper = EmbeddingsWrapper("test-model")
+    def test_legacy_alias_resolves_to_the_same_model(self, fake_fastembed):
+        text_embedding, _ = fake_fastembed
+        EmbeddingsWrapper("huggingface_sentence-transformers-all-mpnet-base-v2")
+        assert text_embedding.call_args.kwargs["model_name"] == MPNET.repo
 
-        assert wrapper.dimension == 768
-        mock_model.get_sentence_embedding_dimension.assert_called_once()
+    def test_dimension_comes_from_registry_without_running_the_model(self, fake_fastembed):
+        _, instance = fake_fastembed
+        wrapper = EmbeddingsWrapper(GRANITE_97M.name)
+        assert wrapper.dimension == 384
+        instance.embed.assert_not_called()
 
-    @patch("application.vectorstore.embeddings_local.SentenceTransformer")
-    def test_init_failure(self, mock_st_cls):
-        mock_st_cls.side_effect = Exception("model not found")
+    def test_unknown_model_is_treated_as_a_hf_repo(self, fake_fastembed):
+        text_embedding, _ = fake_fastembed
+        wrapper = EmbeddingsWrapper("some-org/custom-embedder")
+        assert text_embedding.call_args.kwargs["model_name"] == "some-org/custom-embedder"
+        # No registry entry means no known width, so it must be probed.
+        assert wrapper.dimension == 3
 
-        from application.vectorstore.embeddings_local import EmbeddingsWrapper
+    def test_load_failure_names_the_model_and_the_known_ones(self):
+        with patch("fastembed.TextEmbedding", side_effect=OSError("no such repo")):
+            with pytest.raises(RuntimeError) as excinfo:
+                EmbeddingsWrapper("broken/model")
+        message = str(excinfo.value)
+        assert "broken/model" in message
+        assert MPNET.name in message
 
-        with pytest.raises(Exception, match="model not found"):
-            EmbeddingsWrapper("bad-model")
 
-    @patch("application.vectorstore.embeddings_local.SentenceTransformer")
-    def test_init_none_model(self, mock_st_cls):
-        mock_st_cls.return_value = None
+class TestSettingsPassthrough:
+    def test_threads_forwarded_when_configured(self, fake_fastembed):
+        text_embedding, _ = fake_fastembed
+        with patch.object(embeddings_local.settings, "EMBEDDINGS_THREADS", 2, create=True):
+            EmbeddingsWrapper(MPNET.name)
+        assert text_embedding.call_args.kwargs["threads"] == 2
 
-        from application.vectorstore.embeddings_local import EmbeddingsWrapper
+    def test_threads_omitted_when_unset(self, fake_fastembed):
+        text_embedding, _ = fake_fastembed
+        with patch.object(embeddings_local.settings, "EMBEDDINGS_THREADS", None, create=True):
+            EmbeddingsWrapper(MPNET.name)
+        assert "threads" not in text_embedding.call_args.kwargs
 
-        with pytest.raises((ValueError, AttributeError)):
-            EmbeddingsWrapper("bad-model")
+    def test_cache_dir_forwarded_when_configured(self, fake_fastembed):
+        text_embedding, _ = fake_fastembed
+        with patch.object(embeddings_local.settings, "EMBEDDINGS_CACHE_DIR", "/models", create=True):
+            EmbeddingsWrapper(MPNET.name)
+        assert text_embedding.call_args.kwargs["cache_dir"] == "/models"
 
-    @patch("application.vectorstore.embeddings_local.SentenceTransformer")
-    def test_init_null_first_module(self, mock_st_cls):
-        mock_model = MagicMock()
-        mock_model._first_module.return_value = None
-        mock_st_cls.return_value = mock_model
 
-        from application.vectorstore.embeddings_local import EmbeddingsWrapper
+class TestEmbedding:
+    def test_embed_documents_returns_plain_lists(self, fake_fastembed):
+        _, instance = fake_fastembed
+        instance.embed.return_value = iter([np.array([1.0, 2.0]), np.array([3.0, 4.0])])
+        wrapper = EmbeddingsWrapper(MPNET.name)
+        assert wrapper.embed_documents(["a", "b"]) == [[1.0, 2.0], [3.0, 4.0]]
 
-        with pytest.raises(ValueError, match="failed to load properly"):
-            EmbeddingsWrapper("bad-model")
+    def test_embed_documents_short_circuits_on_empty_input(self, fake_fastembed):
+        _, instance = fake_fastembed
+        wrapper = EmbeddingsWrapper(MPNET.name)
+        instance.embed.reset_mock()
+        assert wrapper.embed_documents([]) == []
+        instance.embed.assert_not_called()
 
-    @patch("application.vectorstore.embeddings_local.SentenceTransformer")
-    def test_embed_query(self, mock_st_cls):
-        mock_model = MagicMock()
-        mock_model._first_module.return_value = MagicMock()
-        mock_model.get_sentence_embedding_dimension.return_value = 3
-        mock_model.encode.return_value = MagicMock(tolist=Mock(return_value=[0.1, 0.2, 0.3]))
-        mock_st_cls.return_value = mock_model
+    def test_embed_query_returns_a_single_vector(self, fake_fastembed):
+        _, instance = fake_fastembed
+        instance.embed.return_value = iter([np.array([0.5, 0.6])])
+        wrapper = EmbeddingsWrapper(MPNET.name)
+        assert wrapper.embed_query("hello") == [0.5, 0.6]
 
-        from application.vectorstore.embeddings_local import EmbeddingsWrapper
+    def test_call_dispatches_on_input_type(self, fake_fastembed):
+        _, instance = fake_fastembed
+        wrapper = EmbeddingsWrapper(MPNET.name)
+        instance.embed.return_value = iter([np.array([1.0])])
+        assert wrapper("text") == [1.0]
+        instance.embed.return_value = iter([np.array([1.0]), np.array([2.0])])
+        assert wrapper(["a", "b"]) == [[1.0], [2.0]]
 
-        wrapper = EmbeddingsWrapper("model")
-        result = wrapper.embed_query("hello world")
+    def test_call_rejects_other_types(self, fake_fastembed):
+        wrapper = EmbeddingsWrapper(MPNET.name)
+        with pytest.raises(ValueError):
+            wrapper(42)
 
-        mock_model.encode.assert_called_once_with("hello world")
-        assert result == [0.1, 0.2, 0.3]
 
-    @patch("application.vectorstore.embeddings_local.SentenceTransformer")
-    def test_embed_documents(self, mock_st_cls):
-        mock_model = MagicMock()
-        mock_model._first_module.return_value = MagicMock()
-        mock_model.get_sentence_embedding_dimension.return_value = 3
-        mock_model.encode.return_value = MagicMock(
-            tolist=Mock(return_value=[[0.1, 0.2], [0.3, 0.4]])
+class TestRegistrationIsIdempotent:
+    def test_model_registered_once_per_process(self, fake_fastembed):
+        text_embedding, _ = fake_fastembed
+        EmbeddingsWrapper(MPNET.name)
+        EmbeddingsWrapper(MPNET.name)
+        assert text_embedding.add_custom_model.call_count == 1
+
+
+class TestLengthSortedBatching:
+    """Grouping by length is a throughput/memory win, but order is a contract."""
+
+    def _wrapper(self, fake_fastembed, batch_size):
+        _, instance = fake_fastembed
+        wrapper = EmbeddingsWrapper(MPNET.name)
+        instance.embed.side_effect = lambda texts, batch_size=None: iter(
+            [np.array([float(len(t))]) for t in texts]
         )
-        mock_st_cls.return_value = mock_model
+        return wrapper, instance
 
-        from application.vectorstore.embeddings_local import EmbeddingsWrapper
+    def test_output_order_matches_input_order(self, fake_fastembed):
+        with patch.object(embeddings_local.settings, "EMBEDDINGS_BATCH_SIZE", 2, create=True):
+            wrapper, _ = self._wrapper(fake_fastembed, 2)
+            texts = ["dddd", "a", "ccc", "bb", "eeeee"]
+            out = wrapper.embed_documents(texts)
+        # Each stub vector encodes its own text length, so a reordered result
+        # is immediately visible.
+        assert out == [[4.0], [1.0], [3.0], [2.0], [5.0]]
 
-        wrapper = EmbeddingsWrapper("model")
-        result = wrapper.embed_documents(["doc1", "doc2"])
+    def test_inputs_are_grouped_by_length_before_batching(self, fake_fastembed):
+        with patch.object(embeddings_local.settings, "EMBEDDINGS_BATCH_SIZE", 2, create=True):
+            wrapper, instance = self._wrapper(fake_fastembed, 2)
+            wrapper.embed_documents(["dddd", "a", "ccc", "bb", "eeeee"])
+        sent = instance.embed.call_args.args[0]
+        assert [len(t) for t in sent] == [1, 2, 3, 4, 5]
 
-        mock_model.encode.assert_called_with(["doc1", "doc2"])
-        assert result == [[0.1, 0.2], [0.3, 0.4]]
+    def test_single_batch_is_not_reordered(self, fake_fastembed):
+        with patch.object(embeddings_local.settings, "EMBEDDINGS_BATCH_SIZE", 32, create=True):
+            wrapper, instance = self._wrapper(fake_fastembed, 32)
+            texts = ["dddd", "a", "ccc"]
+            out = wrapper.embed_documents(texts)
+        assert instance.embed.call_args.args[0] == texts
+        assert out == [[4.0], [1.0], [3.0]]
 
-    @patch("application.vectorstore.embeddings_local.SentenceTransformer")
-    def test_call_with_string(self, mock_st_cls):
-        mock_model = MagicMock()
-        mock_model._first_module.return_value = MagicMock()
-        mock_model.get_sentence_embedding_dimension.return_value = 3
-        mock_model.encode.return_value = MagicMock(tolist=Mock(return_value=[0.1]))
-        mock_st_cls.return_value = mock_model
-
-        from application.vectorstore.embeddings_local import EmbeddingsWrapper
-
-        wrapper = EmbeddingsWrapper("model")
-        result = wrapper("hello")
-        assert result == [0.1]
-
-    @patch("application.vectorstore.embeddings_local.SentenceTransformer")
-    def test_call_with_list(self, mock_st_cls):
-        mock_model = MagicMock()
-        mock_model._first_module.return_value = MagicMock()
-        mock_model.get_sentence_embedding_dimension.return_value = 3
-        mock_model.encode.return_value = MagicMock(
-            tolist=Mock(return_value=[[0.1], [0.2]])
-        )
-        mock_st_cls.return_value = mock_model
-
-        from application.vectorstore.embeddings_local import EmbeddingsWrapper
-
-        wrapper = EmbeddingsWrapper("model")
-        result = wrapper(["a", "b"])
-        assert result == [[0.1], [0.2]]
-
-    @patch("application.vectorstore.embeddings_local.SentenceTransformer")
-    def test_call_with_invalid_type(self, mock_st_cls):
-        mock_model = MagicMock()
-        mock_model._first_module.return_value = MagicMock()
-        mock_model.get_sentence_embedding_dimension.return_value = 3
-        mock_st_cls.return_value = mock_model
-
-        from application.vectorstore.embeddings_local import EmbeddingsWrapper
-
-        wrapper = EmbeddingsWrapper("model")
-        with pytest.raises(ValueError, match="Input must be a string or a list"):
-            wrapper(123)
-
-    @patch("application.vectorstore.embeddings_local.SentenceTransformer")
-    def test_trust_remote_code_default(self, mock_st_cls):
-        mock_model = MagicMock()
-        mock_model._first_module.return_value = MagicMock()
-        mock_model.get_sentence_embedding_dimension.return_value = 768
-        mock_st_cls.return_value = mock_model
-
-        from application.vectorstore.embeddings_local import EmbeddingsWrapper
-
-        EmbeddingsWrapper("model")
-
-        call_kwargs = mock_st_cls.call_args[1]
-        assert call_kwargs["trust_remote_code"] is True
+    def test_duplicate_texts_are_handled(self, fake_fastembed):
+        with patch.object(embeddings_local.settings, "EMBEDDINGS_BATCH_SIZE", 2, create=True):
+            wrapper, _ = self._wrapper(fake_fastembed, 2)
+            out = wrapper.embed_documents(["aa", "b", "aa", "ccc"])
+        assert out == [[2.0], [1.0], [2.0], [3.0]]

tests/vectorstore/test_faiss.py

@@ -185,15 +185,20 @@ class TestFaissStoreAssertEmbeddingDimensions:
             with pytest.raises(ValueError, match="Embedding dimension mismatch"):
                 populated.assert_embedding_dimensions(Mock(dimension=768))
 
-    def test_missing_dimension_attr_raises(self, populated):
+    def test_unknown_dimension_defers_rather_than_raising(self, populated):
+        """A remote model reports no width until its first call.
+
+        Refusing to open the index in that window would break startup for a
+        perfectly valid remote configuration, so an unknown width is deferred,
+        not treated as a mismatch.
+        """
         with patch("application.vectorstore.faiss.settings") as mock_settings:
             mock_settings.EMBEDDINGS_NAME = (
                 "huggingface_sentence-transformers/all-mpnet-base-v2"
             )
             embeddings = Mock()
-            del embeddings.dimension
-            with pytest.raises(AttributeError, match="'dimension' attribute not found"):
-                populated.assert_embedding_dimensions(embeddings)
+            embeddings.dimension = None
+            assert populated.assert_embedding_dimensions(embeddings) is None
 
     def test_dimension_match_passes(self, populated):
         with patch("application.vectorstore.faiss.settings") as mock_settings:
@@ -202,10 +207,22 @@ class TestFaissStoreAssertEmbeddingDimensions:
             )
             assert populated.assert_embedding_dimensions(Mock(dimension=3)) is None
 
-    def test_non_huggingface_skips_dimension_check(self, populated):
+    def test_mismatch_is_caught_for_every_model_not_just_mpnet(self, populated):
+        """The check used to run only when EMBEDDINGS_NAME was mpnet.
+
+        That skipped exactly the case it exists for: an index built with one
+        model being opened under a different one.
+        """
         with patch("application.vectorstore.faiss.settings") as mock_settings:
             mock_settings.EMBEDDINGS_NAME = "openai_text-embedding-ada-002"
-            assert populated.assert_embedding_dimensions(Mock(dimension=1536)) is None
+            with pytest.raises(ValueError, match="Embedding dimension mismatch"):
+                populated.assert_embedding_dimensions(Mock(dimension=1536))
+
+    def test_mismatch_message_points_at_the_reembed_script(self, populated):
+        with patch("application.vectorstore.faiss.settings") as mock_settings:
+            mock_settings.EMBEDDINGS_NAME = "granite-311m"
+            with pytest.raises(ValueError, match="application.scripts.reembed"):
+                populated.assert_embedding_dimensions(Mock(dimension=768))
 
 
 @pytest.mark.unit

tests/vectorstore/test_model_registry.py

@@ -0,0 +1,88 @@
+"""The registry is the single source of truth for embedding-model facts."""
+
+import pytest
+
+from application.vectorstore import model_registry as reg
+
+
+class TestResolve:
+    def test_resolves_canonical_name(self):
+        assert reg.resolve(reg.MPNET.name) is reg.MPNET
+
+    @pytest.mark.parametrize(
+        "alias",
+        [
+            "huggingface_sentence-transformers-all-mpnet-base-v2",
+            "sentence-transformers/all-mpnet-base-v2",
+            "all-mpnet-base-v2",
+        ],
+    )
+    def test_resolves_legacy_spellings_of_mpnet(self, alias):
+        """Every spelling the old factory dict accepted must still work."""
+        assert reg.resolve(alias) is reg.MPNET
+
+    def test_resolution_is_case_insensitive_and_trims(self):
+        assert reg.resolve("  GRANITE-311M  ") is reg.GRANITE_311M
+
+    def test_unknown_name_is_none_not_an_error(self):
+        """Unknown names are a valid configuration: an arbitrary HF repo."""
+        assert reg.resolve("some-org/some-model") is None
+
+    def test_empty_and_none_resolve_to_none(self):
+        assert reg.resolve(None) is None
+        assert reg.resolve("") is None
+
+
+class TestModelFacts:
+    def test_granite_311m_matches_the_existing_column_width(self):
+        """768 is what makes granite a drop-in for an mpnet index."""
+        assert reg.GRANITE_311M.dimension == reg.MPNET.dimension == 768
+
+    def test_granite_97m_is_narrower(self):
+        assert reg.GRANITE_97M.dimension == 384
+
+    def test_pooling_and_normalisation_are_recorded(self):
+        assert reg.MPNET.pooling == "mean"
+        assert reg.GRANITE_311M.pooling == "cls"
+        assert all(m.normalize for m in reg.MODELS)
+
+    def test_granite_context_is_much_wider_than_mpnet(self):
+        assert reg.GRANITE_311M.max_input_tokens == 32768
+        assert reg.MPNET.max_input_tokens == 384
+
+    def test_local_runners_carry_a_repo_and_onnx_file(self):
+        for model in reg.MODELS:
+            if model.provider == "fastembed":
+                assert model.repo, f"{model.name} has no repo"
+                assert model.onnx_file, f"{model.name} has no onnx_file"
+
+    def test_openai_model_needs_no_local_artifacts(self):
+        assert reg.OPENAI_ADA_002.provider == "openai"
+        assert reg.OPENAI_ADA_002.repo is None
+
+
+class TestHelpers:
+    def test_dimension_for_known_and_unknown(self):
+        assert reg.dimension_for(reg.GRANITE_97M.name) == 384
+        assert reg.dimension_for("nope/nope") is None
+
+    def test_max_input_tokens_for_known_and_unknown(self):
+        assert reg.max_input_tokens_for("granite-311m") == 32768
+        assert reg.max_input_tokens_for("nope/nope") is None
+
+    def test_known_names_lists_canonical_spellings(self):
+        names = reg.known_names()
+        assert reg.MPNET.name in names
+        assert reg.GRANITE_311M.name in names
+
+    def test_defaults_point_at_the_intended_models(self):
+        """Existing installs stay on mpnet; new installs get granite."""
+        assert reg.DEFAULT_LEGACY == reg.MPNET.name
+        assert reg.DEFAULT_NEW_INSTALL == reg.GRANITE_311M.name
+
+    def test_no_alias_collisions_between_models(self):
+        seen = {}
+        for model in reg.MODELS:
+            for key in (model.name, *model.aliases):
+                assert key.lower() not in seen, f"{key} claimed twice"
+                seen[key.lower()] = model

tests/vectorstore/test_remote_embeddings_truncation.py

@@ -83,3 +83,48 @@ def test_query_path_is_truncated(monkeypatch):
 
     sent = captured["payload"]["input"]
     assert sent == enc.decode(enc.encode(long_text)[:10])
+
+
+class TestInputLimitResolution:
+    """The cap falls back to the model's own context window."""
+
+    def _remote(self, model_name):
+        from application.vectorstore.base import RemoteEmbeddings
+
+        return RemoteEmbeddings(
+            api_url="http://embeddings", model_name=model_name, api_key=None
+        )
+
+    def test_explicit_setting_wins(self, monkeypatch):
+        from application.vectorstore import base
+
+        monkeypatch.setattr(base.settings, "EMBEDDINGS_MAX_INPUT_TOKENS", 123)
+        assert self._remote("granite-311m")._resolve_input_limit() == 123
+
+    def test_registered_model_supplies_its_own_ceiling(self, monkeypatch):
+        from application.vectorstore import base
+
+        monkeypatch.setattr(base.settings, "EMBEDDINGS_MAX_INPUT_TOKENS", None)
+        assert self._remote("granite-311m")._resolve_input_limit() == 32768
+        # mpnet genuinely stops at 384; sending more is paid for and discarded.
+        assert self._remote("all-mpnet-base-v2")._resolve_input_limit() == 384
+
+    def test_unknown_model_stays_unlimited(self, monkeypatch):
+        from application.vectorstore import base
+
+        monkeypatch.setattr(base.settings, "EMBEDDINGS_MAX_INPUT_TOKENS", None)
+        assert self._remote("some-org/mystery")._resolve_input_limit() is None
+
+    def test_non_positive_setting_falls_through_to_the_registry(self, monkeypatch):
+        from application.vectorstore import base
+
+        monkeypatch.setattr(base.settings, "EMBEDDINGS_MAX_INPUT_TOKENS", 0)
+        assert self._remote("granite-97m")._resolve_input_limit() == 32768
+
+    def test_dimension_is_taken_from_the_registry(self):
+        assert self._remote("granite-97m").dimension == 384
+        assert self._remote("granite-311m").dimension == 768
+
+    def test_unknown_model_dimension_is_probed_not_assumed(self):
+        """The old hardcoded 768 made the probe below unreachable."""
+        assert self._remote("some-org/mystery").dimension is None