Chunked embeddings — no record content invisible to search (#280) #104

Merged
bvandeusen merged 4 commits from dev into main 2026-08-08 23:54:34 -04:00
9 changed files with 277 additions and 41 deletions
Showing only changes of commit 0e70a3896b - Show all commits
@@ -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)")
+18 -2
View File
@@ -1,7 +1,7 @@
from datetime import datetime, timezone from datetime import datetime, timezone
from pgvector.sqlalchemy import Vector 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 sqlalchemy.orm import Mapped, mapped_column
from scribe.models import Base from scribe.models import Base
@@ -14,7 +14,15 @@ EMBEDDING_DIM = 384
class NoteEmbedding(Base): 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" __tablename__ = "note_embeddings"
@@ -23,8 +31,16 @@ class NoteEmbedding(Base):
ForeignKey("notes.id", ondelete="CASCADE"), ForeignKey("notes.id", ondelete="CASCADE"),
primary_key=True, primary_key=True,
) )
chunk_index: Mapped[int] = mapped_column(Integer, primary_key=True)
user_id: Mapped[int] = mapped_column(Integer, nullable=False, index=True) user_id: Mapped[int] = mapped_column(Integer, nullable=False, index=True)
embedding: Mapped[list] = mapped_column(Vector(EMBEDDING_DIM), nullable=False) 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( updated_at: Mapped[datetime] = mapped_column(
DateTime(timezone=True), DateTime(timezone=True),
default=lambda: datetime.now(timezone.utc), default=lambda: datetime.now(timezone.utc),
+73 -20
View File
@@ -75,10 +75,19 @@ async def get_embedding(text: str) -> list[float]:
Raises if the fastembed model fails to load. Callers should catch and Raises if the fastembed model fails to load. Callers should catch and
degrade to keyword search. 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() embedder = await _get_model()
# embed() is synchronous CPU work; offload so we don't block the event loop. # embed() is synchronous CPU work; offload so we don't block the event loop.
vecs = await asyncio.to_thread(lambda: list(embedder.embed([text]))) vecs = await asyncio.to_thread(lambda: list(embedder.embed(texts)))
return vecs[0].tolist() return [v.tolist() for v in vecs]
def _cosine_similarity(a: list[float], b: list[float]) -> float: 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 return chunks
async def upsert_note_embedding(note_id: int, user_id: int, text: str) -> None: async def upsert_note_embedding(
"""Generate and persist an embedding for a note. Safe to fire-and-forget.""" note_id: int, user_id: int, title: str | None, body: str | None
if not text or not text.strip(): ) -> None:
return """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: 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: except Exception:
logger.debug("Skipping embedding for note %d — embedder unavailable", note_id) logger.debug("Skipping embedding for note %d — embedder unavailable", note_id)
return return
@@ -337,9 +370,19 @@ async def upsert_note_embedding(note_id: int, user_id: int, text: str) -> None:
await session.execute( await session.execute(
delete(NoteEmbedding).where(NoteEmbedding.note_id == note_id) 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() 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: except Exception:
logger.warning("Failed to persist embedding for note %d", note_id, exc_info=True) 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: 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 Runs as a background task at startup. Version-awareness is what makes a
so a large backfill doesn't peg CPU. 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: try:
async with async_session() as session: async with async_session() as session:
existing = { current = {
row[0] row[0]
for row in ( for row in (
await session.execute(select(NoteEmbedding.note_id)) await session.execute(
select(NoteEmbedding.note_id).where(
NoteEmbedding.chunker_version == CHUNKER_VERSION
)
)
).fetchall() ).fetchall()
} }
result = await session.execute( result = await session.execute(
select(Note.id, Note.user_id, Note.title, Note.body) select(Note.id, Note.user_id, Note.title, Note.body)
) )
notes_to_embed = [ 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: except Exception:
logger.warning("Embedding backfill: failed to query notes", exc_info=True) logger.warning("Embedding backfill: failed to query notes", exc_info=True)
return return
if not notes_to_embed: 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 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 success = 0
for note_id, user_id, title, body in notes_to_embed: for note_id, user_id, title, body in notes_to_embed:
text = embedding_text(title, body) if not chunk_document(title, body):
if not text:
continue continue
await upsert_note_embedding(note_id, user_id, text) await upsert_note_embedding(note_id, user_id, title, body)
success += 1 success += 1
await asyncio.sleep(0.05) # gentle pacing await asyncio.sleep(0.05) # gentle pacing
+6 -5
View File
@@ -33,11 +33,12 @@ def embed_note(note) -> None:
try: try:
import asyncio import asyncio
from scribe.services.embeddings import embedding_text, upsert_note_embedding from scribe.services.embeddings import upsert_note_embedding
text = embedding_text(note.title, note.body) # Chunking and the empty-record gate live inside upsert_note_embedding —
if not text: # one path for every writer (#280).
return asyncio.create_task(
asyncio.create_task(upsert_note_embedding(note.id, note.user_id, text)) upsert_note_embedding(note.id, note.user_id, note.title, note.body)
)
except RuntimeError: except RuntimeError:
pass # no running loop — a sync caller, not a failure pass # no running loop — a sync caller, not a failure
except Exception: # noqa: BLE001 - never let indexing break a write except Exception: # noqa: BLE001 - never let indexing break a write
+3 -5
View File
@@ -1,5 +1,4 @@
"""Recurring task rule validation, date calculation, and scheduled spawning.""" """Recurring task rule validation, date calculation, and scheduled spawning."""
import asyncio
import calendar import calendar
import logging import logging
from datetime import date, datetime, timedelta, timezone from datetime import date, datetime, timedelta, timezone
@@ -103,7 +102,6 @@ async def spawn_recurring_tasks() -> int:
Returns the number of tasks spawned. Returns the number of tasks spawned.
""" """
from scribe.services.embeddings import embedding_text, upsert_note_embedding
from scribe.services.notes import create_note from scribe.services.notes import create_note
now = datetime.now(timezone.utc) now = datetime.now(timezone.utc)
@@ -139,9 +137,9 @@ async def spawn_recurring_tasks() -> int:
milestone_id=task.milestone_id, milestone_id=task.milestone_id,
recurrence_rule=task.recurrence_rule, recurrence_rule=task.recurrence_rule,
) )
text = embedding_text(child.title, child.body) # No explicit embed here: create_note calls embed_note itself, so
if text: # the spawn path stopped being a second copy of the embedding rule
asyncio.create_task(upsert_note_embedding(child.id, task.user_id, text)) # the moment that moved into the service (#2056, #280).
except Exception: except Exception:
logger.exception("Failed to spawn recurring task %d", task.id) logger.exception("Failed to spawn recurring task %d", task.id)
continue continue
+94
View File
@@ -134,3 +134,97 @@ def test_a_monster_single_paragraph_is_hard_split_not_dropped():
assert len(chunks) > 1 assert len(chunks) > 1
total_words = sum(chunk.count("word") for chunk in chunks) total_words = sum(chunk.count("word") for chunk in chunks)
assert total_words == 2000 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"
+16 -2
View File
@@ -31,6 +31,20 @@ def _vec(*nonzero_first):
return v[:EMBEDDING_DIM] 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) @pytest_asyncio.fixture(autouse=True)
async def _dispose_engine(): async def _dispose_engine():
# Per-loop pool: dispose after each test (see test_integration_db_maintenance). # Per-loop pool: dispose after each test (see test_integration_db_maintenance).
@@ -54,8 +68,8 @@ async def seeded():
await s.flush() await s.flush()
# query vector will be [1,0,0,...]; near ~ identical (sim≈1.0), # query vector will be [1,0,0,...]; near ~ identical (sim≈1.0),
# far is orthogonal (sim≈0.0 -> filtered by the default threshold). # 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(_emb(near.id, user.id, 0, _vec(1.0)))
s.add(NoteEmbedding(note_id=far.id, user_id=user.id, embedding=_vec(0.0, 1.0))) s.add(_emb(far.id, user.id, 0, _vec(0.0, 1.0)))
await s.commit() await s.commit()
ids = (user.id, near.id, far.id) ids = (user.id, near.id, far.id)
yield ids yield ids
-1
View File
@@ -254,7 +254,6 @@ async def test_spawn_recurring_tasks_creates_child():
new_callable=AsyncMock, new_callable=AsyncMock,
return_value=mock_child, return_value=mock_child,
) as mock_create, ) as mock_create,
patch("scribe.services.embeddings.upsert_note_embedding"),
): ):
from scribe.services.recurrence import spawn_recurring_tasks from scribe.services.recurrence import spawn_recurring_tasks
count = await spawn_recurring_tasks() count = await spawn_recurring_tasks()
+13 -6
View File
@@ -29,18 +29,25 @@ def test_embed_note_uses_the_OWNER_not_the_caller():
upsert.assert_called_once() upsert.assert_called_once()
assert upsert.call_args.args[0] == 5 assert upsert.call_args.args[0] == 5
assert upsert.call_args.args[1] == 42 # owner, never the caller 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(): def test_embed_note_hands_even_an_empty_record_to_the_one_path():
"""An empty embedding is worse than none — it is a row that matches nothing """The empty-record decision moved INTO upsert_note_embedding (#280): an
and hides the fact that the record was never indexed.""" 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 from unittest.mock import MagicMock, patch
note = MagicMock(id=5, user_id=42, title="", body="") 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) 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(): def test_embed_note_without_a_running_loop_is_not_an_error():