diff --git a/studio/backend/core/rag/chunking.py b/studio/backend/core/rag/chunking.py index e2a5c06493..4178b47c8a 100644 --- a/studio/backend/core/rag/chunking.py +++ b/studio/backend/core/rag/chunking.py @@ -106,7 +106,22 @@ def chunk_pages( max_tokens: int, overlap_tokens: int, token_counter: TokenCounter | None = None, - separators: tuple[str, ...] = ("\n\n", "\n", ". ", " ", ""), + separators: tuple[str, ...] = ( + # Markdown heading boundaries first — when the parser emits + # layout-aware Markdown (PDF via pymupdf4llm, DOCX via mammoth, + # HTML via markdownify) chunks split at section breaks rather + # than mid-paragraph. Falls back to the original separators on + # plain text input where headings are absent. + "\n# ", + "\n## ", + "\n### ", + "\n#### ", + "\n\n", + "\n", + ". ", + " ", + "", + ), ) -> list[Chunk]: """Split parsed pages into overlapping chunks. diff --git a/studio/backend/core/rag/ingestion.py b/studio/backend/core/rag/ingestion.py index a6aa7c720c..226fd53129 100644 --- a/studio/backend/core/rag/ingestion.py +++ b/studio/backend/core/rag/ingestion.py @@ -60,7 +60,11 @@ def _subprocess_worker( from core.rag.parsers import parse out_queue.put({"type": "progress", "stage": "parse", "progress": 0.05}) - pages = parse(Path(stored_path)) + # want_images stays False for the text-only ingestion path; the + # multimodal path (Phase 3B-multimodal) will flip this based on + # the KB's mode. + parsed = parse(Path(stored_path), want_images = False) + pages = parsed.pages if not pages: out_queue.put({"type": "error", "error": "no extractable text in document"}) return diff --git a/studio/backend/core/rag/parsers/__init__.py b/studio/backend/core/rag/parsers/__init__.py index 105e2a89d8..8b2bcff1ae 100644 --- a/studio/backend/core/rag/parsers/__init__.py +++ b/studio/backend/core/rag/parsers/__init__.py @@ -3,21 +3,67 @@ from __future__ import annotations -from dataclasses import dataclass +from dataclasses import dataclass, field from pathlib import Path @dataclass(frozen = True) class ParsedPage: + """One page (or page-equivalent) of Markdown-rendered text from a source document. + + For PDFs `page_number` is the 1-indexed physical page. For DOCX / HTML / + TXT / MD the whole document is one ParsedPage with `page_number = None`. + + Text is expected to be Markdown — heading markers (`#`, `##`, …), + pipe-tables, and list bullets survive extraction so the chunker can + split on them. Parsers MUST emit Markdown, not bare plain text. + """ text: str page_number: int | None = None +@dataclass(frozen = True) +class ParsedImage: + """One image extracted from a source document. + + Captured only when `parse(..., want_images=True)` is set — the + multimodal ingestion path in Phase 3B-multimodal consumes these. + `nearest_caption` is best-effort paragraph-adjacency; can be empty + when no caption could be paired (the image still ingests, just + without the paired-caption chunk). + """ + image_bytes: bytes + mime_type: str + page_number: int | None = None + nearest_caption: str = "" + + +@dataclass(frozen = True) +class ParseResult: + """Result of parsing a single source document. + + `pages` is always populated; `images` is empty unless the caller + passed `want_images=True`. Iteration aliases for `pages` so legacy + code that did `for page in parse(path)` keeps working. + """ + pages: list[ParsedPage] = field(default_factory = list) + images: list[ParsedImage] = field(default_factory = list) + + def __iter__(self): + return iter(self.pages) + + def __len__(self): + return len(self.pages) + + def __bool__(self): + return bool(self.pages) or bool(self.images) + + class UnsupportedFormatError(ValueError): pass -def parse(path: Path) -> list[ParsedPage]: +def parse(path: Path, *, want_images: bool = False) -> ParseResult: suffix = path.suffix.lower() if suffix == ".pdf": from .pdf import extract @@ -29,4 +75,4 @@ def parse(path: Path) -> list[ParsedPage]: from .html import extract else: raise UnsupportedFormatError(f"Unsupported file type: {suffix}") - return extract(path) + return extract(path, want_images = want_images) diff --git a/studio/backend/core/rag/parsers/docx.py b/studio/backend/core/rag/parsers/docx.py index fedc401773..bf16a9052c 100644 --- a/studio/backend/core/rag/parsers/docx.py +++ b/studio/backend/core/rag/parsers/docx.py @@ -1,27 +1,113 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 +"""DOCX parsing via mammoth. + +Mammoth converts Word documents to Markdown while preserving Heading +styles (`# `, `## `, ...), bullet/numbered lists, tables, and basic +emphasis. This replaces the previous python-docx paragraph-concat +approach that lost all heading metadata. + +Images are captured via the `convert_image` handler when +`want_images=True`, falling back to python-docx for inline image bytes +if mammoth's docx adapter can't reach them. +""" + from __future__ import annotations +import logging +import re from pathlib import Path -from . import ParsedPage +from . import ParsedImage, ParsedPage, ParseResult + +logger = logging.getLogger(__name__) -def extract(path: Path) -> list[ParsedPage]: - from docx import Document +# Mammoth uses some default DOCX-style-name → Markdown mappings, but a +# few common variants ship with non-default names. Map them explicitly +# so we don't lose headings. +_STYLE_MAP = """ +p[style-name='Title'] => h1.title:fresh +p[style-name='Subtitle'] => h2.subtitle:fresh +p[style-name='Heading 1'] => h1:fresh +p[style-name='Heading 2'] => h2:fresh +p[style-name='Heading 3'] => h3:fresh +p[style-name='Heading 4'] => h4:fresh +p[style-name='Heading 5'] => h5:fresh +p[style-name='Heading 6'] => h6:fresh +""" - document = Document(str(path)) - parts: list[str] = [] - for paragraph in document.paragraphs: - if paragraph.text and paragraph.text.strip(): - parts.append(paragraph.text) - for table in document.tables: - for row in table.rows: - cells = [cell.text.strip() for cell in row.cells if cell.text.strip()] - if cells: - parts.append(" | ".join(cells)) - text = "\n\n".join(parts).strip() - if not text: - return [] - return [ParsedPage(text = text, page_number = None)] + +def _html_to_markdown(html: str) -> str: + """Convert mammoth's HTML output to Markdown via markdownify.""" + from markdownify import markdownify + + md = markdownify(html, heading_style = "ATX", strip = ["script", "style"]) + # markdownify can emit excessive blank lines on tables; tighten up. + md = re.sub(r"\n{3,}", "\n\n", md) + return md.strip() + + +def extract(path: Path, *, want_images: bool = False) -> ParseResult: + import mammoth + + images: list[ParsedImage] = [] + + if want_images: + # mammoth's image converter is called for every inline image. + # We capture the bytes here and substitute a stable placeholder + # in the rendered Markdown so the chunker doesn't trip over + # base64 blobs. Caption-pairing is approximate — we use the + # full document text as the caption pool (better than nothing + # for DOCX where heading→figure adjacency isn't reliable). + def _convert(image): + with image.open() as image_bytes: + blob = image_bytes.read() + mime = (image.content_type or "application/octet-stream").lower() + images.append( + ParsedImage( + image_bytes = blob, + mime_type = mime, + page_number = None, + nearest_caption = "", + ) + ) + return {"src": ""} + + convert_image = mammoth.images.img_element(_convert) + else: + # Drop image elements entirely — cheaper and avoids embedding + # base64 in Markdown when the caller doesn't want images. + convert_image = mammoth.images.img_element(lambda _image: {"src": ""}) + + with open(path, "rb") as fp: + result = mammoth.convert_to_html( + fp, + convert_image = convert_image, + style_map = _STYLE_MAP, + ) + for message in result.messages: + logger.debug("mammoth %s: %s", getattr(message, "type", "msg"), message.message) + + markdown = _html_to_markdown(result.value) + if want_images and images: + # Best-effort: every image inherits the whole doc text as a + # caption pool. Phase 3B-multimodal will improve this once + # multimodal embedders consume captions directly. + caption_pool = markdown[:1500] + images = [ + ParsedImage( + image_bytes = img.image_bytes, + mime_type = img.mime_type, + page_number = img.page_number, + nearest_caption = caption_pool, + ) + for img in images + ] + if not markdown: + return ParseResult(pages = [], images = images) + return ParseResult( + pages = [ParsedPage(text = markdown, page_number = None)], + images = images, + ) diff --git a/studio/backend/core/rag/parsers/html.py b/studio/backend/core/rag/parsers/html.py index 121cb6faf3..721e280736 100644 --- a/studio/backend/core/rag/parsers/html.py +++ b/studio/backend/core/rag/parsers/html.py @@ -1,26 +1,98 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 +"""HTML parsing via markdownify. + +Converts HTML to Markdown so headings (`

`…`

`), tables, and lists +arrive at the chunker as Markdown structure. The previous +`BeautifulSoup.get_text()` approach stripped all tags and made every +heading indistinguishable from body text. + +Image extraction (when `want_images=True`) only handles local file +references — remote URLs are skipped to avoid network calls during +ingestion. Phase 3B-multimodal can revisit this if HTML inputs with +remote images become a common pattern. +""" + from __future__ import annotations +import logging +import re from pathlib import Path +from urllib.parse import unquote, urlparse -from . import ParsedPage +from . import ParsedImage, ParsedPage, ParseResult + +logger = logging.getLogger(__name__) + +_SKIP_TAGS = ("script", "style", "noscript", "template") -_SKIP_TAGS = {"script", "style", "noscript", "template"} +def _collect_local_images(soup, html_path: Path) -> list[ParsedImage]: + images: list[ParsedImage] = [] + base_dir = html_path.parent + for tag in soup.find_all("img"): + src = tag.get("src") or "" + parsed = urlparse(src) + if parsed.scheme and parsed.scheme not in ("file", ""): + # Remote / data URLs — skip; we don't fetch over network. + continue + local_path = (base_dir / unquote(parsed.path or src)).resolve() + try: + local_path.relative_to(base_dir.resolve()) + except ValueError: + # Refuse to read outside the source's own directory. + continue + if not local_path.is_file(): + continue + try: + blob = local_path.read_bytes() + except OSError: + continue + suffix = local_path.suffix.lower().lstrip(".") + mime = { + "png": "image/png", + "jpg": "image/jpeg", + "jpeg": "image/jpeg", + "gif": "image/gif", + "webp": "image/webp", + "svg": "image/svg+xml", + }.get(suffix, f"image/{suffix or 'octet-stream'}") + caption = tag.get("alt") or tag.get("title") or "" + images.append( + ParsedImage( + image_bytes = blob, + mime_type = mime, + page_number = None, + nearest_caption = caption, + ) + ) + return images -def extract(path: Path) -> list[ParsedPage]: +def extract(path: Path, *, want_images: bool = False) -> ParseResult: from bs4 import BeautifulSoup + from markdownify import markdownify raw = path.read_bytes() soup = BeautifulSoup(raw, "lxml") - for tag in soup(_SKIP_TAGS): - tag.decompose() - text = soup.get_text(separator = "\n").strip() - lines = [line.strip() for line in text.splitlines() if line.strip()] - cleaned = "\n".join(lines) - if not cleaned: - return [] - return [ParsedPage(text = cleaned, page_number = None)] + for tag_name in _SKIP_TAGS: + for tag in soup.find_all(tag_name): + tag.decompose() + + images: list[ParsedImage] = [] + if want_images: + images = _collect_local_images(soup, path) + + md = markdownify( + str(soup), + heading_style = "ATX", + strip = list(_SKIP_TAGS), + ) + md = re.sub(r"\n{3,}", "\n\n", md).strip() + if not md: + return ParseResult(pages = [], images = images) + return ParseResult( + pages = [ParsedPage(text = md, page_number = None)], + images = images, + ) diff --git a/studio/backend/core/rag/parsers/pdf.py b/studio/backend/core/rag/parsers/pdf.py index 2a3f3fbda1..994b1aaed9 100644 --- a/studio/backend/core/rag/parsers/pdf.py +++ b/studio/backend/core/rag/parsers/pdf.py @@ -1,14 +1,101 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 +"""Layout-aware PDF parsing via pymupdf + pymupdf4llm. + +Produces Markdown per page (headings, pipe-tables, lists survive) so the +recursive chunker can split on heading boundaries. Falls back to pypdf +text-extraction only when pymupdf fails to open the file — keeps the +pipeline alive for malformed PDFs. + +Image extraction is gated behind `want_images=True` so text-only KBs +pay zero cost for images they don't index. +""" + from __future__ import annotations +import logging from pathlib import Path -from . import ParsedPage +from . import ParsedImage, ParsedPage, ParseResult + +logger = logging.getLogger(__name__) -def extract(path: Path) -> list[ParsedPage]: +def _extract_with_pymupdf(path: Path, want_images: bool) -> ParseResult: + import pymupdf + import pymupdf4llm + + doc = pymupdf.open(str(path)) + try: + pages: list[ParsedPage] = [] + for page_index in range(len(doc)): + try: + md = pymupdf4llm.to_markdown( + doc, + pages = [page_index], + write_images = False, + ignore_images = True, + show_progress = False, + ) + except Exception: + # pymupdf4llm can choke on individual pages (rare). Fall + # back to plain text extraction for just that page. + md = doc[page_index].get_text("text") or "" + md = md.strip() + if md: + pages.append(ParsedPage(text = md, page_number = page_index + 1)) + + images: list[ParsedImage] = [] + if want_images: + images = _extract_images_pymupdf(doc, pages) + return ParseResult(pages = pages, images = images) + finally: + doc.close() + + +def _extract_images_pymupdf(doc, pages: list[ParsedPage]) -> list[ParsedImage]: + """Pull embedded images and pair each with the nearest text on the same page.""" + captions_by_page: dict[int, str] = {p.page_number: p.text for p in pages if p.page_number} + out: list[ParsedImage] = [] + for page_index in range(len(doc)): + page_number = page_index + 1 + try: + image_list = doc[page_index].get_images(full = True) + except Exception: + continue + for img_info in image_list: + xref = img_info[0] + try: + extracted = doc.extract_image(xref) + except Exception: + continue + image_bytes = extracted.get("image") + ext = (extracted.get("ext") or "png").lower() + mime = { + "png": "image/png", + "jpg": "image/jpeg", + "jpeg": "image/jpeg", + "gif": "image/gif", + "webp": "image/webp", + "bmp": "image/bmp", + "tiff": "image/tiff", + }.get(ext, f"image/{ext}") + if not image_bytes: + continue + caption = (captions_by_page.get(page_number, "") or "")[:1500] + out.append( + ParsedImage( + image_bytes = image_bytes, + mime_type = mime, + page_number = page_number, + nearest_caption = caption, + ) + ) + return out + + +def _extract_with_pypdf_fallback(path: Path) -> ParseResult: from pypdf import PdfReader reader = PdfReader(str(path)) @@ -21,4 +108,17 @@ def extract(path: Path) -> list[ParsedPage]: text = text.strip() if text: pages.append(ParsedPage(text = text, page_number = index + 1)) - return pages + return ParseResult(pages = pages, images = []) + + +def extract(path: Path, *, want_images: bool = False) -> ParseResult: + try: + return _extract_with_pymupdf(path, want_images) + except Exception as exc: + logger.warning( + "pymupdf failed for %s (%s: %s); falling back to pypdf", + path, + type(exc).__name__, + exc, + ) + return _extract_with_pypdf_fallback(path) diff --git a/studio/backend/core/rag/parsers/text.py b/studio/backend/core/rag/parsers/text.py index 402b5f32f9..439a07440b 100644 --- a/studio/backend/core/rag/parsers/text.py +++ b/studio/backend/core/rag/parsers/text.py @@ -5,10 +5,11 @@ from __future__ import annotations from pathlib import Path -from . import ParsedPage +from . import ParsedPage, ParseResult -def extract(path: Path) -> list[ParsedPage]: +def extract(path: Path, *, want_images: bool = False) -> ParseResult: + # want_images is ignored — plain text / Markdown have no embedded images. raw = path.read_bytes() try: text = raw.decode("utf-8") @@ -23,5 +24,8 @@ def extract(path: Path) -> list[ParsedPage]: text = raw.decode(encoding, errors = "replace") text = text.strip() if not text: - return [] - return [ParsedPage(text = text, page_number = None)] + return ParseResult(pages = [], images = []) + return ParseResult( + pages = [ParsedPage(text = text, page_number = None)], + images = [], + ) diff --git a/studio/backend/requirements/no-torch-runtime.txt b/studio/backend/requirements/no-torch-runtime.txt index 36d1e0f5c4..ce879ccf94 100644 --- a/studio/backend/requirements/no-torch-runtime.txt +++ b/studio/backend/requirements/no-torch-runtime.txt @@ -83,6 +83,15 @@ pillow # server. bm25s persists per-scope indices to disk. qdrant-client>=1.12 bm25s>=0.2 +# RAG parsers (Phase 3A): layout-aware Markdown extraction so the chunker +# can split on real headings instead of running paragraphs together. +# pymupdf4llm preserves headings + pipe-tables; mammoth handles DOCX +# Heading styles; markdownify converts HTML //