fix: embeddings registry edge cases and chunk-size accounting

Follow-up to the embeddings work, from a review pass over the branch.

- Route the OpenAI/Azure key handling through the model registry instead of
  matching the canonical name literally, so the `text-embedding-ada-002`
  alias the registry now accepts also reaches the Azure deployment name
  rather than failing every embed with DeploymentNotFound.
- Fall back to a default width where the embeddings model reports no
  dimension. A model outside the registry returns None rather than no
  attribute, so `getattr` with a default did not catch it and the width
  reached the DDL as `vector(None)` / `list_size=None`.
- Point HF_HUB_CACHE at the prefetch directory. Chunking loads the tokenizer
  through `tokenizers`, which reads the hub cache, so a fresh container
  fetched over the network on first ingest and an offline one silently fell
  back to cl100k.
- Charge a token that collapses a long unbroken run by its character span.
  WordPiece emits one [UNK] for any word over its character limit, which made
  base64 and minified content count as near-zero tokens, so nothing split it
  and oversized chunks reached the embedding server.
- Preserve chunk ids and honour --batch-size when rebuilding a FAISS index.
  Fresh uuids orphaned GraphRAG's graph_node_chunks rows, and the whole index
  went out in a single embed call on remote servers.
- Document that granite runs an int8-quantised graph, and scope the
  SentenceTransformer parity claim to mpnet's fp32 graph, which is where it
  was measured.
- Correct the embeddings docs: a matching dimension is not a matching model,
  so a same-width swap raises nothing and silently degrades retrieval.

application/Dockerfile

@@ -70,7 +70,12 @@ COPY . /app/application
 # 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
+# The prefetch writes hub-layout snapshots (including tokenizer.json) here, so
+# HF_HUB_CACHE has to point at the same directory: chunking loads the tokenizer
+# through ``tokenizers``, which reads the hub cache and would otherwise fetch
+# over the network on first ingest -- and fall back to cl100k when offline.
+ENV EMBEDDINGS_CACHE_DIR=/app/models \
+    HF_HUB_CACHE=/app/models
 ARG EMBEDDINGS_PREFETCH=""
 RUN PYTHONPATH=/app /venv/bin/python -m application.scripts.prefetch_models ${EMBEDDINGS_PREFETCH}
 

application/graphrag/store.py

@@ -157,12 +157,14 @@ class GraphStore:
         """Dimension of the configured embeddings model, matching ``PGVectorStore``.
 
         Falls back to ``DEFAULT_NAME_EMBEDDING_DIM`` so the graph table and the
-        pgvector ``documents`` table always agree on the configured model.
+        pgvector ``documents`` table always agree on the configured model. A
+        model outside the registry reports ``None`` rather than no attribute,
+        so the fallback cannot be left to ``getattr``.
         """
         from application.vectorstore.base import get_embeddings
 
         embedding = get_embeddings()
-        return getattr(embedding, "dimension", DEFAULT_NAME_EMBEDDING_DIM)
+        return getattr(embedding, "dimension", None) or DEFAULT_NAME_EMBEDDING_DIM
 
     @staticmethod
     def create_schema(conn, *, dimension: int = DEFAULT_NAME_EMBEDDING_DIM) -> None:

application/parser/tokenization.py

@@ -49,6 +49,37 @@ def _windows(total: int, first: int, rest: int) -> Iterator[Tuple[int, int]]:
         budget = max(1, rest)
 
 
+# A token spanning more characters than this collapsed a run the tokenizer
+# could not break up: WordPiece emits a single ``[UNK]`` for any word longer
+# than ``max_input_chars_per_word``. Charging that once makes a base64 blob or
+# a minified bundle look tiny, so nothing splits it and an oversized chunk
+# reaches the embedding server. Real tokens are a few characters, so prose
+# never reaches this bound.
+_MAX_CHARS_PER_TOKEN = 16
+
+
+def _token_weight(start: int, end: int) -> int:
+    """Tokens a span costs, charging a collapsed run by its length."""
+    span = end - start
+    return max(1, -(-span // _MAX_CHARS_PER_TOKEN))
+
+
+def _cap_piece_chars(pieces: List[str], first: int, rest: int) -> List[str]:
+    """Cut any piece holding more characters than its budget can cover."""
+    capped: List[str] = []
+    budget = first
+    for piece in pieces:
+        limit = max(1, budget * _MAX_CHARS_PER_TOKEN)
+        while len(piece) > limit:
+            capped.append(piece[:limit])
+            piece = piece[limit:]
+            budget = rest
+            limit = max(1, budget * _MAX_CHARS_PER_TOKEN)
+        capped.append(piece)
+        budget = rest
+    return capped
+
+
 class TokenCounter:
     """Counts and splits text in one tokenizer's units.
 
@@ -123,7 +154,10 @@ class HuggingFaceCounter(TokenCounter):
     def count(self, text: str) -> int:
         if not text:
             return 0
-        return len(self._encode(text).ids)
+        encoding = self._encode(text)
+        if not encoding.offsets:
+            return len(encoding.ids)
+        return sum(_token_weight(start, end) for start, end in encoding.offsets)
 
     def split(
         self, text: str, max_tokens: int, first_max_tokens: Optional[int] = None
@@ -133,7 +167,7 @@ class HuggingFaceCounter(TokenCounter):
         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:
+        if self.count(text) <= first:
             return [text]
 
         pieces: List[str] = []
@@ -155,7 +189,7 @@ class HuggingFaceCounter(TokenCounter):
                 pieces[-1] = pieces[-1] + text[cursor:]
             else:
                 pieces.append(text[cursor:])
-        return pieces
+        return _cap_piece_chars(pieces, first, rest)
 
 
 def _load_hf_counter(repo: str) -> Optional[HuggingFaceCounter]:

application/scripts/reembed.py

@@ -216,11 +216,16 @@ def reembed_faiss(source_id: str, batch_size: int, dry_run: bool) -> Tuple[int,
     # 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.
+    # Keep the existing ids: GraphRAG's ``graph_node_chunks`` rows and any id a
+    # client already holds point at them.
+    ids = [chunk.get("doc_id") for chunk in chunks]
     rebuilt = VectorCreator.create_vectorstore(
         "faiss",
         source_id=source_id,
         embeddings_key=settings.EMBEDDINGS_KEY,
         docs_init=docs,
+        ids=ids if all(ids) else None,
+        batch_size=batch_size,
     )
     rebuilt.save_local()
     return len(chunks), len(docs)

application/vectorstore/base.py

@@ -285,7 +285,11 @@ def get_embeddings(
         )
         return EmbeddingsSingleton._remote_instance(embeddings_name, embeddings_key)
 
-    if embeddings_name == "openai_text-embedding-ada-002":
+    # Match through the registry, not on the canonical string: the bare
+    # ``text-embedding-ada-002`` alias resolves here too, and skipping this
+    # branch would drop the Azure deployment name and the API key.
+    spec = resolve(embeddings_name)
+    if spec is not None and spec.provider == "openai":
         if _azure_configured():
             embedding_instance = EmbeddingsSingleton.get_instance(
                 embeddings_name, model=settings.AZURE_EMBEDDINGS_DEPLOYMENT_NAME

application/vectorstore/embeddings_local.py

@@ -6,9 +6,11 @@ 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.
+text with SentenceTransformer and with FastEmbed's fp32 ONNX graph of mpnet
+returns vectors at cosine 1.0, so an mpnet index built before this change
+keeps working unchanged. That result is specific to the fp32 graph -- the
+granite entries run an int8-quantised one and are not bit-comparable to a
+fp32 index of the same model.
 
 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

application/vectorstore/faiss.py

@@ -81,7 +81,14 @@ class FaissStore(BaseVectorStore):
     # and must not be shown as one.
     score_kind = "l2_distance"
 
-    def __init__(self, source_id: str, embeddings_key: str, docs_init=None):
+    def __init__(
+        self,
+        source_id: str,
+        embeddings_key: str,
+        docs_init=None,
+        ids=None,
+        batch_size=None,
+    ):
         super().__init__()
         self.source_id = source_id
         self.path = get_vectorstore(source_id)
@@ -94,7 +101,7 @@ class FaissStore(BaseVectorStore):
 
         try:
             if docs_init:
-                self._build_from_documents(docs_init)
+                self._build_from_documents(docs_init, ids=ids, batch_size=batch_size)
             else:
                 self._load_from_storage()
         except Exception as e:
@@ -104,17 +111,36 @@ class FaissStore(BaseVectorStore):
 
     # -- Construction ----------------------------------------------------
 
-    def _build_from_documents(self, docs_init) -> None:
-        """Create a fresh index seeded with ``docs_init``."""
+    def _build_from_documents(self, docs_init, ids=None, batch_size=None) -> None:
+        """Create a fresh index seeded with ``docs_init``.
+
+        Args:
+            docs_init: Documents to embed.
+            ids: Chunk ids to keep. Generated when omitted, which renumbers
+                every chunk and orphans anything referencing the old ids.
+            batch_size: Documents per embed call. Without it the whole index
+                goes out in one call, which a remote embeddings server rejects
+                or times out on.
+        """
         texts, metadatas = [], []
         for doc in docs_init:
             texts.append(getattr(doc, "page_content", None) or getattr(doc, "text", "") or "")
             metadatas.append(getattr(doc, "metadata", None) or getattr(doc, "extra_info", None) or {})
 
         faiss = _dependable_faiss_import()
-        vectors = self.embeddings.embed_documents(texts)
-        self.index = faiss.IndexFlatL2(len(vectors[0]))
-        self._append(texts, metadatas, vectors)
+        ids = list(ids) if ids else None
+        step = batch_size if batch_size and batch_size > 0 else len(texts)
+        for start in range(0, len(texts), max(1, step)):
+            stop = start + max(1, step)
+            vectors = self.embeddings.embed_documents(texts[start:stop])
+            if self.index is None:
+                self.index = faiss.IndexFlatL2(len(vectors[0]))
+            self._append(
+                texts[start:stop],
+                metadatas[start:stop],
+                vectors,
+                ids[start:stop] if ids else None,
+            )
 
     def _load_from_storage(self) -> None:
         """Load the index and its sidecar, preferring JSON over the pickle."""

application/vectorstore/lancedb.py

@@ -2,6 +2,7 @@ from typing import List, Optional
 import importlib
 from application.vectorstore.base import BaseVectorStore
 from application.core.settings import settings
+from application.vectorstore.model_registry import DEFAULT_EMBEDDING_DIMENSION
 
 class LanceDBVectorStore(BaseVectorStore):
     """Class for LanceDB Vector Store integration."""
@@ -54,8 +55,11 @@ class LanceDBVectorStore(BaseVectorStore):
         """Ensure the table exists before performing operations."""
         if self.table is None:
             embeddings = self._get_embeddings(settings.EMBEDDINGS_NAME, self.embeddings_key)
+            # A model outside the registry reports no width until it has run;
+            # ``list_size=None`` is a TypeError, not a permissive schema.
+            dimension = getattr(embeddings, "dimension", None) or DEFAULT_EMBEDDING_DIMENSION
             schema = self.pa.schema([
-                self.pa.field("vector", self.pa.list_(self.pa.float32(), list_size=embeddings.dimension)),
+                self.pa.field("vector", self.pa.list_(self.pa.float32(), list_size=dimension)),
                 self.pa.field("text", self.pa.string()),
                 self.pa.field("metadata", self.pa.struct([
                     self.pa.field("key", self.pa.string()),

application/vectorstore/model_registry.py

@@ -70,6 +70,12 @@ MPNET = EmbeddingModel(
 
 #: Default for new installs: same 768 dimensions as mpnet, an 8x wider
 #: effective input, and multilingual retrieval.
+#:
+#: Runs the int8-quantised graph, which is ~4x smaller than the fp32 one
+#: (313 MB against 1247 MB) and keeps the image shippable now that both
+#: defaults are baked in. It costs some numeric fidelity against fp32 --
+#: measured at cosine 0.96 on short texts, with retrieval rank order
+#: unchanged -- so it is not bit-comparable to a fp32 granite index.
 GRANITE_311M = EmbeddingModel(
     name="ibm-granite/granite-embedding-311m-multilingual-r2",
     dimension=768,
@@ -82,6 +88,7 @@ GRANITE_311M = EmbeddingModel(
 )
 
 #: Smaller granite. Half the vector width, roughly three times the speed.
+#: Int8-quantised on the same terms as GRANITE_311M above.
 GRANITE_97M = EmbeddingModel(
     name="ibm-granite/granite-embedding-97m-multilingual-r2",
     dimension=384,