diff --git a/studio/backend/core/chat/__init__.py b/studio/backend/core/chat/__init__.py new file mode 100644 index 0000000000..1c34339412 --- /dev/null +++ b/studio/backend/core/chat/__init__.py @@ -0,0 +1,60 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Chat-surface helpers (not core/inference or core/data_recipe). + +Exposes the document-extraction pipeline for files dropped into the chat +composer: PDF via PyMuPDF4LLM, DOCX via mammoth; PPTX unsupported. +""" + +from __future__ import annotations + +from .document_extractor import ( + DOCUMENT_EXTRACTION_AVAILABLE, + DEFAULT_DOCUMENT_VISUAL_PAYLOADS, + DocumentExtractionBusy, + DocumentExtractionCancelled, + DocumentExtractionEncrypted, + DocumentExtractionTimeout, + DocumentExtractionUnavailable, + ExtractedFigure, + ExtractResult, + _EXTRACT_CONCURRENCY, + MAX_DOCUMENT_VISUAL_PAYLOADS, + SUPPORTED_MIME_TYPES, + SUPPORTED_SUFFIXES, + _EXTRACT_SEMAPHORE, + _drain_future_exception, + document_parser_support, + document_parser_unavailable_reasons, + extract_document, +) +from .vlm_capability import ( + VlmCapability, + detect_loaded_vlm, + extract_self_base_url, +) + +__all__ = [ + "DOCUMENT_EXTRACTION_AVAILABLE", + "DEFAULT_DOCUMENT_VISUAL_PAYLOADS", + "DocumentExtractionBusy", + "DocumentExtractionCancelled", + "DocumentExtractionEncrypted", + "DocumentExtractionTimeout", + "DocumentExtractionUnavailable", + "ExtractedFigure", + "ExtractResult", + "_EXTRACT_CONCURRENCY", + "MAX_DOCUMENT_VISUAL_PAYLOADS", + "SUPPORTED_MIME_TYPES", + "SUPPORTED_SUFFIXES", + "VlmCapability", + "_EXTRACT_SEMAPHORE", + "_drain_future_exception", + "detect_loaded_vlm", + "document_parser_support", + "document_parser_unavailable_reasons", + "extract_document", + "extract_self_base_url", +] diff --git a/studio/backend/core/chat/document_extractor.py b/studio/backend/core/chat/document_extractor.py new file mode 100644 index 0000000000..e1fe8cf902 --- /dev/null +++ b/studio/backend/core/chat/document_extractor.py @@ -0,0 +1,1206 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Document extractor for the Chat composer. + +Converts raw file bytes (PDF via pymupdf4llm, DOCX via mammoth, HTML to +Markdown, text/MD as UTF-8 with replacement) into Markdown to splice into an +outgoing chat message. With a vision-capable model loaded, selected figures +are captioned through the OpenAI-compat ``/v1/chat/completions`` surface. + +No local OCR: scanned PDFs without a text layer yield near-empty Markdown; +``use_vlm_ocr`` instead renders bounded page images for VLM captioning. PPTX +is not advertised. Parser deps are checked per format, and without a vision +model captions are ``None`` with ``describe_skipped_reason`` set. +""" + +from __future__ import annotations + +import asyncio +import base64 +import inspect +import io +import logging +import math +import multiprocessing +import os +import queue +import threading +import time +from dataclasses import dataclass, field, replace +from typing import Any, Awaitable, Callable, Literal, List, Optional + +from .vlm_capability import VlmCapability, detect_loaded_vlm + + +logger = logging.getLogger(__name__) + + +SUPPORTED_MIME_TYPES = frozenset( + { + "application/pdf", + "application/vnd.openxmlformats-officedocument.wordprocessingml.document", + "application/json", + "application/x-ndjson", + "application/xml", + "application/yaml", + "application/javascript", + "text/html", + "text/markdown", + "text/plain", + "text/csv", + "text/css", + "text/javascript", + "text/xml", + "text/yaml", + } +) + +SUPPORTED_SUFFIXES = frozenset( + { + ".pdf", + ".docx", + ".html", + ".htm", + ".md", + ".txt", + ".csv", + ".json", + ".jsonl", + ".yaml", + ".yml", + ".py", + ".js", + ".jsx", + ".ts", + ".tsx", + ".go", + ".rs", + ".java", + ".c", + ".cpp", + ".h", + ".hpp", + ".cs", + ".php", + ".rb", + ".swift", + ".kt", + ".kts", + ".scala", + ".sh", + ".bash", + ".zsh", + ".ps1", + ".sql", + ".toml", + ".ini", + ".cfg", + ".log", + ".xml", + ".css", + ".scss", + } +) + + +_DESCRIBE_PROMPT = ( + "Describe this figure in <=60 words. Focus on factual content " + "(axes, labels, captions, visible text, main objects). Do not " + "speculate beyond what is visible." +) + +# Scanned/image-only pages get a transcription prompt instead of a summary so +# the page's actual text lands in the chat context rather than a description. +_OCR_PAGE_PROMPT = ( + "Transcribe everything on this page to Markdown. Preserve headings, " + "lists, and tables. Output only the transcribed content with no " + "commentary; if the page is blank, output nothing." +) + + +DEFAULT_DOCUMENT_VISUAL_PAYLOADS = 3 +MAX_DOCUMENT_VISUAL_PAYLOADS = 10 +_MAX_ENCODED_VISUALS = DEFAULT_DOCUMENT_VISUAL_PAYLOADS +_EXTRACT_TIMEOUT_SECONDS = 120 +_VLM_CAPTION_TOTAL_TIMEOUT_SECONDS = 180 +_LOCAL_VLM_CAPTION_CONCURRENCY = max( + 1, int(os.environ.get("UNSLOTH_STUDIO_LOCAL_CAPTION_CONCURRENCY", "2")) +) +_DEFAULT_VLM_CAPTION_CONCURRENCY = max( + 1, int(os.environ.get("UNSLOTH_STUDIO_CAPTION_CONCURRENCY", "3")) +) +_EXTRACT_CONCURRENCY = max(1, int(os.environ.get("UNSLOTH_STUDIO_EXTRACT_CONCURRENCY", "2"))) +_EXTRACT_SEMAPHORE = threading.BoundedSemaphore(_EXTRACT_CONCURRENCY) +# Callers park here for a slot instead of failing fast with 503; bursts drain +# naturally while stuck workers still hit _EXTRACT_TIMEOUT_SECONDS. +_EXTRACT_QUEUE_WAIT_SECONDS = max( + 0.0, + float(os.environ.get("UNSLOTH_STUDIO_EXTRACT_QUEUE_WAIT", "60")), +) +_PAGE_RENDER_DPI = max(36, int(os.environ.get("UNSLOTH_STUDIO_PAGE_RENDER_DPI", "120"))) +_MAX_PAGE_RENDER_PIXELS = 4_000_000 +# A PDF page is treated as scanned/image-only when its text layer is below this +# many non-whitespace characters and a single image covers most of the page. +# Born-digital pages already extract via pymupdf4llm, so only scanned pages are +# rendered for VLM OCR, keeping born-digital extraction render-free and fast. +_SCANNED_TEXT_MIN_CHARS = max(0, int(os.environ.get("UNSLOTH_STUDIO_SCANNED_TEXT_MIN_CHARS", "16"))) +_SCANNED_IMAGE_AREA_FRAC = 0.5 +_MIME_TO_SUFFIX = { + "application/pdf": ".pdf", + "application/vnd.openxmlformats-officedocument.wordprocessingml.document": ".docx", + "application/json": ".json", + "application/x-ndjson": ".jsonl", + "application/xml": ".xml", + "application/yaml": ".yaml", + "application/javascript": ".js", + "text/html": ".html", + "text/markdown": ".md", + "text/plain": ".txt", + "text/csv": ".csv", + "text/css": ".css", + "text/javascript": ".js", + "text/xml": ".xml", + "text/yaml": ".yaml", +} + +_PLAIN_TEXT_SUFFIXES = SUPPORTED_SUFFIXES - {".pdf", ".docx", ".html", ".htm"} + + +def _normalized_suffix(filename: str, content_type: str = "") -> str: + suffix = os.path.splitext(filename)[1].lower() + if suffix in SUPPORTED_SUFFIXES: + return suffix + mime = (content_type or "").split(";", 1)[0].strip().lower() + return _MIME_TO_SUFFIX.get(mime, suffix) + + +class DocumentExtractionUnavailable(RuntimeError): + """Extraction backend (PyMuPDF4LLM + mammoth) missing or failed to import.""" + + +class DocumentExtractionTimeout(RuntimeError): + """Raised when document parsing exceeds the 120-second worker limit.""" + + +class DocumentExtractionBusy(RuntimeError): + """Raised when the bounded document extraction worker pool is saturated.""" + + +class DocumentExtractionCancelled(RuntimeError): + """Raised when the caller cancels an in-flight extraction.""" + + +class DocumentExtractionEncrypted(RuntimeError): + """Raised when a PDF is encrypted and cannot be parsed without a password.""" + + +try: # pragma: no cover - presence depends on optional install + import pymupdf # type: ignore + import pymupdf4llm # type: ignore +except Exception as _pdf_extract_exc: # pragma: no cover + pymupdf = None # type: ignore[assignment] + pymupdf4llm = None # type: ignore[assignment] + _PDF_EXTRACTION_IMPORT_ERROR: Optional[BaseException] = _pdf_extract_exc +else: + _PDF_EXTRACTION_IMPORT_ERROR = None + +try: # pragma: no cover - presence depends on optional install + import mammoth # type: ignore +except Exception as _docx_extract_exc: # pragma: no cover + mammoth = None # type: ignore[assignment] + _DOCX_EXTRACTION_IMPORT_ERROR: Optional[BaseException] = _docx_extract_exc +else: + _DOCX_EXTRACTION_IMPORT_ERROR = None + +# Plain text / code / data formats still work when PDF/DOCX parsers are +# missing; helpers raise DocumentExtractionUnavailable only when requested. +DOCUMENT_EXTRACTION_AVAILABLE = True +_DOCUMENT_EXTRACTION_IMPORT_ERROR: Optional[BaseException] = ( + _PDF_EXTRACTION_IMPORT_ERROR or _DOCX_EXTRACTION_IMPORT_ERROR +) + + +def document_parser_support() -> dict[str, bool]: + return { + "pdf": _PDF_EXTRACTION_IMPORT_ERROR is None, + "docx": _DOCX_EXTRACTION_IMPORT_ERROR is None, + "html": True, + "text": True, + "data": True, + "code": True, + } + + +def document_parser_unavailable_reasons() -> dict[str, str]: + reasons: dict[str, str] = {} + if _PDF_EXTRACTION_IMPORT_ERROR is not None: + reasons["pdf"] = "PDF extraction requires pymupdf and pymupdf4llm." + if _DOCX_EXTRACTION_IMPORT_ERROR is not None: + reasons["docx"] = "DOCX extraction requires mammoth." + return reasons + + +@dataclass +class ExtractedFigure: + id: str + page: Optional[int] + caption: Optional[str] + error: Optional[str] = None + kind: Literal["figure", "page"] = "figure" + image_mime: Optional[str] = None + image_base64: Optional[str] = None + image_width: Optional[int] = None + image_height: Optional[int] = None + + +@dataclass +class ExtractResult: + markdown: str + figures: List[ExtractedFigure] = field(default_factory = list) + page_count: int = 0 + tokens_est: int = 0 + describe_skipped_reason: Optional[str] = None + vlm_source: Optional[str] = None + vlm_model: Optional[str] = None + image_input_available: bool = False + warnings: List[str] = field(default_factory = list) + + +ProgressCb = Callable[[dict], Awaitable[None]] + + +def _ensure_pdf_backend() -> None: + if pymupdf is None or pymupdf4llm is None: + if _PDF_EXTRACTION_IMPORT_ERROR is not None: + logger.debug( + "PDF extraction parser import failed: %s", + _PDF_EXTRACTION_IMPORT_ERROR, + ) + raise DocumentExtractionUnavailable( + "PDF extraction requires pymupdf and pymupdf4llm. Re-run Studio " + "setup to install the parser dependencies from " + "studio/backend/requirements/single-env/data-designer-deps.txt" + ) + + +def _ensure_docx_backend() -> None: + if mammoth is None: + if _DOCX_EXTRACTION_IMPORT_ERROR is not None: + logger.debug( + "DOCX extraction parser import failed: %s", + _DOCX_EXTRACTION_IMPORT_ERROR, + ) + raise DocumentExtractionUnavailable( + "DOCX extraction requires mammoth. Re-run Studio setup to install " + "the parser dependencies from " + "studio/backend/requirements/single-env/data-designer-deps.txt" + ) + + +def _estimate_tokens(text: str) -> int: + return max(0, len(text) // 4) + + +def _encode_pil_image_for_chat( + image: Any, +) -> tuple[Optional[str], Optional[int], Optional[int], Optional[str]]: + if image is None: + return None, None, None, None + try: + from PIL import Image as PILImage + + img = image.copy() + img.thumbnail((1600, 1600)) + if img.mode in ("RGBA", "LA"): + background = PILImage.new("RGB", img.size, (255, 255, 255)) + alpha = img.getchannel("A") + background.paste(img.convert("RGB"), mask = alpha) + img = background + elif img.mode != "RGB": + img = img.convert("RGB") + + out = io.BytesIO() + img.save(out, format = "JPEG", quality = 88, optimize = True) + encoded = base64.b64encode(out.getvalue()).decode("ascii") + return encoded, img.width, img.height, "image/jpeg" + except (ImportError, AttributeError, ValueError, OSError) as exc: + logger.warning("Failed to encode extracted document image", exc_info = exc) + return None, None, None, None + + +async def _describe_image_via_vlm( + *, + image_base64: str, + image_mime: str, + endpoint_url: str, + model_name: str, + authorization_header: Optional[str], + timeout_seconds: float, + prompt: str = _DESCRIBE_PROMPT, + max_tokens: int = 512, +) -> tuple[Optional[str], Optional[str]]: + try: + import httpx + except Exception as exc: + return None, f"httpx unavailable: {exc}" + + headers = {"Content-Type": "application/json"} + if authorization_header: + headers["Authorization"] = authorization_header + + data_url = f"data:{image_mime};base64,{image_base64}" + payload = { + "model": model_name, + "stream": False, + "max_tokens": max_tokens, + "temperature": 0.2, + "top_p": 0.9, + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": prompt}, + {"type": "image_url", "image_url": {"url": data_url}}, + ], + } + ], + } + try: + async with httpx.AsyncClient(timeout = timeout_seconds) as client: + response = await client.post( + endpoint_url.rstrip("/") + "/v1/chat/completions", + headers = headers, + json = payload, + ) + if response.status_code >= 400: + return None, (f"VLM caption request failed with HTTP " f"{response.status_code}") + body = response.json() + choice = (body.get("choices") or [{}])[0] + message = choice.get("message") or {} + finish_reason = choice.get("finish_reason") + + # Some templates (Gemma 3/3n GGUF, Qwen3 always-think) put the whole + # reply in reasoning_content and leave content empty; mirror the chat + # UI streaming fallback (llama_cpp._chat_completion) here. + candidates: list[Any] = [ + message.get("content"), + message.get("reasoning_content"), + message.get("text"), + ] + # Content may be OpenAI multimodal parts; join text before the check. + normalized: list[str] = [] + for raw in candidates: + if isinstance(raw, str): + if raw.strip(): + normalized.append(raw.strip()) + elif isinstance(raw, list): + parts = [ + part.get("text", "") + for part in raw + if isinstance(part, dict) and isinstance(part.get("text"), str) + ] + joined = "".join(parts).strip() + if joined: + normalized.append(joined) + + if not normalized: + logger.warning( + "VLM caption empty: finish_reason=%r message_keys=%s", + finish_reason, + list(message.keys()), + ) + return None, (f"VLM caption empty (finish_reason={finish_reason!r})") + # First non-empty of content > reasoning_content > text. + return normalized[0], None + except Exception as exc: + logger.debug("VLM caption request failed", exc_info = True) + return None, f"VLM caption request failed: {type(exc).__name__}" + + +def _build_extract_options( + *, extract_images: bool, use_vlm_ocr: bool, max_visual_payloads: int +) -> tuple[dict, list[str]]: + """Return ``(options, build_warnings)`` for the sync extract dispatcher. + + No local OCR in this build; ``use_vlm_ocr=True`` becomes a bounded + full-page visual extraction fallback for VLM captioning. + """ + build_warnings: list[str] = [] + if use_vlm_ocr: + build_warnings.append( + "Full-page OCR was requested, but this build has no local OCR " + "engine; rendered page images will be sent to the loaded vision " + "model when image description is enabled." + ) + options = { + "extract_images": bool(extract_images), + "use_vlm_ocr": bool(use_vlm_ocr), + "max_visual_payloads": max(0, max_visual_payloads), + } + return options, build_warnings + + +def _pymupdf4llm_markdown_kwargs() -> dict[str, Any]: + """Return kwargs supported by the installed pymupdf4llm.to_markdown().""" + preferred = { + "write_images": False, + "show_progress": False, + "ignore_images": True, + "table_strategy": "lines_strict", + "use_ocr": False, + "force_ocr": False, + } + try: + signature = inspect.signature(pymupdf4llm.to_markdown) + except (TypeError, ValueError): + return { + key: value for key, value in preferred.items() if key not in {"use_ocr", "force_ocr"} + } + params = signature.parameters + if any(param.kind == inspect.Parameter.VAR_KEYWORD for param in params.values()): + return preferred + return {key: value for key, value in preferred.items() if key in params} + + +def _page_is_scanned(page: Any) -> bool: + """True when a PDF page has no usable text layer but is mostly a raster + image (scanner output). Born-digital pages return False so they are never + rendered for VLM OCR, keeping text-PDF extraction render-free.""" + try: + text = page.get_text("text") + except Exception: # pragma: no cover - defensive against PyMuPDF variants + return False + if len("".join(text.split())) >= _SCANNED_TEXT_MIN_CHARS: + return False + try: + rect = getattr(page, "rect", None) + page_area = max( + float(getattr(rect, "width", 0) or 0) * float(getattr(rect, "height", 0) or 0), + 1.0, + ) + for info in page.get_image_info(): + bbox = info.get("bbox") if isinstance(info, dict) else None + if not bbox or len(bbox) < 4: + continue + width = float(bbox[2]) - float(bbox[0]) + height = float(bbox[3]) - float(bbox[1]) + if (width * height) / page_area >= _SCANNED_IMAGE_AREA_FRAC: + return True + except Exception: # pragma: no cover - fall back to a coarse image check + try: + return bool(page.get_images()) + except Exception: + return False + return False + + +def _safe_page_pixmap(page: Any) -> Any: + rect = getattr(page, "rect", None) + width_pt = max(float(getattr(rect, "width", 0) or 0), 1.0) + height_pt = max(float(getattr(rect, "height", 0) or 0), 1.0) + scale = _PAGE_RENDER_DPI / 72.0 + projected_pixels = width_pt * scale * height_pt * scale + if projected_pixels > _MAX_PAGE_RENDER_PIXELS: + scale *= math.sqrt(_MAX_PAGE_RENDER_PIXELS / projected_pixels) + scale = max(scale, 0.05) + matrix = pymupdf.Matrix(scale, scale) # type: ignore[union-attr] + return page.get_pixmap(matrix = matrix, alpha = False) + + +def _append_page_image_figure( + doc: Any, + figures_out: list[ExtractedFigure], + *, + page_index: int, + max_figures: int, + encode_image: bool = True, +) -> bool: + if len(figures_out) >= max_figures: + return False + if not encode_image: + figures_out.append( + ExtractedFigure( + id = f"page-{page_index + 1}", + page = page_index + 1, + caption = None, + error = None, + kind = "page", + ) + ) + return True + try: + from PIL import Image as PILImage + + pix = _safe_page_pixmap(doc[page_index]) + png_bytes = pix.tobytes("png") + page_image = PILImage.open(io.BytesIO(png_bytes)) + image_base64, image_width, image_height, image_mime = _encode_pil_image_for_chat(page_image) + if not image_base64: + return False + figures_out.append( + ExtractedFigure( + id = f"page-{page_index + 1}", + page = page_index + 1, + caption = None, + error = None, + kind = "page", + image_mime = image_mime, + image_base64 = image_base64, + image_width = image_width, + image_height = image_height, + ) + ) + return True + except ( + ImportError, + MemoryError, + OverflowError, + ValueError, + OSError, + RuntimeError, + ) as exc: + logger.warning( + "Failed to render page %d preview for PDF", + page_index + 1, + exc_info = exc, + ) + return False + + +def _extract_pdf( + file_bytes: bytes, max_figures: int, use_vlm_ocr: bool, max_visual_payloads: int +) -> tuple[str, list[ExtractedFigure], int, int, int]: + """PDF -> ``(markdown, figures, page_count, truncated_count, seen)`` via PyMuPDF4LLM.""" + _ensure_pdf_backend() + assert pymupdf is not None and pymupdf4llm is not None # for type-checkers + + doc = pymupdf.open(stream = file_bytes, filetype = "pdf") + try: + # is_encrypted flags any /Encrypt dict (Acrobat-distilled PDFs, scanner + # output, the Orimi test file); needs_pass is the real password signal. + # Refuse only when a password is required. + if getattr(doc, "needs_pass", False): + raise DocumentExtractionEncrypted( + "Encrypted PDF; provide a password before extracting it." + ) + markdown = pymupdf4llm.to_markdown(doc, **_pymupdf4llm_markdown_kwargs()) + + figures_out: list[ExtractedFigure] = [] + encoded_visuals = 0 + seen = 0 + truncated_count = 0 + page_count = len(doc) + + if max_figures > 0 and page_count > 0: + # use_vlm_ocr renders every page (explicit scanned mode); the default + # path renders only pages without a text layer, so born-digital PDFs + # do no page rendering and issue no VLM calls. + if use_vlm_ocr: + pages_to_render = list(range(page_count)) + else: + pages_to_render = [ + index for index in range(page_count) if _page_is_scanned(doc[index]) + ] + rendered_pages: set[int] = set() + for position, page_index in enumerate(pages_to_render): + if len(figures_out) >= max_figures: + truncated_count += len(pages_to_render) - position + break + if _append_page_image_figure( + doc, + figures_out, + page_index = page_index, + max_figures = max_figures, + encode_image = encoded_visuals < max_visual_payloads, + ): + rendered_pages.add(page_index) + if figures_out[-1].image_base64: + encoded_visuals += 1 + seen += 1 + + if not use_vlm_ocr: + try: + from PIL import Image as PILImage + for page_index in range(page_count): + # A scanned page was already emitted whole as a kind="page" + # figure; skip it so its full-page raster is not re-added + # as a duplicate kind="figure". + if page_index in rendered_pages: + continue + page = doc[page_index] + try: + images = page.get_images(full = True) + except (ValueError, RuntimeError) as exc: + logger.debug( + "page.get_images failed on page %d", + page_index + 1, + exc_info = exc, + ) + continue + for img_info in images: + xref = img_info[0] if img_info else 0 + if not xref: + continue + try: + extracted = doc.extract_image(xref) + except (ValueError, RuntimeError) as exc: + logger.debug( + "doc.extract_image failed for xref %s", + xref, + exc_info = exc, + ) + continue + if not extracted: + continue + raw_bytes = extracted.get("image") + if not raw_bytes: + continue + try: + pil_img = PILImage.open(io.BytesIO(raw_bytes)) + pil_img.load() + except (OSError, ValueError) as exc: + logger.debug( + "PIL failed to decode extracted image xref %s", + xref, + exc_info = exc, + ) + continue + if pil_img.width < 50 or pil_img.height < 50: + continue + seen += 1 + if len(figures_out) >= max_figures: + truncated_count += 1 + continue + image_base64 = None + image_width = None + image_height = None + image_mime = None + if encoded_visuals < max_visual_payloads: + ( + image_base64, + image_width, + image_height, + image_mime, + ) = _encode_pil_image_for_chat(pil_img) + if image_base64: + encoded_visuals += 1 + figures_out.append( + ExtractedFigure( + id = f"fig-{len(figures_out)}", + page = page_index + 1, + caption = None, + error = None, + kind = "figure", + image_mime = image_mime, + image_base64 = image_base64, + image_width = image_width, + image_height = image_height, + ) + ) + except ImportError as exc: + logger.warning( + "Pillow is unavailable; skipping embedded-image extraction", + exc_info = exc, + ) + + return markdown, figures_out, page_count, truncated_count, seen + finally: + try: + doc.close() + except Exception: # pragma: no cover - defensive + logger.debug("pymupdf doc.close() raised", exc_info = True) + + +def _extract_docx(file_bytes: bytes) -> tuple[str, list[ExtractedFigure], int, int, int]: + _ensure_docx_backend() + assert mammoth is not None # for type-checkers + stream = io.BytesIO(file_bytes) + result = mammoth.convert_to_markdown(stream) + markdown = result.value or "" + return markdown, [], 0, 0, 0 + + +def _extract_plaintext(file_bytes: bytes) -> tuple[str, list[ExtractedFigure], int, int, int]: + text = file_bytes.decode("utf-8", errors = "replace") + return text, [], 0, 0, 0 + + +def _extract_html(file_bytes: bytes) -> tuple[str, list[ExtractedFigure], int, int, int]: + html = file_bytes.decode("utf-8", errors = "replace") + try: + from core.inference._html_to_md import html_to_markdown + except Exception as exc: + logger.warning( + "HTML-to-Markdown converter unavailable; using raw HTML", + exc_info = exc, + ) + return html, [], 0, 0, 0 + return html_to_markdown(html), [], 0, 0, 0 + + +def _run_extract_sync( + file_bytes: bytes, + filename: str, + options: dict, + content_type: str = "", +) -> tuple[str, list[ExtractedFigure], int, int, int]: + """Sync dispatch by suffix -> ``(markdown, figures, page_count, truncated_count, seen)``.""" + suffix = _normalized_suffix(filename, content_type) + extract_images = bool(options.get("extract_images")) + use_vlm_ocr = bool(options.get("use_vlm_ocr")) + max_figures = int(options.get("max_figures", 0)) if extract_images else 0 + max_visual_payloads = int(options.get("max_visual_payloads", DEFAULT_DOCUMENT_VISUAL_PAYLOADS)) + + if suffix == ".pdf": + return _extract_pdf(file_bytes, max_figures, use_vlm_ocr, max_visual_payloads) + if suffix == ".docx": + return _extract_docx(file_bytes) + if suffix in {".html", ".htm"}: + return _extract_html(file_bytes) + if suffix in _PLAIN_TEXT_SUFFIXES: + return _extract_plaintext(file_bytes) + raise ValueError(f"Unsupported file type: {filename}") + + +_RUN_EXTRACT_SYNC_ORIGINAL = _run_extract_sync + + +def _run_extract_worker( + result_queue: Any, file_bytes: bytes, filename: str, options: dict, content_type: str +) -> None: + try: + result_queue.put(("ok", _run_extract_sync(file_bytes, filename, options, content_type))) + except DocumentExtractionUnavailable as exc: + result_queue.put(("extraction_unavailable", str(exc))) + except DocumentExtractionEncrypted as exc: + result_queue.put(("encrypted", str(exc))) + except ValueError as exc: + result_queue.put(("value_error", str(exc))) + except BaseException as exc: + result_queue.put(("error", type(exc).__name__, str(exc))) + + +def _drain_future_exception(fut: Any) -> None: + """Retrieve a future's exception so asyncio's "Future exception was never + retrieved" warning stays quiet when the awaiting task is cancelled.""" + try: + if fut.cancelled(): + return + fut.exception() + except BaseException: + # Best effort only; a drain hook must never raise. + pass + + +def _terminate_extract_process(proc: multiprocessing.Process) -> None: + if not proc.is_alive(): + return + proc.terminate() + proc.join(5) + if proc.is_alive() and hasattr(proc, "kill"): + proc.kill() + proc.join(2) + + +def _run_extract_process_sync( + file_bytes: bytes, + filename: str, + options: dict, + content_type: str, + timeout_seconds: int, + cancel_event: Optional[threading.Event] = None, +) -> tuple[str, list[ExtractedFigure], int, int, int]: + if cancel_event is not None and cancel_event.is_set(): + raise DocumentExtractionCancelled("document extraction was cancelled") + # Park up to _EXTRACT_QUEUE_WAIT_SECONDS for a slot, polling cancel_event + # so a client disconnect short-circuits instead of holding the request. + deadline = time.monotonic() + _EXTRACT_QUEUE_WAIT_SECONDS + acquired = _EXTRACT_SEMAPHORE.acquire(blocking = False) + while True: + if acquired: + break + if cancel_event is not None and cancel_event.is_set(): + raise DocumentExtractionCancelled("document extraction was cancelled") + remaining = deadline - time.monotonic() + if remaining <= 0: + break + wait = min(remaining, 0.5) + if _EXTRACT_SEMAPHORE.acquire(timeout = wait): + acquired = True + break + if not acquired: + raise DocumentExtractionBusy("document extraction is busy") + + # Everything past the acquire lives in this try/finally so the slot is + # released even if mp context / Queue / Process construction raises + # (fork-resource exhaustion, Windows EAGAIN). + result_queue = None + proc = None + try: + # fork only on Linux: ObjC runtimes (PyMuPDF/Quartz) crash under fork + # on macOS, and Windows has no fork. + import sys as _sys + + if os.name == "nt" or _sys.platform == "darwin": + mp_method = "spawn" + else: + mp_method = "fork" + ctx = multiprocessing.get_context(mp_method) + result_queue = ctx.Queue(maxsize = 1) + proc = ctx.Process( + target = _run_extract_worker, + args = ( + result_queue, + file_bytes, + filename, + options, + content_type, + ), + daemon = True, + ) + if cancel_event is not None and cancel_event.is_set(): + raise DocumentExtractionCancelled("document extraction was cancelled") + proc.start() + deadline = time.monotonic() + timeout_seconds + message = None + while message is None: + try: + message = result_queue.get(timeout = 0.1) + break + except queue.Empty: + if cancel_event is not None and cancel_event.is_set(): + _terminate_extract_process(proc) + raise DocumentExtractionCancelled("document extraction was cancelled") + if not proc.is_alive(): + # Worker may have queued its result and exited between the + # get timeout and is_alive; drain once more so a success is + # not lost. + try: + message = result_queue.get_nowait() + except queue.Empty: + pass + break + if time.monotonic() >= deadline: + _terminate_extract_process(proc) + raise DocumentExtractionTimeout( + "document parsing exceeded the 120-second worker limit" + ) + + proc.join(2) + if proc.is_alive(): + proc.terminate() + proc.join(2) + if message is None: + # Post-join drain: the worker may have exited cleanly with a + # result still queued. + try: + message = result_queue.get_nowait() + except queue.Empty: + pass + if message is None: + raise RuntimeError( + f"document extraction worker exited without a result " f"(exitcode={proc.exitcode})" + ) + + kind = message[0] + if kind == "ok": + return message[1] + if kind == "extraction_unavailable": + raise DocumentExtractionUnavailable(message[1]) + if kind == "encrypted": + raise DocumentExtractionEncrypted(message[1]) + if kind == "value_error": + raise ValueError(message[1]) + if kind == "error": + raise RuntimeError(f"{message[1]}: {message[2]}") + raise RuntimeError(f"unexpected document worker result: {kind!r}") + finally: + if proc is not None: + try: + _terminate_extract_process(proc) + except Exception: + pass + if result_queue is not None: + try: + result_queue.close() + result_queue.join_thread() + except Exception: + pass + _EXTRACT_SEMAPHORE.release() + + +async def extract_document( + file_bytes: bytes, + filename: str, + *, + content_type: str = "", + describe_images: bool = True, + use_vlm_ocr: bool = False, + max_figures: int = 40, + max_visual_payloads: int = DEFAULT_DOCUMENT_VISUAL_PAYLOADS, + vlm_timeout_seconds: float = 60.0, + capability: Optional[VlmCapability] = None, + self_base_url: Optional[str] = None, + authorization_header: Optional[str] = None, + progress_cb: Optional[ProgressCb] = None, + cancel_event: Optional[threading.Event] = None, +) -> ExtractResult: + """Extract layout-aware Markdown plus figure metadata. + + When ``describe_images`` is True and the active model is + vision-capable, the selected visual references are captioned via the + OpenAI-compat ``/v1/chat/completions`` surface after extraction. + Otherwise figures come back with ``caption=None`` and + ``describe_skipped_reason`` carries the human-readable reason. + """ + + async def _emit(**event: Any) -> None: + if cancel_event is not None and cancel_event.is_set(): + raise DocumentExtractionCancelled("document extraction was cancelled") + if progress_cb is not None: + try: + await progress_cb(event) + except Exception: + logger.debug("progress_cb raised; continuing", exc_info = True) + + max_figures = max(0, max_figures) + max_visual_payloads = max( + 0, min(max_visual_payloads, max_figures, MAX_DOCUMENT_VISUAL_PAYLOADS) + ) + cap = capability if capability is not None else detect_loaded_vlm(self_base_url) + image_input_available = bool(cap.is_vlm and cap.endpoint_url and cap.model_name) + describe_available = bool( + describe_images and cap.is_vlm and cap.endpoint_url and cap.model_name + ) + effective_describe = describe_available and max_figures > 0 and max_visual_payloads > 0 + extract_images = max_figures > 0 + + skipped_reason: Optional[str] = None + if describe_images and not effective_describe: + if describe_available and max_figures <= 0: + skipped_reason = "figure description disabled because max_figures is 0" + elif describe_available and max_visual_payloads <= 0: + skipped_reason = "figure description disabled because max_visual_payloads is 0" + else: + skipped_reason = cap.reason or "no_vlm" + + await _emit(stage = "parsing") + + options, build_warnings = _build_extract_options( + extract_images = extract_images, + use_vlm_ocr = use_vlm_ocr, + max_visual_payloads = max_visual_payloads, + ) + options["max_figures"] = max_figures + + try: + if _run_extract_sync is _RUN_EXTRACT_SYNC_ORIGINAL: + # run_in_executor (not asyncio.to_thread) so a done-callback can + # retrieve the exception even if the awaiting task is cancelled. + loop = asyncio.get_running_loop() + extract_future = loop.run_in_executor( + None, + _run_extract_process_sync, + file_bytes, + filename, + options, + content_type, + _EXTRACT_TIMEOUT_SECONDS, + cancel_event, + ) + extract_future.add_done_callback(_drain_future_exception) + ( + markdown, + figures_out, + page_count, + truncated_count, + seen, + ) = await extract_future + else: + # Tests monkeypatch _run_extract_sync; keep patched callables out + # of multiprocessing spawn. + loop = asyncio.get_running_loop() + ( + markdown, + figures_out, + page_count, + truncated_count, + seen, + ) = await asyncio.wait_for( + loop.run_in_executor( + None, + _run_extract_sync, + file_bytes, + filename, + options, + content_type, + ), + timeout = _EXTRACT_TIMEOUT_SECONDS, + ) + except asyncio.TimeoutError: + raise DocumentExtractionTimeout("document parsing exceeded the 120-second worker limit") + except DocumentExtractionTimeout: + raise + except DocumentExtractionBusy: + raise + except DocumentExtractionCancelled: + raise + except DocumentExtractionEncrypted: + raise + except DocumentExtractionUnavailable: + raise + except ValueError: + # Surface unchanged so the route maps it to 415. + raise + except Exception as exc: + logger.exception("document extraction failed for %s", filename) + raise RuntimeError("document extraction failed") from exc + + caption_deadline_hit = False + if effective_describe: + caption_concurrency = ( + _LOCAL_VLM_CAPTION_CONCURRENCY + if cap.source in {"transformers", "unsloth"} + else _DEFAULT_VLM_CAPTION_CONCURRENCY + ) + sem = asyncio.Semaphore(caption_concurrency) + + captionable_total = sum( + 1 for fig in figures_out[:max_figures] if fig.image_base64 and fig.image_mime + ) + captioned_completed = 0 + await _emit( + stage = "captioning", + current = 0, + total = captionable_total, + page = None, + total_pages = page_count, + ) + + async def _describe_one(index: int, figure: ExtractedFigure) -> None: + nonlocal captioned_completed + if figure.caption or not figure.image_base64 or not figure.image_mime: + return + if cancel_event is not None and cancel_event.is_set(): + raise DocumentExtractionCancelled("document extraction was cancelled") + async with sem: + if cancel_event is not None and cancel_event.is_set(): + raise DocumentExtractionCancelled("document extraction was cancelled") + try: + is_page = figure.kind == "page" + caption, error = await _describe_image_via_vlm( + image_base64 = figure.image_base64, + image_mime = figure.image_mime, + endpoint_url = cap.endpoint_url or "", + model_name = cap.model_name or "", + authorization_header = authorization_header, + timeout_seconds = vlm_timeout_seconds, + prompt = _OCR_PAGE_PROMPT if is_page else _DESCRIBE_PROMPT, + max_tokens = 1024 if is_page else 512, + ) + figures_out[index] = replace( + figure, + caption = caption, + error = error, + ) + except asyncio.TimeoutError as exc: + logger.warning("VLM describe timed out for figure %s", figure.id, exc_info = exc) + figures_out[index] = replace( + figure, + error = f"VLM describe timed out: {type(exc).__name__}", + ) + except Exception as exc: + logger.warning("VLM describe failed for figure %s", figure.id, exc_info = exc) + figures_out[index] = replace( + figure, + error = f"VLM describe failed: {type(exc).__name__}", + ) + finally: + captioned_completed += 1 + await _emit( + stage = "captioning", + current = captioned_completed, + total = captionable_total, + page = figure.page, + total_pages = page_count, + ) + + tasks = [ + _describe_one(index, fig) + for index, fig in enumerate(figures_out[:max_figures]) + if fig.image_base64 and fig.image_mime + ] + if tasks: + try: + caption_timeout_seconds = _VLM_CAPTION_TOTAL_TIMEOUT_SECONDS + if cap.source in {"transformers", "unsloth"}: + caption_timeout_seconds = max( + caption_timeout_seconds, + len(tasks) * vlm_timeout_seconds + 15, + ) + results = await asyncio.wait_for( + asyncio.gather(*tasks, return_exceptions = True), + timeout = caption_timeout_seconds, + ) + for result in results: + if isinstance( + result, + (DocumentExtractionCancelled, asyncio.CancelledError), + ): + raise result + except asyncio.TimeoutError: + caption_deadline_hit = True + for index, figure in enumerate(figures_out): + if figure.image_base64 and not figure.caption and not figure.error: + figures_out[index] = replace( + figure, + error = "VLM caption deadline exceeded", + ) + + warnings: List[str] = list(build_warnings) + if truncated_count > 0: + warnings.append( + f"Document has {seen} figures; showing the first {max_figures} " + f"({truncated_count} truncated)." + ) + visual_payload_count = sum(1 for figure in figures_out if figure.image_base64) + if visual_payload_count >= max_visual_payloads and len(figures_out) > visual_payload_count: + warnings.append( + f"Only the first {max_visual_payloads} visual payloads " + "were attached; remaining figure references are text-only." + ) + if effective_describe and figures_out and all(f.caption is None for f in figures_out): + error_samples: list[str] = [] + seen_errors: set[str] = set() + for figure in figures_out: + if not figure.error or figure.error in seen_errors: + continue + seen_errors.add(figure.error) + error_samples.append(f"{figure.id}: {figure.error}") + if len(error_samples) >= 3: + break + sample_suffix = " Examples: " + "; ".join(error_samples) + "." if error_samples else "" + warnings.append( + "Figure descriptions were requested but none were produced — " + "check that the loaded model accepts image inputs via /v1." + f"{sample_suffix}" + ) + if caption_deadline_hit: + warnings.append( + "Figure captioning reached the inline timeout; some image descriptions were skipped." + ) + + await _emit(stage = "done") + + return ExtractResult( + markdown = markdown, + figures = figures_out, + page_count = page_count, + tokens_est = _estimate_tokens(markdown), + describe_skipped_reason = skipped_reason, + vlm_source = cap.source, + vlm_model = cap.model_name, + image_input_available = image_input_available, + warnings = warnings, + ) diff --git a/studio/backend/core/chat/vlm_capability.py b/studio/backend/core/chat/vlm_capability.py new file mode 100644 index 0000000000..f61866c295 --- /dev/null +++ b/studio/backend/core/chat/vlm_capability.py @@ -0,0 +1,193 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Runtime probe: is the loaded model vision-capable, and at which +OpenAI-compatible endpoint? + +Unifies the three Studio backends (embedded llama-server GGUF, transformers, +Unsloth/LoRA) behind one read-only ``VlmCapability`` dataclass. Replaces the +static ``VISION_ARCHITECTURES`` allow-list, which silently excluded new +vision architectures and could not see the actually loaded model. +""" + +from __future__ import annotations + +import logging +from dataclasses import asdict, dataclass +from typing import Any, Literal, Optional +from urllib.parse import urlparse + + +logger = logging.getLogger(__name__) + + +VlmSource = Literal["gguf", "transformers", "unsloth", "none"] + + +@dataclass(frozen = True) +class VlmCapability: + """Immutable snapshot of the loaded model's image-input capability.""" + + is_vlm: bool + endpoint_url: Optional[str] + model_name: Optional[str] + source: VlmSource + reason: Optional[str] = None + + @classmethod + def none(cls, reason: str = "no model loaded") -> "VlmCapability": + return cls( + is_vlm = False, + endpoint_url = None, + model_name = None, + source = "none", + reason = reason, + ) + + def to_dict(self) -> dict: + return asdict(self) + + +def _probe_gguf(llama: Any = None) -> Optional[VlmCapability]: + if llama is None: + try: + from core.inference.llama_cpp import get_llama_cpp_backend + except Exception: # pragma: no cover - older embedding paths + return None + + try: + llama = get_llama_cpp_backend() + except Exception: + return None + + if not getattr(llama, "is_loaded", False): + return None + + base_url = getattr(llama, "base_url", None) + model_id = getattr(llama, "model_identifier", None) + is_vision = bool(getattr(llama, "is_vision", False)) + + if not base_url or not model_id: + # Half-initialised llama-server state: fall through to the + # transformers probe instead of a misleading non-vision GGUF result. + logger.debug("llama-server reports is_loaded=True but base_url / model id missing") + return None + + return VlmCapability( + is_vlm = is_vision, + endpoint_url = base_url, + model_name = model_id, + source = "gguf", + reason = None if is_vision else "gguf: model loaded, is_vision=False (no mmproj clip)", + ) + + +def _probe_transformers(self_base_url: Optional[str]) -> Optional[VlmCapability]: + try: + from core.inference import get_inference_backend + except ModuleNotFoundError as exc: + if exc.name == "core.inference" or (exc.name and exc.name.startswith("core.inference.")): + return None + logger.exception("Failed to import transformers inference backend") + return None + except ImportError: + # Other ImportError variants (circular import) mean backend + # unavailable; NameError/AttributeError propagate so real bugs are + # not masked as "no VLM loaded". + logger.exception("Failed to import transformers inference backend") + return None + + try: + ib = get_inference_backend() + except Exception: + return None + + name: Optional[str] = getattr(ib, "active_model_name", None) + if not name: + return None + + models: dict = getattr(ib, "models", {}) or {} + info: dict = models.get(name) or {} + is_vision = bool(info.get("is_vision", False)) + is_lora = bool(info.get("is_lora", False)) + source: VlmSource = "unsloth" if is_lora else "transformers" + + if not self_base_url: + return VlmCapability( + is_vlm = False, + endpoint_url = None, + model_name = name, + source = source, + reason = f"{source}: self_base_url=None (cannot self-loopback to /v1/chat/completions)", + ) + + return VlmCapability( + is_vlm = is_vision, + endpoint_url = self_base_url.rstrip("/"), + model_name = name, + source = source, + reason = None if is_vision else f"{source}: active model not marked is_vision", + ) + + +def detect_loaded_vlm( + self_base_url: Optional[str] = None, *, llama_backend: Any = None +) -> VlmCapability: + """Identify the active model and whether it can describe images. + + ``self_base_url`` only matters for transformers / Unsloth models, whose + captioning loops back through our own ``/v1/chat/completions``; GGUF + returns llama-server's URL and ignores it. + """ + gguf = _probe_gguf(llama_backend) + if gguf is not None: + return gguf + + tf = _probe_transformers(self_base_url) + if tf is not None: + return tf + + return VlmCapability.none() + + +def extract_self_base_url(request: Any) -> Optional[str]: + """Derive a trusted local base URL for the active Studio server. + + The Host header is attacker-controlled, so the origin is always + ``127.0.0.1``; only the port is discovered (run.py, then the ASGI scope, + then ``request.base_url`` as a test/embedding fallback). + """ + port: Optional[int] = None + + try: + candidate = getattr(getattr(request, "app", None), "state", None) + candidate = getattr(candidate, "server_port", None) + if isinstance(candidate, int) and candidate > 0: + port = candidate + except Exception: + port = None + + if port is None: + try: + server = getattr(request, "scope", {}).get("server") + if ( + isinstance(server, tuple) + and len(server) >= 2 + and isinstance(server[1], int) + and server[1] > 0 + ): + port = server[1] + except Exception: + port = None + + if port is None: + try: + base = str(getattr(request, "base_url", "") or "") + if not base: + return None + parsed = urlparse(base) + port = parsed.port if parsed.port is not None else 8888 + except Exception: + return None + + return f"http://127.0.0.1:{int(port)}" diff --git a/studio/backend/core/inference/__init__.py b/studio/backend/core/inference/__init__.py index 2faf70bb79..ea7d555391 100644 --- a/studio/backend/core/inference/__init__.py +++ b/studio/backend/core/inference/__init__.py @@ -1,23 +1,45 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -""" -Inference submodule - backend for model loading and generation. +"""Inference submodule - backend for model loading and generation. -The default get_inference_backend() returns an InferenceOrchestrator that -delegates to a subprocess. The original InferenceBackend runs inside the -subprocess and can be imported directly from .inference when needed. +get_inference_backend() returns an InferenceOrchestrator that delegates to a +subprocess; the original InferenceBackend runs inside it and can be imported +from .inference directly. + +Symbols are lazy (PEP 562 ``__getattr__``) so importing a stdlib-only helper +(e.g. ``core.inference._html_to_md``) never pulls in the orchestrator or the +GGUF backend - the document-extractor HTML path must work even when the +inference extras are broken or missing. """ -from .orchestrator import InferenceOrchestrator, get_inference_backend -from .llama_cpp import LlamaCppBackend - -# Expose InferenceOrchestrator as InferenceBackend for backward compat. -InferenceBackend = InferenceOrchestrator +from typing import Any __all__ = [ "InferenceBackend", "InferenceOrchestrator", "get_inference_backend", + "get_llama_cpp_backend", "LlamaCppBackend", ] + + +def __getattr__(name: str) -> Any: + if name in ("InferenceOrchestrator", "get_inference_backend", "InferenceBackend"): + from .orchestrator import InferenceOrchestrator, get_inference_backend + + globals()["InferenceOrchestrator"] = InferenceOrchestrator + globals()["get_inference_backend"] = get_inference_backend + globals()["InferenceBackend"] = InferenceOrchestrator + return globals()[name] + if name in ("LlamaCppBackend", "get_llama_cpp_backend"): + from .llama_cpp import LlamaCppBackend, get_llama_cpp_backend + + globals()["LlamaCppBackend"] = LlamaCppBackend + globals()["get_llama_cpp_backend"] = get_llama_cpp_backend + return globals()[name] + raise AttributeError(name) + + +def __dir__() -> list[str]: + return sorted(set(globals()) | set(__all__)) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 152a3f19b2..81b0e14b21 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -1376,6 +1376,10 @@ class LlamaCppBackend: def base_url(self) -> str: return f"http://127.0.0.1:{self._port}" + @property + def api_key(self) -> Optional[str]: + return self._api_key + @property def _auth_headers(self) -> "Optional[dict[str, str]]": """Bearer header matching the --api-key direct-stream mode uses, else @@ -9121,3 +9125,19 @@ class LlamaCppBackend: return LlamaCppBackend._codec_mgr.decode( audio_type, device, token_ids = token_ids, text = data.get("content", "") ) + + +_llama_cpp_backend: Optional[LlamaCppBackend] = None + + +def get_llama_cpp_backend() -> LlamaCppBackend: + """Return the process-wide GGUF llama-server backend. + + Lives in core.inference so core helpers (core.chat.detect_loaded_vlm) + need no route imports; lazy so model-helper imports get no subprocess + cleanup side effects. + """ + global _llama_cpp_backend + if _llama_cpp_backend is None: + _llama_cpp_backend = LlamaCppBackend() + return _llama_cpp_backend diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 26825a472e..77a4de3a52 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -1681,3 +1681,135 @@ class AnthropicMessagesResponse(BaseModel): stop_reason: Optional[str] = None stop_sequence: Optional[str] = None usage: AnthropicUsage = Field(default_factory = AnthropicUsage) + + +# ---------------------------------------------------------------------- # +# Chat document extraction (parsed documents + optional VLM captions) # +# ---------------------------------------------------------------------- # + + +class ExtractedFigureModel(BaseModel): + """An extracted visual reference, optionally captioned by the loaded VLM.""" + + id: str = Field(..., description = "Stable id (e.g. 'fig-0')") + page: Optional[int] = Field(None, description = "1-based page number, if known") + caption: Optional[str] = Field( + None, description = "Short VLM-generated caption, or null if skipped/failed" + ) + error: Optional[str] = Field(None, description = "Reason the describe call failed, if any") + kind: Literal["figure", "page"] = Field( + "figure", + description = "Whether this reference is a detected figure or page image", + ) + image_mime: Optional[str] = Field( + None, description = "MIME type for image_base64 when a visual payload is present" + ) + image_base64: Optional[str] = Field( + None, + description = ( + "Base64-encoded visual payload for this reference. The first visual " + "reference is sent to vision-capable chat models as [Image #1]." + ), + ) + image_width: Optional[int] = Field(None, ge = 1, description = "Width of image_base64 after resize") + image_height: Optional[int] = Field( + None, ge = 1, description = "Height of image_base64 after resize" + ) + + +class ExtractDocumentResponse(BaseModel): + """Sync response of ``POST /chat/extract-document`` (or the final SSE event).""" + + schema_version: int = Field(1, description = "Document extraction payload schema version") + filename: str = Field(..., description = "Original filename uploaded") + markdown: str = Field(..., description = "Layout-aware Markdown extracted from the document") + page_count: int = Field(0, ge = 0, description = "Number of pages in the source") + tokens_est: int = Field(0, ge = 0, description = "Rough char/4 token estimate for the markdown") + truncated: bool = Field( + False, + description = "Whether markdown was clipped to the requested token budget", + ) + figures: List[ExtractedFigureModel] = Field( + default_factory = list, + description = "Figures discovered in the document (captions optional)", + ) + describe_skipped_reason: Optional[str] = Field( + None, + description = ( + "If image description was requested but skipped, the reason " + "(e.g. 'loaded GGUF is not vision-capable'). Mirrors the " + "``reason`` surfaced by /chat/document-support." + ), + ) + vlm_source: Optional[str] = Field( + None, + description = ( + "Which inference backend served the describe calls: 'gguf', " + "'transformers', 'unsloth', or 'none' when no VLM was used." + ), + ) + vlm_model: Optional[str] = Field( + None, + description = "Identifier of the VLM whose captions appear in this document", + ) + image_input_available: bool = Field( + False, + description = ( + "Whether the active model can receive an extracted visual payload " + "alongside the markdown." + ), + ) + warnings: List[str] = Field( + default_factory = list, + description = "Non-fatal warnings surfaced to the UI", + ) + + +class VlmCapabilityModel(BaseModel): + """Runtime probe result for the currently-loaded model.""" + + is_vlm: bool = Field(..., description = "Whether the active model accepts image inputs") + endpoint_url: Optional[str] = Field( + None, + description = "Root URL serving /v1/chat/completions for the active model", + ) + model_name: Optional[str] = Field( + None, description = "Identifier of the active model, if any is loaded" + ) + source: Literal["gguf", "transformers", "unsloth", "none"] = Field( + ..., description = "Which backend currently owns the active model" + ) + reason: Optional[str] = Field( + None, + description = "Populated when is_vlm is false; explains why the UI toggle is disabled", + ) + + +class DocumentSupportResponse(BaseModel): + """GET /chat/document-support response; drives the Chat settings toggles. + ``max_visual_payloads`` is an informational hint, not a hard cap.""" + + schema_version: int = Field(1, description = "Document support payload schema version") + extraction_available: bool = Field( + ..., + description = ("Whether the document extraction backend successfully imported on the server"), + ) + max_visual_payloads: int = Field( + ..., + ge = 0, + description = "Legacy visual-payload hint; not a hard request cap", + ) + max_extract_concurrency: int = Field( + 1, + ge = 1, + description = "Maximum server-side document extraction workers", + ) + format_support: Dict[str, bool] = Field( + default_factory = dict, + description = "Per-format parser availability for document extraction", + ) + unavailable_formats: Dict[str, str] = Field( + default_factory = dict, + description = "Per-format parser unavailability reasons", + ) + vlm: VlmCapabilityModel diff --git a/studio/backend/requirements/studio.txt b/studio/backend/requirements/studio.txt index 1b7f7a668c..e3d0c78789 100644 --- a/studio/backend/requirements/studio.txt +++ b/studio/backend/requirements/studio.txt @@ -14,6 +14,16 @@ huggingface-hub==0.36.2 structlog>=24.1.0 diceware ddgs +pypdf>=6.0.0,<7 +python-multipart>=0.0.26 +# Document extraction relies on pymupdf4llm 1.27+ (installed via +# data-designer-deps.txt), which pulls pymupdf-layout. The bundled ONNX +# models work fine on modern onnxruntime; we require >=1.19 because +# earlier wheels (e.g. 1.17.x) were built against NumPy 1.x and crash +# on import in venvs that have NumPy 2.x installed (pymupdf.layout -> +# onnxruntime -> numpy._multiarray_umath ABI mismatch). Verified +# end-to-end with onnxruntime 1.25.0 + numpy 2.4.x. +onnxruntime>=1.19 cryptography>=42.0.0 boto3>=1.34.0 # optional: S3 dataset loading httpx>=0.27.0 diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 8b0981cd2b..e210836624 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -19,6 +19,8 @@ import httpx from loggers import get_logger import asyncio import threading +from contextlib import suppress +from dataclasses import asdict as _asdict import re as _re @@ -1092,6 +1094,9 @@ from models.inference import ( ListOpenAIContainersResponse, OpenAIContainerRequest, OpenAIContainerSummary, + DocumentSupportResponse, + ExtractDocumentResponse, + ExtractedFigureModel, ) from core.inference.anthropic_compat import ( anthropic_messages_to_openai, @@ -10025,3 +10030,749 @@ async def _openai_passthrough_non_streaming( # redundant parse + re-serialize round-trip. return Response(content = resp.content, media_type = "application/json") return JSONResponse(content = data) + + +# ---------------------------------------------------------------------- # +# Chat document extraction (PyMuPDF4LLM + optional VLM image description)# +# ---------------------------------------------------------------------- # + +try: + from core.chat import ( + DOCUMENT_EXTRACTION_AVAILABLE as _DOCUMENT_EXTRACTION_AVAILABLE, + DEFAULT_DOCUMENT_VISUAL_PAYLOADS as _DEFAULT_DOCUMENT_VISUAL_PAYLOADS, + DocumentExtractionBusy as _DocumentExtractionBusy, + DocumentExtractionCancelled as _DocumentExtractionCancelled, + DocumentExtractionEncrypted as _DocumentExtractionEncrypted, + DocumentExtractionTimeout as _DocumentExtractionTimeout, + DocumentExtractionUnavailable as _DocumentExtractionUnavailable, + _EXTRACT_CONCURRENCY as _DOCUMENT_EXTRACT_CONCURRENCY, + MAX_DOCUMENT_VISUAL_PAYLOADS as _MAX_DOCUMENT_VISUAL_PAYLOADS, + SUPPORTED_MIME_TYPES as _DOC_MIME_OK, + SUPPORTED_SUFFIXES as _DOC_SUFFIX_OK, + VlmCapability as _VlmCapability, + _drain_future_exception as _drain_doc_future_exception, + detect_loaded_vlm as _detect_loaded_vlm, + document_parser_support as _document_parser_support, + document_parser_unavailable_reasons as _document_parser_unavailable_reasons, + extract_document as _extract_document, + extract_self_base_url as _extract_self_base_url, + ) +except ImportError: # pragma: no cover - package always installed alongside + _DOCUMENT_EXTRACTION_AVAILABLE = False + _DEFAULT_DOCUMENT_VISUAL_PAYLOADS = 0 + _DOCUMENT_EXTRACT_CONCURRENCY = 1 + _MAX_DOCUMENT_VISUAL_PAYLOADS = 0 + _DOC_MIME_OK = frozenset() + _DOC_SUFFIX_OK = frozenset() + _detect_loaded_vlm = None # type: ignore[assignment] + _extract_document = None # type: ignore[assignment] + _extract_self_base_url = None # type: ignore[assignment] + _document_parser_support = lambda: {} # type: ignore[assignment] + _document_parser_unavailable_reasons = lambda: {} # type: ignore[assignment] + _VlmCapability = None # type: ignore[assignment] + _drain_doc_future_exception = lambda _f: None # type: ignore[assignment] + + class _DocumentExtractionUnavailable(RuntimeError): # type: ignore[no-redef] + pass + + class _DocumentExtractionTimeout(RuntimeError): # type: ignore[no-redef] + pass + + class _DocumentExtractionBusy(RuntimeError): # type: ignore[no-redef] + pass + + class _DocumentExtractionCancelled(RuntimeError): # type: ignore[no-redef] + pass + + class _DocumentExtractionEncrypted(RuntimeError): # type: ignore[no-redef] + pass + + +_EXTRACT_MAX_BYTES = 100 * 1024 * 1024 +_EXTRACT_MULTIPART_OVERHEAD_BYTES = 1024 * 1024 +_EXTRACT_READ_CHUNK_BYTES = 64 * 1024 +_EXTRACT_MAX_PAGES_INLINE = 200 +_EXTRACT_TOKEN_BUDGET_DEFAULT = 8000 +_EXTRACT_TOKEN_BUDGET_MIN = 0 + +# Caught together by the extract endpoint; dispatched to a status/detail below. +_DOC_EXTRACTION_HTTP_ERRORS = ( + _DocumentExtractionUnavailable, + _DocumentExtractionTimeout, + _DocumentExtractionBusy, + _DocumentExtractionCancelled, + _DocumentExtractionEncrypted, + ValueError, +) + + +def _doc_exc_to_status_detail(exc: BaseException) -> tuple[int, str]: + """Map an extraction failure to (status, detail); shared by the JSON and NDJSON paths.""" + if isinstance(exc, _DocumentExtractionUnavailable): + return 501, str(exc) + if isinstance(exc, _DocumentExtractionTimeout): + return 504, "Document parsing timed out after 120s before image captioning" + if isinstance(exc, _DocumentExtractionBusy): + return 503, "Document extraction is busy" + if isinstance(exc, _DocumentExtractionCancelled): + return 499, "Client closed request" + if isinstance(exc, _DocumentExtractionEncrypted): + return 422, str(exc) + detail = str(exc) # ValueError + return (415 if detail.lower().startswith("unsupported file type") else 400), detail + + +def _ndjson_error(status_code: int, detail: str) -> str: + return json.dumps({"stage": "error", "status_code": status_code, "detail": detail}) + "\n" + + +def _page_limit_detail(page_count: int) -> str: + return ( + f"Document has {page_count} pages; inline extraction " + f"is capped at {_EXTRACT_MAX_PAGES_INLINE}. Split into smaller " + f"documents or reduce the page range." + ) + + +async def _drain_cancelled_extraction(cancel_event, extraction_task) -> None: + """Signal cancel, give the worker 10s to unwind, then force-cancel into the 499 path.""" + cancel_event.set() + with suppress( + _DocumentExtractionCancelled, + asyncio.CancelledError, + asyncio.TimeoutError, + ): + await asyncio.wait_for(asyncio.shield(extraction_task), timeout = 10) + if not extraction_task.done(): + extraction_task.cancel() + raise _DocumentExtractionCancelled("document extraction was cancelled") + + +_DOCX_MIME = "application/vnd.openxmlformats-officedocument.wordprocessingml.document" +_HTML_MIME_TYPES = {"text/html"} +_DATA_MIME_TYPES = { + "application/json", + "application/x-ndjson", + "application/xml", + "application/yaml", + "text/csv", + "text/xml", + "text/yaml", +} +_CODE_MIME_TYPES = { + "application/javascript", + "text/css", + "text/javascript", +} +_DATA_SUFFIXES = {".csv", ".json", ".jsonl", ".yaml", ".yml", ".xml"} +_CODE_SUFFIXES = { + ".py", + ".js", + ".jsx", + ".ts", + ".tsx", + ".go", + ".rs", + ".java", + ".c", + ".cpp", + ".h", + ".hpp", + ".cs", + ".php", + ".rb", + ".swift", + ".kt", + ".kts", + ".scala", + ".sh", + ".bash", + ".zsh", + ".ps1", + ".sql", + ".toml", + ".ini", + ".cfg", + ".css", + ".scss", +} + + +async def _wait_for_document_request_disconnect( + fastapi_request: Request, cancel_event: threading.Event +) -> bool: + while not cancel_event.is_set(): + if await fastapi_request.is_disconnected(): + cancel_event.set() + return True + await asyncio.sleep(0.2) + return False + + +def _extract_ext(filename: str) -> str: + return os.path.splitext(filename or "")[1].lower() + + +def _is_supported_upload(filename: str, content_type: str) -> bool: + if (content_type or "").split(";")[0].strip().lower() in _DOC_MIME_OK: + return True + return _extract_ext(filename) in _DOC_SUFFIX_OK + + +def _document_upload_format(filename: str, content_type: str) -> Optional[str]: + mime = (content_type or "").split(";")[0].strip().lower() + ext = _extract_ext(filename) + if mime == "application/pdf" or ext == ".pdf": + return "pdf" + if mime == _DOCX_MIME or ext == ".docx": + return "docx" + if mime in _HTML_MIME_TYPES or ext in {".html", ".htm"}: + return "html" + if mime in _DATA_MIME_TYPES or ext in _DATA_SUFFIXES: + return "data" + if mime in _CODE_MIME_TYPES or ext in _CODE_SUFFIXES: + return "code" + if mime.startswith("text/") or ext in {".md", ".txt", ".log"}: + return "text" + return None + + +def _raise_if_document_parser_unavailable(filename: str, content_type: str) -> None: + format_key = _document_upload_format(filename, content_type) + if format_key is None: + return + support = _document_parser_support() + if support.get(format_key, True): + return + reason = _document_parser_unavailable_reasons().get( + format_key, + f"{format_key.upper()} extraction is not available on this server.", + ) + raise HTTPException(status_code = 501, detail = reason) + + +def _document_caption_authorization_header( + capability: Any, llama_backend: Any, studio_authorization_header: Optional[str] +) -> Optional[str]: + if getattr(capability, "source", None) != "gguf": + return studio_authorization_header + api_key = getattr(llama_backend, "api_key", None) or getattr(llama_backend, "_api_key", None) + return f"Bearer {api_key}" if api_key else None + + +_FORM_TRUE = {"1", "true", "yes", "on"} +_FORM_FALSE = {"0", "false", "no", "off"} + + +def _parse_bool_form( + value: Any, + *, + default: bool, + field: str = "value", +) -> bool: + if value is None: + return default + norm = str(value).strip().lower() + if not norm: + return default + if norm in _FORM_TRUE: + return True + if norm in _FORM_FALSE: + return False + raise HTTPException( + status_code = 400, + detail = f"Invalid boolean value for {field}: {value!r}", + ) + + +def _parse_int_form( + value: Any, + *, + default: int, + lo: int, + hi: Optional[int] = None, +) -> int: + try: + parsed = int(value) if value is not None else default + except (TypeError, ValueError): + parsed = default + parsed = max(lo, parsed) + return min(parsed, hi) if hi is not None else parsed + + +def _reject_oversized_content_length(request: Request) -> None: + raw = request.headers.get("content-length") + if raw is None: + return + try: + total = int(raw) + except ValueError: + raise HTTPException( + status_code = 400, + detail = "Invalid Content-Length header", + ) + max_request_bytes = _EXTRACT_MAX_BYTES + _EXTRACT_MULTIPART_OVERHEAD_BYTES + if total > max_request_bytes: + raise HTTPException( + status_code = 413, + detail = (f"Request exceeds the {_EXTRACT_MAX_BYTES // (1024*1024)} MB file limit"), + ) + + +async def _iter_request_body_limited(request: Request, *, max_bytes: int): + total = 0 + async for chunk in request.stream(): + if not chunk: + continue + total += len(chunk) + if total > max_bytes: + raise HTTPException( + status_code = 413, + detail = (f"Request exceeds the {_EXTRACT_MAX_BYTES // (1024*1024)} MB file limit"), + ) + yield chunk + + +async def _read_multipart_form_limited(request: Request, *, max_bytes: int): + from starlette.formparsers import MultiPartException, MultiPartParser + try: + parser = MultiPartParser( + request.headers, + _iter_request_body_limited(request, max_bytes = max_bytes), + ) + return await parser.parse() + except HTTPException: + raise + except MultiPartException as exc: + raise HTTPException(status_code = 400, detail = exc.message) from exc + + +async def _read_upload_limited(upload: Any, *, max_bytes: int) -> bytes: + buf = bytearray() + while True: + chunk = await upload.read(_EXTRACT_READ_CHUNK_BYTES) + if not chunk: + break + buf.extend(chunk) + if len(buf) > max_bytes: + raise HTTPException( + status_code = 413, + detail = f"File exceeds the {max_bytes // (1024*1024)} MB limit", + ) + return bytes(buf) + + +def _is_pdf_upload(filename: str, content_type: str) -> bool: + mime = (content_type or "").split(";")[0].strip().lower() + return mime == "application/pdf" or _extract_ext(filename) == ".pdf" + + +def _preflight_pdf_page_count(file_bytes: bytes, filename: str, content_type: str) -> Optional[int]: + if not _is_pdf_upload(filename, content_type): + return None + + pypdf_error: Optional[BaseException] = None + try: + from pypdf import PdfReader + + reader = PdfReader(io.BytesIO(file_bytes), strict = False) + # Many PDFs report is_encrypted=True with only a null user password + # (Acrobat-distilled docs, the Orimi test PDF). Try the empty password + # first; PyMuPDF's needs_pass is the real signal in the fallback. + if getattr(reader, "is_encrypted", False): + try: + if reader.decrypt("") == 0: + raise HTTPException( + status_code = 422, + detail = "Encrypted PDFs are not supported for inline extraction", + ) + except HTTPException: + raise + except Exception: + # decrypt failed (corrupt /Encrypt, unknown algorithm); fall + # through to PyMuPDF rather than declaring it encrypted. + raise RuntimeError("pypdf decrypt probe failed") + return len(reader.pages) + except HTTPException: + raise + except Exception as exc: + pypdf_error = exc + logger.warning( + "pypdf page-count preflight failed for %s; trying PyMuPDF fallback", + filename, + ) + + try: + import pymupdf as _pymupdf # type: ignore + doc = _pymupdf.open(stream = file_bytes, filetype = "pdf") + try: + # needs_pass is True only when a password is actually required; + # is_encrypted also flags the null-password case that opens fine. + # Refuse only on needs_pass. + if getattr(doc, "needs_pass", False): + raise HTTPException( + status_code = 422, + detail = "Encrypted PDFs are not supported for inline extraction", + ) + return len(doc) + finally: + doc.close() + except HTTPException: + raise + except Exception as exc: + if pypdf_error is not None: + logger.warning( + "PyMuPDF page-count fallback also failed for %s: %s", + filename, + exc, + ) + else: + logger.exception("PDF page-count preflight failed for %s", filename) + raise HTTPException( + status_code = 400, + detail = "Unable to read PDF page count before extraction", + ) from exc + + +def _truncate_markdown_to_token_budget( + markdown: str, *, token_budget: int, original_tokens_est: int +) -> tuple[str, int, Optional[str]]: + char_budget = max(_EXTRACT_TOKEN_BUDGET_MIN, token_budget) * 4 + if len(markdown) <= char_budget: + return markdown, original_tokens_est, None + + clipped = markdown[:char_budget] + clipped = _re.sub(r"\s+\S*$", "", clipped).rstrip() or markdown[:char_budget].rstrip() + clipped += f"\n\n[... truncated; original was ~{original_tokens_est} tokens ...]" + warning = ( + f"Extracted markdown was truncated to {token_budget} tokens " + f"(original was ~{original_tokens_est} tokens)." + ) + return clipped, max(0, len(clipped) // 4), warning + + +@studio_router.get("/chat/document-support", response_model = DocumentSupportResponse) +async def document_support_endpoint( + fastapi_request: Request, current_subject: str = Depends(get_current_subject) +): + """Whether document extraction + per-figure captions are available. + + Polled on settings mount and model change; when ``vlm.is_vlm`` is false + the UI disables the describe toggle and shows ``vlm.reason`` as tooltip. + """ + if _extract_document is None or _detect_loaded_vlm is None: + return DocumentSupportResponse( + extraction_available = False, + max_visual_payloads = 0, + max_extract_concurrency = 1, + format_support = {}, + unavailable_formats = {}, + vlm = { + "is_vlm": False, + "endpoint_url": None, + "model_name": None, + "source": "none", + "reason": "document extraction backend is not installed", + }, + ) + + self_base_url = _extract_self_base_url(fastapi_request) if _extract_self_base_url else None + try: + cap = _detect_loaded_vlm( + self_base_url, + llama_backend = get_llama_cpp_backend(), + ) + except Exception as exc: + logger.exception("Document support VLM probe failed") + if _VlmCapability is not None: + cap = _VlmCapability.none(f"document support probe failed: {type(exc).__name__}") + else: # pragma: no cover - only when core.chat import fallback is active + cap = None + return DocumentSupportResponse( + extraction_available = _DOCUMENT_EXTRACTION_AVAILABLE, + max_visual_payloads = _MAX_DOCUMENT_VISUAL_PAYLOADS, + max_extract_concurrency = _DOCUMENT_EXTRACT_CONCURRENCY, + format_support = _document_parser_support(), + unavailable_formats = _document_parser_unavailable_reasons(), + vlm = cap.to_dict() + if cap is not None + else { + "is_vlm": False, + "endpoint_url": None, + "model_name": None, + "source": "none", + "reason": "document support probe failed", + }, + ) + + +@studio_router.post("/chat/extract-document") +async def extract_document_endpoint( + fastapi_request: Request, current_subject: str = Depends(get_current_subject) +): + """Upload a PDF / DOCX / HTML / MD / text file; stream NDJSON progress + events plus a final layout-aware Markdown payload. + + Pre-stream validation errors return standard HTTP 4xx/5xx; after that the + final line is ``{"stage":"result"|"error", ...}``. Documents over 200 + pages are rejected with 413 until the background-job path lands. + """ + if _extract_document is None: + raise HTTPException( + status_code = 501, + detail = ( + "document extraction backend is not installed. Re-run Studio " + "setup to install the parser dependencies." + ), + ) + + _reject_oversized_content_length(fastapi_request) + + try: + try: + form = await _read_multipart_form_limited( + fastapi_request, + max_bytes = _EXTRACT_MAX_BYTES + _EXTRACT_MULTIPART_OVERHEAD_BYTES, + ) + except HTTPException: + raise + except Exception as exc: + logger.exception("Invalid multipart document extraction payload") + raise HTTPException(status_code = 400, detail = "Invalid multipart payload") + + upload = form.get("file") + if upload is None or not hasattr(upload, "read"): + raise HTTPException(status_code = 400, detail = "Missing 'file' field") + + filename = getattr(upload, "filename", None) or "upload" + content_type = getattr(upload, "content_type", "") or "" + if not _is_supported_upload(filename, content_type): + raise HTTPException( + status_code = 415, + detail = f"Unsupported file type: {filename} ({content_type})", + ) + _raise_if_document_parser_unavailable(filename, content_type) + + file_bytes = await _read_upload_limited(upload, max_bytes = _EXTRACT_MAX_BYTES) + if not file_bytes: + raise HTTPException(status_code = 400, detail = "Uploaded file is empty") + + preflight_page_count = _preflight_pdf_page_count(file_bytes, filename, content_type) + if preflight_page_count is not None and preflight_page_count > _EXTRACT_MAX_PAGES_INLINE: + raise HTTPException( + status_code = 413, + detail = _page_limit_detail(preflight_page_count), + ) + + describe_images = _parse_bool_form( + form.get("describe_images"), default = False, field = "describe_images" + ) + use_vlm_ocr = _parse_bool_form(form.get("use_vlm_ocr"), default = False, field = "use_vlm_ocr") + max_figures = _parse_int_form( + form.get("max_figures"), + default = 40, + lo = 0, + ) + max_visual_payloads = _parse_int_form( + form.get("max_visual_payloads"), + default = _DEFAULT_DOCUMENT_VISUAL_PAYLOADS, + lo = 0, + hi = _MAX_DOCUMENT_VISUAL_PAYLOADS, + ) + token_budget = _parse_int_form( + form.get("token_budget"), + default = _EXTRACT_TOKEN_BUDGET_DEFAULT, + lo = 0, + ) + + self_base_url = _extract_self_base_url(fastapi_request) if _extract_self_base_url else None + llama_backend = get_llama_cpp_backend() + capability = ( + _detect_loaded_vlm( + self_base_url, + llama_backend = llama_backend, + ) + if _detect_loaded_vlm + else None + ) + caption_authorization_header = _document_caption_authorization_header( + capability, + llama_backend, + fastapi_request.headers.get("authorization"), + ) + + if await fastapi_request.is_disconnected(): + raise HTTPException(status_code = 499, detail = "Client closed request") + + accept_header = (fastapi_request.headers.get("accept", "") or "").lower() + wants_stream = "application/x-ndjson" in accept_header + + def _build_response_payload(result: Any) -> ExtractDocumentResponse: + markdown_, tokens_est_, truncate_warning_ = _truncate_markdown_to_token_budget( + result.markdown, + token_budget = token_budget, + original_tokens_est = result.tokens_est, + ) + warnings_ = list(result.warnings) + if truncate_warning_: + warnings_.append(truncate_warning_) + return ExtractDocumentResponse( + filename = filename, + markdown = markdown_, + page_count = result.page_count, + tokens_est = tokens_est_, + truncated = truncate_warning_ is not None, + figures = [ExtractedFigureModel(**_asdict(f)) for f in result.figures], + describe_skipped_reason = result.describe_skipped_reason, + vlm_source = result.vlm_source, + vlm_model = result.vlm_model, + image_input_available = getattr(result, "image_input_available", False), + warnings = warnings_, + ) + + def _spawn_extraction(cancel_event, progress_cb = None) -> asyncio.Task: + extra = {"progress_cb": progress_cb} if progress_cb is not None else {} + return asyncio.create_task( + _extract_document( + file_bytes, + filename, + content_type = content_type, + describe_images = describe_images, + use_vlm_ocr = use_vlm_ocr, + max_figures = max_figures, + max_visual_payloads = max_visual_payloads, + capability = capability, + self_base_url = self_base_url, + authorization_header = caption_authorization_header, + cancel_event = cancel_event, + **extra, + ) + ) + + if not wants_stream: + # ---- Legacy JSON path (no progress events) ----------------- + cancel_event = threading.Event() + extraction_task = _spawn_extraction(cancel_event) + disconnect_task = asyncio.create_task( + _wait_for_document_request_disconnect(fastapi_request, cancel_event) + ) + try: + done, _pending = await asyncio.wait( + {extraction_task, disconnect_task}, + return_when = asyncio.FIRST_COMPLETED, + ) + if ( + extraction_task not in done + and disconnect_task in done + and disconnect_task.result() + ): + await _drain_cancelled_extraction(cancel_event, extraction_task) + result = await extraction_task + except _DOC_EXTRACTION_HTTP_ERRORS as exc: + status_code, detail = _doc_exc_to_status_detail(exc) + raise HTTPException(status_code = status_code, detail = detail) + except Exception: + logger.exception("Document extraction failed for %s", filename) + raise HTTPException(status_code = 500, detail = "Extraction failed") + finally: + cancel_event.set() + disconnect_task.cancel() + with suppress(asyncio.CancelledError): + await disconnect_task + + if result.page_count > _EXTRACT_MAX_PAGES_INLINE: + raise HTTPException( + status_code = 413, + detail = _page_limit_detail(result.page_count), + ) + return _build_response_payload(result) + + # ---- Streaming NDJSON path (Accept: application/x-ndjson) ------ + progress_queue: asyncio.Queue = asyncio.Queue() + + async def _progress_cb(event: dict) -> None: + await progress_queue.put(dict(event)) + + async def _ndjson_stream(): + cancel_event = threading.Event() + extraction_task = _spawn_extraction(cancel_event, _progress_cb) + # Drain the task's exception so a busy/cancel race doesn't log + # "Future exception was never retrieved" on early exit. + extraction_task.add_done_callback(_drain_doc_future_exception) + disconnect_task = asyncio.create_task( + _wait_for_document_request_disconnect(fastapi_request, cancel_event) + ) + try: + extract_wait = asyncio.ensure_future(asyncio.shield(extraction_task)) + extract_wait.add_done_callback(_drain_doc_future_exception) + while True: + queue_get = asyncio.ensure_future(progress_queue.get()) + queue_get.add_done_callback(_drain_doc_future_exception) + done, _pending = await asyncio.wait( + {queue_get, extract_wait, disconnect_task}, + return_when = asyncio.FIRST_COMPLETED, + ) + if queue_get in done: + event = queue_get.result() + yield json.dumps(event) + "\n" + else: + queue_get.cancel() + with suppress(asyncio.CancelledError): + await queue_get + + if disconnect_task in done and disconnect_task.result(): + await _drain_cancelled_extraction(cancel_event, extraction_task) + + # The shield wrapper can finish (cancelled) before the real + # task; .result() in that window raises InvalidStateError, + # so wait on the task itself. + if extraction_task.done(): + # Drain any remaining progress events before result. + while not progress_queue.empty(): + try: + event = progress_queue.get_nowait() + except asyncio.QueueEmpty: + break + yield json.dumps(event) + "\n" + result = extraction_task.result() + break + if extract_wait in done: + # Wrapper done, task still running: re-arm a fresh + # shielded future and loop. + extract_wait = asyncio.ensure_future(asyncio.shield(extraction_task)) + extract_wait.add_done_callback(_drain_doc_future_exception) + + if result.page_count > _EXTRACT_MAX_PAGES_INLINE: + yield _ndjson_error(413, _page_limit_detail(result.page_count)) + return + + response = _build_response_payload(result) + yield ( + json.dumps( + { + "stage": "result", + "data": response.model_dump(mode = "json"), + } + ) + + "\n" + ) + except _DOC_EXTRACTION_HTTP_ERRORS as exc: + status_code, detail = _doc_exc_to_status_detail(exc) + yield _ndjson_error(status_code, detail) + except Exception: + logger.exception("Document extraction failed for %s", filename) + yield _ndjson_error(500, "Extraction failed") + finally: + cancel_event.set() + disconnect_task.cancel() + with suppress(asyncio.CancelledError): + await disconnect_task + + return StreamingResponse( + _ndjson_stream(), + media_type = "application/x-ndjson", + ) + finally: + # _EXTRACT_SEMAPHORE is owned by _run_extract_process_sync; a busy + # semaphore becomes DocumentExtractionBusy -> in-stream error above. + pass diff --git a/studio/backend/tests/test_chat_document_extraction.py b/studio/backend/tests/test_chat_document_extraction.py new file mode 100644 index 0000000000..2098fe9313 --- /dev/null +++ b/studio/backend/tests/test_chat_document_extraction.py @@ -0,0 +1,907 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for the chat document extractor + VLM capability probe. + +Probe tests only shape-check core.chat.vlm_capability; backend-backed tests +skip when the optional deps (pymupdf / pymupdf4llm / mammoth) are missing. +""" + +from __future__ import annotations + +import importlib.util +import sys +from types import ModuleType +from typing import Any, Dict, Optional + +import pytest + +from core.chat.vlm_capability import ( + VlmCapability, + detect_loaded_vlm, + extract_self_base_url, +) + + +# ---------------------------------------------------------------------- # +# Shared fakes/factories # +# ---------------------------------------------------------------------- # + + +def install_fake_extract( + monkeypatch, + *, + returns = None, + extract = None, +): + """Mark extraction available and stub _run_extract_sync with a fixed `returns` + tuple (markdown, figures, pages, trunc, seen) or a custom `extract` callable.""" + from core.chat import document_extractor as de + + if extract is None: + + def extract( + _fb, + _fn, + _opts, + _ct = "", + ): + return returns + + monkeypatch.setattr(de, "DOCUMENT_EXTRACTION_AVAILABLE", True) + monkeypatch.setattr(de, "_run_extract_sync", extract) + + +def make_figures( + n, + *, + encoded_until = None, + size_with_payload = False, +): + """Build n ExtractedFigure rows; `encoded_until` (None = all) sets how many + carry image payloads, `size_with_payload` ties width/height to the payload.""" + from core.chat.document_extractor import ExtractedFigure + + figs = [] + for i in range(n): + has_payload = encoded_until is None or i < encoded_until + figs.append( + ExtractedFigure( + id = f"fig-{i}", + page = i + 1, + caption = None, + kind = "figure", + image_mime = "image/jpeg" if has_payload else None, + image_base64 = "b64" if has_payload else None, + image_width = (10 if has_payload else None) if size_with_payload else 10, + image_height = (10 if has_payload else None) if size_with_payload else 10, + ) + ) + return figs + + +def vlm_cap(source = "transformers", *, endpoint_url = "http://127.0.0.1:8000"): + """A loaded vision-capable VlmCapability for ``capability=`` arguments.""" + return VlmCapability( + is_vlm = True, + endpoint_url = endpoint_url, + model_name = "vlm", + source = source, + reason = None, + ) + + +# ---------------------------------------------------------------------- # +# VlmCapability dataclass # +# ---------------------------------------------------------------------- # + + +def test_vlm_capability_none_factory_is_safe_default() -> None: + cap = VlmCapability.none() + assert cap.is_vlm is False + assert cap.endpoint_url is None + assert cap.model_name is None + assert cap.source == "none" + assert cap.reason # non-empty + + +def test_vlm_capability_to_dict_round_trips_fields() -> None: + cap = VlmCapability( + is_vlm = True, + endpoint_url = "http://127.0.0.1:8080", + model_name = "qwen2-vl", + source = "gguf", + reason = None, + ) + assert cap.to_dict() == { + "is_vlm": True, + "endpoint_url": "http://127.0.0.1:8080", + "model_name": "qwen2-vl", + "source": "gguf", + "reason": None, + } + + +# ---------------------------------------------------------------------- # +# detect_loaded_vlm() across backend shapes # +# ---------------------------------------------------------------------- # + + +class _FakeLlama: + def __init__( + self, + *, + loaded: bool, + vision: bool = False, + base_url: str = "http://127.0.0.1:8080", + model_id: str = "fake-gguf", + ) -> None: + self.is_loaded = loaded + self.is_vision = vision + self.base_url = base_url + self.model_identifier = model_id + + +class _FakeInferenceBackend: + def __init__( + self, + *, + active: Optional[str], + info: Optional[Dict[str, Any]] = None, + ) -> None: + self.active_model_name = active + self.models: Dict[str, Dict[str, Any]] = {active: info or {}} if active else {} + + +def _patch_probes( + monkeypatch: pytest.MonkeyPatch, + *, + llama: Optional[_FakeLlama], + inference: Optional[_FakeInferenceBackend], +) -> None: + from core.chat import vlm_capability as vc + if llama is None: + monkeypatch.setattr(vc, "_probe_gguf", lambda _llama = None: None) + else: + + def probe_gguf(llama_backend = None): + backend = llama_backend or llama + if not backend.is_loaded: + return None + is_vision = bool(backend.is_vision) + return VlmCapability( + is_vlm = is_vision, + endpoint_url = backend.base_url, + model_name = backend.model_identifier, + source = "gguf", + reason = None if is_vision else "loaded GGUF is not vision-capable", + ) + + monkeypatch.setattr(vc, "_probe_gguf", probe_gguf) + + if inference is None: + monkeypatch.setattr(vc, "_probe_transformers", lambda _u: None) + else: + + def probe_tf(self_base_url): + name = inference.active_model_name + if not name: + return None + info = inference.models.get(name) or {} + is_vision = bool(info.get("is_vision", False)) + source = "unsloth" if info.get("is_lora") else "transformers" + if not self_base_url: + return VlmCapability( + is_vlm = False, + endpoint_url = None, + model_name = name, + source = source, + reason = "cannot self-loopback: request base URL unavailable", + ) + return VlmCapability( + is_vlm = is_vision, + endpoint_url = self_base_url.rstrip("/"), + model_name = name, + source = source, + reason = None if is_vision else "loaded model is not vision-capable", + ) + + monkeypatch.setattr(vc, "_probe_transformers", probe_tf) + + +def test_detect_returns_none_when_no_model_loaded(monkeypatch: pytest.MonkeyPatch) -> None: + _patch_probes(monkeypatch, llama = None, inference = None) + cap = detect_loaded_vlm() + assert cap.source == "none" + assert cap.is_vlm is False + + +def test_detect_gguf_vision_returns_llama_endpoint(monkeypatch: pytest.MonkeyPatch) -> None: + llama = _FakeLlama(loaded = True, vision = True, base_url = "http://127.0.0.1:9999") + _patch_probes(monkeypatch, llama = llama, inference = None) + cap = detect_loaded_vlm("http://studio.local") + assert cap.source == "gguf" + assert cap.is_vlm is True + assert cap.endpoint_url == "http://127.0.0.1:9999" # GGUF ignores self_base_url + assert cap.reason is None + + +def test_detect_gguf_vision_accepts_injected_backend(monkeypatch: pytest.MonkeyPatch) -> None: + from core.chat import vlm_capability as vc + + llama = _FakeLlama(loaded = True, vision = True, base_url = "http://127.0.0.1:9999") + monkeypatch.setattr(vc, "_probe_transformers", lambda _u: None) + + cap = detect_loaded_vlm( + "http://127.0.0.1:8000", + llama_backend = llama, + ) + + assert cap.source == "gguf" + assert cap.is_vlm is True + assert cap.endpoint_url == "http://127.0.0.1:9999" + + +def test_detect_gguf_vision_uses_core_llama_accessor(monkeypatch: pytest.MonkeyPatch) -> None: + """The implicit GGUF fallback must use the core-owned singleton path.""" + from core.chat import vlm_capability as vc + from core.inference import llama_cpp + + llama = _FakeLlama(loaded = True, vision = True, base_url = "http://127.0.0.1:9999") + assert hasattr(llama_cpp, "get_llama_cpp_backend") + monkeypatch.setattr(llama_cpp, "_llama_cpp_backend", llama) + monkeypatch.setattr(vc, "_probe_transformers", lambda _u: None) + + cap = detect_loaded_vlm("http://127.0.0.1:8000") + + assert cap.source == "gguf" + assert cap.is_vlm is True + assert cap.endpoint_url == "http://127.0.0.1:9999" + + +def test_detect_gguf_non_vision_surfaces_reason(monkeypatch: pytest.MonkeyPatch) -> None: + llama = _FakeLlama(loaded = True, vision = False) + _patch_probes(monkeypatch, llama = llama, inference = None) + cap = detect_loaded_vlm() + assert cap.source == "gguf" + assert cap.is_vlm is False + assert cap.reason and "vision" in cap.reason.lower() + + +def test_detect_transformers_vision_uses_self_loopback(monkeypatch: pytest.MonkeyPatch) -> None: + ib = _FakeInferenceBackend( + active = "Qwen2-VL-7B", + info = {"is_vision": True, "is_lora": False}, + ) + _patch_probes(monkeypatch, llama = None, inference = ib) + cap = detect_loaded_vlm("http://127.0.0.1:8000/") + assert cap.source == "transformers" + assert cap.is_vlm is True + assert cap.endpoint_url == "http://127.0.0.1:8000" + assert cap.model_name == "Qwen2-VL-7B" + + +def test_detect_unsloth_lora_vision_reports_unsloth_source(monkeypatch: pytest.MonkeyPatch) -> None: + ib = _FakeInferenceBackend( + active = "my-qwen-vl-lora", + info = {"is_vision": True, "is_lora": True}, + ) + _patch_probes(monkeypatch, llama = None, inference = ib) + cap = detect_loaded_vlm("http://studio.local:8000") + assert cap.source == "unsloth" + assert cap.is_vlm is True + + +def test_detect_falls_through_when_gguf_is_loaded_but_endpoint_data_missing( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A half-initialised llama-server (is_loaded=True but base_url/model + missing) must not suppress the transformers fallback path — otherwise + a misleading non-vision GGUF result hides an active transformers VLM. + """ + from core.chat import vlm_capability as vc + + fake_llama_cpp = ModuleType("core.inference.llama_cpp") + fake_llama_cpp.get_llama_cpp_backend = lambda: _FakeLlama( + loaded = True, + base_url = "", + model_id = "", + ) + fake_inference = ModuleType("core.inference") + fake_inference.__path__ = [] # type: ignore[attr-defined] + fake_inference.llama_cpp = fake_llama_cpp # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, "core.inference", fake_inference) + monkeypatch.setitem(sys.modules, "core.inference.llama_cpp", fake_llama_cpp) + + ib = _FakeInferenceBackend( + active = "Qwen2-VL-7B", + info = {"is_vision": True, "is_lora": False}, + ) + monkeypatch.setattr( + vc, + "_probe_transformers", + lambda self_base_url: VlmCapability( + is_vlm = True, + endpoint_url = self_base_url.rstrip("/") if self_base_url else None, + model_name = ib.active_model_name, + source = "transformers", + reason = None, + ), + ) + + cap = detect_loaded_vlm("http://127.0.0.1:8000") + assert cap.source == "transformers" + assert cap.is_vlm is True + + +def test_detect_transformers_without_self_url_reports_missing_loopback( + monkeypatch: pytest.MonkeyPatch, +) -> None: + ib = _FakeInferenceBackend( + active = "Qwen2-VL-7B", + info = {"is_vision": True, "is_lora": False}, + ) + _patch_probes(monkeypatch, llama = None, inference = ib) + cap = detect_loaded_vlm(None) + assert cap.is_vlm is False + assert cap.reason and "loopback" in cap.reason.lower() + + +# ---------------------------------------------------------------------- # +# extract_self_base_url — request base-URL extraction # +# ---------------------------------------------------------------------- # + + +class _FakeState: + def __init__(self, server_port: Optional[int] = None) -> None: + if server_port is not None: + self.server_port = server_port + + +class _FakeApp: + def __init__(self, server_port: Optional[int] = None) -> None: + self.state = _FakeState(server_port) + + +class _FakeRequest: + def __init__( + self, + base_url: str, + *, + server_port: Optional[int] = None, + scope_server: Optional[tuple[str, int]] = None, + ) -> None: + self.base_url = base_url + self.app = _FakeApp(server_port) + self.scope = {"server": scope_server} if scope_server else {} + + +def test_extract_self_base_url_strips_trailing_slash() -> None: + assert extract_self_base_url(_FakeRequest("http://127.0.0.1:8000/")) == "http://127.0.0.1:8000" + + +def test_extract_self_base_url_prefers_trusted_server_port() -> None: + assert ( + extract_self_base_url( + _FakeRequest( + "http://attacker.invalid:9999/", + server_port = 7777, + scope_server = ("127.0.0.1", 6666), + ) + ) + == "http://127.0.0.1:7777" + ) + assert ( + extract_self_base_url( + _FakeRequest( + "http://attacker.invalid:9999/", + scope_server = ("127.0.0.1", 6666), + ) + ) + == "http://127.0.0.1:6666" + ) + + +def test_extract_self_base_url_ignores_host_header() -> None: + assert ( + extract_self_base_url(_FakeRequest("http://studio.local:8000/")) == "http://127.0.0.1:8000" + ) + assert ( + extract_self_base_url(_FakeRequest("https://example.com:9443/")) == "http://127.0.0.1:9443" + ) + + +def test_extract_self_base_url_none_when_empty() -> None: + assert extract_self_base_url(_FakeRequest("")) is None + + +def test_extract_self_base_url_none_on_missing_attribute() -> None: + assert extract_self_base_url(object()) is None + + +# ---------------------------------------------------------------------- # +# extract_document orchestration — backend-agnostic (monkey-patched) # +# ---------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_max_figures_zero_sets_describe_skipped_reason( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """max_figures=0 must skip description with a specific diagnostic even + when a VLM is available.""" + from core.chat import document_extractor as de + + install_fake_extract(monkeypatch, returns = ("# Smoke\n", [], 1, 0, 0)) + + result = await de.extract_document( + b"# Smoke\n", + "sample.md", + describe_images = True, + max_figures = 0, + capability = vlm_cap(), + ) + + assert result.describe_skipped_reason == ( + "figure description disabled because max_figures is 0" + ) + assert result.markdown == "# Smoke\n" + assert result.figures == [] + + +@pytest.mark.asyncio +async def test_extract_document_clamps_visual_payloads_to_cap( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Core clamps max_visual_payloads to the advertised cap for any caller.""" + from core.chat import document_extractor as de + + captured: dict[str, object] = {} + + def fake_extract( + _fb, + _fn, + opts, + _ct = "", + ): + captured.update(opts) + return "# Doc\n", [], 1, 0, 0 + + install_fake_extract(monkeypatch, extract = fake_extract) + + await de.extract_document( + b"# Doc\n", + "doc.md", + max_figures = 1000, + max_visual_payloads = 222, + ) + + assert captured["max_visual_payloads"] == de.MAX_DOCUMENT_VISUAL_PAYLOADS + + +@pytest.mark.asyncio +async def test_run_extract_sync_seam_receives_content_type(monkeypatch: pytest.MonkeyPatch) -> None: + """The test seam path (monkeypatched _run_extract_sync) must be invoked + with the content_type so dispatch-by-content-type can be exercised in + tests, not only by filename suffix.""" + from core.chat import document_extractor as de + + received: dict[str, str] = {} + + def fake_extract( + _fb, + _fn, + _opts, + ct = "", + ): + received["content_type"] = ct + return "ok", [], 0, 0, 0 + + install_fake_extract(monkeypatch, extract = fake_extract) + + await de.extract_document( + b"hello", + "no-suffix-file", + content_type = "text/plain", + describe_images = False, + ) + assert received["content_type"] == "text/plain" + + +@pytest.mark.asyncio +async def test_describe_image_via_vlm_sends_auth_header_and_max_tokens( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from core.chat import document_extractor as de + + captured: dict[str, Any] = {} + + class FakeResponse: + status_code = 200 + + def json(self): + return {"choices": [{"message": {"content": "A chart."}}]} + + class FakeAsyncClient: + def __init__(self, *, timeout: float) -> None: + captured["timeout"] = timeout + + async def __aenter__(self): + return self + + async def __aexit__(self, *_args): + return None + + async def post(self, url, *, headers, json): + captured["url"] = url + captured["headers"] = headers + captured["json"] = json + return FakeResponse() + + fake_httpx = ModuleType("httpx") + fake_httpx.AsyncClient = FakeAsyncClient + monkeypatch.setitem(sys.modules, "httpx", fake_httpx) + + caption, error = await de._describe_image_via_vlm( + image_base64 = "abc", + image_mime = "image/jpeg", + endpoint_url = "http://127.0.0.1:8000", + model_name = "vlm", + authorization_header = "Bearer token", + timeout_seconds = 7, + ) + + assert caption == "A chart." + assert error is None + assert captured["url"] == "http://127.0.0.1:8000/v1/chat/completions" + assert captured["headers"]["Authorization"] == "Bearer token" + assert captured["json"]["max_tokens"] == 512 + assert "max_completion_tokens" not in captured["json"] + + +# ---------------------------------------------------------------------- # +# Backend dispatch — real _run_extract_sync (requires pymupdf/mammoth) # +# ---------------------------------------------------------------------- # + + +_BACKEND_INSTALLED = ( + importlib.util.find_spec("pymupdf") is not None + and importlib.util.find_spec("pymupdf4llm") is not None + and importlib.util.find_spec("mammoth") is not None +) + + +def test_run_extract_sync_rejects_pptx_with_value_error() -> None: + """PPTX was dropped in the PyMuPDF4LLM migration. _run_extract_sync + must raise ValueError so the route can map it to HTTP 415.""" + if not _BACKEND_INSTALLED: + pytest.skip("extraction backend not installed") + from core.chat import document_extractor as de + + with pytest.raises(ValueError): + de._run_extract_sync( + b"PK\x03\x04", + "deck.pptx", + {"max_figures": 0, "extract_images": False, "use_vlm_ocr": False}, + ) + + +def test_run_extract_sync_text_path_decodes_utf8() -> None: + """TXT / MD paths must not require PDF/DOCX parser dependencies.""" + from core.chat import document_extractor as de + + md, figs, pages, trunc, seen = de._run_extract_sync( + "# Héllo\n".encode("utf-8"), + "notes.md", + {"max_figures": 0, "extract_images": False, "use_vlm_ocr": False}, + ) + assert md == "# Héllo\n" + assert figs == [] + assert pages == 0 and trunc == 0 and seen == 0 + + +def test_run_extract_sync_html_converts_to_markdown_without_parser_deps() -> None: + """HTML must be cleaned before prompt injection and not depend on PDF/DOCX deps.""" + from core.chat import document_extractor as de + + md, figs, pages, trunc, seen = de._run_extract_sync( + b"

Title

Hello world

", + "page.html", + {"max_figures": 0, "extract_images": False, "use_vlm_ocr": False}, + ) + assert "# Title" in md + assert "**world**" in md + assert "