diff --git a/studio/backend/core/rag/parsers.py b/studio/backend/core/rag/parsers.py index ba248cf9a6..9afddf1d9e 100644 --- a/studio/backend/core/rag/parsers.py +++ b/studio/backend/core/rag/parsers.py @@ -12,6 +12,7 @@ from __future__ import annotations import logging import os +import re from dataclasses import dataclass from html.parser import HTMLParser @@ -69,6 +70,39 @@ def _html(raw: str) -> list[Page]: return [_page("\n".join(parser.out), 1)] +# pymupdf4llm rebuilds text from positioned glyphs, which mangles complex-shaping +# scripts (RTL Arabic/Hebrew emerge as shaped Presentation Forms, Indic matras drop to +# U+FFFD) and can silently drop most of a heavy-RTL page. When Markdown trips these +# signals we fall back to PyMuPDF's logical-order get_text(). Thresholds mirror the chat +# extractor guard (unslothai/unsloth#5351 review). +_SHAPED_PRESENTATION_FORMS = re.compile("[\ufb1d-\ufdff\ufe70-\ufefc]") +_PDF_FALLBACK_MIN_BAD_GLYPHS = 5 +_PDF_FALLBACK_BAD_GLYPH_RATIO = 0.0005 +_PDF_INCOMPLETE_RATIO = 0.75 +_PDF_INCOMPLETE_MIN_LETTERS = 200 + + +def _markdown_corrupted(text: str) -> bool: + """True when pymupdf4llm's glyph reconstruction mangled the text: shaped RTL + Presentation Forms or U+FFFD replacements above a small floor/ratio (so a lone + legitimate shaped glyph does not force the fallback).""" + if not text: + return False + threshold = max(_PDF_FALLBACK_MIN_BAD_GLYPHS, _PDF_FALLBACK_BAD_GLYPH_RATIO * len(text)) + shaped = len(_SHAPED_PRESENTATION_FORMS.findall(text)) + return shaped > threshold or text.count("\ufffd") > threshold + + +def _markdown_incomplete(markdown: str, plain: str) -> bool: + """True when ``markdown`` holds far fewer letters than the raw ``get_text`` layer -- a + coarse guard for heavy-RTL pages pymupdf4llm silently drops without shaped glyphs.""" + plain_letters = sum(1 for c in plain if c.isalnum()) + if plain_letters < _PDF_INCOMPLETE_MIN_LETTERS: + return False + markdown_letters = sum(1 for c in markdown if c.isalnum()) + return markdown_letters < _PDF_INCOMPLETE_RATIO * plain_letters + + def _pdf_markdown(doc) -> list[str] | None: """Per-page layout-aware Markdown (tables, headings, lists) via pymupdf4llm; index i maps to page i+1. Returns None when the lib is missing, extraction fails, or the @@ -100,9 +134,19 @@ def _pdf(path: str, want_images: bool) -> tuple[list[Page], list[ParsedImage]]: try: md = _pdf_markdown(doc) if config.PDF_MARKDOWN else None for i, page in enumerate(doc): - # Prefer layout-aware Markdown (keeps tables/headings legible for retrieval); - # fall back to plain text when Markdown is off, unavailable, or empty here. - text = (md[i] if md else "") or page.get_text("text") or "" + plain = page.get_text("text") or "" + candidate = md[i] if md else "" + # Prefer layout-aware Markdown (keeps tables/headings legible for retrieval), + # but drop to PyMuPDF's logical-order text when Markdown is off/empty or when + # pymupdf4llm mangled it (RTL/Indic) or dropped most of the page. + if ( + candidate + and not _markdown_corrupted(candidate) + and not _markdown_incomplete(candidate, plain) + ): + text = candidate + else: + text = plain pages.append(_page(text, i + 1)) if want_images: for img in page.get_images(full = True): @@ -308,12 +352,61 @@ def render_pdf_pages( doc.close() +def _docx_table_rows(table) -> list[str]: + """Each row as pipe-joined cell text (the locator splits anchors on pipes). + Columns stay aligned to the layout grid (merged cells fill their spanned slots, + skipped leading/trailing grid columns become empty fields). Cells are walked in + document order so a nested table, and any text after it, flattens in place.""" + from docx.table import Table + from docx.text.paragraph import Paragraph + + rows: list[str] = [] + seen: set = set() # already emitted; dedups merges spanning columns or rows + for row in table.rows: + cells: list[str] = [""] * getattr(row, "grid_cols_before", 0) + trailing: list[str] = [] # nested rows + any post-nested text, kept in order + for cell in row.cells: + # A merged cell shares one across the columns and rows it spans: + # emit its text once, then placeholders, so columns and rows stay aligned. + if cell._tc in seen: + cells.append("") + continue + seen.add(cell._tc) + # Paragraph text before the first nested table is the aligned field; the + # nested table and anything after it flatten below the row, in order. + field: list[str] = [] + after_table = False + for item in cell.iter_inner_content(): + if isinstance(item, Table): + after_table = True + trailing.extend(_docx_table_rows(item)) + elif isinstance(item, Paragraph): + text = " ".join(item.text.split()) # collapse in-cell newlines + if text: + (trailing if after_table else field).append(text) + cells.append(" ".join(field)) # empty cells kept so columns line up + cells.extend([""] * getattr(row, "grid_cols_after", 0)) + if any(c.strip() for c in cells): + rows.append(" | ".join(cells)) + rows.extend(trailing) + return rows + + def _docx(path: str) -> list[Page]: import docx + from docx.table import Table + from docx.text.paragraph import Paragraph document = docx.Document(path) - text = "\n".join(p.text for p in document.paragraphs) - return [_page(text, None)] + lines: list[str] = [] + # Walk body content in document order: paragraphs alone drop tables entirely. + for block in document.iter_inner_content(): + if isinstance(block, Paragraph): + if block.text.strip(): + lines.append(block.text) + elif isinstance(block, Table): + lines.extend(_docx_table_rows(block)) + return [_page("\n".join(lines), None)] def parse(path: str, *, want_images: bool = False): diff --git a/studio/backend/tests/test_rag_parsing.py b/studio/backend/tests/test_rag_parsing.py index 4c46f49495..14ab0efe2e 100644 --- a/studio/backend/tests/test_rag_parsing.py +++ b/studio/backend/tests/test_rag_parsing.py @@ -86,3 +86,212 @@ def test_pdf_markdown_falls_back_when_lib_missing(tmp_path, monkeypatch): _table_pdf(pdf) pages = parsers.parse(str(pdf)) assert pages and "Quarter" in pages[0].text + + +def _long_text_pdf(path): + import pymupdf + + doc = pymupdf.open() + page = doc.new_page() + body = "The quick brown fox jumps over the lazy dog. " * 12 # >200 letters + page.insert_textbox(pymupdf.Rect(40, 40, 550, 750), body, fontsize = 11) + doc.save(str(path)) + doc.close() + + +def test_pdf_markdown_corruption_falls_back_to_plain(tmp_path, monkeypatch): + # pymupdf4llm can emit shaped RTL Presentation Forms for Arabic/Hebrew; the parser + # detects that and uses PyMuPDF's logical-order text instead of the mangled Markdown. + from core.rag import config, parsers + + monkeypatch.setattr(config, "PDF_MARKDOWN", True) + shaped = "".join(chr(c) for c in range(0xFE8D, 0xFEA0)) * 20 # heavy shaped forms + monkeypatch.setattr(parsers, "_pdf_markdown", lambda doc: [shaped] * doc.page_count) + pdf = tmp_path / "table.pdf" + _table_pdf(pdf) + text = "\n".join(p.text for p in parsers.parse(str(pdf))) + assert "Quarter" in text # real logical-order text recovered + assert not parsers._markdown_corrupted(text) # shaped garbage not carried through + + +def test_pdf_markdown_incomplete_falls_back_to_plain(tmp_path, monkeypatch): + # If pymupdf4llm silently drops most of a page, the parser prefers the fuller raw layer. + from core.rag import config, parsers + + monkeypatch.setattr(config, "PDF_MARKDOWN", True) + monkeypatch.setattr(parsers, "_pdf_markdown", lambda doc: ["x"] * doc.page_count) + pdf = tmp_path / "long.pdf" + _long_text_pdf(pdf) + text = "\n".join(p.text for p in parsers.parse(str(pdf))) + assert "quick brown fox" in text # fuller raw layer used, not the near-empty Markdown + + +def _docx_with_table(path): + import docx + + document = docx.Document() + document.add_paragraph("Intro before table.") + table = document.add_table(rows = 2, cols = 2) + table.cell(0, 0).text = "NAME" + table.cell(0, 1).text = "SCORE" + table.cell(1, 0).text = "Alice" + table.cell(1, 1).text = "97pts" + document.add_paragraph("Outro after table.") + document.save(str(path)) + + +def test_docx_extracts_table_cells(tmp_path): + # document.paragraphs alone drops tables; the parser walks body content in order so + # table cells survive (pipe-joined, which the preview locator anchors on). + pytest.importorskip("docx") + from core.rag import parsers + + docx_path = tmp_path / "t.docx" + _docx_with_table(docx_path) + text = "\n".join(p.text for p in parsers.parse(str(docx_path))) + assert all(v in text for v in ("NAME", "SCORE", "Alice", "97pts")) # cells kept + assert "Alice | 97pts" in text # row cells joined + assert text.index("Intro") < text.index("NAME") < text.index("Outro") # order kept + + +def test_docx_table_keeps_columns_and_collapses_cell_newlines(tmp_path): + # Empty cells are kept (so columns stay aligned across rows) and a cell's internal + # newlines are collapsed to spaces (so a multi-paragraph cell can't break the row). + pytest.importorskip("docx") + import docx + + from core.rag import parsers + + document = docx.Document() + table = document.add_table(rows = 2, cols = 3) + table.cell(0, 0).text = "A" + table.cell(0, 1).text = "" # empty middle cell + table.cell(0, 2).text = "C" + multiline = table.cell(1, 0) + multiline.text = "line1" + multiline.add_paragraph("line2") # cell now holds an internal newline + table.cell(1, 1).text = "mid" + table.cell(1, 2).text = "end" + path = tmp_path / "aligned.docx" + document.save(str(path)) + + text = "\n".join(p.text for p in parsers.parse(str(path))) + assert "A | | C" in text # empty cell preserved -> columns line up + assert "line1 line2 | mid | end" in text # internal newline collapsed to a space + + +def test_docx_table_merged_cell_keeps_grid_alignment(tmp_path): + # A horizontally merged cell repeats across the spanned columns: emit its text once + # then a placeholder, so the row keeps as many fields as its siblings (columns stay + # aligned) without duplicating the merged text. + pytest.importorskip("docx") + import docx + + from core.rag import parsers + + document = docx.Document() + table = document.add_table(rows = 2, cols = 3) + table.cell(0, 0).text = "WIDE" + table.cell(0, 2).text = "END" + table.cell(0, 0).merge(table.cell(0, 1)) # span the first two columns + table.cell(1, 0).text = "a" + table.cell(1, 1).text = "b" + table.cell(1, 2).text = "c" + path = tmp_path / "merged.docx" + document.save(str(path)) + + text = "\n".join(p.text for p in parsers.parse(str(path))) + assert text.count("WIDE") == 1 # merged cell not duplicated across spanned columns + assert "WIDE | | END" in text # placeholder keeps 3 fields, aligned with "a | b | c" + assert "a | b | c" in text + + +def test_docx_table_pads_omitted_grid_columns(tmp_path): + # A row that skips leading grid columns exposes the gap via grid_cols_before; pad it + # with empty fields so the value stays under the right header instead of shifting left. + pytest.importorskip("docx") + import docx + from docx.oxml.ns import qn + + from core.rag import parsers + + document = docx.Document() + table = document.add_table(rows = 2, cols = 3) + table.cell(0, 0).text = "H1" + table.cell(0, 1).text = "H2" + table.cell(0, 2).text = "H3" + tr = table.rows[1]._tr # drop the first cell and mark it skipped via + tr.remove(tr.tc_lst[0]) + trPr = tr.get_or_add_trPr() + trPr.insert(0, trPr.makeelement(qn("w:gridBefore"), {qn("w:val"): "1"})) + table.rows[1].cells[0].text = "X" # sits in column 2 + path = tmp_path / "gap.docx" + document.save(str(path)) + + text = "\n".join(p.text for p in parsers.parse(str(path))) + assert " | X | " in text # leading gap padded so X lines up under H2, not H1 + + +def test_docx_flattens_nested_table(tmp_path): + # cell.text ignores tables nested inside a cell; walk cell.tables so nested rows are + # not silently dropped from the indexed text. + pytest.importorskip("docx") + import docx + + from core.rag import parsers + + document = docx.Document() + outer = document.add_table(rows = 1, cols = 1).cell(0, 0) + outer.text = "outer" + nested = outer.add_table(rows = 1, cols = 2) + nested.cell(0, 0).text = "NESTED-A" + nested.cell(0, 1).text = "NESTED-B" + path = tmp_path / "nested.docx" + document.save(str(path)) + + text = "\n".join(p.text for p in parsers.parse(str(path))) + assert "NESTED-A | NESTED-B" in text # nested table flattened, not dropped + + +def test_docx_nested_table_keeps_in_cell_order(tmp_path): + # A cell holding paragraph, nested table, paragraph must serialize in that order + # (cell.text alone would emit both paragraphs before the nested rows). + pytest.importorskip("docx") + import docx + + from core.rag import parsers + + document = docx.Document() + cell = document.add_table(rows = 1, cols = 1).cell(0, 0) + cell.text = "before" + nested = cell.add_table(rows = 1, cols = 2) + nested.cell(0, 0).text = "NESTED-A" + nested.cell(0, 1).text = "NESTED-B" + cell.add_paragraph("after") + path = tmp_path / "nested_order.docx" + document.save(str(path)) + + text = "\n".join(p.text for p in parsers.parse(str(path))) + assert text.index("before") < text.index("NESTED-A") < text.index("after") + + +def test_docx_table_vertical_merge_emitted_once(tmp_path): + # A vertically merged cell maps every continuation row back to the origin ; + # emit it once and leave placeholders below so a row-spanning label isn't repeated. + pytest.importorskip("docx") + import docx + + from core.rag import parsers + + document = docx.Document() + table = document.add_table(rows = 3, cols = 2) + table.cell(0, 0).merge(table.cell(1, 0)).merge(table.cell(2, 0)).text = "SECTION" + table.cell(0, 1).text = "r0" + table.cell(1, 1).text = "r1" + table.cell(2, 1).text = "r2" + path = tmp_path / "vmerge.docx" + document.save(str(path)) + + text = "\n".join(p.text for p in parsers.parse(str(path))) + assert text.count("SECTION") == 1 # not repeated on each spanned row + assert "SECTION | r0" in text and " | r1" in text and " | r2" in text