chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
Vendored
+50
@@ -0,0 +1,50 @@
|
||||
"""Adapters for external document parsing services.
|
||||
|
||||
Each subpackage under ``parser/external/`` integrates one external parser
|
||||
(docling, mineru, ...) by handling:
|
||||
|
||||
- request/upload/poll choreography against the parser's HTTP API,
|
||||
- on-disk caching of the raw bundle under ``<base>.<engine>_raw/``,
|
||||
- normalization into LightRAG IR (``IRDoc``) for the sidecar writer.
|
||||
|
||||
Shared cross-engine helpers (size/hash, atomic manifest IO, safe zip
|
||||
extraction, env coercion) live at this package root in private modules
|
||||
prefixed ``_``. Engine-specific cache validation, manifest construction,
|
||||
and IR adaptation live in each subpackage.
|
||||
"""
|
||||
|
||||
from lightrag.parser.external._common import (
|
||||
clear_dir_contents,
|
||||
compute_size_and_hash,
|
||||
env_bool,
|
||||
env_int,
|
||||
env_json,
|
||||
raw_dir_for_parsed_dir,
|
||||
)
|
||||
from lightrag.parser.external._manifest import (
|
||||
MANIFEST_FILENAME,
|
||||
MANIFEST_VERSION,
|
||||
Manifest,
|
||||
ManifestFile,
|
||||
load_manifest,
|
||||
manifest_path,
|
||||
write_manifest,
|
||||
)
|
||||
from lightrag.parser.external._zip import safe_extract_zip
|
||||
|
||||
__all__ = [
|
||||
"MANIFEST_FILENAME",
|
||||
"MANIFEST_VERSION",
|
||||
"Manifest",
|
||||
"ManifestFile",
|
||||
"clear_dir_contents",
|
||||
"compute_size_and_hash",
|
||||
"env_bool",
|
||||
"env_int",
|
||||
"env_json",
|
||||
"load_manifest",
|
||||
"manifest_path",
|
||||
"raw_dir_for_parsed_dir",
|
||||
"safe_extract_zip",
|
||||
"write_manifest",
|
||||
]
|
||||
Vendored
+189
@@ -0,0 +1,189 @@
|
||||
"""Shared template for external (download + raw-bundle cache) parser engines.
|
||||
|
||||
``ExternalParserBase.parse`` fixes the common MinerU/Docling flow once:
|
||||
|
||||
resolve → raw_dir → force-reparse check → cache-hit skip
|
||||
else (mkdir + clear_dir_contents + download_into) → build_ir
|
||||
→ write_sidecar → persist full_docs (lightrag) → archive source
|
||||
|
||||
Subclasses implement three engine-private hooks (``is_bundle_valid`` /
|
||||
``download_into`` / ``build_ir``) and set ``raw_dir_suffix`` /
|
||||
``force_reparse_env``. This is the reshaped #3207 contract — now an
|
||||
*internal* template rather than the top-level parser interface — with its two
|
||||
gaps fixed: ``clear_dir_contents`` runs inside the template (cache-miss only),
|
||||
and the per-engine upload-name divergence is normalised to a single
|
||||
``upload_name`` hook parameter.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from abc import abstractmethod
|
||||
from collections.abc import Mapping
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from lightrag.constants import FULL_DOCS_FORMAT_LIGHTRAG
|
||||
from lightrag.parser.base import BaseParser, ParseContext, ParseResult
|
||||
from lightrag.utils import logger
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from lightrag.sidecar.ir import IRDoc
|
||||
|
||||
|
||||
class ExternalParserBase(BaseParser):
|
||||
"""Base for engines that fetch a raw bundle from an external service."""
|
||||
|
||||
raw_dir_suffix: str
|
||||
force_reparse_env: str
|
||||
|
||||
# --- engine-private hooks ------------------------------------------------
|
||||
@abstractmethod
|
||||
def is_bundle_valid(
|
||||
self,
|
||||
raw_dir: Path,
|
||||
source_path: Path,
|
||||
*,
|
||||
engine_params: Mapping[str, Any] | None = None,
|
||||
) -> bool:
|
||||
"""Cheap cache-hit check against the raw bundle on disk.
|
||||
|
||||
``engine_params`` is the per-file engine-parameter override (decoded
|
||||
from ``parse_engine``); it MUST participate in the cache signature so an
|
||||
overridden document does not spuriously hit a bundle parsed with
|
||||
different params.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def download_into(
|
||||
self,
|
||||
raw_dir: Path,
|
||||
source_path: Path,
|
||||
*,
|
||||
upload_name: str,
|
||||
engine_params: Mapping[str, Any] | None = None,
|
||||
) -> None:
|
||||
"""Fetch the raw bundle into ``raw_dir`` (called on cache miss only).
|
||||
|
||||
``engine_params`` is the per-file engine-parameter override applied to
|
||||
both the request payload and the recorded cache signature.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def build_ir(self, raw_dir: Path, document_name: str) -> "IRDoc":
|
||||
"""Convert the raw bundle to an :class:`IRDoc`."""
|
||||
...
|
||||
|
||||
def validate_ir(self, ir: "IRDoc", *, file_path: str, raw_dir: Path) -> None:
|
||||
"""Optional post-build validation hook (default no-op)."""
|
||||
|
||||
# --- template ------------------------------------------------------------
|
||||
async def parse(self, ctx: ParseContext) -> ParseResult:
|
||||
from lightrag.parser.external._common import (
|
||||
clear_dir_contents,
|
||||
env_bool,
|
||||
raw_dir_for_parsed_dir,
|
||||
)
|
||||
from lightrag.parser.routing import decode_parse_engine, encode_parse_engine
|
||||
from lightrag.sidecar import write_sidecar
|
||||
from lightrag.utils_pipeline import (
|
||||
make_lightrag_doc_content,
|
||||
sidecar_uri_for,
|
||||
)
|
||||
|
||||
# Per-file engine params are encoded in the stored ``parse_engine``
|
||||
# directive (e.g. ``mineru(page_range=1-3)``); decode them once and
|
||||
# thread the SAME dict into both the cache-hit check and the download so
|
||||
# an overridden doc can never hit a bundle parsed with different params.
|
||||
# A malformed/corrupt directive fails this doc loudly rather than
|
||||
# silently parsing with no params.
|
||||
_engine, engine_params, decode_errs = decode_parse_engine(
|
||||
ctx.content_data.get("parse_engine")
|
||||
if isinstance(ctx.content_data, dict)
|
||||
else None
|
||||
)
|
||||
if decode_errs:
|
||||
raise ValueError(
|
||||
f"{self.engine_name}: invalid parse_engine for doc_id={ctx.doc_id}: "
|
||||
+ "; ".join(decode_errs)
|
||||
)
|
||||
engine_params = engine_params or None
|
||||
|
||||
rs = ctx.resolve(self.engine_name)
|
||||
source = rs.source_path
|
||||
if not source.is_file():
|
||||
raise FileNotFoundError(
|
||||
f"{self.engine_name} source file not found: {source}"
|
||||
)
|
||||
raw_dir = raw_dir_for_parsed_dir(rs.parsed_dir, suffix=self.raw_dir_suffix)
|
||||
force_reparse = env_bool(self.force_reparse_env, False)
|
||||
|
||||
parse_stage_skipped = False
|
||||
if not force_reparse and self.is_bundle_valid(
|
||||
raw_dir, source, engine_params=engine_params
|
||||
):
|
||||
# Cache hit: stay purely local so a re-parse still works when the
|
||||
# external endpoint is temporarily unavailable.
|
||||
parse_stage_skipped = True
|
||||
logger.info("[%s] raw cache hit doc_id=%s", self.engine_name, ctx.doc_id)
|
||||
else:
|
||||
if force_reparse and raw_dir.exists():
|
||||
logger.info(
|
||||
"[%s] %s set; discarding bundle at %s",
|
||||
self.engine_name,
|
||||
self.force_reparse_env,
|
||||
raw_dir,
|
||||
)
|
||||
# download_into mkdir's raw_dir; we wipe stale contents first so a
|
||||
# previous bundle cannot leak into the new one (cache-miss only).
|
||||
raw_dir.mkdir(parents=True, exist_ok=True)
|
||||
clear_dir_contents(raw_dir)
|
||||
logger.info(
|
||||
"[%s] Parsing %s %s (may take a few minutes)",
|
||||
self.engine_name,
|
||||
ctx.doc_id,
|
||||
source.name,
|
||||
)
|
||||
await self.download_into(
|
||||
raw_dir,
|
||||
source,
|
||||
upload_name=rs.document_name,
|
||||
engine_params=engine_params,
|
||||
)
|
||||
|
||||
ir = self.build_ir(raw_dir, rs.document_name)
|
||||
self.validate_ir(ir, file_path=ctx.file_path, raw_dir=raw_dir)
|
||||
parsed_data = write_sidecar(
|
||||
ir,
|
||||
parsed_dir=rs.parsed_dir,
|
||||
doc_id=ctx.doc_id,
|
||||
engine=self.engine_name,
|
||||
)
|
||||
|
||||
await ctx.rag._persist_parsed_full_docs(
|
||||
ctx.doc_id,
|
||||
{
|
||||
"content": make_lightrag_doc_content(parsed_data["content"]),
|
||||
"file_path": ctx.file_path,
|
||||
"parse_format": FULL_DOCS_FORMAT_LIGHTRAG,
|
||||
"sidecar_location": sidecar_uri_for(rs.parsed_dir),
|
||||
# Re-encode the engine + params so the persisted directive keeps
|
||||
# the per-file params (the `{**existing, **record}` merge in
|
||||
# _persist_parsed_full_docs would otherwise revert it to the
|
||||
# bare engine name).
|
||||
"parse_engine": encode_parse_engine(self.engine_name, engine_params),
|
||||
"update_time": int(time.time()),
|
||||
},
|
||||
)
|
||||
await ctx.archive_source(str(source))
|
||||
return ParseResult(
|
||||
doc_id=ctx.doc_id,
|
||||
file_path=ctx.file_path,
|
||||
parse_format=FULL_DOCS_FORMAT_LIGHTRAG,
|
||||
content=parsed_data["content"],
|
||||
blocks_path=parsed_data["blocks_path"],
|
||||
parse_engine=self.engine_name,
|
||||
parse_stage_skipped=parse_stage_skipped,
|
||||
)
|
||||
Vendored
+152
@@ -0,0 +1,152 @@
|
||||
"""Shared helpers for ``lightrag/parser/external/<engine>/`` packages.
|
||||
|
||||
Currently consumed by the docling subpackage; expected to be reused when
|
||||
mineru is migrated under ``parser/external/mineru/``.
|
||||
|
||||
These are pure functions with no engine-specific knowledge. Engine-specific
|
||||
logic (endpoint signature, options signature, cache validation policy) lives
|
||||
in each engine's own ``cache.py``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from lightrag.constants import PARSED_DIR_SUFFIX
|
||||
from lightrag.utils import logger
|
||||
|
||||
|
||||
def compute_size_and_hash(path: Path) -> tuple[int, str]:
|
||||
"""Single-read computation of ``(size_bytes, "sha256:<hex>")``.
|
||||
|
||||
Manifest writes use this so the recorded size and hash are guaranteed to
|
||||
describe the same byte stream; using two ``open()`` calls would risk a
|
||||
TOCTOU mismatch if the file changed in between.
|
||||
"""
|
||||
h = hashlib.sha256()
|
||||
size = 0
|
||||
with path.open("rb") as f:
|
||||
for chunk in iter(lambda: f.read(1 << 20), b""):
|
||||
h.update(chunk)
|
||||
size += len(chunk)
|
||||
return size, f"sha256:{h.hexdigest()}"
|
||||
|
||||
|
||||
def clear_dir_contents(directory: Path) -> None:
|
||||
"""Delete everything inside ``directory`` but keep ``directory`` itself."""
|
||||
if not directory.exists():
|
||||
return
|
||||
for entry in directory.iterdir():
|
||||
try:
|
||||
if entry.is_dir() and not entry.is_symlink():
|
||||
shutil.rmtree(entry, ignore_errors=True)
|
||||
else:
|
||||
entry.unlink()
|
||||
except OSError:
|
||||
continue
|
||||
|
||||
|
||||
def raw_dir_for_parsed_dir(parsed_dir: Path, *, suffix: str) -> Path:
|
||||
"""Sibling raw dir for a ``*.parsed`` dir.
|
||||
|
||||
``foo.parsed/`` with ``suffix=".docling_raw"`` → ``foo.docling_raw/``.
|
||||
``suffix`` must start with ``.`` and be engine-specific (the caller
|
||||
binds it via ``functools.partial`` or a thin wrapper).
|
||||
"""
|
||||
if not suffix.startswith("."):
|
||||
raise ValueError(f"raw dir suffix must start with '.', got {suffix!r}")
|
||||
stem = parsed_dir.name
|
||||
if stem.endswith(PARSED_DIR_SUFFIX):
|
||||
stem = stem[: -len(PARSED_DIR_SUFFIX)]
|
||||
return parsed_dir.parent / f"{stem}{suffix}"
|
||||
|
||||
|
||||
def env_bool(name: str, default: bool) -> bool:
|
||||
raw = os.getenv(name, "").strip().lower()
|
||||
if raw in {"1", "true", "yes", "on"}:
|
||||
return True
|
||||
if raw in {"0", "false", "no", "off"}:
|
||||
return False
|
||||
return default
|
||||
|
||||
|
||||
def env_int(name: str, default: int) -> int:
|
||||
raw = os.getenv(name, "").strip()
|
||||
if not raw:
|
||||
return default
|
||||
try:
|
||||
return int(raw)
|
||||
except ValueError:
|
||||
logger.warning(
|
||||
"[external_parser] %s=%r is not an integer; using %s", name, raw, default
|
||||
)
|
||||
return default
|
||||
|
||||
|
||||
def env_json(name: str, default: Any) -> Any:
|
||||
"""Parse a JSON env var; on parse error log a warning and return default."""
|
||||
raw = os.getenv(name, "").strip()
|
||||
if not raw:
|
||||
return default
|
||||
try:
|
||||
return json.loads(raw)
|
||||
except json.JSONDecodeError:
|
||||
logger.warning(
|
||||
"[external_parser] %s=%r is not valid JSON; using default", name, raw
|
||||
)
|
||||
return default
|
||||
|
||||
|
||||
def response_error_detail(resp: Any, *, limit: int = 1000) -> str:
|
||||
"""Return a compact response body snippet for HTTP error reporting."""
|
||||
try:
|
||||
payload = resp.json() if getattr(resp, "text", "") else None
|
||||
except Exception:
|
||||
payload = None
|
||||
|
||||
if payload is not None:
|
||||
try:
|
||||
detail = json.dumps(payload, ensure_ascii=False, sort_keys=True)
|
||||
except TypeError:
|
||||
detail = repr(payload)
|
||||
else:
|
||||
detail = str(getattr(resp, "text", "") or "").strip()
|
||||
|
||||
detail = " ".join(detail.split())
|
||||
if not detail:
|
||||
return "empty response body"
|
||||
if len(detail) > limit:
|
||||
return f"{detail[:limit]}...<truncated>"
|
||||
return detail
|
||||
|
||||
|
||||
def raise_for_status_with_detail(resp: Any, operation: str) -> None:
|
||||
"""Raise an HTTP error that preserves service-provided response details.
|
||||
|
||||
Treats any non-2xx response as an error, matching httpx's
|
||||
``raise_for_status`` status handling (which also raises on 1xx/3xx,
|
||||
not just 4xx/5xx) while attaching a compact response-body snippet to
|
||||
the message for faster diagnosis.
|
||||
"""
|
||||
status_code = int(getattr(resp, "status_code", 0) or 0)
|
||||
if 200 <= status_code < 300:
|
||||
return
|
||||
detail = response_error_detail(resp)
|
||||
raise RuntimeError(f"{operation} failed: HTTP {status_code} {detail}")
|
||||
|
||||
|
||||
__all__ = [
|
||||
"clear_dir_contents",
|
||||
"compute_size_and_hash",
|
||||
"env_bool",
|
||||
"env_int",
|
||||
"env_json",
|
||||
"raise_for_status_with_detail",
|
||||
"raw_dir_for_parsed_dir",
|
||||
"response_error_detail",
|
||||
]
|
||||
+167
@@ -0,0 +1,167 @@
|
||||
"""Shared ``_manifest.json`` schema for ``parser/external/<engine>/`` bundles.
|
||||
|
||||
The manifest is the *atomic success marker* for a raw bundle. Its presence
|
||||
implies "all files in this directory finished downloading"; its content is
|
||||
the cache key for "is this bundle for the same source file, the same engine
|
||||
version, the same endpoint, and the same option signature we are using right
|
||||
now?".
|
||||
|
||||
Write path: :func:`write_manifest` writes a temp file then atomically renames
|
||||
to ``_manifest.json``. A crash mid-download leaves no manifest, so the next
|
||||
parse call cleanly invalidates and re-downloads.
|
||||
|
||||
Read path: :func:`load_manifest` returns ``None`` if absent, malformed, or
|
||||
recorded under a different engine — either way the bundle is treated as
|
||||
stale.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
from dataclasses import asdict, dataclass, field
|
||||
from pathlib import Path
|
||||
|
||||
MANIFEST_FILENAME = "_manifest.json"
|
||||
MANIFEST_VERSION = "1.0"
|
||||
|
||||
|
||||
@dataclass
|
||||
class ManifestFile:
|
||||
"""One file entry inside the bundle. Size always; sha256 only for files
|
||||
where silent corruption would break the adapter (the "critical" file).
|
||||
"""
|
||||
|
||||
path: str # relative to the raw dir
|
||||
size: int
|
||||
sha256: str | None = None # ``"sha256:<hex>"`` or ``None``
|
||||
|
||||
|
||||
@dataclass
|
||||
class Manifest:
|
||||
"""Generic manifest schema. ``engine`` is filled by the caller (docling /
|
||||
mineru / etc.); ``options_signature`` lets per-engine cache layers detect
|
||||
when env-driven request parameters changed without bumping the version.
|
||||
"""
|
||||
|
||||
engine: str
|
||||
source_content_hash: str
|
||||
source_size_bytes: int
|
||||
source_filename_at_parse: str
|
||||
critical_file: ManifestFile
|
||||
files: list[ManifestFile]
|
||||
total_size_bytes: int
|
||||
task_id: str = ""
|
||||
api_mode: str = ""
|
||||
engine_version: str = ""
|
||||
endpoint_signature: str = ""
|
||||
options_signature: str = ""
|
||||
downloaded_at: str = ""
|
||||
extras: dict = field(default_factory=dict)
|
||||
version: str = MANIFEST_VERSION
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
return {
|
||||
"version": self.version,
|
||||
"engine": self.engine,
|
||||
"api_mode": self.api_mode,
|
||||
"engine_version": self.engine_version,
|
||||
"endpoint_signature": self.endpoint_signature,
|
||||
"options_signature": self.options_signature,
|
||||
"source_content_hash": self.source_content_hash,
|
||||
"source_size_bytes": int(self.source_size_bytes),
|
||||
"source_filename_at_parse": self.source_filename_at_parse,
|
||||
"task_id": self.task_id,
|
||||
"downloaded_at": self.downloaded_at,
|
||||
"critical_file": asdict(self.critical_file),
|
||||
"files": [asdict(f) for f in self.files],
|
||||
"total_size_bytes": int(self.total_size_bytes),
|
||||
"extras": dict(self.extras or {}),
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, payload: dict) -> "Manifest":
|
||||
critical_raw = payload.get("critical_file") or {}
|
||||
files_raw = payload.get("files") or []
|
||||
return cls(
|
||||
version=str(payload.get("version") or MANIFEST_VERSION),
|
||||
engine=str(payload.get("engine") or ""),
|
||||
api_mode=str(payload.get("api_mode") or ""),
|
||||
engine_version=str(payload.get("engine_version") or ""),
|
||||
endpoint_signature=str(payload.get("endpoint_signature") or ""),
|
||||
options_signature=str(payload.get("options_signature") or ""),
|
||||
source_content_hash=str(payload.get("source_content_hash") or ""),
|
||||
source_size_bytes=int(payload.get("source_size_bytes") or 0),
|
||||
source_filename_at_parse=str(payload.get("source_filename_at_parse") or ""),
|
||||
task_id=str(payload.get("task_id") or ""),
|
||||
downloaded_at=str(payload.get("downloaded_at") or ""),
|
||||
critical_file=ManifestFile(
|
||||
path=str(critical_raw.get("path") or ""),
|
||||
size=int(critical_raw.get("size") or 0),
|
||||
sha256=(
|
||||
str(critical_raw["sha256"]) if critical_raw.get("sha256") else None
|
||||
),
|
||||
),
|
||||
files=[
|
||||
ManifestFile(
|
||||
path=str(f.get("path") or ""),
|
||||
size=int(f.get("size") or 0),
|
||||
sha256=(str(f["sha256"]) if f.get("sha256") else None),
|
||||
)
|
||||
for f in files_raw
|
||||
if isinstance(f, dict)
|
||||
],
|
||||
total_size_bytes=int(payload.get("total_size_bytes") or 0),
|
||||
extras=dict(payload.get("extras") or {}),
|
||||
)
|
||||
|
||||
|
||||
def manifest_path(raw_dir: Path) -> Path:
|
||||
return raw_dir / MANIFEST_FILENAME
|
||||
|
||||
|
||||
def load_manifest(raw_dir: Path, *, expected_engine: str) -> Manifest | None:
|
||||
"""Return the parsed manifest or ``None`` if absent / malformed / for a
|
||||
different engine. ``expected_engine`` is required so a future shared raw
|
||||
dir cannot serve a bundle that belongs to another engine.
|
||||
"""
|
||||
p = manifest_path(raw_dir)
|
||||
if not p.is_file():
|
||||
return None
|
||||
try:
|
||||
payload = json.loads(p.read_text(encoding="utf-8"))
|
||||
except (OSError, json.JSONDecodeError):
|
||||
return None
|
||||
if not isinstance(payload, dict):
|
||||
return None
|
||||
if payload.get("version") != MANIFEST_VERSION:
|
||||
return None
|
||||
if payload.get("engine") != expected_engine:
|
||||
return None
|
||||
try:
|
||||
return Manifest.from_dict(payload)
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
|
||||
|
||||
def write_manifest(raw_dir: Path, manifest: Manifest) -> None:
|
||||
"""Atomically write the manifest using temp-file + rename."""
|
||||
raw_dir.mkdir(parents=True, exist_ok=True)
|
||||
final = manifest_path(raw_dir)
|
||||
tmp = final.with_suffix(".json.tmp")
|
||||
tmp.write_text(
|
||||
json.dumps(manifest.to_dict(), ensure_ascii=False, indent=2),
|
||||
encoding="utf-8",
|
||||
)
|
||||
os.replace(tmp, final)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"MANIFEST_FILENAME",
|
||||
"MANIFEST_VERSION",
|
||||
"Manifest",
|
||||
"ManifestFile",
|
||||
"load_manifest",
|
||||
"manifest_path",
|
||||
"write_manifest",
|
||||
]
|
||||
Vendored
+67
@@ -0,0 +1,67 @@
|
||||
"""Shared zip-bundle extraction for external parser engines.
|
||||
|
||||
Engines like docling return their full output as a zip archive. This helper
|
||||
extracts it safely (refusing path traversal / absolute paths) into a target
|
||||
directory. Engine-specific post-extraction normalization (e.g. mineru's
|
||||
nested-subdir hoist) is *not* done here — each engine's client handles its
|
||||
own quirks.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import io
|
||||
import os
|
||||
import zipfile
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def safe_extract_zip(
|
||||
payload: bytes,
|
||||
dest_dir: Path,
|
||||
*,
|
||||
max_entries: int | None = None,
|
||||
max_total_bytes: int | None = None,
|
||||
) -> list[str]:
|
||||
"""Extract a zip archive into ``dest_dir``, refusing unsafe paths.
|
||||
|
||||
Raises ``RuntimeError`` if any entry name is absolute or contains ``..``
|
||||
components after normalization. Returns the list of extracted member
|
||||
names (as stored in the zip, prior to OS-specific normalization), so
|
||||
callers can validate the bundle layout without re-walking the directory.
|
||||
|
||||
Optional zip-bomb guards (both default ``None`` = unlimited, preserving the
|
||||
original behaviour for existing callers): ``max_entries`` caps the member
|
||||
count and ``max_total_bytes`` caps the summed *uncompressed* size declared
|
||||
in the archive's central directory. Both are checked from ``infolist()``
|
||||
metadata *before* extraction, so a malicious archive is rejected without
|
||||
being written to disk.
|
||||
"""
|
||||
dest_dir.mkdir(parents=True, exist_ok=True)
|
||||
buf = io.BytesIO(payload)
|
||||
with zipfile.ZipFile(buf) as zf:
|
||||
infos = zf.infolist()
|
||||
if max_entries is not None and len(infos) > max_entries:
|
||||
raise RuntimeError(
|
||||
f"Refusing zip with {len(infos)} entries (max {max_entries})"
|
||||
)
|
||||
if max_total_bytes is not None:
|
||||
total = sum(info.file_size for info in infos)
|
||||
if total > max_total_bytes:
|
||||
raise RuntimeError(
|
||||
f"Refusing zip: uncompressed size {total} bytes "
|
||||
f"exceeds limit {max_total_bytes}"
|
||||
)
|
||||
names = zf.namelist()
|
||||
for name in names:
|
||||
norm = os.path.normpath(name)
|
||||
if (
|
||||
norm.startswith("..")
|
||||
or os.path.isabs(norm)
|
||||
or norm.startswith(("/", os.sep))
|
||||
):
|
||||
raise RuntimeError(f"Refusing zip entry with unsafe path: {name!r}")
|
||||
zf.extractall(dest_dir)
|
||||
return names
|
||||
|
||||
|
||||
__all__ = ["safe_extract_zip"]
|
||||
+41
@@ -0,0 +1,41 @@
|
||||
"""Docling parser integration (raw client, cache, manifest, IR adapter).
|
||||
|
||||
Public surface for the rest of the codebase. ``parse_docling`` imports
|
||||
only from this facade so the inner module layout stays free to evolve.
|
||||
"""
|
||||
|
||||
from lightrag.constants import DOCLING_RAW_DIR_SUFFIX
|
||||
from lightrag.parser.external._common import (
|
||||
clear_dir_contents,
|
||||
raw_dir_for_parsed_dir as _raw_dir_for_parsed_dir,
|
||||
)
|
||||
|
||||
MANIFEST_ENGINE = "docling"
|
||||
|
||||
|
||||
def raw_dir_for_parsed_dir(parsed_dir):
|
||||
"""``foo.parsed/`` → ``foo.docling_raw/`` (docling-specific binding)."""
|
||||
return _raw_dir_for_parsed_dir(parsed_dir, suffix=DOCLING_RAW_DIR_SUFFIX)
|
||||
|
||||
|
||||
# Imported after ``MANIFEST_ENGINE`` / ``DOCLING_RAW_DIR_SUFFIX`` because
|
||||
# the submodules read those constants at import time.
|
||||
from lightrag.parser.external.docling.ir_builder import ( # noqa: E402
|
||||
DoclingIRBuilder,
|
||||
)
|
||||
from lightrag.parser.external.docling.cache import ( # noqa: E402
|
||||
is_bundle_valid,
|
||||
)
|
||||
from lightrag.parser.external.docling.client import ( # noqa: E402
|
||||
DoclingRawClient,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"DOCLING_RAW_DIR_SUFFIX",
|
||||
"MANIFEST_ENGINE",
|
||||
"DoclingIRBuilder",
|
||||
"DoclingRawClient",
|
||||
"clear_dir_contents",
|
||||
"is_bundle_valid",
|
||||
"raw_dir_for_parsed_dir",
|
||||
]
|
||||
+257
@@ -0,0 +1,257 @@
|
||||
"""Cache validation for ``*.docling_raw/`` bundles.
|
||||
|
||||
Validation policy (settled in
|
||||
``docs/DoclingSidecarRefactorPlan-zh.md`` §4.1):
|
||||
|
||||
1. ``_manifest.json`` exists, parses, ``engine="docling"`` ∧ schema version
|
||||
matches.
|
||||
2. **Source size fast-path**: ``source_file.stat().st_size`` matches the
|
||||
manifest; mismatch → miss without hashing.
|
||||
3. **Source content_hash**: full sha256 of the current source file matches
|
||||
the manifest.
|
||||
4. **Engine version**: if ``DOCLING_ENGINE_VERSION`` is set in env and the
|
||||
manifest recorded a non-empty value, they must match.
|
||||
5. **Endpoint signature**: if the active ``DOCLING_ENDPOINT`` differs from
|
||||
what was recorded at parse time, miss (avoids re-using a bundle produced
|
||||
by a different docling-serve instance).
|
||||
6. **Options signature**: covers every env or fixed constant that changes
|
||||
the produced bundle (OCR flags, language list, formula enrichment,
|
||||
target format and pipeline). Any change → miss.
|
||||
7. **Critical file**: the main JSON must exist with matching size **and**
|
||||
sha256 — final tie-breaker against silent corruption affecting the file
|
||||
the adapter depends on.
|
||||
8. **Other files**: size-only verification (cheap; covers most corruption
|
||||
modes for markdown / artifacts).
|
||||
|
||||
Any failed step ⇒ cache miss; the caller wipes the directory contents and
|
||||
re-runs the download.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
from collections.abc import Mapping
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from lightrag.parser.external._common import compute_size_and_hash, env_bool
|
||||
from lightrag.parser.external._manifest import load_manifest
|
||||
from lightrag.parser.external.docling import MANIFEST_ENGINE
|
||||
from lightrag.utils import logger
|
||||
|
||||
# Legacy upload-path suffix. ``env.example`` historically documented
|
||||
# ``DOCLING_ENDPOINT=http://host:5001/v1/convert/file/async`` (the full
|
||||
# upload URL); the current client expects a base URL and appends the path
|
||||
# itself. Strip the suffix so an unmodified pre-refactor ``.env`` keeps
|
||||
# working instead of producing
|
||||
# ``/v1/convert/file/async/v1/convert/file/async`` requests.
|
||||
_LEGACY_UPLOAD_PATH_SUFFIX = "/v1/convert/file/async"
|
||||
_legacy_endpoint_warned = False
|
||||
|
||||
# Envs that change the bytes docling-serve produces. Any change here must
|
||||
# invalidate the bundle cache. ``DOCLING_BBOX_ATTRIBUTES`` is intentionally
|
||||
# NOT in this list: it only affects how the adapter writes IR meta, not the
|
||||
# docling bundle, so flipping it should re-emit the sidecar (which we always
|
||||
# do) without forcing a re-download.
|
||||
DOCLING_TUNABLE_ENVS: tuple[str, ...] = (
|
||||
"DOCLING_DO_OCR",
|
||||
"DOCLING_FORCE_OCR",
|
||||
"DOCLING_OCR_ENGINE",
|
||||
"DOCLING_OCR_PRESET",
|
||||
"DOCLING_OCR_LANG",
|
||||
"DOCLING_DO_FORMULA_ENRICHMENT",
|
||||
)
|
||||
|
||||
|
||||
def current_endpoint_signature() -> str:
|
||||
"""The active docling endpoint, normalized to a base URL.
|
||||
|
||||
Normalization:
|
||||
|
||||
- Trims surrounding whitespace and strips trailing slashes.
|
||||
- Strips the legacy ``/v1/convert/file/async`` upload suffix if present,
|
||||
preserving backwards compatibility with the pre-refactor ``env.example``
|
||||
that documented the full upload URL.
|
||||
|
||||
Returns ``""`` if ``DOCLING_ENDPOINT`` is unset — callers that need a
|
||||
real endpoint (``DoclingRawClient``) raise on empty; callers that only
|
||||
compare against a recorded manifest field (``is_bundle_valid``) silently
|
||||
skip the check when either side is empty.
|
||||
"""
|
||||
global _legacy_endpoint_warned
|
||||
endpoint = os.getenv("DOCLING_ENDPOINT", "").strip().rstrip("/")
|
||||
if endpoint.endswith(_LEGACY_UPLOAD_PATH_SUFFIX):
|
||||
endpoint = endpoint[: -len(_LEGACY_UPLOAD_PATH_SUFFIX)]
|
||||
if not _legacy_endpoint_warned:
|
||||
_legacy_endpoint_warned = True
|
||||
logger.warning(
|
||||
"DOCLING_ENDPOINT still includes the legacy %r upload suffix; "
|
||||
"stripping it. Update your .env to a base URL "
|
||||
"(e.g. http://host:5001).",
|
||||
_LEGACY_UPLOAD_PATH_SUFFIX,
|
||||
)
|
||||
return endpoint
|
||||
|
||||
|
||||
def compute_options_signature(
|
||||
*,
|
||||
tunable_env: dict[str, str],
|
||||
fixed_constants: dict[str, object],
|
||||
) -> str:
|
||||
"""Stable signature over user-tunable env values and fixed pipeline
|
||||
constants.
|
||||
|
||||
Storing the constants in the signature means a future code change that
|
||||
flips e.g. ``image_export_mode`` from ``referenced`` to ``embedded``
|
||||
invalidates every existing cache without anyone having to remember to
|
||||
bump a version.
|
||||
"""
|
||||
payload = json.dumps(
|
||||
{"env": tunable_env, "fixed": fixed_constants},
|
||||
ensure_ascii=False,
|
||||
sort_keys=True,
|
||||
separators=(",", ":"),
|
||||
)
|
||||
return "sha256:" + hashlib.sha256(payload.encode("utf-8")).hexdigest()
|
||||
|
||||
|
||||
def snapshot_tunable_env(
|
||||
overrides: "Mapping[str, Any] | None" = None,
|
||||
) -> dict[str, str]:
|
||||
"""Read effective docling tunables so equivalent requests share a signature.
|
||||
|
||||
``overrides`` carries decoded per-file engine params (Phase 2: ``force_ocr``)
|
||||
and replaces the corresponding env value so the override feeds BOTH the live
|
||||
request and the cache signature.
|
||||
"""
|
||||
overrides = overrides or {}
|
||||
force_ocr = (
|
||||
bool(overrides["force_ocr"])
|
||||
if "force_ocr" in overrides
|
||||
else env_bool("DOCLING_FORCE_OCR", True)
|
||||
)
|
||||
return {
|
||||
"DOCLING_DO_OCR": str(env_bool("DOCLING_DO_OCR", True)).lower(),
|
||||
"DOCLING_FORCE_OCR": str(force_ocr).lower(),
|
||||
"DOCLING_OCR_ENGINE": os.getenv("DOCLING_OCR_ENGINE", "auto").strip() or "auto",
|
||||
"DOCLING_OCR_PRESET": os.getenv("DOCLING_OCR_PRESET", "auto").strip() or "auto",
|
||||
"DOCLING_OCR_LANG": os.getenv("DOCLING_OCR_LANG", "").strip(),
|
||||
"DOCLING_DO_FORMULA_ENRICHMENT": str(
|
||||
env_bool("DOCLING_DO_FORMULA_ENRICHMENT", False)
|
||||
).lower(),
|
||||
}
|
||||
|
||||
|
||||
def is_bundle_valid(
|
||||
raw_dir: Path,
|
||||
source_file: Path,
|
||||
*,
|
||||
overrides: "Mapping[str, Any] | None" = None,
|
||||
) -> bool:
|
||||
"""Return True iff the bundle matches the current source + env state."""
|
||||
if not raw_dir.is_dir():
|
||||
return False
|
||||
|
||||
manifest = load_manifest(raw_dir, expected_engine=MANIFEST_ENGINE)
|
||||
if manifest is None:
|
||||
return False
|
||||
|
||||
# 1. Source size fast-path
|
||||
try:
|
||||
cur_size = source_file.stat().st_size
|
||||
except OSError:
|
||||
return False
|
||||
if cur_size != int(manifest.source_size_bytes):
|
||||
return False
|
||||
|
||||
# 2. Source content_hash
|
||||
_, cur_hash = compute_size_and_hash(source_file)
|
||||
if cur_hash != manifest.source_content_hash:
|
||||
return False
|
||||
|
||||
# 3. Engine version. Skip the comparison when either side is empty so
|
||||
# operators can opt out by unsetting the env, and so bundles from
|
||||
# earlier code that never recorded the field aren't force-invalidated.
|
||||
cur_engine_version = os.getenv("DOCLING_ENGINE_VERSION", "").strip()
|
||||
if (
|
||||
cur_engine_version
|
||||
and manifest.engine_version
|
||||
and cur_engine_version != manifest.engine_version
|
||||
):
|
||||
return False
|
||||
|
||||
# 4. Endpoint signature. Same "both non-empty to compare" rule: a bundle
|
||||
# parsed against a different docling-serve URL must not be reused, but
|
||||
# we don't reject the cache just because the env happens to be unset
|
||||
# at validation time (e.g. CLI tooling that only reads the cache).
|
||||
cur_endpoint = current_endpoint_signature()
|
||||
if (
|
||||
cur_endpoint
|
||||
and manifest.endpoint_signature
|
||||
and cur_endpoint != manifest.endpoint_signature
|
||||
):
|
||||
return False
|
||||
|
||||
# 5. Options signature: only enforced if the manifest recorded one
|
||||
# (manifests written before this commit have it empty — they are
|
||||
# treated as stale and re-downloaded the next time the env changes).
|
||||
#
|
||||
# Compare against the *current* fixed constants from client.py, not
|
||||
# the copy stashed in the manifest — using the manifest's copy would
|
||||
# always reproduce the recorded signature and silently swallow
|
||||
# code-only changes (e.g. flipping image_export_mode or to_formats),
|
||||
# defeating the invalidation this step is supposed to provide.
|
||||
# Lazy import: client.py imports from cache.py.
|
||||
#
|
||||
# When per-file overrides are requested (e.g. docling(force_ocr=true)) but
|
||||
# the manifest predates signature recording, we cannot prove the bundle was
|
||||
# produced with those overrides — accepting it would silently drop the
|
||||
# user's explicit param. Treat that as a miss so the override is honored
|
||||
# (mirrors MinerU, which misses on any absent signature). The no-override
|
||||
# case keeps the deliberate leniency above for legacy bundles.
|
||||
if overrides and not manifest.options_signature:
|
||||
return False
|
||||
if manifest.options_signature:
|
||||
from lightrag.parser.external.docling.client import FIXED_CONSTANTS
|
||||
|
||||
cur_options = compute_options_signature(
|
||||
tunable_env=snapshot_tunable_env(overrides),
|
||||
fixed_constants=FIXED_CONSTANTS,
|
||||
)
|
||||
if cur_options != manifest.options_signature:
|
||||
return False
|
||||
|
||||
# 6. Critical file: size + sha256
|
||||
crit = manifest.critical_file
|
||||
crit_path = raw_dir / crit.path
|
||||
try:
|
||||
if crit_path.stat().st_size != int(crit.size):
|
||||
return False
|
||||
except OSError:
|
||||
return False
|
||||
if crit.sha256:
|
||||
_, crit_actual = compute_size_and_hash(crit_path)
|
||||
if crit_actual != crit.sha256:
|
||||
return False
|
||||
|
||||
# 7. Other files: size only
|
||||
for entry in manifest.files:
|
||||
ep = raw_dir / entry.path
|
||||
try:
|
||||
if ep.stat().st_size != int(entry.size):
|
||||
return False
|
||||
except OSError:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
|
||||
__all__ = [
|
||||
"DOCLING_TUNABLE_ENVS",
|
||||
"compute_options_signature",
|
||||
"current_endpoint_signature",
|
||||
"is_bundle_valid",
|
||||
"snapshot_tunable_env",
|
||||
]
|
||||
+353
@@ -0,0 +1,353 @@
|
||||
"""Docling raw bundle downloader.
|
||||
|
||||
Talks to Docling Serve v1 over HTTP:
|
||||
|
||||
- ``POST /v1/convert/file/async`` — multipart upload, returns ``task_id``,
|
||||
- ``GET /v1/status/poll/{task_id}?wait=5`` — long-poll for terminal state,
|
||||
- ``GET /v1/result/{task_id}`` — zip download (only on ``success``).
|
||||
|
||||
The zip is extracted safely under ``raw_dir/`` (refusing path traversal /
|
||||
absolute entries). A success manifest is written atomically at the very
|
||||
end; mid-run crashes therefore leave the directory in a state the cache
|
||||
layer marks as invalid (no manifest → miss → re-download).
|
||||
|
||||
Pipeline constants (``pipeline``, ``target_type``, ``to_formats``,
|
||||
``image_export_mode``) are intentionally **not** env-driven — the sidecar
|
||||
flow depends on them — and are recorded inside the manifest so a future
|
||||
code change automatically invalidates pre-existing caches.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
from collections.abc import Mapping
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any
|
||||
from urllib.parse import quote
|
||||
|
||||
from lightrag.parser.external._common import (
|
||||
env_bool,
|
||||
env_int,
|
||||
raise_for_status_with_detail,
|
||||
)
|
||||
from lightrag.parser.external._zip import safe_extract_zip
|
||||
from lightrag.parser.external.docling.cache import (
|
||||
compute_options_signature,
|
||||
current_endpoint_signature,
|
||||
snapshot_tunable_env,
|
||||
)
|
||||
from lightrag.parser.external.docling.manifest import (
|
||||
build_and_write_docling_manifest,
|
||||
select_main_json,
|
||||
)
|
||||
from lightrag.utils import logger
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import httpx
|
||||
else:
|
||||
try:
|
||||
import httpx
|
||||
except ImportError: # pragma: no cover
|
||||
httpx = None
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fixed pipeline constants (NOT env-driven)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
PIPELINE = "standard"
|
||||
TARGET_TYPE = "zip"
|
||||
TO_FORMATS: tuple[str, ...] = ("json", "md")
|
||||
IMAGE_EXPORT_MODE = "referenced"
|
||||
|
||||
FIXED_CONSTANTS: dict[str, object] = {
|
||||
"pipeline": PIPELINE,
|
||||
"target_type": TARGET_TYPE,
|
||||
"to_formats": list(TO_FORMATS),
|
||||
"image_export_mode": IMAGE_EXPORT_MODE,
|
||||
}
|
||||
|
||||
CONVERT_PATH = "/v1/convert/file/async"
|
||||
POLL_PATH = "/v1/status/poll/{task_id}"
|
||||
RESULT_PATH = "/v1/result/{task_id}"
|
||||
|
||||
DEFAULT_POLL_WAIT_SECONDS = 5
|
||||
DEFAULT_MAX_POLLS = 240 # 240 * 5s long-poll ≈ 20 min worst case
|
||||
|
||||
# ConversionStatus enum from the docling-serve OpenAPI
|
||||
SUCCESS_STATES = {"success"}
|
||||
FAILURE_STATES = {"failure", "partial_success", "skipped"}
|
||||
IN_PROGRESS_STATES = {"pending", "started"}
|
||||
|
||||
|
||||
class DoclingRawClient:
|
||||
"""Downloads docling-serve bundles into ``raw_dir``.
|
||||
|
||||
Construct once per parse call (cheap). Reads ``DOCLING_*`` envs at
|
||||
``__init__`` time, so callers can flip env between calls and pick up
|
||||
the new values without holding a stale instance.
|
||||
"""
|
||||
|
||||
def __init__(self, *, overrides: "Mapping[str, Any] | None" = None) -> None:
|
||||
self._overrides = overrides or {}
|
||||
self.endpoint = current_endpoint_signature()
|
||||
if not self.endpoint:
|
||||
raise ValueError("DOCLING_ENDPOINT is required")
|
||||
self.engine_version = os.getenv("DOCLING_ENGINE_VERSION", "").strip()
|
||||
|
||||
self.do_ocr = env_bool("DOCLING_DO_OCR", True)
|
||||
self.force_ocr = (
|
||||
bool(self._overrides["force_ocr"])
|
||||
if "force_ocr" in self._overrides
|
||||
else env_bool("DOCLING_FORCE_OCR", True)
|
||||
)
|
||||
self.ocr_engine = os.getenv("DOCLING_OCR_ENGINE", "auto").strip() or "auto"
|
||||
self.ocr_preset = os.getenv("DOCLING_OCR_PRESET", "auto").strip() or "auto"
|
||||
self.ocr_lang_raw = os.getenv("DOCLING_OCR_LANG", "").strip()
|
||||
self.do_formula_enrichment = env_bool("DOCLING_DO_FORMULA_ENRICHMENT", False)
|
||||
|
||||
# Poll cadence: docling-serve's ``?wait=N`` is a server-side long-poll
|
||||
# window. ``DOCLING_POLL_INTERVAL_SECONDS`` sets that window; the
|
||||
# client does NOT add its own sleep between polls. ``DOCLING_MAX_POLLS``
|
||||
# bounds the total polling budget — exceeding it raises ``TimeoutError``.
|
||||
wait = env_int("DOCLING_POLL_INTERVAL_SECONDS", DEFAULT_POLL_WAIT_SECONDS)
|
||||
self.poll_wait_seconds = wait if wait > 0 else DEFAULT_POLL_WAIT_SECONDS
|
||||
max_polls = env_int("DOCLING_MAX_POLLS", DEFAULT_MAX_POLLS)
|
||||
self.max_poll_attempts = max_polls if max_polls > 0 else DEFAULT_MAX_POLLS
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Public API
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def download_into(
|
||||
self,
|
||||
raw_dir: Path,
|
||||
source_file_path: Path,
|
||||
*,
|
||||
upload_filename: str | None = None,
|
||||
):
|
||||
"""Upload, poll, download, extract, and write the manifest.
|
||||
|
||||
``upload_filename`` overrides the multipart filename sent to
|
||||
docling-serve (defaults to ``source_file_path.name``). The pipeline
|
||||
passes the canonical, hint-stripped document name here so the
|
||||
bundle's ``<stem>.json`` ends up canonical too — otherwise a file
|
||||
named ``report.[docling].pdf`` would produce ``report.[docling].json``
|
||||
inside the bundle, and the adapter (which only knows the canonical
|
||||
``report.pdf``) would not be able to locate it via the preferred
|
||||
``<stem>.json`` lookup.
|
||||
|
||||
Pre-condition: caller cleared ``raw_dir`` (e.g. via
|
||||
:func:`lightrag.parser.external.clear_dir_contents`). This method
|
||||
does not clean the directory itself — keeping that explicit at the
|
||||
``parse_docling`` entry point.
|
||||
"""
|
||||
if httpx is None:
|
||||
raise RuntimeError(
|
||||
"httpx is required for Docling parsing but is not installed"
|
||||
)
|
||||
raw_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
effective_filename = upload_filename or source_file_path.name
|
||||
|
||||
timeout = httpx.Timeout(120.0, connect=30.0)
|
||||
async with httpx.AsyncClient(timeout=timeout) as client:
|
||||
task_id = await self._submit(
|
||||
client, source_file_path, filename=effective_filename
|
||||
)
|
||||
await self._poll_until_done(client, task_id)
|
||||
payload = await self._download_zip_bytes(client, task_id)
|
||||
|
||||
safe_extract_zip(payload, raw_dir)
|
||||
# Defensive: confirm the main JSON exists before anyone reads the
|
||||
# bundle. Look it up by the *uploaded* filename's stem — that's
|
||||
# what docling-serve uses to name the JSON inside the zip.
|
||||
select_main_json(raw_dir, Path(effective_filename))
|
||||
|
||||
options_signature = compute_options_signature(
|
||||
tunable_env=snapshot_tunable_env(self._overrides),
|
||||
fixed_constants=FIXED_CONSTANTS,
|
||||
)
|
||||
return build_and_write_docling_manifest(
|
||||
raw_dir,
|
||||
source_file_path=source_file_path,
|
||||
task_id=task_id,
|
||||
endpoint_signature=self.endpoint,
|
||||
engine_version=self.engine_version,
|
||||
options_signature=options_signature,
|
||||
fixed_constants=FIXED_CONSTANTS,
|
||||
recorded_filename=effective_filename,
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Upload + poll + download
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _build_multipart_data(self) -> dict[str, str | list[str]]:
|
||||
"""Form fields (everything except the file payload).
|
||||
|
||||
Returns a ``dict`` (not a list of tuples): httpx ≥ 0.28 short-circuits
|
||||
non-``Mapping`` ``data`` into raw-content encoding and ignores
|
||||
``files=`` entirely, producing a sync-only stream that an
|
||||
``AsyncClient`` then rejects. List-valued entries are emitted as
|
||||
repeated form keys by ``MultipartStream``, matching docling-serve's
|
||||
pydantic ``List[Enum]`` form parsing. ``ocr_lang`` is omitted entirely
|
||||
when empty so the engine uses its own default.
|
||||
"""
|
||||
data: dict[str, str | list[str]] = {
|
||||
"pipeline": PIPELINE,
|
||||
"target_type": TARGET_TYPE,
|
||||
"image_export_mode": IMAGE_EXPORT_MODE,
|
||||
"do_ocr": _bool_form(self.do_ocr),
|
||||
"force_ocr": _bool_form(self.force_ocr),
|
||||
"ocr_engine": self.ocr_engine,
|
||||
"ocr_preset": self.ocr_preset,
|
||||
"do_formula_enrichment": _bool_form(self.do_formula_enrichment),
|
||||
"to_formats": list(TO_FORMATS),
|
||||
}
|
||||
if self.ocr_lang_raw:
|
||||
langs = _parse_ocr_lang(self.ocr_lang_raw)
|
||||
if langs:
|
||||
data["ocr_lang"] = langs
|
||||
return data
|
||||
|
||||
async def _submit(
|
||||
self,
|
||||
client: "httpx.AsyncClient",
|
||||
source_file_path: Path,
|
||||
*,
|
||||
filename: str,
|
||||
) -> str:
|
||||
url = f"{self.endpoint}{CONVERT_PATH}"
|
||||
# Hand httpx a file object so its MultipartStream reads the body in
|
||||
# chunks instead of materializing the whole PDF/PPTX in worker memory.
|
||||
# With ``max_parallel_parse_docling > 1`` a per-doc bytes copy can
|
||||
# OOM the worker before docling-serve ever sees the request.
|
||||
with source_file_path.open("rb") as fh:
|
||||
files = {"files": (filename, fh, "application/octet-stream")}
|
||||
resp = await client.post(
|
||||
url, data=self._build_multipart_data(), files=files
|
||||
)
|
||||
raise_for_status_with_detail(resp, f"Docling upload for {filename!r}")
|
||||
payload = resp.json() if resp.text else {}
|
||||
task_id = str(payload.get("task_id") or payload.get("id") or "").strip()
|
||||
if not task_id:
|
||||
raise RuntimeError(f"Docling upload response missing task_id: {payload!r}")
|
||||
return task_id
|
||||
|
||||
async def _poll_until_done(
|
||||
self,
|
||||
client: "httpx.AsyncClient",
|
||||
task_id: str,
|
||||
) -> None:
|
||||
encoded_task_id = quote(task_id, safe="")
|
||||
url = f"{self.endpoint}{POLL_PATH.format(task_id=encoded_task_id)}"
|
||||
params = {"wait": self.poll_wait_seconds}
|
||||
for _ in range(self.max_poll_attempts):
|
||||
iteration_started = time.monotonic()
|
||||
resp = await client.get(url, params=params)
|
||||
raise_for_status_with_detail(resp, f"Docling task {task_id} poll")
|
||||
payload = resp.json() if resp.text else {}
|
||||
status = str(
|
||||
payload.get("task_status") or payload.get("status") or ""
|
||||
).lower()
|
||||
|
||||
if status in SUCCESS_STATES:
|
||||
return
|
||||
if status in FAILURE_STATES:
|
||||
raise RuntimeError(_format_failure(task_id, status, payload))
|
||||
if status not in IN_PROGRESS_STATES:
|
||||
# Unknown status: keep polling, but surface it so operators notice.
|
||||
logger.warning(
|
||||
"[docling] unknown task status %r for task %s; continuing to poll",
|
||||
status,
|
||||
task_id,
|
||||
)
|
||||
|
||||
# The intended cadence is one poll per ``poll_wait_seconds`` — the
|
||||
# design relies on docling-serve's ``?wait=N`` long-polling for
|
||||
# that. Some deployments return immediately instead, which would
|
||||
# burn through ``max_poll_attempts`` in milliseconds and fail
|
||||
# with a spurious timeout. Cap each iteration at the configured
|
||||
# interval ourselves so the total budget holds either way.
|
||||
elapsed = time.monotonic() - iteration_started
|
||||
remaining = self.poll_wait_seconds - elapsed
|
||||
if remaining > 0:
|
||||
await asyncio.sleep(remaining)
|
||||
|
||||
raise TimeoutError(f"Docling task {task_id} polling timeout")
|
||||
|
||||
async def _download_zip_bytes(
|
||||
self,
|
||||
client: "httpx.AsyncClient",
|
||||
task_id: str,
|
||||
) -> bytes:
|
||||
encoded_task_id = quote(task_id, safe="")
|
||||
url = f"{self.endpoint}{RESULT_PATH.format(task_id=encoded_task_id)}"
|
||||
resp = await client.get(url)
|
||||
raise_for_status_with_detail(resp, f"Docling result {task_id} download")
|
||||
ctype = resp.headers.get("content-type", "")
|
||||
if "zip" not in ctype.lower():
|
||||
raise RuntimeError(
|
||||
f"Docling result {task_id} returned non-zip content-type "
|
||||
f"{ctype!r}; body prefix={resp.text[:400]!r}"
|
||||
)
|
||||
return resp.content
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _bool_form(v: bool) -> str:
|
||||
return "true" if v else "false"
|
||||
|
||||
|
||||
def _parse_ocr_lang(raw: str) -> list[str]:
|
||||
"""Best-effort parser for ``DOCLING_OCR_LANG``.
|
||||
|
||||
Accepts a JSON array (``["en","zh"]``) or a comma-separated list
|
||||
(``en,zh``). Returns a list of stripped non-empty strings; empty in →
|
||||
empty out.
|
||||
"""
|
||||
try:
|
||||
parsed = json.loads(raw)
|
||||
except json.JSONDecodeError:
|
||||
parsed = None
|
||||
if isinstance(parsed, list):
|
||||
return [str(x).strip() for x in parsed if str(x).strip()]
|
||||
return [item.strip() for item in raw.split(",") if item.strip()]
|
||||
|
||||
|
||||
def _format_failure(task_id: str, status: str, payload: Any) -> str:
|
||||
if isinstance(payload, dict):
|
||||
err = (
|
||||
payload.get("error_message")
|
||||
or payload.get("error")
|
||||
or payload.get("message")
|
||||
or "<no error_message>"
|
||||
)
|
||||
else:
|
||||
err = "<no error_message>"
|
||||
truncated = json.dumps(payload, ensure_ascii=False)[:400]
|
||||
return f"Docling task {task_id} ended in {status}: {err}; payload={truncated}"
|
||||
|
||||
|
||||
__all__ = [
|
||||
"DoclingRawClient",
|
||||
"CONVERT_PATH",
|
||||
"DEFAULT_MAX_POLLS",
|
||||
"DEFAULT_POLL_WAIT_SECONDS",
|
||||
"FIXED_CONSTANTS",
|
||||
"IMAGE_EXPORT_MODE",
|
||||
"PIPELINE",
|
||||
"POLL_PATH",
|
||||
"RESULT_PATH",
|
||||
"SUCCESS_STATES",
|
||||
"FAILURE_STATES",
|
||||
"TARGET_TYPE",
|
||||
"TO_FORMATS",
|
||||
]
|
||||
+1072
File diff suppressed because it is too large
Load Diff
+130
@@ -0,0 +1,130 @@
|
||||
"""Helpers for building ``_manifest.json`` for docling raw bundles.
|
||||
|
||||
Wraps the generic :class:`Manifest` schema with docling-specific knowledge:
|
||||
|
||||
- the critical file is the main ``<stem>.json`` produced by docling-serve,
|
||||
- non-critical files are the markdown + every entry under ``artifacts/``,
|
||||
- ``extras`` carries the fixed pipeline constants so the options signature
|
||||
remains reproducible across runs.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
|
||||
from lightrag.parser.external._common import compute_size_and_hash
|
||||
from lightrag.parser.external._manifest import (
|
||||
MANIFEST_FILENAME,
|
||||
Manifest,
|
||||
ManifestFile,
|
||||
write_manifest,
|
||||
)
|
||||
from lightrag.parser.external.docling import MANIFEST_ENGINE
|
||||
|
||||
|
||||
def select_main_json(raw_dir: Path, source_file_path: Path) -> Path:
|
||||
"""Locate the primary docling JSON inside ``raw_dir``.
|
||||
|
||||
Priority: ``<source_stem>.json`` if present, else the single ``*.json``
|
||||
sitting at ``raw_dir`` root (excluding ``_manifest.json``, which always
|
||||
sits in the bundle once a download has completed and would otherwise
|
||||
collide with the bundle JSON in the fallback). Raises ``RuntimeError``
|
||||
if zero or multiple candidates exist.
|
||||
"""
|
||||
preferred = raw_dir / f"{source_file_path.stem}.json"
|
||||
if preferred.is_file():
|
||||
return preferred
|
||||
|
||||
candidates = sorted(
|
||||
p for p in raw_dir.glob("*.json") if p.is_file() and p.name != MANIFEST_FILENAME
|
||||
)
|
||||
if len(candidates) == 1:
|
||||
return candidates[0]
|
||||
if not candidates:
|
||||
raise RuntimeError(f"Docling raw bundle at {raw_dir} contains no .json file")
|
||||
names = ", ".join(p.name for p in candidates)
|
||||
raise RuntimeError(
|
||||
f"Docling raw bundle at {raw_dir} has multiple .json candidates ({names}); "
|
||||
f"expected exactly one to derive the critical file from"
|
||||
)
|
||||
|
||||
|
||||
def select_main_md(raw_dir: Path, source_file_path: Path) -> Path | None:
|
||||
"""Locate the markdown twin of the main JSON. Returns ``None`` if no
|
||||
markdown was produced (defensive — docling-serve always emits one for
|
||||
``to_formats=["json","md"]`` but we don't want to crash if it is
|
||||
missing)."""
|
||||
preferred = raw_dir / f"{source_file_path.stem}.md"
|
||||
if preferred.is_file():
|
||||
return preferred
|
||||
candidates = sorted(p for p in raw_dir.glob("*.md") if p.is_file())
|
||||
return candidates[0] if candidates else None
|
||||
|
||||
|
||||
def build_and_write_docling_manifest(
|
||||
raw_dir: Path,
|
||||
*,
|
||||
source_file_path: Path,
|
||||
task_id: str,
|
||||
endpoint_signature: str,
|
||||
engine_version: str,
|
||||
options_signature: str,
|
||||
fixed_constants: dict[str, object],
|
||||
recorded_filename: str | None = None,
|
||||
) -> Manifest:
|
||||
"""Construct the manifest for a freshly downloaded docling bundle and
|
||||
persist it atomically. Returns the in-memory manifest for callers that
|
||||
need the task_id / signatures for logging.
|
||||
|
||||
``recorded_filename`` is the name passed to docling-serve at upload
|
||||
time (canonical, hint-stripped form when called from the pipeline).
|
||||
It governs both the preferred-path lookup for the bundle JSON and the
|
||||
value persisted as ``source_filename_at_parse``. When ``None``, falls
|
||||
back to ``source_file_path.name`` for backward compatibility.
|
||||
"""
|
||||
lookup_path = Path(recorded_filename) if recorded_filename else source_file_path
|
||||
main_json = select_main_json(raw_dir, lookup_path)
|
||||
crit_size, crit_hash = compute_size_and_hash(main_json)
|
||||
critical = ManifestFile(
|
||||
path=main_json.relative_to(raw_dir).as_posix(),
|
||||
size=crit_size,
|
||||
sha256=crit_hash,
|
||||
)
|
||||
|
||||
others: list[ManifestFile] = []
|
||||
for path in sorted(raw_dir.rglob("*")):
|
||||
if not path.is_file():
|
||||
continue
|
||||
rel = path.relative_to(raw_dir).as_posix()
|
||||
if rel == critical.path or rel.startswith("_manifest"):
|
||||
continue
|
||||
others.append(ManifestFile(path=rel, size=path.stat().st_size))
|
||||
|
||||
source_size, source_hash = compute_size_and_hash(source_file_path)
|
||||
total = crit_size + sum(f.size for f in others)
|
||||
|
||||
manifest = Manifest(
|
||||
engine=MANIFEST_ENGINE,
|
||||
source_content_hash=source_hash,
|
||||
source_size_bytes=source_size,
|
||||
source_filename_at_parse=recorded_filename or source_file_path.name,
|
||||
critical_file=critical,
|
||||
files=others,
|
||||
total_size_bytes=total,
|
||||
task_id=task_id,
|
||||
endpoint_signature=endpoint_signature,
|
||||
engine_version=engine_version,
|
||||
options_signature=options_signature,
|
||||
downloaded_at=datetime.now(timezone.utc).isoformat(timespec="seconds"),
|
||||
extras={"fixed_constants": dict(fixed_constants)},
|
||||
)
|
||||
write_manifest(raw_dir, manifest)
|
||||
return manifest
|
||||
|
||||
|
||||
__all__ = [
|
||||
"build_and_write_docling_manifest",
|
||||
"select_main_json",
|
||||
"select_main_md",
|
||||
]
|
||||
+61
@@ -0,0 +1,61 @@
|
||||
"""Docling engine adapter (implements ExternalParserBase hooks)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from lightrag.constants import DOCLING_RAW_DIR_SUFFIX, PARSER_ENGINE_DOCLING
|
||||
from lightrag.parser.external._base import ExternalParserBase
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from lightrag.sidecar.ir import IRDoc
|
||||
|
||||
|
||||
class DoclingParser(ExternalParserBase):
|
||||
engine_name = PARSER_ENGINE_DOCLING
|
||||
raw_dir_suffix = DOCLING_RAW_DIR_SUFFIX
|
||||
force_reparse_env = "LIGHTRAG_FORCE_REPARSE_DOCLING"
|
||||
|
||||
def is_bundle_valid(
|
||||
self,
|
||||
raw_dir: Path,
|
||||
source_path: Path,
|
||||
*,
|
||||
engine_params: "Mapping[str, Any] | None" = None,
|
||||
) -> bool:
|
||||
from lightrag.parser.external.docling import is_bundle_valid
|
||||
|
||||
return is_bundle_valid(raw_dir, source_path, overrides=engine_params)
|
||||
|
||||
async def download_into(
|
||||
self,
|
||||
raw_dir: Path,
|
||||
source_path: Path,
|
||||
*,
|
||||
upload_name: str,
|
||||
engine_params: "Mapping[str, Any] | None" = None,
|
||||
) -> None:
|
||||
from lightrag.parser.external.docling import DoclingRawClient
|
||||
|
||||
# Map the canonical ``upload_name`` onto docling-serve's multipart
|
||||
# filename so the bundle's main JSON is named ``<canonical_stem>.json``
|
||||
# (the IR builder locates it via that canonical stem).
|
||||
await DoclingRawClient(overrides=engine_params).download_into(
|
||||
raw_dir, source_path, upload_filename=upload_name
|
||||
)
|
||||
|
||||
def build_ir(self, raw_dir: Path, document_name: str) -> "IRDoc":
|
||||
from lightrag.parser.external.docling import DoclingIRBuilder
|
||||
|
||||
return DoclingIRBuilder().normalize_from_workdir(
|
||||
raw_dir, document_name=document_name
|
||||
)
|
||||
|
||||
def validate_ir(self, ir: "IRDoc", *, file_path: str, raw_dir: Path) -> None:
|
||||
if not ir.blocks:
|
||||
raise ValueError(
|
||||
f"Docling IR builder produced zero blocks for {file_path} "
|
||||
f"(raw_dir={raw_dir})"
|
||||
)
|
||||
+31
@@ -0,0 +1,31 @@
|
||||
"""MinerU parser integration (raw client, cache, manifest, IR builder).
|
||||
|
||||
Public surface for the rest of the codebase. ``parse_mineru`` imports
|
||||
only from this facade so the inner module layout stays free to evolve.
|
||||
|
||||
See ``docs/LightRAGSidecarFormat-zh.md`` for sidecar format and
|
||||
``docs/FileProcessingConfiguration-zh.md`` for cache lifecycle.
|
||||
"""
|
||||
|
||||
from lightrag.parser.external.mineru.cache import (
|
||||
MINERU_RAW_DIR_SUFFIX,
|
||||
clear_dir_contents,
|
||||
compute_size_and_hash,
|
||||
is_bundle_valid,
|
||||
raw_dir_for_parsed_dir,
|
||||
)
|
||||
from lightrag.parser.external.mineru.client import MinerURawClient
|
||||
from lightrag.parser.external.mineru.ir_builder import MinerUIRBuilder
|
||||
from lightrag.parser.external.mineru.manifest import Manifest, ManifestFile
|
||||
|
||||
__all__ = [
|
||||
"MINERU_RAW_DIR_SUFFIX",
|
||||
"Manifest",
|
||||
"ManifestFile",
|
||||
"MinerUIRBuilder",
|
||||
"MinerURawClient",
|
||||
"clear_dir_contents",
|
||||
"compute_size_and_hash",
|
||||
"is_bundle_valid",
|
||||
"raw_dir_for_parsed_dir",
|
||||
]
|
||||
+431
@@ -0,0 +1,431 @@
|
||||
"""Cache validation for ``*.mineru_raw/`` bundles.
|
||||
|
||||
Validation policy (settled in design discussion; see
|
||||
``LightRAGSidecarFormat-zh.md`` related notes):
|
||||
|
||||
1. ``_manifest.json`` exists, parses, ``version=1.0`` ∧ ``engine=mineru``.
|
||||
2. **Source size fast-path**: ``source_file.stat().st_size`` matches manifest;
|
||||
mismatch → miss without hashing.
|
||||
3. **Source content_hash**: full sha256 of the current source file matches
|
||||
manifest. The size+hash pair is computed by a single-read helper so the
|
||||
stored manifest is internally self-consistent.
|
||||
4. **API mode**: if the manifest recorded ``api_mode`` and it differs from
|
||||
current ``MINERU_API_MODE``, miss.
|
||||
5. **Parser options**: the manifest must record an ``options_signature`` that
|
||||
matches the current effective MinerU request options. Missing signatures
|
||||
from older manifests are treated as stale.
|
||||
6. **Engine version**: if ``MINERU_ENGINE_VERSION`` is set and the manifest
|
||||
recorded a non-empty one, they must match.
|
||||
7. **Endpoint signature**: if the active MinerU endpoint is set and the
|
||||
manifest recorded a non-empty one, they must match.
|
||||
8. **Critical file**: ``content_list.json`` must exist with matching size
|
||||
**and** sha256 — sha256 here is the final tie-breaker against silent
|
||||
corruption affecting the file the adapter depends on.
|
||||
9. **Other files**: size-only verification (cheap; covers most corruption
|
||||
modes for image / middle.json / layout.pdf).
|
||||
|
||||
Any failed step ⇒ cache miss; the caller wipes the directory contents
|
||||
(preserving the directory itself) and re-runs the download.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import asdict, dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from lightrag.constants import MINERU_RAW_DIR_SUFFIX, PARSED_DIR_SUFFIX
|
||||
from lightrag.parser.external.mineru.manifest import load_manifest
|
||||
from lightrag.utils import logger
|
||||
|
||||
DEFAULT_MINERU_API_MODE = "local"
|
||||
DEFAULT_MINERU_OFFICIAL_ENDPOINT = "https://mineru.net"
|
||||
DEFAULT_MINERU_MODEL_VERSION = "vlm"
|
||||
DEFAULT_MINERU_LANGUAGE = "ch"
|
||||
DEFAULT_MINERU_LOCAL_BACKEND = "hybrid-auto-engine"
|
||||
DEFAULT_MINERU_LOCAL_PARSE_METHOD = "auto"
|
||||
DEFAULT_MINERU_LOCAL_IMAGE_ANALYSIS = False
|
||||
DEFAULT_MINERU_LOCAL_START_PAGE_ID = 0
|
||||
DEFAULT_MINERU_LOCAL_END_PAGE_ID = 99999
|
||||
DEFAULT_MINERU_ENABLE_TABLE = True
|
||||
DEFAULT_MINERU_ENABLE_FORMULA = True
|
||||
DEFAULT_MINERU_IS_OCR = False
|
||||
|
||||
|
||||
def raw_dir_for_parsed_dir(parsed_dir: Path) -> Path:
|
||||
"""Sibling raw dir for a given ``*.parsed`` dir.
|
||||
|
||||
``foo.parsed/`` → ``foo.mineru_raw/``. Used both at download time and at
|
||||
cache check time so the layout is canonical.
|
||||
"""
|
||||
stem = parsed_dir.name
|
||||
if stem.endswith(PARSED_DIR_SUFFIX):
|
||||
stem = stem[: -len(PARSED_DIR_SUFFIX)]
|
||||
return parsed_dir.parent / f"{stem}{MINERU_RAW_DIR_SUFFIX}"
|
||||
|
||||
|
||||
def clear_dir_contents(directory: Path) -> None:
|
||||
"""Delete everything inside ``directory`` but keep ``directory`` itself."""
|
||||
if not directory.exists():
|
||||
return
|
||||
for entry in directory.iterdir():
|
||||
try:
|
||||
if entry.is_dir() and not entry.is_symlink():
|
||||
_rmtree_safe(entry)
|
||||
else:
|
||||
entry.unlink()
|
||||
except OSError:
|
||||
# Best-effort cleanup; subsequent download will overwrite.
|
||||
continue
|
||||
|
||||
|
||||
def _rmtree_safe(directory: Path) -> None:
|
||||
import shutil
|
||||
|
||||
shutil.rmtree(directory, ignore_errors=True)
|
||||
|
||||
|
||||
def compute_size_and_hash(path: Path) -> tuple[int, str]:
|
||||
"""Single-read computation of ``(size_bytes, "sha256:<hex>")``.
|
||||
|
||||
Manifest writes use this so the recorded size and hash are guaranteed to
|
||||
describe the same byte stream; using two ``open()`` calls would risk a
|
||||
TOCTOU mismatch if the file changed in between.
|
||||
"""
|
||||
h = hashlib.sha256()
|
||||
size = 0
|
||||
with path.open("rb") as f:
|
||||
for chunk in iter(lambda: f.read(1 << 20), b""):
|
||||
h.update(chunk)
|
||||
size += len(chunk)
|
||||
return size, f"sha256:{h.hexdigest()}"
|
||||
|
||||
|
||||
def _current_api_mode() -> str:
|
||||
mode = _normalize_api_mode(os.getenv("MINERU_API_MODE", DEFAULT_MINERU_API_MODE))
|
||||
return mode
|
||||
|
||||
|
||||
def _normalize_api_mode(mode: str) -> str:
|
||||
mode = str(mode or "").strip().lower()
|
||||
return mode if mode in {"official", "local"} else DEFAULT_MINERU_API_MODE
|
||||
|
||||
|
||||
def _env_bool(name: str, default: bool) -> bool:
|
||||
raw = os.getenv(name, "").strip().lower()
|
||||
if raw in {"1", "true", "yes", "on"}:
|
||||
return True
|
||||
if raw in {"0", "false", "no", "off"}:
|
||||
return False
|
||||
return default
|
||||
|
||||
|
||||
def _env_int(name: str, default: int) -> int:
|
||||
raw = os.getenv(name, "").strip()
|
||||
if not raw:
|
||||
return default
|
||||
try:
|
||||
return int(raw)
|
||||
except ValueError:
|
||||
logger.warning(
|
||||
"[mineru_raw] %s=%r is not an integer; using %s", name, raw, default
|
||||
)
|
||||
return default
|
||||
|
||||
|
||||
def _current_endpoint_signature() -> str:
|
||||
mode = _current_api_mode()
|
||||
if mode == "official":
|
||||
return (
|
||||
os.getenv("MINERU_OFFICIAL_ENDPOINT", DEFAULT_MINERU_OFFICIAL_ENDPOINT)
|
||||
.strip()
|
||||
.rstrip("/")
|
||||
)
|
||||
if mode == "local":
|
||||
return os.getenv("MINERU_LOCAL_ENDPOINT", "").strip().rstrip("/")
|
||||
return ""
|
||||
|
||||
|
||||
def local_page_bounds(page_ranges: str) -> tuple[int, int]:
|
||||
raw = page_ranges.strip()
|
||||
if not raw:
|
||||
return DEFAULT_MINERU_LOCAL_START_PAGE_ID, DEFAULT_MINERU_LOCAL_END_PAGE_ID
|
||||
if "," in raw:
|
||||
raise ValueError(
|
||||
"MINERU_PAGE_RANGES with MINERU_API_MODE=local supports only a "
|
||||
"single page or simple range such as '1-10'"
|
||||
)
|
||||
if raw.isdigit():
|
||||
page = max(int(raw), 1)
|
||||
return page - 1, page - 1
|
||||
if "-" in raw:
|
||||
left, _, right = raw.partition("-")
|
||||
if left.isdigit() and right.isdigit():
|
||||
start = max(int(left), 1)
|
||||
end = max(int(right), start)
|
||||
return start - 1, end - 1
|
||||
raise ValueError(
|
||||
"MINERU_PAGE_RANGES with MINERU_API_MODE=local must be a single "
|
||||
"positive page number or simple range such as '1-10'"
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class MinerUParserOptions:
|
||||
"""Effective MinerU parser options used both for live requests and the
|
||||
cache signature.
|
||||
|
||||
Constructed once via :meth:`from_env` so the client and the cache
|
||||
validator agree on every defaulting / normalization rule.
|
||||
"""
|
||||
|
||||
api_mode: str
|
||||
model_version: str
|
||||
language: str
|
||||
enable_table: bool
|
||||
enable_formula: bool
|
||||
is_ocr: bool
|
||||
page_ranges: str
|
||||
local_backend: str
|
||||
local_parse_method: str
|
||||
local_image_analysis: bool
|
||||
local_start_page_id: int
|
||||
local_end_page_id: int
|
||||
|
||||
@classmethod
|
||||
def from_env(
|
||||
cls,
|
||||
*,
|
||||
api_mode: str | None = None,
|
||||
overrides: "Mapping[str, Any] | None" = None,
|
||||
) -> "MinerUParserOptions":
|
||||
"""Build the effective options from env, with optional per-file overrides.
|
||||
|
||||
``overrides`` carries decoded per-file engine params (``page_range`` →
|
||||
``page_ranges``, ``language``, ``local_parse_method``). They feed both
|
||||
the live request and the cache signature, so an overridden document gets
|
||||
its own bundle.
|
||||
"""
|
||||
overrides = overrides or {}
|
||||
mode = (
|
||||
_normalize_api_mode(api_mode)
|
||||
if api_mode is not None
|
||||
else _current_api_mode()
|
||||
)
|
||||
page_ranges = str(
|
||||
overrides.get("page_range", os.getenv("MINERU_PAGE_RANGES", ""))
|
||||
).strip()
|
||||
language = (
|
||||
str(
|
||||
overrides.get(
|
||||
"language", os.getenv("MINERU_LANGUAGE", DEFAULT_MINERU_LANGUAGE)
|
||||
)
|
||||
).strip()
|
||||
or DEFAULT_MINERU_LANGUAGE
|
||||
)
|
||||
local_parse_method = (
|
||||
str(
|
||||
overrides.get(
|
||||
"local_parse_method",
|
||||
os.getenv(
|
||||
"MINERU_LOCAL_PARSE_METHOD", DEFAULT_MINERU_LOCAL_PARSE_METHOD
|
||||
),
|
||||
)
|
||||
).strip()
|
||||
or DEFAULT_MINERU_LOCAL_PARSE_METHOD
|
||||
)
|
||||
local_start = _env_int(
|
||||
"MINERU_LOCAL_START_PAGE_ID", DEFAULT_MINERU_LOCAL_START_PAGE_ID
|
||||
)
|
||||
local_end = _env_int(
|
||||
"MINERU_LOCAL_END_PAGE_ID", DEFAULT_MINERU_LOCAL_END_PAGE_ID
|
||||
)
|
||||
if mode == "local" and page_ranges:
|
||||
local_start, local_end = local_page_bounds(page_ranges)
|
||||
return cls(
|
||||
api_mode=mode,
|
||||
model_version=(
|
||||
os.getenv("MINERU_MODEL_VERSION", DEFAULT_MINERU_MODEL_VERSION).strip()
|
||||
or DEFAULT_MINERU_MODEL_VERSION
|
||||
),
|
||||
language=language,
|
||||
enable_table=_env_bool("MINERU_ENABLE_TABLE", DEFAULT_MINERU_ENABLE_TABLE),
|
||||
enable_formula=_env_bool(
|
||||
"MINERU_ENABLE_FORMULA", DEFAULT_MINERU_ENABLE_FORMULA
|
||||
),
|
||||
is_ocr=_env_bool("MINERU_IS_OCR", DEFAULT_MINERU_IS_OCR),
|
||||
page_ranges=page_ranges,
|
||||
local_backend=(
|
||||
os.getenv("MINERU_LOCAL_BACKEND", DEFAULT_MINERU_LOCAL_BACKEND).strip()
|
||||
or DEFAULT_MINERU_LOCAL_BACKEND
|
||||
),
|
||||
local_parse_method=local_parse_method,
|
||||
local_image_analysis=_env_bool(
|
||||
"MINERU_LOCAL_IMAGE_ANALYSIS", DEFAULT_MINERU_LOCAL_IMAGE_ANALYSIS
|
||||
),
|
||||
local_start_page_id=local_start,
|
||||
local_end_page_id=local_end,
|
||||
)
|
||||
|
||||
def signature(self) -> str:
|
||||
return mineru_options_signature(**asdict(self))
|
||||
|
||||
|
||||
def mineru_options_signature(
|
||||
*,
|
||||
api_mode: str,
|
||||
model_version: str = DEFAULT_MINERU_MODEL_VERSION,
|
||||
language: str = DEFAULT_MINERU_LANGUAGE,
|
||||
enable_table: bool = DEFAULT_MINERU_ENABLE_TABLE,
|
||||
enable_formula: bool = DEFAULT_MINERU_ENABLE_FORMULA,
|
||||
is_ocr: bool = DEFAULT_MINERU_IS_OCR,
|
||||
page_ranges: str = "",
|
||||
local_backend: str = DEFAULT_MINERU_LOCAL_BACKEND,
|
||||
local_parse_method: str = DEFAULT_MINERU_LOCAL_PARSE_METHOD,
|
||||
local_image_analysis: bool = DEFAULT_MINERU_LOCAL_IMAGE_ANALYSIS,
|
||||
local_start_page_id: int = DEFAULT_MINERU_LOCAL_START_PAGE_ID,
|
||||
local_end_page_id: int = DEFAULT_MINERU_LOCAL_END_PAGE_ID,
|
||||
) -> str:
|
||||
mode = _normalize_api_mode(api_mode)
|
||||
payload: dict[str, Any] = {
|
||||
"signature_version": 1,
|
||||
"api_mode": mode,
|
||||
"language": str(language or "").strip() or DEFAULT_MINERU_LANGUAGE,
|
||||
"enable_table": bool(enable_table),
|
||||
"enable_formula": bool(enable_formula),
|
||||
}
|
||||
if mode == "official":
|
||||
payload.update(
|
||||
{
|
||||
"model_version": str(model_version or "").strip()
|
||||
or DEFAULT_MINERU_MODEL_VERSION,
|
||||
"is_ocr": bool(is_ocr),
|
||||
"page_ranges": str(page_ranges or "").strip(),
|
||||
}
|
||||
)
|
||||
else:
|
||||
payload.update(
|
||||
{
|
||||
"local_backend": str(local_backend or "").strip()
|
||||
or DEFAULT_MINERU_LOCAL_BACKEND,
|
||||
"local_parse_method": str(local_parse_method or "").strip()
|
||||
or DEFAULT_MINERU_LOCAL_PARSE_METHOD,
|
||||
"local_image_analysis": bool(local_image_analysis),
|
||||
"local_start_page_id": int(local_start_page_id),
|
||||
"local_end_page_id": int(local_end_page_id),
|
||||
}
|
||||
)
|
||||
|
||||
raw = json.dumps(payload, sort_keys=True, separators=(",", ":"))
|
||||
return "sha256:" + hashlib.sha256(raw.encode("utf-8")).hexdigest()
|
||||
|
||||
|
||||
def current_mineru_options_signature(
|
||||
overrides: "Mapping[str, Any] | None" = None,
|
||||
) -> str:
|
||||
return MinerUParserOptions.from_env(overrides=overrides).signature()
|
||||
|
||||
|
||||
def is_bundle_valid(
|
||||
raw_dir: Path,
|
||||
source_file: Path,
|
||||
*,
|
||||
overrides: "Mapping[str, Any] | None" = None,
|
||||
) -> bool:
|
||||
"""Return True iff the bundle is intact and matches the current source.
|
||||
|
||||
See module docstring for the full policy. Returns False on any of:
|
||||
missing manifest, malformed manifest, schema version mismatch, source
|
||||
size/hash mismatch, parser options mismatch, engine/endpoint env mismatch,
|
||||
critical file missing or corrupted, or any non-critical file size mismatch.
|
||||
"""
|
||||
if not raw_dir.is_dir():
|
||||
return False
|
||||
|
||||
manifest = load_manifest(raw_dir)
|
||||
if manifest is None:
|
||||
return False
|
||||
|
||||
# 1. Source size fast-path
|
||||
try:
|
||||
cur_size = source_file.stat().st_size
|
||||
except OSError:
|
||||
return False
|
||||
if cur_size != int(manifest.source_size_bytes):
|
||||
return False
|
||||
|
||||
# 2. Source content_hash
|
||||
_, cur_hash = compute_size_and_hash(source_file)
|
||||
if cur_hash != manifest.source_content_hash:
|
||||
return False
|
||||
|
||||
# 3. API mode (only when manifest had one; old manifests remain compatible)
|
||||
cur_api_mode = _current_api_mode()
|
||||
if manifest.api_mode and cur_api_mode != manifest.api_mode:
|
||||
return False
|
||||
|
||||
# 4. Parser options. Old manifests did not record this and must miss so
|
||||
# changes such as MINERU_LOCAL_BACKEND cannot silently reuse stale output.
|
||||
if not manifest.options_signature:
|
||||
return False
|
||||
if current_mineru_options_signature(overrides) != manifest.options_signature:
|
||||
return False
|
||||
|
||||
# 5. Engine version (only when current env exposes one AND manifest had one)
|
||||
cur_engine_version = os.getenv("MINERU_ENGINE_VERSION", "").strip()
|
||||
if (
|
||||
cur_engine_version
|
||||
and manifest.engine_version
|
||||
and cur_engine_version != manifest.engine_version
|
||||
):
|
||||
return False
|
||||
|
||||
# 6. Endpoint signature
|
||||
cur_endpoint = _current_endpoint_signature()
|
||||
if (
|
||||
cur_endpoint
|
||||
and manifest.endpoint_signature
|
||||
and cur_endpoint != manifest.endpoint_signature
|
||||
):
|
||||
return False
|
||||
|
||||
# 7. Critical file: size + sha256
|
||||
crit = manifest.critical_file
|
||||
crit_path = raw_dir / crit.path
|
||||
try:
|
||||
if crit_path.stat().st_size != int(crit.size):
|
||||
return False
|
||||
except OSError:
|
||||
return False
|
||||
if crit.sha256:
|
||||
_, crit_actual = compute_size_and_hash(crit_path)
|
||||
if crit_actual != crit.sha256:
|
||||
return False
|
||||
|
||||
# 8. Other files: size only
|
||||
for entry in manifest.files:
|
||||
ep = raw_dir / entry.path
|
||||
try:
|
||||
if ep.stat().st_size != int(entry.size):
|
||||
return False
|
||||
except OSError:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
|
||||
__all__ = [
|
||||
"MINERU_RAW_DIR_SUFFIX",
|
||||
"MinerUParserOptions",
|
||||
"clear_dir_contents",
|
||||
"compute_size_and_hash",
|
||||
"current_mineru_options_signature",
|
||||
"is_bundle_valid",
|
||||
"local_page_bounds",
|
||||
"mineru_options_signature",
|
||||
"raw_dir_for_parsed_dir",
|
||||
]
|
||||
+702
@@ -0,0 +1,702 @@
|
||||
"""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"]
|
||||
+785
@@ -0,0 +1,785 @@
|
||||
"""MinerU IR builder: ``content_list.json`` (+ images/) → :class:`IRDoc`.
|
||||
|
||||
Input contract: a ``*.mineru_raw/`` directory containing at least
|
||||
``content_list.json``. Optional sibling resources (``images/``,
|
||||
``middle.json``, ``full.md``, ``layout.pdf``) are kept as-is; this builder
|
||||
only reads the content list and image asset bytes.
|
||||
|
||||
Conversion rules (informed by spec §3-§六):
|
||||
|
||||
- ``text`` items with ``text_level>0`` and ``title`` / ``section_header``
|
||||
start a NEW block. The heading text is rendered with a markdown ``#``
|
||||
prefix matching the level (``# foo``, ``## bar`` …) as the first line of
|
||||
the new block's content.
|
||||
- All other items (``text``, ``list``, ``code``, ``table``, ``image``,
|
||||
``equation``) are MERGED into the current block — their text / placeholder
|
||||
is appended (newline-separated) to the heading's block. This mirrors the
|
||||
native docx parser's "split-by-heading, merge-everything-under-heading"
|
||||
behavior (see ``parser/docx/parse_document.py``).
|
||||
- Content emitted before the first heading lands in a synthetic
|
||||
``Preface/Uncategorized`` block at level 0.
|
||||
- ``list`` items joined with ``\n``; ``code`` body taken from ``code_body``
|
||||
if present.
|
||||
- ``table`` → IRTable + ``{{TBL:k}}`` placeholder. MinerU HTML tables are
|
||||
preserved verbatim on ``IRTable.html`` so merged cells (``rowspan`` /
|
||||
``colspan``) survive in ``tables.json``; the block placeholder receives
|
||||
only the table's inner HTML to avoid nested ``<table>`` wrappers. ``rows``
|
||||
is reserved for explicit 2D-array / non-HTML compatibility inputs. A real
|
||||
HTML ``<thead>`` populates ``table_header`` (per spec §5); otherwise the
|
||||
adapter does not guess a header row.
|
||||
- ``image`` / ``picture`` / ``drawing`` → IRDrawing + ``{{IMG:k}}`` placeholder.
|
||||
Asset bytes are referenced via ``img_path`` relative to the raw dir.
|
||||
- ``equation`` → IREquation. ``is_block`` is decided by whether
|
||||
``text_format=="block"`` (MinerU explicit flag) OR ``text_level==0`` with
|
||||
no inline neighbours; otherwise inline. The latex string is preserved
|
||||
verbatim (including any ``$$``/``$`` wrappers) so ``blocks.jsonl``'s
|
||||
``<equation>`` body matches MinerU's raw output; the writer strips the
|
||||
wrappers when persisting ``equations.json`` content.
|
||||
- ``page_idx`` + ``bbox`` → ``IRPosition(type="bbox", anchor=page, range=[x0,y0,x1,y1])``.
|
||||
Empty/missing bbox is acceptable; positions accumulate on the merged block.
|
||||
- ``IRDoc.split_option`` records the MinerU engine version when available.
|
||||
- ``IRDoc.bbox_attributes`` defaults to ``{"origin":"LEFTTOP","max":1000}``
|
||||
reflecting MinerU's PDF coordinate convention. Operators may override
|
||||
via ``MINERU_BBOX_ATTRIBUTES`` (JSON string).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from urllib.parse import urlparse
|
||||
|
||||
from lightrag.parser._html_table import (
|
||||
HTMLTableInfo,
|
||||
extract_html_table_info,
|
||||
extract_thead_html,
|
||||
html_table_inner_body,
|
||||
looks_like_html_table_payload,
|
||||
unwrap_html_table,
|
||||
)
|
||||
from lightrag.parser._markdown import (
|
||||
render_heading_line,
|
||||
strip_heading_markdown_prefix,
|
||||
)
|
||||
from lightrag.sidecar.ir import (
|
||||
AssetSpec,
|
||||
IRBlock,
|
||||
IRDoc,
|
||||
IRDrawing,
|
||||
IREquation,
|
||||
IRPosition,
|
||||
IRTable,
|
||||
)
|
||||
from lightrag.utils import logger
|
||||
|
||||
|
||||
PREFACE_HEADING = "Preface/Uncategorized"
|
||||
CONTENT_LIST_FILENAME = "content_list.json"
|
||||
|
||||
|
||||
class MinerUIRBuilder:
|
||||
"""Stateless except for env-driven config. Reusable across calls."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.engine_version = os.getenv("MINERU_ENGINE_VERSION", "").strip()
|
||||
# Mirror MinerURawClient.__init__: when this is set, the downloader
|
||||
# stores ALL referenced images (including relative ones) under
|
||||
# ``images/<basename>``. The builder has to look in the same place.
|
||||
self.image_url_template = os.getenv("MINERU_IMAGE_URL_TEMPLATE", "").strip()
|
||||
self.bbox_attributes = self._load_bbox_attributes_env()
|
||||
|
||||
def _load_bbox_attributes_env(self) -> dict[str, Any]:
|
||||
default = {"origin": "LEFTTOP", "max": 1000}
|
||||
raw = os.getenv("MINERU_BBOX_ATTRIBUTES", "").strip()
|
||||
if not raw:
|
||||
return default
|
||||
try:
|
||||
parsed = json.loads(raw)
|
||||
except json.JSONDecodeError as exc:
|
||||
logger.warning(
|
||||
"[mineru_ir_builder] MINERU_BBOX_ATTRIBUTES is not valid JSON "
|
||||
"(%s); falling back to default %s",
|
||||
exc,
|
||||
default,
|
||||
)
|
||||
return default
|
||||
if not isinstance(parsed, dict):
|
||||
logger.warning(
|
||||
"[mineru_ir_builder] MINERU_BBOX_ATTRIBUTES must decode to a JSON "
|
||||
"object, got %s; falling back to default %s",
|
||||
type(parsed).__name__,
|
||||
default,
|
||||
)
|
||||
return default
|
||||
return parsed
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Entry point
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def normalize_from_workdir(
|
||||
self,
|
||||
raw_dir: Path,
|
||||
*,
|
||||
document_name: str,
|
||||
) -> IRDoc:
|
||||
"""Read ``raw_dir/content_list.json`` and emit an IRDoc.
|
||||
|
||||
``document_name`` is the canonical filename (e.g. ``foo.pdf``) used
|
||||
for ``meta.document_name``; resolved by the caller from the parser
|
||||
hint chain.
|
||||
"""
|
||||
content_list_path = raw_dir / "content_list.json"
|
||||
if not content_list_path.is_file():
|
||||
raise FileNotFoundError(
|
||||
f"MinerU raw bundle missing content_list.json at {raw_dir}"
|
||||
)
|
||||
content_list = json.loads(content_list_path.read_text(encoding="utf-8"))
|
||||
if not isinstance(content_list, list):
|
||||
raise ValueError(
|
||||
f"MinerU content_list.json malformed (not a JSON array) at {raw_dir}"
|
||||
)
|
||||
return self._normalize_content_list(
|
||||
content_list, raw_dir, document_name=document_name
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Core
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _normalize_content_list(
|
||||
self,
|
||||
content_list: list[Any],
|
||||
raw_dir: Path,
|
||||
*,
|
||||
document_name: str,
|
||||
) -> IRDoc:
|
||||
document_format = Path(document_name).suffix.lower().lstrip(".")
|
||||
|
||||
blocks: list[IRBlock] = []
|
||||
assets: list[AssetSpec] = []
|
||||
seen_assets: dict[str, str] = {} # ref → suggested_name
|
||||
doc_title = ""
|
||||
placeholder_counter = 0
|
||||
|
||||
def _next_key(prefix: str) -> str:
|
||||
nonlocal placeholder_counter
|
||||
placeholder_counter += 1
|
||||
return f"{prefix}{placeholder_counter}"
|
||||
|
||||
# Heading hierarchy stack — index = level-1 (level 1 lives at [0]).
|
||||
heading_stack: list[str] = []
|
||||
|
||||
# Current-block accumulator. The block is materialized when the next
|
||||
# heading arrives (or at end-of-document). The initial block is the
|
||||
# synthetic "Preface/Uncategorized" container at level 0.
|
||||
cb_lines: list[str] = []
|
||||
cb_tables: list[IRTable] = []
|
||||
cb_drawings: list[IRDrawing] = []
|
||||
cb_equations: list[IREquation] = []
|
||||
# Positions are split into two channels:
|
||||
# - ``cb_page_set`` collects ``page_idx`` of bbox-less items; at flush
|
||||
# each unique page becomes one anchor-only summary ``IRPosition``.
|
||||
# - ``cb_bbox_positions`` keeps one fine-grained position per item that
|
||||
# carried a parseable bbox (anchor + range), in source order, with
|
||||
# no deduplication.
|
||||
cb_page_set: set[str] = set()
|
||||
cb_bbox_positions: list[IRPosition] = []
|
||||
cb_heading = PREFACE_HEADING
|
||||
cb_level = 0
|
||||
cb_parents: list[str] = []
|
||||
|
||||
def _record_position(item: dict) -> None:
|
||||
"""Route an item's positional info into the right channel.
|
||||
|
||||
Items with a parseable ``bbox`` produce one fine-grained
|
||||
IRPosition appended to ``cb_bbox_positions`` (no dedupe).
|
||||
Otherwise, ``page_idx`` (if any) is added to ``cb_page_set``
|
||||
and emitted as a single anchor-only summary entry at flush.
|
||||
"""
|
||||
bbox_pos = _extract_bbox_position(item)
|
||||
if bbox_pos is not None:
|
||||
cb_bbox_positions.append(bbox_pos)
|
||||
return
|
||||
page = _extract_page_anchor(item)
|
||||
if page is not None:
|
||||
cb_page_set.add(page)
|
||||
|
||||
def _flush_block() -> None:
|
||||
"""Emit the in-flight block if it carries any content."""
|
||||
nonlocal cb_lines, cb_tables, cb_drawings, cb_equations
|
||||
nonlocal cb_page_set, cb_bbox_positions
|
||||
has_payload = bool(cb_lines or cb_tables or cb_drawings or cb_equations)
|
||||
if not has_payload:
|
||||
return
|
||||
content = "\n".join(line for line in cb_lines if line)
|
||||
if not content.strip() and not (cb_tables or cb_drawings or cb_equations):
|
||||
# Reset and skip — nothing meaningful to emit.
|
||||
cb_lines = []
|
||||
cb_page_set = set()
|
||||
cb_bbox_positions = []
|
||||
return
|
||||
positions = [
|
||||
IRPosition(type="bbox", anchor=p)
|
||||
for p in _sort_page_anchors(cb_page_set)
|
||||
] + list(cb_bbox_positions)
|
||||
blocks.append(
|
||||
IRBlock(
|
||||
content_template=content,
|
||||
heading=cb_heading,
|
||||
level=cb_level,
|
||||
parent_headings=list(cb_parents),
|
||||
positions=positions,
|
||||
tables=list(cb_tables),
|
||||
drawings=list(cb_drawings),
|
||||
equations=list(cb_equations),
|
||||
)
|
||||
)
|
||||
cb_lines = []
|
||||
cb_tables = []
|
||||
cb_drawings = []
|
||||
cb_equations = []
|
||||
cb_page_set = set()
|
||||
cb_bbox_positions = []
|
||||
|
||||
def _open_block(
|
||||
heading: str, level: int, parents: list[str], raw_heading: str | None = None
|
||||
) -> None:
|
||||
nonlocal cb_heading, cb_level, cb_parents
|
||||
cb_heading = heading
|
||||
cb_level = level
|
||||
cb_parents = parents
|
||||
# Render the heading line into the block body so the merged
|
||||
# text reads like markdown (``# Foo`` / ``## Bar`` / …). Levels
|
||||
# are capped at 6 ``#`` and headings already carrying a markdown
|
||||
# prefix are left untouched (see ``render_heading_line``).
|
||||
cb_lines.append(render_heading_line(level, raw_heading or heading))
|
||||
|
||||
def _append_text(text: str) -> bool:
|
||||
"""Append ``text`` to the current block body and return whether
|
||||
anything was actually written. Callers use the return value to
|
||||
decide whether to also record the item's source position — an
|
||||
empty text item must NOT leak its ``page_idx`` to the block.
|
||||
"""
|
||||
if not text:
|
||||
return False
|
||||
cb_lines.append(text)
|
||||
return True
|
||||
|
||||
for item_index, item in enumerate(content_list):
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
item_type = str(item.get("type") or item.get("label") or "").lower()
|
||||
|
||||
# Page numbers are layout noise, not document body. MinerU emits a
|
||||
# ``page_number`` item per page; skip it entirely so it never enters
|
||||
# the block content nor leaks its page_idx into block positions.
|
||||
# (Empty-text page numbers were already dropped by the fallback's
|
||||
# _append_text guard; this also drops page numbers that carry real
|
||||
# text like "12" / "iii".)
|
||||
if item_type == "page_number":
|
||||
continue
|
||||
|
||||
heading_text, heading_level = _detect_heading(item, item_type)
|
||||
if heading_text:
|
||||
clean_heading = strip_heading_markdown_prefix(heading_text)
|
||||
# Heading hierarchy is updated unconditionally so deeper
|
||||
# parents resolve correctly once the next real body item
|
||||
# opens a fresh block.
|
||||
heading_stack = heading_stack[: max(heading_level - 1, 0)]
|
||||
parents = [h for h in heading_stack if h]
|
||||
heading_stack.append(clean_heading)
|
||||
|
||||
# Every recognized heading starts its own block: flush the
|
||||
# in-flight block (whether it had body or was a bare heading)
|
||||
# and open a fresh one. A heading with no following body thus
|
||||
# becomes a standalone block whose content is just the heading
|
||||
# line, matching the native docx parser's behaviour.
|
||||
_flush_block()
|
||||
_open_block(clean_heading, heading_level, parents, heading_text)
|
||||
_record_position(item)
|
||||
|
||||
if not doc_title and heading_level == 1:
|
||||
doc_title = clean_heading
|
||||
continue
|
||||
|
||||
if item_type == "text":
|
||||
if _append_text(_coerce_text(item)):
|
||||
_record_position(item)
|
||||
continue
|
||||
|
||||
if item_type == "list":
|
||||
items = item.get("list_items")
|
||||
if isinstance(items, list):
|
||||
text = "\n".join(str(x) for x in items if str(x).strip())
|
||||
else:
|
||||
text = _coerce_text(item)
|
||||
if _append_text(text):
|
||||
_record_position(item)
|
||||
continue
|
||||
|
||||
if item_type == "code":
|
||||
if _append_text(item.get("code_body") or _coerce_text(item)):
|
||||
_record_position(item)
|
||||
continue
|
||||
|
||||
if item_type == "equation":
|
||||
latex_raw = _coerce_text(item)
|
||||
if not latex_raw:
|
||||
# Spec compliance fix: empty equation must not enter sidecar.
|
||||
continue
|
||||
# Preserve MinerU's raw latex (including any ``$$``/``$``
|
||||
# wrappers); the writer strips them when emitting
|
||||
# equations.json so blocks.jsonl shows the raw form while
|
||||
# the per-equation sidecar holds clean latex.
|
||||
latex = latex_raw.strip()
|
||||
is_block = _is_block_equation(item)
|
||||
caption = str(item.get("caption") or "")
|
||||
placeholder = _next_key("eq")
|
||||
token = "EQ" if is_block else "EQI"
|
||||
cb_equations.append(
|
||||
IREquation(
|
||||
placeholder_key=placeholder,
|
||||
latex=latex,
|
||||
is_block=is_block,
|
||||
caption=caption,
|
||||
footnotes=_as_str_list(item.get("footnotes")),
|
||||
self_ref=_content_list_self_ref(item_index) if is_block else "",
|
||||
)
|
||||
)
|
||||
cb_lines.append(f"{{{{{token}:{placeholder}}}}}")
|
||||
_record_position(item)
|
||||
continue
|
||||
|
||||
if item_type == "table":
|
||||
table = self._build_ir_table(item)
|
||||
if table is None:
|
||||
# Empty body — _build_ir_table already logged the drop.
|
||||
# Skip placeholder allocation and position recording so
|
||||
# the misidentified item leaves no trace in the IR.
|
||||
continue
|
||||
placeholder = _next_key("tb")
|
||||
table.placeholder_key = placeholder
|
||||
table.self_ref = _content_list_self_ref(item_index)
|
||||
cb_tables.append(table)
|
||||
cb_lines.append(f"{{{{TBL:{placeholder}}}}}")
|
||||
_record_position(item)
|
||||
continue
|
||||
|
||||
if item_type in {"image", "picture", "drawing"}:
|
||||
drawing, asset = self._build_ir_drawing(item, raw_dir, seen_assets)
|
||||
placeholder = _next_key("im")
|
||||
drawing.placeholder_key = placeholder
|
||||
drawing.self_ref = _content_list_self_ref(item_index)
|
||||
if asset is not None and asset.ref not in {a.ref for a in assets}:
|
||||
assets.append(asset)
|
||||
cb_drawings.append(drawing)
|
||||
cb_lines.append(f"{{{{IMG:{placeholder}}}}}")
|
||||
_record_position(item)
|
||||
continue
|
||||
|
||||
# Fallback: serialize unknown items as plain text so we don't
|
||||
# silently drop information. Position only recorded when the
|
||||
# fallback actually contributed text — empty unknown items must
|
||||
# not leak their page_idx into the current block.
|
||||
if _append_text(_coerce_text(item)):
|
||||
_record_position(item)
|
||||
|
||||
_flush_block()
|
||||
|
||||
if not doc_title:
|
||||
doc_title = Path(document_name).stem or document_name
|
||||
|
||||
split_option: dict[str, Any] = {}
|
||||
if self.engine_version:
|
||||
split_option["engine_version"] = self.engine_version
|
||||
# Reserved hook for later: detect OCR flag from middle.json / config.
|
||||
|
||||
return IRDoc(
|
||||
document_name=document_name,
|
||||
document_format=document_format,
|
||||
doc_title=doc_title,
|
||||
split_option=split_option,
|
||||
blocks=blocks,
|
||||
assets=assets,
|
||||
bbox_attributes=dict(self.bbox_attributes),
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Tables / drawings
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _build_ir_table(self, item: dict) -> IRTable | None:
|
||||
rows: list[list[str]] | None = None
|
||||
html: str | None = None
|
||||
body_override: str | None = None
|
||||
body_field = item.get("rows")
|
||||
body = body_field if body_field is not None else item.get("table_body")
|
||||
|
||||
if isinstance(body, list):
|
||||
rows = _normalize_grid(body)
|
||||
elif isinstance(body, str):
|
||||
stripped = body.strip()
|
||||
if looks_like_html_table_payload(stripped):
|
||||
# MinerU's table model sometimes wraps output in a
|
||||
# ``<html><body>…</body></html>`` document; unwrap to the bare
|
||||
# ``<table>…</table>`` so the sidecar ``content`` stays a single
|
||||
# clean table and the writer does not nest ``<table>`` wrappers.
|
||||
html = unwrap_html_table(stripped) or None
|
||||
if html:
|
||||
# ``or None`` so a degenerate ``<table></table>`` (empty
|
||||
# inner body) falls back to rendering ``table.html`` in the
|
||||
# writer instead of emitting an empty ``body_override``.
|
||||
body_override = html_table_inner_body(html) or None
|
||||
elif stripped.startswith("[") and stripped.endswith("]"):
|
||||
try:
|
||||
decoded = json.loads(stripped)
|
||||
if isinstance(decoded, list):
|
||||
rows = _normalize_grid(decoded)
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
if rows is None and html is None:
|
||||
# Non-HTML, non-JSON string (or JSON that failed to parse):
|
||||
# fall back to the raw payload as the html body.
|
||||
html = stripped or None
|
||||
elif isinstance(body, dict):
|
||||
grid = body.get("grid") or body.get("rows")
|
||||
if isinstance(grid, list):
|
||||
rows = _normalize_grid(grid)
|
||||
else:
|
||||
html = json.dumps(body, ensure_ascii=False)
|
||||
|
||||
# MinerU occasionally emits table items with no usable body (e.g. when
|
||||
# a page number or blank region is misidentified as a table). Dropping
|
||||
# them here keeps the sidecar free of items that would later trip the
|
||||
# analyze worker's "missing table content" hard-failure path.
|
||||
if not _ir_table_body_has_content(rows, html):
|
||||
logger.debug(
|
||||
"[mineru_ir_builder] dropping empty table item "
|
||||
"(body type=%s, num_rows=%s, num_cols=%s)",
|
||||
type(body).__name__,
|
||||
item.get("num_rows"),
|
||||
item.get("num_cols"),
|
||||
)
|
||||
return None
|
||||
|
||||
num_rows = int(item.get("num_rows") or (len(rows) if rows else 0) or 0)
|
||||
num_cols_default = max((len(r) for r in rows), default=0) if rows else 0
|
||||
num_cols = int(item.get("num_cols") or num_cols_default or 0)
|
||||
html_table_info: HTMLTableInfo | None = None
|
||||
if html and (num_rows <= 0 or num_cols <= 0):
|
||||
html_table_info = extract_html_table_info(html)
|
||||
if num_rows <= 0:
|
||||
num_rows = html_table_info.num_rows
|
||||
if num_cols <= 0:
|
||||
num_cols = html_table_info.num_cols
|
||||
|
||||
captions = item.get("table_caption")
|
||||
caption = str(item.get("caption") or "")
|
||||
if not caption and isinstance(captions, list) and captions:
|
||||
caption = str(captions[0])
|
||||
|
||||
# The header representation follows the table's format so merged-cell
|
||||
# semantics survive: HTML tables keep the raw ``<thead>…</thead>``
|
||||
# (preserving rowspan/colspan); grid/JSON tables keep a 2-D grid.
|
||||
table_header_raw = item.get("header")
|
||||
table_header: list[list[str]] | str | None = None
|
||||
if html:
|
||||
table_header = extract_thead_html(html)
|
||||
# Fallback: an HTML table whose markup carries no ``<thead>`` but for
|
||||
# which MinerU supplied a separate ``header`` grid keeps that grid —
|
||||
# the writer renders it to a (span-less) ``<thead>`` rather than
|
||||
# silently dropping the recovered header.
|
||||
if (
|
||||
table_header is None
|
||||
and isinstance(table_header_raw, list)
|
||||
and table_header_raw
|
||||
):
|
||||
table_header = _normalize_grid(table_header_raw)
|
||||
elif isinstance(table_header_raw, list) and table_header_raw:
|
||||
table_header = _normalize_grid(table_header_raw)
|
||||
|
||||
return IRTable(
|
||||
placeholder_key="", # filled by caller
|
||||
rows=rows,
|
||||
html=html,
|
||||
num_rows=num_rows,
|
||||
num_cols=num_cols,
|
||||
caption=caption,
|
||||
footnotes=_as_str_list(item.get("table_footnote") or item.get("footnotes")),
|
||||
table_header=table_header,
|
||||
body_override=body_override,
|
||||
)
|
||||
|
||||
def _build_ir_drawing(
|
||||
self,
|
||||
item: dict,
|
||||
raw_dir: Path,
|
||||
seen: dict[str, str],
|
||||
) -> tuple[IRDrawing, AssetSpec | None]:
|
||||
img_path = str(item.get("img_path") or item.get("path") or "")
|
||||
src_val = str(item.get("src") or "")
|
||||
captions = item.get("image_caption") or item.get("captions")
|
||||
caption = str(item.get("caption") or "")
|
||||
if not caption and isinstance(captions, list) and captions:
|
||||
caption = str(captions[0])
|
||||
|
||||
fmt = Path(img_path).suffix.lower().lstrip(".") if img_path else ""
|
||||
if not fmt:
|
||||
fmt = str(item.get("format") or "")
|
||||
|
||||
asset: AssetSpec | None = None
|
||||
ref = ""
|
||||
if img_path:
|
||||
ref = img_path
|
||||
if ref in seen:
|
||||
# Already declared by a previous block; reuse name.
|
||||
pass
|
||||
else:
|
||||
# Asset source: file on disk inside raw_dir. ``img_path`` is
|
||||
# untrusted (it comes from MinerU's content_list.json or a
|
||||
# downloaded zip), so we go through a safe resolver that
|
||||
# refuses to escape ``raw_dir`` and mirrors the downloader's
|
||||
# storage layout for absolute-URL / templated references.
|
||||
local_path = _safe_local_asset_path(
|
||||
raw_dir,
|
||||
img_path,
|
||||
image_url_template=self.image_url_template,
|
||||
)
|
||||
suggested_name = _suggested_asset_name(img_path, fmt, len(seen))
|
||||
asset = AssetSpec(
|
||||
ref=ref,
|
||||
suggested_name=suggested_name,
|
||||
source=local_path
|
||||
if local_path is not None and local_path.is_file()
|
||||
else None,
|
||||
)
|
||||
seen[ref] = suggested_name
|
||||
|
||||
drawing = IRDrawing(
|
||||
placeholder_key="", # filled by caller
|
||||
asset_ref=ref,
|
||||
fmt=fmt,
|
||||
caption=caption,
|
||||
footnotes=_as_str_list(item.get("image_footnote") or item.get("footnotes")),
|
||||
src=src_val,
|
||||
)
|
||||
return drawing, asset
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# helpers
|
||||
# ----------------------------------------------------------------------
|
||||
|
||||
|
||||
def _detect_heading(item: dict, item_type: str) -> tuple[str, int]:
|
||||
"""Return ``(heading_text, level)`` if ``item`` is a heading, else ``("", 0)``.
|
||||
|
||||
A heading is either an explicit ``title``/``section_header`` block, or a
|
||||
``text`` block whose ``text_level`` is positive (MinerU's convention).
|
||||
"""
|
||||
if item_type in {"title", "section_header"}:
|
||||
text = _coerce_text(item).strip()
|
||||
level = max(int(item.get("text_level") or item.get("level") or 1), 1)
|
||||
return text, level
|
||||
if item_type == "text":
|
||||
try:
|
||||
tl = int(item.get("text_level") or 0)
|
||||
except (TypeError, ValueError):
|
||||
tl = 0
|
||||
if tl > 0:
|
||||
return _coerce_text(item).strip(), tl
|
||||
return "", 0
|
||||
|
||||
|
||||
def _coerce_text(item: dict) -> str:
|
||||
for key in ("text", "content", "body", "code_body"):
|
||||
val = item.get(key)
|
||||
if isinstance(val, str) and val.strip():
|
||||
return val
|
||||
return ""
|
||||
|
||||
|
||||
def _as_str_list(value: Any) -> list[str]:
|
||||
if value is None:
|
||||
return []
|
||||
if isinstance(value, list):
|
||||
return [str(x) for x in value if str(x).strip()]
|
||||
s = str(value).strip()
|
||||
return [s] if s else []
|
||||
|
||||
|
||||
def _content_list_self_ref(index: int) -> str:
|
||||
return f"{CONTENT_LIST_FILENAME}#/{index}"
|
||||
|
||||
|
||||
def _normalize_grid(grid: Any) -> list[list[str]]:
|
||||
out: list[list[str]] = []
|
||||
if not isinstance(grid, list):
|
||||
return out
|
||||
for row in grid:
|
||||
if not isinstance(row, list):
|
||||
continue
|
||||
out_row: list[str] = []
|
||||
for cell in row:
|
||||
if isinstance(cell, dict):
|
||||
out_row.append(str(cell.get("text", "")).strip())
|
||||
else:
|
||||
out_row.append(str(cell).strip())
|
||||
out.append(out_row)
|
||||
return out
|
||||
|
||||
|
||||
def _ir_table_body_has_content(rows: list[list[str]] | None, html: str | None) -> bool:
|
||||
"""True iff the parsed table body carries any visible cell text or HTML."""
|
||||
if html and html.strip():
|
||||
return True
|
||||
if rows:
|
||||
for row in rows:
|
||||
for cell in row:
|
||||
if isinstance(cell, str) and cell.strip():
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _is_block_equation(item: dict) -> bool:
|
||||
"""Heuristic: MinerU's ``text_format`` distinguishes block vs inline.
|
||||
|
||||
Fallback when absent: treat as block (most MinerU equation items in
|
||||
PDF context represent display equations); inline equations are usually
|
||||
embedded inside ``text`` items rather than first-class ``equation``
|
||||
items.
|
||||
"""
|
||||
fmt = str(item.get("text_format") or "").lower()
|
||||
if fmt in {"inline", "inline_equation"}:
|
||||
return False
|
||||
if fmt in {"block", "block_equation", "display"}:
|
||||
return True
|
||||
return True
|
||||
|
||||
|
||||
def _extract_page_anchor(item: dict) -> str | None:
|
||||
"""Return a 1-based page anchor from MinerU's ``page_idx`` / ``page``.
|
||||
|
||||
Always returns a string so ``blocks.jsonl`` carries a uniform anchor
|
||||
type across Roman / letter / numeric page labels. Integers are bumped
|
||||
to 1-based (``page_idx=0`` → ``"1"``); strings are stripped and passed
|
||||
through verbatim. Returns ``None`` when no usable page info is present.
|
||||
"""
|
||||
page_raw = item.get("page_idx")
|
||||
if page_raw is None:
|
||||
page_raw = item.get("page")
|
||||
if isinstance(page_raw, bool):
|
||||
# bool is a subclass of int — guard so True/False don't sneak in.
|
||||
return None
|
||||
if isinstance(page_raw, int):
|
||||
return str(page_raw + 1 if page_raw >= 0 else page_raw)
|
||||
if isinstance(page_raw, str) and page_raw.strip():
|
||||
return page_raw.strip()
|
||||
return None
|
||||
|
||||
|
||||
def _sort_page_anchors(pages: set[str]) -> list[str]:
|
||||
"""Order page anchors using book pagination convention.
|
||||
|
||||
Non-numeric labels (Roman preface pages ``i``/``ii``/``iv``…, letter
|
||||
pages like ``A``, ``B-1``) come first in lexical order; numeric labels
|
||||
follow, sorted by their integer value so ``"2"`` precedes ``"10"``.
|
||||
Mixing both kinds is safe — the bucketed key avoids the ``TypeError``
|
||||
that ``sorted({"ii", "1"})`` raises when ints and strings mix.
|
||||
"""
|
||||
non_numeric = sorted(p for p in pages if not p.isdigit())
|
||||
numeric = sorted((p for p in pages if p.isdigit()), key=int)
|
||||
return non_numeric + numeric
|
||||
|
||||
|
||||
def _extract_bbox_position(item: dict) -> IRPosition | None:
|
||||
"""Build a fine-grained ``IRPosition`` when ``bbox`` is parseable.
|
||||
|
||||
Returns ``None`` when ``bbox`` is missing or malformed; the caller then
|
||||
falls back to page-only tracking via :func:`_extract_page_anchor`.
|
||||
"""
|
||||
bbox = item.get("bbox")
|
||||
if not isinstance(bbox, (list, tuple)) or len(bbox) < 4:
|
||||
return None
|
||||
try:
|
||||
coords = [float(x) for x in bbox[:4]]
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
return IRPosition(type="bbox", anchor=_extract_page_anchor(item), range=coords)
|
||||
|
||||
|
||||
def _safe_local_asset_path(
|
||||
raw_dir: Path,
|
||||
img_path: str,
|
||||
*,
|
||||
image_url_template: str = "",
|
||||
) -> Path | None:
|
||||
"""Resolve ``img_path`` to a concrete file location inside ``raw_dir``.
|
||||
|
||||
``img_path`` comes from MinerU's ``content_list.json`` and is therefore
|
||||
untrusted. This resolver mirrors :meth:`MinerURawClient._fetch_one_image`
|
||||
storage rules so the builder always looks where the downloader wrote
|
||||
the file:
|
||||
|
||||
- absolute http(s) URLs and absolute filesystem paths
|
||||
→ ``raw_dir/images/<basename>``;
|
||||
- any ref when ``MINERU_IMAGE_URL_TEMPLATE`` is configured (the
|
||||
downloader routes ALL refs — including relative ones — through
|
||||
:meth:`_image_dest_rel`) → ``raw_dir/images/<basename>``;
|
||||
- otherwise relative paths resolve under ``raw_dir`` with ``..``
|
||||
traversal refused and a final ``Path.relative_to`` check.
|
||||
|
||||
Returns ``None`` when the candidate is unsafe or cannot be expressed
|
||||
inside ``raw_dir``. The caller treats ``None`` the same as "file missing"
|
||||
— the drawing tag still gets written, but no bytes are copied.
|
||||
"""
|
||||
if not img_path:
|
||||
return None
|
||||
|
||||
if img_path.startswith(("http://", "https://")):
|
||||
name = Path(urlparse(img_path).path).name
|
||||
return raw_dir / "images" / name if name else None
|
||||
|
||||
if os.path.isabs(img_path):
|
||||
# Absolute filesystem path in img_path is never trusted to point
|
||||
# outside raw_dir; mirror the downloader's basename rule.
|
||||
name = Path(img_path).name
|
||||
return raw_dir / "images" / name if name else None
|
||||
|
||||
if image_url_template:
|
||||
# Templated mode: downloader stored every ref (incl. relative) at
|
||||
# images/<basename>, so we must look there too.
|
||||
name = Path(img_path).name
|
||||
return raw_dir / "images" / name if name else None
|
||||
|
||||
normalized = os.path.normpath(img_path)
|
||||
if normalized.startswith("..") or os.path.isabs(normalized):
|
||||
return None
|
||||
candidate = (raw_dir / normalized).resolve()
|
||||
try:
|
||||
candidate.relative_to(raw_dir.resolve())
|
||||
except ValueError:
|
||||
return None
|
||||
return candidate
|
||||
|
||||
|
||||
def _suggested_asset_name(img_path: str, fmt: str, seen_count: int) -> str:
|
||||
"""Pick an in-assets-dir filename for an asset.
|
||||
|
||||
For URL refs, use the URL path's basename so we get a useful filename
|
||||
(``foo.png`` rather than the whole URL). For local refs, the regular
|
||||
basename. Falls back to ``image-<n>[.fmt]`` when nothing usable.
|
||||
"""
|
||||
if img_path.startswith(("http://", "https://")):
|
||||
name = Path(urlparse(img_path).path).name
|
||||
else:
|
||||
name = Path(img_path).name
|
||||
if name:
|
||||
return name
|
||||
return f"image-{seen_count + 1}{('.' + fmt) if fmt else ''}"
|
||||
|
||||
|
||||
__all__ = ["MinerUIRBuilder"]
|
||||
+164
@@ -0,0 +1,164 @@
|
||||
"""``_manifest.json`` schema for ``*.mineru_raw/`` bundles.
|
||||
|
||||
The manifest is the *atomic success marker* for a raw bundle. Its presence
|
||||
implies "all files in this directory finished downloading"; its content is
|
||||
the cache key for "is this bundle for the same source file, the same MinerU
|
||||
parser options, engine version, and endpoint we are using right now?".
|
||||
|
||||
Write path: ``write_manifest(path, manifest)`` writes a temp file then
|
||||
atomically renames to ``_manifest.json``. A crash mid-download leaves no
|
||||
manifest, so the next ``parse_mineru`` call cleanly invalidates and
|
||||
re-downloads.
|
||||
|
||||
Read path: ``load_manifest(path)`` returns ``None`` if absent or malformed
|
||||
— either way the bundle is treated as stale.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
from dataclasses import asdict, dataclass
|
||||
from pathlib import Path
|
||||
|
||||
MANIFEST_FILENAME = "_manifest.json"
|
||||
MANIFEST_VERSION = "1.0"
|
||||
MANIFEST_ENGINE = "mineru"
|
||||
|
||||
|
||||
@dataclass
|
||||
class ManifestFile:
|
||||
"""One file entry inside the bundle. Size always; sha256 only for the
|
||||
critical file (content_list.json) — see :class:`Manifest.critical_file`.
|
||||
"""
|
||||
|
||||
path: str # relative to the raw dir
|
||||
size: int
|
||||
sha256: str | None = None # ``"sha256:<hex>"`` form or ``None``
|
||||
|
||||
|
||||
@dataclass
|
||||
class Manifest:
|
||||
"""Schema for ``_manifest.json``. Backward-compat policy: new optional
|
||||
fields can be added without bumping version; **any** mismatch on existing
|
||||
field semantics requires a version bump.
|
||||
"""
|
||||
|
||||
source_content_hash: str # ``"sha256:<hex>"`` of source file
|
||||
source_size_bytes: int
|
||||
source_filename_at_parse: str
|
||||
critical_file: ManifestFile # content_list.json; size + sha256
|
||||
files: list[ManifestFile] # other files; size only
|
||||
total_size_bytes: int
|
||||
task_id: str = ""
|
||||
api_mode: str = ""
|
||||
engine_version: str = ""
|
||||
endpoint_signature: str = ""
|
||||
options_signature: str = ""
|
||||
downloaded_at: str = ""
|
||||
version: str = MANIFEST_VERSION
|
||||
engine: str = MANIFEST_ENGINE
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
return {
|
||||
"version": self.version,
|
||||
"engine": self.engine,
|
||||
"api_mode": self.api_mode,
|
||||
"engine_version": self.engine_version,
|
||||
"endpoint_signature": self.endpoint_signature,
|
||||
"options_signature": self.options_signature,
|
||||
"source_content_hash": self.source_content_hash,
|
||||
"source_size_bytes": int(self.source_size_bytes),
|
||||
"source_filename_at_parse": self.source_filename_at_parse,
|
||||
"task_id": self.task_id,
|
||||
"downloaded_at": self.downloaded_at,
|
||||
"critical_file": asdict(self.critical_file),
|
||||
"files": [asdict(f) for f in self.files],
|
||||
"total_size_bytes": int(self.total_size_bytes),
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, payload: dict) -> "Manifest":
|
||||
critical_raw = payload.get("critical_file") or {}
|
||||
files_raw = payload.get("files") or []
|
||||
return cls(
|
||||
version=str(payload.get("version") or MANIFEST_VERSION),
|
||||
engine=str(payload.get("engine") or MANIFEST_ENGINE),
|
||||
api_mode=str(payload.get("api_mode") or ""),
|
||||
engine_version=str(payload.get("engine_version") or ""),
|
||||
endpoint_signature=str(payload.get("endpoint_signature") or ""),
|
||||
options_signature=str(payload.get("options_signature") or ""),
|
||||
source_content_hash=str(payload.get("source_content_hash") or ""),
|
||||
source_size_bytes=int(payload.get("source_size_bytes") or 0),
|
||||
source_filename_at_parse=str(payload.get("source_filename_at_parse") or ""),
|
||||
task_id=str(payload.get("task_id") or ""),
|
||||
downloaded_at=str(payload.get("downloaded_at") or ""),
|
||||
critical_file=ManifestFile(
|
||||
path=str(critical_raw.get("path") or ""),
|
||||
size=int(critical_raw.get("size") or 0),
|
||||
sha256=(
|
||||
str(critical_raw["sha256"]) if critical_raw.get("sha256") else None
|
||||
),
|
||||
),
|
||||
files=[
|
||||
ManifestFile(
|
||||
path=str(f.get("path") or ""),
|
||||
size=int(f.get("size") or 0),
|
||||
sha256=(str(f["sha256"]) if f.get("sha256") else None),
|
||||
)
|
||||
for f in files_raw
|
||||
if isinstance(f, dict)
|
||||
],
|
||||
total_size_bytes=int(payload.get("total_size_bytes") or 0),
|
||||
)
|
||||
|
||||
|
||||
def manifest_path(raw_dir: Path) -> Path:
|
||||
return raw_dir / MANIFEST_FILENAME
|
||||
|
||||
|
||||
def load_manifest(raw_dir: Path) -> Manifest | None:
|
||||
"""Return the parsed manifest or ``None`` if absent / malformed."""
|
||||
p = manifest_path(raw_dir)
|
||||
if not p.is_file():
|
||||
return None
|
||||
try:
|
||||
payload = json.loads(p.read_text(encoding="utf-8"))
|
||||
except (OSError, json.JSONDecodeError):
|
||||
return None
|
||||
if not isinstance(payload, dict):
|
||||
return None
|
||||
if payload.get("version") != MANIFEST_VERSION:
|
||||
return None
|
||||
if payload.get("engine") != MANIFEST_ENGINE:
|
||||
return None
|
||||
try:
|
||||
return Manifest.from_dict(payload)
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
|
||||
|
||||
def write_manifest(raw_dir: Path, manifest: Manifest) -> None:
|
||||
"""Atomically write the manifest. The temp-file + rename pattern
|
||||
guarantees the manifest never appears in a partially-written state."""
|
||||
raw_dir.mkdir(parents=True, exist_ok=True)
|
||||
final = manifest_path(raw_dir)
|
||||
tmp = final.with_suffix(".json.tmp")
|
||||
tmp.write_text(
|
||||
json.dumps(manifest.to_dict(), ensure_ascii=False, indent=2),
|
||||
encoding="utf-8",
|
||||
)
|
||||
os.replace(tmp, final)
|
||||
|
||||
|
||||
# Re-exported for convenience.
|
||||
__all__ = [
|
||||
"MANIFEST_FILENAME",
|
||||
"MANIFEST_VERSION",
|
||||
"MANIFEST_ENGINE",
|
||||
"Manifest",
|
||||
"ManifestFile",
|
||||
"load_manifest",
|
||||
"manifest_path",
|
||||
"write_manifest",
|
||||
]
|
||||
+51
@@ -0,0 +1,51 @@
|
||||
"""MinerU engine adapter (implements ExternalParserBase hooks)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from lightrag.constants import MINERU_RAW_DIR_SUFFIX, PARSER_ENGINE_MINERU
|
||||
from lightrag.parser.external._base import ExternalParserBase
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from lightrag.sidecar.ir import IRDoc
|
||||
|
||||
|
||||
class MinerUParser(ExternalParserBase):
|
||||
engine_name = PARSER_ENGINE_MINERU
|
||||
raw_dir_suffix = MINERU_RAW_DIR_SUFFIX
|
||||
force_reparse_env = "LIGHTRAG_FORCE_REPARSE_MINERU"
|
||||
|
||||
def is_bundle_valid(
|
||||
self,
|
||||
raw_dir: Path,
|
||||
source_path: Path,
|
||||
*,
|
||||
engine_params: "Mapping[str, Any] | None" = None,
|
||||
) -> bool:
|
||||
from lightrag.parser.external.mineru import is_bundle_valid
|
||||
|
||||
return is_bundle_valid(raw_dir, source_path, overrides=engine_params)
|
||||
|
||||
async def download_into(
|
||||
self,
|
||||
raw_dir: Path,
|
||||
source_path: Path,
|
||||
*,
|
||||
upload_name: str,
|
||||
engine_params: "Mapping[str, Any] | None" = None,
|
||||
) -> None:
|
||||
from lightrag.parser.external.mineru import MinerURawClient
|
||||
|
||||
await MinerURawClient(overrides=engine_params).download_into(
|
||||
raw_dir, source_path, upload_name=upload_name
|
||||
)
|
||||
|
||||
def build_ir(self, raw_dir: Path, document_name: str) -> "IRDoc":
|
||||
from lightrag.parser.external.mineru import MinerUIRBuilder
|
||||
|
||||
return MinerUIRBuilder().normalize_from_workdir(
|
||||
raw_dir, document_name=document_name
|
||||
)
|
||||
Reference in New Issue
Block a user