fix(chunk): carry a token-bounded tail so overlap applies when paragraphs exceed the budget

Author: codebybilal18

app/ingest/chunk.py

@@ -128,6 +128,41 @@ def _split_to_budget(text: str, budget_tokens: int, model: str | None = None) ->
     return [text[i : i + approx_chars] for i in range(0, len(text), approx_chars)]
 
 
+def _tail_within_budget(text: str, budget_tokens: int, model: str | None = None) -> str:
+    """The last `budget_tokens` worth of `text`, cut at a clean boundary.
+
+    Starts from a word estimate (~0.8 words per token) and trims until it fits,
+    rather than tokenising once per candidate length — a paragraph is long
+    enough that a naive scan costs hundreds of tokenizer calls per boundary.
+
+    Prefers to begin at a sentence boundary so the carried context reads as
+    prose rather than starting mid-clause, but only when that keeps most of the
+    budget: a fragment is still better overlap than none at all.
+    """
+    if budget_tokens <= 0:
+        return ""
+
+    words = text.split()
+    if not words:
+        return ""
+
+    take = min(len(words), max(1, int(budget_tokens * 0.8)))
+    tail = " ".join(words[-take:])
+    while take > 1 and count_tokens(tail, model) > budget_tokens:
+        take = max(1, int(take * 0.85))
+        tail = " ".join(words[-take:])
+
+    if count_tokens(tail, model) > budget_tokens:
+        return ""
+
+    boundary = re.search(r"(?<=[.!?])\s+", tail)
+    if boundary:
+        candidate = tail[boundary.end() :]
+        if count_tokens(candidate, model) >= budget_tokens * 0.5:
+            return candidate
+    return tail
+
+
 def _emit(buffer: list[_Segment], index: int, model: str | None) -> BuiltChunk:
     content = "\n\n".join(s.text for s in buffer).strip()
     pages = [s.page for s in buffer]
@@ -174,10 +209,30 @@ def chunk_recursive(
             carried = 0
             for prev in reversed(buffer):
                 prev_tokens = count_tokens(prev.text, model)
-                if carried + prev_tokens > overlap_tokens:
-                    break
-                carry.insert(0, prev)
-                carried += prev_tokens
+                if carried + prev_tokens <= overlap_tokens:
+                    carry.insert(0, prev)
+                    carried += prev_tokens
+                    continue
+
+                # This segment is too big to carry whole, so carry a slice of
+                # its tail instead of nothing.
+                #
+                # Carrying only whole segments looks reasonable and silently
+                # does nothing on real prose: segments are paragraphs, and a
+                # paragraph is routinely larger than the entire overlap budget
+                # (90 tokens at 600/15%). The loop then breaks on its first
+                # iteration and the "15% overlap" config produces no overlap at
+                # all — measured at 4% of boundaries before this fix.
+                remaining = overlap_tokens - carried
+                tail = _tail_within_budget(prev.text, remaining, model)
+                if tail:
+                    carry.insert(
+                        0,
+                        _Segment(text=tail, page=prev.page, heading_path=prev.heading_path),
+                    )
+                    carried += count_tokens(tail, model)
+                break
+
             buffer = carry
             buffer_tokens = carried
 

tests/test_chunking.py

@@ -12,6 +12,7 @@
 import pytest
 
 from app.config import ChunkStrategy, RetrievalConfig
+from app.generation.providers import count_tokens
 from app.db.models import ExtractionMethod
 from app.ingest.chunk import build_chunks, chunk_fixed, chunk_recursive
 from app.ingest.parse import Heading, ParsedPage
@@ -204,3 +205,81 @@ def test_ligatures_are_folded_to_ascii():
     assert clean_extracted_text("effectiveness") == "effectiveness"
     assert clean_extracted_text("difficult") == "difficult"
     assert clean_extracted_text("flow") == "flow"
+
+
+# --- overlap actually happening ---------------------------------------------
+
+
+def prose(paragraphs: int, words_per: int = 140) -> str:
+    """Paragraphs sized like real academic prose — comfortably larger than the
+    overlap budget, which is the case that used to produce no overlap."""
+    return "\n\n".join(
+        " ".join(f"p{p}w{i}" for i in range(words_per)) + "." for p in range(paragraphs)
+    )
+
+
+def carried_words(previous: str, current: str) -> int:
+    """Longest suffix of `previous` that is a literal prefix of `current`."""
+    pw, cw = previous.split(), current.split()
+    for n in range(min(len(pw), len(cw), 400), 0, -1):
+        if pw[-n:] == cw[:n]:
+            return n
+    return 0
+
+
+def test_overlap_happens_even_when_paragraphs_exceed_the_overlap_budget(recursive_config):
+    """The regression this guards.
+
+    Carrying only *whole* segments silently produced no overlap on real prose:
+    a paragraph routinely exceeds the entire overlap budget (90 tokens at
+    600/15%), so the carry loop broke on its first iteration. Config B was
+    labelled "recursive + 15% overlap" while behaving identically to no
+    overlap, which would have made the A-to-B comparison measure nothing.
+    """
+    # A production-sized budget, so the 15% overlap is ~90 tokens rather than
+    # the 18 the small test fixture would give.
+    config = recursive_config.model_copy(update={"chunk_size_tokens": 600})
+    overlap_budget = int(600 * config.chunk_overlap_ratio)
+
+    chunks = chunk_recursive([make_page(1, prose(12))], config)
+    assert len(chunks) > 2
+
+    carried = [carried_words(a.content, b.content) for a, b in pairwise(chunks)]
+    with_overlap = [c for c in carried if c > 0]
+
+    assert len(with_overlap) >= 0.8 * len(carried), (
+        f"only {len(with_overlap)}/{len(carried)} boundaries carried text — "
+        "the overlap setting is not taking effect"
+    )
+
+    # Carried spans should use a real share of the budget, not a token or two,
+    # and must never exceed it.
+    for a, b in pairwise(chunks):
+        n = carried_words(a.content, b.content)
+        if not n:
+            continue
+        tokens = count_tokens(" ".join(b.content.split()[:n]))
+        assert tokens <= overlap_budget * 1.2, f"carried {tokens} tokens > budget {overlap_budget}"
+        assert tokens >= overlap_budget * 0.25, (
+            f"carried only {tokens} tokens of a {overlap_budget}-token budget"
+        )
+
+
+def test_zero_overlap_config_carries_nothing(recursive_config):
+    no_overlap = recursive_config.model_copy(update={"chunk_overlap_ratio": 0.0})
+    chunks = chunk_recursive([make_page(1, prose(12))], no_overlap)
+
+    carried = [carried_words(a.content, b.content) for a, b in pairwise(chunks)]
+    assert not any(carried), f"overlap=0 still carried text: {carried}"
+
+
+def test_carried_overlap_respects_the_token_budget(recursive_config):
+    from app.ingest.chunk import _tail_within_budget
+    from app.generation.providers import count_tokens
+
+    text = " ".join(f"word{i}" for i in range(500)) + "."
+    for budget in (10, 45, 90):
+        tail = _tail_within_budget(text, budget)
+        assert tail, f"no tail produced for budget {budget}"
+        assert count_tokens(tail) <= budget, f"tail exceeded budget {budget}"
+        assert text.endswith(tail), "the tail must come from the end of the text"