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()