diff --git a/alembic/versions/0077_chunked_embeddings.py b/alembic/versions/0077_chunked_embeddings.py new file mode 100644 index 0000000..f7a8bd8 --- /dev/null +++ b/alembic/versions/0077_chunked_embeddings.py @@ -0,0 +1,54 @@ +"""Chunked embeddings: one note_embeddings row per chunk (#280) + +Revision ID: 0077 +Revises: 0076 +Create Date: 2026-08-09 + +The embedding model reads at most 512 tokens and fastembed truncates the rest +silently, so the old one-row-per-note shape permanently lost everything past +~400 words of a record. A note now stores one row per chunk of +`embeddings.chunk_document`: PK (note_id, chunk_index), plus the chunk's text +(inspectability + future "matched section" surfacing) and the chunker version +that produced it (so later shape changes re-embed by version comparison +instead of repeating this wipe). + +Embeddings are DERIVED data (0067 precedent): rows are cleared here and the +startup backfill regenerates the whole corpus at the new shape on next boot. +The HNSW index is untouched — it indexes chunk rows exactly as it indexed +note rows. +""" +from alembic import op + + +revision = "0077" +down_revision = "0076" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + # Derived data — the version-aware startup backfill re-embeds everything + # at the chunked shape. Old whole-document rows would be indistinguishable + # from properly-chunked single-chunk notes, so they cannot be carried over. + op.execute("DELETE FROM note_embeddings") + + # Empty table, so NOT NULL columns need no defaults and the PK swap is + # instant. + op.execute("ALTER TABLE note_embeddings ADD COLUMN chunk_index integer NOT NULL") + op.execute("ALTER TABLE note_embeddings ADD COLUMN chunk_text text NOT NULL") + op.execute("ALTER TABLE note_embeddings ADD COLUMN chunker_version integer NOT NULL") + op.execute("ALTER TABLE note_embeddings DROP CONSTRAINT note_embeddings_pkey") + op.execute( + "ALTER TABLE note_embeddings ADD PRIMARY KEY (note_id, chunk_index)" + ) + + +def downgrade() -> None: + # Same reasoning in reverse: chunk rows make no sense to a whole-document + # reader, so clear and let the old backfill regenerate. + op.execute("DELETE FROM note_embeddings") + op.execute("ALTER TABLE note_embeddings DROP CONSTRAINT note_embeddings_pkey") + op.execute("ALTER TABLE note_embeddings DROP COLUMN chunk_index") + op.execute("ALTER TABLE note_embeddings DROP COLUMN chunk_text") + op.execute("ALTER TABLE note_embeddings DROP COLUMN chunker_version") + op.execute("ALTER TABLE note_embeddings ADD PRIMARY KEY (note_id)") diff --git a/src/scribe/models/embedding.py b/src/scribe/models/embedding.py index 5de4fd9..9cf4f50 100644 --- a/src/scribe/models/embedding.py +++ b/src/scribe/models/embedding.py @@ -1,7 +1,7 @@ from datetime import datetime, timezone from pgvector.sqlalchemy import Vector -from sqlalchemy import DateTime, ForeignKey, Integer +from sqlalchemy import DateTime, ForeignKey, Integer, Text from sqlalchemy.orm import Mapped, mapped_column from scribe.models import Base @@ -14,7 +14,15 @@ EMBEDDING_DIM = 384 class NoteEmbedding(Base): - """Stores the embedding vector for a note, used for semantic search.""" + """One embedding vector per CHUNK of a note (#280, migration 0077). + + The model reads at most 512 tokens, so a single whole-document vector + permanently lost everything past ~400 words. A note now stores one row per + chunk of `embeddings.chunk_document`, and a query matches the note if it + matches ANY chunk — retrieval collapses rows to best-chunk-per-note. + A short note has exactly one row (chunk_index 0) whose text is the + historical `title\\nbody` shape. + """ __tablename__ = "note_embeddings" @@ -23,8 +31,16 @@ class NoteEmbedding(Base): ForeignKey("notes.id", ondelete="CASCADE"), primary_key=True, ) + chunk_index: Mapped[int] = mapped_column(Integer, primary_key=True) user_id: Mapped[int] = mapped_column(Integer, nullable=False, index=True) embedding: Mapped[list] = mapped_column(Vector(EMBEDDING_DIM), nullable=False) + # Exactly what this vector encodes — inspectable when a ranking surprises, + # and the hook for surfacing WHICH section matched, later. + chunk_text: Mapped[str] = mapped_column(Text, nullable=False) + # embeddings.CHUNKER_VERSION at write time. The startup backfill re-embeds + # 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) 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 5fa03f9..1a46c89 100644 --- a/src/scribe/services/embeddings.py +++ b/src/scribe/services/embeddings.py @@ -75,10 +75,19 @@ async def get_embedding(text: str) -> list[float]: Raises if the fastembed model fails to load. Callers should catch and degrade to keyword search. """ + return (await get_embeddings([text]))[0] + + +async def get_embeddings(texts: list[str]) -> list[list[float]]: + """Embed several texts in one model call (the chunked write path). + + fastembed batches internally, so N chunks cost far less than N single + calls. Raises like get_embedding; callers catch and degrade. + """ embedder = await _get_model() # embed() is synchronous CPU work; offload so we don't block the event loop. - vecs = await asyncio.to_thread(lambda: list(embedder.embed([text]))) - return vecs[0].tolist() + vecs = await asyncio.to_thread(lambda: list(embedder.embed(texts))) + return [v.tolist() for v in vecs] def _cosine_similarity(a: list[float], b: list[float]) -> float: @@ -322,12 +331,36 @@ def chunk_document(title: str | None, body: str | None) -> list[str]: return chunks -async def upsert_note_embedding(note_id: int, user_id: int, text: str) -> None: - """Generate and persist an embedding for a note. Safe to fire-and-forget.""" - if not text or not text.strip(): - return +async def upsert_note_embedding( + note_id: int, user_id: int, title: str | None, body: str | None +) -> None: + """Chunk, embed and persist a note's vectors. Safe to fire-and-forget. + + Takes title/body rather than pre-built text so the chunking happens HERE — + one path for the write path, the recurrence spawn and the startup backfill, + which is the same single-definition discipline embedding_text existed for. + + Replacement is atomic per note: old rows are deleted and the new chunk set + inserted in one transaction, so a concurrent read sees the old shape or the + new one, never a mixture. + """ + chunks = chunk_document(title, body) try: - embedding = await get_embedding(text) + if not chunks: + # A record emptied of content should stop being findable by its + # old content — clear stale vectors rather than leaving them. + async with async_session() as session: + await session.execute( + delete(NoteEmbedding).where(NoteEmbedding.note_id == note_id) + ) + await session.commit() + return + except Exception: + logger.warning("Failed to clear embedding for note %d", note_id, exc_info=True) + return + + try: + vectors = await get_embeddings(chunks) except Exception: logger.debug("Skipping embedding for note %d — embedder unavailable", note_id) return @@ -337,9 +370,19 @@ async def upsert_note_embedding(note_id: int, user_id: int, text: str) -> None: await session.execute( delete(NoteEmbedding).where(NoteEmbedding.note_id == note_id) ) - session.add(NoteEmbedding(note_id=note_id, user_id=user_id, embedding=embedding)) + for index, (chunk, vector) in enumerate(zip(chunks, vectors)): + session.add( + NoteEmbedding( + note_id=note_id, + chunk_index=index, + user_id=user_id, + embedding=vector, + chunk_text=chunk, + chunker_version=CHUNKER_VERSION, + ) + ) await session.commit() - logger.debug("Upserted embedding for note %d", note_id) + logger.debug("Upserted %d chunk embedding(s) for note %d", len(chunks), note_id) except Exception: logger.warning("Failed to persist embedding for note %d", note_id, exc_info=True) @@ -500,40 +543,50 @@ async def semantic_search_notes( async def backfill_note_embeddings() -> None: - """Generate embeddings for all notes that don't have one yet. + """(Re-)embed every note that is missing vectors OR whose stored vectors + were produced by an older chunker. - Runs as a background task at startup. Adds a small sleep between notes - so a large backfill doesn't peg CPU. + Runs as a background task at startup. Version-awareness is what makes a + document-shape change deployable: migration 0077 cleared the table once, + and every later CHUNKER_VERSION bump re-embeds the stale notes here — a + version comparison instead of another wipe. Adds a small sleep between + notes so a large backfill doesn't peg CPU. """ try: async with async_session() as session: - existing = { + current = { row[0] for row in ( - await session.execute(select(NoteEmbedding.note_id)) + await session.execute( + select(NoteEmbedding.note_id).where( + NoteEmbedding.chunker_version == CHUNKER_VERSION + ) + ) ).fetchall() } result = await session.execute( select(Note.id, Note.user_id, Note.title, Note.body) ) notes_to_embed = [ - row for row in result.fetchall() if row[0] not in existing + row for row in result.fetchall() if row[0] not in current ] except Exception: logger.warning("Embedding backfill: failed to query notes", exc_info=True) return if not notes_to_embed: - logger.info("Embedding backfill: all notes already have embeddings") + logger.info("Embedding backfill: all notes current at chunker v%d", CHUNKER_VERSION) return - logger.info("Embedding backfill: generating embeddings for %d notes", len(notes_to_embed)) + logger.info( + "Embedding backfill: embedding %d notes at chunker v%d", + len(notes_to_embed), CHUNKER_VERSION, + ) success = 0 for note_id, user_id, title, body in notes_to_embed: - text = embedding_text(title, body) - if not text: + if not chunk_document(title, body): continue - await upsert_note_embedding(note_id, user_id, text) + await upsert_note_embedding(note_id, user_id, title, body) success += 1 await asyncio.sleep(0.05) # gentle pacing diff --git a/src/scribe/services/notes.py b/src/scribe/services/notes.py index ee9b130..4a05e75 100644 --- a/src/scribe/services/notes.py +++ b/src/scribe/services/notes.py @@ -33,11 +33,12 @@ def embed_note(note) -> None: try: import asyncio - from scribe.services.embeddings import embedding_text, upsert_note_embedding - text = embedding_text(note.title, note.body) - if not text: - return - asyncio.create_task(upsert_note_embedding(note.id, note.user_id, text)) + from scribe.services.embeddings import upsert_note_embedding + # Chunking and the empty-record gate live inside upsert_note_embedding — + # one path for every writer (#280). + asyncio.create_task( + upsert_note_embedding(note.id, note.user_id, note.title, note.body) + ) except RuntimeError: pass # no running loop — a sync caller, not a failure except Exception: # noqa: BLE001 - never let indexing break a write diff --git a/src/scribe/services/recurrence.py b/src/scribe/services/recurrence.py index 5490b79..39b7858 100644 --- a/src/scribe/services/recurrence.py +++ b/src/scribe/services/recurrence.py @@ -1,5 +1,4 @@ """Recurring task rule validation, date calculation, and scheduled spawning.""" -import asyncio import calendar import logging from datetime import date, datetime, timedelta, timezone @@ -103,7 +102,6 @@ async def spawn_recurring_tasks() -> int: Returns the number of tasks spawned. """ - from scribe.services.embeddings import embedding_text, upsert_note_embedding from scribe.services.notes import create_note now = datetime.now(timezone.utc) @@ -139,9 +137,9 @@ async def spawn_recurring_tasks() -> int: milestone_id=task.milestone_id, recurrence_rule=task.recurrence_rule, ) - text = embedding_text(child.title, child.body) - if text: - asyncio.create_task(upsert_note_embedding(child.id, task.user_id, text)) + # No explicit embed here: create_note calls embed_note itself, so + # the spawn path stopped being a second copy of the embedding rule + # the moment that moved into the service (#2056, #280). except Exception: logger.exception("Failed to spawn recurring task %d", task.id) continue diff --git a/tests/test_chunking.py b/tests/test_chunking.py index 044ea68..1269862 100644 --- a/tests/test_chunking.py +++ b/tests/test_chunking.py @@ -134,3 +134,97 @@ def test_a_monster_single_paragraph_is_hard_split_not_dropped(): assert len(chunks) > 1 total_words = sum(chunk.count("word") for chunk in chunks) assert total_words == 2000 + + +# --- the write path: one row per chunk (#280 step 3) ------------------------- + + +def _session_ctx(): + from unittest.mock import AsyncMock, MagicMock + + session = MagicMock() + session.execute = AsyncMock() + session.commit = AsyncMock() + ctx = MagicMock() + ctx.__aenter__ = AsyncMock(return_value=session) + ctx.__aexit__ = AsyncMock(return_value=False) + return session, ctx + + +async def test_upsert_stores_one_versioned_row_per_chunk(): + from unittest.mock import AsyncMock, patch + + from scribe.services import embeddings as emb + + body = "\n\n".join( + f"## Section {i}\n\n{_long_section(f'sec-{i}')}" for i in range(6) + ) + chunks = chunk_document("T", body) + assert len(chunks) > 1 + + session, ctx = _session_ctx() + with ( + patch.object(emb, "async_session", return_value=ctx), + patch.object( + emb, "get_embeddings", + AsyncMock(return_value=[[0.0] * 384 for _ in chunks]), + ), + ): + await emb.upsert_note_embedding(7, 42, "T", body) + + rows = [call.args[0] for call in session.add.call_args_list] + 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.user_id for r in rows} == {42} + session.execute.assert_awaited() # the delete that makes replacement atomic + + +async def test_upsert_of_an_emptied_record_clears_rows_instead_of_embedding(): + """An empty embedding is worse than none, and a STALE one is worse than + that — a record emptied of content must stop being findable by what it no + longer says.""" + from unittest.mock import patch + + from scribe.services import embeddings as emb + + session, ctx = _session_ctx() + with ( + patch.object(emb, "async_session", return_value=ctx), + patch.object(emb, "get_embeddings") as embedder, + ): + await emb.upsert_note_embedding(7, 42, "", "") + + embedder.assert_not_called() + session.execute.assert_awaited() # the delete + session.add.assert_not_called() + session.commit.assert_awaited() + + +async def test_backfill_reembeds_notes_with_a_stale_chunker_version(): + """The reason chunker_version exists: a shape change becomes a version + bump that re-embeds exactly the stale notes, instead of another 0077-style + table wipe. Only rows AT the current version count as done.""" + from unittest.mock import AsyncMock, MagicMock, patch + + from scribe.services import embeddings as emb + + current_rows = MagicMock() + current_rows.fetchall.return_value = [(1,)] # note 1 is current + note_rows = MagicMock() + note_rows.fetchall.return_value = [ + (1, 42, "current", "body"), + (2, 42, "stale-version", "body"), + ] + session, ctx = _session_ctx() + session.execute = AsyncMock(side_effect=[current_rows, note_rows]) + + with ( + patch.object(emb, "async_session", return_value=ctx), + patch.object(emb, "upsert_note_embedding", AsyncMock()) as upsert, + patch.object(emb.asyncio, "sleep", AsyncMock()), + ): + await emb.backfill_note_embeddings() + + embedded = [call.args[0] for call in upsert.call_args_list] + assert embedded == [2], "only the stale note is re-embedded" diff --git a/tests/test_integration_pgvector_search.py b/tests/test_integration_pgvector_search.py index d9938e6..efb1349 100644 --- a/tests/test_integration_pgvector_search.py +++ b/tests/test_integration_pgvector_search.py @@ -31,6 +31,20 @@ def _vec(*nonzero_first): return v[:EMBEDDING_DIM] +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 + + return NoteEmbedding( + note_id=note_id, + chunk_index=chunk_index, + user_id=user_id, + embedding=vec, + chunk_text=f"chunk {chunk_index} of note {note_id}", + chunker_version=CHUNKER_VERSION, + ) + + @pytest_asyncio.fixture(autouse=True) async def _dispose_engine(): # Per-loop pool: dispose after each test (see test_integration_db_maintenance). @@ -54,8 +68,8 @@ async def seeded(): await s.flush() # query vector will be [1,0,0,...]; near ~ identical (sim≈1.0), # far is orthogonal (sim≈0.0 -> filtered by the default threshold). - s.add(NoteEmbedding(note_id=near.id, user_id=user.id, embedding=_vec(1.0))) - s.add(NoteEmbedding(note_id=far.id, user_id=user.id, embedding=_vec(0.0, 1.0))) + s.add(_emb(near.id, user.id, 0, _vec(1.0))) + s.add(_emb(far.id, user.id, 0, _vec(0.0, 1.0))) await s.commit() ids = (user.id, near.id, far.id) yield ids diff --git a/tests/test_recurrence.py b/tests/test_recurrence.py index fad807d..d11d3f3 100644 --- a/tests/test_recurrence.py +++ b/tests/test_recurrence.py @@ -254,7 +254,6 @@ async def test_spawn_recurring_tasks_creates_child(): new_callable=AsyncMock, return_value=mock_child, ) as mock_create, - patch("scribe.services.embeddings.upsert_note_embedding"), ): from scribe.services.recurrence import spawn_recurring_tasks count = await spawn_recurring_tasks() diff --git a/tests/test_services_notes.py b/tests/test_services_notes.py index da74997..e4b9de9 100644 --- a/tests/test_services_notes.py +++ b/tests/test_services_notes.py @@ -29,18 +29,25 @@ def test_embed_note_uses_the_OWNER_not_the_caller(): upsert.assert_called_once() assert upsert.call_args.args[0] == 5 assert upsert.call_args.args[1] == 42 # owner, never the caller - assert upsert.call_args.args[2] == "T\nB" + # Title and body travel separately since #280 — chunking happens inside + # upsert_note_embedding, the one path every writer shares. + assert upsert.call_args.args[2] == "T" + assert upsert.call_args.args[3] == "B" -def test_embed_note_skips_a_record_with_no_text(): - """An empty embedding is worse than none — it is a row that matches nothing - and hides the fact that the record was never indexed.""" +def test_embed_note_hands_even_an_empty_record_to_the_one_path(): + """The empty-record decision moved INTO upsert_note_embedding (#280): an + emptied record must have its stale vectors CLEARED, not merely skipped — + so embed_note schedules the call unconditionally rather than deciding + here. The clearing behaviour itself is pinned in test_chunking.py.""" from unittest.mock import MagicMock, patch note = MagicMock(id=5, user_id=42, title="", body="") - with patch("asyncio.create_task") as create_task: + with patch("scribe.services.embeddings.upsert_note_embedding") as upsert, \ + patch("asyncio.create_task") as create_task: notes_svc.embed_note(note) - assert not create_task.called + assert create_task.called + upsert.assert_called_once() def test_embed_note_without_a_running_loop_is_not_an_error():