diff --git a/backend/app/memorybank.py b/backend/app/memorybank.py index dba48b3..b955a4b 100644 --- a/backend/app/memorybank.py +++ b/backend/app/memorybank.py @@ -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) diff --git a/backend/app/migrations.py b/backend/app/migrations.py index 1339317..413cc11 100644 --- a/backend/app/migrations.py +++ b/backend/app/migrations.py @@ -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) diff --git a/backend/app/models.py b/backend/app/models.py index 7979e24..625ed7a 100644 --- a/backend/app/models.py +++ b/backend/app/models.py @@ -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) diff --git a/backend/app/routers/settings.py b/backend/app/routers/settings.py index 637093a..58c5b25 100644 --- a/backend/app/routers/settings.py +++ b/backend/app/routers/settings.py @@ -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 diff --git a/backend/tests/test_embedding_blob.py b/backend/tests/test_embedding_blob.py index 0ab4c52..47d8914 100644 --- a/backend/tests/test_embedding_blob.py +++ b/backend/tests/test_embedding_blob.py @@ -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 diff --git a/backend/tests/test_embedding_model_switch.py b/backend/tests/test_embedding_model_switch.py new file mode 100644 index 0000000..ab55714 --- /dev/null +++ b/backend/tests/test_embedding_model_switch.py @@ -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() diff --git a/backend/tests/test_memory_retrieval.py b/backend/tests/test_memory_retrieval.py index e24046f..5f0a593 100644 --- a/backend/tests/test_memory_retrieval.py +++ b/backend/tests/test_memory_retrieval.py @@ -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):