"""Which embedding model an installation is pinned to.""" from unittest.mock import MagicMock, patch import pytest from docsgpt.storage.db import embeddings_pin from docsgpt.storage.db.embeddings_pin import ( NOTICE_KEY, PIN_KEY, resolve_embeddings_pin, ) from docsgpt.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 "docsgpt.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 "docsgpt.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()