mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-14 18:18:08 +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
@@ -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:
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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.
|
||||
|
||||
|
||||
Reference in new issue
Block a user