feat: pin the embedding model to the installation, not the release

EMBEDDINGS_NAME has a code-level default, and moving that default re-points an
existing index at a different vector space without anything noticing: mpnet and
granite are both 768-dimensional, so no width check fires and retrieval simply
gets worse. Which model an index was built with is a property of the
installation, not of the release it happens to be running.

So it is resolved once at boot and stored in app_metadata, which already exists
for exactly this kind of one-off state and needs no migration. An installation
that already has sources is pinned to the legacy model it has been using all
along and told, once, how to move; an empty one is pinned to the current
recommendation. An explicit EMBEDDINGS_NAME in the environment still wins over
both, and an unreachable database falls back to the code default.

This also closes a gap in what "new installs get granite" meant: it held for
setup.sh and for anyone copying .env-template, but a hand-written .env plus
docker compose fell through to the settings default and quietly got mpnet --
English-only, and a 384-token window against a 1250-token chunk default.

Emptiness is counted in sources rather than vector rows so the answer is the
same for every vector store, including the FAISS ones whose vectors are not in
this database at all. Concurrent workers converge through the same
INSERT ... ON CONFLICT DO NOTHING the instance id already uses.
This commit is contained in:
Alex committed 2026-08-28 13:14:35 +01:00
1 parent 47e53a71c4
commit 3242d68fd2
6 files changed
+380 -4

No files matched your search

+10
View File
@@ -35,6 +35,10 @@ from application.storage.db.bootstrap import ( # noqa: E402
ensure_database_ready,
ensure_vector_schema,
)
from application.storage.db.embeddings_pin import ( # noqa: E402
resolve_embeddings_pin,
warn_on_source_model_mismatch,
)
from application.stt.upload_limits import ( # noqa: E402
build_stt_file_size_limit_message,
should_reject_stt_request,
@@ -62,6 +66,12 @@ ensure_database_ready(
logger=logging.getLogger("application.app"),
)
# Which embedding model this installation uses is a property of its index, not of
# the release. Resolve it before the vector schema hook below, which sizes the
# table from EMBEDDINGS_NAME, and before anything embeds.
resolve_embeddings_pin(logging.getLogger("application.app"))
warn_on_source_model_mismatch(logging.getLogger("application.app"))
# Own the vector DB's schema here too, so the retrieval hot path is pure reads
# instead of re-running DDL for every source of every request.
if settings.AUTO_VECTOR_SCHEMA:
+13 -4
View File
@@ -142,12 +142,21 @@ def ensure_vector_schema(*, logger: Optional[logging.Logger] = None) -> None:
try:
from application.vectorstore.base import build_local_embeddings
dim = getattr(build_local_embeddings(), "dimension", None)
embedding = build_local_embeddings()
dim = getattr(embedding, "dimension", None)
if not dim:
# A remote client knows nothing about its server until it has
# called it, so ask once. Without this the table is sized at the
# default and the check below is skipped -- which is how a remote
# model of any other width silently got a vector(768) column, the
# exact failure this hook exists to catch. Milvus and Qdrant probe
# the same way.
dim = len(embedding.embed_query("dimension probe"))
except Exception as exc: # noqa: BLE001 — never block boot on the model
log.warning(
"ensure_vector_schema: could not load the embeddings model (%s); "
"creating the table with %d dimensions and skipping the dimension "
"check.",
"ensure_vector_schema: could not determine the embedding width "
"(%s); creating the table with %d dimensions and skipping the "
"dimension check.",
exc,
DEFAULT_EMBEDDING_DIM,
)
+150
View File
@@ -0,0 +1,150 @@
"""The embedding model an installation is pinned to.
``EMBEDDINGS_NAME`` has a code-level default, and moving that default would
re-point an existing index at a different vector space without anything
noticing: mpnet and granite are both 768-dimensional, so no width check fires
and retrieval simply gets worse. Which model an index was built with is a
property of the installation, not of the release it happens to be running.
So it is resolved once at boot and stored in ``app_metadata``. A fresh install
is pinned to the current recommendation; an install that already has sources is
pinned to the legacy model it has been using all along, and told how to move.
An explicit ``EMBEDDINGS_NAME`` in the environment always wins over both.
"""
from __future__ import annotations
import logging
from typing import Optional
from sqlalchemy import text
from application.core.settings import settings
from application.storage.db.repositories.app_metadata import AppMetadataRepository
from application.storage.db.session import db_session
from application.vectorstore.model_registry import (
DEFAULT_LEGACY,
DEFAULT_NEW_INSTALL,
resolve,
)
logger = logging.getLogger(__name__)
PIN_KEY = "embeddings_name"
NOTICE_KEY = "embeddings_legacy_notice_shown"
def _has_sources(conn) -> bool:
"""True when this installation already has an index to protect.
Counts ``sources`` rather than vector rows so the answer is the same for
every vector store, including the FAISS ones whose vectors are not in this
database at all.
"""
if conn.execute(text("SELECT to_regclass('public.sources')")).scalar() is None:
return False
return bool(conn.execute(text("SELECT EXISTS (SELECT 1 FROM sources)")).scalar())
def _legacy_notice(model: str) -> str:
return (
f"Embeddings: this installation is using {model}, the model its index was "
"built with, and will keep using it.\n"
"To move to granite (multilingual, a 32k-token context, same 768 dimensions):\n"
" 1. EMBEDDINGS_NAME=ibm-granite/granite-embedding-311m-multilingual-r2\n"
" 2. python -m application.scripts.reembed\n"
"Changing the model without step 2 leaves queries searching a different "
"vector space than the stored vectors, which fails silently."
)
def resolve_embeddings_pin(log: Optional[logging.Logger] = None) -> None:
"""Point ``settings.EMBEDDINGS_NAME`` at this installation's pinned model.
Called once at boot, before anything embeds and before the vector schema
hook reads the width. Consumers read ``settings.EMBEDDINGS_NAME`` lazily, so
assigning it here is enough; no caller needs to know the pin exists.
Does nothing when the environment pins the name, and degrades to the
code-level default when the database cannot be reached.
Args:
log: Logger for the first-run notice; module logger when omitted.
"""
out = log or logger
if "EMBEDDINGS_NAME" in settings.model_fields_set:
return
try:
with db_session() as conn:
repo = AppMetadataRepository(conn)
stored = repo.get(PIN_KEY)
if stored:
settings.EMBEDDINGS_NAME = stored
return
upgrading = _has_sources(conn)
pinned = repo.setdefault(
PIN_KEY, DEFAULT_LEGACY if upgrading else DEFAULT_NEW_INSTALL
)
settings.EMBEDDINGS_NAME = pinned
out.info("Embeddings: pinned this installation to %s.", pinned)
if upgrading and repo.get(NOTICE_KEY) is None:
print(_legacy_notice(pinned), flush=True)
repo.set(NOTICE_KEY, "1")
except Exception as exc: # noqa: BLE001 — never block boot on the database
out.debug(
"Embeddings: could not resolve the pinned model (%s); using %s.",
exc,
settings.EMBEDDINGS_NAME,
exc_info=True,
)
def warn_on_source_model_mismatch(log: Optional[logging.Logger] = None) -> None:
"""Report sources whose vectors were built by a different model.
The width check cannot see this: mpnet and granite are both 768, so an
index queried by the wrong model returns worse answers and no error. This
is the only signal, so it names the sources and the command that fixes
them.
A source with no recorded model pre-dates the column and is therefore the
legacy model, not unknown.
Args:
log: Logger to warn on; module logger when omitted.
"""
out = log or logger
active = resolve(settings.EMBEDDINGS_NAME)
try:
with db_session() as conn:
if conn.execute(text("SELECT to_regclass('public.sources')")).scalar() is None:
return
rows = conn.execute(
text("SELECT COALESCE(model, :legacy) AS model, count(*) FROM sources GROUP BY 1"),
{"legacy": DEFAULT_LEGACY},
).fetchall()
except Exception as exc: # noqa: BLE001 — never block boot on the database
out.debug("Embeddings: could not check source models (%s)", exc, exc_info=True)
return
stale = [
(name, count)
for name, count in rows
# Compare through the registry so an alias is not read as a different
# model. An unregistered name resolves to None, which only matches
# another unregistered name if the strings agree.
for a, b in [(resolve(name), active)]
if (a or name) != (b or settings.EMBEDDINGS_NAME)
]
if not stale:
return
detail = ", ".join(f"{count} built with {name}" for name, count in stale)
out.warning(
"Embeddings: %s, but queries are embedded with %s. Retrieval against those "
"sources is degraded and will not raise. Re-embed them with "
"`python -m application.scripts.reembed`, or set EMBEDDINGS_NAME back.",
detail,
settings.EMBEDDINGS_NAME,
)
@@ -37,6 +37,21 @@ class AppMetadataRepository:
{"key": key, "value": value},
)
def setdefault(self, key: str, value: str) -> str:
"""Store ``value`` under ``key`` only if absent, returning what stands.
Two workers racing on first startup converge on one value instead of
each overwriting the other, the way ``get_or_create_instance_id`` does.
"""
self._conn.execute(
text(
"INSERT INTO app_metadata (key, value) VALUES (:key, :value) "
"ON CONFLICT (key) DO NOTHING"
),
{"key": key, "value": value},
)
return self.get(key) or value
def get_or_create_instance_id(self) -> str:
"""Return the anonymous instance UUID, generating one if absent.
@@ -265,3 +265,64 @@ class TestBootDoesNotLoadTheModel:
vector_schema = self._run(None, loader)
loader.assert_called_once()
assert vector_schema.call_args.kwargs["dimension"] == 1024
@pytest.mark.unit
class TestUnknownWidthIsProbed:
"""A remote server's width is only knowable by asking it.
``RemoteEmbeddings`` reports ``None`` until its first call, so sizing the
table from the attribute alone fell back to 768 and skipped the mismatch
check — the silent ``vector(768)`` column this hook exists to prevent.
"""
def _run(self, remote, table_dimension=1024):
conn = MagicMock()
conn.cursor.return_value = MagicMock()
with patch("psycopg.connect", return_value=conn), patch(
"application.vectorstore.model_registry.dimension_for", return_value=None
), patch(
"application.vectorstore.base.build_local_embeddings", return_value=remote
), patch(
"application.vectorstore.pgvector.PGVectorStore.create_schema"
) as vector_schema, patch(
"application.vectorstore.pgvector.PGVectorStore.table_dimension",
return_value=table_dimension,
):
try:
ensure_vector_schema()
raised = False
except RuntimeError:
raised = True
return vector_schema, raised
@staticmethod
def _remote(width=None, error=None):
remote = MagicMock()
remote.dimension = None
remote.embed_query.side_effect = error or (lambda _text: [0.0] * width)
return remote
def test_the_table_is_sized_from_the_probe(self, vector_settings):
remote = self._remote(width=1024)
vector_schema, _ = self._run(remote, table_dimension=1024)
remote.embed_query.assert_called_once()
assert vector_schema.call_args.kwargs["dimension"] == 1024
def test_the_probe_restores_the_mismatch_check(self, vector_settings):
_, raised = self._run(self._remote(width=768), table_dimension=1024)
assert raised, "a 768-dim model against a vector(1024) table must fail loudly"
def test_an_unreachable_server_does_not_block_boot(self, vector_settings):
vector_schema, raised = self._run(
self._remote(error=ConnectionError("server down")), table_dimension=1024
)
assert not raised
assert vector_schema.call_args.kwargs["dimension"] == 768
def test_a_model_that_knows_its_width_is_not_probed(self, vector_settings):
local = MagicMock()
local.dimension = 384
vector_schema, _ = self._run(local, table_dimension=384)
local.embed_query.assert_not_called()
assert vector_schema.call_args.kwargs["dimension"] == 384
+131
View File
@@ -0,0 +1,131 @@
"""Which embedding model an installation is pinned to."""
from unittest.mock import MagicMock, patch
import pytest
from application.storage.db import embeddings_pin
from application.storage.db.embeddings_pin import (
NOTICE_KEY,
PIN_KEY,
resolve_embeddings_pin,
)
from application.vectorstore.model_registry import DEFAULT_LEGACY, DEFAULT_NEW_INSTALL
@pytest.fixture
def store():
"""An in-memory stand-in for the ``app_metadata`` key/value table."""
data = {}
repo = MagicMock()
repo.get.side_effect = data.get
repo.set.side_effect = lambda k, v: data.__setitem__(k, v)
repo.setdefault.side_effect = lambda k, v: data.setdefault(k, v)
repo._data = data
return repo
def _run(store, *, has_sources, env_pinned=False, name="unset"):
session = MagicMock()
session.__enter__ = MagicMock(return_value=MagicMock())
session.__exit__ = MagicMock(return_value=False)
fields = {"EMBEDDINGS_NAME"} if env_pinned else set()
with patch.object(embeddings_pin.settings, "EMBEDDINGS_NAME", name), patch.object(
type(embeddings_pin.settings), "model_fields_set", property(lambda self: fields)
), patch.object(embeddings_pin, "db_session", return_value=session), patch.object(
embeddings_pin, "AppMetadataRepository", return_value=store
), patch.object(
embeddings_pin, "_has_sources", return_value=has_sources
):
resolve_embeddings_pin()
return embeddings_pin.settings.EMBEDDINGS_NAME
class TestFreshInstall:
def test_an_empty_installation_is_pinned_to_the_current_model(self, store):
assert _run(store, has_sources=False) == DEFAULT_NEW_INSTALL
assert store._data[PIN_KEY] == DEFAULT_NEW_INSTALL
def test_no_legacy_notice_is_printed(self, store, capsys):
_run(store, has_sources=False)
assert "reembed" not in capsys.readouterr().out
assert NOTICE_KEY not in store._data
class TestExistingInstall:
"""An index built by the old model must keep being read by it."""
def test_an_installation_with_sources_is_pinned_to_the_legacy_model(self, store):
assert _run(store, has_sources=True) == DEFAULT_LEGACY
assert store._data[PIN_KEY] == DEFAULT_LEGACY
def test_the_notice_names_the_migration_command(self, store, capsys):
_run(store, has_sources=True)
out = capsys.readouterr().out
assert "application.scripts.reembed" in out
assert DEFAULT_NEW_INSTALL in out
def test_the_notice_is_shown_only_once(self, store, capsys):
_run(store, has_sources=True)
capsys.readouterr()
store._data.pop(PIN_KEY) # force the decision again
_run(store, has_sources=True)
assert "reembed" not in capsys.readouterr().out
class TestPrecedence:
def test_a_stored_pin_survives_new_sources(self, store):
store._data[PIN_KEY] = DEFAULT_NEW_INSTALL
assert _run(store, has_sources=True) == DEFAULT_NEW_INSTALL
def test_the_environment_wins_and_nothing_is_stored(self, store):
assert _run(store, has_sources=True, env_pinned=True, name="my/model") == "my/model"
assert store._data == {}
def test_an_unreachable_database_leaves_the_default_alone(self):
with patch.object(embeddings_pin.settings, "EMBEDDINGS_NAME", "fallback"), patch.object(
type(embeddings_pin.settings), "model_fields_set", property(lambda self: set())
), patch.object(embeddings_pin, "db_session", side_effect=OSError("no db")):
resolve_embeddings_pin()
assert embeddings_pin.settings.EMBEDDINGS_NAME == "fallback"
class TestSourceModelMismatch:
"""The only signal that an index is being queried by the wrong model."""
def _run(self, rows, active, log):
conn = MagicMock()
conn.execute.side_effect = [
MagicMock(scalar=MagicMock(return_value="sources")),
MagicMock(fetchall=MagicMock(return_value=rows)),
]
session = MagicMock()
session.__enter__ = MagicMock(return_value=conn)
session.__exit__ = MagicMock(return_value=False)
with patch.object(embeddings_pin.settings, "EMBEDDINGS_NAME", active), patch.object(
embeddings_pin, "db_session", return_value=session
):
embeddings_pin.warn_on_source_model_mismatch(log)
def test_a_different_model_is_reported_with_counts(self):
log = MagicMock()
self._run([(DEFAULT_LEGACY, 28)], DEFAULT_NEW_INSTALL, log)
message = log.warning.call_args.args[0] % log.warning.call_args.args[1:]
assert "28 built with" in message
assert "application.scripts.reembed" in message
def test_an_alias_is_not_a_mismatch(self):
"""A stored alias and the canonical name are the same model."""
log = MagicMock()
self._run([("sentence-transformers/all-mpnet-base-v2", 28)], DEFAULT_LEGACY, log)
log.warning.assert_not_called()
def test_a_matching_model_says_nothing(self):
log = MagicMock()
self._run([(DEFAULT_NEW_INSTALL, 5)], DEFAULT_NEW_INSTALL, log)
log.warning.assert_not_called()
def test_two_unregistered_names_that_differ_are_a_mismatch(self):
log = MagicMock()
self._run([("some/other-model", 3)], "my/custom-model", log)
log.warning.assert_called_once()