diff --git a/apps/worker/app/services/document_parser/providers/mineru/pdf_service.py b/apps/worker/app/services/document_parser/providers/mineru/pdf_service.py index ffaae7af..ce160fc7 100644 --- a/apps/worker/app/services/document_parser/providers/mineru/pdf_service.py +++ b/apps/worker/app/services/document_parser/providers/mineru/pdf_service.py @@ -23,6 +23,7 @@ UnavailableException, ) from shared.services.storage.job_file_storage import JobFileStorage +from shared.utils.zip_download import download_and_extract_zip from app.services.common.file_loading import is_remote MINERU_UPLOAD_TIMEOUT = ( @@ -386,12 +387,264 @@ def _submit_url_task(presigned_url: str, filename: str) -> tuple[str, str]: return batch_id, lease.token_id +def _flatten_extracted_zip(output_dir: str) -> None: + """Flatten a local MinerU ZIP layout into the shape downstream code expects. + + Local MinerU extracts to ``{stem}/auto/{stem}.md`` plus + ``{stem}/auto/images/*``. Downstream code expects ``full.md`` and an + ``images/`` directory at the output dir root. This lifts the contents + of the single ``{stem}/auto/`` directory to the output root (preserving + the ``images/`` subdir), removes the ``{stem}/`` wrapper, drops files + outside the keep set, and renames the single markdown file to + ``full.md``. Hard-fails on zero or multiple ``.md`` files so we never + silently pick the wrong one. + """ + import shutil + from pathlib import Path + + destination = Path(output_dir).resolve() + keep_exts = (".md", ".jpg", ".jpeg", ".png", ".gif", ".json") + exclude_patterns = ("content_list", "middle.json", "model.json") + + auto_dirs = [p for p in destination.glob("*/auto") if p.is_dir()] + if not auto_dirs: + raise MinerUServiceException( + internal_message=( + "Local MinerU ZIP did not contain a {stem}/auto/ directory; " + "layout has changed or the response was not a parse result." + ), + ) + if len(auto_dirs) > 1: + relative_paths = ", ".join(str(p.relative_to(destination)) for p in auto_dirs) + raise MinerUServiceException( + internal_message=( + f"Local MinerU ZIP contained {len(auto_dirs)} */auto directories; " + f"expected exactly one: {relative_paths}" + ), + ) + + auto_dir = auto_dirs[0] + for source_path in auto_dir.rglob("*"): + if source_path.is_dir(): + continue + relative = source_path.relative_to(auto_dir) + if any(pattern in source_path.name for pattern in exclude_patterns): + continue + if source_path.suffix.lower() not in keep_exts: + continue + target = destination / relative + target.parent.mkdir(parents=True, exist_ok=True) + if target.exists(): + continue + shutil.move(str(source_path), str(target)) + + shutil.rmtree(auto_dir.parent, ignore_errors=True) + + markdown_files = sorted(destination.glob("*.md")) + if len(markdown_files) == 0: + raise MinerUServiceException( + internal_message=( + "Local MinerU ZIP contained no markdown file under {stem}/auto/." + ), + ) + if len(markdown_files) > 1: + relative_paths = ", ".join(str(p.relative_to(destination)) for p in markdown_files) + raise MinerUServiceException( + internal_message=( + f"Local MinerU ZIP contained {len(markdown_files)} markdown files; " + f"expected exactly one: {relative_paths}" + ), + ) + + markdown_files[0].rename(destination / "full.md") + + +def _get_local_mineru_session() -> requests.Session: + """Build a session for local MinerU's synchronous /file_parse endpoint. + + Local MinerU is single-concurrency by default + (``max_concurrent_requests=1``). A ``ReadTimeout`` from queue wait must + not cascade into urllib3 retries that push the request to the back of + the same queue, so ``read`` retries are disabled. 429s surface as + ``UnavailableException`` for upstream retry handling and do not go + through the cloud quota manager (local mode has no API key). + """ + from requests.adapters import HTTPAdapter + from urllib3.util.retry import Retry + + session = requests.Session() + retry_strategy = Retry( + total=settings.MINERU_UPLOAD_RETRY_TOTAL, + backoff_factor=settings.MINERU_UPLOAD_RETRY_BACKOFF_FACTOR, + status_forcelist=[502, 503, 504], + allowed_methods=["GET", "POST"], + raise_on_status=False, + read=0, + ) + adapter = HTTPAdapter( + max_retries=retry_strategy, + pool_connections=1, + pool_maxsize=settings.MINERU_POOL_MAXSIZE, + ) + session.mount("https://", adapter) + session.mount("http://", adapter) + return session + + +_local_mineru_session: Optional[requests.Session] = None + + +def _get_local_mineru_session_cached() -> requests.Session: + global _local_mineru_session + if _local_mineru_session is None: + _local_mineru_session = _get_local_mineru_session() + return _local_mineru_session + + +def parse_via_local( + pdf_url: str, + filename: str, + output_dir: str, +) -> None: + """Parse a PDF via a local MinerU instance's synchronous /file_parse. + + Self-hosted deployments that run MinerU on their own network cannot use + the cloud-only batch APIs (``/file-urls/batch``, ``/extract/task/batch``, + ``/extract-results/batch``). Local MinerU exposes a single synchronous + ``/file_parse`` endpoint that accepts the PDF as multipart form data and + returns a ZIP with a different layout (``{stem}/auto/{stem}.md``). + """ + base_url = settings.MINERU_URL.rstrip("/") + endpoint = f"{base_url}/file_parse" + local_logger = mineru_logger( + "local_file_parse", + operation="local_file_parse", + filename=filename, + endpoint=endpoint, + lang_list=settings.MINERU_LOCAL_LANG_LIST, + backend=settings.MINERU_LOCAL_BACKEND, + ) + + form_fields = { + "lang_list": settings.MINERU_LOCAL_LANG_LIST, + "backend": settings.MINERU_LOCAL_BACKEND, + "return_images": "true", + } + + if is_remote(pdf_url): + import tempfile + + local_logger.info("Downloading remote source file before local MinerU parse") + try: + download_response = requests.get( + pdf_url, + stream=True, + timeout=APIConstants.S3_FILE_DOWNLOAD_TIMEOUT, + ) + download_response.raise_for_status() + with tempfile.NamedTemporaryFile( + delete=False, suffix=os.path.splitext(filename)[1] + ) as temp_file: + for chunk in download_response.iter_content(chunk_size=8192): + temp_file.write(chunk) + temp_path = temp_file.name + except requests.RequestException as exc: + local_logger.bind(error_message=str(exc)).error( + "Failed to stage remote source file for local MinerU" + ) + raise StorageServiceException( + internal_message=f"Failed to download remote file: {exc}" + ) + local_file_path = temp_path + cleanup_temp = True + else: + local_logger.bind(local_path=pdf_url).info( + "Uploading local file to local MinerU" + ) + local_file_path = pdf_url + cleanup_temp = False + + try: + with open(local_file_path, "rb") as file_obj: + files = {"file": (filename, file_obj, "application/pdf")} + local_logger.info("Posting PDF to local MinerU /file_parse") + try: + response = _get_local_mineru_session_cached().post( + endpoint, + data=form_fields, + files=files, + timeout=settings.MINERU_LOCAL_TIMEOUT, + ) + except requests.RequestException as exc: + local_logger.bind(error_type=type(exc).__name__).error( + "Local MinerU /file_parse request failed" + ) + raise MinerUServiceException( + internal_message=f"Local MinerU /file_parse failed: {exc}", + original_exception=exc, + ) from exc + finally: + if cleanup_temp: + os.unlink(local_file_path) + + if response.status_code == 429: + retry_after = 60 + local_logger.bind( + status_code=response.status_code, + retry_after=retry_after, + ).warning("Local MinerU /file_parse rate-limited") + raise UnavailableException( + internal_message="Local MinerU rate limited during /file_parse", + retry_after=retry_after, + limit=1, + period="minute", + user_message="Document processing is busy right now. Please retry shortly.", + ) + if response.status_code != 200: + local_logger.bind(status_code=response.status_code).error( + "Local MinerU /file_parse failed" + ) + raise MinerUServiceException( + internal_message=( + f"Local MinerU /file_parse returned {response.status_code}: " + f"{response.text[:500]}" + ), + status_code=response.status_code, + ) + + import io + import zipfile + + local_logger.info("Local MinerU /file_parse completed, extracting ZIP") + try: + with zipfile.ZipFile(io.BytesIO(response.content)) as extracted_zip: + extracted_zip.extractall(output_dir) + except zipfile.BadZipFile as exc: + local_logger.bind(error_type=type(exc).__name__).error( + "Local MinerU response was not a valid ZIP" + ) + raise MinerUServiceException( + internal_message=f"Local MinerU returned a non-ZIP body: {exc}", + original_exception=exc, + ) from exc + + _flatten_extracted_zip(output_dir) + local_logger.info("Local MinerU parse completed and ZIP flattened") + + def parse_via_full( pdf_url: str, filename: str, output_dir: str, s3_key: Optional[str] = None, ) -> None: + if settings.MINERU_LOCAL_MODE: + mineru_logger("ingestion_mode", mode="local").info( + "Using local MinerU mode for ingestion" + ) + parse_via_local(pdf_url=pdf_url, filename=filename, output_dir=output_dir) + return + batch_id: str | None = None token_id: str | None = None resolved_s3_key = resolve_mineru_source_s3_key( diff --git a/apps/worker/tests/unit/test_pdf_service_flatten.py b/apps/worker/tests/unit/test_pdf_service_flatten.py new file mode 100644 index 00000000..50ad5979 --- /dev/null +++ b/apps/worker/tests/unit/test_pdf_service_flatten.py @@ -0,0 +1,154 @@ +"""Unit tests for the local MinerU ZIP flattener. + +The cloud MinerU flow returns a flat ZIP layout (``full.md`` + ``images/*`` +at the root). Local MinerU returns a nested layout +(``{stem}/auto/{stem}.md`` + ``{stem}/auto/images/*``), which downstream +code cannot consume. ``_flatten_extracted_zip`` rewrites the extracted +tree into the expected flat shape and hard-fails on ambiguous markdown +counts so we never silently pick the wrong file. +""" + +from __future__ import annotations + +import io +import os +import zipfile +from pathlib import Path + +import pytest + +os.environ.setdefault("DATABASE_URL", "postgresql+asyncpg://test:test@localhost/test") +os.environ.setdefault("TMP_PATH", "/tmp/knowhere-test") +os.environ.setdefault("S3_BUCKET_NAME", "test-uploads") +os.environ.setdefault("S3_ACCESS_KEY_ID", "test") +os.environ.setdefault("S3_SECRET_ACCESS_KEY", "test") +os.environ.setdefault("S3_TEMP_PATH", "/tmp") + +from app.services.document_parser.providers.mineru.pdf_service import ( # noqa: E402 + _flatten_extracted_zip, +) +from shared.core.exceptions.domain_exceptions import MinerUServiceException # noqa: E402 + + +def _write_zip(tree: dict[str, bytes], output_dir: Path) -> None: + """Extract a synthetic local-MinerU ZIP tree into ``output_dir``.""" + buffer = io.BytesIO() + with zipfile.ZipFile(buffer, "w", zipfile.ZIP_DEFLATED) as archive: + for relative_path, content in tree.items(): + archive.writestr(relative_path, content) + buffer.seek(0) + with zipfile.ZipFile(buffer, "r") as archive: + archive.extractall(output_dir) + + +def test_flatten_single_markdown_and_images(tmp_path: Path) -> None: + _write_zip( + { + "report/auto/report.md": b"# Heading\n\nBody text", + "report/auto/images/page-1-fig-1.jpg": b"\x00img1", + "report/auto/images/page-2-fig-2.png": b"\x00img2", + "report/auto/content_list.json": b"[]", + "report/auto/middle.json": b"{}", + }, + tmp_path, + ) + + _flatten_extracted_zip(str(tmp_path)) + + assert (tmp_path / "full.md").read_text() == "# Heading\n\nBody text" + image_names = sorted(p.name for p in (tmp_path / "images").iterdir()) + assert image_names == ["page-1-fig-1.jpg", "page-2-fig-2.png"] + assert not (tmp_path / "report").exists() + assert not (tmp_path / "content_list.json").exists() + assert not (tmp_path / "middle.json").exists() + +def test_flatten_raises_when_no_markdown(tmp_path: Path) -> None: + _write_zip( + { + "report/auto/images/page-1-fig-1.jpg": b"\x00img1", + "report/auto/content_list.json": b"[]", + }, + tmp_path, + ) + + with pytest.raises(MinerUServiceException, match="no markdown"): + _flatten_extracted_zip(str(tmp_path)) + + +def test_flatten_raises_when_no_auto_dir(tmp_path: Path) -> None: + _write_zip( + { + "report/report.md": b"# Heading", + "report/images/page-1-fig-1.jpg": b"\x00img1", + }, + tmp_path, + ) + + with pytest.raises(MinerUServiceException, match=r"\{stem\}/auto/"): + _flatten_extracted_zip(str(tmp_path)) + + +def test_flatten_raises_when_multiple_markdown(tmp_path: Path) -> None: + _write_zip( + { + "report/auto/report.md": b"# First", + "report/auto/second.md": b"# Second", + "report/auto/images/page-1-fig-1.jpg": b"\x00img1", + }, + tmp_path, + ) + + with pytest.raises(MinerUServiceException, match="2 markdown"): + _flatten_extracted_zip(str(tmp_path)) + + +def test_flatten_raises_when_multiple_auto_dirs(tmp_path: Path) -> None: + _write_zip( + { + "report1/auto/report1.md": b"# First", + "report2/auto/report2.md": b"# Second", + }, + tmp_path, + ) + + with pytest.raises(MinerUServiceException, match=r"\*/auto directories"): + _flatten_extracted_zip(str(tmp_path)) + + +def test_flatten_preserves_images_when_only_images_present(tmp_path: Path) -> None: + _write_zip( + { + "report/auto/images/page-1-fig-1.jpg": b"\x00img1", + "report/auto/images/page-2-fig-2.png": b"\x00img2", + }, + tmp_path, + ) + + with pytest.raises(MinerUServiceException, match="no markdown"): + _flatten_extracted_zip(str(tmp_path)) + + image_names = sorted(p.name for p in (tmp_path / "images").iterdir()) + assert image_names == ["page-1-fig-1.jpg", "page-2-fig-2.png"] + + +def test_flatten_excludes_metadata_json(tmp_path: Path) -> None: + _write_zip( + { + "report/auto/report.md": b"# Heading", + "report/auto/images/page-1-fig-1.jpg": b"\x00img1", + "report/auto/content_list.json": b"[]", + "report/auto/middle.json": b"{}", + "report/auto/model.json": b"{}", + "report/auto/keep.json": b"{}", + }, + tmp_path, + ) + + _flatten_extracted_zip(str(tmp_path)) + + assert (tmp_path / "full.md").read_text() == "# Heading" + assert (tmp_path / "keep.json").exists() + assert not (tmp_path / "content_list.json").exists() + assert not (tmp_path / "middle.json").exists() + assert not (tmp_path / "model.json").exists() + assert not (tmp_path / "report").exists() diff --git a/packages/shared-python/shared/core/config/mineru.py b/packages/shared-python/shared/core/config/mineru.py index 1b1a4389..45b48163 100644 --- a/packages/shared-python/shared/core/config/mineru.py +++ b/packages/shared-python/shared/core/config/mineru.py @@ -68,3 +68,36 @@ class MineruConfig(BaseModel): default=3600, description="Presigned URL TTL in seconds for S3 URL mode ingestion.", ) + MINERU_LOCAL_MODE: bool = Field( + default=False, + description=( + "Use a local (self-hosted) MinerU instance via the synchronous " + "/file_parse endpoint instead of the cloud batch APIs. When true, " + "MINERU_URL should point at the local MinerU base URL " + "(e.g. http://host.docker.internal:8000)." + ), + ) + MINERU_LOCAL_LANG_LIST: str = Field( + default="ch", + description=( + "Language code passed to local MinerU's /file_parse lang_list " + "parameter. Local MinerU does not accept 'auto'; the default 'ch' " + "covers Chinese, English, Japanese, Traditional Chinese, and Latin." + ), + ) + MINERU_LOCAL_BACKEND: str = Field( + default="pipeline", + description=( + "Backend passed to local MinerU's /file_parse backend parameter. " + "Use 'pipeline' for CPU-only parsing or 'vlm-engine' for GPU VLM." + ), + ) + MINERU_LOCAL_TIMEOUT: int = Field( + default=3600, + description=( + "Per-shard timeout in seconds for local MinerU's synchronous " + "/file_parse call. Includes queue wait when " + "MINERU_SHARD_CONCURRENCY > 1, because local MinerU is " + "single-concurrency by default." + ), + )