From 6b5043a69c721045c932fb32de0f3a1961b3a158 Mon Sep 17 00:00:00 2001 From: Bryan Van Deusen Date: Sat, 8 Aug 2026 23:43:51 -0400 Subject: [PATCH] =?UTF-8?q?feat(embeddings):=20chunk=5Fdocument=20?= =?UTF-8?q?=E2=80=94=20the=20chunked=20document=20shape=20(#280=20step=201?= =?UTF-8?q?)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit bge-small reads 512 tokens and fastembed truncates silently, so a single whole-document vector permanently lost everything past ~400 words. The new shape: split at markdown headings (fence-aware), merge small sections, split oversize ones at paragraph boundaries, title-anchor every chunk, repeat the section heading on continuation pieces. A record that fits the window yields exactly one chunk identical to the historical title\nbody shape, so the corpus's sharpest records are byte-for-byte unaffected. CHUNKER_VERSION added so later shape changes re-embed by version comparison instead of a table wipe. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01UaYUaouG9jjhATyuxCKrQs --- src/scribe/services/embeddings.py | 159 ++++++++++++++++++++++++++++-- tests/test_chunking.py | 136 +++++++++++++++++++++++++ 2 files changed, 285 insertions(+), 10 deletions(-) create mode 100644 tests/test_chunking.py diff --git a/src/scribe/services/embeddings.py b/src/scribe/services/embeddings.py index 21e1483..5fa03f9 100644 --- a/src/scribe/services/embeddings.py +++ b/src/scribe/services/embeddings.py @@ -159,7 +159,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,21 +168,160 @@ 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 +# --- 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, text: str) -> None: """Generate and persist an embedding for a note. Safe to fire-and-forget.""" if not text or not text.strip(): diff --git a/tests/test_chunking.py b/tests/test_chunking.py new file mode 100644 index 0000000..044ea68 --- /dev/null +++ b/tests/test_chunking.py @@ -0,0 +1,136 @@ +"""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