From 8145f1d5274d12ca3bdedddc3d983e8b47eb194b Mon Sep 17 00:00:00 2001
From: Roland Tannous <rolandtannous@gravityq.ai>
Date: Wed, 27 May 2026 17:02:36 +0400
Subject: [PATCH] Studio: splice VLM figure captions next to their 'Figure N:'
 line
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Captions were appended at the bottom of the page text, so the chunk
containing 'Figure 1: Asymmetries ...' got chunked separately from
'**Figure**: Flowchart with ...' on the same page. Retrieval surfaced
the caption-text chunk but the VLM description landed in a different
chunk, leaving the LLM without the visual content right next to the
figure label.

Splice each VLM caption right after the matching 'Figure N:' (or
'Table N:') line as '**Figure N description**: ...', so:

  - The figure-boundary chunker now keeps both the original in-PDF
    caption AND the VLM description in the same chunk (which starts
    with 'Figure N:').
  - Multi-figure pages get per-figure attribution — the prefix
    'Figure N description' lets the LLM tell two figures on the same
    page apart, even though the bbox renderer still emits one image
    per page today (multi-figure clustering is a follow-up).
  - When the page text has no figure lines (DOCX/HTML/TXT or rare
    PDF layouts) the old end-of-page appendix is kept as a fallback.
---
 studio/backend/core/rag/parsers/__init__.py | 88 ++++++++++++++++++---
 1 file changed, 77 insertions(+), 11 deletions(-)

diff --git a/studio/backend/core/rag/parsers/__init__.py b/studio/backend/core/rag/parsers/__init__.py
index bc0b613b1e..99e565a123 100644
--- a/studio/backend/core/rag/parsers/__init__.py
+++ b/studio/backend/core/rag/parsers/__init__.py
@@ -3,9 +3,20 @@
 
 from __future__ import annotations
 
+import re
 from dataclasses import dataclass, field
 from pathlib import Path
 
+# Same shape as the chunker's figure-boundary regex but captures the
+# figure label (Figure / Fig. / Table / Tab.) AND the number so we can
+# attribute a VLM caption back to a specific figure on a multi-figure
+# page.
+_FIGURE_LINE_RE = re.compile(
+    r"^(?P<lead>\**)(?P<label>Figure|Fig\.|Table|Tab\.)\s+"
+    r"(?P<num>[A-Z]?\.?\d+(?:\.\d+)?)(?P<trail>\**[\.:])",
+    re.MULTILINE | re.IGNORECASE,
+)
+
 
 @dataclass(frozen = True)
 class ParsedPage:
@@ -44,6 +55,63 @@ class UnsupportedFormatError(ValueError):
     pass
 
 
+def _normalize_label(raw: str) -> str:
+    head = raw.lower()
+    if head.startswith("fig"):
+        return "Figure"
+    if head.startswith("tab"):
+        return "Table"
+    return raw.capitalize()
+
+
+def _splice_inline_at_figure_lines(text: str, captions: list[str]) -> str:
+    """Splice each caption right after the matching 'Figure N:' line.
+
+    Captions are consumed in order against the figure caption lines
+    appearing in the page text. Each spliced block is prefixed with
+    "**Figure N description**:" so a retrieved chunk lets the LLM
+    distinguish between multiple figures on the same page.
+
+    Fallback when no figure caption lines exist (or fewer than we have
+    captions): leftover captions are appended at the end of the page
+    text as generic "**Figure**: ..." entries.
+    """
+    if not captions:
+        return text
+    matches = list(_FIGURE_LINE_RE.finditer(text))
+    if not matches:
+        appendix = "\n\n".join(f"**Figure**: {c}" for c in captions)
+        return f"{text}\n\n{appendix}"
+
+    parts: list[str] = []
+    cursor = 0
+    caps = iter(captions)
+    used = 0
+    for m in matches:
+        # Insertion point = end of the line containing the figure label.
+        line_end = text.find("\n", m.end())
+        if line_end == -1:
+            line_end = len(text)
+        parts.append(text[cursor:line_end])
+        try:
+            cap = next(caps)
+        except StopIteration:
+            cursor = line_end
+            continue
+        used += 1
+        label = f"{_normalize_label(m.group('label'))} {m.group('num')}"
+        parts.append(f"\n\n**{label} description**: {cap}")
+        cursor = line_end
+    parts.append(text[cursor:])
+
+    leftover = list(caps)
+    body = "".join(parts)
+    if leftover:
+        appendix = "\n\n".join(f"**Figure**: {c}" for c in leftover)
+        body = f"{body}\n\n{appendix}"
+    return body
+
+
 def inline_image_captions(
     pages: list[ParsedPage],
     images: list[ParsedImage],
@@ -51,10 +119,13 @@ def inline_image_captions(
 ) -> list[ParsedPage]:
     """Splice per-image captions into the markdown of the pages they came from.
 
-    Mirrors PR #5351's chat-composer pattern: figure captions become
-    inline text in the page markdown so the chunker indexes them like
-    any other content. Captions appear at the end of the page's text
-    block as ``**Figure**: …`` lines.
+    Each caption lands right after the page's matching ``Figure N:`` or
+    ``Table N:`` line as ``**Figure N description**: …`` so the chunker
+    keeps the VLM description adjacent to the figure's existing in-PDF
+    caption text. When a page has more figure caption lines than we have
+    captioned images, the extra figure lines are left alone; when there
+    are more captions than figure lines (or no figure lines at all),
+    leftovers fall back to an end-of-page ``**Figure**: …`` appendix.
 
     ``captions`` is parallel to ``images`` (same length, same order).
     Empty or whitespace-only captions are skipped. Images without a
@@ -64,11 +135,8 @@ def inline_image_captions(
     if not images or not captions:
         return list(pages)
     if len(captions) != len(images):
-        # Defensive: caller mismatch shouldn't happen but we don't want
-        # to lose pages over it.
         return list(pages)
 
-    # Bucket captions per page_number (None bucket → single-page docs).
     per_page: dict[int | None, list[str]] = {}
     for img, cap in zip(images, captions):
         cleaned = (cap or "").strip()
@@ -83,17 +151,15 @@ def inline_image_captions(
     null_bucket = per_page.get(None, [])
     for page in pages:
         captions_for_this = per_page.get(page.page_number, [])
-        # If this is the single-page case (no page_number) also flush
-        # the null-bucket so DOCX/HTML/TXT pick up captions correctly.
         if page.page_number is None and null_bucket:
             captions_for_this = captions_for_this + null_bucket
         if not captions_for_this:
             out.append(page)
             continue
-        appendix = "\n\n".join(f"**Figure**: {cap}" for cap in captions_for_this)
+        new_text = _splice_inline_at_figure_lines(page.text, captions_for_this)
         out.append(
             ParsedPage(
-                text = f"{page.text}\n\n{appendix}",
+                text = new_text,
                 page_number = page.page_number,
             )
         )