From 3242d68fd2a7dfaf23dbcd94cf5f675a8e089a23 Mon Sep 17 00:00:00 2001 From: Alex Date: Fri, 28 Aug 2026 13:14:35 +0100 Subject: [PATCH] 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. --- application/app.py | 10 ++ application/storage/db/bootstrap.py | 17 +- application/storage/db/embeddings_pin.py | 150 ++++++++++++++++++ .../storage/db/repositories/app_metadata.py | 15 ++ .../db/test_bootstrap_vector_schema.py | 61 +++++++ tests/storage/db/test_embeddings_pin.py | 131 +++++++++++++++ 6 files changed, 380 insertions(+), 4 deletions(-) create mode 100644 application/storage/db/embeddings_pin.py create mode 100644 tests/storage/db/test_embeddings_pin.py diff --git a/application/app.py b/application/app.py index 466939c1..b412b31e 100644 --- a/application/app.py +++ b/application/app.py @@ -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: diff --git a/application/storage/db/bootstrap.py b/application/storage/db/bootstrap.py index 758f5227..48dc6264 100644 --- a/application/storage/db/bootstrap.py +++ b/application/storage/db/bootstrap.py @@ -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, ) diff --git a/application/storage/db/embeddings_pin.py b/application/storage/db/embeddings_pin.py new file mode 100644 index 00000000..c0a17652 --- /dev/null +++ b/application/storage/db/embeddings_pin.py @@ -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, + ) diff --git a/application/storage/db/repositories/app_metadata.py b/application/storage/db/repositories/app_metadata.py index c2d8e1a8..ea1b4645 100644 --- a/application/storage/db/repositories/app_metadata.py +++ b/application/storage/db/repositories/app_metadata.py @@ -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. diff --git a/tests/storage/db/test_bootstrap_vector_schema.py b/tests/storage/db/test_bootstrap_vector_schema.py index 64d50f3c..e9761b03 100644 --- a/tests/storage/db/test_bootstrap_vector_schema.py +++ b/tests/storage/db/test_bootstrap_vector_schema.py @@ -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 diff --git a/tests/storage/db/test_embeddings_pin.py b/tests/storage/db/test_embeddings_pin.py new file mode 100644 index 00000000..c2275c28 --- /dev/null +++ b/tests/storage/db/test_embeddings_pin.py @@ -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()