From a090bca96b6df27eb84f3acb5be3c8b11a2bacdd Mon Sep 17 00:00:00 2001 From: medlouaynjima Date: Wed, 19 Aug 2026 17:20:42 +0100 Subject: [PATCH 1/3] fix: validate embedding model dimensions --- README.md | 16 +++++++- tests/test_embeddings.py | 55 +++++++++++++++++++++++++ vaultrag/embeddings.py | 86 +++++++++++++++++++++++++++------------- 3 files changed, 127 insertions(+), 30 deletions(-) create mode 100644 tests/test_embeddings.py diff --git a/README.md b/README.md index 8ffdd39..23d6b2f 100644 --- a/README.md +++ b/README.md @@ -264,8 +264,20 @@ EMBEDDER=fake LLM=fake ./.venv/bin/python -m uvicorn vaultrag.main:app --reload Tests need no API keys: they use a deterministic offline embedder and a scripted LLM, because access control is not a semantic question and shouldn't need a 90MB model download to verify. -For real use, set `EMBEDDER=local` (sentence-transformers, free, no key) and `LLM=groq` with a -`GROQ_API_KEY`. Zero cost either way. +For real use, set `EMBEDDER=local` (sentence-transformers, free, no key) and +`LLM=groq` with a `GROQ_API_KEY`. Zero cost either way. + +### Changing the embedding model + +`EMBED_MODEL` defaults to `sentence-transformers/all-MiniLM-L6-v2`, which produces +384-dimensional embeddings. + +If you change `EMBED_MODEL`, the model's embedding dimension must match the +`vector(384)` column in `vaultrag/schema.sql`. Changing to a model with a different +dimension requires updating the database schema and re-ingesting the corpus. + +VaultRAG validates the embedding dimension when the local model is loaded and raises +a clear error if it does not match the database schema. ## Status diff --git a/tests/test_embeddings.py b/tests/test_embeddings.py new file mode 100644 index 0000000..2b8ad7e --- /dev/null +++ b/tests/test_embeddings.py @@ -0,0 +1,55 @@ +from types import SimpleNamespace + +import pytest + +from vaultrag.embeddings import DIMS, LocalEmbedder + + +class FakeSentenceTransformer: + def __init__(self, model_name: str, dims: int) -> None: + self.model_name = model_name + self._dims = dims + + def get_sentence_embedding_dimension(self) -> int: + return self._dims + + +def install_fake_sentence_transformers(monkeypatch, dims: int): + def constructor(model_name: str): + return FakeSentenceTransformer(model_name, dims) + + monkeypatch.setitem( + __import__("sys").modules, + "sentence_transformers", + SimpleNamespace(SentenceTransformer=constructor), + ) + + +def test_local_embedder_detects_model_dimensions(monkeypatch): + install_fake_sentence_transformers(monkeypatch, DIMS) + + embedder = LocalEmbedder("fake-model") + + assert embedder.dims == DIMS + + +def test_local_embedder_rejects_wrong_dimensions(monkeypatch): + install_fake_sentence_transformers(monkeypatch, 768) + + model_name = "sentence-transformers/all-mpnet-base-v2" + embedder = LocalEmbedder(model_name) + + with pytest.raises( + ValueError, + match=r"all-mpnet-base-v2.*768.*384.*vector\(384\)", + ): + embedder.dims + + +def test_local_embedder_error_mentions_schema(monkeypatch): + install_fake_sentence_transformers(monkeypatch, 1024) + + embedder = LocalEmbedder("another-model") + + with pytest.raises(ValueError, match=r"vector\(384\)"): + embedder.dims \ No newline at end of file diff --git a/vaultrag/embeddings.py b/vaultrag/embeddings.py index 71afb74..223fbbc 100644 --- a/vaultrag/embeddings.py +++ b/vaultrag/embeddings.py @@ -2,14 +2,11 @@ Two implementations, one interface: -- LocalEmbedder: sentence-transformers, runs on CPU, no API key, free forever. This is the - default and the one that matters. Paying an API per embedding for a document corpus is a - choice, not a requirement. -- FakeEmbedder: deterministic hash-based vectors. Not semantically meaningful, but stable and - instant, which is exactly what the ACL tests need. Those tests are about who can see what, and - they should not need a 90MB model download or a network call to run. - -The interface is one method so swapping providers is a config change, not a refactor. +- LocalEmbedder: sentence-transformers, runs on CPU, no API key, free forever. +- FakeEmbedder: deterministic hash-based vectors. Used for tests. + +The local embedder detects the actual model dimension and verifies that it +matches the database schema dimension. """ from __future__ import annotations @@ -18,21 +15,19 @@ import math from typing import Protocol -DIMS = 384 # all-MiniLM-L6-v2. If you change the model, change the schema's vector(384) too. + +DIMS = 384 class Embedder(Protocol): + @property + def dims(self) -> int: ... + def embed(self, texts: list[str]) -> list[list[float]]: ... class FakeEmbedder: - """Deterministic, offline, instant. For tests. - - Hashes text into a fixed-dimension unit vector. Same text always gives the same vector, and - different text gives a different one, which is all a retrieval test needs. It is NOT semantic: - "cat" and "kitten" are unrelated here. Any test that depends on semantic similarity should use - the real embedder or, better, not be a unit test. - """ + """Deterministic, offline, instant embedder for tests.""" dims = DIMS @@ -40,41 +35,69 @@ def embed(self, texts: list[str]) -> list[list[float]]: return [self._one(t) for t in texts] def _one(self, text: str) -> list[float]: - # Expand a digest into DIMS floats by rehashing with a counter. vec: list[float] = [] counter = 0 + while len(vec) < DIMS: h = hashlib.sha256(f"{text}:{counter}".encode()).digest() vec.extend(b / 255.0 - 0.5 for b in h) counter += 1 + vec = vec[:DIMS] return _normalize(vec) class LocalEmbedder: - """sentence-transformers on CPU. Free, no key, no network after the first download. + """sentence-transformers on CPU. - Loaded lazily so that importing this module (which the tests do) does not pull in torch. + The model is loaded lazily and its actual embedding dimension is checked + against the dimension expected by the database schema. """ - dims = DIMS - - def __init__(self, model_name: str = "sentence-transformers/all-MiniLM-L6-v2") -> None: + def __init__( + self, + model_name: str = "sentence-transformers/all-MiniLM-L6-v2", + ) -> None: self._model_name = model_name self._model = None + self._dims: int | None = None + + @property + def dims(self) -> int: + if self._dims is None: + self._load() + + assert self._dims is not None + return self._dims def _load(self): if self._model is None: - from sentence_transformers import SentenceTransformer # imported lazily on purpose + from sentence_transformers import SentenceTransformer self._model = SentenceTransformer(self._model_name) + self._dims = self._model.get_sentence_embedding_dimension() + + if self._dims != DIMS: + raise ValueError( + f"Embedding model {self._model_name!r} produces " + f"{self._dims}-dimensional vectors, but VaultRAG expects " + f"{DIMS} dimensions. The embedding column in " + f"vaultrag/schema.sql is vector({DIMS}). Changing " + f"EMBED_MODEL requires updating the schema and " + f"re-ingesting the corpus." + ) + return self._model def embed(self, texts: list[str]) -> list[list[float]]: model = self._load() - # normalize_embeddings=True so cosine distance in pgvector behaves, and so the schema's - # vector_cosine_ops index is the right choice. - arr = model.encode(texts, normalize_embeddings=True, show_progress_bar=False) + + arr = model.encode( + texts, + normalize_embeddings=True, + show_progress_bar=False, + ) + return [list(map(float, row)) for row in arr] @@ -83,9 +106,16 @@ def _normalize(vec: list[float]) -> list[float]: return [v / norm for v in vec] -def get_embedder(kind: str = "local", model: str = "sentence-transformers/all-MiniLM-L6-v2") -> Embedder: +def get_embedder( + kind: str = "local", + model: str = "sentence-transformers/all-MiniLM-L6-v2", +) -> Embedder: if kind == "fake": return FakeEmbedder() + if kind == "local": return LocalEmbedder(model) - raise ValueError(f"unknown embedder: {kind!r} (expected 'local' or 'fake')") + + raise ValueError( + f"unknown embedder: {kind!r} (expected 'local' or 'fake')" + ) \ No newline at end of file From 11c0444b7a610df81c011f210e5a92adcee31d70 Mon Sep 17 00:00:00 2001 From: medlouaynjima Date: Sat, 22 Aug 2026 15:13:12 +0100 Subject: [PATCH 2/3] fix: validate embedding dimensions before caching --- tests/test_embeddings.py | 16 ++++++++-------- vaultrag/embeddings.py | 11 ++++++----- 2 files changed, 14 insertions(+), 13 deletions(-) diff --git a/tests/test_embeddings.py b/tests/test_embeddings.py index 2b8ad7e..3568162 100644 --- a/tests/test_embeddings.py +++ b/tests/test_embeddings.py @@ -33,17 +33,17 @@ def test_local_embedder_detects_model_dimensions(monkeypatch): assert embedder.dims == DIMS -def test_local_embedder_rejects_wrong_dimensions(monkeypatch): +def test_local_embedder_rejects_wrong_dimensions_on_repeated_access(monkeypatch): install_fake_sentence_transformers(monkeypatch, 768) - model_name = "sentence-transformers/all-mpnet-base-v2" - embedder = LocalEmbedder(model_name) + embedder = LocalEmbedder("sentence-transformers/all-mpnet-base-v2") - with pytest.raises( - ValueError, - match=r"all-mpnet-base-v2.*768.*384.*vector\(384\)", - ): - embedder.dims + for _ in range(2): + with pytest.raises( + ValueError, + match=r"all-mpnet-base-v2.*768.*384.*vector\(384\)", + ): + embedder.dims def test_local_embedder_error_mentions_schema(monkeypatch): diff --git a/vaultrag/embeddings.py b/vaultrag/embeddings.py index 223fbbc..161c190 100644 --- a/vaultrag/embeddings.py +++ b/vaultrag/embeddings.py @@ -74,20 +74,21 @@ def _load(self): if self._model is None: from sentence_transformers import SentenceTransformer - self._model = SentenceTransformer(self._model_name) - self._dims = self._model.get_sentence_embedding_dimension() + model = SentenceTransformer(self._model_name) + dims = model.get_sentence_embedding_dimension() - if self._dims != DIMS: + if dims != DIMS: raise ValueError( f"Embedding model {self._model_name!r} produces " - f"{self._dims}-dimensional vectors, but VaultRAG expects " + f"{dims}-dimensional vectors, but VaultRAG expects " f"{DIMS} dimensions. The embedding column in " f"vaultrag/schema.sql is vector({DIMS}). Changing " f"EMBED_MODEL requires updating the schema and " f"re-ingesting the corpus." ) - return self._model + self._model = model + self._dims = dims def embed(self, texts: list[str]) -> list[list[float]]: model = self._load() From ce10d323b435102cbab8aed57e1869a9b54134ac Mon Sep 17 00:00:00 2001 From: medlouaynjima Date: Mon, 24 Aug 2026 10:09:10 +0100 Subject: [PATCH 3/3] fix: validate embedding dimensions before caching --- tests/test_embeddings.py | 19 +++++++++++++++++++ vaultrag/embeddings.py | 3 +++ 2 files changed, 22 insertions(+) diff --git a/tests/test_embeddings.py b/tests/test_embeddings.py index 3568162..18a5c5e 100644 --- a/tests/test_embeddings.py +++ b/tests/test_embeddings.py @@ -13,6 +13,14 @@ def __init__(self, model_name: str, dims: int) -> None: def get_sentence_embedding_dimension(self) -> int: return self._dims + def encode( + self, + texts: list[str], + normalize_embeddings: bool, + show_progress_bar: bool, + ) -> list[list[float]]: + return [[0.0] * self._dims for _ in texts] + def install_fake_sentence_transformers(monkeypatch, dims: int): def constructor(model_name: str): @@ -33,6 +41,17 @@ def test_local_embedder_detects_model_dimensions(monkeypatch): assert embedder.dims == DIMS +def test_local_embedder_embed_works_with_valid_dimensions(monkeypatch): + install_fake_sentence_transformers(monkeypatch, DIMS) + + embedder = LocalEmbedder("fake-model") + + result = embedder.embed(["hello", "world"]) + + assert len(result) == 2 + assert all(len(vector) == DIMS for vector in result) + + def test_local_embedder_rejects_wrong_dimensions_on_repeated_access(monkeypatch): install_fake_sentence_transformers(monkeypatch, 768) diff --git a/vaultrag/embeddings.py b/vaultrag/embeddings.py index 161c190..71f478b 100644 --- a/vaultrag/embeddings.py +++ b/vaultrag/embeddings.py @@ -89,6 +89,9 @@ def _load(self): self._model = model self._dims = dims + + return self._model + def embed(self, texts: list[str]) -> list[list[float]]: model = self._load()