From f4e9cd429b1eeb8e3b64a19df8e0bbbe4140f29d Mon Sep 17 00:00:00 2001 From: Bryan Van Deusen Date: Wed, 23 Sep 2026 19:02:30 -0400 Subject: [PATCH] feat(embeddings): every vector records the model whose space it lives in (#4132) 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 --- .../0109_embeddings_record_their_model.py | 42 +++++++++++++++++++ src/scribe/models/embedding.py | 11 +++++ src/scribe/services/embeddings.py | 42 ++++++++++++++++--- src/scribe/services/retrieval_migration.py | 30 +++++++++++++ tests/test_chunking.py | 15 +++++++ tests/test_integration_lesson_reach.py | 3 +- tests/test_integration_milestone_search.py | 5 ++- tests/test_integration_pgvector_search.py | 3 +- tests/test_integration_rule_scope.py | 3 +- tests/test_retrieval_migration.py | 34 +++++++++++++++ 10 files changed, 178 insertions(+), 10 deletions(-) create mode 100644 alembic/versions/0109_embeddings_record_their_model.py diff --git a/alembic/versions/0109_embeddings_record_their_model.py b/alembic/versions/0109_embeddings_record_their_model.py new file mode 100644 index 0000000..334b298 --- /dev/null +++ b/alembic/versions/0109_embeddings_record_their_model.py @@ -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") diff --git a/src/scribe/models/embedding.py b/src/scribe/models/embedding.py index de5797c..4691a30 100644 --- a/src/scribe/models/embedding.py +++ b/src/scribe/models/embedding.py @@ -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), diff --git a/src/scribe/services/embeddings.py b/src/scribe/services/embeddings.py index fe1675e..ea14d71 100644 --- a/src/scribe/services/embeddings.py +++ b/src/scribe/services/embeddings.py @@ -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 = { diff --git a/src/scribe/services/retrieval_migration.py b/src/scribe/services/retrieval_migration.py index ba20473..00881b6 100644 --- a/src/scribe/services/retrieval_migration.py +++ b/src/scribe/services/retrieval_migration.py @@ -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( diff --git a/tests/test_chunking.py b/tests/test_chunking.py index 10d881a..5836164 100644 --- a/tests/test_chunking.py +++ b/tests/test_chunking.py @@ -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 diff --git a/tests/test_integration_lesson_reach.py b/tests/test_integration_lesson_reach.py index 24c7c40..7f996d5 100644 --- a/tests/test_integration_lesson_reach.py +++ b/tests/test_integration_lesson_reach.py @@ -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 diff --git a/tests/test_integration_milestone_search.py b/tests/test_integration_milestone_search.py index 94fe9b4..b786a92 100644 --- a/tests/test_integration_milestone_search.py +++ b/tests/test_integration_milestone_search.py @@ -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() diff --git a/tests/test_integration_pgvector_search.py b/tests/test_integration_pgvector_search.py index 4748c1e..d36ee63 100644 --- a/tests/test_integration_pgvector_search.py +++ b/tests/test_integration_pgvector_search.py @@ -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, ) diff --git a/tests/test_integration_rule_scope.py b/tests/test_integration_rule_scope.py index 27ca02d..41b928b 100644 --- a/tests/test_integration_rule_scope.py +++ b/tests/test_integration_rule_scope.py @@ -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) diff --git a/tests/test_retrieval_migration.py b/tests/test_retrieval_migration.py index e46820e..cc41aa4 100644 --- a/tests/test_retrieval_migration.py +++ b/tests/test_retrieval_migration.py @@ -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