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
11 changed files with 727 additions and 103 deletions
@@ -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),
+45 -27
View File
@@ -69,6 +69,12 @@ _SEMANTIC_THRESHOLD = 0.90
# structural signals cannot see. # structural signals cannot see.
_SNIPPET_SEMANTIC_THRESHOLD = 0.96 _SNIPPET_SEMANTIC_THRESHOLD = 0.96
# The gate queries per CHUNK of the candidate (#280) — this caps how many
# searches one save may cost. Eight chunks ≈ five thousand words of candidate;
# a duplicate hiding past that is the duplicate report's job to find, not a
# reason to stall the write path.
_GATE_MAX_CHUNKS = 8
@dataclass @dataclass
class DuplicateMatch: class DuplicateMatch:
@@ -258,39 +264,43 @@ async def find_duplicate_note(
# --- Signal 3: semantic similarity (only with a substantial body) --- # --- Signal 3: semantic similarity (only with a substantial body) ---
if body and len(body.strip()) >= _MIN_BODY_FOR_SEMANTIC: if body and len(body.strip()) >= _MIN_BODY_FOR_SEMANTIC:
# Built by the SAME function the corpus was embedded with. This one is # Query with the SAME chunker the corpus was embedded with (#280). This
# the copy that mattered most and was easiest to miss: it is a QUERY # was the copy that mattered most and was easiest to miss: these are
# document, compared against embedded ones. Shaped differently from the # QUERY documents, compared against embedded ones — shaped differently
# corpus it searches, the gate degrades silently — it still returns # from the corpus, the gate degrades silently. Chunking also makes the
# neighbours, just less apt ones, and no signal says the query and the # gate see what the whole-document query diluted: a long candidate that
# index stopped agreeing (found by the guard in test_embedding_text). # duplicates an existing record IN ONE SECTION now matches on that
query = embeddings_svc.embedding_text(title, body) # section. Capped so one pathological paste can't turn a save into
# Scope the semantic check the same way as the title check: a record in # dozens of searches — a duplicate past the cap is the duplicate
# project P compares only to P; a project-less (orphan) record compares # report's job, not the gate's.
# only to other orphans (orphan_only), NOT across every project — without for query in embeddings_svc.chunk_document(title, body)[:_GATE_MAX_CHUNKS]:
# this, semantic_search_notes applies no project filter when project_id # Scope the semantic check the same way as the title check: a record
# is None and would match an orphan note against any project's notes. # in project P compares only to P; a project-less (orphan) record
# compares only to other orphans (orphan_only), NOT across every
# project — without this, semantic_search_notes applies no project
# filter when project_id is None and would match an orphan note
# against any project's notes.
hits = await embeddings_svc.semantic_search_notes( hits = await embeddings_svc.semantic_search_notes(
user_id, query, project_id=project_id, is_task=is_task, user_id, query, project_id=project_id, is_task=is_task,
orphan_only=(project_id is None), orphan_only=(project_id is None),
limit=3, limit=3,
threshold=(_SNIPPET_SEMANTIC_THRESHOLD threshold=(_SNIPPET_SEMANTIC_THRESHOLD
if note_type == SNIPPET_NOTE_TYPE else _SEMANTIC_THRESHOLD), if note_type == SNIPPET_NOTE_TYPE else _SEMANTIC_THRESHOLD),
# Owner-only, deliberately: this gate BLOCKS a create and tells the # Owner-only, deliberately: this gate BLOCKS a create and tells
# caller to update the match instead. Matching someone else's record # the caller to update the match instead. Matching someone
# would refuse their write and point them at something they may not # else's record would refuse their write and point them at
# be able to edit. # something they may not be able to edit.
scope="own", scope="own",
# NOT demoted by supersession (#278). A superseded record is still a # NOT demoted by supersession (#278). A superseded record is
# duplicate of what you are about to write — the claim is that it is # still a duplicate of what you are about to write — the claim
# no longer CURRENT, not that it is gone. Demoting it here would let # is that it is no longer CURRENT, not that it is gone. Demoting
# the same note be recorded a second time, and the second copy would # it here would let the same note be recorded a second time, and
# be the one nothing warns about. # the second copy would be the one nothing warns about.
demote_superseded=False, demote_superseded=False,
) )
for score, note in hits: for score, note in hits:
# semantic_search_notes doesn't filter note_type — enforce it here so # semantic_search_notes doesn't filter note_type — enforce it
# a note doesn't shadow a task of the same wording, etc. # here so a note doesn't shadow a task of the same wording, etc.
if note.note_type == note_type: if note.note_type == note_type:
return DuplicateMatch(note.id, note.title, round(score, 3), "semantic") return DuplicateMatch(note.id, note.title, round(score, 3), "semantic")
@@ -502,15 +512,22 @@ async def find_duplicate_records(
left_note = aliased(Note, name="left_note") left_note = aliased(Note, name="left_note")
right_note = aliased(Note, name="right_note") right_note = aliased(Note, name="right_note")
distance = left.embedding.cosine_distance(right.embedding) distance = left.embedding.cosine_distance(right.embedding)
# Chunk grain (#280): a note-pair's similarity is its closest CHUNK pair —
# two records duplicate each other where their most similar sections do,
# which is the honest definition when one section of a long note restates
# another record. GROUP BY collapses the chunk cross-product to one row
# per note pair.
best = func.min(distance)
pairs: list[tuple[int, int, float]] = [] pairs: list[tuple[int, int, float]] = []
try: try:
async with async_session() as session: async with async_session() as session:
stmt = ( stmt = (
select(left.note_id, right.note_id, distance.label("distance")) select(left.note_id, right.note_id, best.label("distance"))
.select_from(left) .select_from(left)
# `<` not `!=`: each unordered pair exactly once, and it drops # `<` not `!=`: each unordered pair exactly once, and it drops
# the self-pair (distance 0) that would otherwise dominate. # the self-pairs (including cross-chunk self-pairs, which would
# otherwise flag every multi-chunk note against itself).
.join(right, left.note_id < right.note_id) .join(right, left.note_id < right.note_id)
.join(left_note, left_note.id == left.note_id) .join(left_note, left_note.id == left.note_id)
.join(right_note, right_note.id == right.note_id) .join(right_note, right_note.id == right.note_id)
@@ -523,9 +540,10 @@ async def find_duplicate_records(
# report is bounded by what merge can actually act on. # report is bounded by what merge can actually act on.
left_note.user_id == user_id, left_note.user_id == user_id,
right_note.user_id == user_id, right_note.user_id == user_id,
distance <= max_distance,
) )
.order_by(distance.asc()) .group_by(left.note_id, right.note_id)
.having(best <= max_distance)
.order_by(best.asc())
.limit(max(1, limit)) .limit(max(1, limit))
) )
rows = list((await session.execute(stmt)).all()) rows = list((await session.execute(stmt)).all())
+246 -33
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:
@@ -115,6 +124,14 @@ _SUPERSESSION_PENALTY = 0.05
# whose neighbours sit ~0.01-0.02 apart. # whose neighbours sit ~0.01-0.02 apart.
_SUPERSESSION_OVERFETCH = 3 _SUPERSESSION_OVERFETCH = 3
# Chunk rows fetched per requested result (#280). The HNSW top-k runs at CHUNK
# grain — several chunks of one strong note can occupy consecutive ranks, and
# each collapses into a single result. Four ranks of headroom per result keeps
# the top-k indexed while making it effectively impossible for collapsing to
# starve the result list: that would need every requested note to be shadowed
# by four chunks of notes ranked above it.
_CHUNK_OVERFETCH = 4
async def _apply_supersession_penalty( async def _apply_supersession_penalty(
scored: list[tuple[float, "Note"]], limit: int scored: list[tuple[float, "Note"]], limit: int
@@ -159,7 +176,7 @@ async def _apply_supersession_penalty(
def embedding_text(title: str | None, body: str | None) -> str: def embedding_text(title: str | None, body: str | None) -> str:
"""The document a record is embedded AS. """`title\\n{body}` — the atomic join every embedded document is built from.
One definition, deliberately. This was written out three times — the write One definition, deliberately. This was written out three times — the write
path (`notes.embed_note`), the recurring-task spawn, and the startup path (`notes.embed_note`), the recurring-task spawn, and the startup
@@ -168,27 +185,190 @@ def embedding_text(title: str | None, body: str | None) -> str:
recurring task embedded to a different shape than everything else would be recurring task embedded to a different shape than everything else would be
ranked against a corpus it doesn't match, and nothing would report it. ranked against a corpus it doesn't match, and nothing would report it.
It is also a PRECONDITION for changing the shape at all (#2486). Measured, Since the chunking build (#280) this is a BUILDING BLOCK, not the whole
a dev-log's vector separates from five unrelated dev-logs by 0.023 while a story: the document shape a record is embedded as is `chunk_document`
snippet's separates by 0.153 — the difference being that a snippet states below, which calls this once per chunk. Callers that want "the text this
its purpose twice in a short document, so the purpose dominates. Testing an note is embedded as" want `chunk_document`; this stays public because the
alternative shape against three copies would mean testing a shape that is two functions are one contract and the guard in test_embedding_text pins
not the one in production. both.
Whether `title\\n{body}` is the RIGHT shape is the open question. That it is
one shape is what makes the question answerable.
""" """
title = title or "" title = title or ""
body = body or "" body = body or ""
return f"{title}\n{body}".strip() if body else title return f"{title}\n{body}".strip() if body else title
async def upsert_note_embedding(note_id: int, user_id: int, text: str) -> None: # --- chunking (#280): the document shape ------------------------------------
"""Generate and persist an embedding for a note. Safe to fire-and-forget.""" #
if not text or not text.strip(): # bge-small reads at most 512 tokens and fastembed silently truncates the rest,
return # so a single whole-document vector loses everything past ~400 words — for a
# long dev-log, three quarters of the record was PERMANENTLY invisible to
# search. The fix is the document shape: one vector per meaningful chunk, and a
# record is as findable as its best-matching section.
# Bumped whenever chunk_document's output can change for the same input. Stored
# on every note_embeddings row so the startup backfill can re-embed exactly the
# notes whose stored shape is stale — a version comparison instead of the table
# wipe migrations 0067/0077 had to do.
CHUNKER_VERSION = 1
# 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
# to every chunk. Deliberately conservative: our own measurement (#2485) says
# shorter, single-topic documents embed SHARPER, so the cost of over-splitting
# is a few extra cheap vectors while the cost of under-splitting is truncation —
# the exact data loss this exists to end.
_CHUNK_CHAR_BUDGET = 1400
_HEADING_RE = None # compiled lazily below to keep re import local
def _split_sections(body: str) -> list[str]:
"""Split a markdown body at heading lines, fence-aware.
Each section is a heading line plus everything under it; text before the
first heading is its own section. Heading-looking lines inside ``` / ~~~
code fences do not split — a commented `# step` in a recorded shell snippet
is content, not structure.
"""
import re
global _HEADING_RE
if _HEADING_RE is None:
_HEADING_RE = re.compile(r"^#{1,6}\s")
sections: list[list[str]] = [[]]
in_fence = False
for line in body.splitlines():
if line.lstrip().startswith(("```", "~~~")):
in_fence = not in_fence
if not in_fence and _HEADING_RE.match(line) and sections[-1]:
sections.append([line])
else:
sections[-1].append(line)
return ["\n".join(chunk).strip() for chunk in sections if any(s.strip() for s in chunk)]
def _split_paragraphs(section: str, budget: int) -> list[str]:
"""Break one oversize section into budget-sized pieces at paragraph
boundaries, hard-splitting only a single paragraph that alone exceeds the
budget (a monster table or code block — split at line boundaries so no
content is dropped, which is the entire point of this module)."""
pieces: list[str] = []
current = ""
for para in section.split("\n\n"):
while len(para) > budget:
# Hard split: prefer the last newline inside the budget so lines
# stay whole, then the last space so words do; a clean char cut is
# the final resort for one enormous unbroken token.
cut = para.rfind("\n", 0, budget)
if cut <= 0:
cut = para.rfind(" ", 0, budget)
if cut <= 0:
cut = budget
head, para = para[:cut], para[cut:].lstrip("\n ")
if current:
pieces.append(current)
current = ""
pieces.append(head.strip())
if not para.strip():
continue
candidate = f"{current}\n\n{para}" if current else para
if len(candidate) > budget and current:
pieces.append(current)
current = para
else:
current = candidate
if current:
pieces.append(current)
return pieces
def chunk_document(title: str | None, body: str | None) -> list[str]:
"""The document(s) a record is embedded AS — one string per chunk.
The contract every retrieval surface builds on:
- A record that fits the model window yields EXACTLY ONE chunk, identical
to the historical `title\\nbody` shape — snippets and reference notes,
the corpus's sharpest records, are byte-for-byte unaffected.
- A longer record is split at markdown heading boundaries (fence-aware),
small neighbouring sections merged, oversize sections split at paragraph
boundaries, so every chunk fits the window. NOTHING is dropped: every
line of the body lands in some chunk.
- Every chunk is prefixed with the record's title — each vector carries
its own topical anchor, the property that makes snippets discriminative
(#2485). Pieces sub-split from one section also repeat that section's
heading line, so "which part of which topic" survives the split.
- An empty record yields [] (callers gate on falsiness to skip embedding).
Bump CHUNKER_VERSION when changing anything observable here.
"""
single = embedding_text(title, body)
if not single:
return []
if len(single) <= _CHUNK_CHAR_BUDGET:
return [single]
title = title or ""
# Budget for section content, net of the title prefix added to every chunk.
budget = max(200, _CHUNK_CHAR_BUDGET - len(title) - 1)
# Merge small adjacent sections upward so tiny sections don't each spend a
# vector, then split anything still over budget at paragraph boundaries.
merged: list[str] = []
for section in _split_sections(body or ""):
if merged and len(merged[-1]) + 2 + len(section) <= budget:
merged[-1] = f"{merged[-1]}\n\n{section}"
else:
merged.append(section)
chunks: list[str] = []
for section in merged:
if len(section) <= budget:
chunks.append(embedding_text(title, section))
continue
pieces = _split_paragraphs(section, budget)
first_line = section.split("\n", 1)[0]
heading = first_line if first_line.lstrip().startswith("#") else ""
for i, piece in enumerate(pieces):
# Repeat the section heading on continuation pieces so each vector
# still knows what topic it is part of.
if i > 0 and heading and not piece.startswith(heading):
piece = f"{heading}\n{piece}"
chunks.append(embedding_text(title, piece))
return chunks
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: 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
@@ -198,9 +378,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)
@@ -342,7 +532,9 @@ async def semantic_search_notes(
# penalty far smaller than the window's score spread, that case # penalty far smaller than the window's score spread, that case
# needs the true answer to be more than _SUPERSESSION_OVERFETCH # needs the true answer to be more than _SUPERSESSION_OVERFETCH
# ranks down, which no observed query comes close to. # ranks down, which no observed query comes close to.
fetch = limit * _SUPERSESSION_OVERFETCH if demote_superseded else limit fetch = limit * _CHUNK_OVERFETCH * (
_SUPERSESSION_OVERFETCH if demote_superseded else 1
)
stmt = ( stmt = (
stmt.where(distance <= max_distance) stmt.where(distance <= max_distance)
.order_by(distance.asc()) .order_by(distance.asc())
@@ -353,48 +545,69 @@ async def semantic_search_notes(
logger.warning("Failed to query note embeddings", exc_info=True) logger.warning("Failed to query note embeddings", exc_info=True)
return [] return []
# Recover similarity (1 - distance) and preserve the highest-first contract. # Collapse chunk rows to BEST-CHUNK-PER-NOTE (#280): rows arrive ordered by
scored = [(1.0 - float(dist), note) for note, dist in rows] # distance, so the first appearance of a note is its best chunk and later
# appearances are the same note matched less well. A note's relevance IS
# its best section's relevance — a query about one topic of a long record
# must find that record as strongly as if the topic were the whole record.
# Recover similarity (1 - distance); order stays highest-first.
scored: list[tuple[float, Note]] = []
seen: set[int] = set()
for note, dist in rows:
if int(note.id) in seen:
continue
seen.add(int(note.id))
scored.append((1.0 - float(dist), note))
if not demote_superseded: if not demote_superseded:
return scored[:limit] return scored[:limit]
return await _apply_supersession_penalty(scored, limit) return await _apply_supersession_penalty(scored, limit)
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
+25 -14
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
@@ -235,15 +236,25 @@ async def list_notes(
if query_vec is not None: if query_vec is not None:
from scribe.models.embedding import NoteEmbedding from scribe.models.embedding import NoteEmbedding
from scribe.services.embeddings import INTERACTIVE_SEARCH_THRESHOLD from scribe.services.embeddings import INTERACTIVE_SEARCH_THRESHOLD
distance = NoteEmbedding.embedding.cosine_distance(query_vec) # Best-chunk-per-note as a correlated MIN, not a join (#280):
sem_filter = distance <= (1.0 - INTERACTIVE_SEARCH_THRESHOLD) # a note stores one embedding row PER CHUNK, so the plain join
query = query.join( # this used to be would repeat a long note once per matching
NoteEmbedding, NoteEmbedding.note_id == Note.id # chunk — duplicated list rows and a total that counts chunks.
).where(sem_filter) # This query is filter-heavy and paginated, never HNSW-bound,
count_query = count_query.join( # so the scalar subquery costs what the join did.
NoteEmbedding, NoteEmbedding.note_id == Note.id best_distance = (
).where(sem_filter) select(
semantic_order = distance.asc() func.min(
NoteEmbedding.embedding.cosine_distance(query_vec)
)
)
.where(NoteEmbedding.note_id == Note.id)
.scalar_subquery()
)
sem_filter = best_distance <= (1.0 - INTERACTIVE_SEARCH_THRESHOLD)
query = query.where(sem_filter)
count_query = count_query.where(sem_filter)
semantic_order = best_distance.asc()
else: else:
terms = _strip_type_nouns(q) terms = _strip_type_nouns(q)
for term in terms: for term in terms:
+4 -6
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)
@@ -127,7 +125,7 @@ async def spawn_recurring_tasks() -> int:
base_date = task.due_date or now.date() base_date = task.due_date or now.date()
next_due = calculate_next_due(task.recurrence_rule, base_date) next_due = calculate_next_due(task.recurrence_rule, base_date)
try: try:
child = await create_note( await create_note(
task.user_id, task.user_id,
title=task.title, title=task.title,
body=task.body or "", body=task.body or "",
@@ -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
+260
View File
@@ -0,0 +1,260 @@
"""chunk_document — the document shape a record is embedded as (#280).
The model window is 512 tokens and fastembed truncates silently, so before
chunking, everything past ~400 words of a record was PERMANENTLY invisible to
semantic search. These tests pin the two halves of the fix's contract: short
records keep the exact historical shape (the corpus's sharpest vectors are
untouched), and long records lose NOTHING — every line of the body lands in
some chunk, each chunk inside the window budget, each carrying the title as
its topical anchor.
"""
from scribe.services.embeddings import (
_CHUNK_CHAR_BUDGET,
chunk_document,
embedding_text,
)
def _long_section(tag: str, paragraphs: int = 6, sentence: str = None) -> str:
sentence = sentence or f"This paragraph discusses {tag} in useful detail."
para = " ".join([sentence] * 6)
return "\n\n".join(f"{para} (p{i})" for i in range(paragraphs))
# --- the identity half: short records are byte-for-byte unaffected -----------
def test_a_short_record_yields_exactly_the_historical_shape():
"""Snippets and reference notes are the sharpest records in the corpus
(#2485) precisely because of this shape — chunking must not touch them."""
assert chunk_document("A title", "A short body") == [
embedding_text("A title", "A short body")
]
def test_a_bodyless_record_is_one_chunk_of_its_title():
assert chunk_document("Just a title", "") == ["Just a title"]
def test_an_empty_record_yields_no_chunks():
"""Callers gate on falsiness to skip embedding entirely."""
assert chunk_document("", "") == []
assert chunk_document(None, None) == []
def test_a_record_exactly_at_budget_stays_whole():
body = "x" * (_CHUNK_CHAR_BUDGET - len("T\n"))
assert chunk_document("T", body) == [embedding_text("T", body)]
# --- the no-lost-data half: this is what the build is FOR --------------------
def test_every_line_of_a_long_body_lands_in_some_chunk():
"""The point of #280. Before chunking, a 2,000-word dev-log's last three
quarters could not influence retrieval at all. Nothing may be dropped."""
sections = [
f"## Topic {i}\n\n{_long_section(f'topic-{i}')}" for i in range(8)
]
body = "Intro paragraph before any heading.\n\n" + "\n\n".join(sections)
chunks = chunk_document("A very long dev-log", body)
assert len(chunks) > 1
joined = "\n".join(chunks)
for line in body.splitlines():
if line.strip():
assert line.strip() in joined, f"content dropped: {line[:60]!r}"
def test_every_chunk_fits_the_window_budget():
body = "\n\n".join(_long_section(f"t{i}") for i in range(10))
for chunk in chunk_document("T", body):
assert len(chunk) <= _CHUNK_CHAR_BUDGET + len("T") + 1
def test_every_chunk_is_anchored_by_the_title():
"""Each vector must carry its own topical anchor — the property that makes
snippets discriminative. An unanchored mid-document chunk would embed as
free-floating prose about nothing in particular."""
body = "\n\n".join(
f"## Section {i}\n\n{_long_section(f'sec-{i}')}" for i in range(6)
)
chunks = chunk_document("Retrieval reference", body)
assert len(chunks) > 1
for chunk in chunks:
assert chunk.startswith("Retrieval reference\n")
# --- boundary behaviour ------------------------------------------------------
def test_sections_split_at_markdown_headings_and_stay_whole_when_they_fit():
a = "## Alpha\n\nShort alpha content."
b = "## Beta\n\n" + _long_section("beta")
c = "## Gamma\n\n" + _long_section("gamma")
chunks = chunk_document("T", f"{a}\n\n{b}\n\n{c}")
# Beta's content never shares a chunk with Gamma's heading-onward content:
# heading boundaries are chunk boundaries unless merging small sections.
for chunk in chunks:
assert not ("(p5)" in chunk and "## Gamma" in chunk and "beta" in chunk)
def test_small_adjacent_sections_merge_instead_of_each_spending_a_vector():
body = (
"\n\n".join(f"## S{i}\n\nTiny." for i in range(4))
+ "\n\n## Big\n\n"
+ _long_section("big", paragraphs=10)
)
chunks = chunk_document("T", body)
tiny_chunks = [c for c in chunks if "Tiny." in c]
assert len(tiny_chunks) == 1, "four tiny sections should share one chunk"
def test_a_heading_inside_a_code_fence_does_not_split():
"""A commented `# step` in a recorded shell snippet is content, not
structure."""
body = "Intro.\n\n```bash\n# not a heading\necho hi\n```\n\nOutro."
chunks = chunk_document("T", body)
assert chunks == [embedding_text("T", body)]
def test_pieces_subsplit_from_one_section_repeat_its_heading():
"""'Which part of which topic' must survive the split — a continuation
piece without its heading embeds as context-free prose."""
body = "## The Only Topic\n\n" + _long_section("only", paragraphs=40)
chunks = chunk_document("T", body)
assert len(chunks) > 1
for chunk in chunks:
assert "## The Only Topic" in chunk
def test_a_monster_single_paragraph_is_hard_split_not_dropped():
body = "word " * 2000 # one paragraph, no newlines to split at
chunks = chunk_document("T", body)
assert len(chunks) > 1
total_words = sum(chunk.count("word") for chunk in chunks)
assert total_words == 2000
# --- the read path: best chunk wins (#280 step 4) ----------------------------
async def test_search_collapses_chunk_rows_to_best_chunk_per_note():
"""Rows arrive at CHUNK grain ordered by distance; a note appearing via
several chunks must come back ONCE, scored by its best chunk — otherwise a
long record fills the top-k with copies of itself."""
from unittest.mock import AsyncMock, MagicMock, patch
from scribe.services import embeddings as emb
note_a, note_b = MagicMock(id=1), MagicMock(id=2)
rows = [(note_a, 0.10), (note_b, 0.20), (note_a, 0.25), (note_a, 0.30)]
result = MagicMock()
result.all.return_value = rows
session, ctx = _session_ctx()
session.execute = AsyncMock(return_value=result)
with (
patch.object(emb, "async_session", return_value=ctx),
patch.object(emb, "get_embedding", AsyncMock(return_value=[0.0] * 384)),
):
out = await emb.semantic_search_notes(
1, "a query", limit=8, demote_superseded=False
)
assert [note.id for _s, note in out] == [1, 2]
assert out[0][0] == 1.0 - 0.10 # the BEST chunk's score, not a later one
# --- 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"
+23 -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,12 @@ 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))) # near gets a SECOND, weaker chunk (sim≈0.6) — the collapse to
s.add(NoteEmbedding(note_id=far.id, user_id=user.id, embedding=_vec(0.0, 1.0))) # best-chunk-per-note (#280) is under test: near must come back once,
# at its best chunk's score, not twice.
s.add(_emb(near.id, user.id, 0, _vec(1.0)))
s.add(_emb(near.id, user.id, 1, _vec(0.6, 0.8)))
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
@@ -82,6 +100,9 @@ async def test_semantic_search_ranks_and_thresholds_via_pgvector(seeded):
assert near_id in ids assert near_id in ids
assert far_id not in ids assert far_id not in ids
assert ids[0] == near_id assert ids[0] == near_id
# Chunk collapse (#280): near has TWO chunk rows above the floor (sim≈1.0
# and ≈0.6) and must appear exactly once, at its best chunk's score.
assert ids.count(near_id) == 1
top_score = results[0][0] top_score = results[0][0]
assert top_score == pytest.approx(1.0, abs=1e-3) assert top_score == pytest.approx(1.0, abs=1e-3)
-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()
+27
View File
@@ -68,6 +68,33 @@ async def test_semantic_match_when_body_substantial():
assert dup.similarity == 0.93 assert dup.similarity == 0.93
@pytest.mark.asyncio
async def test_gate_catches_a_duplicate_hiding_in_a_later_chunk():
"""The capability #280 adds to the gate: a long candidate that duplicates
an existing record in ONE SECTION is caught, where the whole-document
query this replaces diluted exactly the section that mattered. The gate
queries once per chunk and any chunk's hit blocks."""
para = ("This section restates an existing decision in enough words to be "
"a real paragraph of content for the chunker to keep. ") * 4
body = "\n\n".join(f"## Topic {i}\n\n{para} (t{i})" for i in range(8))
from scribe.services.embeddings import chunk_document
n_chunks = len(chunk_document("Title", body))
assert n_chunks > 1, "test body must actually chunk"
hit = _fake_note(id=30, title="The existing decision", note_type="note")
# Every chunk misses except the LAST one the gate will ask about.
sem = AsyncMock(side_effect=[[] for _ in range(n_chunks - 1)] + [[(0.94, hit)]])
with patch("scribe.services.dedup.async_session",
return_value=_session_returning(None)), \
patch("scribe.services.dedup.embeddings_svc.semantic_search_notes", sem):
dup = await find_duplicate_note(
7, "Title", body=body, project_id=2, is_task=False, note_type="note",
)
assert dup is not None and dup.id == 30
assert sem.await_count == n_chunks
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_semantic_match_of_other_note_type_is_ignored(): async def test_semantic_match_of_other_note_type_is_ignored():
other = _fake_note(id=21, title="X", note_type="process") other = _fake_note(id=21, title="X", note_type="process")
+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():