Files
hkuds--lightrag/lightrag/parser/external/mineru/client.py
T
2026-07-13 12:08:54 +08:00

703 lines
26 KiB
Python

"""MinerU raw bundle downloader.
Supports MinerU's official cloud and self-hosted API protocols and lands the
final parser bundle on disk under ``raw_dir/``:
- ``official`` — MinerU precision API v4: apply for signed upload URL, PUT the
local file, poll batch results, download ``full_zip_url``.
- ``local`` — self-hosted ``mineru-api`` / ``mineru-router``: submit
``POST /tasks``, poll ``GET /tasks/{task_id}``, download
``GET /tasks/{task_id}/result``.
Both protocols request a zip result bundle. Archives are extracted under
``raw_dir/`` and normalized so the adapter can read a root-level
``content_list.json``.
"""
from __future__ import annotations
import asyncio
import io
import json
import os
import shutil
import zipfile
from collections.abc import AsyncIterator, Mapping
from datetime import datetime, timezone
from pathlib import Path
from typing import TYPE_CHECKING, Any
from urllib.parse import quote, urlparse
from lightrag.parser.external._common import raise_for_status_with_detail
from lightrag.parser.external.mineru.cache import (
MinerUParserOptions,
compute_size_and_hash,
)
from lightrag.parser.external.mineru.manifest import (
Manifest,
ManifestFile,
write_manifest,
)
from lightrag.utils import logger
if TYPE_CHECKING:
import httpx
else:
try:
import httpx
except ImportError: # pragma: no cover
httpx = None
CONTENT_LIST_FILENAME = "content_list.json"
DEFAULT_MINERU_API_MODE = "local"
DEFAULT_MINERU_OFFICIAL_ENDPOINT = "https://mineru.net"
VALID_MINERU_API_MODES = {"official", "local"}
OFFICIAL_DONE_STATES = {"done"}
OFFICIAL_FAILED_STATES = {"failed"}
LOCAL_DONE_STATES = {"completed"}
LOCAL_FAILED_STATES = {"failed"}
UPLOAD_CHUNK_SIZE = 1024 * 1024
def _get_by_path(payload: Any, path: str) -> Any:
"""Walk a dotted path through a nested dict; returns None if any segment
is missing or non-dict."""
if not path:
return None
cur = payload
for part in path.split("."):
if isinstance(cur, dict) and part in cur:
cur = cur[part]
else:
return None
return cur
def _strip_trailing_slash(url: str) -> str:
return url.rstrip("/")
def _resolve_upload_name(upload_name: str | None, source_file_path: Path) -> str:
candidate = Path(str(upload_name or "")).name
return candidate or source_file_path.name
async def _iter_file_bytes(path: Path) -> AsyncIterator[bytes]:
with path.open("rb") as fh:
while True:
chunk = await asyncio.to_thread(fh.read, UPLOAD_CHUNK_SIZE)
if not chunk:
break
yield chunk
def _validate_base_url(
name: str, endpoint: str, forbidden_segments: tuple[str, ...]
) -> None:
parsed = urlparse(endpoint)
path = (parsed.path or "").rstrip("/")
for segment in forbidden_segments:
if path.endswith(segment) or f"{segment}/" in path:
raise ValueError(
f"{name} must be a base URL, not an API path: {endpoint!r}"
)
class MinerURawClient:
"""Downloads MinerU bundles into ``raw_dir``.
Construct once per call (cheap). Reads ``MINERU_*`` env vars at
construction time. Methods are async and use a single shared httpx
client across all calls in :meth:`download_into`.
Implements the MinerU-specific upload + poll + zip download flow
inline; bundle handling needs the ``result_url`` *and* the
``Content-Type`` of the response, which a generic protocol helper
cannot expose without leaking abstractions.
"""
def __init__(self, *, overrides: "Mapping[str, Any] | None" = None) -> None:
self._overrides = overrides or {}
self.api_mode = (
os.getenv("MINERU_API_MODE", DEFAULT_MINERU_API_MODE).strip().lower()
)
if self.api_mode not in VALID_MINERU_API_MODES:
allowed = ", ".join(sorted(VALID_MINERU_API_MODES))
raise ValueError(
f"MINERU_API_MODE must be one of {allowed}, got {self.api_mode!r}"
)
self.official_endpoint = _strip_trailing_slash(
os.getenv(
"MINERU_OFFICIAL_ENDPOINT", DEFAULT_MINERU_OFFICIAL_ENDPOINT
).strip()
or DEFAULT_MINERU_OFFICIAL_ENDPOINT
)
self.local_endpoint = _strip_trailing_slash(
os.getenv("MINERU_LOCAL_ENDPOINT", "").strip()
)
self.api_token = os.getenv("MINERU_API_TOKEN", "").strip()
if self.api_mode == "official":
if not self.api_token:
raise ValueError(
"MINERU_API_TOKEN is required when MINERU_API_MODE=official"
)
_validate_base_url(
"MINERU_OFFICIAL_ENDPOINT",
self.official_endpoint,
("/api/v4", "/api/v4/file-urls/batch", "/api/v4/extract/task"),
)
self.endpoint = self.official_endpoint
elif self.api_mode == "local":
if not self.local_endpoint:
raise ValueError(
"MINERU_LOCAL_ENDPOINT is required when MINERU_API_MODE=local"
)
_validate_base_url(
"MINERU_LOCAL_ENDPOINT",
self.local_endpoint,
("/tasks", "/file_parse", "/health"),
)
self.endpoint = self.local_endpoint
self.poll_interval = float(os.getenv("MINERU_POLL_INTERVAL_SECONDS", "2"))
# 600 * 2s client-side sleep ≈ 20 min worst case; raise for very large PDFs.
self.max_polls = int(os.getenv("MINERU_MAX_POLLS", "600"))
self.engine_version = os.getenv("MINERU_ENGINE_VERSION", "").strip()
options = MinerUParserOptions.from_env(
api_mode=self.api_mode, overrides=self._overrides
)
self._parser_options = options
self.model_version = options.model_version
self.language = options.language
self.enable_table = options.enable_table
self.enable_formula = options.enable_formula
self.is_ocr = options.is_ocr
self.page_ranges = options.page_ranges
self.local_backend = options.local_backend
self.local_parse_method = options.local_parse_method
self.local_image_analysis = options.local_image_analysis
self.local_start_page_id = options.local_start_page_id
self.local_end_page_id = options.local_end_page_id
# ------------------------------------------------------------------
# Public API
# ------------------------------------------------------------------
async def download_into(
self,
raw_dir: Path,
source_file_path: Path,
*,
upload_name: str | None = None,
) -> Manifest:
"""Download a fresh bundle and write the manifest.
Pre-condition: caller cleared ``raw_dir`` contents (recommended via
:func:`clear_dir_contents`). This method does NOT clean the
directory itself — leaving that to the caller keeps cache miss
semantics explicit at the parse_mineru entry point.
Returns the :class:`Manifest` describing the bundle.
"""
if httpx is None:
raise RuntimeError("httpx is required for MinerU parsing but not installed")
raw_dir.mkdir(parents=True, exist_ok=True)
resolved_upload_name = _resolve_upload_name(upload_name, source_file_path)
timeout = httpx.Timeout(120.0, connect=30.0)
try:
async with httpx.AsyncClient(timeout=timeout) as client:
if self.api_mode == "official":
task_id = await self._download_official(
client, source_file_path, raw_dir, resolved_upload_name
)
else:
task_id = await self._download_local(
client, source_file_path, raw_dir, resolved_upload_name
)
except httpx.RequestError as exc:
# Transport-level failures (connection refused/reset, server
# disconnect, read/connect timeout) bubble up from httpx with an
# opaque, sometimes empty message like "All connection attempts
# failed" that gives no hint the parse engine was MinerU. HTTP
# status errors and protocol errors already raise context-rich
# RuntimeErrors via raise_for_status_with_detail, so they stay
# untouched. Re-raise with the engine + endpoint and the exception
# class name so the doc_status error_msg is always non-empty and
# clearly attributable to the MinerU backend.
raise RuntimeError(
f"MinerU {self.api_mode} backend request failed "
f"(endpoint={self.endpoint}): {type(exc).__name__}: {exc}"
) from exc
self._normalize_raw_bundle(raw_dir, source_file_path, resolved_upload_name)
return self._build_and_write_manifest(
raw_dir, source_file_path, task_id, resolved_upload_name
)
# ------------------------------------------------------------------
# Upload + poll
# ------------------------------------------------------------------
def _official_headers(self) -> dict[str, str]:
return {
"Content-Type": "application/json",
"Authorization": f"Bearer {self.api_token}",
}
def _official_payload(self, upload_name: str) -> dict[str, Any]:
file_entry: dict[str, Any] = {"name": upload_name}
if self.is_ocr:
file_entry["is_ocr"] = True
if self.page_ranges:
file_entry["page_ranges"] = self.page_ranges
return {
"files": [file_entry],
"model_version": self.model_version,
"language": self.language,
"enable_table": self.enable_table,
"enable_formula": self.enable_formula,
}
async def _download_official(
self,
client: "httpx.AsyncClient",
source_file_path: Path,
raw_dir: Path,
upload_name: str,
) -> str:
apply_url = f"{self.official_endpoint}/api/v4/file-urls/batch"
resp = await client.post(
apply_url,
headers=self._official_headers(),
json=self._official_payload(upload_name),
)
raise_for_status_with_detail(resp, "MinerU official upload URL request")
payload = resp.json() if resp.text else {}
self._raise_if_official_error(payload, "MinerU official upload URL request")
data = payload.get("data") if isinstance(payload, dict) else {}
batch_id = str((data or {}).get("batch_id") or "")
file_urls = (data or {}).get("file_urls") or []
if not batch_id or not isinstance(file_urls, list) or not file_urls:
raise RuntimeError(
f"MinerU official upload URL response missing batch_id/file_urls: "
f"{payload}"
)
first_file_url = file_urls[0]
if isinstance(first_file_url, dict):
upload_url = str(
first_file_url.get("url") or first_file_url.get("file_url") or ""
)
else:
upload_url = str(first_file_url)
if not upload_url:
raise RuntimeError(
f"MinerU official upload URL response had an empty upload URL: "
f"{payload}"
)
upload_resp = await client.put(
upload_url,
content=_iter_file_bytes(source_file_path),
headers={"Content-Length": str(source_file_path.stat().st_size)},
)
raise_for_status_with_detail(upload_resp, "MinerU official file upload")
result_url = await self._poll_official_batch(client, batch_id, upload_name)
await self._download_zip(client, result_url, raw_dir)
return batch_id
async def _poll_official_batch(
self,
client: "httpx.AsyncClient",
batch_id: str,
upload_name: str,
) -> str:
encoded_batch_id = quote(batch_id, safe="")
poll_url = (
f"{self.official_endpoint}/api/v4/extract-results/batch/{encoded_batch_id}"
)
for _ in range(self.max_polls):
await asyncio.sleep(self.poll_interval)
resp = await client.get(poll_url, headers=self._official_headers())
raise_for_status_with_detail(resp, "MinerU official batch poll")
payload = resp.json() if resp.text else {}
self._raise_if_official_error(payload, "MinerU official batch poll")
results = _get_by_path(payload, "data.extract_result")
if isinstance(results, dict):
results = [results]
if not isinstance(results, list):
continue
selected = _select_official_extract_result(results, upload_name)
if selected is None:
continue
state = str(selected.get("state") or "").lower()
if state in OFFICIAL_DONE_STATES:
full_zip_url = str(selected.get("full_zip_url") or "")
if not full_zip_url:
raise RuntimeError(
f"MinerU official batch {batch_id} is done but has no "
f"full_zip_url: {selected}"
)
return full_zip_url
if state in OFFICIAL_FAILED_STATES:
err = selected.get("err_msg") or selected.get("error") or selected
raise RuntimeError(
f"MinerU official parse failed for batch {batch_id}: {err}"
)
raise TimeoutError(f"MinerU official batch polling timeout: {batch_id}")
def _raise_if_official_error(self, payload: Any, operation: str) -> None:
if not isinstance(payload, dict):
raise RuntimeError(f"{operation} returned non-object payload: {payload!r}")
code = payload.get("code", 0)
if code not in (0, "0", None):
raise RuntimeError(
f"{operation} failed: code={code} msg={payload.get('msg')!r}"
)
def _local_form_data(self) -> dict[str, str]:
return {
"lang_list": self.language,
"backend": self.local_backend,
"parse_method": self.local_parse_method,
"formula_enable": _bool_form(self.enable_formula),
"table_enable": _bool_form(self.enable_table),
"image_analysis": _bool_form(self.local_image_analysis),
"return_md": "true",
"return_middle_json": "true",
"return_model_output": "true",
"return_content_list": "true",
"return_images": "true",
"response_format_zip": "true",
"return_original_file": "true",
"start_page_id": str(self.local_start_page_id),
"end_page_id": str(self.local_end_page_id),
}
async def _download_local(
self,
client: "httpx.AsyncClient",
source_file_path: Path,
raw_dir: Path,
upload_name: str,
) -> str:
submit_url = f"{self.local_endpoint}/tasks"
# Keep data as a Mapping so httpx 0.28 builds an async MultipartStream
# and reads the file handle in chunks instead of buffering the payload.
with source_file_path.open("rb") as fh:
files = {"files": (upload_name, fh, "application/octet-stream")}
resp = await client.post(
submit_url,
data=self._local_form_data(),
files=files,
)
raise_for_status_with_detail(
resp,
f"MinerU local task submission for {upload_name!r}",
)
payload = resp.json() if resp.text else {}
task_id = str(payload.get("task_id") or "")
if not task_id:
raise RuntimeError(
f"MinerU local /tasks response missing task_id: {payload}"
)
await self._poll_local_task(client, task_id)
await self._download_zip(
client,
f"{self.local_endpoint}/tasks/{quote(task_id, safe='')}/result",
raw_dir,
)
return task_id
async def _poll_local_task(
self,
client: "httpx.AsyncClient",
task_id: str,
) -> None:
# ``task_id`` is service-returned; encode it as a single path segment so
# a crafted value can't break out of ``/tasks/{id}``. The raw value is
# kept for the error/timeout messages below where the real ID matters.
poll_url = f"{self.local_endpoint}/tasks/{quote(task_id, safe='')}"
for _ in range(self.max_polls):
await asyncio.sleep(self.poll_interval)
resp = await client.get(poll_url)
raise_for_status_with_detail(resp, "MinerU local task poll")
payload = resp.json() if resp.text else {}
status = str(payload.get("status") or "").lower()
if status in LOCAL_DONE_STATES:
return
if status in LOCAL_FAILED_STATES:
err = payload.get("error") or payload.get("message") or payload
raise RuntimeError(
f"MinerU local parse failed for task {task_id}: {err}"
)
raise TimeoutError(f"MinerU local task polling timeout: {task_id}")
async def _download_zip(
self,
client: "httpx.AsyncClient",
result_url: str,
raw_dir: Path,
resp: Any = None,
) -> None:
"""Download (or re-use already-fetched response) and extract."""
if resp is None or not hasattr(resp, "content"):
resp = await client.get(result_url)
raise_for_status_with_detail(resp, "MinerU result bundle download")
buf = io.BytesIO(resp.content)
with zipfile.ZipFile(buf) as zf:
# Safe-extract: refuse absolute paths and ``..`` traversal.
for name in zf.namelist():
norm = os.path.normpath(name)
if norm.startswith("..") or os.path.isabs(norm):
raise RuntimeError(f"Refusing zip entry with unsafe path: {name!r}")
zf.extractall(raw_dir)
# Normalize: if the zip nested everything under a single top-level
# dir, hoist its contents up so content_list.json sits at raw_dir
# root. This matches the common MinerU bundle layout.
self._maybe_hoist_single_subdir(raw_dir)
def _maybe_hoist_single_subdir(self, raw_dir: Path) -> None:
entries = [p for p in raw_dir.iterdir() if p.name != "_manifest.json"]
if len(entries) != 1 or not entries[0].is_dir():
return
sub = entries[0]
for child in list(sub.iterdir()):
child.rename(raw_dir / child.name)
try:
sub.rmdir()
except OSError:
pass
def _normalize_raw_bundle(
self,
raw_dir: Path,
source_file_path: Path,
upload_name: str | None = None,
) -> None:
"""Ensure a downloaded bundle has root-level ``content_list.json``.
Official and local MinerU zip archives commonly place parser outputs at
``<doc>/<parse_method>/<doc>_content_list.json``. The adapter consumes a
canonical root ``content_list.json`` plus optional root ``images/``.
After hoisting we delete the nested originals so the manifest does not
bookkeep two copies (and disk usage doesn't double for big bundles).
Sibling artifacts of the parse subdir (``*.md``, ``middle.json`` etc.)
are also hoisted to ``raw_dir`` root for easier diagnostics.
"""
if (raw_dir / CONTENT_LIST_FILENAME).is_file():
return
candidate = _select_content_list_candidate(
raw_dir, source_file_path, upload_name
)
if candidate is None:
return
source_dir = candidate.parent
target_root = raw_dir.resolve()
# Guard: never hoist from above raw_dir (defensive — candidate already
# comes from rglob inside raw_dir, but cheap to verify).
try:
source_dir.resolve().relative_to(target_root)
except ValueError:
shutil.copy2(candidate, raw_dir / CONTENT_LIST_FILENAME)
return
# Move the critical file first; then hoist sibling files/dirs that
# don't already exist at raw_dir root.
shutil.move(str(candidate), str(raw_dir / CONTENT_LIST_FILENAME))
for entry in list(source_dir.iterdir()):
target = raw_dir / entry.name
if target.exists():
continue
shutil.move(str(entry), str(target))
# Best-effort cleanup of the now-empty parse subtree.
cursor = source_dir
while cursor != raw_dir and cursor.is_dir():
try:
cursor.rmdir()
except OSError:
break
cursor = cursor.parent
# ------------------------------------------------------------------
# Manifest construction
# ------------------------------------------------------------------
def _build_and_write_manifest(
self,
raw_dir: Path,
source_file_path: Path,
task_id: str,
upload_name: str,
) -> Manifest:
source_size, source_hash = compute_size_and_hash(source_file_path)
# Critical file — required.
crit_path = raw_dir / CONTENT_LIST_FILENAME
if not crit_path.is_file():
raise RuntimeError(
f"MinerU bundle missing required {CONTENT_LIST_FILENAME} "
f"after download (raw_dir={raw_dir})"
)
crit_size, crit_hash = compute_size_and_hash(crit_path)
# Other files.
others: list[ManifestFile] = []
total = crit_size
for p in sorted(raw_dir.rglob("*")):
if not p.is_file():
continue
if p.name == "_manifest.json":
continue
rel = p.relative_to(raw_dir).as_posix()
if rel == CONTENT_LIST_FILENAME:
continue
size = p.stat().st_size
others.append(ManifestFile(path=rel, size=size))
total += size
manifest = Manifest(
source_content_hash=source_hash,
source_size_bytes=source_size,
source_filename_at_parse=upload_name,
critical_file=ManifestFile(
path=CONTENT_LIST_FILENAME,
size=crit_size,
sha256=crit_hash,
),
files=others,
total_size_bytes=total,
task_id=task_id,
api_mode=self.api_mode,
engine_version=self.engine_version,
endpoint_signature=self.endpoint,
options_signature=self._options_signature(),
downloaded_at=datetime.now(timezone.utc).isoformat(),
)
write_manifest(raw_dir, manifest)
return manifest
def _options_signature(self) -> str:
return self._parser_options.signature()
def _find_content_list(payload: Any, content_field: str) -> list[dict] | None:
"""Heuristic content_list extractor.
Tries (in order):
1. The provided dotted path if it lands on a list of dicts.
2. Direct ``content_list`` / ``content`` / ``items`` / ``result`` keys.
3. Recursive descent.
"""
if isinstance(payload, list):
if payload and all(isinstance(x, dict) for x in payload):
return payload
return None
if not isinstance(payload, dict):
return None
via_field = _get_by_path(payload, content_field)
candidate = _find_content_list(via_field, content_field)
if candidate is not None:
return candidate
for key in ("content_list", "content", "items", "result"):
value = payload.get(key)
candidate = _find_content_list(value, content_field)
if candidate is not None:
return candidate
for value in payload.values():
candidate = _find_content_list(value, content_field)
if candidate is not None:
return candidate
return None
def _bool_form(value: bool) -> str:
return "true" if value else "false"
def _select_official_extract_result(
results: list[Any],
source_filename: str,
) -> dict[str, Any] | None:
"""Pick the extract_result entry that matches the file we uploaded.
Invariant: :meth:`MinerURawClient._download_official` always submits a
single-file batch, so a non-matching ``file_name`` from the API would
indicate either a server response we don't understand or a future
multi-file extension. We fall back to ``dict_results[0]`` to remain
forward-compatible but log a warning so the mismatch is visible.
"""
dict_results = [item for item in results if isinstance(item, dict)]
if not dict_results:
return None
source_name = Path(source_filename).name
source_stem = Path(source_filename).stem
for item in dict_results:
file_name = str(item.get("file_name") or item.get("name") or "")
if Path(file_name).name == source_name or Path(file_name).stem == source_stem:
return item
logger.warning(
"[mineru_raw] official extract_result did not contain a match for "
"%r; falling back to the first entry (%r). This is unexpected for "
"a single-file batch.",
source_name,
str(dict_results[0].get("file_name") or dict_results[0].get("name") or ""),
)
return dict_results[0]
def _select_content_list_candidate(
raw_dir: Path,
source_file_path: Path,
upload_name: str | None = None,
) -> Path | None:
source_stem = Path(upload_name or source_file_path.name).stem
candidates: list[tuple[int, int, str, Path]] = []
for path in raw_dir.rglob("*.json"):
if not path.is_file():
continue
if path.name != CONTENT_LIST_FILENAME and not path.name.endswith(
"_content_list.json"
):
continue
try:
payload = json.loads(path.read_text(encoding="utf-8"))
except (OSError, json.JSONDecodeError):
continue
content_list = _find_content_list(payload, "content")
if content_list is None:
continue
score = 10
if path.name == CONTENT_LIST_FILENAME:
score = 0
elif path.name == f"{source_stem}_content_list.json":
score = 1
elif path.stem.endswith("_content_list"):
score = 2
depth = len(path.relative_to(raw_dir).parts)
candidates.append((score, depth, path.as_posix(), path))
if not candidates:
return None
candidates.sort()
return candidates[0][3]
__all__ = ["MinerURawClient", "CONTENT_LIST_FILENAME"]