feat: local MinerU mode via synchronous /file_parse

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(