From 1069b28c43d7f6375e946cd8632a8448f259de0d Mon Sep 17 00:00:00 2001 From: Michael Han <107991372+shimmyshimmer@users.noreply.github.com> Date: Mon, 29 Jun 2026 07:12:44 -0700 Subject: [PATCH] Studio: name the missing extractor when a Recipes file upload fails (#6642) * Studio: name the missing extractor when a Recipes upload fails A missing optional dependency (pymupdf4llm for PDF, mammoth for DOCX) was reported as a generic "Text extraction failed", which gives the user nothing to act on. Catch ImportError and surface the package name instead. * Studio: narrow missing extractor error handling * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: wasimysaid Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/routes/data_recipe/seed.py | 31 +++++ studio/backend/tests/test_data_recipe_seed.py | 120 +++++++++++++++++- 2 files changed, 148 insertions(+), 3 deletions(-) diff --git a/studio/backend/routes/data_recipe/seed.py b/studio/backend/routes/data_recipe/seed.py index 8fb034ea4e..57a291291e 100644 --- a/studio/backend/routes/data_recipe/seed.py +++ b/studio/backend/routes/data_recipe/seed.py @@ -481,6 +481,37 @@ async def upload_unstructured_file( error = "No extractable text found in file", ) extracted_path.write_text(extracted_text, encoding = "utf-8") + except ImportError as e: + raw_path.unlink(missing_ok = True) + extracted_path.unlink(missing_ok = True) + missing = getattr(e, "name", None) + expected_missing = {".pdf": "pymupdf4llm", ".docx": "mammoth"}.get(ext) + if isinstance(e, ModuleNotFoundError) and missing == expected_missing: + logger.error( + "data_recipe.seed.text_extraction_dependency_missing", + error = str(e), + missing = missing, + exc_info = True, + ) + return UnstructuredFileUploadResponse( + file_id = file_id, + filename = original_filename, + size_bytes = size_bytes, + status = "error", + error = f"Cannot read {ext} files: the '{missing}' package is not installed.", + ) + logger.error( + "data_recipe.seed.text_extraction_failed", + error = str(e), + exc_info = True, + ) + return UnstructuredFileUploadResponse( + file_id = file_id, + filename = original_filename, + size_bytes = size_bytes, + status = "error", + error = "Text extraction failed.", + ) except Exception as e: raw_path.unlink(missing_ok = True) extracted_path.unlink(missing_ok = True) diff --git a/studio/backend/tests/test_data_recipe_seed.py b/studio/backend/tests/test_data_recipe_seed.py index 601df8bbfe..09e22116ed 100644 --- a/studio/backend/tests/test_data_recipe_seed.py +++ b/studio/backend/tests/test_data_recipe_seed.py @@ -1,12 +1,126 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 +import asyncio +import importlib.util from pathlib import Path +import pytest -def test_seed_inspect_load_kwargs_disables_remote_code_execution(): - seed_route = ( + +def _seed_route_source() -> str: + return ( Path(__file__).resolve().parent.parent / "routes" / "data_recipe" / "seed.py" ).read_text() - assert '"trust_remote_code": False' in seed_route + +def test_seed_inspect_load_kwargs_disables_remote_code_execution(): + assert '"trust_remote_code": False' in _seed_route_source() + + +class _FakeUpload: + def __init__(self, filename: str, content: bytes): + self.filename = filename + self._content = content + + async def read(self) -> bytes: + return self._content + + +def _load_seed_route(monkeypatch: pytest.MonkeyPatch, tmp_path: Path): + pytest.importorskip("fastapi") + pytest.importorskip("multipart") + pytest.importorskip("structlog") + + backend_root = Path(__file__).resolve().parent.parent + monkeypatch.syspath_prepend(str(backend_root)) + route_path = backend_root / "routes" / "data_recipe" / "seed.py" + spec = importlib.util.spec_from_file_location("seed_under_test", route_path) + assert spec is not None and spec.loader is not None + seed_route = importlib.util.module_from_spec(spec) + spec.loader.exec_module(seed_route) + seed_route.UNSTRUCTURED_UPLOAD_ROOT = tmp_path / "unstructured-uploads" + return seed_route + + +def _run_upload( + seed_route, + filename: str, + content: bytes, + block_id: str = "block", +): + return asyncio.run( + seed_route.upload_unstructured_file(_FakeUpload(filename, content), block_id) + ) + + +def _block_files(seed_route, block_id: str = "block") -> list[str]: + block_dir = seed_route.UNSTRUCTURED_UPLOAD_ROOT / block_id + if not block_dir.exists(): + return [] + return sorted(path.name for path in block_dir.iterdir()) + + +def _raise(exc: BaseException): + def raise_exc(*args, **kwargs): + raise exc + + return raise_exc + + +@pytest.mark.parametrize( + ("filename", "package"), + [ + ("paper.pdf", "pymupdf4llm"), + ("notes.docx", "mammoth"), + ], +) +def test_unstructured_upload_names_missing_extractor_dependency( + monkeypatch, tmp_path, filename, package +): + seed_route = _load_seed_route(monkeypatch, tmp_path) + monkeypatch.setattr( + seed_route, + "_extract_text_from_file", + _raise(ModuleNotFoundError(f"No module named {package!r}", name = package)), + ) + + result = _run_upload(seed_route, filename, b"%PDF-1.7") + + assert result.status == "error" + assert ( + result.error + == f"Cannot read {Path(filename).suffix} files: the '{package}' package is not installed." + ) + assert _block_files(seed_route) == [] + + +def test_unstructured_upload_keeps_txt_path_working(monkeypatch, tmp_path): + seed_route = _load_seed_route(monkeypatch, tmp_path) + + result = _run_upload(seed_route, "notes.txt", b"hello") + + assert result.status == "ok" + assert result.error is None + assert any(name.endswith(".txt") for name in _block_files(seed_route)) + assert any(name.endswith(".extracted.txt") for name in _block_files(seed_route)) + + +@pytest.mark.parametrize( + "exc", + [ + ImportError("cannot import internal symbol"), + ModuleNotFoundError( + "No module named 'missing_transitive_pkg'", + name = "missing_transitive_pkg", + ), + ], +) +def test_unstructured_upload_import_errors_stay_generic(monkeypatch, tmp_path, exc): + seed_route = _load_seed_route(monkeypatch, tmp_path) + monkeypatch.setattr(seed_route, "_extract_text_from_file", _raise(exc)) + result = _run_upload(seed_route, "paper.pdf", b"%PDF-1.7") + + assert result.status == "error" + assert result.error == "Text extraction failed." + assert _block_files(seed_route) == []