mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-14 22:17:00 +00:00
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:
1 parent
47e53a71c4
commit
3242d68fd2
6 files changed
+380
-4
No files matched your search
@@ -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
|
||||
@@ -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()
|
||||
Reference in new issue
Block a user