feat(embeddings): every vector records the model whose space it lives in (#4132)
CI & Build / Python lint (push) Successful in 3s
CI & Build / Plugin hooks (push) Successful in 14s
CI & Build / TypeScript typecheck (push) Successful in 54s
CI & Build / integration (push) Successful in 1m1s
CI & Build / Python tests (push) Successful in 1m40s
CI & Build / Build & push image (push) Canceled after 27s
CI & Build / Python lint (push) Successful in 3s
CI & Build / Plugin hooks (push) Successful in 14s
CI & Build / TypeScript typecheck (push) Successful in 54s
CI & Build / integration (push) Successful in 1m1s
CI & Build / Python tests (push) Successful in 1m40s
CI & Build / Build & push image (push) Canceled after 27s
The four embedding tables stamped chunker_version but not the model, and vector(384) is a width, not an identity: a same-width model swap would write a second geometry beside the first with no error. - embedding_model on note/rule/milestone/system embeddings (0109; existing rows stamped with the only model any install has ever run). - Every write stamps EMBEDDING_MODEL; every backfill's "current" test is is_current_stamp(), both halves of calibration_stamp(). - migrate_floor refuses while any row its surface searches is off the live model, before sampling: re-embed, then migrate. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,42 @@
|
||||
"""embeddings_record_their_model — a vector says whose space it lives in (#4132)
|
||||
|
||||
Revision ID: 0109
|
||||
Revises: 0108
|
||||
Create Date: 2026-09-23
|
||||
|
||||
Every embedding table stamps `chunker_version`, so a change to the document
|
||||
shape is caught per row and re-embedded. None stamped the MODEL. The column is
|
||||
`vector(384)` — a width, not an identity — so swapping bge-small for any other
|
||||
384-dim model would write a second geometry beside the first with no error,
|
||||
and search would go on ranking by cosines between the two, which mean nothing.
|
||||
|
||||
`embedding_model` is the other half of `calibration_stamp()`, stored per row on
|
||||
all four tables. The startup backfill now re-embeds on either half moving.
|
||||
|
||||
THE BACKFILL LITERAL IS SAFE BECAUSE NO INSTALL HAS EVER CHANGED MODELS. The
|
||||
name is hardcoded in `services/embeddings.py`, and every vector ever written was
|
||||
written by it — so stamping every existing row with that name states a fact,
|
||||
not a guess. It is frozen here rather than imported: a migration records what
|
||||
was true when it ran, and a later model change must not rewrite history.
|
||||
"""
|
||||
from alembic import op
|
||||
|
||||
revision = "0109"
|
||||
down_revision = "0108"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
_TABLES = ("note_embeddings", "rule_embeddings", "milestone_embeddings", "system_embeddings")
|
||||
_MODEL = "BAAI/bge-small-en-v1.5"
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
for table in _TABLES:
|
||||
op.execute(f"ALTER TABLE {table} ADD COLUMN embedding_model text")
|
||||
op.execute(f"UPDATE {table} SET embedding_model = '{_MODEL}'")
|
||||
op.execute(f"ALTER TABLE {table} ALTER COLUMN embedding_model SET NOT NULL")
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
for table in _TABLES:
|
||||
op.execute(f"ALTER TABLE {table} DROP COLUMN embedding_model")
|
||||
@@ -41,6 +41,14 @@ class NoteEmbedding(Base):
|
||||
# any note whose rows carry a stale version — shape changes become a
|
||||
# version bump instead of a table wipe.
|
||||
chunker_version: Mapped[int] = mapped_column(Integer, nullable=False)
|
||||
# embeddings.EMBEDDING_MODEL at write time — the SPACE the vector lives in
|
||||
# (#4132). The column is `vector(384)`, a width and not an identity, so a
|
||||
# swap to another 384-dim model writes a second geometry beside the first
|
||||
# with no error, and cosine across the two is a number that means nothing.
|
||||
# The version above says what text was embedded; this says in whose
|
||||
# geometry. Either one moving makes the row stale, and the backfill
|
||||
# re-embeds on both.
|
||||
embedding_model: Mapped[str] = mapped_column(Text, nullable=False)
|
||||
updated_at: Mapped[datetime] = mapped_column(
|
||||
DateTime(timezone=True),
|
||||
default=lambda: datetime.now(timezone.utc),
|
||||
@@ -91,6 +99,7 @@ class RuleEmbedding(Base):
|
||||
# centroid (measured in note 2485).
|
||||
chunk_text: Mapped[str] = mapped_column(Text, nullable=False)
|
||||
chunker_version: Mapped[int] = mapped_column(Integer, nullable=False)
|
||||
embedding_model: Mapped[str] = mapped_column(Text, nullable=False) # see NoteEmbedding
|
||||
updated_at: Mapped[datetime] = mapped_column(
|
||||
DateTime(timezone=True),
|
||||
default=lambda: datetime.now(timezone.utc),
|
||||
@@ -124,6 +133,7 @@ class MilestoneEmbedding(Base):
|
||||
embedding: Mapped[list] = mapped_column(Vector(EMBEDDING_DIM), nullable=False)
|
||||
chunk_text: Mapped[str] = mapped_column(Text, nullable=False)
|
||||
chunker_version: Mapped[int] = mapped_column(Integer, nullable=False)
|
||||
embedding_model: Mapped[str] = mapped_column(Text, nullable=False) # see NoteEmbedding
|
||||
updated_at: Mapped[datetime] = mapped_column(
|
||||
DateTime(timezone=True),
|
||||
default=lambda: datetime.now(timezone.utc),
|
||||
@@ -163,6 +173,7 @@ class SystemEmbedding(Base):
|
||||
embedding: Mapped[list] = mapped_column(Vector(EMBEDDING_DIM), nullable=False)
|
||||
chunk_text: Mapped[str] = mapped_column(Text, nullable=False)
|
||||
chunker_version: Mapped[int] = mapped_column(Integer, nullable=False)
|
||||
embedding_model: Mapped[str] = mapped_column(Text, nullable=False) # see NoteEmbedding
|
||||
updated_at: Mapped[datetime] = mapped_column(
|
||||
DateTime(timezone=True),
|
||||
default=lambda: datetime.now(timezone.utc),
|
||||
|
||||
@@ -18,7 +18,7 @@ from collections.abc import Sequence
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from sqlalchemy import delete, func, or_, select
|
||||
from sqlalchemy import and_, delete, func, or_, select
|
||||
|
||||
from scribe.models import async_session
|
||||
from scribe.models.embedding import NoteEmbedding, RuleEmbedding
|
||||
@@ -367,6 +367,34 @@ def calibration_stamp() -> dict:
|
||||
"""
|
||||
return {"embedding_model": EMBEDDING_MODEL, "shape_version": CHUNKER_VERSION}
|
||||
|
||||
|
||||
def is_current_stamp(table):
|
||||
"""The rows written by THIS chunker in THIS model's space (#4132).
|
||||
|
||||
Every embedding table stores both halves of `calibration_stamp()` per row,
|
||||
and a row is current only when both match. One predicate for all four
|
||||
tables, because a backfill that checked the version alone is exactly how a
|
||||
same-width model swap would have gone unnoticed.
|
||||
"""
|
||||
return and_(
|
||||
table.chunker_version == CHUNKER_VERSION,
|
||||
table.embedding_model == EMBEDDING_MODEL,
|
||||
)
|
||||
|
||||
|
||||
async def rows_off_the_live_model(table) -> int:
|
||||
"""How many rows of an embedding table were NOT written in the live space.
|
||||
|
||||
Nonzero means the corpus is part-way through a model change: a search over
|
||||
it compares vectors from two geometries, and a statistic computed from it
|
||||
(`retrieval_migration.migrate_floor`) is a blend of both.
|
||||
"""
|
||||
async with async_session() as session:
|
||||
return int((await session.execute(
|
||||
select(func.count()).select_from(table)
|
||||
.where(table.embedding_model != EMBEDDING_MODEL)
|
||||
)).scalar_one())
|
||||
|
||||
# Character budget approximating the model window. Tokens-per-char varies by
|
||||
# content — ~4 chars/token for prose, closer to 3 for code and tables — so 1400
|
||||
# chars sits at roughly 350-470 tokens, leaving headroom for the title prefixed
|
||||
@@ -633,6 +661,7 @@ async def upsert_note_embedding(
|
||||
embedding=vector,
|
||||
chunk_text=chunk,
|
||||
chunker_version=CHUNKER_VERSION,
|
||||
embedding_model=EMBEDDING_MODEL,
|
||||
)
|
||||
)
|
||||
await session.commit()
|
||||
@@ -1092,7 +1121,7 @@ async def backfill_note_embeddings() -> None:
|
||||
for row in (
|
||||
await session.execute(
|
||||
select(NoteEmbedding.note_id).where(
|
||||
NoteEmbedding.chunker_version == CHUNKER_VERSION
|
||||
is_current_stamp(NoteEmbedding)
|
||||
)
|
||||
)
|
||||
).fetchall()
|
||||
@@ -1294,6 +1323,7 @@ async def upsert_rule_embedding(
|
||||
embedding=vector,
|
||||
chunk_text=chunk,
|
||||
chunker_version=CHUNKER_VERSION,
|
||||
embedding_model=EMBEDDING_MODEL,
|
||||
)
|
||||
)
|
||||
await session.commit()
|
||||
@@ -1464,7 +1494,7 @@ async def backfill_rule_embeddings() -> None:
|
||||
try:
|
||||
async with async_session() as session:
|
||||
current = select(RuleEmbedding.rule_id).where(
|
||||
RuleEmbedding.chunker_version == CHUNKER_VERSION
|
||||
is_current_stamp(RuleEmbedding)
|
||||
)
|
||||
# IDS ONLY — the text is re-read per rule below (#4262).
|
||||
by_version = {
|
||||
@@ -1563,6 +1593,7 @@ async def upsert_milestone_embedding(
|
||||
session.add(MilestoneEmbedding(
|
||||
milestone_id=milestone_id, chunk_index=index, embedding=vector,
|
||||
chunk_text=chunk, chunker_version=CHUNKER_VERSION,
|
||||
embedding_model=EMBEDDING_MODEL,
|
||||
))
|
||||
await session.commit()
|
||||
except Exception:
|
||||
@@ -1725,6 +1756,7 @@ async def upsert_system_embedding(
|
||||
session.add(SystemEmbedding(
|
||||
system_id=system_id, chunk_index=index, embedding=vector,
|
||||
chunk_text=chunk, chunker_version=CHUNKER_VERSION,
|
||||
embedding_model=EMBEDDING_MODEL,
|
||||
))
|
||||
await session.commit()
|
||||
except Exception:
|
||||
@@ -1841,7 +1873,7 @@ async def backfill_system_embeddings() -> None:
|
||||
try:
|
||||
async with async_session() as session:
|
||||
current = select(SystemEmbedding.system_id).where(
|
||||
SystemEmbedding.chunker_version == CHUNKER_VERSION
|
||||
is_current_stamp(SystemEmbedding)
|
||||
)
|
||||
# IDS ONLY — the charter is re-read per System below (#4262).
|
||||
by_version = {
|
||||
@@ -1884,7 +1916,7 @@ async def backfill_milestone_embeddings() -> None:
|
||||
try:
|
||||
async with async_session() as session:
|
||||
current = select(MilestoneEmbedding.milestone_id).where(
|
||||
MilestoneEmbedding.chunker_version == CHUNKER_VERSION
|
||||
is_current_stamp(MilestoneEmbedding)
|
||||
)
|
||||
# IDS ONLY — the plan is re-read per milestone below (#4262).
|
||||
by_version = {
|
||||
|
||||
@@ -58,9 +58,11 @@ import logging
|
||||
from sqlalchemy import select
|
||||
|
||||
from scribe.models import async_session
|
||||
from scribe.models.embedding import NoteEmbedding, RuleEmbedding
|
||||
from scribe.models.retrieval_log import RetrievalLog
|
||||
from scribe.services.embeddings import (
|
||||
calibration_stamp,
|
||||
rows_off_the_live_model,
|
||||
semantic_search_notes,
|
||||
semantic_search_rules,
|
||||
)
|
||||
@@ -117,6 +119,19 @@ _RESCORERS = {
|
||||
"report_preference": lambda u, q, p: _rescore_rules(u, q, p, "preference"),
|
||||
}
|
||||
|
||||
# The embedding table each surface's re-scorer reads. A migration from a corpus
|
||||
# that is still part-way through a model change re-scores against a blend of
|
||||
# two geometries and writes a floor computed from it, with a confident reason
|
||||
# attached (#4132) — so the corpus has to be wholly in the live space first.
|
||||
_CORPUS = {
|
||||
"auto_inject": NoteEmbedding,
|
||||
"write_path": NoteEmbedding,
|
||||
"write_path_rule": RuleEmbedding,
|
||||
"pre_tool_rule": RuleEmbedding,
|
||||
"prompt_rule": RuleEmbedding,
|
||||
"report_preference": RuleEmbedding,
|
||||
}
|
||||
|
||||
|
||||
def _floor_admitting(scores: list[float], fraction: float) -> float:
|
||||
"""The floor that admits `fraction` of `scores`, on this scale.
|
||||
@@ -161,6 +176,21 @@ async def migrate_floor(
|
||||
"arm's own corpus filters."
|
||||
)
|
||||
|
||||
off_model = await rows_off_the_live_model(_CORPUS[surface])
|
||||
if off_model:
|
||||
# Refused before sampling anything. The order of operations — re-embed,
|
||||
# THEN migrate — used to be held only in the operator's memory.
|
||||
stamp = calibration_stamp()
|
||||
return {
|
||||
"surface": surface, "migrated": False,
|
||||
"why": f"{off_model} embedding row(s) this surface searches are not "
|
||||
f"yet in {stamp['embedding_model']}'s space. Re-scoring now "
|
||||
"would measure a blend of two models; let the startup "
|
||||
"backfill finish re-embedding, then migrate",
|
||||
"rows_off_model": off_model,
|
||||
"calibration": stamp,
|
||||
}
|
||||
|
||||
old_floor = await floor_for(user_id, surface)
|
||||
async with async_session() as session:
|
||||
rows = (await session.execute(
|
||||
|
||||
@@ -228,6 +228,7 @@ async def test_upsert_stores_one_versioned_row_per_chunk():
|
||||
assert [r.chunk_index for r in rows] == list(range(len(chunks)))
|
||||
assert [r.chunk_text for r in rows] == chunks
|
||||
assert {r.chunker_version for r in rows} == {emb.CHUNKER_VERSION}
|
||||
assert {r.embedding_model for r in rows} == {emb.EMBEDDING_MODEL}
|
||||
assert {r.user_id for r in rows} == {42}
|
||||
session.execute.assert_awaited() # the delete that makes replacement atomic
|
||||
|
||||
@@ -314,3 +315,17 @@ async def test_backfill_reembeds_notes_with_a_stale_chunker_version():
|
||||
|
||||
embedded = [call.args[0] for call in upsert.call_args_list]
|
||||
assert embedded == [2], "only the stale note is re-embedded"
|
||||
|
||||
|
||||
def test_a_row_is_current_only_in_the_live_models_space():
|
||||
"""#4132: the column is a width, not an identity, so a same-width model
|
||||
swap writes a second geometry with no error. The backfill's "current"
|
||||
predicate has to name BOTH halves of the stamp, on every embedding table."""
|
||||
from scribe.models.embedding import (
|
||||
MilestoneEmbedding, NoteEmbedding, RuleEmbedding, SystemEmbedding,
|
||||
)
|
||||
from scribe.services import embeddings as emb
|
||||
|
||||
for table in (NoteEmbedding, RuleEmbedding, MilestoneEmbedding, SystemEmbedding):
|
||||
clause = str(emb.is_current_stamp(table))
|
||||
assert "chunker_version" in clause and "embedding_model" in clause, table
|
||||
|
||||
@@ -23,7 +23,7 @@ from scribe.models.embedding import EMBEDDING_DIM, NoteEmbedding
|
||||
from scribe.models.note import Note
|
||||
from scribe.models.project import Project
|
||||
from scribe.services import lessons as lessons_svc
|
||||
from scribe.services.embeddings import CHUNKER_VERSION, semantic_search_notes
|
||||
from scribe.services.embeddings import CHUNKER_VERSION, EMBEDDING_MODEL, semantic_search_notes
|
||||
from tests.helpers import ensure_user
|
||||
|
||||
pytestmark = [pytest.mark.integration, pytest.mark.usefixtures("_dispose_engine")]
|
||||
@@ -71,6 +71,7 @@ async def corpus():
|
||||
note_id=note.id, chunk_index=0, user_id=owner.id,
|
||||
embedding=QUERY_VEC, chunk_text=note.title,
|
||||
chunker_version=CHUNKER_VERSION,
|
||||
embedding_model=EMBEDDING_MODEL,
|
||||
))
|
||||
ids = {k: n.id for k, n in rows.items()}
|
||||
ids["owner"], ids["a"], ids["b"] = owner.id, a.id, b.id
|
||||
|
||||
@@ -17,7 +17,7 @@ from scribe.models.embedding import EMBEDDING_DIM, MilestoneEmbedding
|
||||
from scribe.models.milestone import Milestone
|
||||
from scribe.models.project import Project
|
||||
from scribe.services import dedup as dedup_svc
|
||||
from scribe.services.embeddings import CHUNKER_VERSION, semantic_search_milestones
|
||||
from scribe.services.embeddings import CHUNKER_VERSION, EMBEDDING_MODEL, semantic_search_milestones
|
||||
from tests.helpers import ensure_user
|
||||
|
||||
pytestmark = [pytest.mark.integration, pytest.mark.usefixtures("_dispose_engine", "_no_embedding")]
|
||||
@@ -48,7 +48,8 @@ async def roadmap():
|
||||
await s.flush()
|
||||
for ms, vec in ((m3, NEAR), (done, NEAR), (unrelated, FAR), (foreign, NEAR)):
|
||||
s.add(MilestoneEmbedding(milestone_id=ms.id, chunk_index=0, embedding=vec,
|
||||
chunk_text=ms.title, chunker_version=CHUNKER_VERSION))
|
||||
chunk_text=ms.title, chunker_version=CHUNKER_VERSION,
|
||||
embedding_model=EMBEDDING_MODEL))
|
||||
ids = {"owner": owner.id, "stranger": stranger.id, "mine": mine.id,
|
||||
"m3": m3.id, "done": done.id, "unrelated": unrelated.id, "foreign": foreign.id}
|
||||
await s.commit()
|
||||
|
||||
@@ -33,7 +33,7 @@ def _vec(*nonzero_first):
|
||||
|
||||
def _emb(note_id, user_id, chunk_index, vec):
|
||||
"""A chunk row at the current chunker version (#280, migration 0077)."""
|
||||
from scribe.services.embeddings import CHUNKER_VERSION
|
||||
from scribe.services.embeddings import CHUNKER_VERSION, EMBEDDING_MODEL
|
||||
|
||||
return NoteEmbedding(
|
||||
note_id=note_id,
|
||||
@@ -42,6 +42,7 @@ def _emb(note_id, user_id, chunk_index, vec):
|
||||
embedding=vec,
|
||||
chunk_text=f"chunk {chunk_index} of note {note_id}",
|
||||
chunker_version=CHUNKER_VERSION,
|
||||
embedding_model=EMBEDDING_MODEL,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -18,7 +18,7 @@ from scribe.models.embedding import EMBEDDING_DIM, RuleEmbedding
|
||||
from scribe.models.project import Project
|
||||
from scribe.models.share import ProjectShare
|
||||
from scribe.services import rulebooks as rulebooks_svc
|
||||
from scribe.services.embeddings import CHUNKER_VERSION, semantic_search_rules
|
||||
from scribe.services.embeddings import CHUNKER_VERSION, EMBEDDING_MODEL, semantic_search_rules
|
||||
from tests.helpers import ensure_user
|
||||
|
||||
pytestmark = [pytest.mark.integration, pytest.mark.usefixtures("_dispose_engine")]
|
||||
@@ -66,6 +66,7 @@ async def homes():
|
||||
s.add(RuleEmbedding(
|
||||
rule_id=rule.id, chunk_index=0, embedding=QUERY_VEC,
|
||||
chunk_text=rule.title, chunker_version=CHUNKER_VERSION,
|
||||
embedding_model=EMBEDDING_MODEL,
|
||||
))
|
||||
await s.commit()
|
||||
ids.update(glob=glob.id, on_a=on_a.id, on_b=on_b.id)
|
||||
|
||||
@@ -15,6 +15,9 @@ WHAT THIS PINS
|
||||
mid-backfill would end up with every bar at zero.
|
||||
4. **Every registry surface can be migrated.** A seventh arm that nobody adds
|
||||
a re-scorer for is one whose floor silently cannot survive a model change.
|
||||
5. **A half-migrated corpus is a refusal (#4132).** While any row the surface
|
||||
searches is stamped with another model, a re-score measures a blend of two
|
||||
geometries — so nothing is sampled until the backfill has finished.
|
||||
"""
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
@@ -43,11 +46,42 @@ def _session_with(rows):
|
||||
return session
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _corpus_on_the_live_model():
|
||||
"""Every test below starts from a wholly re-embedded corpus; the one that
|
||||
is about a half-migrated corpus patches this again."""
|
||||
with patch.object(rm, "rows_off_the_live_model", AsyncMock(return_value=0)):
|
||||
yield
|
||||
|
||||
|
||||
def test_every_surface_has_a_rescorer():
|
||||
"""Otherwise a surface's floor cannot cross a model change at all."""
|
||||
assert set(rm._RESCORERS) == set(SURFACES)
|
||||
|
||||
|
||||
def test_every_surface_names_the_corpus_it_searches():
|
||||
"""Otherwise the half-migrated check has no table to count for it."""
|
||||
assert set(rm._CORPUS) == set(SURFACES)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_corpus_part_way_through_a_model_change_refuses_before_sampling():
|
||||
rows = _logs([("q1", None, 0.9), ("q2", None, 0.1)])
|
||||
rescore = AsyncMock(return_value=0.4)
|
||||
with patch.object(rm, "rows_off_the_live_model", AsyncMock(return_value=17)) as off, \
|
||||
patch.object(rm, "async_session", MagicMock(return_value=_session_with(rows))), \
|
||||
patch.object(rm, "floor_for", AsyncMock(return_value=0.5)), \
|
||||
patch.object(rm, "set_dial", AsyncMock()) as set_dial, \
|
||||
patch.dict(rm._RESCORERS, {"prompt_rule": rescore}):
|
||||
out = await rm.migrate_floor(1, "prompt_rule", apply=True)
|
||||
|
||||
assert out["migrated"] is False
|
||||
assert out["rows_off_model"] == 17
|
||||
assert off.await_args.args[0] is rm.RuleEmbedding
|
||||
rescore.assert_not_called()
|
||||
set_dial.assert_not_called()
|
||||
|
||||
|
||||
def test_the_floor_that_admits_a_fraction_is_an_observed_score():
|
||||
scores = [0.9, 0.8, 0.7, 0.6, 0.5, 0.4, 0.3, 0.2, 0.1, 0.05]
|
||||
# 30% of ten is three; the third-best score is the bar that admits exactly
|
||||
|
||||
Reference in New Issue
Block a user