Drop the JSON vector column, and fix what was hiding behind it

Migration 38 left memories.embedding in place so a rollback could still find
the vectors. Production has since been verified reading from embedding_blob,
so migration 42 drops it: 4 MB of a 99.6 MB database holding nothing anyone
reads.

Removing it surfaced a live bug. Changing your embedding model is supposed to
throw the bank's vectors away and let the post-turn pass rebuild them, because
two models' vectors are not comparable. The settings route did that by nulling
memories.embedding -- correct until 38 moved the vectors, after which it
cleared the dead column and left the blob intact with `embedded` still true.
_embed_pending filters on `embedded IS FALSE`, so it never saw those rows and
the bank went on ranking against the old model's vectors permanently.

Nothing would have reported it. cosine returns 0.0 on a width mismatch, so a
different-width model scores every memory zero and retrieval returns whichever
rows happen to sort first; a same-width model scores plausible garbage.

The bulk clear now sets both columns. It stays a bulk UPDATE rather than going
through set_vector -- loading the rows is the cost that whole path exists to
avoid -- so set_vector's docstring now names it as the one caller that
legitimately writes those columns by hand. No cache invalidation is added:
clearing `embedded` drops the rows out of the catalogue query, and set_vector
evicts each entry as the re-embed puts it back.

test_embedding_blob.py now rebuilds the pre-38 schema by hand where it tests
the backfill, since create_all no longer produces the column it converts from,
and asserts 42 removes it at the end of a full bootstrap -- 38 reads that
column and 42 drops it, so an upgrade that reordered them would arrive with an
empty bank.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_017Dvvqn9ZDR4ixeFPHNbww7
This commit is contained in:
parththakkar106
2026-08-17 13:44:59 +05:30
co-authored by Claude Opus 5
parent 85b188977e
commit 2c5909a268
7 changed files with 237 additions and 27 deletions
+8 -6
View File
@@ -86,13 +86,15 @@ def set_vector(memory: models.Memory, vector: list[float] | None) -> None:
"""Store (or clear) a memory's embedding.
Every column that describes the vector moves together: `embedding_blob` is
what the ranking reads, `embedded` is the flag everything else reads, and
the JSON `embedding` stays correct behind both until the follow-up
migration drops it. Going through one function is what keeps them in step —
and it is also the only place a stored vector can change, which is what
makes the cache below safe to invalidate here and nowhere else.
what the ranking reads and `embedded` is the flag everything else reads.
Going through one function is what keeps them in step — and it is also the
only place a stored vector can change, which is what makes the cache below
safe to invalidate here and nowhere else.
The one caller that legitimately cannot come through here is the bulk
clear in `routers/settings.py` when the embedding model changes. It has to
set the same two columns by hand; see the note there.
"""
memory.embedding = vector
memory.embedding_blob = None if vector is None else vectors.pack(vector)
memory.embedded = vector is not None
cached = _vector_cache.get(memory.adventure_id)
+6
View File
@@ -144,6 +144,12 @@ MIGRATIONS: list[tuple[int, str | dict[str, str]]] = [
# so anyone who picked a value keeps it — same rule as migration 29.
# Adventures already over 80 evict down on their next turn.
(41, "UPDATE settings SET memory_bank_capacity = 80 WHERE memory_bank_capacity = 200"),
# The JSON vectors, gone. Migration 38 left them in place so a rollback
# could still find them; production has since been verified reading from
# embedding_blob (schema_version 41, 134/134 backfilled), so the column is
# now 4 MB of a 99.6 MB database holding nothing anyone reads. DROP COLUMN
# is spelled the same on both dialects — SQLite has had it since 3.35.
(42, "ALTER TABLE memories DROP COLUMN embedding"),
]
LATEST_VERSION = max((v for v, _ in MIGRATIONS), default=1)
-4
View File
@@ -158,10 +158,6 @@ class Memory(Base):
id: Mapped[int] = mapped_column(primary_key=True)
adventure_id: Mapped[int] = mapped_column(ForeignKey("adventures.id", ondelete="CASCADE"))
text: Mapped[str] = mapped_column(Text, default="")
# Superseded by embedding_blob and still written alongside it, so a
# rollback finds the vectors, until a follow-up migration drops it. Nothing
# reads it.
embedding: Mapped[list | None] = mapped_column(JSON, nullable=True, deferred=True)
# The vector, little-endian float32. Deferred because it is wider than the
# rest of the row put together and exactly one code path wants it: anything
# bulk-loading memories (the Memories drawer, eviction, the embed queue)
+13 -1
View File
@@ -51,14 +51,26 @@ def update_settings(
# Vectors from the old model have a different dimensionality/space;
# clear them so the post-turn task re-embeds with the new model.
# (This user's adventures only — settings are per-user now.)
#
# Both columns, and the flag. This is the one place that clears vectors
# in bulk rather than through memorybank.set_vector, and when the
# vectors moved to embedding_blob it kept nulling the old JSON column
# alone: the blob survived, `embedded` stayed true, and _embed_pending
# — which looks for embedded IS FALSE — never picked the rows up. The
# bank went on ranking against the previous model's vectors forever.
owned = (
db.query(models.Adventure.id)
.filter(models.Adventure.user_id == user.id)
.scalar_subquery()
)
db.query(models.Memory).filter(models.Memory.adventure_id.in_(owned)).update(
{"embedding": None}, synchronize_session=False
{"embedding_blob": None, "embedded": False}, synchronize_session=False
)
# No cache invalidation needed, and deliberately none added: clearing
# `embedded` drops these rows out of the catalogue query, so retrieval
# stops asking for them, and by the time _embed_pending puts one back
# it has gone through set_vector, which evicts that entry. The rule
# holds — anything that removes a memory from play self-corrects.
db.commit()
return settings
+32 -8
View File
@@ -111,9 +111,9 @@ def test_cosine_moved_but_still_reachable_from_memorybank():
# -------------------------------------------------------------- set_vector
def test_set_vector_writes_both_columns(db, adventure):
"""Until the follow-up migration drops the JSON column, it has to stay
correct — a rollback reads it."""
def test_set_vector_writes_the_blob_and_the_flag(db, adventure):
"""The two columns that describe a vector move together, or a reader that
trusts `embedded` gets a NULL blob."""
memory = models.Memory(adventure_id=adventure.id, text="a fact")
db.add(memory)
db.commit()
@@ -123,7 +123,6 @@ def test_set_vector_writes_both_columns(db, adventure):
db.commit()
db.expire_all()
assert memory.embedding == vector
assert list(vectors.unpack(memory.embedding_blob)) == vector
assert memory.embedded is True
@@ -140,24 +139,40 @@ def test_set_vector_none_clears_both(db, adventure):
db.commit()
db.expire_all()
assert memory.embedding is None
assert memory.embedding_blob is None
assert memory.embedded is False
# --------------------------------------------------------------- the backfill
def add_legacy_json_column(db) -> None:
"""Put `memories.embedding` back for the length of a test.
Migration 42 dropped it and the model no longer declares it, so
`create_all` does not produce it — but everything below is testing the
upgrade *from* a database that still has it, which is the only state in
which the backfill has any work to do. Re-adding it by hand is what keeps
these tests honest about the schema they claim to be starting from.
"""
db.execute(text("ALTER TABLE memories ADD COLUMN embedding JSON"))
db.commit()
def seed_json_only(db, adventure, count: int, dims: int = 64) -> dict[int, list[float]]:
"""Memories as they exist before the migration: JSON vector, no blob."""
add_legacy_json_column(db)
rng = random.Random(count)
expected = {}
for i in range(count):
vector = sample_vector(rng, dims)
memory = models.Memory(
adventure_id=adventure.id, text=f"fact {i}", embedding=vector
)
memory = models.Memory(adventure_id=adventure.id, text=f"fact {i}")
db.add(memory)
db.flush()
# Raw, because the ORM no longer knows this column exists.
db.execute(
text("UPDATE memories SET embedding = :v WHERE id = :id"),
{"v": json.dumps(vector), "id": memory.id},
)
expected[memory.id] = vector
db.commit()
db.execute(text("UPDATE memories SET embedding_blob = NULL, embedded = false"))
@@ -193,6 +208,7 @@ def test_backfill_reaches_past_one_batch(db, adventure):
def test_backfill_leaves_unembedded_memories_alone(db, adventure):
add_legacy_json_column(db)
db.add(models.Memory(adventure_id=adventure.id, text="not embedded yet"))
db.commit()
@@ -274,6 +290,14 @@ def test_bootstrap_adds_the_columns_and_backfills_them(db, adventure):
# embedded must still read as not embedded afterwards.
assert by_id[unembedded_id] == (None, False)
# ...and migration 42, at the end of the same run, takes the JSON column
# away. Ordering matters: 38 reads it, 42 drops it, and an upgrade that
# ran them the other way round would arrive with an empty bank.
with engine.begin() as conn:
columns = {row[1] for row in conn.execute(text("PRAGMA table_info(memories)"))}
assert "embedding" not in columns
assert {"embedding_blob", "embedded"} <= columns
def test_migration_38_is_spelled_for_both_dialects():
"""Every Postgres deploy replays migrations from 24 on, so a SQLite-only
@@ -0,0 +1,169 @@
"""Switching embedding models must re-embed the bank.
Vectors from two different models are not comparable — different space, often
different width — so changing the model has to throw the stored ones away and
let the post-turn pass rebuild them.
That worked while the vectors lived in `memories.embedding`: the settings
route nulled that column and the embed queue picked the rows up. Migration 38
moved the vectors to `embedding_blob` with an `embedded` flag beside them, and
the bulk clear kept nulling the old column alone. The blob survived, the flag
stayed true, `_embed_pending` (which looks for `embedded IS FALSE`) never saw
the rows, and the bank went on ranking against the previous model's vectors
for good.
Nothing reports this. `cosine` returns 0.0 on a width mismatch, so a
different-width model scores every memory zero and retrieval quietly returns
whichever rows sort first; a same-width model scores plausible-looking
garbage.
python -m pytest tests/test_embedding_model_switch.py -v
"""
import os
import tempfile
_tmp = tempfile.NamedTemporaryFile(suffix=".db", delete=False)
_tmp.close()
os.environ["AIDND_DB_PATH"] = _tmp.name
os.environ.pop("AIDND_DATABASE_URL", None)
os.environ.pop("DATABASE_URL", None)
import asyncio
import pytest
from fastapi import Depends
from fastapi.testclient import TestClient
from app import auth, limits, memorybank, models
from app.database import Base, SessionLocal, engine, get_db
from app.main import app
DIMS = 8
@pytest.fixture()
def client(monkeypatch):
Base.metadata.create_all(bind=engine)
memorybank._vector_cache.clear()
setup = SessionLocal()
user = models.User(is_guest=False, email="switch@example.com")
setup.add(user)
setup.flush()
setup.add(models.Settings(
user_id=user.id, api_key="enc:dummy", model="test-model",
embedding_model="model-a",
))
adventure = models.Adventure(
user_id=user.id, title="Cave", script_state={}, memory_bank_enabled=True
)
setup.add(adventure)
setup.flush()
for i in range(5):
memory = models.Memory(adventure_id=adventure.id, text=f"Memory {i}")
memorybank.set_vector(memory, [float(i)] + [0.0] * (DIMS - 1))
setup.add(memory)
setup.commit()
adv_id, user_id = adventure.id, user.id
setup.close()
monkeypatch.setattr(limits, "rate_limit", lambda *a, **k: None)
monkeypatch.setattr(limits, "check_row_cap", lambda *a, **k: None)
def _current_user(db=Depends(get_db)):
return db.get(models.User, user_id)
app.dependency_overrides[auth.get_current_user] = _current_user
c = TestClient(app)
c.adv_id = adv_id
try:
yield c
finally:
app.dependency_overrides.clear()
memorybank._vector_cache.clear()
Base.metadata.drop_all(bind=engine)
def memories(db):
return db.query(models.Memory).order_by(models.Memory.id).all()
def test_the_bank_starts_embedded(client):
db = SessionLocal()
try:
rows = memories(db)
assert len(rows) == 5
assert all(m.embedded for m in rows)
assert all(m.embedding_blob for m in rows)
finally:
db.close()
def test_changing_the_model_clears_every_vector(client):
r = client.put("/api/settings", json={"embedding_model": "model-b"})
assert r.status_code == 200, r.text
db = SessionLocal()
try:
rows = memories(db)
assert [m.embedding_blob for m in rows] == [None] * 5, \
"the blob survived the model change"
assert not any(m.embedded for m in rows), \
"`embedded` stayed true, so nothing will ever re-embed these"
finally:
db.close()
def test_cleared_memories_are_queued_for_re_embedding(client):
"""The flag is not cosmetic: it is the only thing `_embed_pending` filters
on, so this is the assertion that the bank actually recovers."""
client.put("/api/settings", json={"embedding_model": "model-b"})
db = SessionLocal()
try:
pending = (
db.query(models.Memory)
.filter(models.Memory.embedded.is_(False),
models.Memory.forgotten.is_(False))
.all()
)
assert len(pending) == 5
finally:
db.close()
def test_retrieval_uses_no_stale_vector_after_the_switch(client, monkeypatch):
"""Until the re-embed runs, the bank must return nothing rather than
ranking against the old model's vectors."""
client.put("/api/settings", json={"embedding_model": "model-b"})
class Embedder:
async def embed(self, texts):
return [[1.0] + [0.0] * (DIMS - 1) for _ in texts]
monkeypatch.setattr(memorybank, "embedding_provider", lambda s: Embedder())
db = SessionLocal()
try:
adventure = db.get(models.Adventure, client.adv_id)
settings = db.query(models.Settings).first()
result = asyncio.run(
memorybank.retrieve_memories(adventure, settings, update_stats=False)
)
assert result["used"] == []
finally:
db.close()
def test_an_unrelated_settings_change_keeps_the_vectors(client):
"""Only an embedding-model change may clear the bank — re-embedding costs
an API call per memory."""
r = client.put("/api/settings", json={"model": "some-other-chat-model"})
assert r.status_code == 200, r.text
db = SessionLocal()
try:
rows = memories(db)
assert all(m.embedded for m in rows)
assert all(m.embedding_blob for m in rows)
finally:
db.close()
+9 -8
View File
@@ -26,7 +26,7 @@ import asyncio
from datetime import timedelta
import pytest
from sqlalchemy import event
from sqlalchemy import event, inspect as sa_inspect
from app import memorybank, models
from app.database import Base, SessionLocal, engine
@@ -206,13 +206,14 @@ def memory_selects(statements):
]
def test_the_json_column_is_never_selected(db, adventure, settings, bank, sql_log):
"""`memories.embedding` is dead weight kept only until a follow-up
migration drops it. If anything still reads it, dropping it breaks."""
retrieve(adventure, settings, StubEmbedder())
offenders = [s for s in memory_selects(sql_log) if "memories.embedding " in s
or s.rstrip().endswith("memories.embedding")]
assert offenders == [], f"the JSON column was read:\n{offenders[0][:300]}"
def test_the_json_column_is_gone(db):
"""`memories.embedding` held the vectors before migration 38 and nothing
read it afterwards; migration 42 dropped it. Bringing it back would restore
4 MB of dead weight and a second place vectors can be written from — which
is how the model-switch bug happened (test_embedding_model_switch.py)."""
columns = {c["name"] for c in sa_inspect(engine).get_columns("memories")}
assert "embedding" not in columns
assert {"embedding_blob", "embedded"} <= columns
def test_the_catalogue_query_carries_no_vectors(db, adventure, settings, bank, sql_log):