diff --git a/application/Dockerfile b/application/Dockerfile
index 44e7a530..590d1537 100644
--- a/application/Dockerfile
+++ b/application/Dockerfile
@@ -63,9 +63,6 @@ RUN groupadd -r appuser && \
# Copy the virtual environment and model from the builder stage
COPY --from=builder /venv /venv
-# 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
@@ -76,9 +73,20 @@ COPY . /app/application
# over the network on first ingest -- and fall back to cl100k when offline.
ENV EMBEDDINGS_CACHE_DIR=/app/models \
HF_HUB_CACHE=/app/models
+
+# Only the modules the prefetch imports are copied first. It reaches nothing
+# beyond model_registry, which is stdlib-only, so keeping the full source copy
+# below this layer stops an unrelated edit from re-downloading ~780 MB of model
+# artifacts on every build.
+COPY __init__.py /app/application/__init__.py
+COPY scripts/__init__.py scripts/prefetch_models.py /app/application/scripts/
+COPY vectorstore/__init__.py vectorstore/model_registry.py /app/application/vectorstore/
ARG EMBEDDINGS_PREFETCH=""
RUN PYTHONPATH=/app /venv/bin/python -m application.scripts.prefetch_models ${EMBEDDINGS_PREFETCH}
+# Copy your application code
+COPY . /app/application
+
# Change the ownership of the /app directory to the appuser
RUN mkdir -p /app/application/inputs/local
diff --git a/application/parser/chunking.py b/application/parser/chunking.py
index 8c150338..099337df 100644
--- a/application/parser/chunking.py
+++ b/application/parser/chunking.py
@@ -7,6 +7,13 @@ from application.parser.tokenization import get_token_counter
logger = logging.getLogger(__name__)
+# Smallest share of ``max_tokens`` a chunk must keep for body text when the
+# header is duplicated onto every chunk. A header that leaves less than this
+# makes each chunk mostly repeated header and multiplies the chunk count -- at
+# a budget of 32 out of 1250 a document splits into 39x more chunks than it
+# needs -- so duplication is dropped rather than honoured.
+_MIN_BODY_BUDGET_RATIO = 0.25
+
class Chunker:
"""Classic token-window chunker (registered as ``classic_chunk``).
@@ -24,7 +31,9 @@ class Chunker:
duplicate_headers: bool = False,
):
self.chunking_strategy = chunking_strategy
- self.max_tokens = max_tokens
+ # A budget below 1 would ask for a chunk per token; the strategy
+ # chunkers clamp the same way.
+ self.max_tokens = max(1, int(max_tokens))
self.min_tokens = min_tokens
self.duplicate_headers = duplicate_headers
# Counted in the embedding model's tokenizer, not cl100k: ``max_tokens``
@@ -55,18 +64,51 @@ class Chunker:
header, body = self.separate_header_and_body(doc.text)
header_tokens = self.counter.count(header) if header else 0
+ if header and header_tokens >= self.max_tokens:
+ # The header alone fills the budget, so no cut of the body can keep
+ # a chunk within it and duplicating it would leave a one-token body
+ # budget -- a chunk per body token. It is only the first three
+ # lines, not something worth preserving at that cost, so it goes
+ # back to being ordinary text and the document splits evenly.
+ logger.warning(
+ "Header of %s is %d token(s), at or over the %d-token chunk "
+ "budget; treating it as body text.",
+ doc.doc_id,
+ header_tokens,
+ self.max_tokens,
+ )
+ body = f"{header}{body}"
+ header, header_tokens = "", 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:
+ duplicate_headers = self.duplicate_headers
+ if duplicate_headers and with_header_budget < self.max_tokens * _MIN_BODY_BUDGET_RATIO:
+ logger.warning(
+ "Header of %s leaves only %d of %d tokens for body text; "
+ "carrying it on the first chunk only.",
+ doc.doc_id,
+ with_header_budget,
+ self.max_tokens,
+ )
+ duplicate_headers = False
+
+ if 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
)
+ if not body_pieces and header:
+ # Nothing but a header: the loop below only ever emits the header
+ # attached to a body piece, so without this the document is dropped
+ # from the index entirely.
+ body_pieces = [""]
+
split_docs = []
for part_index, piece in enumerate(body_pieces):
- include_header = bool(header) and (self.duplicate_headers or part_index == 0)
+ include_header = bool(header) and (duplicate_headers or part_index == 0)
chunk_text = f"{header}{piece}" if include_header else piece
split_docs.append(
Document(
diff --git a/application/parser/tokenization.py b/application/parser/tokenization.py
index 3af02994..88059644 100644
--- a/application/parser/tokenization.py
+++ b/application/parser/tokenization.py
@@ -135,10 +135,34 @@ class TiktokenCounter(TokenCounter):
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)
- ]
+
+ # Cut the original string at the character offsets the tokenizer
+ # reports, the way :class:`HuggingFaceCounter` does. Decoding each
+ # window on its own instead splits any multi-byte character that
+ # straddles a boundary across two byte sequences, and each half decodes
+ # to U+FFFD -- roughly one boundary in five on CJK text, silently
+ # destroying a character per cut.
+ offsets = self._encoding.decode_with_offsets(tokens)[1]
+ end_of_text = len(text)
+ pieces: List[str] = []
+ cursor = 0
+ for _, end in _windows(len(tokens), first, rest):
+ # Offsets index into the decoded string. ``encode_ordinary``
+ # round-trips for anything cl100k can represent, so it is ``text``;
+ # clamping keeps the cuts in range if it ever is not, and slicing
+ # ``text`` monotonically keeps the pieces reassembling exactly
+ # either way.
+ end_char = end_of_text if end >= len(offsets) else min(offsets[end], end_of_text)
+ if end_char <= cursor:
+ continue
+ pieces.append(text[cursor:end_char])
+ cursor = end_char
+ if cursor < end_of_text:
+ if pieces:
+ pieces[-1] = pieces[-1] + text[cursor:]
+ else:
+ pieces.append(text[cursor:])
+ return pieces
class HuggingFaceCounter(TokenCounter):
diff --git a/application/scripts/reembed.py b/application/scripts/reembed.py
index 7d716619..0434d21d 100644
--- a/application/scripts/reembed.py
+++ b/application/scripts/reembed.py
@@ -8,7 +8,9 @@ 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.
+crawl, no docling. It works against pgvector and FAISS. With ``GRAPHRAG_ENABLED``
+it also rewrites ``graph_nodes.name_embedding``, which seeds graph traversal and
+would otherwise be left behind in the previous model's space.
Usage::
@@ -122,6 +124,79 @@ def _faiss_source_ids() -> List[str]:
return sorted(set(ids))
+def _graph_nodes_exist(conn) -> bool:
+ """True when the GraphRAG node table is present in this database."""
+ cursor = conn.cursor()
+ try:
+ cursor.execute("SELECT to_regclass('public.graph_nodes')")
+ row = cursor.fetchone()
+ return bool(row and row[0])
+ finally:
+ cursor.close()
+
+
+def reembed_graph_nodes(store, conn, source_id: str, batch_size: int, dry_run: bool) -> int:
+ """Rewrite one source's GraphRAG node name vectors.
+
+ ``graph_nodes.name_embedding`` is what every graph traversal starts from:
+ the retriever embeds the query and takes the nearest node names as seeds.
+ It is written once at extraction time and never revisited, so re-embedding
+ only the chunk table leaves the graph seeding against the previous model.
+ Nothing errors -- mpnet and granite-311m are both 768-dimensional, so the
+ column accepts the mismatch -- the graph just walks from the wrong nodes.
+
+ The embedded text is ``graph_nodes.name``, which is already stored, so this
+ needs no LLM re-extraction.
+
+ Args:
+ store: The pgvector store, for its embeddings client.
+ conn: Open connection to the same database.
+ source_id: Source whose nodes to re-embed.
+ batch_size: Names per embed call and per transaction.
+ dry_run: When true, count the work and change nothing.
+
+ Returns:
+ Number of node rows seen (dry run) or rewritten.
+ """
+ if not _graph_nodes_exist(conn):
+ return 0
+
+ cursor = conn.cursor()
+ try:
+ cursor.execute(
+ "SELECT id, name FROM graph_nodes "
+ "WHERE source_id = %s AND name_embedding IS NOT NULL ORDER BY id",
+ (source_id,),
+ )
+ rows = cursor.fetchall()
+ finally:
+ cursor.close()
+
+ if dry_run or not rows:
+ return len(rows)
+
+ written = 0
+ 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(
+ "UPDATE graph_nodes SET name_embedding = %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)
+ return written
+
+
def reembed_pgvector(source_id: str, batch_size: int, dry_run: bool) -> Tuple[int, int]:
"""Rewrite one source's vectors in place.
@@ -152,35 +227,45 @@ def reembed_pgvector(source_id: str, batch_size: int, dry_run: bool) -> Tuple[in
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(
- sql.SQL(
- "UPDATE {table} SET {column} = %s::vector WHERE id = %s"
- ).format(
- table=sql.Identifier(table),
- column=sql.Identifier(vector_column),
- ),
- [
- (str(list(vector)), row[0])
- for vector, row in zip(vectors, batch)
- ],
+ if not dry_run:
+ 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(
+ sql.SQL(
+ "UPDATE {table} SET {column} = %s::vector WHERE id = %s"
+ ).format(
+ table=sql.Identifier(table),
+ column=sql.Identifier(vector_column),
+ ),
+ [
+ (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)
+
+ # The graph seeds every traversal from its own vectors, so leaving them
+ # in the old model's space is the same silent mismatch this script
+ # exists to remove -- and at equal widths nothing would report it.
+ if getattr(settings, "GRAPHRAG_ENABLED", False):
+ nodes = reembed_graph_nodes(store, conn, source_id, batch_size, dry_run)
+ if nodes:
+ logger.info(
+ " %s: %d graph node name(s)%s",
+ source_id,
+ nodes,
+ "" if dry_run else " re-embedded",
)
- conn.commit()
- except Exception:
- conn.rollback()
- raise
- finally:
- cursor.close()
- written += len(batch)
finally:
store.close()
return len(rows), written
@@ -201,8 +286,15 @@ def reembed_faiss(source_id: str, batch_size: int, dry_run: bool) -> Tuple[int,
Returns:
``(chunks_seen, chunks_written)``.
"""
+ # A width change is the main reason to run this, and it is exactly what
+ # ``assert_embedding_dimensions`` refuses to open. The chunk text lives in
+ # the sidecar rather than the index, so reading it needs no matching width,
+ # and the index this reads is replaced below.
store = VectorCreator.create_vectorstore(
- "faiss", source_id=source_id, embeddings_key=settings.EMBEDDINGS_KEY
+ "faiss",
+ source_id=source_id,
+ embeddings_key=settings.EMBEDDINGS_KEY,
+ skip_dimension_check=True,
)
chunks: List[Dict[str, Any]] = store.get_chunks() or []
if dry_run or not chunks:
diff --git a/application/vectorstore/base.py b/application/vectorstore/base.py
index 0270cda4..f503d9fd 100644
--- a/application/vectorstore/base.py
+++ b/application/vectorstore/base.py
@@ -14,8 +14,21 @@ from application.vectorstore.model_registry import (
def _embeddings_name_is_explicit() -> bool:
- """True when ``EMBEDDINGS_NAME`` was configured rather than defaulted."""
- return "EMBEDDINGS_NAME" in getattr(settings, "model_fields_set", set())
+ """True when ``EMBEDDINGS_NAME`` names a model somebody actually chose.
+
+ Not ``model_fields_set``: pydantic marks a field as set for any value that
+ reached it, including one read from ``.env``, and every setup script has
+ always written ``EMBEDDINGS_NAME`` unconditionally. An install carrying the
+ legacy name a script wrote years ago would read as a deliberate choice and
+ lend a remote server that model's context window.
+ """
+ fields = getattr(type(settings), "model_fields", None)
+ if not isinstance(fields, dict):
+ return False
+ field = fields.get("EMBEDDINGS_NAME")
+ if field is None:
+ return False
+ return settings.EMBEDDINGS_NAME != field.default
class RemoteEmbeddings:
diff --git a/application/vectorstore/embeddings_local.py b/application/vectorstore/embeddings_local.py
index fff6f875..817db488 100644
--- a/application/vectorstore/embeddings_local.py
+++ b/application/vectorstore/embeddings_local.py
@@ -44,6 +44,17 @@ def _pooling_type(pooling: str):
return {"cls": PoolingType.CLS, "mean": PoolingType.MEAN}[pooling]
+def _is_builtin(repo: str) -> bool:
+ """True when FastEmbed already ships a description for ``repo``."""
+ from fastembed import TextEmbedding
+
+ lowered = repo.lower()
+ return any(
+ str(entry.get("model", "")).lower() == lowered
+ for entry in TextEmbedding.list_supported_models()
+ )
+
+
def _register(model: EmbeddingModel) -> None:
"""Teach FastEmbed about a model, exactly once per process."""
from fastembed import TextEmbedding
@@ -52,6 +63,14 @@ def _register(model: EmbeddingModel) -> None:
with _register_lock:
if model.repo in _registered:
return
+ if _is_builtin(model.repo):
+ # ``add_custom_model`` refuses a name FastEmbed already ships, and
+ # its own description carries the pooling, width and graph file we
+ # would be supplying, so there is nothing to add. Without this,
+ # configuring any of FastEmbed's ~30 built-in models (bge, e5,
+ # MiniLM, gte, ...) fails every embed call.
+ _registered.add(model.repo)
+ return
TextEmbedding.add_custom_model(
model=model.repo,
pooling=_pooling_type(model.pooling),
@@ -162,11 +181,6 @@ class EmbeddingsWrapper:
"""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]
diff --git a/application/vectorstore/faiss.py b/application/vectorstore/faiss.py
index 86526c92..1db5a53e 100644
--- a/application/vectorstore/faiss.py
+++ b/application/vectorstore/faiss.py
@@ -88,7 +88,22 @@ class FaissStore(BaseVectorStore):
docs_init=None,
ids=None,
batch_size=None,
+ skip_dimension_check: bool = False,
):
+ """Open or build one source's FAISS index.
+
+ Args:
+ source_id: Source whose index to open.
+ embeddings_key: API key handed to the embeddings provider.
+ docs_init: Documents to build a fresh index from. Loads the stored
+ index instead when omitted.
+ ids: Chunk ids to keep when building. Generated when omitted.
+ batch_size: Documents per embed call when building.
+ skip_dimension_check: Open an index whose width does not match the
+ configured model. Only for a caller that is about to replace
+ that index, such as the re-embed script, which otherwise cannot
+ read the chunks it needs to rebuild from.
+ """
super().__init__()
self.source_id = source_id
self.path = get_vectorstore(source_id)
@@ -107,7 +122,8 @@ class FaissStore(BaseVectorStore):
except Exception as e:
raise Exception(f"Error loading FAISS index: {str(e)}")
- self.assert_embedding_dimensions(self.embeddings)
+ if not skip_dimension_check:
+ self.assert_embedding_dimensions(self.embeddings)
# -- Construction ----------------------------------------------------
diff --git a/docs/content/Models/embeddings.md b/docs/content/Models/embeddings.md
index 8fcc5c59..634eaaee 100644
--- a/docs/content/Models/embeddings.md
+++ b/docs/content/Models/embeddings.md
@@ -111,7 +111,11 @@ Switching between same-width models therefore still requires re-embedding:
python -m application.scripts.reembed
```
-Run it after changing `EMBEDDINGS_NAME` and before serving queries. See [Upgrading](/upgrading) for the granite migration specifically. Changing to a model of a *different* width is not supported by the script — re-ingest those sources instead.
+Run it after changing `EMBEDDINGS_NAME` and before serving queries. See [Upgrading](/upgrading) for the granite migration specifically.
+
+Changing to a model of a *different* width is supported for FAISS: the script rebuilds the index at the new width and keeps the existing chunk ids. For `pgvector` the vector column is sized at creation time, so a width change there still means re-ingesting those sources.
+
+With `GRAPHRAG_ENABLED`, the script also rewrites `graph_nodes.name_embedding` on `pgvector`. Those vectors seed every graph traversal, and they share the chunk vectors' width, so leaving them behind degrades graph retrieval just as silently.
## Adding Support for Other Embedding Models
diff --git a/docs/content/upgrading.mdx b/docs/content/upgrading.mdx
index 70f07ead..866c848a 100644
--- a/docs/content/upgrading.mdx
+++ b/docs/content/upgrading.mdx
@@ -46,6 +46,10 @@ Useful flags:
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.
+
+ Running [GraphRAG](/Sources/GraphRAG)? The script also rewrites `graph_nodes.name_embedding`, which seeds every graph traversal. Those vectors are written once at extraction time and share the chunk vectors' width, so leaving them in the old model's space degrades graph retrieval just as silently as the chunk vectors would — and needs no LLM re-extraction to fix.
+
+
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.
diff --git a/setup.ps1 b/setup.ps1
index 8d94a2e6..40943515 100644
--- a/setup.ps1
+++ b/setup.ps1
@@ -377,17 +377,18 @@ function Configure-Embeddings {
Write-Host ""
Write-ColorText "Embeddings Configuration" -ForegroundColor "White" -Bold
Write-ColorText "Choose your embeddings provider:" -ForegroundColor "White"
- Write-ColorText "1) HuggingFace (default, local)" -ForegroundColor "Yellow"
+ Write-ColorText "1) Granite multilingual (default, local)" -ForegroundColor "Yellow"
Write-ColorText "2) OpenAI Embeddings" -ForegroundColor "Yellow"
Write-ColorText "3) Custom Remote Embeddings (OpenAI-compatible API)" -ForegroundColor "Yellow"
+ Write-ColorText "4) all-mpnet-base-v2 (legacy local, English-only)" -ForegroundColor "Yellow"
Write-ColorText "b) Back" -ForegroundColor "Yellow"
Write-Host ""
- $emb_choice = Read-Host "Choose option (1-3, or b)"
+ $emb_choice = Read-Host "Choose option (1-4, or b)"
switch ($emb_choice) {
"1" {
- "EMBEDDINGS_NAME=huggingface_sentence-transformers/all-mpnet-base-v2" | Add-Content -Path $ENV_FILE -Encoding utf8
- Write-ColorText "Embeddings set to HuggingFace (local)." -ForegroundColor "Green"
+ "EMBEDDINGS_NAME=ibm-granite/granite-embedding-311m-multilingual-r2" | Add-Content -Path $ENV_FILE -Encoding utf8
+ Write-ColorText "Embeddings set to granite-311m-multilingual-r2 (local)." -ForegroundColor "Green"
}
"2" {
"EMBEDDINGS_NAME=openai_text-embedding-ada-002" | Add-Content -Path $ENV_FILE -Encoding utf8
@@ -404,6 +405,10 @@ function Configure-Embeddings {
if ($emb_key) { "EMBEDDINGS_KEY=$emb_key" | Add-Content -Path $ENV_FILE -Encoding utf8 }
Write-ColorText "Custom remote embeddings configured." -ForegroundColor "Green"
}
+ "4" {
+ "EMBEDDINGS_NAME=huggingface_sentence-transformers/all-mpnet-base-v2" | Add-Content -Path $ENV_FILE -Encoding utf8
+ Write-ColorText "Embeddings set to all-mpnet-base-v2 (legacy local)." -ForegroundColor "Green"
+ }
{$_ -eq "b" -or $_ -eq "B"} { return }
default {
Write-Host ""
@@ -716,7 +721,7 @@ function Serve-LocalOllama {
"LLM_NAME=$model_name" | Add-Content -Path $ENV_FILE -Encoding utf8
"VITE_API_STREAMING=true" | Add-Content -Path $ENV_FILE -Encoding utf8
"OPENAI_BASE_URL=http://ollama:11434/v1" | Add-Content -Path $ENV_FILE -Encoding utf8
- "EMBEDDINGS_NAME=huggingface_sentence-transformers/all-mpnet-base-v2" | Add-Content -Path $ENV_FILE -Encoding utf8
+ "EMBEDDINGS_NAME=ibm-granite/granite-embedding-311m-multilingual-r2" | Add-Content -Path $ENV_FILE -Encoding utf8
Write-ColorText ".env file configured for Ollama ($($docker_compose_file_suffix.ToUpper()))." -ForegroundColor "Green"
Write-ColorText "Note: MODEL_NAME is set to '$model_name'. You can change it later in the .env file." -ForegroundColor "Yellow"
@@ -881,7 +886,7 @@ function Connect-LocalInferenceEngine {
"LLM_NAME=$model_name" | Add-Content -Path $ENV_FILE -Encoding utf8
"VITE_API_STREAMING=true" | Add-Content -Path $ENV_FILE -Encoding utf8
"OPENAI_BASE_URL=$openai_base_url" | Add-Content -Path $ENV_FILE -Encoding utf8
- "EMBEDDINGS_NAME=huggingface_sentence-transformers/all-mpnet-base-v2" | Add-Content -Path $ENV_FILE -Encoding utf8
+ "EMBEDDINGS_NAME=ibm-granite/granite-embedding-311m-multilingual-r2" | Add-Content -Path $ENV_FILE -Encoding utf8
Write-ColorText ".env file configured for $engine_name with OpenAI API format." -ForegroundColor "Green"
Write-ColorText "Note: MODEL_NAME is set to '$model_name'. You can change it later in the .env file." -ForegroundColor "Yellow"
diff --git a/tests/parser/test_chunking.py b/tests/parser/test_chunking.py
index 57d038b1..7c817dfc 100644
--- a/tests/parser/test_chunking.py
+++ b/tests/parser/test_chunking.py
@@ -127,6 +127,71 @@ class TestSplitDocument:
assert len(result) > 1
assert all("h1" in chunk.text for chunk in result)
+ def test_oversized_header_does_not_multiply_chunks(self):
+ """A header past the budget used to leave a one-token body budget.
+
+ ``max(1, max_tokens - header_tokens)`` bottomed out at 1, so the body
+ was cut into one chunk per token -- each still over ``max_tokens``,
+ since the whole header was prepended to it.
+ """
+ chunker = Chunker(max_tokens=50, min_tokens=5, duplicate_headers=True)
+ header = ("verylongheaderword " * 40) + "h2\nh3\n"
+ body = "word " * 100
+ doc = Document(text=f"{header}\n{body}", doc_id="doc1")
+
+ result = chunker.split_document(doc)
+
+ body_tokens = chunker.counter.count(body)
+ assert len(result) <= body_tokens // 10, "chunk count must track the budget"
+ for chunk in result:
+ assert chunk.extra_info["token_count"] <= chunker.max_tokens
+ assert "".join(chunk.text for chunk in result) == doc.text
+
+ def test_header_taking_most_of_the_budget_is_not_duplicated(self):
+ """Duplication is dropped when it would leave almost no room for body.
+
+ Below a quarter of the budget every chunk is mostly repeated header,
+ which multiplies the chunk count without adding retrievable text.
+ """
+ chunker = Chunker(max_tokens=50, min_tokens=5, duplicate_headers=True)
+ header = ("headerword " * 20) + "h2\nh3\n"
+ doc = Document(text=f"{header}\n" + "word " * 200, doc_id="doc1")
+
+ result = chunker.split_document(doc)
+
+ assert len(result) > 1
+ assert "headerword" in result[0].text
+ assert all("headerword" not in chunk.text for chunk in result[1:])
+
+ def test_header_only_document_is_not_dropped(self):
+ """The loop only ever emitted the header attached to a body piece.
+
+ With no body there was no piece, so an entire document disappeared
+ from the index with no error and no log line.
+ """
+ chunker = Chunker(max_tokens=50, min_tokens=1, duplicate_headers=False)
+ text = "h1\nh2\nh3\n"
+ doc = Document(text=text, doc_id="doc1")
+
+ result = chunker.split_document(doc)
+
+ assert len(result) == 1
+ assert result[0].text == text
+
+ def test_header_only_document_keeps_its_text_through_chunk(self):
+ """The reachable shape: three lines that alone exceed the budget."""
+ chunker = Chunker(max_tokens=5, min_tokens=1, duplicate_headers=False)
+ text = "h1\nh2\nh3\n"
+
+ result = chunker.chunk([Document(text=text, doc_id="doc1")])
+
+ assert result, "the document must not vanish"
+ assert "".join(chunk.text for chunk in result) == text
+
+ def test_zero_max_tokens_does_not_split_per_token(self):
+ chunker = Chunker(max_tokens=0, min_tokens=1)
+ assert chunker.max_tokens == 1
+
def test_split_preserves_embedding(self):
chunker = Chunker(max_tokens=50, min_tokens=5)
doc = Document(
diff --git a/tests/parser/test_tokenization.py b/tests/parser/test_tokenization.py
index df51847a..afa7bd0f 100644
--- a/tests/parser/test_tokenization.py
+++ b/tests/parser/test_tokenization.py
@@ -29,6 +29,10 @@ class _StubEncoding:
def decode(self, ids):
return "".join(chr(i) for i in ids)
+ def decode_with_offsets(self, ids):
+ # One token per character, so each token starts where the last ended.
+ return self.decode(ids), list(range(len(ids)))
+
@pytest.fixture(autouse=True)
def _clear_cache():
@@ -107,6 +111,51 @@ class TestSelection:
assert get_token_counter("granite-311m") is first
+class TestTiktokenSplitAgainstRealCl100k:
+ """The stub above is one token per character, so it can never place a cut
+ inside a character. Real cl100k can, and that is the case that corrupted
+ text: decoding each window on its own turns a straddled multi-byte
+ character into U+FFFD on both sides of the cut."""
+
+ @pytest.fixture
+ def real_counter(self):
+ try:
+ counter = TiktokenCounter()
+ counter.count("probe")
+ except Exception as exc: # offline CI, same policy as the HF fixture
+ pytest.skip(f"cl100k encoding unavailable: {exc}")
+ return counter
+
+ # 2000 is the shipped default max_tokens, 384 mpnet's window; the small
+ # values place many more cuts per unit of text.
+ @pytest.mark.parametrize("window", [1, 2, 3, 7, 128, 384, 2000])
+ @pytest.mark.parametrize(
+ "text",
+ [
+ "日本語のテキストです。絵文字も🎉あります。",
+ "検索は自然言語でできます。" * 200,
+ "Здравствуйте, как настроить аутентификацию?",
+ "🎉🎊✨🚀🔥💡📚🧠" * 50,
+ "Ünïcödé — em-dashes, curly “quotes”, and 日本語 text.",
+ ],
+ ids=["ja-short", "ja-long", "ru", "emoji", "mixed"],
+ )
+ def test_split_reassembles_exactly(self, real_counter, text, window):
+ pieces = real_counter.split(text, window)
+ assert "".join(pieces) == text
+
+ @pytest.mark.parametrize("window", [1, 3, 128, 2000])
+ def test_split_never_emits_a_replacement_character(self, real_counter, window):
+ text = "検索は自然言語でできます。絵文字も🎉あります。" * 100
+ assert "�" not in "".join(real_counter.split(text, window))
+
+ def test_first_window_budget_is_honoured_and_lossless(self, real_counter):
+ text = "日本語のテキストです。" * 50
+ pieces = real_counter.split(text, 20, first_max_tokens=5)
+ assert "".join(pieces) == text
+ assert real_counter.count(pieces[0]) <= 5
+
+
class TestTiktokenCounterEdges:
"""The cl100k path is the fallback, so its edges matter as much."""
diff --git a/tests/scripts/test_reembed.py b/tests/scripts/test_reembed.py
index 71ed28c6..14d8e374 100644
--- a/tests/scripts/test_reembed.py
+++ b/tests/scripts/test_reembed.py
@@ -117,6 +117,20 @@ class TestFaissRebuild:
reembed.reembed_faiss("s1", batch_size=8, dry_run=False)
assert factory.call_args_list[1].kwargs["ids"] == ["1", "2"]
+ def test_existing_index_is_opened_past_the_dimension_check(self, stores):
+ """A width change is the main reason to run this script.
+
+ ``assert_embedding_dimensions`` refuses to open an index whose width
+ differs from the configured model -- and its error message recommends
+ this script, so without the opt-out the advice failed on every source.
+ The chunk text lives in the sidecar, so reading it needs no match.
+ """
+ _, _, factory = stores
+ reembed.reembed_faiss("s1", batch_size=8, dry_run=False)
+ assert factory.call_args_list[0].kwargs["skip_dimension_check"] is True
+ # The rebuild writes the new width, so it must still be checked.
+ assert "skip_dimension_check" not in factory.call_args_list[1].kwargs
+
def test_batch_size_is_forwarded_to_the_rebuild(self, stores):
"""On a remote embeddings server the whole index is otherwise one POST."""
_, _, factory = stores
@@ -171,7 +185,7 @@ class TestPgvectorWithoutTheExtension:
]
with patch.object(
reembed.VectorCreator, "create_vectorstore", return_value=store
- ):
+ ), patch.object(reembed.settings, "GRAPHRAG_ENABLED", False):
yield store, conn, cursor
def test_reads_and_rewrites_every_chunk(self, store):
@@ -255,3 +269,108 @@ class TestFaissSourceDiscovery:
def test_unsupported_store_error_names_the_alternatives(self):
with patch.object(reembed.settings, "VECTOR_STORE", "milvus", create=True):
assert reembed.main([]) == 2
+
+
+class TestGraphNodeReembedding:
+ """``graph_nodes.name_embedding`` seeds every graph traversal.
+
+ It is written once at extraction time and never revisited, so rewriting
+ only the chunk table leaves the graph seeding from the previous model --
+ and since mpnet and granite-311m are both 768-dimensional, the column
+ accepts the mismatch and nothing reports it.
+ """
+
+ @pytest.fixture
+ def graph(self):
+ store = MagicMock()
+ store._embedding.embed_documents.side_effect = lambda texts: [
+ [0.5] * 4 for _ in texts
+ ]
+ cursor = MagicMock()
+ cursor.fetchone.return_value = ("graph_nodes",)
+ cursor.fetchall.return_value = [
+ ("n1", "Alpha"),
+ ("n2", "Beta"),
+ ("n3", "Gamma"),
+ ]
+ conn = MagicMock()
+ conn.cursor.return_value = cursor
+ return store, conn, cursor
+
+ def test_rewrites_every_node_name(self, graph):
+ store, conn, cursor = graph
+ written = reembed.reembed_graph_nodes(
+ store, conn, "s1", batch_size=64, dry_run=False
+ )
+ assert written == 3
+ store._embedding.embed_documents.assert_called_once_with(
+ ["Alpha", "Beta", "Gamma"]
+ )
+ statement = cursor.executemany.call_args.args[0]
+ assert "graph_nodes" in statement and "name_embedding" in statement
+ conn.commit.assert_called()
+
+ def test_dry_run_counts_without_embedding(self, graph):
+ store, conn, cursor = graph
+ assert reembed.reembed_graph_nodes(store, conn, "s1", 64, dry_run=True) == 3
+ store._embedding.embed_documents.assert_not_called()
+ cursor.executemany.assert_not_called()
+
+ def test_missing_table_is_a_no_op(self, graph):
+ store, conn, cursor = graph
+ cursor.fetchone.return_value = (None,)
+ assert reembed.reembed_graph_nodes(store, conn, "s1", 64, dry_run=False) == 0
+ store._embedding.embed_documents.assert_not_called()
+
+ def test_batches_commit_separately(self, graph):
+ store, conn, cursor = graph
+ reembed.reembed_graph_nodes(store, conn, "s1", batch_size=2, dry_run=False)
+ assert cursor.executemany.call_count == 2
+ assert conn.commit.call_count == 2
+
+ def test_failed_batch_rolls_back_and_raises(self, graph):
+ store, conn, cursor = graph
+ cursor.executemany.side_effect = RuntimeError("write failed")
+ with pytest.raises(RuntimeError, match="write failed"):
+ reembed.reembed_graph_nodes(store, conn, "s1", 64, dry_run=False)
+ conn.rollback.assert_called_once()
+
+ def test_pgvector_run_skips_the_graph_when_disabled(self):
+ store, conn, cursor = self._pgvector_mocks()
+ with patch.object(
+ reembed.VectorCreator, "create_vectorstore", return_value=store
+ ), patch.object(reembed.settings, "GRAPHRAG_ENABLED", False):
+ reembed.reembed_pgvector("s1", batch_size=64, dry_run=False)
+ assert not self._graph_statements(cursor)
+
+ def test_pgvector_run_reembeds_the_graph_when_enabled(self):
+ store, conn, cursor = self._pgvector_mocks()
+ with patch.object(
+ reembed.VectorCreator, "create_vectorstore", return_value=store
+ ), patch.object(reembed.settings, "GRAPHRAG_ENABLED", True):
+ reembed.reembed_pgvector("s1", batch_size=64, dry_run=False)
+ assert self._graph_statements(cursor)
+
+ @staticmethod
+ def _pgvector_mocks():
+ store = MagicMock()
+ store._table_name = "documents"
+ store._vector_column = "embedding"
+ cursor = MagicMock()
+ cursor.fetchone.return_value = ("graph_nodes",)
+ cursor.fetchall.return_value = [(1, "alpha"), (2, "beta")]
+ conn = MagicMock()
+ conn.cursor.return_value = cursor
+ store._get_connection.return_value = conn
+ store._embedding.embed_documents.side_effect = lambda texts: [
+ [0.5] * 4 for _ in texts
+ ]
+ return store, conn, cursor
+
+ @staticmethod
+ def _graph_statements(cursor):
+ return [
+ call
+ for call in cursor.executemany.call_args_list
+ if "graph_nodes" in str(call.args[0])
+ ]
diff --git a/tests/scripts/test_reembed_pgvector_live.py b/tests/scripts/test_reembed_pgvector_live.py
index 165e311d..8f44bc14 100644
--- a/tests/scripts/test_reembed_pgvector_live.py
+++ b/tests/scripts/test_reembed_pgvector_live.py
@@ -189,3 +189,109 @@ class TestReembedPgvectorLive:
_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"]
+
+
+GRAPH_SOURCE = "11111111-2222-3333-4444-555555555555"
+
+
+def _seed_graph_node(dsn, source_id, name, seed):
+ """Insert one graph node carrying a name embedding at ``seed``."""
+ from application.graphrag.store import GraphStore
+
+ store = PGVectorStore(source_id=source_id, connection_string=dsn)
+ conn = store._get_connection()
+ GraphStore.create_schema(conn, dimension=DIM)
+ cursor = conn.cursor()
+ try:
+ cursor.execute(
+ "INSERT INTO graph_nodes (id, source_id, name, normalized_name, type, "
+ "description, degree, doc_freq, name_embedding) "
+ "VALUES (gen_random_uuid(), %s, %s, %s, 'ENTITY', '', 0, 0, %s::vector)",
+ (source_id, name, name.lower(), str([seed] + [0.0] * (DIM - 1))),
+ )
+ conn.commit()
+ finally:
+ cursor.close()
+ store.close()
+
+
+def _node_vectors(dsn, source_id):
+ store = PGVectorStore(source_id=source_id, connection_string=dsn)
+ conn = store._get_connection()
+ cursor = conn.cursor()
+ try:
+ cursor.execute(
+ "SELECT name, name_embedding FROM graph_nodes "
+ "WHERE source_id = %s ORDER BY name",
+ (source_id,),
+ )
+ return [(name, _as_list(vector)) for name, vector in cursor.fetchall()]
+ finally:
+ cursor.close()
+ store.close()
+
+
+class TestGraphNodeReembedLive:
+ """``graph_nodes.name_embedding`` seeds every graph traversal.
+
+ Rewriting only the chunk table leaves it in the previous model's space,
+ and because mpnet and granite-311m share a width the column accepts the
+ mismatch silently -- exactly the failure the script exists to prevent.
+ """
+
+ def test_node_names_are_re_embedded(self, live_dsn, monkeypatch):
+ from application.core import settings as settings_module
+
+ monkeypatch.setattr(
+ settings_module.settings, "GRAPHRAG_ENABLED", True, raising=False
+ )
+ _seed(live_dsn, GRAPH_SOURCE, TEXTS, _Embeddings(1.0))
+ _seed_graph_node(live_dsn, GRAPH_SOURCE, "Alpha", 1.0)
+
+ assert all(v[0] == pytest.approx(1.0) for _, v in _node_vectors(live_dsn, GRAPH_SOURCE))
+
+ with patch(
+ "application.vectorstore.base.BaseVectorStore._get_embeddings",
+ return_value=_Embeddings(9.0),
+ ):
+ reembed.reembed_pgvector(GRAPH_SOURCE, batch_size=64, dry_run=False)
+
+ after = _node_vectors(live_dsn, GRAPH_SOURCE)
+ assert [name for name, _ in after] == ["Alpha"]
+ assert all(v[0] == pytest.approx(9.0) for _, v in after), (
+ "graph node names must move with the chunk vectors"
+ )
+
+ def test_graph_is_left_alone_when_graphrag_is_off(self, live_dsn, monkeypatch):
+ from application.core import settings as settings_module
+
+ monkeypatch.setattr(
+ settings_module.settings, "GRAPHRAG_ENABLED", False, raising=False
+ )
+ _seed(live_dsn, GRAPH_SOURCE, TEXTS, _Embeddings(1.0))
+ _seed_graph_node(live_dsn, GRAPH_SOURCE, "Alpha", 1.0)
+
+ with patch(
+ "application.vectorstore.base.BaseVectorStore._get_embeddings",
+ return_value=_Embeddings(9.0),
+ ):
+ reembed.reembed_pgvector(GRAPH_SOURCE, batch_size=64, dry_run=False)
+
+ assert all(v[0] == pytest.approx(1.0) for _, v in _node_vectors(live_dsn, GRAPH_SOURCE))
+
+ def test_dry_run_leaves_node_vectors_untouched(self, live_dsn, monkeypatch):
+ from application.core import settings as settings_module
+
+ monkeypatch.setattr(
+ settings_module.settings, "GRAPHRAG_ENABLED", True, raising=False
+ )
+ _seed(live_dsn, GRAPH_SOURCE, TEXTS, _Embeddings(1.0))
+ _seed_graph_node(live_dsn, GRAPH_SOURCE, "Alpha", 1.0)
+
+ with patch(
+ "application.vectorstore.base.BaseVectorStore._get_embeddings",
+ return_value=_Embeddings(9.0),
+ ):
+ reembed.reembed_pgvector(GRAPH_SOURCE, batch_size=64, dry_run=True)
+
+ assert all(v[0] == pytest.approx(1.0) for _, v in _node_vectors(live_dsn, GRAPH_SOURCE))
diff --git a/tests/vectorstore/test_embeddings_local.py b/tests/vectorstore/test_embeddings_local.py
index 68e84dca..36b20754 100644
--- a/tests/vectorstore/test_embeddings_local.py
+++ b/tests/vectorstore/test_embeddings_local.py
@@ -25,10 +25,57 @@ def fake_fastembed():
instance = MagicMock()
instance.embed.return_value = iter([np.array([0.1, 0.2, 0.3])])
text_embedding.return_value = instance
+ # Registration checks this before calling ``add_custom_model``; an empty
+ # list means "no built-in collides", which is the case for every name in
+ # our registry.
+ text_embedding.list_supported_models.return_value = []
with patch("fastembed.TextEmbedding", text_embedding):
yield text_embedding, instance
+class TestBuiltinModelRegistration:
+ """FastEmbed ships ~30 models of its own and refuses to re-register any of
+ them, so registering unconditionally broke every natively-supported name."""
+
+ def test_builtin_name_is_not_re_registered(self, fake_fastembed):
+ text_embedding, _ = fake_fastembed
+ text_embedding.list_supported_models.return_value = [
+ {"model": "BAAI/bge-small-en-v1.5"}
+ ]
+ EmbeddingsWrapper("BAAI/bge-small-en-v1.5")
+ text_embedding.add_custom_model.assert_not_called()
+ assert text_embedding.call_args.kwargs["model_name"] == "BAAI/bge-small-en-v1.5"
+
+ def test_builtin_match_ignores_case(self, fake_fastembed):
+ text_embedding, _ = fake_fastembed
+ text_embedding.list_supported_models.return_value = [
+ {"model": "baai/BGE-Small-EN-v1.5"}
+ ]
+ EmbeddingsWrapper("BAAI/bge-small-en-v1.5")
+ text_embedding.add_custom_model.assert_not_called()
+
+ def test_unknown_name_is_still_registered(self, fake_fastembed):
+ text_embedding, _ = fake_fastembed
+ text_embedding.list_supported_models.return_value = [
+ {"model": "BAAI/bge-small-en-v1.5"}
+ ]
+ EmbeddingsWrapper("some-org/custom-embedder")
+ text_embedding.add_custom_model.assert_called_once()
+
+ def test_real_fastembed_accepts_its_own_builtin(self):
+ """Runs against the installed FastEmbed, not the MagicMock.
+
+ The mocked tests above cannot catch this: the failure was
+ ``add_custom_model`` raising, and a MagicMock never raises.
+ """
+ fastembed = pytest.importorskip("fastembed")
+ builtins = [m["model"] for m in fastembed.TextEmbedding.list_supported_models()]
+ assert builtins, "expected FastEmbed to ship built-in models"
+ spec = embeddings_local._spec_for(builtins[0])
+ # Must not raise ValueError("... is already registered ...").
+ embeddings_local._register(spec)
+
+
class TestRegistryDrivenLoading:
def test_registered_model_loads_by_repo_not_by_configured_name(self, fake_fastembed):
text_embedding, _ = fake_fastembed
diff --git a/tests/vectorstore/test_remote_embeddings_truncation.py b/tests/vectorstore/test_remote_embeddings_truncation.py
index 8b7e2d64..cb439a65 100644
--- a/tests/vectorstore/test_remote_embeddings_truncation.py
+++ b/tests/vectorstore/test_remote_embeddings_truncation.py
@@ -150,3 +150,48 @@ class TestInputLimitResolution:
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
+
+
+class TestEmbeddingsNameIsExplicit:
+ """Which names count as chosen decides whether a remote server inherits a
+ context window it may not have.
+
+ The tests above monkeypatch the predicate, so they cannot see it being
+ wrong. These drive the real one.
+ """
+
+ def _default(self):
+ from application.core.settings import Settings
+
+ return Settings.model_fields["EMBEDDINGS_NAME"].default
+
+ def test_the_default_name_is_not_a_choice(self, monkeypatch):
+ """Every setup script has always written this value unconditionally.
+
+ Reading it as deliberate lends the server mpnet's 384-token window and
+ clips ~80% off every chunk, silently, on upgrade. ``model_fields_set``
+ could not tell the difference: pydantic marks a field set for anything
+ that reached it, ``.env`` included.
+ """
+ monkeypatch.setattr(base.settings, "EMBEDDINGS_NAME", self._default())
+ assert base._embeddings_name_is_explicit() is False
+
+ def test_a_different_name_is_a_choice(self, monkeypatch):
+ monkeypatch.setattr(
+ base.settings,
+ "EMBEDDINGS_NAME",
+ "ibm-granite/granite-embedding-311m-multilingual-r2",
+ )
+ assert base._embeddings_name_is_explicit() is True
+
+ def test_dotenv_written_default_does_not_clip(self, monkeypatch):
+ """End to end: the common upgrade path must send the full chunk."""
+ monkeypatch.setattr(base.settings, "EMBEDDINGS_MAX_INPUT_TOKENS", None)
+ monkeypatch.setattr(base.settings, "EMBEDDINGS_NAME", self._default())
+ captured = _capture_post(monkeypatch)
+
+ long_text = " ".join(["word"] * 1000)
+ emb = RemoteEmbeddings(api_url="http://embeddings", model_name=self._default())
+ emb.embed_documents([long_text])
+
+ assert captured["payload"]["input"][0] == long_text