Studio: break RAG chunks at figure/table caption boundaries

Dense embedders mean-pool over a whole chunk, so a 'Figure 1:' caption
buried at the end of a 500-token body chunk gets washed out by the
surrounding theory text and never surfaces for queries about that
figure. Pre-split each page's markdown at the start of every
Figure/Table caption line so the caption anchors its own chunk, which
gives both BM25 and the dense vector a focused, figure-dominated
target. Handles numbered, decimal, and appendix-style labels
(Figure 1, Figure 1.2, Figure B.1, Table 4, Fig./Tab. abbreviations).

studio/backend/core/rag/chunking.py

@@ -3,11 +3,45 @@
 
 from __future__ import annotations
 
+import re
 from dataclasses import dataclass
 from typing import Callable
 
 from .parsers import ParsedPage
 
+# Match "Figure 1:", "Figure 1.2:", "Fig. 3.", "Table 4:" etc. at line-start,
+# tolerating leading bold markers. Used to break chunks BEFORE such captions
+# so the caption ends up at the start of its own chunk — dense embeddings
+# pool over the whole chunk, so figure references buried at the end get
+# diluted by surrounding body text.
+_FIGURE_BOUNDARY_RE = re.compile(
+    # Number forms covered: "1", "12", "1.2", "B.1" (appendix-style),
+    # tolerating bold wrappers around either the label or the number.
+    r"^\**(?:Figure|Fig\.|Table|Tab\.)\s+[A-Z]?\.?\d+(?:\.\d+)?\**[\.:]",
+    re.MULTILINE | re.IGNORECASE,
+)
+
+
+def _split_at_figure_boundaries(text: str) -> list[str]:
+    """Split markdown at the start of each figure / table caption.
+
+    Each segment starts with either the original text head or a "Figure N:" /
+    "Table N:" line, so the caption anchors the embedding of the chunk it
+    lands in. Returns the original text as a single-element list when no
+    captions are found.
+    """
+    matches = list(_FIGURE_BOUNDARY_RE.finditer(text))
+    if not matches:
+        return [text]
+    segments: list[str] = []
+    last = 0
+    for m in matches:
+        if m.start() > last:
+            segments.append(text[last : m.start()])
+        last = m.start()
+    segments.append(text[last:])
+    return [s for s in segments if s.strip()]
+
 
 @dataclass(frozen = True)
 class Chunk:
@@ -124,16 +158,17 @@ def chunk_pages(
     count = token_counter or _char_token_estimate
     out: list[Chunk] = []
     for page in pages:
-        atomic = _atomic_split(page.text, separators, max_tokens, count)
-        merged = _merge(atomic, max_tokens, overlap_tokens, count)
-        for piece in merged:
-            out.append(
-                Chunk(
-                    text = piece,
-                    token_count = count(piece),
-                    page_number = page.page_number,
+        for segment in _split_at_figure_boundaries(page.text):
+            atomic = _atomic_split(segment, separators, max_tokens, count)
+            merged = _merge(atomic, max_tokens, overlap_tokens, count)
+            for piece in merged:
+                out.append(
+                    Chunk(
+                        text = piece,
+                        token_count = count(piece),
+                        page_number = page.page_number,
+                    )
                 )
-            )
     return out
 
 
@@ -169,7 +204,9 @@ def chunk_pages_with_spans(
             cursor += len(_PAGE_SEPARATOR)
     full_doc = _PAGE_SEPARATOR.join(parts)
 
-    atomic = _atomic_split(full_doc, separators, max_tokens, count)
+    atomic: list[str] = []
+    for segment in _split_at_figure_boundaries(full_doc):
+        atomic.extend(_atomic_split(segment, separators, max_tokens, count))
     merged = _merge(atomic, max_tokens, overlap_tokens, count)
 
     chunks: list[Chunk] = []