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 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),
+45 -27
View File
@@ -69,6 +69,12 @@ _SEMANTIC_THRESHOLD = 0.90
# structural signals cannot see.
_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
class DuplicateMatch:
@@ -258,39 +264,43 @@ async def find_duplicate_note(
# --- Signal 3: semantic similarity (only with a substantial body) ---
if body and len(body.strip()) >= _MIN_BODY_FOR_SEMANTIC:
# Built by the SAME function the corpus was embedded with. This one is
# the copy that mattered most and was easiest to miss: it is a QUERY
# document, compared against embedded ones. Shaped differently from the
# corpus it searches, the gate degrades silently — it still returns
# neighbours, just less apt ones, and no signal says the query and the
# index stopped agreeing (found by the guard in test_embedding_text).
query = embeddings_svc.embedding_text(title, body)
# Scope the semantic check the same way as the title check: a record 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.
# Query with the SAME chunker the corpus was embedded with (#280). This
# was the copy that mattered most and was easiest to miss: these are
# QUERY documents, compared against embedded ones — shaped differently
# from the corpus, the gate degrades silently. Chunking also makes the
# gate see what the whole-document query diluted: a long candidate that
# duplicates an existing record IN ONE SECTION now matches on that
# section. Capped so one pathological paste can't turn a save into
# dozens of searches — a duplicate past the cap is the duplicate
# report's job, not the gate's.
for query in embeddings_svc.chunk_document(title, body)[:_GATE_MAX_CHUNKS]:
# Scope the semantic check the same way as the title check: a record
# 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(
user_id, query, project_id=project_id, is_task=is_task,
orphan_only=(project_id is None),
limit=3,
threshold=(_SNIPPET_SEMANTIC_THRESHOLD
if note_type == SNIPPET_NOTE_TYPE else _SEMANTIC_THRESHOLD),
# Owner-only, deliberately: this gate BLOCKS a create and tells the
# caller to update the match instead. Matching someone else's record
# would refuse their write and point them at something they may not
# be able to edit.
# Owner-only, deliberately: this gate BLOCKS a create and tells
# the caller to update the match instead. Matching someone
# else's record would refuse their write and point them at
# something they may not be able to edit.
scope="own",
# NOT demoted by supersession (#278). A superseded record is still a
# duplicate of what you are about to write — the claim is that it is
# no longer CURRENT, not that it is gone. Demoting it here would let
# the same note be recorded a second time, and the second copy would
# be the one nothing warns about.
# NOT demoted by supersession (#278). A superseded record is
# still a duplicate of what you are about to write — the claim
# is that it is no longer CURRENT, not that it is gone. Demoting
# it here would let the same note be recorded a second time, and
# the second copy would be the one nothing warns about.
demote_superseded=False,
)
for score, note in hits:
# semantic_search_notes doesn't filter note_type — enforce it here so
# a note doesn't shadow a task of the same wording, etc.
# semantic_search_notes doesn't filter note_type — enforce it
# here so a note doesn't shadow a task of the same wording, etc.
if note.note_type == note_type:
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")
right_note = aliased(Note, name="right_note")
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]] = []
try:
async with async_session() as session:
stmt = (
select(left.note_id, right.note_id, distance.label("distance"))
select(left.note_id, right.note_id, best.label("distance"))
.select_from(left)
# `<` 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(left_note, left_note.id == left.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.
left_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))
)
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
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:
@@ -115,6 +124,14 @@ _SUPERSESSION_PENALTY = 0.05
# whose neighbours sit ~0.01-0.02 apart.
_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(
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:
"""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
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
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,
a dev-log's vector separates from five unrelated dev-logs by 0.023 while a
snippet's separates by 0.153 — the difference being that a snippet states
its purpose twice in a short document, so the purpose dominates. Testing an
alternative shape against three copies would mean testing a shape that is
not the one in production.
Whether `title\\n{body}` is the RIGHT shape is the open question. That it is
one shape is what makes the question answerable.
Since the chunking build (#280) this is a BUILDING BLOCK, not the whole
story: the document shape a record is embedded as is `chunk_document`
below, which calls this once per chunk. Callers that want "the text this
note is embedded as" want `chunk_document`; this stays public because the
two functions are one contract and the guard in test_embedding_text pins
both.
"""
title = title or ""
body = body or ""
return f"{title}\n{body}".strip() if body else title
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
# --- chunking (#280): the document shape ------------------------------------
#
# bge-small reads at most 512 tokens and fastembed silently truncates the rest,
# 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:
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
@@ -198,9 +378,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)
@@ -342,7 +532,9 @@ async def semantic_search_notes(
# penalty far smaller than the window's score spread, that case
# needs the true answer to be more than _SUPERSESSION_OVERFETCH
# 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.where(distance <= max_distance)
.order_by(distance.asc())
@@ -353,48 +545,69 @@ async def semantic_search_notes(
logger.warning("Failed to query note embeddings", exc_info=True)
return []
# Recover similarity (1 - distance) and preserve the highest-first contract.
scored = [(1.0 - float(dist), note) for note, dist in rows]
# Collapse chunk rows to BEST-CHUNK-PER-NOTE (#280): rows arrive ordered by
# 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:
return scored[:limit]
return await _apply_supersession_penalty(scored, limit)
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
+25 -14
View File
@@ -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
@@ -235,15 +236,25 @@ async def list_notes(
if query_vec is not None:
from scribe.models.embedding import NoteEmbedding
from scribe.services.embeddings import INTERACTIVE_SEARCH_THRESHOLD
distance = NoteEmbedding.embedding.cosine_distance(query_vec)
sem_filter = distance <= (1.0 - INTERACTIVE_SEARCH_THRESHOLD)
query = query.join(
NoteEmbedding, NoteEmbedding.note_id == Note.id
).where(sem_filter)
count_query = count_query.join(
NoteEmbedding, NoteEmbedding.note_id == Note.id
).where(sem_filter)
semantic_order = distance.asc()
# Best-chunk-per-note as a correlated MIN, not a join (#280):
# a note stores one embedding row PER CHUNK, so the plain join
# this used to be would repeat a long note once per matching
# chunk — duplicated list rows and a total that counts chunks.
# This query is filter-heavy and paginated, never HNSW-bound,
# so the scalar subquery costs what the join did.
best_distance = (
select(
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:
terms = _strip_type_nouns(q)
for term in terms:
+4 -6
View File
@@ -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)
@@ -127,7 +125,7 @@ async def spawn_recurring_tasks() -> int:
base_date = task.due_date or now.date()
next_due = calculate_next_due(task.recurrence_rule, base_date)
try:
child = await create_note(
await create_note(
task.user_id,
title=task.title,
body=task.body or "",
@@ -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
+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]
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,12 @@ 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)))
# near gets a SECOND, weaker chunk (sim≈0.6) — the collapse to
# 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()
ids = (user.id, near.id, far.id)
yield ids
@@ -82,6 +100,9 @@ async def test_semantic_search_ranks_and_thresholds_via_pgvector(seeded):
assert near_id in ids
assert far_id not in ids
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]
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,
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()
+27
View File
@@ -68,6 +68,33 @@ async def test_semantic_match_when_body_substantial():
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
async def test_semantic_match_of_other_note_type_is_ignored():
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()
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():