chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,13 @@
|
||||
"""Document parsing layer.
|
||||
|
||||
Sub-packages and modules:
|
||||
|
||||
- ``routing``: parser engine selection and per-file parser directives.
|
||||
- ``debug``: minimal ``LightRAG`` stand-in for offline parser debugging.
|
||||
- ``cli``: ``python -m lightrag.parser.cli`` entry point for single-file
|
||||
parser debugging across all engines.
|
||||
- ``docx``: native ``.docx`` parser. Additional native format parsers
|
||||
should live as sibling sub-packages here (e.g. ``parser/pdf/``).
|
||||
- ``external``: adapters for external parsing services (``mineru``,
|
||||
``docling``) that post to a remote API and cache the raw bundle.
|
||||
"""
|
||||
@@ -0,0 +1,217 @@
|
||||
"""Shared HTML-table parsing helpers for parser engines.
|
||||
|
||||
Pure functions that recover structural facts from an HTML ``<table>`` string
|
||||
without a heavy dependency: row/column counts (colspan-aware), the verbatim
|
||||
``<thead>`` substring, table-payload detection, and ``<html>/<body>`` wrapper
|
||||
stripping. Originally private to the mineru IR builder; lifted into a leaf
|
||||
module so the native markdown IR builder can reuse the exact same logic
|
||||
(merged-cell semantics must survive identically across engines).
|
||||
|
||||
Leaf module with no parser-layer imports — safe for any engine to import.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
from html.parser import HTMLParser
|
||||
|
||||
from lightrag.utils import logger
|
||||
|
||||
|
||||
@dataclass
|
||||
class HTMLTableInfo:
|
||||
num_rows: int = 0
|
||||
num_cols: int = 0
|
||||
|
||||
|
||||
class _HTMLTableInfoParser(HTMLParser):
|
||||
"""Count ``<tr>`` rows and their (colspan-aware) column widths.
|
||||
|
||||
Used only to recover ``num_rows`` / ``num_cols`` when the engine did not
|
||||
supply them; the ``<thead>`` header itself is preserved verbatim by
|
||||
:func:`extract_thead_html`, not reconstructed here.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__(convert_charrefs=True)
|
||||
# ``col_count`` (sum of colspans) for each completed top-level ``<tr>``.
|
||||
self.row_col_counts: list[int] = []
|
||||
self._tr_depth = 0
|
||||
self._cell_depth = 0
|
||||
self._row_col_count = 0
|
||||
self._cell_colspan = 1
|
||||
|
||||
def handle_starttag(self, tag: str, attrs: list[tuple[str, str | None]]) -> None:
|
||||
tag = tag.lower()
|
||||
if tag == "tr":
|
||||
if self._tr_depth == 0:
|
||||
self._row_col_count = 0
|
||||
self._tr_depth += 1
|
||||
return
|
||||
if tag in {"td", "th"} and self._tr_depth > 0:
|
||||
if self._cell_depth == 0:
|
||||
self._cell_colspan = _cell_span(attrs, "colspan")
|
||||
self._cell_depth += 1
|
||||
|
||||
def handle_endtag(self, tag: str) -> None:
|
||||
tag = tag.lower()
|
||||
if tag in {"td", "th"} and self._cell_depth > 0:
|
||||
self._cell_depth -= 1
|
||||
if self._cell_depth == 0:
|
||||
self._row_col_count += self._cell_colspan
|
||||
self._cell_colspan = 1
|
||||
return
|
||||
if tag == "tr" and self._tr_depth > 0:
|
||||
self._tr_depth -= 1
|
||||
# ``col_count > 0`` ⇔ the row held at least one cell (colspan ≥ 1),
|
||||
# so an empty ``<tr></tr>`` is skipped exactly as before.
|
||||
if self._tr_depth == 0 and self._row_col_count > 0:
|
||||
self.row_col_counts.append(self._row_col_count)
|
||||
self._row_col_count = 0
|
||||
|
||||
|
||||
def _cell_span(attrs: list[tuple[str, str | None]], name: str) -> int:
|
||||
"""Read a ``colspan``/``rowspan`` attribute as an int ``>= 1`` (default 1)."""
|
||||
for key, value in attrs:
|
||||
if key.lower() != name:
|
||||
continue
|
||||
try:
|
||||
return max(int(value or "1"), 1)
|
||||
except ValueError:
|
||||
return 1
|
||||
return 1
|
||||
|
||||
|
||||
def extract_html_table_info(html: str) -> HTMLTableInfo:
|
||||
parser = _HTMLTableInfoParser()
|
||||
try:
|
||||
parser.feed(html or "")
|
||||
parser.close()
|
||||
except Exception as exc: # pragma: no cover - HTMLParser is forgiving.
|
||||
logger.debug("[html_table] failed to parse table HTML: %s", exc)
|
||||
return HTMLTableInfo()
|
||||
return HTMLTableInfo(
|
||||
num_rows=len(parser.row_col_counts),
|
||||
num_cols=max(parser.row_col_counts, default=0),
|
||||
)
|
||||
|
||||
|
||||
def extract_thead_html(html: str) -> str | None:
|
||||
"""Return the first top-level ``<thead …>…</thead>`` substring verbatim.
|
||||
|
||||
The raw markup is kept so merged-cell semantics (``rowspan`` / ``colspan``)
|
||||
survive into ``tables.json`` and, later, into every repeated header chunk
|
||||
of a split table. Returns ``None`` when the table has no ``<thead>`` or the
|
||||
``<thead>`` carries no visible text (a blank spacer row, which would
|
||||
otherwise emit empty ``<th>`` headers).
|
||||
"""
|
||||
stripped = (html or "").strip()
|
||||
lower = stripped.lower()
|
||||
start = find_html_tag(lower, "thead")
|
||||
if start < 0:
|
||||
return None
|
||||
close = lower.find("</thead>", start)
|
||||
if close < 0:
|
||||
return None
|
||||
thead = stripped[start : close + len("</thead>")]
|
||||
# Blank check: drop a header whose cells hold no non-whitespace text.
|
||||
if not re.sub(r"<[^>]+>", "", thead).strip():
|
||||
return None
|
||||
return thead
|
||||
|
||||
|
||||
def looks_like_html_table_payload(body: str) -> bool:
|
||||
lower = (body or "").lstrip().lower()
|
||||
return any(
|
||||
starts_with_html_tag(lower, tag)
|
||||
for tag in ("table", "thead", "tbody", "tfoot", "tr", "html", "body")
|
||||
)
|
||||
|
||||
|
||||
def unwrap_html_table(payload: str) -> str:
|
||||
"""Strip a ``<html>/<body>`` document wrapper that a table model
|
||||
sometimes emits, returning the outermost ``<table…>…</table>`` span. Keeps
|
||||
a single clean ``<table>`` so the writer does not nest tables and the
|
||||
non-greedy ``TABLE_TAG_RE`` is not truncated at an inner ``</table>``.
|
||||
Falls back to the stripped payload when no ``<table>`` element exists."""
|
||||
stripped = (payload or "").strip()
|
||||
lower = stripped.lower()
|
||||
start = _find_table_open(lower)
|
||||
if start < 0:
|
||||
return stripped
|
||||
close = lower.rfind("</table>")
|
||||
if close < start:
|
||||
return stripped
|
||||
return stripped[start : close + len("</table>")]
|
||||
|
||||
|
||||
def _find_table_open(lower: str) -> int:
|
||||
"""First index of a real ``<table`` start tag (not e.g. ``<tablefoo``).
|
||||
Returns -1 when none is present."""
|
||||
return find_html_tag(lower, "table")
|
||||
|
||||
|
||||
def find_html_tag(lower: str, tag: str) -> int:
|
||||
"""First index of a real ``<tag`` start tag (not e.g. ``<tablefoo`` for
|
||||
``tag="table"``). ``lower`` must already be lower-cased. Returns -1 when
|
||||
none is present."""
|
||||
needle = f"<{tag}"
|
||||
idx = 0
|
||||
while True:
|
||||
idx = lower.find(needle, idx)
|
||||
if idx < 0:
|
||||
return -1
|
||||
nxt = idx + len(needle)
|
||||
if nxt >= len(lower) or lower[nxt] in {" ", "\t", "\r", "\n", ">", "/"}:
|
||||
return idx
|
||||
idx = nxt
|
||||
|
||||
|
||||
def starts_with_html_tag(lower: str, tag: str) -> bool:
|
||||
prefix = f"<{tag}"
|
||||
if not lower.startswith(prefix):
|
||||
return False
|
||||
if len(lower) == len(prefix):
|
||||
return True
|
||||
return lower[len(prefix)] in {" ", "\t", "\r", "\n", ">", "/"}
|
||||
|
||||
|
||||
def html_table_inner_body(html: str) -> str:
|
||||
stripped = (html or "").strip()
|
||||
lower = stripped.lower()
|
||||
if not starts_with_html_tag(lower, "table"):
|
||||
return stripped
|
||||
open_end = _open_tag_end(stripped)
|
||||
close_start = lower.rfind("</table>")
|
||||
if open_end < 0 or close_start <= open_end:
|
||||
return stripped
|
||||
return stripped[open_end + 1 : close_start].strip()
|
||||
|
||||
|
||||
def _open_tag_end(html: str) -> int:
|
||||
"""Index of the ``>`` closing the leading tag, skipping quoted attribute
|
||||
values so a ``>`` inside an attribute (e.g. ``<table data-x="a>b">``) does
|
||||
not terminate the tag early. Returns -1 when no closing ``>`` is found."""
|
||||
quote: str | None = None
|
||||
for idx, ch in enumerate(html):
|
||||
if quote is not None:
|
||||
if ch == quote:
|
||||
quote = None
|
||||
elif ch in {'"', "'"}:
|
||||
quote = ch
|
||||
elif ch == ">":
|
||||
return idx
|
||||
return -1
|
||||
|
||||
|
||||
__all__ = [
|
||||
"HTMLTableInfo",
|
||||
"extract_html_table_info",
|
||||
"extract_thead_html",
|
||||
"looks_like_html_table_payload",
|
||||
"unwrap_html_table",
|
||||
"find_html_tag",
|
||||
"starts_with_html_tag",
|
||||
"html_table_inner_body",
|
||||
]
|
||||
@@ -0,0 +1,72 @@
|
||||
"""Shared markdown rendering helpers for parser engines.
|
||||
|
||||
Used by the native docx parser and the external (mineru / docling) IR
|
||||
builders so heading-line rendering stays identical across engines. This is
|
||||
a leaf module with no heavy imports — ``lightrag/parser/__init__.py`` only
|
||||
carries a docstring — so all three engines can import it without risking a
|
||||
circular dependency.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
|
||||
# Markdown caps heading levels at 6 (``######``); deeper outline levels are
|
||||
# clamped to 6 rather than emitting an illegal 7+ run of ``#``.
|
||||
MAX_HEADING_LEVEL = 6
|
||||
|
||||
# A heading line that is ALREADY markdown: 1-6 ``#`` followed by one or more
|
||||
# spaces. Used to avoid double-prefixing text that an upstream engine emitted
|
||||
# with its own markdown heading marker (e.g. mineru/docling extracting
|
||||
# ``# Foo``).
|
||||
#
|
||||
# Ambiguity note: this is a heuristic, not a parse. A heading whose text
|
||||
# genuinely starts with ``#`` plus a space (an author literally writing
|
||||
# ``"# Note"`` as heading content) is indistinguishable from an
|
||||
# already-rendered markdown heading and will be treated as the latter. This
|
||||
# is accepted as a rare edge case in exchange for engine-agnostic dedup.
|
||||
_MD_HEADING_RE = re.compile(r"^#{1,6} +")
|
||||
|
||||
|
||||
def strip_heading_markdown_prefix(text: str) -> str:
|
||||
"""Return heading metadata without an existing markdown heading prefix.
|
||||
|
||||
The content renderer may keep a source line such as ``"# Foo"`` verbatim
|
||||
to avoid double-prefixing, but structured metadata (``heading``,
|
||||
``parent_headings``, doc title) must stay clean. The whole ``#`` run and
|
||||
all following spaces are removed, so ``"# Extra"`` yields ``"Extra"``
|
||||
(no leading space leaks into the metadata).
|
||||
|
||||
See the module-level ambiguity note: text that genuinely begins with a
|
||||
``#`` + space run is stripped here too.
|
||||
"""
|
||||
return _MD_HEADING_RE.sub("", text, count=1)
|
||||
|
||||
|
||||
def render_heading_line(level: int, text: str) -> str:
|
||||
"""Render a heading as a markdown-prefixed content line.
|
||||
|
||||
Args:
|
||||
level: 1-based heading level (1 = H1). Values < 1 are treated as 1;
|
||||
values > :data:`MAX_HEADING_LEVEL` are clamped so a level >= 7
|
||||
heading still gets ``######``.
|
||||
text: The heading text.
|
||||
|
||||
Returns:
|
||||
``text`` unchanged when it already starts with a markdown heading
|
||||
prefix (``^#{1,6} +`` — 1-6 ``#`` then one or more spaces); otherwise
|
||||
``"#" * clamped_level + " " + text``. See the module-level ambiguity
|
||||
note: a heading whose content genuinely begins with such a run is
|
||||
kept verbatim rather than re-prefixed.
|
||||
"""
|
||||
if _MD_HEADING_RE.match(text):
|
||||
return text
|
||||
hashes = "#" * min(max(level, 1), MAX_HEADING_LEVEL)
|
||||
return f"{hashes} {text}"
|
||||
|
||||
|
||||
__all__ = [
|
||||
"MAX_HEADING_LEVEL",
|
||||
"render_heading_line",
|
||||
"strip_heading_markdown_prefix",
|
||||
]
|
||||
@@ -0,0 +1,155 @@
|
||||
"""Unified parser contract for native + external parser engines.
|
||||
|
||||
Every engine (native docx, mineru, docling, legacy, plus the internal
|
||||
``reuse``/``passthrough`` format handlers) implements :class:`BaseParser`.
|
||||
The pipeline dispatches through the registry
|
||||
(:mod:`lightrag.parser.registry`) instead of a growing ``if engine == …``
|
||||
chain.
|
||||
|
||||
Design notes:
|
||||
|
||||
- ``BaseParser`` carries *behaviour only* (the ``parse`` coroutine and any
|
||||
engine-private hooks). Capability metadata (supported suffixes, queue
|
||||
group, endpoint requirements) lives in the registry's lightweight
|
||||
``ParserSpec`` table so capability queries never import a parser
|
||||
implementation.
|
||||
- ``ParseResult.to_dict()`` emits only semantically-present fields so the
|
||||
returned dict stays byte-for-byte compatible with the pre-refactor
|
||||
``parse_native``/``parse_mineru``/``parse_docling`` return shapes (the
|
||||
worker reads these by ``.get(...)``).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from lightrag.sidecar.ir import IRDoc # noqa: F401
|
||||
|
||||
|
||||
@dataclass
|
||||
class ResolvedSource:
|
||||
"""The common parse preamble shared by every engine."""
|
||||
|
||||
source_path: Path
|
||||
document_name: str
|
||||
parsed_dir: Path
|
||||
|
||||
|
||||
@dataclass
|
||||
class ParseContext:
|
||||
"""Inputs handed to :meth:`BaseParser.parse`.
|
||||
|
||||
Wraps the LightRAG instance plus the per-document handles so a parser
|
||||
does not receive a bare ``self``. Convenience helpers lazily import
|
||||
pipeline-layer functions at call time, keeping this module import-cheap
|
||||
and free of import cycles.
|
||||
"""
|
||||
|
||||
rag: Any
|
||||
doc_id: str
|
||||
file_path: str
|
||||
content_data: dict[str, Any]
|
||||
|
||||
def source_path(self, parser_engine: str) -> Path:
|
||||
"""Resolve the on-disk source file for this document."""
|
||||
from lightrag.pipeline import (
|
||||
_call_source_file_resolver,
|
||||
_read_source_file,
|
||||
)
|
||||
|
||||
return Path(
|
||||
_call_source_file_resolver(
|
||||
self.rag,
|
||||
self.file_path,
|
||||
source_file=_read_source_file(self.content_data),
|
||||
parser_engine=parser_engine,
|
||||
)
|
||||
)
|
||||
|
||||
def resolve(self, parser_engine: str) -> ResolvedSource:
|
||||
"""Resolve ``(source_path, document_name, parsed_dir)``.
|
||||
|
||||
Mirrors the preamble shared by ``parse_mineru``/``parse_docling``/
|
||||
``parse_native``: canonicalize the document name defensively (so
|
||||
direct callers may pass absolute or hint-bearing paths) and derive
|
||||
the ``__parsed__/<base>.parsed/`` output directory.
|
||||
"""
|
||||
from lightrag.utils_pipeline import (
|
||||
normalize_document_file_path,
|
||||
parsed_artifact_dir_for,
|
||||
)
|
||||
|
||||
source_path = self.source_path(parser_engine)
|
||||
document_name = normalize_document_file_path(self.file_path)
|
||||
if document_name == "unknown_source":
|
||||
document_name = source_path.name or f"{self.doc_id}.bin"
|
||||
parsed_dir = parsed_artifact_dir_for(
|
||||
document_name, parent_hint=source_path.parent
|
||||
)
|
||||
return ResolvedSource(source_path, document_name, parsed_dir)
|
||||
|
||||
async def archive_source(self, source_path: str) -> str | None:
|
||||
"""Archive the source after a successful parse + full_docs sync.
|
||||
|
||||
Resolved through the pipeline module's namespace (where the function
|
||||
is imported) so existing tests that patch
|
||||
``lightrag.pipeline.archive_docx_source_after_full_docs_sync`` keep
|
||||
intercepting it now that the call site lives in the parser layer.
|
||||
"""
|
||||
import lightrag.pipeline as _pipeline
|
||||
|
||||
return await _pipeline.archive_docx_source_after_full_docs_sync(source_path)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ParseResult:
|
||||
"""Structured parser output.
|
||||
|
||||
``to_dict`` only emits fields that carry meaning so the dict matches the
|
||||
pre-refactor return shapes exactly (no spurious ``None``/``False`` keys).
|
||||
"""
|
||||
|
||||
doc_id: str
|
||||
file_path: str
|
||||
parse_format: str
|
||||
content: str
|
||||
blocks_path: str = ""
|
||||
parse_engine: str | None = None
|
||||
parse_stage_skipped: bool = False
|
||||
parse_warnings: dict[str, Any] | None = None
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
out: dict[str, Any] = {
|
||||
"doc_id": self.doc_id,
|
||||
"file_path": self.file_path,
|
||||
"parse_format": self.parse_format,
|
||||
"content": self.content,
|
||||
"blocks_path": self.blocks_path,
|
||||
}
|
||||
if self.parse_engine is not None:
|
||||
out["parse_engine"] = self.parse_engine
|
||||
if self.parse_stage_skipped:
|
||||
out["parse_stage_skipped"] = True
|
||||
if self.parse_warnings:
|
||||
out["parse_warnings"] = self.parse_warnings
|
||||
return out
|
||||
|
||||
|
||||
class BaseParser(ABC):
|
||||
"""Abstract base for every parser engine.
|
||||
|
||||
Subclasses set ``engine_name`` and implement :meth:`parse`. Capability
|
||||
metadata (suffixes/queue group/endpoint) is declared in the registry
|
||||
``ParserSpec``, not here.
|
||||
"""
|
||||
|
||||
engine_name: str
|
||||
|
||||
@abstractmethod
|
||||
async def parse(self, ctx: ParseContext) -> ParseResult:
|
||||
"""Parse one document and return its :class:`ParseResult`."""
|
||||
...
|
||||
@@ -0,0 +1,314 @@
|
||||
"""Unified sidecar debug CLI for any registered parser engine.
|
||||
|
||||
Dispatches one source file through the parser registry
|
||||
(``get_parser(engine).parse(...)``) and writes the resulting sidecar (and
|
||||
raw bundle, for external engines) into a flat layout — no ``__parsed__/``
|
||||
middle layer, source file never archived — so the artifacts can be
|
||||
inspected next to the input file. Because dispatch goes through the
|
||||
registry, a third-party engine registered via ``register_parser`` is a
|
||||
valid ``--engine`` choice with no CLI changes.
|
||||
|
||||
Invocation::
|
||||
|
||||
python -m lightrag.parser.cli path/to/sample.docx --engine native
|
||||
python -m lightrag.parser.cli path/to/sample.pdf --engine mineru
|
||||
python -m lightrag.parser.cli path/to/sample.pdf --engine docling --force-reparse
|
||||
|
||||
See ``docs/ParserDebugCLI-zh.md`` for the full reference.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import json
|
||||
import sys
|
||||
from contextlib import ExitStack
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from unittest import mock
|
||||
|
||||
|
||||
def _normalize_direct_script_sys_path() -> None:
|
||||
if __package__:
|
||||
return
|
||||
|
||||
parser_dir = Path(__file__).resolve().parent
|
||||
repo_root = parser_dir.parent.parent
|
||||
|
||||
# Direct execution adds lightrag/parser to sys.path, which makes the
|
||||
# native parser's third-party ``docx`` import resolve to parser/docx.
|
||||
sys.path[:] = [
|
||||
entry for entry in sys.path if Path(entry or ".").resolve() != parser_dir
|
||||
]
|
||||
repo_root_str = str(repo_root)
|
||||
if repo_root_str in sys.path:
|
||||
sys.path.remove(repo_root_str)
|
||||
sys.path.insert(0, repo_root_str)
|
||||
|
||||
|
||||
_normalize_direct_script_sys_path()
|
||||
|
||||
|
||||
def _build_parser() -> argparse.ArgumentParser:
|
||||
# Registry import is cheap by design (no parser impl is pulled in), so
|
||||
# deriving --engine choices here keeps third-party engines selectable
|
||||
# without making --help pay for a heavy import.
|
||||
from lightrag.parser.registry import supported_parser_engines
|
||||
|
||||
engines = sorted(supported_parser_engines())
|
||||
|
||||
parser = argparse.ArgumentParser(
|
||||
prog="parse_sidecar",
|
||||
description=(
|
||||
"Run a registered parser engine on a single file and emit sidecar "
|
||||
"artifacts (plus a raw bundle for external engines) into a flat "
|
||||
"layout alongside the source. No __parsed__/ middle layer; the "
|
||||
"source file is never moved."
|
||||
),
|
||||
)
|
||||
parser.add_argument("input_file", type=Path, help="Source file to parse.")
|
||||
parser.add_argument(
|
||||
"--engine",
|
||||
required=True,
|
||||
choices=engines,
|
||||
help="Parser engine to drive.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"-o",
|
||||
"--sidecar-parent-dir",
|
||||
type=Path,
|
||||
default=None,
|
||||
help=(
|
||||
"Parent directory for <name>.parsed/ and <name>.<engine>_raw/. "
|
||||
"Default: the source file's parent directory."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--doc-id",
|
||||
default=None,
|
||||
help="Override the doc id. Default: doc-<md5(absolute input path)>.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--force-reparse",
|
||||
action="store_true",
|
||||
help=(
|
||||
"Only affects mineru/docling. By default a non-empty raw_dir is "
|
||||
"treated as a valid cache and reused without manifest checks; "
|
||||
"this flag clears raw_dir and forces a fresh download/parse."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--preview",
|
||||
type=int,
|
||||
default=5,
|
||||
metavar="N",
|
||||
help="Number of block rows to preview after parsing (0 disables).",
|
||||
)
|
||||
return parser
|
||||
|
||||
|
||||
def _print_summary(blocks_path: Path, raw_dir: Path | None, preview: int) -> None:
|
||||
with blocks_path.open("r", encoding="utf-8") as fh:
|
||||
meta_line = fh.readline().strip()
|
||||
if not meta_line:
|
||||
raise SystemExit(f"empty blocks file at {blocks_path}")
|
||||
meta = json.loads(meta_line)
|
||||
rows = [json.loads(line) for line in fh if line.strip()]
|
||||
parsed_dir = blocks_path.parent
|
||||
print(f"parsed dir : {parsed_dir} (exists={parsed_dir.exists()})")
|
||||
if raw_dir is not None:
|
||||
print(f"raw dir : {raw_dir} (exists={raw_dir.exists()})")
|
||||
print(f"document : {meta.get('document_name')}")
|
||||
print(f"doc_id : {meta.get('doc_id')}")
|
||||
print(f"engine : {meta.get('parse_engine')}")
|
||||
print(f"blocks : {meta.get('blocks')}")
|
||||
print(
|
||||
f"sidecars : tables={meta.get('table_file')} "
|
||||
f"drawings={meta.get('drawing_file')} "
|
||||
f"equations={meta.get('equation_file')} "
|
||||
f"asset_dir={meta.get('asset_dir')}"
|
||||
)
|
||||
if preview > 0 and rows:
|
||||
shown = min(preview, len(rows))
|
||||
print(f"--- preview (first {shown} of {len(rows)} blocks) ---")
|
||||
for row in rows[:preview]:
|
||||
heading = row.get("heading") or ""
|
||||
content = (row.get("content") or "").replace("\n", " ")
|
||||
snippet = content if len(content) <= 80 else content[:77] + "..."
|
||||
print(f" [{row.get('blockid', '')[:8]}] heading={heading!r} :: {snippet}")
|
||||
|
||||
|
||||
def _print_raw_summary(result: dict, preview: int) -> None:
|
||||
"""Summary for engines that emit plain content with no sidecar (legacy /
|
||||
any non-sidecar third-party engine: ``blocks_path`` is empty)."""
|
||||
content = result.get("content") or ""
|
||||
print(f"engine : {result.get('parse_engine')}")
|
||||
print(f"format : {result.get('parse_format')}")
|
||||
print(f"content : {len(content)} chars")
|
||||
if preview > 0 and content:
|
||||
snippet = content[:400]
|
||||
print(f"--- preview (first {len(snippet)} of {len(content)} chars) ---")
|
||||
print(snippet + ("..." if len(content) > len(snippet) else ""))
|
||||
|
||||
|
||||
async def _run(args: argparse.Namespace) -> int:
|
||||
# Pipeline + heavy parser imports are deferred so ``--help`` and the
|
||||
# input-file existence check don't pay for them.
|
||||
from lightrag.constants import FULL_DOCS_FORMAT_PENDING_PARSE
|
||||
from lightrag.parser.base import ParseContext
|
||||
from lightrag.parser.external._base import ExternalParserBase
|
||||
from lightrag.parser.registry import get_parser, suffix_capabilities
|
||||
from lightrag.parser.debug import build_debug_rag
|
||||
from lightrag.parser.docx.parse_document import DocxContentError
|
||||
from lightrag.utils import compute_mdhash_id
|
||||
import lightrag.pipeline as pipeline_mod
|
||||
import lightrag.utils_pipeline as utils_pipeline_mod
|
||||
|
||||
source = args.input_file.resolve()
|
||||
if not source.is_file():
|
||||
print(f"error: input file does not exist: {source}", file=sys.stderr)
|
||||
return 1
|
||||
|
||||
# Reject suffix/engine mismatches up-front: the pipeline would otherwise
|
||||
# fail deep inside the IR builder with a less helpful message.
|
||||
suffix = source.suffix.lstrip(".").lower()
|
||||
supported = suffix_capabilities(args.engine)
|
||||
if suffix not in supported:
|
||||
print(
|
||||
f"error: engine '{args.engine}' does not support .{suffix or '<no suffix>'} "
|
||||
f"files (supported: {', '.join(sorted(supported))})",
|
||||
file=sys.stderr,
|
||||
)
|
||||
return 1
|
||||
|
||||
parser = get_parser(args.engine)
|
||||
if parser is None:
|
||||
print(f"error: engine '{args.engine}' is not registered", file=sys.stderr)
|
||||
return 1
|
||||
is_external = isinstance(parser, ExternalParserBase)
|
||||
|
||||
sidecar_parent = (args.sidecar_parent_dir or source.parent).resolve()
|
||||
sidecar_parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
parsed_dir = sidecar_parent / f"{source.name}.parsed"
|
||||
# External engines preserve a raw bundle next to the sidecar; its name is
|
||||
# derived from the engine's own raw_dir_suffix (mirrors the flattened
|
||||
# raw_dir_for_parsed_dir layout), so any external engine works here.
|
||||
raw_dir = (
|
||||
sidecar_parent / f"{source.name}{parser.raw_dir_suffix}"
|
||||
if is_external
|
||||
else None
|
||||
)
|
||||
|
||||
doc_id = args.doc_id or compute_mdhash_id(str(source), prefix="doc-")
|
||||
|
||||
def _patched_artifact_dir(
|
||||
file_path: str | None = None,
|
||||
*,
|
||||
parent_hint: Any | None = None,
|
||||
) -> Path:
|
||||
# Flatten the production "<INPUT_DIR>/__parsed__/<base>.parsed/"
|
||||
# layout to "<sidecar_parent>/<source.name>.parsed/" so the sidecar
|
||||
# and the source file sit side by side.
|
||||
return parsed_dir
|
||||
|
||||
def _lenient_bundle(
|
||||
raw_dir_arg: Path, _source_file: Path | None = None, *_args: Any, **_kwargs: Any
|
||||
) -> bool:
|
||||
# Tolerate the ExternalParserBase hook's keyword-only ``engine_params``
|
||||
# (and any future args) — this stub replaces ``is_bundle_valid``.
|
||||
return raw_dir_arg.exists() and any(raw_dir_arg.iterdir())
|
||||
|
||||
def _force_miss(*_args: Any, **_kwargs: Any) -> bool:
|
||||
return False
|
||||
|
||||
bundle_check = _force_miss if args.force_reparse else _lenient_bundle
|
||||
|
||||
async def _noop_archive(*_args: Any, **_kwargs: Any) -> None:
|
||||
return None
|
||||
|
||||
rag = build_debug_rag()
|
||||
|
||||
with ExitStack() as stack:
|
||||
# Patch 1: redirect sidecar output to the flat layout.
|
||||
# parsed_artifact_dir_for is from-imported into pipeline at
|
||||
# module load, so patch both namespaces.
|
||||
stack.enter_context(
|
||||
mock.patch.object(
|
||||
utils_pipeline_mod,
|
||||
"parsed_artifact_dir_for",
|
||||
_patched_artifact_dir,
|
||||
)
|
||||
)
|
||||
stack.enter_context(
|
||||
mock.patch.object(
|
||||
pipeline_mod,
|
||||
"parsed_artifact_dir_for",
|
||||
_patched_artifact_dir,
|
||||
)
|
||||
)
|
||||
|
||||
# Patch 2: raw cache strategy. ExternalParserBase.parse calls
|
||||
# ``self.is_bundle_valid(...)``, so patch the resolved instance method
|
||||
# directly — works for any external engine without knowing its module.
|
||||
if is_external:
|
||||
stack.enter_context(
|
||||
mock.patch.object(parser, "is_bundle_valid", bundle_check)
|
||||
)
|
||||
|
||||
# Patch 3: keep the source file in place. Every engine archives via
|
||||
# ctx.archive_source -> lightrag.pipeline.archive_docx_source_after_...
|
||||
stack.enter_context(
|
||||
mock.patch.object(
|
||||
pipeline_mod,
|
||||
"archive_docx_source_after_full_docs_sync",
|
||||
_noop_archive,
|
||||
)
|
||||
)
|
||||
|
||||
try:
|
||||
result = (
|
||||
await parser.parse(
|
||||
ParseContext(
|
||||
rag,
|
||||
doc_id,
|
||||
str(source),
|
||||
{
|
||||
"parse_format": FULL_DOCS_FORMAT_PENDING_PARSE,
|
||||
"content": "",
|
||||
},
|
||||
)
|
||||
)
|
||||
).to_dict()
|
||||
except DocxContentError as exc:
|
||||
# The native DOCX parser now raises (instead of sys.exit) on a
|
||||
# content-limit violation so the server pipeline can fail just the
|
||||
# offending document. Preserve the friendly CLI UX: print the
|
||||
# formatted message and return a non-zero exit code, no traceback.
|
||||
print(str(exc), file=sys.stderr)
|
||||
return 1
|
||||
|
||||
blocks_path = result.get("blocks_path") or ""
|
||||
if blocks_path:
|
||||
_print_summary(Path(blocks_path), raw_dir, args.preview)
|
||||
else:
|
||||
# No sidecar (legacy / non-sidecar third-party engine): the result
|
||||
# carries plain content, so summarize that instead of a blocks file.
|
||||
_print_raw_summary(result, args.preview)
|
||||
return 0
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> int:
|
||||
# Discover third-party parser engines before --engine choices are built,
|
||||
# so a `lightrag.parsers` entry-point engine is selectable here exactly
|
||||
# like in the server (which loads plugins in create_app).
|
||||
from lightrag.parser.plugins import load_third_party_parsers
|
||||
|
||||
load_third_party_parsers()
|
||||
args = _build_parser().parse_args(argv)
|
||||
return asyncio.run(_run(args))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -0,0 +1,100 @@
|
||||
"""Shared debug LightRAG stand-in for registry-dispatched parsing.
|
||||
|
||||
A minimal ``LightRAG`` stand-in plus a deterministic ``datetime`` shim,
|
||||
shared by the unified parser debug CLI (``lightrag/parser/cli.py``),
|
||||
the golden-fixture regen script (``scripts/regen_native_docx_golden.py``),
|
||||
and the byte-equivalence golden tests
|
||||
(``tests/parser/docx/test_native_docx_golden.py``).
|
||||
|
||||
Every engine is driven the same way — ``get_parser(engine).parse(
|
||||
ParseContext(rag, ...))`` — and ``ParseContext`` reads the same ``rag``
|
||||
surface (``_persist_parsed_full_docs``, ``_resolve_source_file_for_parser``,
|
||||
``self.full_docs``, ``self.doc_status``), so a single stand-in covers every
|
||||
engine — when a parser grows a new dependency, extend this module rather
|
||||
than copy-pasting parallel stubs into each call site.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
|
||||
|
||||
class DebugFullDocs:
|
||||
"""In-memory ``full_docs`` shim — captures the persisted record."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.data: dict[str, Any] = {}
|
||||
|
||||
async def upsert(self, payload: dict[str, Any]) -> None:
|
||||
self.data.update(payload)
|
||||
|
||||
async def get_by_id(self, doc_id: str) -> Any:
|
||||
return self.data.get(doc_id)
|
||||
|
||||
async def index_done_callback(self) -> None:
|
||||
return None
|
||||
|
||||
|
||||
class DebugDocStatus:
|
||||
"""No-op ``doc_status`` shim — the parse_* methods never read/write it."""
|
||||
|
||||
async def get_by_id(self, doc_id: str) -> Any:
|
||||
return None
|
||||
|
||||
async def upsert(self, data: dict[str, Any]) -> None:
|
||||
return None
|
||||
|
||||
|
||||
def build_debug_rag():
|
||||
"""Build a minimal LightRAG stand-in that exposes what a parser reads.
|
||||
|
||||
The import of ``LightRAG`` is intentionally function-local: deferring
|
||||
it avoids a circular import when this helper is loaded during package
|
||||
init (the parser CLI resolves ``lightrag.parser.debug`` before
|
||||
``lightrag`` itself is fully bound).
|
||||
|
||||
A parser is driven via ``get_parser(engine).parse(ParseContext(rag, ...))``
|
||||
(CLI / golden tests / regen script). ``ParseContext`` reads these off the
|
||||
``rag`` it is handed — every entry MUST be provided by this stand-in:
|
||||
|
||||
- **methods** (rebound from :class:`LightRAG`):
|
||||
- ``_persist_parsed_full_docs(doc_id, payload)`` — async; touches
|
||||
``self.full_docs``.
|
||||
- ``_resolve_source_file_for_parser(file_path)`` — returns the
|
||||
on-disk source path. Stubbed to identity here since the CLI / tests
|
||||
feed an already-resolved path.
|
||||
- **storages**:
|
||||
- ``self.full_docs.upsert(...)`` / ``.get_by_id(...)`` /
|
||||
``.index_done_callback()`` — :class:`DebugFullDocs` covers all three.
|
||||
- ``self.doc_status.get_by_id(...)`` / ``.upsert(...)`` —
|
||||
:class:`DebugDocStatus` covers both.
|
||||
|
||||
When a parser grows a new dependency on the ``rag`` handed to
|
||||
``ParseContext``, extend this stand-in (and update the list above) rather
|
||||
than copy-pasting a parallel stub into the call sites.
|
||||
"""
|
||||
from lightrag import LightRAG
|
||||
|
||||
class _DebugRag:
|
||||
_persist_parsed_full_docs = LightRAG._persist_parsed_full_docs
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.full_docs = DebugFullDocs()
|
||||
self.doc_status = DebugDocStatus()
|
||||
|
||||
def _resolve_source_file_for_parser(self, file_path: str) -> str:
|
||||
return file_path
|
||||
|
||||
return _DebugRag()
|
||||
|
||||
|
||||
_FROZEN_NOW = datetime(2026, 1, 1, tzinfo=timezone.utc)
|
||||
|
||||
|
||||
class FrozenDateTime(datetime):
|
||||
"""Pin ``datetime.now`` so ``write_sidecar`` stamps a deterministic time."""
|
||||
|
||||
@classmethod
|
||||
def now(cls, tz=None): # noqa: D401
|
||||
return _FROZEN_NOW if tz is None else _FROZEN_NOW.astimezone(tz)
|
||||
@@ -0,0 +1,14 @@
|
||||
"""LightRAG native DOCX parser package.
|
||||
|
||||
The :mod:`parse_document` / :mod:`numbering_resolver` / :mod:`table_extractor` /
|
||||
:mod:`drawing_image_extractor` / :mod:`utils` / :mod:`omml` modules ship the
|
||||
upstream DOCX extraction logic verbatim (with imports localized for the new
|
||||
package path).
|
||||
|
||||
The pipeline-side orchestration (extract → IR → sidecar) now lives in
|
||||
:meth:`lightrag.pipeline._PipelineMixin.parse_native` so the native and
|
||||
MinerU engines share one shape; see :mod:`lightrag.parser.docx.ir_builder`
|
||||
for the engine IR builder.
|
||||
"""
|
||||
|
||||
__all__: list[str] = []
|
||||
@@ -0,0 +1,445 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
ABOUTME: Shared drawing/image extraction utilities for DOCX parsing and editing
|
||||
ABOUTME: Resolves w:drawing -> a:blip relationships, exports embedded images, builds placeholders
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import posixpath
|
||||
import re
|
||||
import shutil
|
||||
import zipfile
|
||||
from dataclasses import dataclass, field
|
||||
from html import escape, unescape
|
||||
from pathlib import Path, PurePosixPath
|
||||
from typing import Dict, Optional, Tuple
|
||||
from urllib.parse import urlparse
|
||||
|
||||
try:
|
||||
from defusedxml import ElementTree as ET
|
||||
except ImportError: # pragma: no cover
|
||||
from xml.etree import ElementTree as ET
|
||||
|
||||
|
||||
NS = {
|
||||
"w": "http://schemas.openxmlformats.org/wordprocessingml/2006/main",
|
||||
"wp": "http://schemas.openxmlformats.org/drawingml/2006/wordprocessingDrawing",
|
||||
"a": "http://schemas.openxmlformats.org/drawingml/2006/main",
|
||||
"r": "http://schemas.openxmlformats.org/officeDocument/2006/relationships",
|
||||
"v": "urn:schemas-microsoft-com:vml",
|
||||
}
|
||||
|
||||
REL_NS = "http://schemas.openxmlformats.org/package/2006/relationships"
|
||||
CONTENT_TYPE_NS = "http://schemas.openxmlformats.org/package/2006/content-types"
|
||||
IMAGE_REL_TYPE = (
|
||||
"http://schemas.openxmlformats.org/officeDocument/2006/relationships/image"
|
||||
)
|
||||
SOURCE_DOCUMENT_PART = "/word/document.xml"
|
||||
|
||||
# Match old and new drawing placeholders (requires id/name, allows extra attributes)
|
||||
DRAWING_PATTERN = re.compile(
|
||||
r'<drawing\b(?=[^>]*\bid="[^"]*")(?=[^>]*\bname="[^"]*")[^>]*/>'
|
||||
)
|
||||
DRAWING_TAG_PATTERN = re.compile(r"<drawing\b[^>]*/>")
|
||||
DRAWING_ATTR_PATTERN = re.compile(r'([a-zA-Z_][\w:.-]*)="([^"]*)"')
|
||||
|
||||
|
||||
@dataclass
|
||||
class DrawingRelationship:
|
||||
"""Relationship metadata for a single relationship ID."""
|
||||
|
||||
rel_id: str
|
||||
target: str
|
||||
target_mode: str
|
||||
rel_type: str
|
||||
part_name: Optional[str] = None
|
||||
content_type: Optional[str] = None
|
||||
image_format: Optional[str] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class DrawingExtractionContext:
|
||||
"""Context used to resolve and export drawing images for one DOCX file."""
|
||||
|
||||
docx_path: Path
|
||||
blocks_output_path: Optional[Path] = None
|
||||
export_dir_name: Optional[str] = None
|
||||
export_dir_path: Optional[Path] = None
|
||||
relationships: Dict[str, DrawingRelationship] = field(default_factory=dict)
|
||||
_exported_part_to_relpath: Dict[str, str] = field(default_factory=dict)
|
||||
_used_filenames: Dict[str, str] = field(default_factory=dict)
|
||||
|
||||
def resolve_relationship(self, rel_id: str) -> Optional[DrawingRelationship]:
|
||||
return self.relationships.get(rel_id)
|
||||
|
||||
def export_embedded_image(self, rel: DrawingRelationship) -> Optional[str]:
|
||||
"""
|
||||
Export an embedded image relationship target to export_dir.
|
||||
|
||||
Returns:
|
||||
Relative path like "<blocks_stem>.image/image1.png" if exported,
|
||||
or None when export is not applicable.
|
||||
"""
|
||||
if not self.export_dir_path or not self.export_dir_name:
|
||||
return None
|
||||
if rel.target_mode.lower() == "external":
|
||||
return None
|
||||
if not rel.part_name:
|
||||
return None
|
||||
if rel.part_name in self._exported_part_to_relpath:
|
||||
return self._exported_part_to_relpath[rel.part_name]
|
||||
|
||||
zip_member = rel.part_name.lstrip("/")
|
||||
try:
|
||||
with zipfile.ZipFile(self.docx_path, "r") as zf:
|
||||
blob = zf.read(zip_member)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
filename = self._dedupe_filename(PurePosixPath(rel.part_name).name or "image")
|
||||
output_file = self.export_dir_path / filename
|
||||
output_file.write_bytes(blob)
|
||||
|
||||
rel_path = str(PurePosixPath(self.export_dir_name) / filename)
|
||||
self._exported_part_to_relpath[rel.part_name] = rel_path
|
||||
return rel_path
|
||||
|
||||
def _dedupe_filename(self, base_name: str) -> str:
|
||||
if base_name not in self._used_filenames:
|
||||
self._used_filenames[base_name] = base_name
|
||||
return base_name
|
||||
|
||||
stem = Path(base_name).stem
|
||||
suffix = Path(base_name).suffix
|
||||
index = 2
|
||||
while True:
|
||||
candidate = f"{stem}_{index}{suffix}"
|
||||
if candidate not in self._used_filenames:
|
||||
self._used_filenames[candidate] = candidate
|
||||
return candidate
|
||||
index += 1
|
||||
|
||||
|
||||
def _normalize_image_format(ext_or_type: str) -> Optional[str]:
|
||||
if not ext_or_type:
|
||||
return None
|
||||
value = ext_or_type.strip().lower()
|
||||
|
||||
# Content-Type
|
||||
if value.startswith("image/"):
|
||||
value = value.split("/", 1)[1]
|
||||
if "+" in value:
|
||||
value = value.split("+", 1)[0]
|
||||
if value.startswith("x-"):
|
||||
value = value[2:]
|
||||
|
||||
# Extension (with or without leading dot)
|
||||
value = value.lstrip(".")
|
||||
if value == "jpg":
|
||||
return "jpeg"
|
||||
if value in {"jpeg", "png", "gif", "bmp", "tiff", "webp", "svg", "emf", "wmf"}:
|
||||
return value
|
||||
return value or None
|
||||
|
||||
|
||||
def _infer_format_from_target(target: str) -> Optional[str]:
|
||||
if not target:
|
||||
return None
|
||||
parsed = urlparse(target)
|
||||
path = parsed.path if parsed.scheme else target
|
||||
suffix = PurePosixPath(path).suffix
|
||||
return _normalize_image_format(suffix)
|
||||
|
||||
|
||||
def _resolve_part_name(source_part_name: str, target: str) -> str:
|
||||
if target.startswith("/"):
|
||||
return posixpath.normpath(target)
|
||||
source_dir = posixpath.dirname(source_part_name)
|
||||
joined = posixpath.join(source_dir, target)
|
||||
normalized = posixpath.normpath(joined)
|
||||
if not normalized.startswith("/"):
|
||||
normalized = "/" + normalized
|
||||
return normalized
|
||||
|
||||
|
||||
def create_drawing_context(
|
||||
docx_path: str,
|
||||
blocks_output_path: Optional[str] = None,
|
||||
) -> DrawingExtractionContext:
|
||||
"""
|
||||
Create extraction context for a DOCX file.
|
||||
|
||||
If blocks_output_path is provided, this also prepares `<blocks_stem>.image/`
|
||||
beside the blocks file and clears any previous content.
|
||||
"""
|
||||
docx_file = Path(docx_path)
|
||||
ctx = DrawingExtractionContext(docx_path=docx_file)
|
||||
|
||||
if blocks_output_path:
|
||||
output_path = Path(blocks_output_path)
|
||||
export_dir_name = f"{output_path.stem}.image"
|
||||
export_dir_path = output_path.parent / export_dir_name
|
||||
if export_dir_path.exists():
|
||||
shutil.rmtree(export_dir_path)
|
||||
export_dir_path.mkdir(parents=True, exist_ok=True)
|
||||
ctx.blocks_output_path = output_path
|
||||
ctx.export_dir_name = export_dir_name
|
||||
ctx.export_dir_path = export_dir_path
|
||||
|
||||
load_relationships(ctx)
|
||||
return ctx
|
||||
|
||||
|
||||
def load_relationships(ctx: DrawingExtractionContext) -> None:
|
||||
rels_xml = "word/_rels/document.xml.rels"
|
||||
content_types_xml = "[Content_Types].xml"
|
||||
|
||||
overrides: Dict[str, str] = {}
|
||||
defaults: Dict[str, str] = {}
|
||||
|
||||
try:
|
||||
with zipfile.ZipFile(ctx.docx_path, "r") as zf:
|
||||
if content_types_xml in zf.namelist():
|
||||
ct_root = ET.parse(zf.open(content_types_xml)).getroot()
|
||||
for node in ct_root.findall(f".//{{{CONTENT_TYPE_NS}}}Override"):
|
||||
part_name = node.get("PartName")
|
||||
content_type = node.get("ContentType")
|
||||
if part_name and content_type:
|
||||
overrides[part_name] = content_type
|
||||
for node in ct_root.findall(f".//{{{CONTENT_TYPE_NS}}}Default"):
|
||||
ext = node.get("Extension")
|
||||
content_type = node.get("ContentType")
|
||||
if ext and content_type:
|
||||
defaults[ext.lower()] = content_type
|
||||
|
||||
if rels_xml not in zf.namelist():
|
||||
return
|
||||
rels_root = ET.parse(zf.open(rels_xml)).getroot()
|
||||
except Exception:
|
||||
return
|
||||
|
||||
for rel in rels_root.findall(f".//{{{REL_NS}}}Relationship"):
|
||||
rel_id = rel.get("Id")
|
||||
target = rel.get("Target", "")
|
||||
target_mode = rel.get("TargetMode", "")
|
||||
rel_type = rel.get("Type", "")
|
||||
if not rel_id:
|
||||
continue
|
||||
|
||||
part_name = None
|
||||
content_type = None
|
||||
image_format = None
|
||||
|
||||
if target_mode.lower() != "external":
|
||||
part_name = _resolve_part_name(SOURCE_DOCUMENT_PART, target)
|
||||
if part_name:
|
||||
content_type = overrides.get(part_name)
|
||||
if not content_type:
|
||||
ext = PurePosixPath(part_name).suffix.lower().lstrip(".")
|
||||
content_type = defaults.get(ext)
|
||||
image_format = _normalize_image_format(
|
||||
content_type or _infer_format_from_target(part_name)
|
||||
)
|
||||
else:
|
||||
image_format = _normalize_image_format(_infer_format_from_target(target))
|
||||
|
||||
ctx.relationships[rel_id] = DrawingRelationship(
|
||||
rel_id=rel_id,
|
||||
target=target,
|
||||
target_mode=target_mode,
|
||||
rel_type=rel_type,
|
||||
part_name=part_name,
|
||||
content_type=content_type,
|
||||
image_format=image_format,
|
||||
)
|
||||
|
||||
|
||||
def _extract_blip_relationship(drawing_elem) -> Optional[Tuple[str, str]]:
|
||||
for blip in drawing_elem.findall(".//a:blip", NS):
|
||||
# Prefer explicit external links when both link/embed are present on one blip.
|
||||
# Word may keep an embedded cache for linked pictures.
|
||||
rel_link = blip.get(f"{{{NS['r']}}}link")
|
||||
if rel_link:
|
||||
return "link", rel_link
|
||||
rel_embed = blip.get(f"{{{NS['r']}}}embed")
|
||||
if rel_embed:
|
||||
return "embed", rel_embed
|
||||
return None
|
||||
|
||||
|
||||
def _extract_imagedata_relationship(container_elem) -> Optional[str]:
|
||||
"""Find an image relationship id from a w:pict / w:object via v:imagedata.
|
||||
|
||||
These legacy VML containers are how Word references EMF/WMF metafiles
|
||||
(and the rendered preview of any embedded OLE object). v:imagedata uses
|
||||
``r:id`` to point at the image part for both embedded and externally
|
||||
linked images — the relationship's ``TargetMode`` is what disambiguates
|
||||
the two cases, so the caller must inspect the resolved relationship.
|
||||
"""
|
||||
r_id_attr = f"{{{NS['r']}}}id"
|
||||
for imgdata in container_elem.findall(".//v:imagedata", NS):
|
||||
rel_id = imgdata.get(r_id_attr)
|
||||
if rel_id:
|
||||
return rel_id
|
||||
return None
|
||||
|
||||
|
||||
def _build_placeholder(attrs: Dict[str, str]) -> str:
|
||||
ordered_keys = ["id", "name", "path", "format"]
|
||||
pieces = []
|
||||
for key in ordered_keys:
|
||||
if key in attrs and attrs[key] is not None:
|
||||
pieces.append(f'{key}="{escape(str(attrs[key]), quote=True)}"')
|
||||
|
||||
# Preserve extra attributes deterministically (sorted by name)
|
||||
for key in sorted(k for k in attrs.keys() if k not in ordered_keys):
|
||||
value = attrs[key]
|
||||
if value is not None:
|
||||
pieces.append(f'{key}="{escape(str(value), quote=True)}"')
|
||||
|
||||
return f"<drawing {' '.join(pieces)} />"
|
||||
|
||||
|
||||
def extract_drawing_placeholder_from_element(
|
||||
drawing_elem,
|
||||
context: Optional[DrawingExtractionContext] = None,
|
||||
include_extended_attrs: bool = True,
|
||||
) -> str:
|
||||
"""
|
||||
Build a <drawing ... /> placeholder from a w:drawing element.
|
||||
|
||||
Behavior:
|
||||
- Always emits id/name from wp:docPr when present.
|
||||
- For embedded images (a:blip@r:embed): exports image and sets path/format.
|
||||
- For linked images (a:blip@r:link): does not download; path is original link target.
|
||||
- When no image reference exists (e.g. chart drawing): keeps id/name only.
|
||||
"""
|
||||
doc_pr = drawing_elem.find(".//wp:docPr", NS)
|
||||
attrs = {
|
||||
"id": doc_pr.get("id", "") if doc_pr is not None else "",
|
||||
"name": doc_pr.get("name", "") if doc_pr is not None else "",
|
||||
}
|
||||
|
||||
if include_extended_attrs:
|
||||
rel_ref = _extract_blip_relationship(drawing_elem)
|
||||
if rel_ref is not None and context is not None:
|
||||
rel_kind, rel_id = rel_ref
|
||||
rel = context.resolve_relationship(rel_id)
|
||||
if rel is not None:
|
||||
if rel_kind == "embed" and rel.rel_type == IMAGE_REL_TYPE:
|
||||
rel_path = context.export_embedded_image(rel)
|
||||
if rel_path:
|
||||
attrs["path"] = rel_path
|
||||
if rel.image_format:
|
||||
attrs["format"] = rel.image_format
|
||||
elif rel_kind == "link":
|
||||
if rel.target:
|
||||
attrs["path"] = rel.target
|
||||
if rel.image_format:
|
||||
attrs["format"] = rel.image_format
|
||||
|
||||
return _build_placeholder(attrs)
|
||||
|
||||
|
||||
def extract_vml_image_placeholder_from_element(
|
||||
container_elem,
|
||||
context: Optional[DrawingExtractionContext] = None,
|
||||
include_extended_attrs: bool = True,
|
||||
) -> str:
|
||||
"""
|
||||
Build a <drawing ... /> placeholder from a w:pict or w:object element.
|
||||
|
||||
Legacy Word documents and OLE-embedded objects (Visio diagrams, equation
|
||||
editor previews, etc.) expose their rendered image via VML rather than
|
||||
DrawingML. The image is referenced through ``<v:imagedata r:id="..."/>``
|
||||
inside ``<v:shape>``, and the underlying bytes are commonly EMF/WMF
|
||||
metafiles. This function exports those bytes through the same context as
|
||||
DrawingML images so EMF/WMF assets land in the blocks.assets directory
|
||||
alongside PNG/JPEG ones.
|
||||
|
||||
The output placeholder format matches
|
||||
``extract_drawing_placeholder_from_element`` so downstream consumers
|
||||
treat both paths uniformly.
|
||||
"""
|
||||
shape = container_elem.find(".//v:shape", NS)
|
||||
attrs = {
|
||||
"id": shape.get("id", "") if shape is not None else "",
|
||||
"name": shape.get("alt", "") if shape is not None else "",
|
||||
}
|
||||
|
||||
if include_extended_attrs:
|
||||
rel_id = _extract_imagedata_relationship(container_elem)
|
||||
if rel_id and context is not None:
|
||||
rel = context.resolve_relationship(rel_id)
|
||||
if rel is not None and rel.rel_type == IMAGE_REL_TYPE:
|
||||
# VML reuses r:id for both embedded image parts and externally
|
||||
# linked images; only the resolved TargetMode tells us which.
|
||||
# Treating an external relationship as embedded would call
|
||||
# export_embedded_image() (which short-circuits on external)
|
||||
# and silently drop the linked path.
|
||||
if rel.target_mode.lower() == "external":
|
||||
if rel.target:
|
||||
attrs["path"] = rel.target
|
||||
if rel.image_format:
|
||||
attrs["format"] = rel.image_format
|
||||
else:
|
||||
rel_path = context.export_embedded_image(rel)
|
||||
if rel_path:
|
||||
attrs["path"] = rel_path
|
||||
if rel.image_format:
|
||||
attrs["format"] = rel.image_format
|
||||
|
||||
return _build_placeholder(attrs)
|
||||
|
||||
|
||||
def parse_drawing_attributes(placeholder: str) -> Dict[str, str]:
|
||||
"""Parse attributes from a <drawing ... /> placeholder."""
|
||||
return {
|
||||
name: unescape(value)
|
||||
for name, value in DRAWING_ATTR_PATTERN.findall(placeholder)
|
||||
}
|
||||
|
||||
|
||||
def normalize_drawing_placeholder(
|
||||
placeholder: str,
|
||||
include_extended_attrs: bool = False,
|
||||
) -> str:
|
||||
"""
|
||||
Normalize one drawing placeholder into canonical attribute order.
|
||||
|
||||
Args:
|
||||
placeholder: Input placeholder string
|
||||
include_extended_attrs: If False, keeps only id/name.
|
||||
"""
|
||||
attrs = parse_drawing_attributes(placeholder)
|
||||
normalized = {
|
||||
"id": attrs.get("id", ""),
|
||||
"name": attrs.get("name", ""),
|
||||
}
|
||||
if include_extended_attrs:
|
||||
if "path" in attrs:
|
||||
normalized["path"] = attrs["path"]
|
||||
if "format" in attrs:
|
||||
normalized["format"] = attrs["format"]
|
||||
for key, value in attrs.items():
|
||||
if key not in {"id", "name", "path", "format"}:
|
||||
normalized[key] = value
|
||||
return _build_placeholder(normalized)
|
||||
|
||||
|
||||
def normalize_drawing_placeholders_in_text(
|
||||
text: str,
|
||||
include_extended_attrs: bool = False,
|
||||
) -> str:
|
||||
"""Normalize all drawing placeholders inside a text blob."""
|
||||
if not text:
|
||||
return text
|
||||
|
||||
def _replace(match: re.Match) -> str:
|
||||
return normalize_drawing_placeholder(
|
||||
match.group(0),
|
||||
include_extended_attrs=include_extended_attrs,
|
||||
)
|
||||
|
||||
return DRAWING_TAG_PATTERN.sub(_replace, text)
|
||||
@@ -0,0 +1,339 @@
|
||||
"""Native DOCX IR builder: ``extract_docx_blocks`` output → :class:`IRDoc`.
|
||||
|
||||
Input contract: a list of block dicts as produced by
|
||||
``lightrag.parser.docx.parse_document.extract_docx_blocks``. Each
|
||||
block carries ``content`` text in which ``<table>``, ``<equation>`` and
|
||||
``<drawing …/>`` placeholders are already embedded by the upstream parser.
|
||||
The builder rewrites those placeholders into IR placeholder tokens
|
||||
(``{{TBL:k}} / {{EQ:k}} / {{EQI:k}} / {{IMG:k}}``) and builds the matching
|
||||
``IRTable`` / ``IREquation`` / ``IRDrawing`` items.
|
||||
|
||||
Asset bytes are extracted to disk by the upstream parser *before* this
|
||||
builder runs (via ``DrawingExtractionContext`` passed to
|
||||
``extract_docx_blocks``). The builder therefore declares assets with
|
||||
``AssetSpec.source=None`` — the writer records each entry's size without
|
||||
copying.
|
||||
|
||||
Block-vs-inline equation distinction follows the legacy native rule: an
|
||||
``<equation>…</equation>`` tag is *block* iff each side is either the
|
||||
content boundary or a ``\\n`` character. Anything else stays inline,
|
||||
keeps its tag in block text without an id, and never enters
|
||||
``equations.json``.
|
||||
|
||||
Positions are always emitted as ``IRPosition(type="paraid", range=[start,
|
||||
end])`` where each side may be ``None`` (legacy / non-Word docx authors
|
||||
sometimes omit ``w14:paraId``). The writer's ``to_jsonable`` faithfully
|
||||
preserves the per-side null so consumers can distinguish "start missing"
|
||||
vs "both missing".
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import itertools
|
||||
import json
|
||||
import re
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path, PurePosixPath
|
||||
from typing import Any
|
||||
|
||||
from lightrag.parser.docx.drawing_image_extractor import (
|
||||
DRAWING_TAG_PATTERN,
|
||||
parse_drawing_attributes,
|
||||
)
|
||||
from lightrag.sidecar.ir import (
|
||||
AssetSpec,
|
||||
IRBlock,
|
||||
IRDoc,
|
||||
IRDrawing,
|
||||
IREquation,
|
||||
IRPosition,
|
||||
IRTable,
|
||||
)
|
||||
|
||||
|
||||
_TABLE_TAG_RE = re.compile(r"<table>(.*?)</table>", re.DOTALL)
|
||||
_EQUATION_TAG_RE = re.compile(r"<equation>(.*?)</equation>", re.DOTALL)
|
||||
|
||||
|
||||
def _normalize_dimension(rows_value: Any) -> tuple[int, int]:
|
||||
if not isinstance(rows_value, list):
|
||||
return 0, 0
|
||||
num_rows = len(rows_value)
|
||||
num_cols = max((len(r) for r in rows_value if isinstance(r, list)), default=0)
|
||||
return num_rows, num_cols
|
||||
|
||||
|
||||
def _placeholder_keyspace() -> Callable[[str], str]:
|
||||
"""Return a fresh counter producing ``{prefix}{N}`` keys (1-indexed)."""
|
||||
counter = itertools.count(1)
|
||||
return lambda prefix: f"{prefix}{next(counter)}"
|
||||
|
||||
|
||||
def _safe_asset_ref_from_path(path_val: str, asset_prefix: str) -> str | None:
|
||||
"""Return the path inside ``asset_prefix`` only when it is safe.
|
||||
|
||||
Native DOCX images are pre-extracted into ``<base>.blocks.assets/``.
|
||||
Treat a drawing path as local only when the suffix is a clean POSIX
|
||||
relative path. Unsafe local-looking paths are dropped instead of being
|
||||
registered as assets or preserved as linked references.
|
||||
"""
|
||||
if not asset_prefix or not path_val.startswith(asset_prefix):
|
||||
return None
|
||||
|
||||
rel_raw = path_val[len(asset_prefix) :]
|
||||
if not rel_raw or "\\" in rel_raw:
|
||||
return None
|
||||
|
||||
rel_path = PurePosixPath(rel_raw)
|
||||
if rel_path.is_absolute():
|
||||
return None
|
||||
if any(part == ".." for part in rel_path.parts):
|
||||
return None
|
||||
|
||||
rel = rel_path.as_posix()
|
||||
if rel in {"", "."}:
|
||||
return None
|
||||
return rel
|
||||
|
||||
|
||||
@dataclass
|
||||
class _BlockBuilder:
|
||||
"""Per-block scratch state for the three ``re.sub`` rewrite passes.
|
||||
|
||||
Keeping the replacer routines as bound methods (rather than closures
|
||||
redefined inside the per-block loop) means they're compiled once at
|
||||
class-load and the state they mutate — ``tables`` / ``drawings`` /
|
||||
``equations`` / ``table_position`` — is held explicitly rather than
|
||||
captured implicitly from the enclosing frame.
|
||||
"""
|
||||
|
||||
next_key: Callable[[str], str]
|
||||
assets: list[AssetSpec]
|
||||
seen_asset_refs: set[str]
|
||||
asset_prefix: str
|
||||
block_table_headers: list[Any]
|
||||
tables: list[IRTable] = field(default_factory=list)
|
||||
drawings: list[IRDrawing] = field(default_factory=list)
|
||||
equations: list[IREquation] = field(default_factory=list)
|
||||
# Position of the *next* ``<table>`` placeholder within this block,
|
||||
# used to look up the matching entry in ``block_table_headers``.
|
||||
table_position: int = 0
|
||||
|
||||
def replace_table(self, match: "re.Match[str]") -> str:
|
||||
table_body_raw = match.group(1)
|
||||
try:
|
||||
rows = json.loads(table_body_raw)
|
||||
if not isinstance(rows, list):
|
||||
rows = None
|
||||
except json.JSONDecodeError:
|
||||
rows = None
|
||||
|
||||
if rows is not None:
|
||||
parsed_rows: list[list[str]] | None = [
|
||||
[str(c) for c in r] if isinstance(r, list) else [str(r)] for r in rows
|
||||
]
|
||||
html: str | None = None
|
||||
else:
|
||||
parsed_rows = None
|
||||
html = table_body_raw
|
||||
|
||||
num_rows, num_cols = _normalize_dimension(parsed_rows)
|
||||
|
||||
header_pos = self.table_position
|
||||
self.table_position += 1
|
||||
header_rows = (
|
||||
self.block_table_headers[header_pos]
|
||||
if header_pos < len(self.block_table_headers)
|
||||
else None
|
||||
)
|
||||
# Treat empty list / explicit None identically: no header
|
||||
# entry on the sidecar item.
|
||||
table_header = header_rows if header_rows else None
|
||||
|
||||
placeholder = self.next_key("tb")
|
||||
self.tables.append(
|
||||
IRTable(
|
||||
placeholder_key=placeholder,
|
||||
rows=parsed_rows,
|
||||
html=html,
|
||||
num_rows=num_rows,
|
||||
num_cols=num_cols,
|
||||
caption="",
|
||||
footnotes=[],
|
||||
table_header=table_header,
|
||||
body_override=table_body_raw,
|
||||
)
|
||||
)
|
||||
return f"{{{{TBL:{placeholder}}}}}"
|
||||
|
||||
def replace_equation(self, match: "re.Match[str]") -> str:
|
||||
latex = match.group(1)
|
||||
source = match.string
|
||||
start, end = match.start(), match.end()
|
||||
is_block = (start == 0 or source[start - 1] == "\n") and (
|
||||
end == len(source) or source[end] == "\n"
|
||||
)
|
||||
placeholder = self.next_key("eq")
|
||||
self.equations.append(
|
||||
IREquation(
|
||||
placeholder_key=placeholder,
|
||||
latex=latex,
|
||||
is_block=is_block,
|
||||
caption="",
|
||||
footnotes=[],
|
||||
)
|
||||
)
|
||||
token = "EQ" if is_block else "EQI"
|
||||
return f"{{{{{token}:{placeholder}}}}}"
|
||||
|
||||
def replace_drawing(self, match: "re.Match[str]") -> str:
|
||||
attrs = parse_drawing_attributes(match.group(0))
|
||||
path_val = attrs.get("path", "") or ""
|
||||
src_val = attrs.get("src", "") or ""
|
||||
fmt = attrs.get("format", "") or ""
|
||||
if not fmt and path_val:
|
||||
fmt = Path(path_val).suffix.lower().lstrip(".")
|
||||
|
||||
# Two flavours of <drawing path="…">:
|
||||
# 1. Local asset under <base>.blocks.assets/ — already
|
||||
# extracted to disk by DrawingExtractionContext;
|
||||
# register as AssetSpec(source=None) and let the
|
||||
# writer resolve the path via asset_paths.
|
||||
# 2. External/linked path (URL, or any path that does
|
||||
# not live under asset_prefix) — pass through
|
||||
# verbatim via IRDrawing.path_override; do NOT emit
|
||||
# an AssetSpec (no on-disk bytes to materialize).
|
||||
rel_inside_assets = _safe_asset_ref_from_path(path_val, self.asset_prefix)
|
||||
if rel_inside_assets is not None:
|
||||
asset_ref = rel_inside_assets
|
||||
suggested_name = Path(rel_inside_assets).name or rel_inside_assets
|
||||
if asset_ref and asset_ref not in self.seen_asset_refs:
|
||||
self.assets.append(
|
||||
AssetSpec(
|
||||
ref=asset_ref,
|
||||
suggested_name=suggested_name,
|
||||
source=None, # already extracted to disk
|
||||
)
|
||||
)
|
||||
self.seen_asset_refs.add(asset_ref)
|
||||
path_override: str | None = None
|
||||
else:
|
||||
asset_ref = ""
|
||||
# Only mark as an external/linked reference when the
|
||||
# upstream parser actually emitted a path. An empty
|
||||
# ``path=""`` should fall back to the regular asset-
|
||||
# resolution path (which will also produce ``path=""``
|
||||
# downstream) rather than masquerading as an explicit
|
||||
# builder override.
|
||||
path_override = (
|
||||
None
|
||||
if self.asset_prefix and path_val.startswith(self.asset_prefix)
|
||||
else path_val or None
|
||||
)
|
||||
|
||||
placeholder = self.next_key("im")
|
||||
self.drawings.append(
|
||||
IRDrawing(
|
||||
placeholder_key=placeholder,
|
||||
asset_ref=asset_ref,
|
||||
fmt=fmt,
|
||||
caption="",
|
||||
footnotes=[],
|
||||
src=src_val,
|
||||
path_override=path_override,
|
||||
)
|
||||
)
|
||||
return f"{{{{IMG:{placeholder}}}}}"
|
||||
|
||||
|
||||
class NativeDocxIRBuilder:
|
||||
"""Translate ``extract_docx_blocks`` output into an :class:`IRDoc`.
|
||||
|
||||
The builder is stateless — instantiate per call. ``asset_dir_name`` is
|
||||
the relative name (without trailing slash) of ``<base>.blocks.assets/``
|
||||
that the upstream parser used when emitting ``<drawing path>``
|
||||
attributes; the builder strips that prefix when building
|
||||
:attr:`AssetSpec.ref` so the writer's ref↔filename mapping has
|
||||
predictable keys.
|
||||
"""
|
||||
|
||||
def normalize(
|
||||
self,
|
||||
blocks: list[dict[str, Any]],
|
||||
*,
|
||||
document_name: str,
|
||||
asset_dir_name: str,
|
||||
parse_metadata: dict[str, Any] | None = None,
|
||||
) -> IRDoc:
|
||||
next_key = _placeholder_keyspace()
|
||||
ir_blocks: list[IRBlock] = []
|
||||
assets: list[AssetSpec] = []
|
||||
seen_asset_refs: set[str] = set()
|
||||
|
||||
asset_prefix = f"{asset_dir_name}/" if asset_dir_name else ""
|
||||
|
||||
for block in blocks:
|
||||
raw_content = block.get("content") or ""
|
||||
heading = block.get("heading") or ""
|
||||
level = int(block.get("level", 0) or 0)
|
||||
parent_headings = list(block.get("parent_headings") or [])
|
||||
# Preserve per-side nulls in [start, end].
|
||||
uuid_start = block.get("uuid") or None
|
||||
uuid_end = block.get("uuid_end") or None
|
||||
|
||||
builder = _BlockBuilder(
|
||||
next_key=next_key,
|
||||
assets=assets,
|
||||
seen_asset_refs=seen_asset_refs,
|
||||
asset_prefix=asset_prefix,
|
||||
block_table_headers=list(block.get("table_headers") or []),
|
||||
)
|
||||
|
||||
# Rewrite order matches the legacy native flow: tables, then
|
||||
# equations, then drawings — each ``re.sub`` operates on the
|
||||
# output of the previous pass.
|
||||
content_template = _TABLE_TAG_RE.sub(builder.replace_table, raw_content)
|
||||
content_template = _EQUATION_TAG_RE.sub(
|
||||
builder.replace_equation, content_template
|
||||
)
|
||||
content_template = DRAWING_TAG_PATTERN.sub(
|
||||
builder.replace_drawing, content_template
|
||||
)
|
||||
|
||||
positions = [
|
||||
IRPosition(type="paraid", range=[uuid_start, uuid_end]),
|
||||
]
|
||||
|
||||
ir_blocks.append(
|
||||
IRBlock(
|
||||
content_template=content_template,
|
||||
heading=heading,
|
||||
level=level,
|
||||
parent_headings=parent_headings,
|
||||
positions=positions,
|
||||
tables=builder.tables,
|
||||
drawings=builder.drawings,
|
||||
equations=builder.equations,
|
||||
)
|
||||
)
|
||||
|
||||
# doc_title: parse_metadata["first_heading"] when present, else file
|
||||
# stem fallback (resolved here so the writer doesn't have to know).
|
||||
first_heading = ""
|
||||
if isinstance(parse_metadata, dict):
|
||||
first_heading = str(parse_metadata.get("first_heading") or "")
|
||||
doc_title = first_heading or (Path(document_name).stem or document_name)
|
||||
|
||||
return IRDoc(
|
||||
document_name=document_name,
|
||||
document_format=Path(document_name).suffix.lower().lstrip("."),
|
||||
doc_title=doc_title,
|
||||
split_option={},
|
||||
blocks=ir_blocks,
|
||||
assets=assets,
|
||||
bbox_attributes=None,
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["NativeDocxIRBuilder"]
|
||||
@@ -0,0 +1,423 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
ABOUTME: Resolves automatic numbering labels from DOCX documents
|
||||
ABOUTME: Parses numbering.xml and computes rendered number strings
|
||||
"""
|
||||
|
||||
import zipfile
|
||||
from defusedxml import ElementTree as ET
|
||||
from typing import Dict
|
||||
|
||||
NSMAP = {"w": "http://schemas.openxmlformats.org/wordprocessingml/2006/main"}
|
||||
|
||||
|
||||
class NumberingResolver:
|
||||
"""
|
||||
Resolves paragraph numbering to rendered label strings.
|
||||
|
||||
DOCX stores numbering definitions in numbering.xml:
|
||||
- abstractNum: Defines format templates (lvlText like "%1.%2.")
|
||||
- num: Links numId to abstractNumId
|
||||
|
||||
Each paragraph references: numId (which definition) + ilvl (which level)
|
||||
"""
|
||||
|
||||
# Number format converters
|
||||
FORMAT_CONVERTERS = {
|
||||
"decimal": lambda n: str(n),
|
||||
"lowerLetter": lambda n: chr(ord("a") + (n - 1) % 26),
|
||||
"upperLetter": lambda n: chr(ord("A") + (n - 1) % 26),
|
||||
"lowerRoman": lambda n: NumberingResolver._to_roman(n).lower(),
|
||||
"upperRoman": lambda n: NumberingResolver._to_roman(n),
|
||||
"chineseCountingThousand": lambda n: NumberingResolver._to_chinese(n),
|
||||
"ideographTraditional": lambda n: "甲乙丙丁戊己庚辛壬癸"[(n - 1) % 10],
|
||||
"bullet": lambda n: "•",
|
||||
"none": lambda n: "",
|
||||
}
|
||||
|
||||
def __init__(self, docx_path: str):
|
||||
self.abstract_nums: Dict[str, dict] = {} # abstractNumId -> level definitions
|
||||
self.num_to_abstract: Dict[str, str] = {} # numId -> abstractNumId
|
||||
self.counters: Dict[
|
||||
str, Dict[int, int]
|
||||
] = {} # numId -> {ilvl -> current_count}
|
||||
self.start_overrides: Dict[
|
||||
str, Dict[int, int]
|
||||
] = {} # numId -> {ilvl -> start_value}
|
||||
self.style_numpr: Dict[
|
||||
str, dict
|
||||
] = {} # styleId -> {numId, ilvl} from styles.xml
|
||||
self.style_based_on: Dict[str, str] = {} # styleId -> basedOn styleId
|
||||
# Smart numbering merge state (Word's rendering behavior)
|
||||
self.last_numId: str = None # Previous paragraph's numId
|
||||
self.last_abstract_id: str = None # Previous paragraph's abstractNumId
|
||||
self.last_style_id: str = None # Previous paragraph's style ID
|
||||
self._parse_numbering_xml(docx_path)
|
||||
self._parse_styles_xml(docx_path)
|
||||
|
||||
def _parse_numbering_xml(self, docx_path: str):
|
||||
"""Parse numbering.xml from DOCX archive"""
|
||||
try:
|
||||
with zipfile.ZipFile(docx_path, "r") as zf:
|
||||
if "word/numbering.xml" not in zf.namelist():
|
||||
return
|
||||
|
||||
tree = ET.parse(zf.open("word/numbering.xml"))
|
||||
root = tree.getroot()
|
||||
|
||||
# Parse abstractNum definitions
|
||||
for abstract in root.findall(".//w:abstractNum", NSMAP):
|
||||
abstract_id = abstract.get(f"{{{NSMAP['w']}}}abstractNumId")
|
||||
levels = {}
|
||||
|
||||
for lvl in abstract.findall("w:lvl", NSMAP):
|
||||
ilvl = int(lvl.get(f"{{{NSMAP['w']}}}ilvl"))
|
||||
|
||||
start_elem = lvl.find("w:start", NSMAP)
|
||||
start = (
|
||||
int(start_elem.get(f"{{{NSMAP['w']}}}val"))
|
||||
if start_elem is not None
|
||||
else 1
|
||||
)
|
||||
|
||||
num_fmt_elem = lvl.find("w:numFmt", NSMAP)
|
||||
num_fmt = (
|
||||
num_fmt_elem.get(f"{{{NSMAP['w']}}}val")
|
||||
if num_fmt_elem is not None
|
||||
else "decimal"
|
||||
)
|
||||
|
||||
lvl_text_elem = lvl.find("w:lvlText", NSMAP)
|
||||
lvl_text = (
|
||||
lvl_text_elem.get(f"{{{NSMAP['w']}}}val")
|
||||
if lvl_text_elem is not None
|
||||
else "%1."
|
||||
)
|
||||
|
||||
is_lgl_elem = lvl.find("w:isLgl", NSMAP)
|
||||
is_lgl = False
|
||||
if is_lgl_elem is not None:
|
||||
val = is_lgl_elem.get(f"{{{NSMAP['w']}}}val")
|
||||
is_lgl = val is None or val not in ("0", "false")
|
||||
|
||||
levels[ilvl] = {
|
||||
"start": start,
|
||||
"numFmt": num_fmt,
|
||||
"lvlText": lvl_text,
|
||||
"isLgl": is_lgl,
|
||||
}
|
||||
|
||||
self.abstract_nums[abstract_id] = levels
|
||||
|
||||
# Parse num -> abstractNum mapping and startOverride
|
||||
for num in root.findall(".//w:num", NSMAP):
|
||||
num_id = num.get(f"{{{NSMAP['w']}}}numId")
|
||||
abstract_ref = num.find("w:abstractNumId", NSMAP)
|
||||
if abstract_ref is not None:
|
||||
self.num_to_abstract[num_id] = abstract_ref.get(
|
||||
f"{{{NSMAP['w']}}}val"
|
||||
)
|
||||
|
||||
# Parse lvlOverride/startOverride for this num
|
||||
for lvl_override in num.findall("w:lvlOverride", NSMAP):
|
||||
ilvl = int(lvl_override.get(f"{{{NSMAP['w']}}}ilvl"))
|
||||
start_override = lvl_override.find("w:startOverride", NSMAP)
|
||||
if start_override is not None:
|
||||
start_val = int(start_override.get(f"{{{NSMAP['w']}}}val"))
|
||||
if num_id not in self.start_overrides:
|
||||
self.start_overrides[num_id] = {}
|
||||
self.start_overrides[num_id][ilvl] = start_val
|
||||
except Exception:
|
||||
# Silently ignore parsing errors - document may not have numbering
|
||||
pass
|
||||
|
||||
def _parse_styles_xml(self, docx_path: str):
|
||||
"""Parse styles.xml to get style-inherited numbering definitions"""
|
||||
try:
|
||||
with zipfile.ZipFile(docx_path, "r") as zf:
|
||||
if "word/styles.xml" not in zf.namelist():
|
||||
return
|
||||
|
||||
tree = ET.parse(zf.open("word/styles.xml"))
|
||||
root = tree.getroot()
|
||||
|
||||
# Parse style definitions
|
||||
for style in root.findall(".//w:style", NSMAP):
|
||||
style_id = style.get(f"{{{NSMAP['w']}}}styleId")
|
||||
if not style_id:
|
||||
continue
|
||||
|
||||
# Check for basedOn (style inheritance)
|
||||
based_on = style.find("w:basedOn", NSMAP)
|
||||
if based_on is not None:
|
||||
parent_id = based_on.get(f"{{{NSMAP['w']}}}val")
|
||||
if parent_id:
|
||||
self.style_based_on[style_id] = parent_id
|
||||
|
||||
# Check for numPr in style's pPr
|
||||
pPr = style.find("w:pPr", NSMAP)
|
||||
if pPr is not None:
|
||||
numPr = pPr.find("w:numPr", NSMAP)
|
||||
if numPr is not None:
|
||||
num_id_elem = numPr.find("w:numId", NSMAP)
|
||||
ilvl_elem = numPr.find("w:ilvl", NSMAP)
|
||||
|
||||
if num_id_elem is not None:
|
||||
num_id = num_id_elem.get(f"{{{NSMAP['w']}}}val")
|
||||
ilvl = (
|
||||
int(ilvl_elem.get(f"{{{NSMAP['w']}}}val"))
|
||||
if ilvl_elem is not None
|
||||
else 0
|
||||
)
|
||||
self.style_numpr[style_id] = {
|
||||
"numId": num_id,
|
||||
"ilvl": ilvl,
|
||||
}
|
||||
except Exception:
|
||||
# Silently ignore parsing errors
|
||||
pass
|
||||
|
||||
def _get_numbering_from_style(self, style_id: str, visited=None) -> dict:
|
||||
"""
|
||||
Get numbering definition from style, following inheritance chain.
|
||||
|
||||
Args:
|
||||
style_id: Style ID to look up
|
||||
visited: Set of visited style IDs (to prevent circular references)
|
||||
|
||||
Returns:
|
||||
dict with 'numId' and 'ilvl', or None
|
||||
"""
|
||||
if visited is None:
|
||||
visited = set()
|
||||
|
||||
# Prevent circular references
|
||||
if style_id in visited:
|
||||
return None
|
||||
visited.add(style_id)
|
||||
|
||||
# Check if this style has numPr
|
||||
if style_id in self.style_numpr:
|
||||
return self.style_numpr[style_id]
|
||||
|
||||
# Check parent style
|
||||
if style_id in self.style_based_on:
|
||||
parent_id = self.style_based_on[style_id]
|
||||
return self._get_numbering_from_style(parent_id, visited)
|
||||
|
||||
return None
|
||||
|
||||
def reset_tracking_state(self):
|
||||
"""
|
||||
Reset numbering tracking state.
|
||||
|
||||
Call this when encountering structural breaks that should
|
||||
interrupt numbering continuity:
|
||||
- Section breaks (sectPr)
|
||||
- Table boundaries (before and after tables)
|
||||
|
||||
This prevents incorrect numbering continuation across
|
||||
document structure boundaries.
|
||||
"""
|
||||
self.last_numId = None
|
||||
self.last_abstract_id = None
|
||||
self.last_style_id = None
|
||||
|
||||
def get_label(self, para_element) -> str:
|
||||
"""
|
||||
Get rendered numbering label for a paragraph.
|
||||
|
||||
Checks both direct numPr and style-inherited numbering. Direct numPr
|
||||
is a paragraph-local override and applies only to the current
|
||||
paragraph; subsequent paragraphs that carry only pStyle fall back to
|
||||
the style's numPr declared in styles.xml.
|
||||
|
||||
Args:
|
||||
para_element: lxml Element for <w:p>
|
||||
|
||||
Returns:
|
||||
Rendered label string (e.g., "1.1", "a)", "第一章") or empty string
|
||||
"""
|
||||
try:
|
||||
pPr = para_element.find(f"{{{NSMAP['w']}}}pPr")
|
||||
if pPr is None:
|
||||
return ""
|
||||
|
||||
num_id = None
|
||||
ilvl = 0
|
||||
style_id = None
|
||||
|
||||
# Get pStyle (if present)
|
||||
pStyle = pPr.find(f"{{{NSMAP['w']}}}pStyle")
|
||||
if pStyle is not None:
|
||||
style_id = pStyle.get(f"{{{NSMAP['w']}}}val")
|
||||
|
||||
# Check for direct numPr in paragraph
|
||||
numPr = pPr.find(f"{{{NSMAP['w']}}}numPr")
|
||||
if numPr is not None:
|
||||
num_id_elem = numPr.find(f"{{{NSMAP['w']}}}numId")
|
||||
ilvl_elem = numPr.find(f"{{{NSMAP['w']}}}ilvl")
|
||||
|
||||
if num_id_elem is not None:
|
||||
num_id = num_id_elem.get(f"{{{NSMAP['w']}}}val")
|
||||
ilvl = (
|
||||
int(ilvl_elem.get(f"{{{NSMAP['w']}}}val"))
|
||||
if ilvl_elem is not None
|
||||
else 0
|
||||
)
|
||||
|
||||
# If no direct numPr, fall back to style-inherited numbering.
|
||||
# Direct numPr is a paragraph-local override in Word; it must not
|
||||
# persist as a runtime default for the style, otherwise subsequent
|
||||
# paragraphs that only carry pStyle will keep following the local
|
||||
# override instead of the style's declared numPr.
|
||||
if num_id is None and style_id:
|
||||
style_num = self._get_numbering_from_style(style_id)
|
||||
if style_num:
|
||||
num_id = style_num["numId"]
|
||||
ilvl = style_num["ilvl"]
|
||||
|
||||
# If still no numbering found, clear state and return empty
|
||||
if num_id is None:
|
||||
# We should use list structure breaking logic to reset last_numId, last_abstract_id and last_style_id
|
||||
return ""
|
||||
|
||||
# Get abstract definition
|
||||
abstract_id = self.num_to_abstract.get(num_id)
|
||||
if abstract_id is None or abstract_id not in self.abstract_nums:
|
||||
# Clear state for invalid numbering
|
||||
self.last_numId = None
|
||||
self.last_abstract_id = None
|
||||
return ""
|
||||
|
||||
levels = self.abstract_nums[abstract_id]
|
||||
if ilvl not in levels:
|
||||
# Clear state for invalid level
|
||||
self.last_numId = None
|
||||
self.last_abstract_id = None
|
||||
return ""
|
||||
|
||||
# Smart numbering merge: (Word's rendering behavior)
|
||||
# When consecutive paragraphs have different numId but same abstractNumId,
|
||||
# Word continues the numbering sequence rather than restarting.
|
||||
# This happens regardless of whether the numId is new or style matches.
|
||||
|
||||
if (
|
||||
self.last_numId is not None
|
||||
and self.last_numId != num_id
|
||||
and self.last_abstract_id == abstract_id
|
||||
and self.last_numId in self.counters
|
||||
):
|
||||
# Merge: copy previous numId's counter to current numId
|
||||
self.counters[num_id] = self.counters[self.last_numId].copy()
|
||||
|
||||
# Initialize/update counter
|
||||
if num_id not in self.counters:
|
||||
self.counters[num_id] = {}
|
||||
|
||||
# Initialize all parent levels if not present (for deep nested numbering)
|
||||
for i in range(ilvl):
|
||||
if i not in self.counters[num_id] and i in levels:
|
||||
# Use startOverride if exists, otherwise use abstractNum's start value
|
||||
if (
|
||||
num_id in self.start_overrides
|
||||
and i in self.start_overrides[num_id]
|
||||
):
|
||||
self.counters[num_id][i] = self.start_overrides[num_id][i]
|
||||
else:
|
||||
self.counters[num_id][i] = levels[i]["start"]
|
||||
|
||||
# Reset lower levels when higher level increments
|
||||
for i in range(ilvl + 1, 10):
|
||||
if i in self.counters[num_id]:
|
||||
del self.counters[num_id][i]
|
||||
|
||||
# Initialize current level if needed
|
||||
if ilvl not in self.counters[num_id]:
|
||||
# Use startOverride if exists, otherwise use abstractNum's start value
|
||||
if (
|
||||
num_id in self.start_overrides
|
||||
and ilvl in self.start_overrides[num_id]
|
||||
):
|
||||
self.counters[num_id][ilvl] = self.start_overrides[num_id][ilvl]
|
||||
else:
|
||||
self.counters[num_id][ilvl] = levels[ilvl]["start"]
|
||||
else:
|
||||
self.counters[num_id][ilvl] += 1
|
||||
|
||||
# Format the label using lvlText template
|
||||
label = self._format_label(num_id, ilvl, levels)
|
||||
|
||||
# Update tracking state for next paragraph
|
||||
self.last_numId = num_id
|
||||
self.last_abstract_id = abstract_id
|
||||
self.last_style_id = style_id
|
||||
|
||||
return label
|
||||
except Exception:
|
||||
# Return empty on any error to avoid breaking document parsing
|
||||
return ""
|
||||
|
||||
def _format_label(self, num_id: str, ilvl: int, levels: dict) -> str:
|
||||
"""Format label string by replacing %1, %2, etc."""
|
||||
try:
|
||||
lvl_text = levels[ilvl]["lvlText"]
|
||||
result = lvl_text
|
||||
current_is_lgl = levels[ilvl].get("isLgl", False)
|
||||
|
||||
for i in range(ilvl + 1):
|
||||
if i in levels and i in self.counters.get(num_id, {}):
|
||||
num_fmt = levels[i]["numFmt"]
|
||||
if current_is_lgl and i < ilvl:
|
||||
num_fmt = "decimal"
|
||||
count = self.counters[num_id][i]
|
||||
converter = self.FORMAT_CONVERTERS.get(num_fmt, lambda n: str(n))
|
||||
formatted = converter(count)
|
||||
result = result.replace(f"%{i + 1}", formatted)
|
||||
|
||||
return result
|
||||
except Exception:
|
||||
return ""
|
||||
|
||||
@staticmethod
|
||||
def _to_roman(n: int) -> str:
|
||||
"""Convert integer to Roman numeral"""
|
||||
if n <= 0 or n >= 4000:
|
||||
return str(n)
|
||||
values = [
|
||||
(1000, "M"),
|
||||
(900, "CM"),
|
||||
(500, "D"),
|
||||
(400, "CD"),
|
||||
(100, "C"),
|
||||
(90, "XC"),
|
||||
(50, "L"),
|
||||
(40, "XL"),
|
||||
(10, "X"),
|
||||
(9, "IX"),
|
||||
(5, "V"),
|
||||
(4, "IV"),
|
||||
(1, "I"),
|
||||
]
|
||||
result = ""
|
||||
for value, numeral in values:
|
||||
while n >= value:
|
||||
result += numeral
|
||||
n -= value
|
||||
return result
|
||||
|
||||
@staticmethod
|
||||
def _to_chinese(n: int) -> str:
|
||||
"""Convert integer to Chinese numeral"""
|
||||
digits = "零一二三四五六七八九"
|
||||
if n <= 0 or n > 99:
|
||||
return str(n)
|
||||
if n < 10:
|
||||
return digits[n]
|
||||
if n < 20:
|
||||
return "十" + (digits[n % 10] if n % 10 else "")
|
||||
if n < 100:
|
||||
tens = n // 10
|
||||
ones = n % 10
|
||||
return digits[tens] + "十" + (digits[ones] if ones else "")
|
||||
return str(n)
|
||||
@@ -0,0 +1,10 @@
|
||||
"""
|
||||
ABOUTME: OMML (Office Math Markup Language) to LaTeX conversion
|
||||
"""
|
||||
|
||||
from .ommlparser import OMMLParser
|
||||
|
||||
|
||||
def convert_omml_to_latex(omml_element) -> str:
|
||||
"""Convert an m:oMath XML element to a LaTeX string."""
|
||||
return OMMLParser().parse(omml_element)
|
||||
@@ -0,0 +1,38 @@
|
||||
"""
|
||||
Postprocessing functions for cleaning up latex equations in linear format which don't give valid LaTeX.
|
||||
"""
|
||||
|
||||
import re
|
||||
|
||||
clean_exps = {
|
||||
r"\\degf": "°F",
|
||||
r"\\degc": "°C",
|
||||
r"(\\cbrt)(\w+)": r"\\sqrt[3]{\2}",
|
||||
r"(\\qdrt)(\w+)": r"\\sqrt[4]{\2}",
|
||||
r"\\sfrac": r"\\frac",
|
||||
r"(\\o[i]+nt)(\w+)": r"\1{\2}",
|
||||
r"\\bullet(\w+)": r"\\bullet \1",
|
||||
r"\\sum([a-zA-Z0-9]+)": r"\\sum{\1}",
|
||||
r"\\prod([a-zA-Z0-9]+)": r"\\prod{\1}",
|
||||
r"\\amalg([a-zA-Z0-9]+)": r"\\amalg{\1}",
|
||||
r"\\bigcup([a-zA-Z0-9]+)": r"\\bigcup{\1}",
|
||||
r"\\bigcap([a-zA-Z0-9]+)": r"\\bigcap{\1}",
|
||||
r"\\bigvee([a-zA-Z0-9]+)": r"\\bigvee{\1}",
|
||||
r"\\bigwedge([a-zA-Z0-9]+)": r"\\bigwedge{\1}",
|
||||
r"\\lfloor([a-zA-Z0-9]+)": r"\\lfloor{\1}",
|
||||
r"\\lceil([a-zA-Z0-9]+)": r"\\lceil{\1}",
|
||||
r"\\lim\\below\{(.+)\}\{(.+)\}": r"\\lim_{\1}{\2}",
|
||||
r"\\min\\below\{(.+)\}\{(.+)\}": r"\\min_{\1}{\2}",
|
||||
r"\\max\\below\{(.+)\}\{(.+)\}": r"\\max_{\1}{\2}",
|
||||
}
|
||||
|
||||
|
||||
def clean_exp(exp):
|
||||
"""
|
||||
Takes in a linear expression and converts known invalid LaTeX equations to valid LaTeX
|
||||
:param exp:str - An equation in invalid syntax
|
||||
:return :str - A valid equation
|
||||
"""
|
||||
for e in clean_exps:
|
||||
exp = re.sub(e, clean_exps[e], exp)
|
||||
return exp
|
||||
@@ -0,0 +1,511 @@
|
||||
from xml.etree.cElementTree import Element
|
||||
|
||||
from .utils import qn
|
||||
|
||||
|
||||
class OMMLParser:
|
||||
"""
|
||||
Parser class for reading OMML and converting it into LaTeX.
|
||||
"""
|
||||
|
||||
FUNCTION_MAP = {
|
||||
"sin": "\\sin",
|
||||
"cos": "\\cos",
|
||||
"tan": "\\tan",
|
||||
"cot": "\\cot",
|
||||
"sec": "\\sec",
|
||||
"csc": "\\csc",
|
||||
"sinh": "\\sinh",
|
||||
"cosh": "\\cosh",
|
||||
"tanh": "\\tanh",
|
||||
"coth": "\\coth",
|
||||
"sech": "\\operatorname{sech}",
|
||||
"csch": "\\operatorname{csch}",
|
||||
"log": "\\log",
|
||||
"ln": "\\ln",
|
||||
"min": "\\min",
|
||||
"max": "\\max",
|
||||
"lim": "\\lim",
|
||||
}
|
||||
|
||||
def _normalize_func_name(self, content: str) -> str:
|
||||
if not content:
|
||||
return content
|
||||
if content.startswith("\\"):
|
||||
return content
|
||||
key = content.strip()
|
||||
mapped = self.FUNCTION_MAP.get(key)
|
||||
return mapped if mapped else content
|
||||
|
||||
def parse(self, root: Element) -> str:
|
||||
"""
|
||||
Parses an m:oMath OMML tag into LaTeX.
|
||||
:param root: An m:oMath OMML tag
|
||||
:return: The LaTeX representation of the OMML input
|
||||
"""
|
||||
text = ""
|
||||
try:
|
||||
if root.tag == qn("m:t"):
|
||||
return self.parse_t(root)
|
||||
for child in root:
|
||||
if child.tag in self.parsers:
|
||||
text += self.parsers[child.tag](self, child)
|
||||
except AttributeError:
|
||||
# In case of missing attributes on OMML tags,
|
||||
# we return an empty string (ref:issue_14)
|
||||
return ""
|
||||
return text
|
||||
|
||||
def parse_e(self, root: Element) -> str:
|
||||
text = ""
|
||||
for child in root:
|
||||
text += self.parse(child)
|
||||
return text
|
||||
|
||||
def parse_r(self, root: Element) -> str:
|
||||
# TODO: Add support for m:rPr and m:scr to support different character styles
|
||||
# For now, we just parse the text content of m:r
|
||||
text = ""
|
||||
for child in root:
|
||||
text += self.parse(child)
|
||||
return text
|
||||
|
||||
def parse_t(self, root: Element):
|
||||
symbol_map = {
|
||||
"≜": "\\triangleq",
|
||||
"≝": "\\stackrel{\\tiny def}{=}",
|
||||
"≞": "\\stackrel{\\tiny m}{=}",
|
||||
}
|
||||
replacements = {
|
||||
"<": "\\lt ",
|
||||
">": "\\gt ",
|
||||
"≤": "\\leq ",
|
||||
"≥": "\\geq ",
|
||||
"∞": "\\infty ",
|
||||
"<": "\\lt ",
|
||||
">": "\\gt ",
|
||||
"≤": "\\leq ",
|
||||
"≥": "\\geq ",
|
||||
}
|
||||
text = root.text.split()
|
||||
if not text:
|
||||
return " "
|
||||
for i, t in enumerate(text):
|
||||
if t in symbol_map:
|
||||
text[i] = symbol_map[t]
|
||||
for key, value in replacements.items():
|
||||
for i, t in enumerate(text):
|
||||
text[i] = t.replace(key, value)
|
||||
return " ".join(text)
|
||||
|
||||
def parse_acc(self, root: Element) -> str:
|
||||
character_map = {
|
||||
768: "\\grave",
|
||||
769: "\\acute",
|
||||
770: "\\hat",
|
||||
771: "\\tilde",
|
||||
773: "\\bar",
|
||||
774: "\\breve",
|
||||
775: "\\dot",
|
||||
776: "\\ddot",
|
||||
780: "\\check",
|
||||
831: "\\overline{\\overline",
|
||||
8400: "\\overset\\leftharpoonup",
|
||||
8401: "\\overset\\rightharpoonup",
|
||||
8406: "\\overleftarrow",
|
||||
8407: "\\overrightarrow",
|
||||
8411: "\\dddot",
|
||||
8417: "\\overset\\leftrightarrow",
|
||||
}
|
||||
text = ""
|
||||
accent = 770
|
||||
for child in root:
|
||||
if child.tag == qn("m:accPr"):
|
||||
for child2 in child:
|
||||
if child2.tag == qn("m:chr"):
|
||||
val = child2.attrib.get(qn("m:val"))
|
||||
if val:
|
||||
try:
|
||||
accent = ord(val)
|
||||
except TypeError:
|
||||
pass
|
||||
|
||||
accent_cmd = character_map.get(accent)
|
||||
if accent_cmd is None:
|
||||
accent_cmd = character_map.get(770, "\\hat")
|
||||
text += accent_cmd + "{"
|
||||
for child in root:
|
||||
if child.tag == qn("m:e"):
|
||||
text += self.parse(child)
|
||||
text += "}"
|
||||
if accent == 831:
|
||||
text += "}"
|
||||
return text
|
||||
|
||||
def parse_bar(self, root: Element) -> str:
|
||||
text = "\\overline{"
|
||||
for child in root:
|
||||
if child.tag == qn("m:barPr"):
|
||||
for child2 in child:
|
||||
if child2.tag == qn("m:pos"):
|
||||
if child2.attrib.get(qn("m:val")) == "bot":
|
||||
text = "\\underline{"
|
||||
|
||||
for child in root:
|
||||
if child.tag == qn("m:e"):
|
||||
text += self.parse(child)
|
||||
text += "}"
|
||||
return text
|
||||
|
||||
def parse_border_box(self, root: Element) -> str:
|
||||
text = "\\boxed{"
|
||||
for child in root:
|
||||
if child.tag == qn("m:e"):
|
||||
text += self.parse(child)
|
||||
text += "}"
|
||||
return text
|
||||
|
||||
def parse_box(self, root: Element) -> str:
|
||||
text = ""
|
||||
for child in root:
|
||||
text += self.parse(child)
|
||||
return text
|
||||
|
||||
def parse_group_chr(self, root: Element) -> str:
|
||||
character_map = {
|
||||
"←": "\\leftarrow",
|
||||
"→": "\\rightarrow",
|
||||
"↔": "\\leftrightarrow",
|
||||
"⇐": "\\Leftarrow",
|
||||
"⇒": "\\Rightarrow",
|
||||
"⇔": "\\Leftrightarrow",
|
||||
}
|
||||
text = "\\underbrace{"
|
||||
bottom = False
|
||||
for child in root:
|
||||
if child.tag == qn("m:groupChrPr"):
|
||||
for child2 in child:
|
||||
if child2.tag == qn("m:chr"):
|
||||
char = child2.attrib.get(qn("m:val"))
|
||||
if char in character_map:
|
||||
text = character_map[char]
|
||||
for child2 in child:
|
||||
if (
|
||||
child2.tag == qn("m:pos")
|
||||
and child2.attrib.get(qn("m:val")) == "top"
|
||||
):
|
||||
# If m:pos is set to "top", the symbol is supposed to
|
||||
# be on top and the text is actually supposed to be under
|
||||
bottom = True
|
||||
|
||||
content = ""
|
||||
for child in root:
|
||||
if child.tag == qn("m:e"):
|
||||
content = self.parse(child)
|
||||
if text == "\\underbrace{":
|
||||
if bottom:
|
||||
text = "\\overbrace{" + content + "}"
|
||||
else:
|
||||
text += content + "}"
|
||||
else:
|
||||
if not bottom:
|
||||
text = "\\overset{" + content + "}" + "{" + text + "}"
|
||||
else:
|
||||
text = "\\underset{" + content + "}" + "{" + text + "}"
|
||||
return text
|
||||
|
||||
def parse_d(self, root: Element) -> str:
|
||||
bracket_map = {
|
||||
"(": "\\left(",
|
||||
")": "\\right)",
|
||||
"[": "\\left[",
|
||||
"]": "\\right]",
|
||||
"{": "\\left{",
|
||||
"}": "\\right}",
|
||||
"〈": "\\left\\langle",
|
||||
"〉": "\\right\\rangle",
|
||||
"⟨": "\\left\\langle",
|
||||
"⟩": "\\right\\rangle",
|
||||
"⌊": "\\left\\lfloor",
|
||||
"⌋": "\\right\\rfloor",
|
||||
"⌈": "\\left\\lceil",
|
||||
"⌉": "\\right\\rceil",
|
||||
"|": "\\left|",
|
||||
"‖": "\\left\\|",
|
||||
"⟦": "[\\![",
|
||||
"⟧": "]\\!]",
|
||||
}
|
||||
text = ""
|
||||
start_bracket = "("
|
||||
end_bracket = ")"
|
||||
seperator = "|"
|
||||
is_matrix = False
|
||||
for child in root:
|
||||
for child2 in child:
|
||||
if child.tag == qn("m:dPr"):
|
||||
if child2.tag == qn("m:begChr"):
|
||||
start_bracket = child2.attrib.get(qn("m:val"))
|
||||
if child2.tag == qn("m:endChr"):
|
||||
end_bracket = child2.attrib.get(qn("m:val"))
|
||||
if child2.tag == qn("m:sepChr"):
|
||||
seperator = child2.attrib.get(qn("m:val"))
|
||||
if child2.tag == qn("m:m"):
|
||||
is_matrix = True
|
||||
|
||||
for child in root:
|
||||
if child.tag == qn("m:e"):
|
||||
if text:
|
||||
text += seperator
|
||||
text += self.parse(child)
|
||||
end_bracket_replacements = {
|
||||
"|": "\\right|",
|
||||
"‖": "\\right\\|",
|
||||
"[": "\\right[",
|
||||
}
|
||||
start_bracket_replacements = {
|
||||
"]": "\\left]",
|
||||
}
|
||||
start = ""
|
||||
end = ""
|
||||
if start_bracket:
|
||||
if start_bracket in start_bracket_replacements:
|
||||
start = start_bracket_replacements[start_bracket] + " "
|
||||
elif start_bracket in bracket_map:
|
||||
start = bracket_map[start_bracket] + " "
|
||||
else:
|
||||
start = "\\left(" + " "
|
||||
if end_bracket:
|
||||
if end_bracket in end_bracket_replacements:
|
||||
end = " " + end_bracket_replacements[end_bracket]
|
||||
elif end_bracket in bracket_map:
|
||||
end = " " + bracket_map[end_bracket]
|
||||
else:
|
||||
end = " " + "\\right)"
|
||||
# If there is no end bracket and this tag contains an m:eqArr tag as a
|
||||
# child, we assume that the eqArr should be translated to a cases environment
|
||||
# instead of an eqnarray* environment.
|
||||
else:
|
||||
for child in root:
|
||||
if child.tag == qn("m:e"):
|
||||
for child2 in child:
|
||||
if child2.tag == qn("m:eqArr"):
|
||||
text = text.replace("\\begin{eqnarray*}", "")
|
||||
text = text.replace("\\end{eqnarray*}", "")
|
||||
return "\\begin{cases} " + text + " \\end{cases}"
|
||||
if is_matrix:
|
||||
if start_bracket == "(" and end_bracket == ")":
|
||||
return text.replace("{matrix}", "{pmatrix}")
|
||||
elif start_bracket == "|" and end_bracket == "|":
|
||||
return text.replace("{matrix}", "{vmatrix}")
|
||||
elif start_bracket == "‖" and end_bracket == "‖":
|
||||
return text.replace("{matrix}", "{Vmatrix}")
|
||||
else:
|
||||
return text.replace("{matrix}", "{bmatrix}")
|
||||
return start + text + end
|
||||
|
||||
def parse_eq_arr(self, root: Element) -> str:
|
||||
text = "\\begin{eqnarray*}"
|
||||
for child in root:
|
||||
if child.tag == qn("m:e"):
|
||||
text += self.parse(child) + " \\\\"
|
||||
text += "\\end{eqnarray*}"
|
||||
return text
|
||||
|
||||
def parse_f(self, root: Element) -> str:
|
||||
text = "\\frac{"
|
||||
num = ""
|
||||
den = ""
|
||||
is_binom = False
|
||||
for child in root:
|
||||
if child.tag == qn("m:fPr"):
|
||||
for child2 in child:
|
||||
if (
|
||||
child2.tag == qn("m:type")
|
||||
and child2.attrib.get(qn("m:val")) == "noBar"
|
||||
):
|
||||
is_binom = True
|
||||
if child.tag == qn("m:num"):
|
||||
num = self.parse(child)
|
||||
if child.tag == qn("m:den"):
|
||||
den = self.parse(child)
|
||||
if is_binom:
|
||||
text = "\\genfrac{}{}{0pt}{}{"
|
||||
text += num + "}{" + den + "}"
|
||||
return text
|
||||
|
||||
def parse_m(self, root: Element) -> str:
|
||||
text = "\\begin{matrix} "
|
||||
text += self.parse(root)[:-3] # Remove the last ' \\'
|
||||
text += "\\end{matrix}"
|
||||
return text
|
||||
|
||||
def parse_mr(self, root: Element) -> str:
|
||||
text = ""
|
||||
for child in root:
|
||||
if child.tag == qn("m:e"):
|
||||
text += self.parse(child) + " & "
|
||||
return text[:-2] + "\\\\ " # Remove the last ' & '
|
||||
|
||||
def parse_func(self, root: Element) -> str:
|
||||
subscript = ""
|
||||
superscript = ""
|
||||
text = ""
|
||||
func_name = "sin"
|
||||
for child in root:
|
||||
if child.tag == qn("m:fName"):
|
||||
for child2 in child:
|
||||
if child2.tag in [qn("m:sSup"), qn("m:sSub"), qn("m:r")]:
|
||||
for child3 in child2:
|
||||
if child3.tag == qn("m:sub"):
|
||||
subscript = self.parse(child3)
|
||||
if child3.tag == qn("m:sup"):
|
||||
superscript = self.parse(child3)
|
||||
if child3.tag == qn("m:t") or child3.tag == qn("m:e"):
|
||||
func_name = self.parse(child3)
|
||||
elif child2.tag == qn("m:limLow"):
|
||||
for child3 in child2:
|
||||
if child3.tag == qn("m:lim"):
|
||||
for child4 in child3:
|
||||
subscript += self.parse(child4)
|
||||
if child3.tag == qn("m:e"):
|
||||
func_name = self.parse(child3)
|
||||
|
||||
if child.tag == qn("m:e"):
|
||||
text += self.parse(child)
|
||||
if func_name in ["lim", "max", "min"]:
|
||||
return f"\\{func_name}\\limits_{{{subscript}}}^{{{superscript}}}{{{text}}}"
|
||||
if func_name not in self.FUNCTION_MAP:
|
||||
return f"{{{func_name}}}^{{{superscript}}}_{{{subscript}}}{{{text}}}"
|
||||
return (
|
||||
self.FUNCTION_MAP[func_name]
|
||||
+ f"_{{{subscript}}}^{{{superscript}}}{{{text}}}"
|
||||
)
|
||||
|
||||
def parse_s_sup(self, root: Element) -> str:
|
||||
content = ""
|
||||
exp_content = ""
|
||||
for child in root:
|
||||
if child.tag == qn("m:e"):
|
||||
content = self.parse(child)
|
||||
if child.tag == qn("m:sup"):
|
||||
exp_content = self.parse(child)
|
||||
content = self._normalize_func_name(content)
|
||||
return f"{{{content}}}^{{{exp_content}}}"
|
||||
|
||||
def parse_s_sub(self, root: Element) -> str:
|
||||
content = ""
|
||||
sub_content = ""
|
||||
for child in root:
|
||||
if child.tag == qn("m:e"):
|
||||
content = self.parse(child)
|
||||
if child.tag == qn("m:sub"):
|
||||
sub_content = self.parse(child)
|
||||
content = self._normalize_func_name(content)
|
||||
return f"{{{content}}}_{{{sub_content}}}"
|
||||
|
||||
def parse_s_sub_sup(self, root: Element) -> str:
|
||||
content = ""
|
||||
sub_content = ""
|
||||
exp_content = ""
|
||||
for child in root:
|
||||
if child.tag == qn("m:e"):
|
||||
content = self.parse(child)
|
||||
if child.tag == qn("m:sub"):
|
||||
sub_content = self.parse(child)
|
||||
if child.tag == qn("m:sup"):
|
||||
exp_content = self.parse(child)
|
||||
content = self._normalize_func_name(content)
|
||||
return f"{{{content}}}_{{{sub_content}}}^{{{exp_content}}}"
|
||||
|
||||
def parse_s_pre(self, root: Element) -> str:
|
||||
content = ""
|
||||
sub_content = ""
|
||||
exp_content = ""
|
||||
for child in root:
|
||||
if child.tag == qn("m:e"):
|
||||
content = self.parse(child)
|
||||
if child.tag == qn("m:sub"):
|
||||
sub_content = self.parse(child)
|
||||
if child.tag == qn("m:sup"):
|
||||
exp_content = self.parse(child)
|
||||
return "{}^{" + exp_content + "}_{" + sub_content + "}{" + content + "}"
|
||||
|
||||
def parse_rad(self, root: Element) -> str:
|
||||
content = ""
|
||||
order = ""
|
||||
for child in root:
|
||||
if child.tag == qn("m:deg"):
|
||||
order = self.parse(child)
|
||||
if child.tag == qn("m:e"):
|
||||
content += self.parse(child)
|
||||
if order:
|
||||
return f"\\sqrt[{order}]{{{content}}}"
|
||||
return f"\\sqrt{{{content}}}"
|
||||
|
||||
def parse_nary(self, root: Element) -> str:
|
||||
character_map = {
|
||||
8719: "\\prod",
|
||||
8720: "\\coprod",
|
||||
8721: "\\sum",
|
||||
8747: "\\int",
|
||||
8748: "\\iint",
|
||||
8749: "\\iiint",
|
||||
8750: "\\oint",
|
||||
8751: "\\oiint",
|
||||
8752: "\\oiiint",
|
||||
8896: "\\bigwedge",
|
||||
8897: "\\bigvee",
|
||||
8898: "\\bigcap",
|
||||
8899: "\\bigcup",
|
||||
}
|
||||
char = 8747
|
||||
for child in root:
|
||||
if child.tag == qn("m:naryPr"):
|
||||
for child2 in child:
|
||||
if child2.tag == qn("m:chr"):
|
||||
val = child2.attrib.get(qn("m:val"))
|
||||
if val:
|
||||
try:
|
||||
char = ord(val)
|
||||
except TypeError:
|
||||
pass
|
||||
text = character_map.get(char, character_map[8721])
|
||||
sub = ""
|
||||
sup = ""
|
||||
content = ""
|
||||
for child in root:
|
||||
if child.tag == qn("m:sub"):
|
||||
sub = self.parse(child)
|
||||
if child.tag == qn("m:sup"):
|
||||
sup = self.parse(child)
|
||||
if child.tag == qn("m:e"):
|
||||
content = self.parse(child)
|
||||
if sub:
|
||||
text += f"_{{{sub}}}"
|
||||
if sup:
|
||||
text += f"^{{{sup}}}"
|
||||
text += "{" + content + "}"
|
||||
return text
|
||||
|
||||
parsers = {
|
||||
qn("m:r"): parse_r,
|
||||
qn("m:acc"): parse_acc,
|
||||
qn("m:borderBox"): parse_border_box,
|
||||
qn("m:bar"): parse_bar,
|
||||
qn("m:box"): parse_box,
|
||||
qn("m:d"): parse_d,
|
||||
qn("m:e"): parse_e,
|
||||
qn("m:groupChr"): parse_group_chr,
|
||||
qn("m:f"): parse_f,
|
||||
qn("m:sSup"): parse_s_sup,
|
||||
qn("m:sSub"): parse_s_sub,
|
||||
qn("m:sSubSup"): parse_s_sub_sup,
|
||||
qn("m:sPre"): parse_s_pre,
|
||||
qn("m:t"): parse_t,
|
||||
qn("m:rad"): parse_rad,
|
||||
qn("m:nary"): parse_nary,
|
||||
qn("m:eqArr"): parse_eq_arr,
|
||||
qn("m:func"): parse_func,
|
||||
qn("m:m"): parse_m,
|
||||
qn("m:mr"): parse_mr,
|
||||
}
|
||||
@@ -0,0 +1,40 @@
|
||||
"""
|
||||
Utility functions to extract text from the supported mathematical equations from xml tags and
|
||||
convert them into LaTeX
|
||||
"""
|
||||
|
||||
from .cleaners import clean_exp
|
||||
|
||||
ns_map = {
|
||||
"w": "http://schemas.openxmlformats.org/wordprocessingml/2006/main",
|
||||
"m": "http://schemas.openxmlformats.org/officeDocument/2006/math",
|
||||
}
|
||||
|
||||
|
||||
def linear_expression(tag):
|
||||
"""
|
||||
Just returns the text contained in the given tag while setting docxlatex_skip_iteration flags
|
||||
for all its children.
|
||||
:param tag:defusedxml.Element - An xml element which contains a math equation in linear form
|
||||
:return text:str - The equation in valid LaTeX syntax
|
||||
"""
|
||||
text = ""
|
||||
for child in tag.iter():
|
||||
child.set("docxlatex_skip_iteration", True)
|
||||
text += child.text if child.text is not None else ""
|
||||
text = clean_exp(text)
|
||||
return text
|
||||
|
||||
|
||||
def qn(tag):
|
||||
"""
|
||||
A utility function to turn a namespace
|
||||
prefixed tag name into a Clark-notation qualified tag name for lxml. For
|
||||
example, qn('m:oMath') returns '{http://schemas.openxmlformats.org/officeDocument/2006/math}oMath'
|
||||
|
||||
:param tag:str - A namespace-prefixed tag name
|
||||
:return qn:str - A Clark-notation qualified name tag for lxml.
|
||||
"""
|
||||
prefix, tag_root = tag.split(":")
|
||||
uri = ns_map[prefix]
|
||||
return "{{{}}}{}".format(uri, tag_root)
|
||||
Executable
+932
@@ -0,0 +1,932 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
ABOUTME: Parses DOCX documents into text blocks using python-docx
|
||||
ABOUTME: Extracts automatic numbering, splits by headings, converts tables to JSON
|
||||
"""
|
||||
|
||||
import json
|
||||
import sys
|
||||
|
||||
try:
|
||||
from docx import Document
|
||||
from docx.opc.exceptions import PackageNotFoundError
|
||||
except ImportError as exc:
|
||||
# Raise instead of sys.exit: this module is imported in-process by the
|
||||
# gunicorn/uvicorn worker, where a SystemExit would tear down the whole
|
||||
# worker rather than surfacing a normal, catchable error.
|
||||
raise ImportError(
|
||||
"python-docx not installed. Run: pip install python-docx"
|
||||
) from exc
|
||||
|
||||
from lightrag.parser._markdown import (
|
||||
render_heading_line,
|
||||
strip_heading_markdown_prefix,
|
||||
)
|
||||
from .numbering_resolver import NumberingResolver
|
||||
from .table_extractor import TableExtractor
|
||||
from .drawing_image_extractor import (
|
||||
DrawingExtractionContext,
|
||||
extract_drawing_placeholder_from_element,
|
||||
extract_vml_image_placeholder_from_element,
|
||||
)
|
||||
|
||||
|
||||
# Constants for content validation (character-based for UI/display)
|
||||
MAX_HEADING_LENGTH = 200 # Maximum heading length in characters (UI constraint)
|
||||
|
||||
# OOXML tracked-change/comment tags whose subtree must be dropped so we only
|
||||
# surface the *final* revised text. w:ins / w:moveTo are kept via default
|
||||
# recursion so inserted/moved-in content survives.
|
||||
_SKIP_REVISION_TAGS = frozenset({"del", "moveFrom"})
|
||||
_SKIP_COMMENT_TAGS = frozenset(
|
||||
{"commentRangeStart", "commentRangeEnd", "commentReference", "annotationRef"}
|
||||
)
|
||||
_SKIP_PARAGRAPH_TAGS = _SKIP_REVISION_TAGS | _SKIP_COMMENT_TAGS
|
||||
|
||||
|
||||
class DocxContentError(ValueError):
|
||||
"""DOCX content violates a parsing constraint (heading/table/anchor limits).
|
||||
|
||||
Raised instead of calling ``sys.exit`` so the pipeline's per-document
|
||||
``except Exception`` handler marks just that document FAILED while the
|
||||
gunicorn/uvicorn worker process keeps running. Subclasses ``ValueError``
|
||||
(i.e. an ``Exception``, not ``BaseException``) so the existing pipeline
|
||||
handlers catch it.
|
||||
"""
|
||||
|
||||
|
||||
def format_error(title: str, details: str, solution: str) -> str:
|
||||
"""
|
||||
Build a friendly, formatted error message (title / details / SOLUTION).
|
||||
|
||||
Args:
|
||||
title: Error title
|
||||
details: Detailed error information
|
||||
solution: Suggested solution steps
|
||||
|
||||
Returns:
|
||||
str: The formatted multi-line message.
|
||||
"""
|
||||
return (
|
||||
"\n"
|
||||
+ "=" * 68
|
||||
+ f"\nERROR: {title}\n"
|
||||
+ "=" * 68
|
||||
+ f"\n\n{details}"
|
||||
+ "\n\nSOLUTION:\n"
|
||||
+ solution
|
||||
+ "\n\n"
|
||||
+ "=" * 68
|
||||
+ "\n"
|
||||
)
|
||||
|
||||
|
||||
def print_error(title: str, details: str, solution: str):
|
||||
"""Print a friendly, formatted error message to stderr."""
|
||||
print(format_error(title, details, solution), file=sys.stderr)
|
||||
|
||||
|
||||
def _diagnose_invalid_docx(file_path: str) -> tuple[str, str]:
|
||||
"""Diagnose why a ``.docx`` file is not a valid OOXML/ZIP package.
|
||||
|
||||
python-docx raises ``PackageNotFoundError("Package not found at '...'")``
|
||||
both when the path is missing AND when the file exists but is not a valid
|
||||
zip. By the time this runs the file has already been confirmed to exist
|
||||
(the native worker validates ``p.exists()`` first), so the real cause is a
|
||||
corrupt file or a non-DOCX payload wearing a ``.docx`` extension. Sniff the
|
||||
magic bytes to name the actual format so the error message reflects the
|
||||
real problem instead of an empty "not found".
|
||||
|
||||
Returns a ``(details, solution)`` tuple for :func:`format_error`. Reads only
|
||||
the file header and never raises — any IO failure degrades to a generic
|
||||
"cannot read" diagnosis.
|
||||
"""
|
||||
import zipfile
|
||||
|
||||
convert_solution = (
|
||||
" 1. Open the file in Microsoft Word or WPS\n"
|
||||
' 2. Use "Save As" and choose "Word Document (*.docx)"\n'
|
||||
" 3. Re-upload the converted .docx to LightRAG"
|
||||
)
|
||||
|
||||
try:
|
||||
with open(file_path, "rb") as f:
|
||||
head = f.read(8)
|
||||
except OSError as exc:
|
||||
return (
|
||||
f"The file at '{file_path}' could not be read: {exc}",
|
||||
" 1. Verify the file exists and is readable\n"
|
||||
" 2. Re-upload it to LightRAG",
|
||||
)
|
||||
|
||||
if not head:
|
||||
return (
|
||||
f"The file at '{file_path}' is empty (0 bytes). The upload was "
|
||||
"likely truncated or the source file is corrupt.",
|
||||
" 1. Check the original document opens correctly\n"
|
||||
" 2. Re-upload a complete copy to LightRAG",
|
||||
)
|
||||
|
||||
if head.startswith(b"\xd0\xcf\x11\xe0\xa1\xb1\x1a\xe1"):
|
||||
# OLE2 Compound File — the legacy binary Word 97-2003 .doc format.
|
||||
return (
|
||||
f"The file at '{file_path}' is a legacy Word 97-2003 (.doc) "
|
||||
"document saved with a .docx extension. The .doc binary format is "
|
||||
"not a ZIP/OOXML package and cannot be parsed by the native engine.",
|
||||
convert_solution,
|
||||
)
|
||||
|
||||
if head.startswith(b"{\\rtf"):
|
||||
return (
|
||||
f"The file at '{file_path}' is an RTF document saved with a .docx "
|
||||
"extension. RTF is not a ZIP/OOXML package.",
|
||||
convert_solution,
|
||||
)
|
||||
|
||||
if head.startswith(b"%PDF"):
|
||||
return (
|
||||
f"The file at '{file_path}' is a PDF saved with a .docx extension. "
|
||||
"It is not a ZIP/OOXML package.",
|
||||
" 1. Convert the PDF to .docx, or upload it through a PDF-capable "
|
||||
"parser engine (e.g. mineru/docling)\n"
|
||||
" 2. Re-upload to LightRAG",
|
||||
)
|
||||
|
||||
stripped = head.lstrip()
|
||||
if stripped.startswith(b"<"):
|
||||
# <?xml ...>, <html ...>, or Word 2003 "<w:wordDocument>" flat XML.
|
||||
return (
|
||||
f"The file at '{file_path}' is an HTML or XML document saved with a "
|
||||
".docx extension, not a ZIP/OOXML package.",
|
||||
convert_solution,
|
||||
)
|
||||
|
||||
if head.startswith(b"PK\x03\x04") and not zipfile.is_zipfile(file_path):
|
||||
# Has the ZIP local-file-header magic but the archive is unreadable.
|
||||
return (
|
||||
f"The file at '{file_path}' starts like a ZIP archive but is "
|
||||
"truncated or corrupt, so it cannot be opened as a DOCX package.",
|
||||
" 1. Check the original document opens correctly\n"
|
||||
" 2. Re-upload a complete, uncorrupted copy to LightRAG",
|
||||
)
|
||||
|
||||
return (
|
||||
f"The file at '{file_path}' is not a valid DOCX (ZIP/OOXML) package. "
|
||||
"It is either corrupt or a different file format saved with a .docx "
|
||||
"extension.",
|
||||
convert_solution,
|
||||
)
|
||||
|
||||
|
||||
def truncate_heading(heading_text: str, para_id: str = None) -> str:
|
||||
"""
|
||||
Truncate heading if it exceeds MAX_HEADING_LENGTH.
|
||||
|
||||
Args:
|
||||
heading_text: The heading text to check
|
||||
para_id: Optional paragraph ID for warning message
|
||||
|
||||
Returns:
|
||||
str: Original heading if within limit, truncated heading with "..." if too long
|
||||
"""
|
||||
if len(heading_text) > MAX_HEADING_LENGTH:
|
||||
truncated = heading_text[: MAX_HEADING_LENGTH - 3] + "..."
|
||||
location = f" (para_id: {para_id})" if para_id else ""
|
||||
print(
|
||||
f"Warning: Heading truncated (length {len(heading_text)} > max {MAX_HEADING_LENGTH}){location}: "
|
||||
f'"{truncated}"',
|
||||
file=sys.stderr,
|
||||
)
|
||||
return truncated
|
||||
return heading_text
|
||||
|
||||
|
||||
def validate_heading_length(heading_text: str, para_id: str):
|
||||
"""
|
||||
Validate that heading length does not exceed MAX_HEADING_LENGTH.
|
||||
|
||||
Args:
|
||||
heading_text: The heading text to validate
|
||||
para_id: The paragraph ID for error reporting
|
||||
|
||||
Raises:
|
||||
DocxContentError: if heading exceeds maximum length
|
||||
"""
|
||||
if len(heading_text) > MAX_HEADING_LENGTH:
|
||||
preview = (
|
||||
heading_text[:100] + "..." if len(heading_text) > 100 else heading_text
|
||||
)
|
||||
raise DocxContentError(
|
||||
format_error(
|
||||
f"Heading too long ({len(heading_text)} characters, max {MAX_HEADING_LENGTH})",
|
||||
f"The following heading exceeds the maximum allowed length:\n\n{preview}\n\n"
|
||||
f"Location(para_id): {para_id}\n"
|
||||
f"Actual length: {len(heading_text)} characters",
|
||||
" 1. Open the document in Microsoft Word\n"
|
||||
f" 2. Shorten this heading to {MAX_HEADING_LENGTH} characters or less\n"
|
||||
" 3. Re-upload it to LightRAG",
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def find_first_valid_para_id(para_ids: list) -> str | None:
|
||||
"""
|
||||
Find the first valid paraId in a 2D array of paraIds.
|
||||
|
||||
Args:
|
||||
para_ids: 2D list of paraIds from table cells
|
||||
|
||||
Returns:
|
||||
First non-None paraId found, or None when every cell lacks a paraId.
|
||||
Callers must tolerate ``None`` and treat it as a tracking gap rather
|
||||
than a fatal error (legacy / non-Word docx authors omit ``w14:paraId``
|
||||
attributes and we want to keep parsing).
|
||||
"""
|
||||
for row in para_ids:
|
||||
for para_id in row:
|
||||
if para_id:
|
||||
return para_id
|
||||
return None
|
||||
|
||||
|
||||
def find_last_valid_para_id(para_ids: list) -> str | None:
|
||||
"""
|
||||
Find the last valid paraId in a 2D array of paraIds.
|
||||
|
||||
Returns the last non-None paraId, falling back to the first valid one
|
||||
when reverse-iteration does not yield anything (single-paraId tables),
|
||||
and finally ``None`` when every cell lacks a paraId.
|
||||
"""
|
||||
for row in reversed(para_ids):
|
||||
for para_id in reversed(row):
|
||||
if para_id:
|
||||
return para_id
|
||||
|
||||
return find_first_valid_para_id(para_ids)
|
||||
|
||||
|
||||
def _table_has_any_paraid(para_ids: list) -> bool:
|
||||
"""True when at least one cell in the 2D paraId grid carries an id."""
|
||||
return find_first_valid_para_id(para_ids) is not None
|
||||
|
||||
|
||||
def extract_para_id(para_element) -> str:
|
||||
"""
|
||||
Extract w14:paraId attribute from paragraph element.
|
||||
|
||||
Args:
|
||||
para_element: lxml paragraph element
|
||||
|
||||
Returns:
|
||||
8-character hex paraId, or ``None`` when the paragraph carries no
|
||||
``w14:paraId`` attribute (legacy / non-Word docx authors). Callers
|
||||
propagate the ``None`` upward — the LightRAG adapter counts these
|
||||
and surfaces a single warning per document.
|
||||
"""
|
||||
return para_element.get(
|
||||
"{http://schemas.microsoft.com/office/word/2010/wordml}paraId"
|
||||
)
|
||||
|
||||
|
||||
def parse_styles_outline_levels(docx_path: str) -> dict:
|
||||
"""
|
||||
Parse styles.xml to extract outlineLvl definitions for each style,
|
||||
following style inheritance chain (basedOn).
|
||||
|
||||
Args:
|
||||
docx_path: Path to DOCX file
|
||||
|
||||
Returns:
|
||||
dict: styleId -> outlineLvl (0-8 for headings, 9 for body text)
|
||||
"""
|
||||
import zipfile
|
||||
|
||||
try:
|
||||
from defusedxml import ElementTree as ET
|
||||
except ImportError:
|
||||
from xml.etree import ElementTree as ET
|
||||
|
||||
styles_outline = {} # styleId -> outlineLvl (directly defined)
|
||||
style_based_on = {} # styleId -> parent styleId
|
||||
|
||||
try:
|
||||
with zipfile.ZipFile(docx_path, "r") as zf:
|
||||
if "word/styles.xml" not in zf.namelist():
|
||||
return styles_outline
|
||||
|
||||
tree = ET.parse(zf.open("word/styles.xml"))
|
||||
root = tree.getroot()
|
||||
|
||||
ns = "http://schemas.openxmlformats.org/wordprocessingml/2006/main"
|
||||
|
||||
# First pass: collect outlineLvl and basedOn for all styles
|
||||
for style in root.findall(f".//{{{ns}}}style"):
|
||||
style_id = style.get(f"{{{ns}}}styleId")
|
||||
if not style_id:
|
||||
continue
|
||||
|
||||
# Check for basedOn (style inheritance)
|
||||
based_on = style.find(f"{{{ns}}}basedOn")
|
||||
if based_on is not None:
|
||||
parent_id = based_on.get(f"{{{ns}}}val")
|
||||
if parent_id:
|
||||
style_based_on[style_id] = parent_id
|
||||
|
||||
# Check for outlineLvl in style's pPr
|
||||
pPr = style.find(f"{{{ns}}}pPr")
|
||||
if pPr is not None:
|
||||
outline_lvl_elem = pPr.find(f"{{{ns}}}outlineLvl")
|
||||
if outline_lvl_elem is not None:
|
||||
level = int(outline_lvl_elem.get(f"{{{ns}}}val"))
|
||||
styles_outline[style_id] = level
|
||||
|
||||
# Second pass: resolve inheritance chain for styles without direct outlineLvl
|
||||
def get_outline_level(style_id: str, visited: set = None) -> int:
|
||||
if visited is None:
|
||||
visited = set()
|
||||
if style_id in visited:
|
||||
return None # Prevent circular references
|
||||
visited.add(style_id)
|
||||
|
||||
# If this style directly defines outlineLvl, return it
|
||||
if style_id in styles_outline:
|
||||
return styles_outline[style_id]
|
||||
|
||||
# Otherwise check parent style
|
||||
if style_id in style_based_on:
|
||||
parent_id = style_based_on[style_id]
|
||||
return get_outline_level(parent_id, visited)
|
||||
|
||||
return None
|
||||
|
||||
# Fill in missing outlineLvl from inheritance chain
|
||||
all_style_ids = set(styles_outline.keys()) | set(style_based_on.keys())
|
||||
for style_id in all_style_ids:
|
||||
if style_id not in styles_outline:
|
||||
level = get_outline_level(style_id)
|
||||
if level is not None:
|
||||
styles_outline[style_id] = level
|
||||
except Exception:
|
||||
# Silently ignore parsing errors
|
||||
pass
|
||||
|
||||
return styles_outline
|
||||
|
||||
|
||||
def get_heading_level(para_element, styles_outline_map: dict) -> int:
|
||||
"""
|
||||
Get heading level from paragraph, checking both direct format and style.
|
||||
|
||||
Priority: paragraph outlineLvl > style outlineLvl
|
||||
|
||||
Args:
|
||||
para_element: lxml paragraph element
|
||||
styles_outline_map: dict of styleId -> outlineLvl from styles.xml
|
||||
|
||||
Returns:
|
||||
int: 0-8 for heading levels (0=level 1, 1=level 2, etc.), None for non-heading
|
||||
"""
|
||||
# 1. Check paragraph direct format
|
||||
pPr = para_element.find(
|
||||
"{http://schemas.openxmlformats.org/wordprocessingml/2006/main}pPr"
|
||||
)
|
||||
if pPr is not None:
|
||||
outline_elem = pPr.find(
|
||||
"{http://schemas.openxmlformats.org/wordprocessingml/2006/main}outlineLvl"
|
||||
)
|
||||
if outline_elem is not None:
|
||||
level = int(
|
||||
outline_elem.get(
|
||||
"{http://schemas.openxmlformats.org/wordprocessingml/2006/main}val"
|
||||
)
|
||||
)
|
||||
# Only 0-8 are true heading levels (9 is body text)
|
||||
if level < 9:
|
||||
return level
|
||||
else:
|
||||
return None # Level 9 is body text
|
||||
|
||||
# 2. Check style definition's outlineLvl
|
||||
if pPr is not None:
|
||||
pStyle_elem = pPr.find(
|
||||
"{http://schemas.openxmlformats.org/wordprocessingml/2006/main}pStyle"
|
||||
)
|
||||
if pStyle_elem is not None:
|
||||
style_id = pStyle_elem.get(
|
||||
"{http://schemas.openxmlformats.org/wordprocessingml/2006/main}val"
|
||||
)
|
||||
if style_id and style_id in styles_outline_map:
|
||||
level = styles_outline_map[style_id]
|
||||
if level < 9:
|
||||
return level
|
||||
else:
|
||||
return None
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def extract_text_from_run(
|
||||
run,
|
||||
ns: dict,
|
||||
drawing_context: DrawingExtractionContext = None,
|
||||
) -> str:
|
||||
"""
|
||||
Extract text from a run element, preserving superscript/subscript with markup.
|
||||
|
||||
Converts Word formatting to HTML-like tags:
|
||||
- Superscript: <sup>text</sup>
|
||||
- Subscript: <sub>text</sub>
|
||||
- Normal text: unchanged
|
||||
|
||||
Args:
|
||||
run: lxml run element (w:r)
|
||||
ns: XML namespace dictionary
|
||||
|
||||
Returns:
|
||||
Text string with <sup>/<sub> markup for formatted portions
|
||||
"""
|
||||
text = ""
|
||||
|
||||
# Check for vertAlign in rPr (superscript/subscript)
|
||||
vert_align = None
|
||||
rPr = run.find("w:rPr", ns)
|
||||
if rPr is not None:
|
||||
vert_elem = rPr.find("w:vertAlign", ns)
|
||||
if vert_elem is not None:
|
||||
vert_align = vert_elem.get(
|
||||
"{http://schemas.openxmlformats.org/wordprocessingml/2006/main}val"
|
||||
)
|
||||
|
||||
# Extract text content from run children
|
||||
for child in run:
|
||||
tag = child.tag.split("}")[-1] # Remove namespace
|
||||
if tag == "t" and child.text:
|
||||
text += child.text
|
||||
elif tag == "tab":
|
||||
text += "\t"
|
||||
elif tag == "br":
|
||||
# Handle line breaks - textWrapping or no type = soft line break
|
||||
br_type = child.get(
|
||||
"{http://schemas.openxmlformats.org/wordprocessingml/2006/main}type"
|
||||
)
|
||||
if br_type in (None, "textWrapping"):
|
||||
text += "\n"
|
||||
# Skip page and column breaks (layout elements)
|
||||
elif tag == "drawing":
|
||||
text += extract_drawing_placeholder_from_element(
|
||||
child,
|
||||
context=drawing_context,
|
||||
include_extended_attrs=True,
|
||||
)
|
||||
elif tag in ("pict", "object"):
|
||||
text += extract_vml_image_placeholder_from_element(
|
||||
child,
|
||||
context=drawing_context,
|
||||
include_extended_attrs=True,
|
||||
)
|
||||
|
||||
# Apply superscript/subscript markup if needed
|
||||
if text and vert_align == "superscript":
|
||||
return f"<sup>{text}</sup>"
|
||||
elif text and vert_align == "subscript":
|
||||
return f"<sub>{text}</sub>"
|
||||
|
||||
return text
|
||||
|
||||
|
||||
def extract_paragraph_content(
|
||||
element,
|
||||
ns,
|
||||
drawing_context: DrawingExtractionContext = None,
|
||||
) -> str:
|
||||
"""
|
||||
Extract text and equations from a paragraph element in document order.
|
||||
|
||||
Handles w:r (text runs), m:oMath (inline equations), and m:oMathPara
|
||||
(block equations). Recurses into container elements (e.g., w:hyperlink,
|
||||
w:ins, w:sdt, w:fldSimple, w:smartTag) to avoid dropping content.
|
||||
|
||||
Args:
|
||||
element: lxml paragraph element (w:p)
|
||||
ns: XML namespace dictionary
|
||||
|
||||
Returns:
|
||||
Text string with equations wrapped in <equation> tags
|
||||
"""
|
||||
parts = []
|
||||
|
||||
def append_from(node) -> None:
|
||||
tag = node.tag.split("}")[-1]
|
||||
# Drop tracked-change deletions (w:del/w:moveFrom) and comment markers
|
||||
# (w:commentRangeStart/End, w:commentReference, w:annotationRef) so the
|
||||
# output only contains the final revised text without annotation glyphs.
|
||||
if tag in _SKIP_PARAGRAPH_TAGS:
|
||||
return
|
||||
if tag == "r":
|
||||
parts.append(
|
||||
extract_text_from_run(node, ns, drawing_context=drawing_context)
|
||||
)
|
||||
return
|
||||
if tag == "oMath":
|
||||
from .omml import convert_omml_to_latex
|
||||
|
||||
latex = convert_omml_to_latex(node)
|
||||
if latex:
|
||||
parts.append(f"<equation>{latex}</equation>")
|
||||
return
|
||||
if tag == "oMathPara":
|
||||
from .omml import convert_omml_to_latex
|
||||
|
||||
for omath in node:
|
||||
if omath.tag.split("}")[-1] == "oMath":
|
||||
latex = convert_omml_to_latex(omath)
|
||||
if latex:
|
||||
parts.append(f"<equation>{latex}</equation>")
|
||||
return
|
||||
for child in node:
|
||||
append_from(child)
|
||||
|
||||
for child in element:
|
||||
append_from(child)
|
||||
|
||||
return "".join(parts)
|
||||
|
||||
|
||||
def _is_table_empty(rows: list) -> bool:
|
||||
"""Return True iff every cell in ``rows`` is whitespace-only."""
|
||||
return all(not (cell or "").strip() for row in rows for cell in row)
|
||||
|
||||
|
||||
def _collect_table_headers(paragraphs: list) -> list:
|
||||
"""Collect per-table cross-page header rows from ``is_table`` paragraphs.
|
||||
|
||||
The returned list is aligned 1:1 with the order of ``<table>`` placeholder
|
||||
tags emitted into the block's content; entries are either the list of
|
||||
header rows captured from ``w:tblHeader`` or ``None`` when the table has
|
||||
no cross-page repeating header.
|
||||
"""
|
||||
return [p.get("_table_header") for p in paragraphs if p.get("is_table")]
|
||||
|
||||
|
||||
def _build_unsplit_block(
|
||||
heading: str, paragraphs: list, parent_headings: list, level: int
|
||||
) -> dict:
|
||||
"""Build a single block from paragraphs without size-based splitting."""
|
||||
last_para = paragraphs[-1]
|
||||
block = {
|
||||
"uuid": paragraphs[0]["para_id"],
|
||||
"uuid_end": last_para.get("para_id_end") or last_para.get("para_id"),
|
||||
"heading": heading,
|
||||
"content": "\n".join(p["text"] for p in paragraphs),
|
||||
"type": "text",
|
||||
"parent_headings": parent_headings,
|
||||
"level": level,
|
||||
}
|
||||
table_headers = _collect_table_headers(paragraphs)
|
||||
if table_headers:
|
||||
block["table_headers"] = table_headers
|
||||
return block
|
||||
|
||||
|
||||
def _flush_current_block(
|
||||
blocks: list,
|
||||
heading: str,
|
||||
paragraphs: list,
|
||||
parent_headings: list,
|
||||
level: int,
|
||||
) -> None:
|
||||
"""Flush accumulated paragraphs into a single heading-scoped block.
|
||||
|
||||
The native parser performs only heading-driven structural splitting; block
|
||||
sizing (long-block anchor splitting, table row splitting, small-block
|
||||
merging) is the downstream paragraph-semantic chunker's responsibility.
|
||||
"""
|
||||
if not paragraphs:
|
||||
return
|
||||
|
||||
blocks.append(_build_unsplit_block(heading, paragraphs, parent_headings, level))
|
||||
|
||||
|
||||
def extract_docx_blocks(
|
||||
file_path: str,
|
||||
drawing_context: DrawingExtractionContext = None,
|
||||
parse_warnings: dict | None = None,
|
||||
parse_metadata: dict | None = None,
|
||||
) -> list:
|
||||
"""
|
||||
Extract heading-scoped text blocks from a DOCX file.
|
||||
|
||||
Uses python-docx with a custom numbering resolver to:
|
||||
1. Capture automatic numbering (list labels)
|
||||
2. Split the document into one block per heading (structural splitting)
|
||||
3. Convert tables to JSON (2D array) and emit them as <table> placeholders
|
||||
4. Preserve superscript/subscript formatting with <sup>/<sub> markup
|
||||
|
||||
Block sizing — long-block anchor splitting, table row splitting, and
|
||||
small-block merging — is intentionally NOT done here; it is the downstream
|
||||
paragraph-semantic chunker's responsibility. Blocks emitted here may
|
||||
therefore be arbitrarily large.
|
||||
|
||||
Args:
|
||||
file_path: Path to the DOCX file
|
||||
parse_warnings: Optional out-dict that this function mutates with
|
||||
non-fatal warnings observed during parsing. Currently used for
|
||||
``missing_paraid_count`` — incremented once per body-level
|
||||
paragraph (heading or text) that lacks a ``w14:paraId`` and once
|
||||
per table whose every cell lacks one. Callers (the LightRAG
|
||||
adapter / debug CLI) read this to surface a one-line warning per
|
||||
document instead of crashing.
|
||||
parse_metadata: Optional out-dict that this function mutates with
|
||||
document-level metadata derived during parsing. Currently used
|
||||
for ``first_heading`` — the text of the first heading encountered
|
||||
in document order (regardless of level). Used by the LightRAG
|
||||
adapter to populate ``meta.doc_title`` in ``.blocks.jsonl``.
|
||||
|
||||
Returns:
|
||||
List of block dictionaries with heading, content, type, and metadata
|
||||
"""
|
||||
try:
|
||||
doc = Document(file_path)
|
||||
except PackageNotFoundError as exc:
|
||||
# python-docx surfaces a misleading "Package not found at '...'" for any
|
||||
# file it cannot open as a ZIP/OOXML package — including files that
|
||||
# exist but are corrupt or a different format wearing a .docx extension.
|
||||
# Diagnose the real cause from the magic bytes and raise a DocxContentError
|
||||
# (a ValueError) so the pipeline's per-document handler marks just this
|
||||
# document FAILED with an accurate, actionable message.
|
||||
details, solution = _diagnose_invalid_docx(file_path)
|
||||
raise DocxContentError(
|
||||
format_error("File is not a valid DOCX document", details, solution)
|
||||
) from exc
|
||||
resolver = NumberingResolver(file_path)
|
||||
styles_outline = parse_styles_outline_levels(file_path)
|
||||
|
||||
blocks = []
|
||||
current_heading = "Preface/Uncategorized"
|
||||
current_heading_level = 1 # Default level for "Preface/Uncategorized"
|
||||
current_heading_stack = {} # {level: heading_text} - Use dict to correctly track heading hierarchy
|
||||
current_parent_headings = [] # Parent headings for current block
|
||||
current_paragraphs = [] # Track paragraphs with metadata for splitting
|
||||
first_heading_recorded = (
|
||||
False # Track whether the document's first heading has been captured
|
||||
)
|
||||
|
||||
# Iterate through document body elements (paragraphs and tables)
|
||||
body = doc._element.body
|
||||
|
||||
for element in body:
|
||||
tag = element.tag.split("}")[-1] # Remove namespace
|
||||
|
||||
if tag == "sectPr": # Document-level section break
|
||||
resolver.reset_tracking_state()
|
||||
continue
|
||||
|
||||
if tag == "p": # Paragraph
|
||||
# Get paragraph text with superscript/subscript markup and equations
|
||||
para_text = ""
|
||||
ns = {
|
||||
"w": "http://schemas.openxmlformats.org/wordprocessingml/2006/main",
|
||||
"wp": "http://schemas.openxmlformats.org/drawingml/2006/wordprocessingDrawing",
|
||||
"m": "http://schemas.openxmlformats.org/officeDocument/2006/math",
|
||||
}
|
||||
para_text = extract_paragraph_content(
|
||||
element,
|
||||
ns,
|
||||
drawing_context=drawing_context,
|
||||
)
|
||||
|
||||
para_text = para_text.strip()
|
||||
if not para_text:
|
||||
continue
|
||||
|
||||
# Get numbering label using our resolver
|
||||
label = resolver.get_label(element)
|
||||
full_text = f"{label} {para_text}".strip() if label else para_text
|
||||
|
||||
# Check if this is a heading using the new function
|
||||
outline_level = get_heading_level(element, styles_outline)
|
||||
|
||||
# A "heading" longer than MAX_HEADING_LENGTH is not a real heading.
|
||||
# The common cause (WPS/Word): the author set an outline level on a
|
||||
# paragraph but typed the body with soft line breaks (Shift+Enter →
|
||||
# <w:br/> → '\n') instead of starting a new paragraph, so heading
|
||||
# text + body live in one <w:p>. Split at the first soft break: the
|
||||
# first line stays the heading, the remainder becomes body text. If
|
||||
# there is no usable soft break (a genuine single-line over-long
|
||||
# heading), demote the whole paragraph to body text. Either way we
|
||||
# avoid crashing via validate_heading_length() and never drop content.
|
||||
demoted_body_text = None
|
||||
if outline_level is not None and len(full_text) > MAX_HEADING_LENGTH:
|
||||
head, sep, rest = full_text.partition("\n")
|
||||
if sep and len(head) <= MAX_HEADING_LENGTH:
|
||||
full_text = head
|
||||
demoted_body_text = rest.strip() or None
|
||||
if parse_warnings is not None:
|
||||
parse_warnings["heading_softbreak_split_count"] = (
|
||||
parse_warnings.get("heading_softbreak_split_count", 0) + 1
|
||||
)
|
||||
print(
|
||||
f"Warning: heading paragraph exceeded {MAX_HEADING_LENGTH} "
|
||||
"chars; split at soft line break — kept first line as "
|
||||
"heading, rest as body.",
|
||||
file=sys.stderr,
|
||||
)
|
||||
else:
|
||||
outline_level = None
|
||||
if parse_warnings is not None:
|
||||
parse_warnings["demoted_oversize_heading_count"] = (
|
||||
parse_warnings.get("demoted_oversize_heading_count", 0) + 1
|
||||
)
|
||||
print(
|
||||
f"Warning: paragraph has outline level but is "
|
||||
f"{len(full_text)} chars (> {MAX_HEADING_LENGTH}); treating "
|
||||
"as body text, not a heading.",
|
||||
file=sys.stderr,
|
||||
)
|
||||
|
||||
if outline_level is not None:
|
||||
# This is a heading (outline level 0-8)
|
||||
# Convert 0-based to 1-based level
|
||||
level = outline_level + 1
|
||||
|
||||
# Extract paraId for this heading
|
||||
heading_para_id = extract_para_id(element)
|
||||
if parse_warnings is not None and not heading_para_id:
|
||||
parse_warnings["missing_paraid_count"] = (
|
||||
parse_warnings.get("missing_paraid_count", 0) + 1
|
||||
)
|
||||
|
||||
# Validate heading length
|
||||
validate_heading_length(full_text, heading_para_id)
|
||||
|
||||
# Truncate heading if needed before storing
|
||||
truncated_text = truncate_heading(full_text, heading_para_id)
|
||||
clean_heading_text = strip_heading_markdown_prefix(truncated_text)
|
||||
|
||||
# Record the document's first heading (any level) for meta.doc_title.
|
||||
if not first_heading_recorded:
|
||||
if parse_metadata is not None:
|
||||
parse_metadata["first_heading"] = clean_heading_text
|
||||
first_heading_recorded = True
|
||||
|
||||
# Every recognized heading starts its own block. Always flush the
|
||||
# accumulated paragraphs so a heading with no body becomes a
|
||||
# standalone block whose content is just the heading text,
|
||||
# instead of being folded into the next heading's block.
|
||||
if current_paragraphs:
|
||||
_flush_current_block(
|
||||
blocks,
|
||||
current_heading,
|
||||
current_paragraphs,
|
||||
current_parent_headings,
|
||||
current_heading_level,
|
||||
)
|
||||
|
||||
# Reset for new block
|
||||
current_paragraphs = []
|
||||
|
||||
# Add heading to current_paragraphs. The content line gets
|
||||
# a markdown ``#`` prefix (capped at 6) via
|
||||
# render_heading_line; ``clean_heading_text`` is kept
|
||||
# for the heading field / stack / parent_headings below.
|
||||
current_paragraphs.append(
|
||||
{
|
||||
"text": render_heading_line(level, truncated_text),
|
||||
"para_id": heading_para_id,
|
||||
"is_table": False,
|
||||
}
|
||||
)
|
||||
|
||||
# Update current_heading and parent_headings for the FIRST heading in a block
|
||||
# (when current_paragraphs just had this heading added as its first element)
|
||||
if len(current_paragraphs) == 1:
|
||||
current_heading = clean_heading_text
|
||||
current_heading_level = level # Only set level when setting heading
|
||||
# Parent headings = all headings from levels strictly less than current level
|
||||
# Sort by level to maintain hierarchy order
|
||||
current_parent_headings = [
|
||||
current_heading_stack[lvl]
|
||||
for lvl in sorted(current_heading_stack.keys())
|
||||
if lvl < level
|
||||
]
|
||||
|
||||
# Update heading stack: remove current level and all lower levels, then add current
|
||||
current_heading_stack = {
|
||||
k: v for k, v in current_heading_stack.items() if k < level
|
||||
}
|
||||
current_heading_stack[level] = clean_heading_text
|
||||
|
||||
# Carry the body text that followed a soft break in an over-long
|
||||
# heading paragraph as a regular body paragraph in the same block.
|
||||
if demoted_body_text:
|
||||
current_paragraphs.append(
|
||||
{
|
||||
"text": demoted_body_text,
|
||||
"para_id": heading_para_id,
|
||||
"is_table": False,
|
||||
}
|
||||
)
|
||||
else:
|
||||
# Regular paragraph content
|
||||
para_id = extract_para_id(element)
|
||||
if parse_warnings is not None and not para_id:
|
||||
parse_warnings["missing_paraid_count"] = (
|
||||
parse_warnings.get("missing_paraid_count", 0) + 1
|
||||
)
|
||||
|
||||
# Store paragraph with metadata for potential splitting
|
||||
current_paragraphs.append(
|
||||
{"text": full_text, "para_id": para_id, "is_table": False}
|
||||
)
|
||||
|
||||
# Check for paragraph-level section break (after processing paragraph)
|
||||
# sectPr in pPr means this paragraph ends a section
|
||||
pPr = element.find(
|
||||
"{http://schemas.openxmlformats.org/wordprocessingml/2006/main}pPr"
|
||||
)
|
||||
if pPr is not None:
|
||||
sectPr = pPr.find(
|
||||
"{http://schemas.openxmlformats.org/wordprocessingml/2006/main}sectPr"
|
||||
)
|
||||
if sectPr is not None:
|
||||
# Section break after this paragraph - reset tracking
|
||||
resolver.reset_tracking_state()
|
||||
|
||||
elif tag == "tbl": # Table
|
||||
# Reset numbering tracking before table (table start boundary)
|
||||
resolver.reset_tracking_state()
|
||||
|
||||
# Directly create Table object from XML element to avoid index mismatch
|
||||
# (doc.tables may have different order due to nested tables)
|
||||
from docx.table import Table
|
||||
|
||||
table = Table(element, doc)
|
||||
table_metadata = TableExtractor.extract_with_metadata(
|
||||
table,
|
||||
numbering_resolver=resolver,
|
||||
drawing_context=drawing_context,
|
||||
)
|
||||
|
||||
table_rows = table_metadata["rows"]
|
||||
para_ids = table_metadata["para_ids"]
|
||||
para_ids_end = table_metadata["para_ids_end"] # Last paraId in each cell
|
||||
header_indices = table_metadata["header_indices"]
|
||||
|
||||
# Skip tables whose every cell is whitespace-only — otherwise an
|
||||
# empty `<table>[[""]]</table>` placeholder would leak into block
|
||||
# content and a useless IRTable would appear in tables.json.
|
||||
if _is_table_empty(table_rows):
|
||||
resolver.reset_tracking_state()
|
||||
continue
|
||||
|
||||
# Count tables whose cells carry no w14:paraId. Legacy / non-Word
|
||||
# docx authors omit these attributes; we no longer fail-fast, but
|
||||
# the adapter surfaces a single warning so the user knows the edit
|
||||
# range hints will be missing for these tables.
|
||||
if parse_warnings is not None and not _table_has_any_paraid(para_ids):
|
||||
parse_warnings["missing_paraid_count"] = (
|
||||
parse_warnings.get("missing_paraid_count", 0) + 1
|
||||
)
|
||||
|
||||
# Convert table to JSON
|
||||
table_json = json.dumps(table_rows, ensure_ascii=False)
|
||||
|
||||
# Extract cross-page repeating header rows (w:tblHeader) once per
|
||||
# table so we can surface them to the sidecar via the block-level
|
||||
# ``table_headers`` list.
|
||||
header_rows = []
|
||||
if header_indices:
|
||||
header_rows = [
|
||||
table_rows[idx] for idx in header_indices if idx < len(table_rows)
|
||||
]
|
||||
header_rows_or_none = header_rows if header_rows else None
|
||||
|
||||
# Emit the whole table as a single <table> placeholder. Token-based
|
||||
# table row splitting is the downstream chunker's responsibility.
|
||||
# Use first valid paraId from table, and last valid paraId (from
|
||||
# para_ids_end) for uuid_end.
|
||||
table_para_id = find_first_valid_para_id(para_ids)
|
||||
table_para_id_end = find_last_valid_para_id(para_ids_end)
|
||||
current_paragraphs.append(
|
||||
{
|
||||
"text": f"<table>{table_json}</table>",
|
||||
"para_id": table_para_id,
|
||||
"para_id_end": table_para_id_end, # Store end paraId for uuid_end calculation
|
||||
"is_table": True,
|
||||
"_table_header": header_rows_or_none,
|
||||
}
|
||||
)
|
||||
|
||||
# Reset numbering tracking after table (table end boundary)
|
||||
resolver.reset_tracking_state()
|
||||
|
||||
# Save final block
|
||||
_flush_current_block(
|
||||
blocks,
|
||||
current_heading,
|
||||
current_paragraphs,
|
||||
current_parent_headings,
|
||||
current_heading_level,
|
||||
)
|
||||
|
||||
return blocks
|
||||
@@ -0,0 +1,96 @@
|
||||
"""Native DOCX engine adapter (implements NativeParserBase hooks)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from lightrag.constants import PARSER_ENGINE_NATIVE
|
||||
from lightrag.parser.native_base import NativeParserBase
|
||||
from lightrag.utils import logger
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from lightrag.sidecar.ir import IRDoc
|
||||
|
||||
|
||||
class NativeDocxParser(NativeParserBase):
|
||||
"""Native DOCX parser for LightRAG's production parsing path.
|
||||
|
||||
``extract_docx_blocks`` performs only heading-driven structural splitting
|
||||
(one block per DOCX heading). Block sizing is intentionally left to the
|
||||
downstream paragraph-semantic chunker, so this parser emits the
|
||||
one-heading-one-block sidecar contract that chunking consumes.
|
||||
"""
|
||||
|
||||
engine_name = PARSER_ENGINE_NATIVE
|
||||
sidecar_path_style = "basename_only" # legacy native docx convention
|
||||
empty_content_label = "DOCX"
|
||||
|
||||
def validate_source(self, source: Path, file_path: str) -> None:
|
||||
if not (
|
||||
source.exists() and source.is_file() and source.suffix.lower() == ".docx"
|
||||
):
|
||||
raise ValueError(
|
||||
f"Native parser does not support pending file: {file_path}"
|
||||
)
|
||||
|
||||
def extract(
|
||||
self, source: Path, *, parsed_dir: Path, asset_dir: Path, base_name: str
|
||||
) -> tuple[list[dict[str, Any]], dict[str, Any], dict[str, Any]]:
|
||||
"""Extract heading-scoped DOCX blocks (sizing left to the chunker)."""
|
||||
|
||||
from lightrag.parser.docx.drawing_image_extractor import (
|
||||
DrawingExtractionContext,
|
||||
load_relationships,
|
||||
)
|
||||
from lightrag.parser.docx.parse_document import extract_docx_blocks
|
||||
|
||||
ctx = DrawingExtractionContext(
|
||||
docx_path=source,
|
||||
blocks_output_path=parsed_dir / f"{base_name}.blocks.jsonl",
|
||||
export_dir_name=asset_dir.name,
|
||||
export_dir_path=asset_dir,
|
||||
)
|
||||
load_relationships(ctx)
|
||||
warnings: dict[str, Any] = {}
|
||||
metadata: dict[str, Any] = {}
|
||||
blocks = extract_docx_blocks(
|
||||
str(source),
|
||||
drawing_context=ctx,
|
||||
parse_warnings=warnings,
|
||||
parse_metadata=metadata,
|
||||
)
|
||||
return blocks, warnings, metadata
|
||||
|
||||
def build_ir(
|
||||
self,
|
||||
blocks: list[dict[str, Any]],
|
||||
*,
|
||||
document_name: str,
|
||||
asset_dir_name: str,
|
||||
metadata: dict[str, Any],
|
||||
) -> "IRDoc":
|
||||
from lightrag.parser.docx.ir_builder import NativeDocxIRBuilder
|
||||
|
||||
return NativeDocxIRBuilder().normalize(
|
||||
blocks,
|
||||
document_name=document_name,
|
||||
asset_dir_name=asset_dir_name,
|
||||
parse_metadata=metadata,
|
||||
)
|
||||
|
||||
def surface_warnings(
|
||||
self, warnings: dict[str, Any], source: Path
|
||||
) -> dict[str, Any] | None:
|
||||
missing = int(warnings.get("missing_paraid_count", 0) or 0)
|
||||
if missing > 0:
|
||||
# Surface once per document; affected blocks emit
|
||||
# ``positions: [{"type": "paraid", "range": null}]``.
|
||||
logger.warning(
|
||||
"[parse_native] %s: %d paragraphs lack paraId; "
|
||||
"Re-saving file in Word 2013+ to regenerate ids.",
|
||||
source.name,
|
||||
missing,
|
||||
)
|
||||
return {"missing_paraid_count": missing}
|
||||
return None
|
||||
@@ -0,0 +1,419 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
ABOUTME: Extracts tables from DOCX with proper merged cell handling
|
||||
ABOUTME: Vertically merged cells: content repeated in all rows with shared paraId
|
||||
ABOUTME: Horizontally merged cells: content in first cell only
|
||||
ABOUTME: Preserves superscript/subscript formatting with <sup>/<sub> markup
|
||||
"""
|
||||
|
||||
from docx.table import Table
|
||||
from docx.oxml.ns import qn
|
||||
from typing import List
|
||||
|
||||
from .drawing_image_extractor import (
|
||||
DrawingExtractionContext,
|
||||
extract_drawing_placeholder_from_element,
|
||||
extract_vml_image_placeholder_from_element,
|
||||
)
|
||||
|
||||
# Keep in sync with parse_document._SKIP_PARAGRAPH_TAGS — duplicated here to
|
||||
# avoid a circular import between parse_document and table_extractor.
|
||||
_SKIP_PARAGRAPH_TAGS = frozenset(
|
||||
{
|
||||
"del",
|
||||
"moveFrom",
|
||||
"commentRangeStart",
|
||||
"commentRangeEnd",
|
||||
"commentReference",
|
||||
"annotationRef",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def extract_text_from_run_table(
|
||||
run_elem,
|
||||
qn_func,
|
||||
drawing_context: DrawingExtractionContext = None,
|
||||
) -> str:
|
||||
"""
|
||||
Extract text from a run element in table cell, preserving superscript/subscript with markup.
|
||||
|
||||
Converts Word formatting to HTML-like tags:
|
||||
- Superscript: <sup>text</sup>
|
||||
- Subscript: <sub>text</sub>
|
||||
- Normal text: unchanged
|
||||
|
||||
Args:
|
||||
run_elem: lxml run element (w:r)
|
||||
qn_func: qn function for namespace handling
|
||||
|
||||
Returns:
|
||||
Text string with <sup>/<sub> markup for formatted portions
|
||||
"""
|
||||
text = ""
|
||||
|
||||
# Check for vertAlign in rPr (superscript/subscript)
|
||||
vert_align = None
|
||||
rPr = run_elem.find(qn_func("w:rPr"))
|
||||
if rPr is not None:
|
||||
vert_elem = rPr.find(qn_func("w:vertAlign"))
|
||||
if vert_elem is not None:
|
||||
vert_align = vert_elem.get(qn_func("w:val"))
|
||||
|
||||
# Extract text content from run children
|
||||
for child in run_elem:
|
||||
tag = child.tag.split("}")[-1] # Remove namespace
|
||||
if tag == "t" and child.text:
|
||||
text += child.text
|
||||
elif tag == "tab":
|
||||
text += "\t"
|
||||
elif tag == "br":
|
||||
# Handle line breaks - textWrapping or no type = soft line break
|
||||
br_type = child.get(qn_func("w:type"))
|
||||
if br_type in (None, "textWrapping"):
|
||||
text += "\n"
|
||||
# Skip page and column breaks (layout elements)
|
||||
elif tag == "drawing":
|
||||
text += extract_drawing_placeholder_from_element(
|
||||
child,
|
||||
context=drawing_context,
|
||||
include_extended_attrs=True,
|
||||
)
|
||||
elif tag in ("pict", "object"):
|
||||
text += extract_vml_image_placeholder_from_element(
|
||||
child,
|
||||
context=drawing_context,
|
||||
include_extended_attrs=True,
|
||||
)
|
||||
|
||||
# Apply superscript/subscript markup if needed
|
||||
if text and vert_align == "superscript":
|
||||
return f"<sup>{text}</sup>"
|
||||
elif text and vert_align == "subscript":
|
||||
return f"<sub>{text}</sub>"
|
||||
|
||||
return text
|
||||
|
||||
|
||||
def extract_paragraph_content_table(
|
||||
para_elem,
|
||||
qn_func,
|
||||
drawing_context: DrawingExtractionContext = None,
|
||||
) -> str:
|
||||
"""
|
||||
Extract text and equations from a table cell paragraph in document order.
|
||||
|
||||
Handles w:r (text runs), m:oMath (inline equations), and m:oMathPara
|
||||
(block equations). Recurses into container elements (e.g., w:hyperlink,
|
||||
w:ins, w:sdt, w:fldSimple, w:smartTag) to avoid dropping content.
|
||||
|
||||
Args:
|
||||
para_elem: lxml paragraph element (w:p)
|
||||
qn_func: qn function for namespace handling
|
||||
|
||||
Returns:
|
||||
Text string with equations wrapped in <equation> tags
|
||||
"""
|
||||
parts = []
|
||||
|
||||
def append_from(node) -> None:
|
||||
tag = node.tag.split("}")[-1]
|
||||
# Drop tracked-change deletions (w:del/w:moveFrom) and comment markers
|
||||
# (w:commentRangeStart/End, w:commentReference, w:annotationRef) so the
|
||||
# output only contains the final revised text without annotation glyphs.
|
||||
if tag in _SKIP_PARAGRAPH_TAGS:
|
||||
return
|
||||
if tag == "r":
|
||||
parts.append(
|
||||
extract_text_from_run_table(
|
||||
node,
|
||||
qn_func,
|
||||
drawing_context=drawing_context,
|
||||
)
|
||||
)
|
||||
return
|
||||
if tag == "oMath":
|
||||
from omml import convert_omml_to_latex
|
||||
|
||||
latex = convert_omml_to_latex(node)
|
||||
if latex:
|
||||
parts.append(f"<equation>{latex}</equation>")
|
||||
return
|
||||
if tag == "oMathPara":
|
||||
from omml import convert_omml_to_latex
|
||||
|
||||
for omath in node:
|
||||
if omath.tag.split("}")[-1] == "oMath":
|
||||
latex = convert_omml_to_latex(omath)
|
||||
if latex:
|
||||
parts.append(f"<equation>{latex}</equation>")
|
||||
return
|
||||
for child in node:
|
||||
append_from(child)
|
||||
|
||||
for child in para_elem:
|
||||
append_from(child)
|
||||
|
||||
return "".join(parts)
|
||||
|
||||
|
||||
class TableExtractor:
|
||||
"""
|
||||
Extract table content handling merged cells correctly.
|
||||
|
||||
Merged cells in DOCX:
|
||||
- Horizontal: w:gridSpan specifies how many columns cell spans
|
||||
- Vertical: w:vMerge with val="restart" starts merge, subsequent cells continue
|
||||
|
||||
Output format:
|
||||
- 2D list of strings
|
||||
- Vertically merged cells: content repeated in all rows, all rows use the same paraId (from start cell)
|
||||
- Horizontally merged cells: content in left-most position only, other positions empty
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def extract(
|
||||
table: Table,
|
||||
numbering_resolver=None,
|
||||
drawing_context: DrawingExtractionContext = None,
|
||||
) -> List[List[str]]:
|
||||
"""
|
||||
Extract table to 2D string array.
|
||||
|
||||
Args:
|
||||
table: python-docx Table object
|
||||
numbering_resolver: Optional NumberingResolver for extracting numbering
|
||||
|
||||
Returns:
|
||||
List of rows, each row is list of cell text strings
|
||||
"""
|
||||
result = TableExtractor.extract_with_metadata(
|
||||
table,
|
||||
numbering_resolver=numbering_resolver,
|
||||
drawing_context=drawing_context,
|
||||
)
|
||||
return result["rows"]
|
||||
|
||||
@staticmethod
|
||||
def extract_with_metadata(
|
||||
table: Table,
|
||||
numbering_resolver=None,
|
||||
drawing_context: DrawingExtractionContext = None,
|
||||
) -> dict:
|
||||
"""
|
||||
Extract table to 2D string array with metadata (paraIds, header info).
|
||||
|
||||
Vertical merge behavior:
|
||||
- All rows in a vertically merged region share the same content
|
||||
- All rows use the paraId from the merge start cell (for precise edit targeting)
|
||||
|
||||
Args:
|
||||
table: python-docx Table object
|
||||
numbering_resolver: Optional NumberingResolver for extracting numbering
|
||||
|
||||
Returns:
|
||||
Dict with:
|
||||
- rows: 2D list of cell text strings
|
||||
- para_ids: 2D list of paraIds (first paraId in each cell, or None)
|
||||
For vertically merged cells, all rows share the start cell's paraId
|
||||
- para_ids_end: 2D list of paraIds (last paraId in each cell, or None)
|
||||
For vertically merged cells, all rows share the start cell's paraId
|
||||
- header_indices: List of row indices marked as table headers
|
||||
"""
|
||||
tbl = table._tbl
|
||||
|
||||
# Get number of columns from tblGrid
|
||||
tbl_grid = tbl.find(qn("w:tblGrid"))
|
||||
num_cols = 0
|
||||
if tbl_grid is not None:
|
||||
num_cols = len(tbl_grid.findall(qn("w:gridCol")))
|
||||
|
||||
if num_cols == 0:
|
||||
return {
|
||||
"rows": [],
|
||||
"para_ids": [],
|
||||
"para_ids_end": [],
|
||||
"header_indices": [],
|
||||
}
|
||||
|
||||
# Detect header rows using w:tblHeader attribute
|
||||
header_indices = []
|
||||
for idx, tr in enumerate(tbl.findall(qn("w:tr"))):
|
||||
trPr = tr.find(qn("w:trPr"))
|
||||
if trPr is not None:
|
||||
tbl_header = trPr.find(qn("w:tblHeader"))
|
||||
if tbl_header is not None:
|
||||
header_indices.append(idx)
|
||||
|
||||
# Process each row by directly iterating <w:tr> elements
|
||||
grid = []
|
||||
para_ids_grid = []
|
||||
para_ids_end_grid = [] # Track last paraId in each cell
|
||||
vmerge_content = {} # Track vertical merge by column: {col: {'text': str, 'para_id': str, 'para_id_end': str}}
|
||||
|
||||
for tr in tbl.findall(qn("w:tr")):
|
||||
row_data = [""] * num_cols # Pre-fill with empty strings
|
||||
row_para_ids = [None] * num_cols # Pre-fill with None
|
||||
row_para_ids_end = [None] * num_cols # Pre-fill with None for last paraId
|
||||
grid_col = 0
|
||||
|
||||
# Iterate actual <w:tc> elements (each physical cell appears once)
|
||||
for tc in tr.findall(qn("w:tc")):
|
||||
# Reset numbering state when cell changes to prevent incorrect continuation
|
||||
if numbering_resolver is not None:
|
||||
numbering_resolver.reset_tracking_state()
|
||||
|
||||
tcPr = tc.find(qn("w:tcPr"))
|
||||
|
||||
# Check gridSpan (horizontal merge)
|
||||
grid_span = 1
|
||||
if tcPr is not None:
|
||||
gs = tcPr.find(qn("w:gridSpan"))
|
||||
if gs is not None:
|
||||
grid_span = int(gs.get(qn("w:val")))
|
||||
|
||||
# Check vMerge (vertical merge)
|
||||
vmerge_elem = None
|
||||
vmerge_val = None
|
||||
if tcPr is not None:
|
||||
vmerge_elem = tcPr.find(qn("w:vMerge"))
|
||||
if vmerge_elem is not None:
|
||||
vmerge_val = vmerge_elem.get(
|
||||
qn("w:val")
|
||||
) # 'restart' or None (means 'continue')
|
||||
|
||||
# Determine vMerge status
|
||||
is_vmerge_restart = vmerge_elem is not None and vmerge_val == "restart"
|
||||
is_vmerge_continue = vmerge_elem is not None and vmerge_val in (
|
||||
None,
|
||||
"continue",
|
||||
)
|
||||
is_normal_cell = vmerge_elem is None
|
||||
|
||||
cell_text = ""
|
||||
cell_para_id = None
|
||||
cell_para_id_end = None # Track last paraId in cell
|
||||
|
||||
# Handle different vMerge cases
|
||||
if is_vmerge_restart or is_normal_cell:
|
||||
# Extract content for restart or normal cells
|
||||
# Get cell text with numbering support and format preservation
|
||||
if numbering_resolver is not None:
|
||||
# Extract text with numbering labels and superscript/subscript markup
|
||||
cell_paragraphs = []
|
||||
for para_elem in tc.findall(qn("w:p")):
|
||||
# Capture paraId from each paragraph
|
||||
para_id_attr = para_elem.get(
|
||||
"{http://schemas.microsoft.com/office/word/2010/wordml}paraId"
|
||||
)
|
||||
if para_id_attr:
|
||||
if cell_para_id is None:
|
||||
cell_para_id = para_id_attr # First paraId
|
||||
cell_para_id_end = (
|
||||
para_id_attr # Always update to get last
|
||||
)
|
||||
|
||||
# Get text content with format preservation (superscript/subscript/equations)
|
||||
para_text = extract_paragraph_content_table(
|
||||
para_elem,
|
||||
qn,
|
||||
drawing_context=drawing_context,
|
||||
)
|
||||
|
||||
# Get numbering label
|
||||
label = numbering_resolver.get_label(para_elem)
|
||||
|
||||
# Combine label and text
|
||||
if label:
|
||||
full_text = f"{label} {para_text}".strip()
|
||||
else:
|
||||
full_text = para_text.strip()
|
||||
|
||||
if full_text:
|
||||
cell_paragraphs.append(full_text)
|
||||
|
||||
cell_text = "\n".join(cell_paragraphs).replace("\x07", "")
|
||||
else:
|
||||
# Fallback to simple text extraction with format preservation
|
||||
# Cannot use cell.text here, must extract from XML
|
||||
para_texts = []
|
||||
for para_elem in tc.findall(qn("w:p")):
|
||||
# Capture paraId from each paragraph
|
||||
para_id_attr = para_elem.get(
|
||||
"{http://schemas.microsoft.com/office/word/2010/wordml}paraId"
|
||||
)
|
||||
if para_id_attr:
|
||||
if cell_para_id is None:
|
||||
cell_para_id = para_id_attr # First paraId
|
||||
cell_para_id_end = (
|
||||
para_id_attr # Always update to get last
|
||||
)
|
||||
|
||||
# Extract text with format preservation (superscript/subscript/equations)
|
||||
para_text = extract_paragraph_content_table(
|
||||
para_elem,
|
||||
qn,
|
||||
drawing_context=drawing_context,
|
||||
)
|
||||
|
||||
if para_text:
|
||||
para_texts.append(para_text.strip())
|
||||
cell_text = "\n".join(para_texts).replace("\x07", "")
|
||||
|
||||
# Store content and paraIds for vMerge restart
|
||||
if is_vmerge_restart:
|
||||
vmerge_content[grid_col] = {
|
||||
"text": cell_text,
|
||||
"para_id": cell_para_id,
|
||||
"para_id_end": cell_para_id_end,
|
||||
}
|
||||
elif is_normal_cell:
|
||||
# For normal cells: if empty and we have active vMerge, copy all from start
|
||||
# If non-empty, this ends the vMerge region
|
||||
if not cell_text and grid_col in vmerge_content:
|
||||
# Empty cell in vMerge region - copy content and paraIds from start
|
||||
cell_text = vmerge_content[grid_col]["text"]
|
||||
cell_para_id = vmerge_content[grid_col]["para_id"]
|
||||
cell_para_id_end = vmerge_content[grid_col]["para_id_end"]
|
||||
elif cell_text:
|
||||
# Non-empty cell - this ends the vMerge for this column
|
||||
vmerge_content.pop(grid_col, None)
|
||||
|
||||
elif is_vmerge_continue:
|
||||
# Copy content and para_id from previous merge start
|
||||
# But extract actual para_id_end from this continue cell for range boundary
|
||||
if grid_col in vmerge_content:
|
||||
cell_text = vmerge_content[grid_col]["text"]
|
||||
cell_para_id = vmerge_content[grid_col][
|
||||
"para_id"
|
||||
] # Use restart's paraId for edit targeting
|
||||
|
||||
# Extract actual paraId from this continue cell for uuid_end (range boundary)
|
||||
for para_elem in tc.findall(qn("w:p")):
|
||||
para_id_attr = para_elem.get(
|
||||
"{http://schemas.microsoft.com/office/word/2010/wordml}paraId"
|
||||
)
|
||||
if para_id_attr:
|
||||
cell_para_id_end = (
|
||||
para_id_attr # Use actual paraId for range boundary
|
||||
)
|
||||
|
||||
# Place content at starting grid position only
|
||||
if grid_col < num_cols:
|
||||
row_data[grid_col] = cell_text
|
||||
row_para_ids[grid_col] = cell_para_id
|
||||
row_para_ids_end[grid_col] = cell_para_id_end
|
||||
|
||||
# Move grid position by gridSpan
|
||||
grid_col += grid_span
|
||||
|
||||
grid.append(row_data)
|
||||
para_ids_grid.append(row_para_ids)
|
||||
para_ids_end_grid.append(row_para_ids_end)
|
||||
|
||||
return {
|
||||
"rows": grid,
|
||||
"para_ids": para_ids_grid,
|
||||
"para_ids_end": para_ids_end_grid,
|
||||
"header_indices": header_indices,
|
||||
}
|
||||
@@ -0,0 +1,791 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
ABOUTME: Shared token estimation utilities for audit scripts
|
||||
ABOUTME: XML sanitization helpers for document processing
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
|
||||
try:
|
||||
from google import genai
|
||||
from google.genai import types
|
||||
|
||||
HAS_GEMINI = True
|
||||
except ImportError: # pragma: no cover - optional dependency
|
||||
genai = None
|
||||
types = None
|
||||
HAS_GEMINI = False
|
||||
|
||||
try:
|
||||
import openai
|
||||
|
||||
HAS_OPENAI = True
|
||||
except ImportError: # pragma: no cover - optional dependency
|
||||
openai = None
|
||||
HAS_OPENAI = False
|
||||
|
||||
|
||||
def estimate_tokens(text: str) -> int:
|
||||
"""
|
||||
Estimate token count for LLM context management.
|
||||
|
||||
Uses a weighted formula based on character types:
|
||||
- Chinese characters: ~0.75 tokens per character (subword tokenization)
|
||||
- JSON structural characters (brackets, quotes, commas): ~1 tokens per character
|
||||
- Other characters (English, numbers, symbols): ~0.4 tokens per character (~3 chars/token)
|
||||
|
||||
Includes 5% buffer and safety offset for special formatting and system prompt overhead.
|
||||
|
||||
Args:
|
||||
text: Input text to estimate tokens for
|
||||
|
||||
Returns:
|
||||
int: Estimated token count
|
||||
"""
|
||||
if not text:
|
||||
return 0
|
||||
|
||||
chinese_count = len(re.findall(r"[\u4e00-\u9fa5]", text))
|
||||
json_chars_count = len(re.findall(r'[\[\]",{}]', text))
|
||||
other_count = len(text) - chinese_count - json_chars_count
|
||||
|
||||
base_estimate = (
|
||||
(chinese_count * 0.75) + (json_chars_count * 1) + (other_count * 0.4)
|
||||
)
|
||||
final_tokens = int(base_estimate * 1.05) + 2
|
||||
return final_tokens
|
||||
|
||||
|
||||
def sanitize_xml_string(text: str) -> str:
|
||||
"""
|
||||
Remove control characters that are illegal in XML 1.0.
|
||||
|
||||
XML 1.0 allows: #x9 (tab), #xA (LF), #xD (CR), and #x20-#xD7FF, #xE000-#xFFFD, #x10000-#x10FFFF
|
||||
This function removes all other control characters (0x00-0x08, 0x0B, 0x0C, 0x0E-0x1F).
|
||||
|
||||
Args:
|
||||
text: Text that may contain control characters
|
||||
|
||||
Returns:
|
||||
Sanitized text safe for XML. Returns input unchanged if not a non-empty string.
|
||||
"""
|
||||
if not text or not isinstance(text, str):
|
||||
return text
|
||||
# Build a translation table to remove illegal control characters
|
||||
# Keep: \t (0x09), \n (0x0A), \r (0x0D)
|
||||
# Remove: 0x00-0x08, 0x0B, 0x0C, 0x0E-0x1F
|
||||
illegal_chars = "".join(chr(c) for c in range(0x20) if c not in (0x09, 0x0A, 0x0D))
|
||||
return text.translate(str.maketrans("", "", illegal_chars))
|
||||
|
||||
|
||||
def is_vertex_ai_mode() -> bool:
|
||||
"""
|
||||
Check if Vertex AI mode is enabled via environment variable.
|
||||
|
||||
Returns:
|
||||
True if GOOGLE_GENAI_USE_VERTEXAI is set to 'true', False otherwise
|
||||
"""
|
||||
return os.getenv("GOOGLE_GENAI_USE_VERTEXAI", "").lower() == "true"
|
||||
|
||||
|
||||
def create_gemini_client(use_async: bool = False):
|
||||
"""
|
||||
Create Gemini client for AI Studio or Vertex AI.
|
||||
|
||||
Supports two modes:
|
||||
- AI Studio (default): Uses GOOGLE_API_KEY for authentication
|
||||
- Vertex AI: Uses ADC (GOOGLE_APPLICATION_CREDENTIALS or gcloud auth)
|
||||
|
||||
Environment variables for Vertex AI mode:
|
||||
- GOOGLE_GENAI_USE_VERTEXAI: Set to 'true' to enable Vertex AI mode
|
||||
- GOOGLE_CLOUD_PROJECT: Required GCP project ID
|
||||
- GOOGLE_CLOUD_LOCATION: Optional region (default: us-central1)
|
||||
- GOOGLE_VERTEX_BASE_URL: Optional custom API endpoint (for API gateway proxies)
|
||||
- GOOGLE_APPLICATION_CREDENTIALS: Path to service account JSON (or use gcloud auth)
|
||||
|
||||
Args:
|
||||
use_async: If True, return the async client (.aio), otherwise return sync client
|
||||
|
||||
Returns:
|
||||
Gemini client instance (sync or async based on use_async parameter)
|
||||
|
||||
Raises:
|
||||
ValueError: If required environment variables are not set
|
||||
"""
|
||||
use_vertex = is_vertex_ai_mode()
|
||||
|
||||
if use_vertex:
|
||||
# Vertex AI mode - uses ADC (GOOGLE_APPLICATION_CREDENTIALS or gcloud auth)
|
||||
project = os.getenv("GOOGLE_CLOUD_PROJECT")
|
||||
location = os.getenv("GOOGLE_CLOUD_LOCATION", "us-central1")
|
||||
base_url = os.getenv("GOOGLE_VERTEX_BASE_URL")
|
||||
|
||||
if not project:
|
||||
raise ValueError(
|
||||
"GOOGLE_CLOUD_PROJECT is required for Vertex AI mode. "
|
||||
"Set GOOGLE_GENAI_USE_VERTEXAI=false to use AI Studio mode instead."
|
||||
)
|
||||
|
||||
# Build http_options only if custom base_url is specified
|
||||
http_options = None
|
||||
if base_url:
|
||||
http_options = {"base_url": base_url}
|
||||
|
||||
# Note: ADC handles authentication automatically
|
||||
# via GOOGLE_APPLICATION_CREDENTIALS env var or gcloud auth
|
||||
client = genai.Client(
|
||||
vertexai=True, project=project, location=location, http_options=http_options
|
||||
)
|
||||
else:
|
||||
# AI Studio mode - requires API key
|
||||
api_key = os.getenv("GOOGLE_API_KEY")
|
||||
if not api_key:
|
||||
raise ValueError(
|
||||
"GOOGLE_API_KEY is required for AI Studio mode. "
|
||||
"Set GOOGLE_GENAI_USE_VERTEXAI=true and configure GCP credentials for Vertex AI mode."
|
||||
)
|
||||
|
||||
client = genai.Client(api_key=api_key)
|
||||
|
||||
# Return async or sync client based on parameter
|
||||
return client.aio if use_async else client
|
||||
|
||||
|
||||
def get_gemini_provider_name() -> str:
|
||||
"""
|
||||
Get the Gemini provider name based on current mode.
|
||||
|
||||
Returns:
|
||||
Provider name string for display purposes
|
||||
"""
|
||||
if is_vertex_ai_mode():
|
||||
project = os.getenv("GOOGLE_CLOUD_PROJECT", "unknown")
|
||||
location = os.getenv("GOOGLE_CLOUD_LOCATION", "us-central1")
|
||||
return f"Google Gemini (Vertex AI: {project}/{location})"
|
||||
return "Google Gemini (AI Studio)"
|
||||
|
||||
|
||||
def create_openai_client(use_async: bool = True):
|
||||
"""
|
||||
Create OpenAI client with optional custom base URL.
|
||||
|
||||
Environment variables:
|
||||
- OPENAI_API_KEY: Required API key
|
||||
- OPENAI_BASE_URL: Optional custom API endpoint (for proxies, Azure, etc.)
|
||||
|
||||
Args:
|
||||
use_async: If True, return AsyncOpenAI, otherwise return OpenAI
|
||||
|
||||
Returns:
|
||||
OpenAI client instance (async or sync based on use_async parameter)
|
||||
|
||||
Raises:
|
||||
ValueError: If OPENAI_API_KEY is not set
|
||||
"""
|
||||
if not HAS_OPENAI:
|
||||
raise ValueError("openai library is not installed.")
|
||||
api_key = os.getenv("OPENAI_API_KEY")
|
||||
if not api_key:
|
||||
raise ValueError("OPENAI_API_KEY is required for OpenAI mode.")
|
||||
|
||||
base_url = os.getenv("OPENAI_BASE_URL")
|
||||
|
||||
if use_async:
|
||||
return openai.AsyncOpenAI(base_url=base_url)
|
||||
return openai.OpenAI(base_url=base_url)
|
||||
|
||||
|
||||
def get_openai_provider_name() -> str:
|
||||
"""
|
||||
Get the OpenAI provider name, including custom endpoint if configured.
|
||||
|
||||
Returns:
|
||||
Provider name string for display purposes
|
||||
"""
|
||||
base_url = os.getenv("OPENAI_BASE_URL")
|
||||
if base_url:
|
||||
return f"OpenAI (Custom: {base_url})"
|
||||
return "OpenAI"
|
||||
|
||||
|
||||
def is_openai_reasoning_model(model_name: str) -> bool:
|
||||
"""
|
||||
Check if the OpenAI model supports reasoning_effort parameter.
|
||||
|
||||
Models that support reasoning_effort:
|
||||
- o-series: o1, o3, o4 and their variants (o1-mini, o1-2024-12-17, etc.)
|
||||
- gpt-5 series: gpt-5, gpt-5.2, gpt-5-turbo, etc.
|
||||
|
||||
Non-reasoning models like gpt-4.1, gpt-4o, etc. will reject this parameter.
|
||||
|
||||
Handles proxy/router prefixes like "openai/o1-mini" or "openrouter/gpt-5.2".
|
||||
|
||||
Args:
|
||||
model_name: The OpenAI model name (may include path prefix)
|
||||
|
||||
Returns:
|
||||
True if the model supports reasoning_effort, False otherwise
|
||||
"""
|
||||
model_lower = model_name.lower()
|
||||
|
||||
# Handle proxy/router prefixes like "openai/o1-mini", "openrouter/gpt-5.2"
|
||||
# Extract the base model name after the last "/"
|
||||
if "/" in model_lower:
|
||||
model_lower = model_lower.rsplit("/", 1)[-1]
|
||||
|
||||
# Match o-series and gpt-5 series
|
||||
return model_lower.startswith(("o1", "o3", "o4", "gpt-5"))
|
||||
|
||||
|
||||
def is_openai_retryable(error: Exception) -> bool:
|
||||
"""
|
||||
Determine if an OpenAI error should be retried.
|
||||
|
||||
Non-retryable errors:
|
||||
- AuthenticationError (401): Invalid API key
|
||||
- PermissionDeniedError (403): No access to resource
|
||||
- BadRequestError (400): Invalid request format
|
||||
- NotFoundError (404): Model or resource not found
|
||||
|
||||
Retryable errors:
|
||||
- RateLimitError (429): Rate limit exceeded
|
||||
- APIConnectionError: Network issues
|
||||
- InternalServerError (500): Server errors
|
||||
- APIStatusError with 502, 503, 504: Gateway/service errors
|
||||
|
||||
Args:
|
||||
error: The exception from OpenAI API call
|
||||
|
||||
Returns:
|
||||
True if the error should be retried, False otherwise
|
||||
"""
|
||||
if not HAS_OPENAI:
|
||||
return True
|
||||
|
||||
# Authentication error - invalid API key (401)
|
||||
if isinstance(error, openai.AuthenticationError):
|
||||
return False
|
||||
|
||||
# Permission denied - no access to resource (403)
|
||||
if isinstance(error, openai.PermissionDeniedError):
|
||||
return False
|
||||
|
||||
# Bad request - invalid request format (400)
|
||||
if isinstance(error, openai.BadRequestError):
|
||||
return False
|
||||
|
||||
# Not found - model or resource doesn't exist (404)
|
||||
if isinstance(error, openai.NotFoundError):
|
||||
return False
|
||||
|
||||
# Rate limit exceeded - should retry with backoff (429)
|
||||
if isinstance(error, openai.RateLimitError):
|
||||
return True
|
||||
|
||||
# API connection error - network issues, should retry
|
||||
if isinstance(error, openai.APIConnectionError):
|
||||
return True
|
||||
|
||||
# Internal server error - should retry (500)
|
||||
if isinstance(error, openai.InternalServerError):
|
||||
return True
|
||||
|
||||
# For other APIStatusError, check HTTP status code
|
||||
if isinstance(error, openai.APIStatusError):
|
||||
# Retryable server-side errors
|
||||
return error.status_code in (429, 500, 502, 503, 504)
|
||||
|
||||
# For unknown errors, default to retry (network issues, timeouts, etc.)
|
||||
return True
|
||||
|
||||
|
||||
def is_gemini_retryable(error: Exception) -> bool:
|
||||
"""
|
||||
Determine if a Gemini error should be retried.
|
||||
|
||||
Uses string matching on error messages since google-genai may not have
|
||||
well-defined exception types for all error cases.
|
||||
|
||||
Non-retryable errors:
|
||||
- API key errors
|
||||
- Authentication/permission errors
|
||||
- Invalid request errors
|
||||
- Model not found errors
|
||||
- Billing/quota permanently exceeded
|
||||
|
||||
Retryable errors:
|
||||
- Rate limit (429)
|
||||
- Server errors (500, 502, 503, 504)
|
||||
- Timeout/connection errors
|
||||
|
||||
Args:
|
||||
error: The exception from Gemini API call
|
||||
|
||||
Returns:
|
||||
True if the error should be retried, False otherwise
|
||||
"""
|
||||
error_str = str(error).lower()
|
||||
|
||||
# API key / authentication errors - do not retry
|
||||
if "api_key" in error_str or "api key" in error_str:
|
||||
return False
|
||||
if "authentication" in error_str or "authenticate" in error_str:
|
||||
return False
|
||||
if "invalid_api_key" in error_str or "invalid api key" in error_str:
|
||||
return False
|
||||
|
||||
# Permission / forbidden errors - do not retry
|
||||
if "permission" in error_str and "denied" in error_str:
|
||||
return False
|
||||
if "forbidden" in error_str or "403" in error_str:
|
||||
return False
|
||||
|
||||
# Invalid request errors - do not retry
|
||||
if "invalid" in error_str and ("request" in error_str or "argument" in error_str):
|
||||
return False
|
||||
if "400" in error_str and "bad request" in error_str:
|
||||
return False
|
||||
|
||||
# Model not found - do not retry
|
||||
if "model" in error_str and ("not found" in error_str or "not exist" in error_str):
|
||||
return False
|
||||
if "404" in error_str:
|
||||
return False
|
||||
|
||||
# Billing / permanent quota errors - do not retry
|
||||
if "billing" in error_str:
|
||||
return False
|
||||
if "quota" in error_str and ("exceeded" in error_str or "exhausted" in error_str):
|
||||
# Check if it mentions billing which indicates permanent quota issue
|
||||
if "billing" in error_str or "payment" in error_str:
|
||||
return False
|
||||
# Temporary quota (rate limit) - should retry
|
||||
return True
|
||||
|
||||
# Rate limit errors - should retry (429)
|
||||
if "rate" in error_str and "limit" in error_str:
|
||||
return True
|
||||
if "429" in error_str or "resource_exhausted" in error_str:
|
||||
return True
|
||||
|
||||
# Server errors - should retry (500, 502, 503, 504)
|
||||
if any(code in error_str for code in ["500", "502", "503", "504"]):
|
||||
return True
|
||||
if "internal" in error_str and ("error" in error_str or "server" in error_str):
|
||||
return True
|
||||
if "service" in error_str and "unavailable" in error_str:
|
||||
return True
|
||||
if "gateway" in error_str:
|
||||
return True
|
||||
|
||||
# Timeout / connection errors - should retry
|
||||
if "timeout" in error_str or "timed out" in error_str:
|
||||
return True
|
||||
if "connection" in error_str:
|
||||
return True
|
||||
if "network" in error_str:
|
||||
return True
|
||||
|
||||
# Unknown errors - default to retry with limited attempts
|
||||
return True
|
||||
|
||||
|
||||
# JSON Schema for LLM structured output
|
||||
AUDIT_RESULT_SCHEMA = {
|
||||
"type": "object",
|
||||
"additionalProperties": False,
|
||||
"properties": {
|
||||
"is_violation": {
|
||||
"type": "boolean",
|
||||
"description": "Whether any violations were found",
|
||||
},
|
||||
"violations": {
|
||||
"type": "array",
|
||||
"description": "List of violations found",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"additionalProperties": False,
|
||||
"properties": {
|
||||
"rule_id": {
|
||||
"type": "string",
|
||||
"description": "ID of the violated rule (e.g., R001)",
|
||||
},
|
||||
"violation_text": {
|
||||
"type": "string",
|
||||
"description": "The problematic text directly verbatim quote from the source content, and not span multiple cells",
|
||||
},
|
||||
"violation_reason": {
|
||||
"type": "string",
|
||||
"description": "Explanation of why this violates the rule",
|
||||
},
|
||||
"fix_action": {
|
||||
"type": "string",
|
||||
"enum": ["replace", "manual"],
|
||||
"description": "Action type: replace substitutes text (including deletion-via-replace), manual requires human review",
|
||||
},
|
||||
"revised_text": {
|
||||
"type": "string",
|
||||
"description": "For replace: complete replacement text (including deletion-via-replace). For manual: additional guidance for human reviewer",
|
||||
},
|
||||
},
|
||||
"required": [
|
||||
"rule_id",
|
||||
"violation_text",
|
||||
"violation_reason",
|
||||
"fix_action",
|
||||
"revised_text",
|
||||
],
|
||||
},
|
||||
},
|
||||
},
|
||||
"required": ["is_violation", "violations"],
|
||||
}
|
||||
|
||||
# JSON Schema for global extraction output
|
||||
GLOBAL_EXTRACT_SCHEMA = {
|
||||
"type": "object",
|
||||
"additionalProperties": False,
|
||||
"properties": {
|
||||
"results": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"additionalProperties": False,
|
||||
"properties": {
|
||||
"rule_id": {"type": "string"},
|
||||
"extracted_results": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"additionalProperties": False,
|
||||
"properties": {
|
||||
"entity": {"type": "string"},
|
||||
"fields": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"additionalProperties": False,
|
||||
"properties": {
|
||||
"name": {"type": "string"},
|
||||
"value": {"type": "string"},
|
||||
"evidence": {"type": "string"},
|
||||
},
|
||||
"required": ["name", "value", "evidence"],
|
||||
},
|
||||
},
|
||||
},
|
||||
"required": ["entity", "fields"],
|
||||
},
|
||||
},
|
||||
},
|
||||
"required": ["rule_id", "extracted_results"],
|
||||
},
|
||||
}
|
||||
},
|
||||
"required": ["results"],
|
||||
}
|
||||
|
||||
# JSON Schema for global verification output
|
||||
GLOBAL_VERIFY_SCHEMA = {
|
||||
"type": "object",
|
||||
"additionalProperties": False,
|
||||
"properties": {
|
||||
"violations": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"additionalProperties": False,
|
||||
"properties": {
|
||||
"rule_id": {"type": "string"},
|
||||
"uuid": {"type": "string"},
|
||||
"uuid_end": {"type": "string"},
|
||||
"violation_text": {"type": "string"},
|
||||
"violation_reason": {"type": "string"},
|
||||
"fix_action": {"type": "string", "enum": ["replace", "manual"]},
|
||||
"revised_text": {"type": "string"},
|
||||
},
|
||||
"required": [
|
||||
"rule_id",
|
||||
"uuid",
|
||||
"uuid_end",
|
||||
"violation_text",
|
||||
"violation_reason",
|
||||
"fix_action",
|
||||
"revised_text",
|
||||
],
|
||||
},
|
||||
}
|
||||
},
|
||||
"required": ["violations"],
|
||||
}
|
||||
|
||||
|
||||
async def global_extract_gemini_async(
|
||||
user_prompt: str,
|
||||
system_prompt: str,
|
||||
model_name: str,
|
||||
client,
|
||||
thinking_level: str = None,
|
||||
thinking_budget: int = None,
|
||||
) -> dict:
|
||||
thinking_config = None
|
||||
if thinking_level and thinking_level.upper() in (
|
||||
"MINIMAL",
|
||||
"LOW",
|
||||
"MEDIUM",
|
||||
"HIGH",
|
||||
):
|
||||
level_map = {
|
||||
"MINIMAL": types.ThinkingLevel.MINIMAL,
|
||||
"LOW": types.ThinkingLevel.LOW,
|
||||
"MEDIUM": types.ThinkingLevel.MEDIUM,
|
||||
"HIGH": types.ThinkingLevel.HIGH,
|
||||
}
|
||||
thinking_config = types.ThinkingConfig(
|
||||
thinking_level=level_map[thinking_level.upper()]
|
||||
)
|
||||
elif thinking_budget is not None:
|
||||
thinking_config = types.ThinkingConfig(thinking_budget=int(thinking_budget))
|
||||
|
||||
config_params = {
|
||||
"system_instruction": system_prompt,
|
||||
"response_mime_type": "application/json",
|
||||
"response_schema": GLOBAL_EXTRACT_SCHEMA,
|
||||
}
|
||||
if thinking_config:
|
||||
config_params["thinking_config"] = thinking_config
|
||||
|
||||
response = await client.models.generate_content(
|
||||
model=model_name,
|
||||
contents=user_prompt,
|
||||
config=types.GenerateContentConfig(**config_params),
|
||||
)
|
||||
return json.loads(response.text)
|
||||
|
||||
|
||||
async def global_extract_openai_async(
|
||||
user_prompt: str,
|
||||
system_prompt: str,
|
||||
model_name: str,
|
||||
client,
|
||||
reasoning_effort: str = None,
|
||||
) -> dict:
|
||||
request_params = {
|
||||
"model": model_name,
|
||||
"messages": [
|
||||
{"role": "system", "content": system_prompt},
|
||||
{"role": "user", "content": user_prompt},
|
||||
],
|
||||
"response_format": {
|
||||
"type": "json_schema",
|
||||
"json_schema": {
|
||||
"name": "global_extract",
|
||||
"strict": True,
|
||||
"schema": GLOBAL_EXTRACT_SCHEMA,
|
||||
},
|
||||
},
|
||||
}
|
||||
if (
|
||||
reasoning_effort
|
||||
and reasoning_effort.lower() in ("low", "medium", "high")
|
||||
and is_openai_reasoning_model(model_name)
|
||||
):
|
||||
request_params["reasoning_effort"] = reasoning_effort.lower()
|
||||
|
||||
response = await client.chat.completions.create(**request_params)
|
||||
return json.loads(response.choices[0].message.content)
|
||||
|
||||
|
||||
async def global_verify_gemini_async(
|
||||
user_prompt: str,
|
||||
system_prompt: str,
|
||||
model_name: str,
|
||||
client,
|
||||
thinking_level: str = None,
|
||||
thinking_budget: int = None,
|
||||
) -> dict:
|
||||
thinking_config = None
|
||||
if thinking_level and thinking_level.upper() in (
|
||||
"MINIMAL",
|
||||
"LOW",
|
||||
"MEDIUM",
|
||||
"HIGH",
|
||||
):
|
||||
level_map = {
|
||||
"MINIMAL": types.ThinkingLevel.MINIMAL,
|
||||
"LOW": types.ThinkingLevel.LOW,
|
||||
"MEDIUM": types.ThinkingLevel.MEDIUM,
|
||||
"HIGH": types.ThinkingLevel.HIGH,
|
||||
}
|
||||
thinking_config = types.ThinkingConfig(
|
||||
thinking_level=level_map[thinking_level.upper()]
|
||||
)
|
||||
elif thinking_budget is not None:
|
||||
thinking_config = types.ThinkingConfig(thinking_budget=int(thinking_budget))
|
||||
|
||||
config_params = {
|
||||
"system_instruction": system_prompt,
|
||||
"response_mime_type": "application/json",
|
||||
"response_schema": GLOBAL_VERIFY_SCHEMA,
|
||||
}
|
||||
if thinking_config:
|
||||
config_params["thinking_config"] = thinking_config
|
||||
|
||||
response = await client.models.generate_content(
|
||||
model=model_name,
|
||||
contents=user_prompt,
|
||||
config=types.GenerateContentConfig(**config_params),
|
||||
)
|
||||
return json.loads(response.text)
|
||||
|
||||
|
||||
async def global_verify_openai_async(
|
||||
user_prompt: str,
|
||||
system_prompt: str,
|
||||
model_name: str,
|
||||
client,
|
||||
reasoning_effort: str = None,
|
||||
) -> dict:
|
||||
request_params = {
|
||||
"model": model_name,
|
||||
"messages": [
|
||||
{"role": "system", "content": system_prompt},
|
||||
{"role": "user", "content": user_prompt},
|
||||
],
|
||||
"response_format": {
|
||||
"type": "json_schema",
|
||||
"json_schema": {
|
||||
"name": "global_verify",
|
||||
"strict": True,
|
||||
"schema": GLOBAL_VERIFY_SCHEMA,
|
||||
},
|
||||
},
|
||||
}
|
||||
if (
|
||||
reasoning_effort
|
||||
and reasoning_effort.lower() in ("low", "medium", "high")
|
||||
and is_openai_reasoning_model(model_name)
|
||||
):
|
||||
request_params["reasoning_effort"] = reasoning_effort.lower()
|
||||
|
||||
response = await client.chat.completions.create(**request_params)
|
||||
return json.loads(response.choices[0].message.content)
|
||||
|
||||
|
||||
async def audit_block_gemini_async(
|
||||
user_prompt: str,
|
||||
system_prompt: str,
|
||||
model_name: str,
|
||||
client,
|
||||
thinking_level: str = None,
|
||||
thinking_budget: int = None,
|
||||
) -> dict:
|
||||
"""
|
||||
Audit a text block using Google Gemini with strict JSON mode (async version).
|
||||
|
||||
Args:
|
||||
user_prompt: User prompt to audit
|
||||
system_prompt: Cached system prompt with rules and instructions
|
||||
model_name: Gemini model to use
|
||||
client: Gemini async client instance (client.aio)
|
||||
thinking_level: Thinking level for Gemini 3 models (MINIMAL, LOW, MEDIUM, HIGH)
|
||||
thinking_budget: Thinking token budget for Gemini 2.5 models (integer)
|
||||
|
||||
Returns:
|
||||
Audit result dictionary
|
||||
"""
|
||||
# Build thinking config based on model and parameters
|
||||
thinking_config = None
|
||||
|
||||
if thinking_level and thinking_level.upper() in (
|
||||
"MINIMAL",
|
||||
"LOW",
|
||||
"MEDIUM",
|
||||
"HIGH",
|
||||
):
|
||||
# For Gemini 3 models
|
||||
level_map = {
|
||||
"MINIMAL": types.ThinkingLevel.MINIMAL,
|
||||
"LOW": types.ThinkingLevel.LOW,
|
||||
"MEDIUM": types.ThinkingLevel.MEDIUM,
|
||||
"HIGH": types.ThinkingLevel.HIGH,
|
||||
}
|
||||
thinking_config = types.ThinkingConfig(
|
||||
thinking_level=level_map[thinking_level.upper()]
|
||||
)
|
||||
elif thinking_budget is not None:
|
||||
# For Gemini 2.5 models
|
||||
thinking_config = types.ThinkingConfig(thinking_budget=int(thinking_budget))
|
||||
|
||||
config_params = {
|
||||
"system_instruction": system_prompt,
|
||||
"response_mime_type": "application/json",
|
||||
"response_schema": AUDIT_RESULT_SCHEMA,
|
||||
}
|
||||
|
||||
# Only add thinking_config if it's configured
|
||||
if thinking_config:
|
||||
config_params["thinking_config"] = thinking_config
|
||||
|
||||
response = await client.models.generate_content(
|
||||
model=model_name,
|
||||
contents=user_prompt,
|
||||
config=types.GenerateContentConfig(**config_params),
|
||||
)
|
||||
|
||||
# With structured output, response is guaranteed to be valid JSON
|
||||
result = json.loads(response.text)
|
||||
return result
|
||||
|
||||
|
||||
async def audit_block_openai_async(
|
||||
user_prompt: str,
|
||||
system_prompt: str,
|
||||
model_name: str,
|
||||
client,
|
||||
reasoning_effort: str = None,
|
||||
) -> dict:
|
||||
"""
|
||||
Audit a text block using OpenAI with strict JSON mode (async version).
|
||||
|
||||
Args:
|
||||
user_prompt: User prompt to audit
|
||||
system_prompt: Cached system prompt with rules and instructions
|
||||
model_name: OpenAI model to use
|
||||
client: AsyncOpenAI client instance
|
||||
reasoning_effort: Reasoning effort for o-series models (low, medium, high)
|
||||
|
||||
Returns:
|
||||
Audit result dictionary
|
||||
"""
|
||||
request_params = {
|
||||
"model": model_name,
|
||||
"messages": [
|
||||
{"role": "system", "content": system_prompt},
|
||||
{"role": "user", "content": user_prompt},
|
||||
],
|
||||
"response_format": {
|
||||
"type": "json_schema",
|
||||
"json_schema": {
|
||||
"name": "audit_result",
|
||||
"strict": True,
|
||||
"schema": AUDIT_RESULT_SCHEMA,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
# Add reasoning_effort only for o-series models that support it
|
||||
if (
|
||||
reasoning_effort
|
||||
and reasoning_effort.lower() in ("low", "medium", "high")
|
||||
and is_openai_reasoning_model(model_name)
|
||||
):
|
||||
request_params["reasoning_effort"] = reasoning_effort.lower()
|
||||
|
||||
response = await client.chat.completions.create(**request_params)
|
||||
|
||||
# With structured output, response is guaranteed to be valid JSON
|
||||
result = json.loads(response.choices[0].message.content)
|
||||
return result
|
||||
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
|
||||
)
|
||||
@@ -0,0 +1,14 @@
|
||||
"""Legacy parser engine: simple in-process text extraction (no sidecar).
|
||||
|
||||
Produces ``raw``-format plain text. The extraction helpers
|
||||
(:func:`extract_text` and the per-format ``_extract_*`` functions) were moved
|
||||
here from the API layer so the core parser owns them (the API layer imports
|
||||
from here instead of the other way round).
|
||||
"""
|
||||
|
||||
from lightrag.parser.legacy.extractors import (
|
||||
LegacyExtractionError,
|
||||
extract_text,
|
||||
)
|
||||
|
||||
__all__ = ["LegacyExtractionError", "extract_text"]
|
||||
@@ -0,0 +1,185 @@
|
||||
"""Legacy text extractors (moved from the API layer).
|
||||
|
||||
``extract_text`` dispatches on file suffix: binary office/pdf formats use the
|
||||
dedicated ``_extract_*`` helpers; everything else is decoded as UTF-8 text
|
||||
with the same validation (empty / binary-looking / non-UTF-8) the API upload
|
||||
path used to enforce — now raised as :class:`LegacyExtractionError` so a bad
|
||||
file fails the parse stage instead of silently yielding an empty document.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from io import BytesIO
|
||||
|
||||
|
||||
class LegacyExtractionError(ValueError):
|
||||
"""Raised when legacy extraction cannot produce usable text."""
|
||||
|
||||
|
||||
def _extract_pdf_pypdf(file_bytes: bytes, password: str | None = None) -> str:
|
||||
"""Extract PDF content using pypdf (synchronous)."""
|
||||
from pypdf import PdfReader # type: ignore
|
||||
|
||||
pdf_file = BytesIO(file_bytes)
|
||||
reader = PdfReader(pdf_file)
|
||||
|
||||
if reader.is_encrypted:
|
||||
# Try empty password first (covers permission-only encrypted PDFs)
|
||||
decrypt_result = reader.decrypt(password or "")
|
||||
if decrypt_result == 0:
|
||||
if password:
|
||||
raise Exception("Incorrect PDF password")
|
||||
else:
|
||||
raise Exception("PDF is encrypted but no password provided")
|
||||
|
||||
content = ""
|
||||
for page in reader.pages:
|
||||
content += page.extract_text() + "\n"
|
||||
return content
|
||||
|
||||
|
||||
def _extract_docx(file_bytes: bytes) -> str:
|
||||
"""Extract DOCX content including tables in document order (synchronous)."""
|
||||
from docx import Document # type: ignore
|
||||
from docx.table import Table # type: ignore
|
||||
from docx.text.paragraph import Paragraph # type: ignore
|
||||
|
||||
docx_file = BytesIO(file_bytes)
|
||||
doc = Document(docx_file)
|
||||
|
||||
def escape_cell(cell_value: str | None) -> str:
|
||||
if cell_value is None:
|
||||
return ""
|
||||
text = str(cell_value)
|
||||
return (
|
||||
text.replace("\\", "\\\\")
|
||||
.replace("\t", "  ")
|
||||
.replace("\r\n", "<br>")
|
||||
.replace("\r", "<br>")
|
||||
.replace("\n", "<br>")
|
||||
)
|
||||
|
||||
content_parts = []
|
||||
in_table = False
|
||||
for element in doc.element.body:
|
||||
if element.tag.endswith("p"):
|
||||
if in_table:
|
||||
content_parts.append("")
|
||||
in_table = False
|
||||
paragraph = Paragraph(element, doc)
|
||||
content_parts.append(paragraph.text)
|
||||
elif element.tag.endswith("tbl"):
|
||||
if content_parts and not in_table:
|
||||
content_parts.append("")
|
||||
in_table = True
|
||||
table = Table(element, doc)
|
||||
for row in table.rows:
|
||||
row_text = [escape_cell(cell.text) for cell in row.cells]
|
||||
if any(cell for cell in row_text):
|
||||
content_parts.append("\t".join(row_text))
|
||||
return "\n".join(content_parts)
|
||||
|
||||
|
||||
def _extract_pptx(file_bytes: bytes) -> str:
|
||||
"""Extract PPTX content (synchronous)."""
|
||||
from pptx import Presentation # type: ignore
|
||||
|
||||
pptx_file = BytesIO(file_bytes)
|
||||
prs = Presentation(pptx_file)
|
||||
content = ""
|
||||
for slide in prs.slides:
|
||||
for shape in slide.shapes:
|
||||
if hasattr(shape, "text"):
|
||||
content += shape.text + "\n"
|
||||
return content
|
||||
|
||||
|
||||
def _extract_xlsx(file_bytes: bytes) -> str:
|
||||
"""Extract XLSX content in tab-delimited format with sheet separators."""
|
||||
from openpyxl import load_workbook # type: ignore
|
||||
|
||||
xlsx_file = BytesIO(file_bytes)
|
||||
wb = load_workbook(xlsx_file)
|
||||
|
||||
def escape_cell(cell_value: str | int | float | None) -> str:
|
||||
if cell_value is None:
|
||||
return ""
|
||||
text = str(cell_value)
|
||||
return (
|
||||
text.replace("\\", "\\\\")
|
||||
.replace("\t", "\\t")
|
||||
.replace("\r\n", "\\n")
|
||||
.replace("\r", "\\n")
|
||||
.replace("\n", "\\n")
|
||||
)
|
||||
|
||||
def escape_sheet_title(title: str) -> str:
|
||||
return str(title).replace("\n", " ").replace("\t", " ").replace("\r", " ")
|
||||
|
||||
content_parts: list[str] = []
|
||||
sheet_separator = "=" * 20
|
||||
|
||||
for idx, sheet in enumerate(wb):
|
||||
if idx > 0:
|
||||
content_parts.append("")
|
||||
safe_title = escape_sheet_title(sheet.title)
|
||||
content_parts.append(f"{sheet_separator} Sheet: {safe_title} {sheet_separator}")
|
||||
max_columns = sheet.max_column if sheet.max_column else 0
|
||||
for row in sheet.iter_rows(values_only=True):
|
||||
row_parts = []
|
||||
for col_idx in range(max_columns):
|
||||
if col_idx < len(row):
|
||||
row_parts.append(escape_cell(row[col_idx]))
|
||||
else:
|
||||
row_parts.append("")
|
||||
if all(part == "" for part in row_parts):
|
||||
content_parts.append("")
|
||||
else:
|
||||
content_parts.append("\t".join(row_parts))
|
||||
content_parts.append(sheet_separator)
|
||||
return "\n".join(content_parts)
|
||||
|
||||
|
||||
# Suffixes (without dot) routed to dedicated binary extractors.
|
||||
_BINARY_EXTRACTORS = {
|
||||
"pdf": _extract_pdf_pypdf,
|
||||
"docx": _extract_docx,
|
||||
"pptx": _extract_pptx,
|
||||
"xlsx": _extract_xlsx,
|
||||
}
|
||||
|
||||
|
||||
def _decode_text(file_bytes: bytes) -> str:
|
||||
"""UTF-8 decode with the upload-path validation, raised on failure."""
|
||||
try:
|
||||
content = file_bytes.decode("utf-8")
|
||||
except UnicodeDecodeError as e:
|
||||
raise LegacyExtractionError(
|
||||
"File is not valid UTF-8 encoded text. Please convert it to "
|
||||
f"UTF-8 before processing: {e}"
|
||||
) from e
|
||||
if not content or len(content.strip()) == 0:
|
||||
raise LegacyExtractionError("File contains no content or only whitespace")
|
||||
if content.startswith("b'") or content.startswith('b"'):
|
||||
raise LegacyExtractionError(
|
||||
"File appears to contain binary data representation instead of text"
|
||||
)
|
||||
return content
|
||||
|
||||
|
||||
def extract_text(
|
||||
file_bytes: bytes, suffix: str, *, pdf_password: str | None = None
|
||||
) -> str:
|
||||
"""Extract plain text from ``file_bytes`` based on ``suffix`` (no dot).
|
||||
|
||||
Synchronous; callers run it in a thread. Raises
|
||||
:class:`LegacyExtractionError` (or the extractor's own exception) on
|
||||
failure.
|
||||
"""
|
||||
suffix = suffix.lower().lstrip(".")
|
||||
extractor = _BINARY_EXTRACTORS.get(suffix)
|
||||
if extractor is _extract_pdf_pypdf:
|
||||
return _extract_pdf_pypdf(file_bytes, pdf_password)
|
||||
if extractor is not None:
|
||||
return extractor(file_bytes)
|
||||
return _decode_text(file_bytes)
|
||||
@@ -0,0 +1,78 @@
|
||||
"""Legacy engine adapter: worker-stage plain-text extraction (RAW output)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import time
|
||||
|
||||
from lightrag.constants import FULL_DOCS_FORMAT_RAW, PARSER_ENGINE_LEGACY
|
||||
from lightrag.parser.base import BaseParser, ParseContext, ParseResult
|
||||
|
||||
|
||||
class LegacyParser(BaseParser):
|
||||
"""Extract plain text in-process and store it as a ``raw`` document.
|
||||
|
||||
Also serves as the dispatch fallback engine: an unknown/unsupported
|
||||
suffix raises (caught at the parse stage → doc_status FAILED) rather than
|
||||
silently producing empty content.
|
||||
"""
|
||||
|
||||
engine_name = PARSER_ENGINE_LEGACY
|
||||
|
||||
async def parse(self, ctx: ParseContext) -> ParseResult:
|
||||
from lightrag.parser.legacy.extractors import (
|
||||
LegacyExtractionError,
|
||||
extract_text,
|
||||
)
|
||||
from lightrag.parser.registry import suffix_capabilities
|
||||
|
||||
rs = ctx.resolve(self.engine_name)
|
||||
source = rs.source_path
|
||||
if not source.is_file():
|
||||
raise FileNotFoundError(f"legacy source file not found: {source}")
|
||||
|
||||
suffix = source.suffix.lower().lstrip(".")
|
||||
if suffix not in suffix_capabilities(self.engine_name):
|
||||
raise ValueError(
|
||||
f"legacy parser does not support .{suffix or '<no suffix>'}: "
|
||||
f"doc_id={ctx.doc_id} file={ctx.file_path}"
|
||||
)
|
||||
|
||||
file_bytes = await asyncio.to_thread(source.read_bytes)
|
||||
# The PDF password is sourced from env only: the parser layer reads
|
||||
# ``PDF_DECRYPT_PASSWORD`` directly rather than from the API layer's
|
||||
# ``global_args`` (which would invert the parser -> API dependency
|
||||
# direction).
|
||||
pdf_password = os.getenv("PDF_DECRYPT_PASSWORD") or None
|
||||
text = await asyncio.to_thread(
|
||||
extract_text, file_bytes, suffix, pdf_password=pdf_password
|
||||
)
|
||||
# The binary extractors (pdf/docx/pptx/xlsx) return whatever the
|
||||
# library yields — a scanned PDF with no text layer extracts to pure
|
||||
# whitespace. Fail the parse (like the text-decode path already does)
|
||||
# instead of persisting an empty document into chunking.
|
||||
if not text.strip():
|
||||
raise LegacyExtractionError(
|
||||
f"extracted no usable text from {ctx.file_path} (doc_id={ctx.doc_id})"
|
||||
)
|
||||
|
||||
await ctx.rag._persist_parsed_full_docs(
|
||||
ctx.doc_id,
|
||||
{
|
||||
"content": text,
|
||||
"file_path": ctx.file_path,
|
||||
"parse_format": FULL_DOCS_FORMAT_RAW,
|
||||
"parse_engine": self.engine_name,
|
||||
"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_RAW,
|
||||
content=text,
|
||||
blocks_path="",
|
||||
parse_engine=self.engine_name,
|
||||
)
|
||||
@@ -0,0 +1 @@
|
||||
"""Native Markdown parser engine (.md / .textpack) for the native engine."""
|
||||
@@ -0,0 +1,391 @@
|
||||
"""Pure markdown → block-list extraction for the native markdown engine.
|
||||
|
||||
This is the engine-private counterpart of ``parser/docx/parse_document.py``:
|
||||
it turns raw markdown text into the same shape the native IR builder consumes
|
||||
(a list of heading-split block dicts whose ``content`` carries placeholder
|
||||
markers), plus side tables describing the tables / equations / images those
|
||||
markers stand for.
|
||||
|
||||
Placeholder protocol (two-stage, mirroring the docx parser):
|
||||
|
||||
- ``extract_markdown`` embeds **self-closing temporary markers** in block
|
||||
``content`` — ``<mdtable ref="t0"/>`` / ``<mdequation ref="e0"/>`` /
|
||||
``<mddrawing ref="d0"/>`` — and never builds IR objects itself.
|
||||
- :class:`lightrag.parser.markdown.ir_builder.NativeMarkdownIRBuilder` later
|
||||
rewrites each marker into the IR placeholder token (``{{TBL:k}}`` /
|
||||
``{{EQ:k}}`` / ``{{IMG:k}}``) and builds the matching IR item from the side
|
||||
tables.
|
||||
|
||||
The actual table/equation/image payloads live in side tables keyed by the
|
||||
marker ref (NOT inside the marker string). This keeps a captured HTML
|
||||
``<table>`` out of the content stream — so its inner ``</table>`` can never
|
||||
truncate a naive ``<mdtable>…</mdtable>`` wrapper — and keeps image bytes off
|
||||
the content string entirely.
|
||||
|
||||
Supported subset (NOT full CommonMark/GFM, by design — see the parser plan):
|
||||
ATX headings, simple pipe tables (with a header row), block-level ``$$`` math,
|
||||
inline ```` images, and HTML ``<table>`` blocks. Reference-style
|
||||
images, escaped pipes, nested tables, setext headings and list/quote-nested
|
||||
structures are left as verbatim text rather than misrecognised.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Protocol
|
||||
|
||||
from lightrag.parser._html_table import starts_with_html_tag
|
||||
from lightrag.parser._markdown import render_heading_line
|
||||
|
||||
PREFACE_HEADING = "Preface/Uncategorized"
|
||||
|
||||
# --- placeholder marker protocol (shared with the IR builder) --------------
|
||||
_TABLE_MARKER = '<mdtable ref="{ref}"/>'
|
||||
_EQUATION_MARKER = '<mdequation ref="{ref}"/>'
|
||||
_DRAWING_MARKER = '<mddrawing ref="{ref}"/>'
|
||||
|
||||
TABLE_MARKER_RE = re.compile(r'<mdtable ref="([^"]+)"/>')
|
||||
EQUATION_MARKER_RE = re.compile(r'<mdequation ref="([^"]+)"/>')
|
||||
DRAWING_MARKER_RE = re.compile(r'<mddrawing ref="([^"]+)"/>')
|
||||
|
||||
# --- markdown token patterns ----------------------------------------------
|
||||
_HEADING_RE = re.compile(r"^(#{1,6})\s+(.*?)\s*$")
|
||||
# Trailing closing-hash run of an ATX heading (``## Foo ##`` → ``Foo``).
|
||||
_HEADING_TRAILING_HASHES_RE = re.compile(r"\s+#+\s*$")
|
||||
_FENCE_RE = re.compile(r"^(`{3,}|~{3,})(.*)$")
|
||||
# A GFM delimiter row: one or more ``---`` cells with optional ``:`` alignment.
|
||||
_DELIMITER_ROW_RE = re.compile(r"^\s*\|?\s*:?-+:?\s*(\|\s*:?-+:?\s*)*\|?\s*$")
|
||||
# A single delimiter cell (after splitting on ``|``): ``---`` with optional
|
||||
# ``:`` alignment markers, nothing else.
|
||||
_DELIMITER_CELL_RE = re.compile(r"^:?-+:?$")
|
||||
# Inline image: ````. ``src`` may be wrapped in
|
||||
# angle brackets. base64 data URLs and bare URLs/paths (no spaces, no ``)``).
|
||||
_IMAGE_RE = re.compile(
|
||||
r'!\[(?P<alt>[^\]]*)\]\(\s*(?P<src><[^>]*>|[^)\s]+)(?:\s+"[^"]*")?\s*\)'
|
||||
)
|
||||
|
||||
|
||||
def table_marker(ref: str) -> str:
|
||||
return _TABLE_MARKER.format(ref=ref)
|
||||
|
||||
|
||||
def equation_marker(ref: str) -> str:
|
||||
return _EQUATION_MARKER.format(ref=ref)
|
||||
|
||||
|
||||
def drawing_marker(ref: str) -> str:
|
||||
return _DRAWING_MARKER.format(ref=ref)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ResolvedImage:
|
||||
"""Outcome of resolving one ```` reference.
|
||||
|
||||
``kind``:
|
||||
* ``"local"`` — bytes available; ``asset_ref`` is a stable identity used
|
||||
to deduplicate (same identity ⇒ one on-disk asset shared by every
|
||||
occurrence). ``data`` / ``suggested_name`` / ``fmt`` describe it.
|
||||
* ``"external"`` — keep as an external link; ``url`` is rendered verbatim
|
||||
into the drawing's ``path_override`` (no bytes materialized).
|
||||
* ``"skip"`` — drop the image (resolver already logged / counted it).
|
||||
"""
|
||||
|
||||
kind: str
|
||||
asset_ref: str = ""
|
||||
data: bytes | None = None
|
||||
suggested_name: str = ""
|
||||
fmt: str = ""
|
||||
url: str = ""
|
||||
|
||||
|
||||
class ImageResolver(Protocol):
|
||||
"""Resolves a markdown image ``src`` to bytes / link / skip.
|
||||
|
||||
Implementations own all I/O (base64 decode, HTTP download, textpack asset
|
||||
read) and any deduplication caching, plus bumping warning counters for
|
||||
skipped / failed images. ``extract_markdown`` stays pure and only records
|
||||
what the resolver returns.
|
||||
"""
|
||||
|
||||
def resolve(self, src: str) -> ResolvedImage: ...
|
||||
|
||||
|
||||
@dataclass
|
||||
class MarkdownExtraction:
|
||||
"""Result of :func:`extract_markdown`.
|
||||
|
||||
``blocks`` mirrors the docx block-dict shape (``heading`` / ``level`` /
|
||||
``parent_headings`` / ``content``). The side tables are keyed by marker
|
||||
ref; ``assets`` is keyed by :attr:`ResolvedImage.asset_ref` (deduped).
|
||||
"""
|
||||
|
||||
blocks: list[dict] = field(default_factory=list)
|
||||
tables: dict[str, dict] = field(default_factory=dict)
|
||||
equations: dict[str, str] = field(default_factory=dict)
|
||||
drawings: dict[str, dict] = field(default_factory=dict)
|
||||
assets: dict[str, dict] = field(default_factory=dict)
|
||||
|
||||
|
||||
def _clean_heading(text: str) -> str:
|
||||
return _HEADING_TRAILING_HASHES_RE.sub("", text).strip()
|
||||
|
||||
|
||||
def _split_pipe_row(line: str) -> list[str]:
|
||||
"""Split a pipe-table row into trimmed cells (no escaped-pipe handling)."""
|
||||
s = line.strip()
|
||||
if s.startswith("|"):
|
||||
s = s[1:]
|
||||
if s.endswith("|"):
|
||||
s = s[:-1]
|
||||
return [cell.strip() for cell in s.split("|")]
|
||||
|
||||
|
||||
def _is_pipe_table_delimiter(header_line: str, delim_line: str) -> bool:
|
||||
"""True iff ``delim_line`` is a GFM delimiter row matching ``header_line``.
|
||||
|
||||
Beyond the row-shape regex, every delimiter cell must be a bare ``---`` (no
|
||||
stray text) and the column count must equal the header's. This rejects a
|
||||
bare ``---`` (a thematic break or setext underline) sitting under a
|
||||
pipe-containing paragraph — that has one column versus the header's many, so
|
||||
it is not a table, matching GFM's column-count rule."""
|
||||
if not _DELIMITER_ROW_RE.match(delim_line):
|
||||
return False
|
||||
delim_cells = _split_pipe_row(delim_line)
|
||||
if not all(_DELIMITER_CELL_RE.match(cell) for cell in delim_cells):
|
||||
return False
|
||||
return len(delim_cells) == len(_split_pipe_row(header_line))
|
||||
|
||||
|
||||
def extract_markdown(
|
||||
text: str,
|
||||
*,
|
||||
image_resolver: ImageResolver,
|
||||
) -> MarkdownExtraction:
|
||||
"""Extract markdown ``text`` into heading-split blocks + side tables.
|
||||
|
||||
Image I/O and any warning counting are delegated to ``image_resolver``;
|
||||
this function stays pure (no filesystem / network).
|
||||
"""
|
||||
out = MarkdownExtraction()
|
||||
lines = text.splitlines()
|
||||
|
||||
heading_stack: list[tuple[int, str]] = [] # (level, clean_text)
|
||||
counters = {"t": 0, "e": 0, "d": 0}
|
||||
|
||||
cur_heading = PREFACE_HEADING
|
||||
cur_level = 0
|
||||
cur_parents: list[str] = []
|
||||
cur_lines: list[str] = []
|
||||
has_block_payload = False # any marker emitted in the current block
|
||||
|
||||
def _flush() -> None:
|
||||
nonlocal cur_lines, has_block_payload
|
||||
content = "\n".join(cur_lines).strip()
|
||||
if not content and not has_block_payload:
|
||||
cur_lines = []
|
||||
has_block_payload = False
|
||||
return
|
||||
out.blocks.append(
|
||||
{
|
||||
"heading": cur_heading,
|
||||
"level": cur_level,
|
||||
"parent_headings": list(cur_parents),
|
||||
"content": "\n".join(cur_lines).rstrip(),
|
||||
}
|
||||
)
|
||||
cur_lines = []
|
||||
has_block_payload = False
|
||||
|
||||
def _open(level: int, clean: str, raw: str, parents: list[str]) -> None:
|
||||
nonlocal cur_heading, cur_level, cur_parents
|
||||
cur_heading = clean
|
||||
cur_level = level
|
||||
cur_parents = parents
|
||||
cur_lines.append(render_heading_line(level, raw))
|
||||
|
||||
def _next_ref(kind: str) -> str:
|
||||
counters[kind] += 1
|
||||
return f"{kind}{counters[kind]}"
|
||||
|
||||
def _resolve_image(match: re.Match[str]) -> str:
|
||||
src = match.group("src").strip()
|
||||
if src.startswith("<") and src.endswith(">"):
|
||||
src = src[1:-1].strip()
|
||||
resolved = image_resolver.resolve(src)
|
||||
if resolved.kind == "skip":
|
||||
# Resolver already warned/counted; drop the image, keep nothing.
|
||||
return ""
|
||||
# A base64 data URL carries no meaningful reference name and would
|
||||
# bloat the sidecar (the bytes are already materialized as an asset),
|
||||
# so it is not echoed into ``src``.
|
||||
display_src = "" if src.lower().startswith("data:") else src
|
||||
ref = _next_ref("d")
|
||||
if resolved.kind == "local":
|
||||
asset_ref = resolved.asset_ref
|
||||
if asset_ref not in out.assets:
|
||||
out.assets[asset_ref] = {
|
||||
"suggested_name": resolved.suggested_name,
|
||||
"data": resolved.data,
|
||||
"fmt": resolved.fmt,
|
||||
}
|
||||
out.drawings[ref] = {
|
||||
"kind": "local",
|
||||
"asset_ref": asset_ref,
|
||||
"fmt": resolved.fmt,
|
||||
"src": display_src,
|
||||
}
|
||||
else: # external
|
||||
out.drawings[ref] = {
|
||||
"kind": "external",
|
||||
"url": resolved.url or src,
|
||||
"fmt": resolved.fmt,
|
||||
"src": display_src,
|
||||
}
|
||||
return drawing_marker(ref)
|
||||
|
||||
def _emit_inline(line: str) -> None:
|
||||
nonlocal has_block_payload
|
||||
before = len(out.drawings)
|
||||
new_line = _IMAGE_RE.sub(_resolve_image, line)
|
||||
if len(out.drawings) != before:
|
||||
has_block_payload = True
|
||||
cur_lines.append(new_line)
|
||||
|
||||
n = len(lines)
|
||||
i = 0
|
||||
fence: tuple[str, int, bool] | None = None # (char, length, is_open)
|
||||
while i < n:
|
||||
line = lines[i]
|
||||
stripped = line.strip()
|
||||
|
||||
# --- fenced code blocks: verbatim, suppress all detection ----------
|
||||
fence_match = _FENCE_RE.match(stripped)
|
||||
if fence is not None:
|
||||
cur_lines.append(line)
|
||||
if fence_match:
|
||||
ch, run = fence_match.group(1)[0], len(fence_match.group(1))
|
||||
# A closing fence is the same char, length >= opener, no info.
|
||||
if ch == fence[0] and run >= fence[1] and not fence_match.group(2):
|
||||
fence = None
|
||||
i += 1
|
||||
continue
|
||||
if fence_match:
|
||||
fence = (fence_match.group(1)[0], len(fence_match.group(1)), True)
|
||||
cur_lines.append(line)
|
||||
i += 1
|
||||
continue
|
||||
|
||||
# --- ATX heading ---------------------------------------------------
|
||||
heading_match = _HEADING_RE.match(line)
|
||||
if heading_match:
|
||||
level = len(heading_match.group(1))
|
||||
raw = heading_match.group(2)
|
||||
clean = _clean_heading(raw)
|
||||
heading_stack[:] = heading_stack[: max(level - 1, 0)]
|
||||
parents = [h for _, h in heading_stack if h]
|
||||
heading_stack.append((level, clean))
|
||||
_flush()
|
||||
_open(level, clean, raw, parents)
|
||||
i += 1
|
||||
continue
|
||||
|
||||
# --- block equation ($$ … $$) --------------------------------------
|
||||
if stripped.startswith("$$"):
|
||||
consumed, latex = _consume_block_equation(lines, i)
|
||||
if consumed > 0:
|
||||
ref = _next_ref("e")
|
||||
out.equations[ref] = latex
|
||||
cur_lines.append(equation_marker(ref))
|
||||
has_block_payload = True
|
||||
i += consumed
|
||||
continue
|
||||
|
||||
# --- HTML <table> block --------------------------------------------
|
||||
if starts_with_html_tag(stripped.lower(), "table"):
|
||||
consumed, html = _consume_html_table(lines, i)
|
||||
if consumed > 0:
|
||||
ref = _next_ref("t")
|
||||
out.tables[ref] = {"kind": "html", "html": html}
|
||||
cur_lines.append(table_marker(ref))
|
||||
has_block_payload = True
|
||||
i += consumed
|
||||
continue
|
||||
|
||||
# --- pipe table ----------------------------------------------------
|
||||
if "|" in line and i + 1 < n and _is_pipe_table_delimiter(line, lines[i + 1]):
|
||||
consumed, rows, header = _consume_pipe_table(lines, i)
|
||||
if consumed > 0:
|
||||
ref = _next_ref("t")
|
||||
out.tables[ref] = {"kind": "pipe", "rows": rows, "header": header}
|
||||
cur_lines.append(table_marker(ref))
|
||||
has_block_payload = True
|
||||
i += consumed
|
||||
continue
|
||||
|
||||
# --- plain text line (inline images resolved here) -----------------
|
||||
_emit_inline(line)
|
||||
i += 1
|
||||
|
||||
_flush()
|
||||
return out
|
||||
|
||||
|
||||
def _consume_block_equation(lines: list[str], start: int) -> tuple[int, str]:
|
||||
"""Parse a ``$$``-delimited block equation starting at ``lines[start]``.
|
||||
|
||||
Returns ``(lines_consumed, latex)`` or ``(0, "")`` when the block is not
|
||||
closed (treated as plain text by the caller). Only paragraph-level math is
|
||||
recognised: the opening line's stripped text must start with ``$$``.
|
||||
"""
|
||||
first = lines[start].strip()
|
||||
inner_first = first[2:]
|
||||
# Single-line ``$$ … $$``.
|
||||
if inner_first.rstrip().endswith("$$") and len(inner_first.rstrip()) >= 2:
|
||||
latex = inner_first.rstrip()[:-2].strip()
|
||||
return 1, latex
|
||||
# Multi-line: collect until a line whose stripped text ends with ``$$``.
|
||||
body: list[str] = []
|
||||
if inner_first.strip():
|
||||
body.append(inner_first.strip())
|
||||
j = start + 1
|
||||
while j < len(lines):
|
||||
s = lines[j].strip()
|
||||
if s.endswith("$$"):
|
||||
tail = s[:-2].strip()
|
||||
if tail:
|
||||
body.append(tail)
|
||||
return (j - start + 1), "\n".join(body).strip()
|
||||
body.append(lines[j])
|
||||
j += 1
|
||||
return 0, ""
|
||||
|
||||
|
||||
def _consume_html_table(lines: list[str], start: int) -> tuple[int, str]:
|
||||
"""Collect a ``<table>…</table>`` block (line-spanning). ``(consumed, html)``
|
||||
or ``(0, "")`` when no closing ``</table>`` is found."""
|
||||
buf: list[str] = []
|
||||
j = start
|
||||
while j < len(lines):
|
||||
buf.append(lines[j])
|
||||
if "</table>" in lines[j].lower():
|
||||
return (j - start + 1), "\n".join(buf).strip()
|
||||
j += 1
|
||||
return 0, ""
|
||||
|
||||
|
||||
def _consume_pipe_table(
|
||||
lines: list[str], start: int
|
||||
) -> tuple[int, list[list[str]], list[list[str]] | None]:
|
||||
"""Parse a GFM pipe table whose header is ``lines[start]`` and delimiter is
|
||||
``lines[start+1]``. Returns ``(consumed, body_rows, header_grid)``."""
|
||||
header = _split_pipe_row(lines[start])
|
||||
body: list[list[str]] = []
|
||||
j = start + 2 # skip header + delimiter
|
||||
while j < len(lines):
|
||||
s = lines[j].strip()
|
||||
if not s or "|" not in s:
|
||||
break
|
||||
body.append(_split_pipe_row(lines[j]))
|
||||
j += 1
|
||||
return (j - start), body, [header] if header else None
|
||||
@@ -0,0 +1,227 @@
|
||||
"""Native Markdown IR builder: ``extract_markdown`` output → :class:`IRDoc`.
|
||||
|
||||
Input contract: the block dicts and side tables produced by
|
||||
:func:`lightrag.parser.markdown.extract.extract_markdown`, threaded through
|
||||
``NativeMarkdownParser.extract`` as ``(blocks, _, metadata)``. ``metadata``
|
||||
carries the marker→payload side tables under the ``md_*`` keys.
|
||||
|
||||
Each block's ``content`` holds self-closing markers (``<mdtable ref=…/>`` /
|
||||
``<mdequation ref=…/>`` / ``<mddrawing ref=…/>``). This builder rewrites them
|
||||
into IR placeholder tokens (``{{TBL:k}}`` / ``{{EQ:k}}`` / ``{{IMG:k}}``) and
|
||||
builds the matching :class:`IRTable` / :class:`IREquation` / :class:`IRDrawing`
|
||||
from the side tables.
|
||||
|
||||
Image bytes are carried in ``metadata["md_assets"]`` (keyed by a stable
|
||||
identity), so assets are declared as ``AssetSpec(source=bytes)`` and the writer
|
||||
materializes + deduplicates them. External-link images carry no bytes and are
|
||||
rendered verbatim through ``IRDrawing.path_override``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import itertools
|
||||
import re
|
||||
from collections.abc import Callable
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from lightrag.parser._html_table import (
|
||||
extract_html_table_info,
|
||||
extract_thead_html,
|
||||
html_table_inner_body,
|
||||
unwrap_html_table,
|
||||
)
|
||||
from lightrag.parser.markdown.extract import (
|
||||
DRAWING_MARKER_RE,
|
||||
EQUATION_MARKER_RE,
|
||||
PREFACE_HEADING,
|
||||
TABLE_MARKER_RE,
|
||||
)
|
||||
from lightrag.sidecar.ir import (
|
||||
AssetSpec,
|
||||
IRBlock,
|
||||
IRDoc,
|
||||
IRDrawing,
|
||||
IREquation,
|
||||
IRPosition,
|
||||
IRTable,
|
||||
)
|
||||
|
||||
|
||||
def _placeholder_keyspace() -> Callable[[str], str]:
|
||||
counter = itertools.count(1)
|
||||
return lambda prefix: f"{prefix}{next(counter)}"
|
||||
|
||||
|
||||
class NativeMarkdownIRBuilder:
|
||||
"""Translate ``extract_markdown`` output into an :class:`IRDoc`.
|
||||
|
||||
Stateless — instantiate per call. ``asset_dir_name`` is the relative name
|
||||
of ``<base>.blocks.assets/``; the writer applies it as the on-disk prefix
|
||||
via ``block_drawing_path_style="with_prefix"``.
|
||||
"""
|
||||
|
||||
def normalize(
|
||||
self,
|
||||
blocks: list[dict[str, Any]],
|
||||
*,
|
||||
document_name: str,
|
||||
asset_dir_name: str,
|
||||
parse_metadata: dict[str, Any] | None = None,
|
||||
) -> IRDoc:
|
||||
meta = parse_metadata or {}
|
||||
tables_meta: dict[str, dict] = meta.get("md_tables") or {}
|
||||
equations_meta: dict[str, str] = meta.get("md_equations") or {}
|
||||
drawings_meta: dict[str, dict] = meta.get("md_drawings") or {}
|
||||
assets_meta: dict[str, dict] = meta.get("md_assets") or {}
|
||||
|
||||
next_key = _placeholder_keyspace()
|
||||
ir_blocks: list[IRBlock] = []
|
||||
assets: list[AssetSpec] = []
|
||||
seen_asset_refs: set[str] = set()
|
||||
doc_title = ""
|
||||
|
||||
for block in blocks:
|
||||
content = block.get("content") or ""
|
||||
heading = block.get("heading") or ""
|
||||
level = int(block.get("level", 0) or 0)
|
||||
parent_headings = list(block.get("parent_headings") or [])
|
||||
|
||||
tables: list[IRTable] = []
|
||||
equations: list[IREquation] = []
|
||||
drawings: list[IRDrawing] = []
|
||||
|
||||
def _replace_table(match: "re.Match[str]") -> str:
|
||||
ref = match.group(1)
|
||||
spec = tables_meta.get(ref) or {}
|
||||
placeholder = next_key("tb")
|
||||
tables.append(_build_table(placeholder, spec))
|
||||
return f"{{{{TBL:{placeholder}}}}}"
|
||||
|
||||
def _replace_equation(match: "re.Match[str]") -> str:
|
||||
ref = match.group(1)
|
||||
latex = equations_meta.get(ref, "")
|
||||
placeholder = next_key("eq")
|
||||
equations.append(
|
||||
IREquation(placeholder_key=placeholder, latex=latex, is_block=True)
|
||||
)
|
||||
return f"{{{{EQ:{placeholder}}}}}"
|
||||
|
||||
def _replace_drawing(match: "re.Match[str]") -> str:
|
||||
ref = match.group(1)
|
||||
spec = drawings_meta.get(ref) or {}
|
||||
placeholder = next_key("im")
|
||||
drawings.append(
|
||||
_build_drawing(
|
||||
placeholder, spec, assets_meta, assets, seen_asset_refs
|
||||
)
|
||||
)
|
||||
return f"{{{{IMG:{placeholder}}}}}"
|
||||
|
||||
content = TABLE_MARKER_RE.sub(_replace_table, content)
|
||||
content = EQUATION_MARKER_RE.sub(_replace_equation, content)
|
||||
content = DRAWING_MARKER_RE.sub(_replace_drawing, content)
|
||||
|
||||
anchor = heading if heading and heading != PREFACE_HEADING else None
|
||||
positions = [IRPosition(type="heading", anchor=anchor)]
|
||||
|
||||
ir_blocks.append(
|
||||
IRBlock(
|
||||
content_template=content,
|
||||
heading=heading,
|
||||
level=level,
|
||||
parent_headings=parent_headings,
|
||||
positions=positions,
|
||||
tables=tables,
|
||||
drawings=drawings,
|
||||
equations=equations,
|
||||
)
|
||||
)
|
||||
|
||||
if not doc_title and level == 1 and heading and heading != PREFACE_HEADING:
|
||||
doc_title = heading
|
||||
|
||||
if not doc_title:
|
||||
doc_title = Path(document_name).stem or document_name
|
||||
|
||||
return IRDoc(
|
||||
document_name=document_name,
|
||||
document_format=Path(document_name).suffix.lower().lstrip("."),
|
||||
doc_title=doc_title,
|
||||
split_option={},
|
||||
blocks=ir_blocks,
|
||||
assets=assets,
|
||||
bbox_attributes=None,
|
||||
)
|
||||
|
||||
|
||||
def _build_table(placeholder: str, spec: dict) -> IRTable:
|
||||
if spec.get("kind") == "html":
|
||||
raw = str(spec.get("html") or "")
|
||||
html = unwrap_html_table(raw) or None
|
||||
body_override = html_table_inner_body(html) or None if html else None
|
||||
info = extract_html_table_info(html or "")
|
||||
return IRTable(
|
||||
placeholder_key=placeholder,
|
||||
rows=None,
|
||||
html=html,
|
||||
num_rows=info.num_rows,
|
||||
num_cols=info.num_cols,
|
||||
table_header=extract_thead_html(html or ""),
|
||||
body_override=body_override,
|
||||
)
|
||||
# pipe table
|
||||
rows = [[str(c) for c in row] for row in (spec.get("rows") or [])]
|
||||
header = spec.get("header")
|
||||
num_rows = len(rows)
|
||||
num_cols = max((len(r) for r in rows), default=0)
|
||||
if header:
|
||||
num_cols = max(num_cols, max((len(h) for h in header), default=0))
|
||||
return IRTable(
|
||||
placeholder_key=placeholder,
|
||||
rows=rows,
|
||||
html=None,
|
||||
num_rows=num_rows,
|
||||
num_cols=num_cols,
|
||||
table_header=header if header else None,
|
||||
)
|
||||
|
||||
|
||||
def _build_drawing(
|
||||
placeholder: str,
|
||||
spec: dict,
|
||||
assets_meta: dict[str, dict],
|
||||
assets: list[AssetSpec],
|
||||
seen_asset_refs: set[str],
|
||||
) -> IRDrawing:
|
||||
fmt = str(spec.get("fmt") or "")
|
||||
src = str(spec.get("src") or "")
|
||||
if spec.get("kind") == "external":
|
||||
return IRDrawing(
|
||||
placeholder_key=placeholder,
|
||||
asset_ref="",
|
||||
fmt=fmt,
|
||||
src=src,
|
||||
path_override=str(spec.get("url") or src),
|
||||
)
|
||||
asset_ref = str(spec.get("asset_ref") or "")
|
||||
if asset_ref and asset_ref not in seen_asset_refs:
|
||||
asset = assets_meta.get(asset_ref) or {}
|
||||
assets.append(
|
||||
AssetSpec(
|
||||
ref=asset_ref,
|
||||
suggested_name=str(asset.get("suggested_name") or asset_ref),
|
||||
source=asset.get("data"),
|
||||
)
|
||||
)
|
||||
seen_asset_refs.add(asset_ref)
|
||||
return IRDrawing(
|
||||
placeholder_key=placeholder,
|
||||
asset_ref=asset_ref,
|
||||
fmt=fmt,
|
||||
src=src,
|
||||
path_override=None,
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["NativeMarkdownIRBuilder"]
|
||||
@@ -0,0 +1,673 @@
|
||||
"""Native Markdown engine adapter (.md / .textpack).
|
||||
|
||||
Implements the :class:`NativeParserBase` hooks for markdown input. ``.textpack``
|
||||
is a zipped TextBundle (a ``text.markdown`` plus an ``assets/`` resource dir;
|
||||
Bear / Ulysses export format); a plain ``.md`` file is parsed directly.
|
||||
|
||||
Image handling (see the parser plan):
|
||||
|
||||
- base64 data-URL images are decoded and materialized into the sidecar assets.
|
||||
- file-reference images are resolved ONLY inside a ``.textpack`` bundle, from
|
||||
the safely-extracted bundle directory (with a read-side path-traversal
|
||||
guard); a relative reference in a standalone ``.md`` is skipped + warned.
|
||||
- external ``http(s)`` images are downloaded + embedded by default
|
||||
(``NATIVE_MD_IMAGE_DOWNLOAD_ENABLED`` defaults to ``true``), SSRF/size guarded;
|
||||
a drawing is always emitted — the fetched asset on success, or an external-link
|
||||
fallback on failure. Set the flag to ``false`` to instead DROP external images
|
||||
entirely (no drawing emitted, so a doc whose only images are external links
|
||||
produces no drawings.json).
|
||||
- SVG images (base64 / textpack file / downloaded) are rasterized to PNG via
|
||||
cairosvg before entering the sidecar; if cairosvg is unavailable or rendering
|
||||
fails the image is skipped + warned.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import binascii
|
||||
import hashlib
|
||||
import http.client
|
||||
import os
|
||||
import re
|
||||
import socket
|
||||
import tempfile
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
from ipaddress import ip_address, ip_network
|
||||
from math import ceil
|
||||
from pathlib import Path
|
||||
from shutil import rmtree
|
||||
from typing import TYPE_CHECKING, Any
|
||||
from urllib.parse import unquote, urlparse
|
||||
|
||||
from lightrag.constants import NATIVE_RAW_DIR_SUFFIX, PARSER_ENGINE_NATIVE
|
||||
from lightrag.parser.external._common import raw_dir_for_parsed_dir
|
||||
from lightrag.parser.markdown.extract import ResolvedImage, extract_markdown
|
||||
from lightrag.parser.markdown.raw_cache import (
|
||||
NativeImageRawCache,
|
||||
native_md_options_signature,
|
||||
)
|
||||
from lightrag.parser.native_base import NativeParserBase
|
||||
from lightrag.utils import logger
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from lightrag.sidecar.ir import IRDoc
|
||||
|
||||
# Zip-bomb guards for .textpack extraction.
|
||||
_TEXTPACK_MAX_ENTRIES = 10_000
|
||||
_TEXTPACK_MAX_TOTAL_BYTES = 512 * 1024 * 1024 # 512 MiB uncompressed
|
||||
|
||||
# Magic-byte signatures → file extension. Authoritative for download
|
||||
# validation; ``None`` means "not a supported image".
|
||||
_PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n"
|
||||
_JPEG_SIGNATURE = b"\xff\xd8\xff"
|
||||
_GIF_SIGNATURES = (b"GIF87a", b"GIF89a")
|
||||
_WEBP_RIFF = b"RIFF"
|
||||
_WEBP_TAG = b"WEBP"
|
||||
|
||||
# Content-Type values that carry no usable type signal — fall back to magic.
|
||||
_GENERIC_CONTENT_TYPES = {"", "application/octet-stream", "binary/octet-stream"}
|
||||
|
||||
|
||||
def _image_ext_from_magic(raw: bytes) -> str | None:
|
||||
"""Return the file extension for ``raw`` raster image bytes, or ``None`` if
|
||||
the bytes are not a recognised raster image (unlike
|
||||
``_vision_utils._detect_mime``, which always falls back to PNG). SVG is text,
|
||||
not a raster format, and is handled separately via :func:`_looks_like_svg`
|
||||
/ :func:`_rasterize_svg`."""
|
||||
if raw.startswith(_PNG_SIGNATURE):
|
||||
return "png"
|
||||
if raw.startswith(_JPEG_SIGNATURE):
|
||||
return "jpg"
|
||||
if any(raw.startswith(sig) for sig in _GIF_SIGNATURES):
|
||||
return "gif"
|
||||
if len(raw) >= 12 and raw[0:4] == _WEBP_RIFF and raw[8:12] == _WEBP_TAG:
|
||||
return "webp"
|
||||
return None
|
||||
|
||||
|
||||
def _looks_like_svg(raw: bytes) -> bool:
|
||||
"""True iff ``raw`` sniffs as SVG markup (an ``<svg`` tag near the start,
|
||||
possibly after an XML declaration / DOCTYPE / comments)."""
|
||||
head = raw[:4096].lstrip(b"\xef\xbb\xbf").lstrip() # drop UTF-8 BOM + space
|
||||
return b"<svg" in head[:2048].lower()
|
||||
|
||||
|
||||
# CSS absolute-length units → pixels (96 px/in). Relative units (``%``, ``em``,
|
||||
# ``ex``) depend on a viewport we do not have, so a dimension in those units is
|
||||
# treated as unresolvable and falls back to the ``viewBox``.
|
||||
_SVG_UNIT_TO_PX = {
|
||||
"": 1.0,
|
||||
"px": 1.0,
|
||||
"pt": 96.0 / 72.0,
|
||||
"pc": 16.0,
|
||||
"in": 96.0,
|
||||
"cm": 96.0 / 2.54,
|
||||
"mm": 96.0 / 25.4,
|
||||
"q": 96.0 / 25.4 / 4.0,
|
||||
}
|
||||
_SVG_LENGTH_RE = re.compile(r"^([0-9]*\.?[0-9]+)\s*([a-z%]*)$")
|
||||
|
||||
|
||||
def _svg_length_px(value: str) -> float | None:
|
||||
"""Parse a CSS/SVG length to pixels, or ``None`` for an unresolvable unit
|
||||
(``%`` / ``em`` / ``ex``) or unparseable text."""
|
||||
match = _SVG_LENGTH_RE.match(value.strip().lower())
|
||||
if not match:
|
||||
return None
|
||||
factor = _SVG_UNIT_TO_PX.get(match.group(2))
|
||||
if factor is None:
|
||||
return None
|
||||
return float(match.group(1)) * factor
|
||||
|
||||
|
||||
def _svg_pixel_dimensions(raw: bytes) -> tuple[int, int] | None:
|
||||
"""Best-effort render dimensions (px) for an SVG, from ``width``/``height``
|
||||
or, failing that, the ``viewBox``. Returns ``None`` when the root is not an
|
||||
``<svg>``, dimensions are missing/unresolvable, or the XML does not parse
|
||||
(parsed with defusedxml, which also blocks XML entity-expansion bombs)."""
|
||||
try:
|
||||
from defusedxml.ElementTree import fromstring
|
||||
|
||||
root = fromstring(raw)
|
||||
except Exception as exc: # noqa: BLE001 - malformed / hostile XML
|
||||
logger.debug("[native_md] SVG XML parse failed: %s", exc)
|
||||
return None
|
||||
if root.tag.rsplit("}", 1)[-1].lower() != "svg":
|
||||
return None
|
||||
width = _svg_length_px(root.get("width") or "")
|
||||
height = _svg_length_px(root.get("height") or "")
|
||||
if width is None or height is None:
|
||||
view_box = (root.get("viewBox") or "").replace(",", " ").split()
|
||||
if len(view_box) != 4:
|
||||
return None
|
||||
try:
|
||||
width, height = float(view_box[2]), float(view_box[3])
|
||||
except ValueError:
|
||||
return None
|
||||
if width <= 0 or height <= 0:
|
||||
return None
|
||||
return ceil(width), ceil(height)
|
||||
|
||||
|
||||
def _rasterize_svg(raw: bytes, *, max_pixels: int) -> bytes | None:
|
||||
"""Render SVG bytes to PNG bytes via cairosvg. Returns ``None`` (caller
|
||||
skips the image) when cairosvg is unavailable, rendering fails, or the SVG's
|
||||
declared canvas exceeds ``max_pixels`` — the dimension check runs *before*
|
||||
rendering so a tiny SVG declaring a huge canvas cannot blow up memory/CPU
|
||||
inside cairosvg before the output-size cap would catch it. A single bad SVG
|
||||
must not abort the whole document."""
|
||||
dims = _svg_pixel_dimensions(raw)
|
||||
if dims is None:
|
||||
logger.warning("[native_md] SVG dimensions missing/unparseable; skipping")
|
||||
return None
|
||||
width, height = dims
|
||||
if width * height > max_pixels:
|
||||
logger.warning(
|
||||
"[native_md] SVG canvas %dx%d exceeds %d px budget; skipping",
|
||||
width,
|
||||
height,
|
||||
max_pixels,
|
||||
)
|
||||
return None
|
||||
try:
|
||||
import cairosvg
|
||||
except Exception as exc: # noqa: BLE001 - optional native dep
|
||||
logger.warning(
|
||||
"[native_md] cairosvg unavailable, cannot rasterize SVG: %s", exc
|
||||
)
|
||||
return None
|
||||
try:
|
||||
return cairosvg.svg2png(bytestring=raw)
|
||||
except Exception as exc: # noqa: BLE001 - malformed / hostile SVG
|
||||
logger.warning("[native_md] SVG rasterization failed: %s", exc)
|
||||
return None
|
||||
|
||||
|
||||
def _image_bytes_and_ext(
|
||||
raw: bytes, *, max_bytes: int, max_svg_pixels: int
|
||||
) -> tuple[bytes, str] | None:
|
||||
"""Coerce ``raw`` image bytes into a sidecar-ready raster image.
|
||||
|
||||
A recognised raster image passes through unchanged; an SVG is rasterized to
|
||||
PNG — bounded *before* rendering by ``max_svg_pixels`` (declared canvas) and
|
||||
*after* by ``max_bytes`` (output size), so a hostile SVG that explodes into a
|
||||
giant bitmap is rejected rather than embedded. Returns ``(bytes, ext)`` or
|
||||
``None`` when the bytes are neither a raster image nor a convertible SVG.
|
||||
"""
|
||||
ext = _image_ext_from_magic(raw)
|
||||
if ext is not None:
|
||||
return raw, ext
|
||||
if _looks_like_svg(raw):
|
||||
png = _rasterize_svg(raw, max_pixels=max_svg_pixels)
|
||||
if png is not None and len(png) <= max_bytes and _image_ext_from_magic(png):
|
||||
return png, "png"
|
||||
return None
|
||||
|
||||
|
||||
def _env_bool(key: str, default: bool) -> bool:
|
||||
raw = os.getenv(key)
|
||||
if raw is None:
|
||||
return default
|
||||
return raw.strip().lower() in ("1", "true", "yes", "on", "t", "y")
|
||||
|
||||
|
||||
def _env_int(key: str, default: int) -> int:
|
||||
raw = os.getenv(key)
|
||||
if raw is None or not raw.strip():
|
||||
return default
|
||||
try:
|
||||
return int(raw.strip())
|
||||
except ValueError:
|
||||
return default
|
||||
|
||||
|
||||
def _allowed_non_public_networks() -> list:
|
||||
"""Parse ``NATIVE_MD_IMAGE_ALLOWED_NON_PUBLIC_CIDRS`` (comma-separated
|
||||
CIDRs / IPs) into networks. Invalid tokens are warned and dropped."""
|
||||
raw = os.getenv("NATIVE_MD_IMAGE_ALLOWED_NON_PUBLIC_CIDRS", "")
|
||||
nets = []
|
||||
for token in raw.split(","):
|
||||
token = token.strip()
|
||||
if not token:
|
||||
continue
|
||||
try:
|
||||
nets.append(ip_network(token, strict=False))
|
||||
except ValueError:
|
||||
logger.warning("[native_md] ignoring invalid allowed CIDR: %s", token)
|
||||
return nets
|
||||
|
||||
|
||||
def _validated_addresses(host: str) -> list[str]:
|
||||
"""Resolve ``host`` and return its addresses iff ALL are safe to fetch.
|
||||
|
||||
Default-deny: an address is accepted only when ``ip.is_global`` (so SSRF to
|
||||
loopback / private / link-local / reserved / CGNAT ``100.64.0.0/10`` /
|
||||
TEST-NET and any other non-globally-routable range is blocked) or it matches
|
||||
the operator-configured ``NATIVE_MD_IMAGE_ALLOWED_NON_PUBLIC_CIDRS`` escape
|
||||
hatch. Returns ``[]`` when resolution fails or *any* resolved address is
|
||||
non-public — so a single poisoned A/AAAA record rejects the whole host.
|
||||
|
||||
The returned addresses are what the connection actually dials (see
|
||||
:class:`_GuardedHTTPConnection`): validating and connecting share one
|
||||
resolution, closing the DNS-rebinding TOCTOU window between check and
|
||||
connect.
|
||||
"""
|
||||
allow = _allowed_non_public_networks()
|
||||
try:
|
||||
infos = socket.getaddrinfo(host, None)
|
||||
except socket.gaierror:
|
||||
return []
|
||||
addrs: list[str] = []
|
||||
for info in infos:
|
||||
addr = str(info[4][0]).split("%", 1)[0]
|
||||
try:
|
||||
ip = ip_address(addr)
|
||||
except ValueError:
|
||||
return []
|
||||
if not (ip.is_global or any(ip in net for net in allow)):
|
||||
return []
|
||||
addrs.append(addr)
|
||||
return addrs
|
||||
|
||||
|
||||
def _host_is_public(host: str) -> bool:
|
||||
"""True iff every resolved address for ``host`` is safe to fetch."""
|
||||
return bool(_validated_addresses(host))
|
||||
|
||||
|
||||
def _pin_socket(host: str, port: int, timeout, source_address):
|
||||
"""Open a TCP socket to a *validated* resolved address of ``host``.
|
||||
|
||||
The address comes from the same :func:`_validated_addresses` resolution
|
||||
that authorised the host, so urllib never independently re-resolves the
|
||||
name (which a DNS-rebinding attacker could flip to an internal IP between
|
||||
our check and the actual connect)."""
|
||||
addrs = _validated_addresses(host)
|
||||
if not addrs:
|
||||
raise OSError(f"refusing connection to non-public host: {host!r}")
|
||||
sock = socket.create_connection((addrs[0], port), timeout, source_address)
|
||||
sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
|
||||
return sock
|
||||
|
||||
|
||||
# No proxy/CONNECT-tunnel handling below: the opener disables proxies, so the
|
||||
# connection always dials the origin directly (``_tunnel_host`` is never set).
|
||||
class _GuardedHTTPConnection(http.client.HTTPConnection):
|
||||
def connect(self) -> None:
|
||||
self.sock = _pin_socket(self.host, self.port, self.timeout, self.source_address)
|
||||
|
||||
|
||||
class _GuardedHTTPSConnection(http.client.HTTPSConnection):
|
||||
def connect(self) -> None:
|
||||
sock = _pin_socket(self.host, self.port, self.timeout, self.source_address)
|
||||
# Wrap with the original hostname so SNI / certificate validation still
|
||||
# runs against the domain, not the pinned IP.
|
||||
self.sock = self._context.wrap_socket(sock, server_hostname=self.host)
|
||||
|
||||
|
||||
class _GuardedHTTPHandler(urllib.request.HTTPHandler):
|
||||
def http_open(self, req):
|
||||
return self.do_open(_GuardedHTTPConnection, req)
|
||||
|
||||
|
||||
class _GuardedHTTPSHandler(urllib.request.HTTPSHandler):
|
||||
def https_open(self, req):
|
||||
# Mirror the stdlib handler's own arguments, which differ across
|
||||
# versions: Python <3.12 stores ``_check_hostname`` on the handler and
|
||||
# forwards it; 3.12+ folds it into ``_context`` and drops the attribute.
|
||||
# Forwarding a non-existent ``_check_hostname`` raised AttributeError on
|
||||
# 3.12, silently failing every download into the external-link fallback.
|
||||
kwargs = {"context": self._context}
|
||||
if hasattr(self, "_check_hostname"):
|
||||
kwargs["check_hostname"] = self._check_hostname
|
||||
return self.do_open(_GuardedHTTPSConnection, req, **kwargs)
|
||||
|
||||
|
||||
class _GuardedRedirectHandler(urllib.request.HTTPRedirectHandler):
|
||||
"""Re-validate the host on every redirect so a redirect cannot bounce the
|
||||
request to an internal address. (Defense in depth: the pinned connection
|
||||
re-validates at connect time too.)"""
|
||||
|
||||
def redirect_request(self, req, fp, code, msg, headers, newurl): # type: ignore[override]
|
||||
host = urlparse(newurl).hostname or ""
|
||||
if not _host_is_public(host):
|
||||
raise urllib.error.HTTPError(
|
||||
newurl, code, "redirect to non-public host blocked", headers, fp
|
||||
)
|
||||
return super().redirect_request(req, fp, code, msg, headers, newurl)
|
||||
|
||||
|
||||
def _build_guarded_opener() -> urllib.request.OpenerDirector:
|
||||
"""Opener whose connections pin to a validated IP and that ignores any
|
||||
ambient ``HTTP(S)_PROXY`` (an env proxy would otherwise fetch the blocked
|
||||
URL on our behalf, bypassing the IP guard)."""
|
||||
return urllib.request.build_opener(
|
||||
urllib.request.ProxyHandler({}),
|
||||
_GuardedHTTPHandler(),
|
||||
_GuardedHTTPSHandler(),
|
||||
_GuardedRedirectHandler(),
|
||||
)
|
||||
|
||||
|
||||
class _MarkdownImageResolver:
|
||||
"""Concrete :class:`~lightrag.parser.markdown.extract.ImageResolver`.
|
||||
|
||||
Caches by ``src`` so a repeated reference downloads / decodes once and
|
||||
every occurrence shares one on-disk asset.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
bundle_root: Path | None,
|
||||
warnings: dict[str, int],
|
||||
download_enabled: bool,
|
||||
download_required: bool,
|
||||
timeout: int,
|
||||
max_bytes: int,
|
||||
max_svg_pixels: int,
|
||||
raw_cache: NativeImageRawCache | None = None,
|
||||
) -> None:
|
||||
self._bundle_root = bundle_root.resolve() if bundle_root else None
|
||||
self._warnings = warnings
|
||||
self._download_enabled = download_enabled
|
||||
self._download_required = download_required
|
||||
self._timeout = timeout
|
||||
self._max_bytes = max_bytes
|
||||
self._max_svg_pixels = max_svg_pixels
|
||||
self._raw_cache = raw_cache
|
||||
self._cache: dict[str, ResolvedImage] = {}
|
||||
|
||||
def resolve(self, src: str) -> ResolvedImage:
|
||||
cached = self._cache.get(src)
|
||||
if cached is not None:
|
||||
return cached
|
||||
result = self._resolve_uncached(src)
|
||||
self._cache[src] = result
|
||||
return result
|
||||
|
||||
def _bump(self, key: str) -> None:
|
||||
self._warnings[key] = self._warnings.get(key, 0) + 1
|
||||
|
||||
def _skip(self, reason: str, src: str) -> ResolvedImage:
|
||||
logger.warning("[native_md] skipping image (%s): %s", reason, src[:120])
|
||||
self._bump("images_skipped")
|
||||
return ResolvedImage(kind="skip")
|
||||
|
||||
def _local(self, data: bytes, ext: str, suggested_name: str) -> ResolvedImage:
|
||||
asset_ref = "sha256:" + hashlib.sha256(data).hexdigest()
|
||||
return ResolvedImage(
|
||||
kind="local",
|
||||
asset_ref=asset_ref,
|
||||
data=data,
|
||||
suggested_name=suggested_name,
|
||||
fmt=ext,
|
||||
)
|
||||
|
||||
def _resolve_uncached(self, src: str) -> ResolvedImage:
|
||||
lower = src.lower()
|
||||
if lower.startswith("data:"):
|
||||
return self._resolve_data_url(src)
|
||||
if lower.startswith(("http://", "https://")):
|
||||
return self._resolve_remote(src)
|
||||
return self._resolve_relative(src)
|
||||
|
||||
def _resolve_data_url(self, src: str) -> ResolvedImage:
|
||||
# ``data:[<mime>][;base64],<payload>`` — only base64 image payloads.
|
||||
if ";base64," not in src:
|
||||
return self._skip("data url not base64", src)
|
||||
_, _, payload = src.partition(";base64,")
|
||||
try:
|
||||
data = base64.b64decode("".join(payload.split()), validate=True)
|
||||
except (ValueError, binascii.Error):
|
||||
return self._skip("invalid base64", src)
|
||||
# Magic bytes are authoritative — the declared MIME type is not trusted
|
||||
# for validation (matching the remote-download path). SVG is rasterized
|
||||
# to PNG here.
|
||||
coerced = _image_bytes_and_ext(
|
||||
data, max_bytes=self._max_bytes, max_svg_pixels=self._max_svg_pixels
|
||||
)
|
||||
if coerced is None:
|
||||
return self._skip("unrecognised image", src)
|
||||
data, ext = coerced
|
||||
digest = hashlib.sha256(data).hexdigest()[:12]
|
||||
return self._local(data, ext, f"image-{digest}.{ext}")
|
||||
|
||||
def _resolve_relative(self, src: str) -> ResolvedImage:
|
||||
if self._bundle_root is None:
|
||||
# Standalone .md: file references are unresolved by design.
|
||||
return self._skip("file reference outside .textpack", src)
|
||||
rel = unquote(src.split("#", 1)[0].split("?", 1)[0])
|
||||
if not rel or "\\" in rel or rel.startswith("/") or os.path.isabs(rel):
|
||||
return self._skip("unsafe image path", src)
|
||||
if any(part == ".." for part in Path(rel).parts):
|
||||
return self._skip("unsafe image path", src)
|
||||
candidate = (self._bundle_root / rel).resolve()
|
||||
if not candidate.is_relative_to(self._bundle_root):
|
||||
return self._skip("image path escapes bundle", src)
|
||||
if not candidate.is_file():
|
||||
return self._skip("image file missing", src)
|
||||
# Cap a single bundled asset at the same ceiling as a remote download,
|
||||
# so one oversized file inside the (zip-bomb-bounded) bundle cannot pull
|
||||
# hundreds of MB into memory.
|
||||
if candidate.stat().st_size > self._max_bytes:
|
||||
return self._skip("image exceeds size limit", src)
|
||||
data = candidate.read_bytes()
|
||||
# Magic bytes are authoritative; the filename suffix is not trusted for
|
||||
# validation. SVG is rasterized to PNG (so the on-disk name takes the
|
||||
# resolved extension, e.g. ``logo.svg`` -> ``logo.png``).
|
||||
coerced = _image_bytes_and_ext(
|
||||
data, max_bytes=self._max_bytes, max_svg_pixels=self._max_svg_pixels
|
||||
)
|
||||
if coerced is None:
|
||||
return self._skip("unrecognised image", src)
|
||||
data, ext = coerced
|
||||
return self._local(data, ext, f"{candidate.stem}.{ext}")
|
||||
|
||||
def _resolve_remote(self, src: str) -> ResolvedImage:
|
||||
ext_hint = Path(urlparse(src).path).suffix.lower().lstrip(".")
|
||||
if not self._download_enabled:
|
||||
# Downloading is opt-in: with it disabled, external images are
|
||||
# dropped entirely (no drawing emitted), so a doc whose only images
|
||||
# are external links produces no drawings.json. This is expected
|
||||
# configuration, not a problem, so it is logged at debug level and
|
||||
# counted under a dedicated key rather than warned per image.
|
||||
logger.debug(
|
||||
"[native_md] dropping external image (download disabled): %s", src[:120]
|
||||
)
|
||||
self._bump("images_external_dropped")
|
||||
return ResolvedImage(kind="skip")
|
||||
if self._raw_cache is not None:
|
||||
hit = self._raw_cache.get(src)
|
||||
if hit is not None:
|
||||
# Reuse the cached bytes (already post-SVG-rasterization), so a
|
||||
# re-parse skips both the network fetch and the rasterization.
|
||||
data, ext = hit
|
||||
self._bump("images_cache_hit")
|
||||
digest = hashlib.sha256(data).hexdigest()[:12]
|
||||
return self._local(data, ext, f"image-{digest}.{ext}")
|
||||
try:
|
||||
data, ext = self._download(src)
|
||||
except Exception as exc: # noqa: BLE001 - best-effort network fetch
|
||||
if self._download_required:
|
||||
raise
|
||||
logger.warning("[native_md] image download failed (%s): %s", exc, src[:120])
|
||||
self._bump("images_download_failed")
|
||||
return ResolvedImage(kind="external", url=src, fmt=ext_hint)
|
||||
if self._raw_cache is not None:
|
||||
self._raw_cache.put(src, data, ext)
|
||||
digest = hashlib.sha256(data).hexdigest()[:12]
|
||||
return self._local(data, ext, f"image-{digest}.{ext}")
|
||||
|
||||
def _download(self, src: str) -> tuple[bytes, str]:
|
||||
parsed = urlparse(src)
|
||||
if parsed.scheme not in ("http", "https"):
|
||||
raise ValueError(f"unsupported scheme: {parsed.scheme!r}")
|
||||
if not _host_is_public(parsed.hostname or ""):
|
||||
raise ValueError("non-public host blocked")
|
||||
opener = _build_guarded_opener()
|
||||
req = urllib.request.Request(src, headers={"User-Agent": "lightrag-native-md"})
|
||||
with opener.open(req, timeout=self._timeout) as resp:
|
||||
content_type = (resp.headers.get_content_type() or "").lower()
|
||||
# Read at most max_bytes + 1 so we can detect an over-limit body.
|
||||
data = resp.read(self._max_bytes + 1)
|
||||
if len(data) > self._max_bytes:
|
||||
raise ValueError(f"image exceeds {self._max_bytes} bytes")
|
||||
# Magic bytes are authoritative; SVG is rasterized to PNG here. Reject
|
||||
# only an *explicit* non-image Content-Type — but ``image/svg+xml`` is a
|
||||
# valid image type, so the check stays correct for converted SVGs.
|
||||
if (
|
||||
content_type
|
||||
and content_type not in _GENERIC_CONTENT_TYPES
|
||||
and not content_type.startswith("image/")
|
||||
):
|
||||
raise ValueError(f"non-image Content-Type: {content_type!r}")
|
||||
coerced = _image_bytes_and_ext(
|
||||
data, max_bytes=self._max_bytes, max_svg_pixels=self._max_svg_pixels
|
||||
)
|
||||
if coerced is None:
|
||||
raise ValueError("downloaded bytes are not a supported image")
|
||||
return coerced
|
||||
|
||||
|
||||
class NativeMarkdownParser(NativeParserBase):
|
||||
engine_name = PARSER_ENGINE_NATIVE
|
||||
empty_content_label = "Markdown"
|
||||
|
||||
def validate_source(self, source: Path, file_path: str) -> None:
|
||||
if not (
|
||||
source.exists()
|
||||
and source.is_file()
|
||||
and source.suffix.lower() in (".md", ".textpack")
|
||||
):
|
||||
raise ValueError(
|
||||
f"Native markdown parser does not support pending file: {file_path}"
|
||||
)
|
||||
|
||||
def extract(
|
||||
self, source: Path, *, parsed_dir: Path, asset_dir: Path, base_name: str
|
||||
) -> tuple[list[dict[str, Any]], dict[str, Any], dict[str, Any]]:
|
||||
# Per-document downloaded-image cache. Lives in a ``<file>.native_raw/``
|
||||
# sibling of ``parsed_dir`` so it survives the ``rmtree(parsed_dir)`` the
|
||||
# base parser runs before each re-extraction; reused across re-parses
|
||||
# unless the source content or download config changed.
|
||||
raw_cache = NativeImageRawCache(
|
||||
raw_dir_for_parsed_dir(parsed_dir, suffix=NATIVE_RAW_DIR_SUFFIX),
|
||||
source_path=source,
|
||||
options_signature=native_md_options_signature(),
|
||||
force_reparse=_env_bool("LIGHTRAG_FORCE_REPARSE_NATIVE", False),
|
||||
)
|
||||
raw_cache.load()
|
||||
if source.suffix.lower() == ".textpack":
|
||||
tmp_dir = Path(tempfile.mkdtemp(prefix="textpack-"))
|
||||
try:
|
||||
md_text, bundle_root = self._open_textpack(source, tmp_dir)
|
||||
result = self._extract_text(
|
||||
md_text, bundle_root=bundle_root, raw_cache=raw_cache
|
||||
)
|
||||
finally:
|
||||
rmtree(tmp_dir, ignore_errors=True)
|
||||
else:
|
||||
md_text = source.read_bytes().decode("utf-8-sig")
|
||||
result = self._extract_text(md_text, bundle_root=None, raw_cache=raw_cache)
|
||||
# Flush only on successful extraction so a transient failure cannot prune
|
||||
# a previously-valid bundle (an exception propagates before this line).
|
||||
raw_cache.flush()
|
||||
return result
|
||||
|
||||
def _open_textpack(self, source: Path, tmp_dir: Path) -> tuple[str, Path]:
|
||||
from lightrag.parser.external._zip import safe_extract_zip
|
||||
|
||||
safe_extract_zip(
|
||||
source.read_bytes(),
|
||||
tmp_dir,
|
||||
max_entries=_TEXTPACK_MAX_ENTRIES,
|
||||
max_total_bytes=_TEXTPACK_MAX_TOTAL_BYTES,
|
||||
)
|
||||
text_file = self._find_text_file(tmp_dir, source.name)
|
||||
return text_file.read_bytes().decode("utf-8-sig"), text_file.parent
|
||||
|
||||
@staticmethod
|
||||
def _find_text_file(root: Path, source_name: str) -> Path:
|
||||
# The body is located by extension, not a fixed ``text.markdown`` name,
|
||||
# so any zip tool can produce a valid textpack. Layout rules:
|
||||
# * If the archive holds a ``*.textbundle`` directory, exactly one is
|
||||
# allowed and the body must live directly inside it.
|
||||
# * Otherwise the body must live directly in the archive root.
|
||||
# * The chosen directory must hold exactly one ``*.md``/``*.markdown``.
|
||||
# ``bundle_root`` (the returned file's parent) anchors asset resolution.
|
||||
bundles = sorted(
|
||||
p
|
||||
for p in root.iterdir()
|
||||
if p.is_dir() and p.suffix.lower() == ".textbundle"
|
||||
)
|
||||
if len(bundles) > 1:
|
||||
names = ", ".join(p.name for p in bundles)
|
||||
raise ValueError(
|
||||
f"multiple .textbundle directories in textpack: {source_name} ({names})"
|
||||
)
|
||||
search_dir = bundles[0] if bundles else root
|
||||
markdown = sorted(
|
||||
p
|
||||
for p in search_dir.iterdir()
|
||||
if p.is_file() and p.suffix.lower() in (".md", ".markdown")
|
||||
)
|
||||
if len(markdown) > 1:
|
||||
names = ", ".join(p.name for p in markdown)
|
||||
raise ValueError(
|
||||
f"multiple markdown files in textpack: {source_name} ({names})"
|
||||
)
|
||||
if not markdown:
|
||||
raise ValueError(f"no markdown text file found in textpack: {source_name}")
|
||||
return markdown[0]
|
||||
|
||||
def _extract_text(
|
||||
self,
|
||||
md_text: str,
|
||||
*,
|
||||
bundle_root: Path | None,
|
||||
raw_cache: NativeImageRawCache | None = None,
|
||||
) -> tuple[list[dict[str, Any]], dict[str, Any], dict[str, Any]]:
|
||||
warnings: dict[str, int] = {}
|
||||
resolver = _MarkdownImageResolver(
|
||||
bundle_root=bundle_root,
|
||||
warnings=warnings,
|
||||
download_enabled=_env_bool("NATIVE_MD_IMAGE_DOWNLOAD_ENABLED", True),
|
||||
download_required=_env_bool("NATIVE_MD_IMAGE_DOWNLOAD_REQUIRED", False),
|
||||
timeout=_env_int("NATIVE_MD_IMAGE_DOWNLOAD_TIMEOUT", 30),
|
||||
max_bytes=_env_int("NATIVE_MD_IMAGE_MAX_BYTES", 25 * 1024 * 1024),
|
||||
max_svg_pixels=_env_int("NATIVE_MD_IMAGE_MAX_SVG_PIXELS", 16_000_000),
|
||||
raw_cache=raw_cache,
|
||||
)
|
||||
extraction = extract_markdown(md_text, image_resolver=resolver)
|
||||
metadata: dict[str, Any] = {
|
||||
"md_tables": extraction.tables,
|
||||
"md_equations": extraction.equations,
|
||||
"md_drawings": extraction.drawings,
|
||||
"md_assets": extraction.assets,
|
||||
}
|
||||
return extraction.blocks, dict(warnings), metadata
|
||||
|
||||
def build_ir(
|
||||
self,
|
||||
blocks: list[dict[str, Any]],
|
||||
*,
|
||||
document_name: str,
|
||||
asset_dir_name: str,
|
||||
metadata: dict[str, Any],
|
||||
) -> "IRDoc":
|
||||
from lightrag.parser.markdown.ir_builder import NativeMarkdownIRBuilder
|
||||
|
||||
return NativeMarkdownIRBuilder().normalize(
|
||||
blocks,
|
||||
document_name=document_name,
|
||||
asset_dir_name=asset_dir_name,
|
||||
parse_metadata=metadata,
|
||||
)
|
||||
|
||||
def surface_warnings(
|
||||
self, warnings: dict[str, Any], source: Path
|
||||
) -> dict[str, Any] | None:
|
||||
relevant = {k: v for k, v in warnings.items() if v}
|
||||
return relevant or None
|
||||
@@ -0,0 +1,226 @@
|
||||
"""Per-document download cache for native-markdown external images.
|
||||
|
||||
Mirrors the ``.mineru_raw`` / ``.docling_raw`` raw-bundle pattern, scoped to the
|
||||
one expensive native-markdown step: downloading external ``http(s)`` images.
|
||||
The bundle lives in a ``<file>.native_raw/`` directory that is a **sibling** of
|
||||
``<file>.parsed/`` so it survives the ``rmtree(parsed_dir)`` that
|
||||
:meth:`NativeParserBase.parse` performs before every re-extraction.
|
||||
|
||||
Bundle layout::
|
||||
|
||||
<file>.native_raw/
|
||||
_manifest.json # atomic success marker + cache key
|
||||
<sha256(url)[:16]>.png # one file per cached image (final bytes)
|
||||
...
|
||||
|
||||
The manifest is the cache key: a bundle is reused only when the source file
|
||||
content hash AND the download-options signature both still match, mirroring the
|
||||
external engines. On a hit the resolver reuses the stored bytes (already
|
||||
post-SVG-rasterization), skipping both the network fetch and the rasterization.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
from lightrag.parser.external._common import (
|
||||
clear_dir_contents,
|
||||
compute_size_and_hash,
|
||||
)
|
||||
from lightrag.utils import logger
|
||||
|
||||
_MANIFEST_FILENAME = "_manifest.json"
|
||||
_MANIFEST_VERSION = "1"
|
||||
_CACHE_ENGINE = "native_md"
|
||||
|
||||
|
||||
def native_md_options_signature() -> str:
|
||||
"""A ``sha256`` over the download knobs that change an image's bytes.
|
||||
|
||||
Deliberately excludes ``NATIVE_MD_IMAGE_DOWNLOAD_ENABLED`` /
|
||||
``..._TIMEOUT`` / ``..._REQUIRED`` — those gate *whether* a fetch happens,
|
||||
not the resulting bytes — and includes the size / SVG-pixel ceilings and the
|
||||
SSRF allowlist (which govern what bytes are accepted at all)."""
|
||||
payload = {
|
||||
"signature_version": 1,
|
||||
"max_bytes": os.getenv("NATIVE_MD_IMAGE_MAX_BYTES", ""),
|
||||
"max_svg_pixels": os.getenv("NATIVE_MD_IMAGE_MAX_SVG_PIXELS", ""),
|
||||
"allowed_non_public_cidrs": os.getenv(
|
||||
"NATIVE_MD_IMAGE_ALLOWED_NON_PUBLIC_CIDRS", ""
|
||||
),
|
||||
}
|
||||
raw = json.dumps(payload, sort_keys=True, separators=(",", ":"))
|
||||
return "sha256:" + hashlib.sha256(raw.encode("utf-8")).hexdigest()
|
||||
|
||||
|
||||
def _url_filename(url: str, fmt: str) -> str:
|
||||
digest = hashlib.sha256(url.encode("utf-8")).hexdigest()[:16]
|
||||
ext = fmt or "bin"
|
||||
return f"{digest}.{ext}"
|
||||
|
||||
|
||||
class NativeImageRawCache:
|
||||
"""Reuse already-downloaded external images across re-parses.
|
||||
|
||||
Per-document and single-writer: each parse owns one ``raw_dir`` and the
|
||||
parse queue serializes work per document, so no locking is needed.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
raw_dir: Path,
|
||||
*,
|
||||
source_path: Path,
|
||||
options_signature: str,
|
||||
force_reparse: bool,
|
||||
) -> None:
|
||||
self._raw_dir = raw_dir
|
||||
self._source_path = source_path
|
||||
self._options_signature = options_signature
|
||||
self._force_reparse = force_reparse
|
||||
self._source_hash = ""
|
||||
self._valid = False
|
||||
self._cleared = False
|
||||
# ``_dirty`` gates the manifest write: it is set only when a real
|
||||
# download is stored (:meth:`put`). A run that only reuses cached images
|
||||
# (a pure hit) leaves it ``False`` so :meth:`flush` touches nothing on
|
||||
# disk — the bundle's mtimes then reveal whether the cache was hit.
|
||||
self._dirty = False
|
||||
# Index of reusable entries from a valid prior bundle (url -> entry).
|
||||
self._index: dict[str, dict] = {}
|
||||
# Entries referenced this run (reused or freshly put) -> manifest output.
|
||||
self._entries: dict[str, dict] = {}
|
||||
|
||||
def load(self) -> None:
|
||||
"""Compute the current source hash and decide whether the on-disk
|
||||
bundle is a cache hit (valid) or must be rebuilt."""
|
||||
try:
|
||||
_, self._source_hash = compute_size_and_hash(self._source_path)
|
||||
except OSError as exc:
|
||||
logger.debug("[native_md_cache] source hash failed: %s", exc)
|
||||
self._source_hash = ""
|
||||
if self._force_reparse:
|
||||
return
|
||||
manifest = self._read_manifest()
|
||||
if manifest is None:
|
||||
return
|
||||
if (
|
||||
manifest.get("source_content_hash") != self._source_hash
|
||||
or manifest.get("options_signature") != self._options_signature
|
||||
):
|
||||
return
|
||||
images = manifest.get("images")
|
||||
if not isinstance(images, dict):
|
||||
return
|
||||
self._index = {k: v for k, v in images.items() if isinstance(v, dict)}
|
||||
self._valid = True
|
||||
|
||||
def get(self, url: str) -> tuple[bytes, str] | None:
|
||||
"""Return ``(bytes, fmt)`` for a cached image, or ``None`` on a miss /
|
||||
integrity failure (corrupt or tampered cache file)."""
|
||||
if not self._valid:
|
||||
return None
|
||||
entry = self._index.get(url)
|
||||
if not entry:
|
||||
return None
|
||||
file_name = str(entry.get("file") or "")
|
||||
if not file_name:
|
||||
return None
|
||||
path = self._raw_dir / file_name
|
||||
if not path.is_file():
|
||||
return None
|
||||
try:
|
||||
data = path.read_bytes()
|
||||
except OSError:
|
||||
return None
|
||||
if "sha256:" + hashlib.sha256(data).hexdigest() != entry.get("sha256"):
|
||||
logger.warning("[native_md_cache] cached file integrity mismatch: %s", url)
|
||||
return None
|
||||
fmt = str(entry.get("fmt") or "")
|
||||
self._entries[url] = entry
|
||||
return data, fmt
|
||||
|
||||
def put(self, url: str, data: bytes, fmt: str) -> None:
|
||||
"""Store freshly-downloaded image bytes and record them for the manifest."""
|
||||
self._ensure_writable_dir()
|
||||
file_name = _url_filename(url, fmt)
|
||||
try:
|
||||
(self._raw_dir / file_name).write_bytes(data)
|
||||
except OSError as exc:
|
||||
logger.warning("[native_md_cache] failed to write cache file: %s", exc)
|
||||
return
|
||||
self._entries[url] = {
|
||||
"file": file_name,
|
||||
"sha256": "sha256:" + hashlib.sha256(data).hexdigest(),
|
||||
"size": len(data),
|
||||
"fmt": fmt,
|
||||
}
|
||||
self._dirty = True
|
||||
|
||||
def flush(self) -> None:
|
||||
"""Persist the bundle only if it actually changed this run.
|
||||
|
||||
Skipped entirely on a **pure cache hit** (every image reused, nothing
|
||||
downloaded) — the source is byte-identical, so the referenced image set
|
||||
equals the on-disk one and a rewrite would only be an idempotent write.
|
||||
Leaving the manifest and image files untouched means their mtimes flag a
|
||||
miss/update vs. a hit. Also a no-op for an image-less or
|
||||
download-disabled run, so a pre-existing valid bundle is left intact."""
|
||||
if not self._dirty:
|
||||
return
|
||||
self._ensure_writable_dir()
|
||||
referenced = {e["file"] for e in self._entries.values()}
|
||||
for child in self._raw_dir.iterdir():
|
||||
if child.name == _MANIFEST_FILENAME:
|
||||
continue
|
||||
if child.is_file() and child.name not in referenced:
|
||||
try:
|
||||
child.unlink()
|
||||
except OSError:
|
||||
pass
|
||||
manifest = {
|
||||
"version": _MANIFEST_VERSION,
|
||||
"engine": _CACHE_ENGINE,
|
||||
"source_content_hash": self._source_hash,
|
||||
"options_signature": self._options_signature,
|
||||
"images": self._entries,
|
||||
}
|
||||
final = self._raw_dir / _MANIFEST_FILENAME
|
||||
tmp = final.with_suffix(".json.tmp")
|
||||
try:
|
||||
tmp.write_text(
|
||||
json.dumps(manifest, ensure_ascii=False, indent=2), encoding="utf-8"
|
||||
)
|
||||
os.replace(tmp, final)
|
||||
except OSError as exc:
|
||||
logger.warning("[native_md_cache] failed to write manifest: %s", exc)
|
||||
|
||||
def _ensure_writable_dir(self) -> None:
|
||||
"""Create ``raw_dir`` and, on the first write of an invalidated bundle,
|
||||
drop the stale contents so reused entries never mingle with old ones."""
|
||||
self._raw_dir.mkdir(parents=True, exist_ok=True)
|
||||
if not self._valid and not self._cleared:
|
||||
clear_dir_contents(self._raw_dir)
|
||||
self._cleared = True
|
||||
|
||||
def _read_manifest(self) -> dict | None:
|
||||
path = self._raw_dir / _MANIFEST_FILENAME
|
||||
if not path.is_file():
|
||||
return None
|
||||
try:
|
||||
payload = json.loads(path.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") != _CACHE_ENGINE:
|
||||
return None
|
||||
return payload
|
||||
|
||||
|
||||
__all__ = ["NativeImageRawCache", "native_md_options_signature"]
|
||||
@@ -0,0 +1,157 @@
|
||||
"""Shared template for native (local, in-process) parser engines.
|
||||
|
||||
``NativeParserBase.parse`` fixes the common local-parse flow once:
|
||||
|
||||
resolve + validate source → compute parsed_dir/asset_dir
|
||||
→ pre-clean (rmtree parsed_dir + mkdir + mkdir asset_dir, with rollback)
|
||||
→ extract() in a thread → build_ir() → write_sidecar(clean_parsed_dir=False)
|
||||
→ persist full_docs (lightrag) → archive source
|
||||
|
||||
Subclasses implement ``extract`` (sync, runs in a thread) and ``build_ir``.
|
||||
Currently only :class:`NativeDocxParser`; xlsx/pptx/md land later as new
|
||||
subclasses implementing the same two hooks.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import shutil
|
||||
import time
|
||||
from abc import abstractmethod
|
||||
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
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from lightrag.sidecar.ir import IRDoc
|
||||
|
||||
|
||||
class NativeParserBase(BaseParser):
|
||||
"""Base for engines that parse a file locally into a sidecar."""
|
||||
|
||||
# ``write_sidecar`` block_drawing_path_style; docx keeps the legacy
|
||||
# "basename_only" shape for byte-equivalence.
|
||||
sidecar_path_style: str = "with_prefix"
|
||||
# Prefix used in the "empty content" error message.
|
||||
empty_content_label: str = "Native"
|
||||
|
||||
# --- engine-private hooks ------------------------------------------------
|
||||
def validate_source(self, source: Path, file_path: str) -> None:
|
||||
"""Validate the resolved source (default: must be an existing file)."""
|
||||
if not (source.exists() and source.is_file()):
|
||||
raise FileNotFoundError(
|
||||
f"{self.engine_name} source file not found: {source}"
|
||||
)
|
||||
|
||||
@abstractmethod
|
||||
def extract(
|
||||
self, source: Path, *, parsed_dir: Path, asset_dir: Path, base_name: str
|
||||
) -> tuple[list[dict[str, Any]], dict[str, Any], dict[str, Any]]:
|
||||
"""Extract ``(blocks, warnings, metadata)`` (sync; runs in a thread).
|
||||
|
||||
``parsed_dir`` and ``asset_dir`` are pre-created by the template; the
|
||||
hook may write side artifacts (e.g. image bytes) into ``asset_dir``
|
||||
before :func:`write_sidecar` runs with ``clean_parsed_dir=False``.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def build_ir(
|
||||
self,
|
||||
blocks: list[dict[str, Any]],
|
||||
*,
|
||||
document_name: str,
|
||||
asset_dir_name: str,
|
||||
metadata: dict[str, Any],
|
||||
) -> "IRDoc": ...
|
||||
|
||||
def surface_warnings(
|
||||
self, warnings: dict[str, Any], source: Path
|
||||
) -> dict[str, Any] | None:
|
||||
"""Map parser warnings to the ``parse_warnings`` result field (opt)."""
|
||||
return None
|
||||
|
||||
# --- template ------------------------------------------------------------
|
||||
async def parse(self, ctx: ParseContext) -> ParseResult:
|
||||
from lightrag.sidecar import write_sidecar
|
||||
from lightrag.utils_pipeline import (
|
||||
make_lightrag_doc_content,
|
||||
sidecar_uri_for,
|
||||
)
|
||||
|
||||
rs = ctx.resolve(self.engine_name)
|
||||
source = rs.source_path
|
||||
self.validate_source(source, ctx.file_path)
|
||||
|
||||
document_name = rs.document_name
|
||||
base_name = Path(document_name).stem or document_name
|
||||
parsed_dir = rs.parsed_dir
|
||||
asset_dir = parsed_dir / f"{base_name}.blocks.assets"
|
||||
|
||||
def _extract_sync():
|
||||
# Pre-clean parsed_dir and pre-create asset_dir so the extractor
|
||||
# can write image bytes BEFORE write_sidecar (clean_parsed_dir=False
|
||||
# then keeps them). parsed_artifact_dir_for returns a unique dir per
|
||||
# source, so this rmtree only clobbers a prior attempt's artifacts.
|
||||
if parsed_dir.exists():
|
||||
shutil.rmtree(parsed_dir)
|
||||
parsed_dir.mkdir(parents=True, exist_ok=True)
|
||||
asset_dir.mkdir(parents=True, exist_ok=True)
|
||||
return self.extract(
|
||||
source, parsed_dir=parsed_dir, asset_dir=asset_dir, base_name=base_name
|
||||
)
|
||||
|
||||
try:
|
||||
blocks, warnings, metadata = await asyncio.to_thread(_extract_sync)
|
||||
except BaseException:
|
||||
# Roll back the pre-created (possibly partial) dirs on any failure.
|
||||
if parsed_dir.exists():
|
||||
shutil.rmtree(parsed_dir, ignore_errors=True)
|
||||
raise
|
||||
if not blocks:
|
||||
if parsed_dir.exists():
|
||||
shutil.rmtree(parsed_dir, ignore_errors=True)
|
||||
raise ValueError(
|
||||
f"{self.empty_content_label} parser returned empty content "
|
||||
f"for {ctx.file_path}"
|
||||
)
|
||||
|
||||
parse_warnings = self.surface_warnings(warnings, source)
|
||||
ir = self.build_ir(
|
||||
blocks,
|
||||
document_name=document_name,
|
||||
asset_dir_name=asset_dir.name,
|
||||
metadata=metadata,
|
||||
)
|
||||
parsed_data = write_sidecar(
|
||||
ir,
|
||||
parsed_dir=parsed_dir,
|
||||
doc_id=ctx.doc_id,
|
||||
engine=self.engine_name,
|
||||
clean_parsed_dir=False, # asset dir pre-populated above
|
||||
block_drawing_path_style=self.sidecar_path_style,
|
||||
)
|
||||
|
||||
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(parsed_dir),
|
||||
"parse_engine": self.engine_name,
|
||||
"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_warnings=parse_warnings,
|
||||
)
|
||||
@@ -0,0 +1,59 @@
|
||||
"""Native engine entry point — dispatches by source suffix.
|
||||
|
||||
The registry exposes a single ``native`` engine with one ``impl``; routing
|
||||
yields the engine key and ``get_parser("native")`` returns this one instance.
|
||||
Suffix capabilities (``docx`` / ``md`` / ``textpack``) are declared on the
|
||||
registry spec but do NOT pick an implementation — that is this dispatcher's
|
||||
job.
|
||||
|
||||
It delegates the **entire** ``parse(ctx)`` to the matching concrete parser
|
||||
(rather than subclassing :class:`NativeParserBase` and only forwarding the
|
||||
``extract`` / ``build_ir`` hooks) so each concrete parser keeps its own
|
||||
``sidecar_path_style`` — docx stays ``basename_only`` (byte-equivalent golden
|
||||
output) while markdown uses ``with_prefix``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
from lightrag.constants import PARSER_ENGINE_NATIVE
|
||||
from lightrag.parser.base import BaseParser, ParseContext, ParseResult
|
||||
|
||||
_MARKDOWN_SUFFIXES = {".md", ".textpack"}
|
||||
|
||||
|
||||
class NativeParser(BaseParser):
|
||||
"""Routes a document to the docx or markdown native parser by suffix."""
|
||||
|
||||
engine_name = PARSER_ENGINE_NATIVE
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._docx: BaseParser | None = None
|
||||
self._markdown: BaseParser | None = None
|
||||
|
||||
def _docx_parser(self) -> BaseParser:
|
||||
parser = self._docx
|
||||
if parser is None:
|
||||
from lightrag.parser.docx.parser import NativeDocxParser
|
||||
|
||||
parser = NativeDocxParser()
|
||||
self._docx = parser
|
||||
return parser
|
||||
|
||||
def _markdown_parser(self) -> BaseParser:
|
||||
parser = self._markdown
|
||||
if parser is None:
|
||||
from lightrag.parser.markdown.parser import NativeMarkdownParser
|
||||
|
||||
parser = NativeMarkdownParser()
|
||||
self._markdown = parser
|
||||
return parser
|
||||
|
||||
async def parse(self, ctx: ParseContext) -> ParseResult:
|
||||
suffix = Path(ctx.file_path).suffix.lower()
|
||||
if suffix in _MARKDOWN_SUFFIXES:
|
||||
return await self._markdown_parser().parse(ctx)
|
||||
# Default to docx; its validate_source raises a clear error for any
|
||||
# other suffix the native engine was (mis)routed for.
|
||||
return await self._docx_parser().parse(ctx)
|
||||
@@ -0,0 +1,67 @@
|
||||
"""No-op format handlers for the two "nothing to parse" document formats.
|
||||
|
||||
These are internal (``user_selectable=False``) parsers dispatched by *format*,
|
||||
not by a user engine choice:
|
||||
|
||||
- :class:`ReuseParser` handles ``lightrag`` rows — a document already parsed
|
||||
by some engine whose sidecar exists. Reached only on resume/retry (a doc
|
||||
that finished parsing but failed/interrupted at a later stage and is pulled
|
||||
back through the parse queue). Re-uses the stored content + sidecar
|
||||
instead of re-parsing (the original source may already be archived).
|
||||
- :class:`PassthroughParser` handles ``raw`` rows — content supplied verbatim
|
||||
at insert time (e.g. ``ainsert`` of a plain string).
|
||||
|
||||
Both mirror the corresponding branches of the former ``parse_native`` exactly
|
||||
(no disk writes, ``parse_stage_skipped=True``, no ``parse_engine`` key).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from lightrag.constants import FULL_DOCS_FORMAT_LIGHTRAG, FULL_DOCS_FORMAT_RAW
|
||||
from lightrag.parser.base import BaseParser, ParseContext, ParseResult
|
||||
|
||||
|
||||
class ReuseParser(BaseParser):
|
||||
"""Reuse an already-parsed (``lightrag``-format) document's sidecar."""
|
||||
|
||||
engine_name = "reuse"
|
||||
|
||||
async def parse(self, ctx: ParseContext) -> ParseResult:
|
||||
from lightrag.utils_pipeline import (
|
||||
sidecar_blocks_path,
|
||||
strip_lightrag_doc_prefix,
|
||||
)
|
||||
|
||||
doc_format = ctx.content_data.get("parse_format", FULL_DOCS_FORMAT_LIGHTRAG)
|
||||
merged_text = strip_lightrag_doc_prefix(
|
||||
ctx.content_data.get("content"), doc_format
|
||||
)
|
||||
# ``sidecar_location`` may be absent on historical/abnormal rows; tolerate
|
||||
# it (blocks_path="") rather than failing or re-routing to extraction.
|
||||
blocks_path = (
|
||||
sidecar_blocks_path(ctx.content_data.get("sidecar_location")) or ""
|
||||
)
|
||||
return ParseResult(
|
||||
doc_id=ctx.doc_id,
|
||||
file_path=ctx.file_path,
|
||||
parse_format=doc_format,
|
||||
content=merged_text,
|
||||
blocks_path=blocks_path,
|
||||
parse_stage_skipped=True,
|
||||
)
|
||||
|
||||
|
||||
class PassthroughParser(BaseParser):
|
||||
"""Pass ``raw``-format content through verbatim (no parser ran)."""
|
||||
|
||||
engine_name = "passthrough"
|
||||
|
||||
async def parse(self, ctx: ParseContext) -> ParseResult:
|
||||
return ParseResult(
|
||||
doc_id=ctx.doc_id,
|
||||
file_path=ctx.file_path,
|
||||
parse_format=FULL_DOCS_FORMAT_RAW,
|
||||
content=ctx.content_data.get("content", ""),
|
||||
blocks_path="",
|
||||
parse_stage_skipped=True,
|
||||
)
|
||||
@@ -0,0 +1,642 @@
|
||||
"""Parameter schema + parenthesis-aware scanners for hint parameters.
|
||||
|
||||
This module is the single source of truth for the *parameters* that may be
|
||||
attached to chunk-strategy selectors inside a parser hint or a
|
||||
``LIGHTRAG_PARSER`` rule, e.g. ``[-R(chunk_ts=800,chunk_ol=80)]`` or
|
||||
``pdf:legacy-R(chunk_ts=800)``.
|
||||
|
||||
Design (see ``docs/FileProcessingPipeline.md``):
|
||||
|
||||
* Inside ``(...)`` a comma is **only** a parameter separator; a single
|
||||
parameter value never contains ``,`` ``(`` ``)`` or ``]``.
|
||||
* Parameter names use a readable ``canonical`` form with optional short
|
||||
``alias`` forms; everything is normalised to canonical before use, so the
|
||||
internal structures never carry an alias.
|
||||
* Parameters are declared per *target* (a chunk selector char ``F``/``R``/``V``/
|
||||
``P``) with a type and the set of selectors they apply to.
|
||||
|
||||
Chunk-parameter scope: the integer ``chunk_token_size`` (alias ``chunk_ts``)
|
||||
and ``chunk_overlap_token_size`` (alias ``chunk_ol``), plus the boolean
|
||||
``drop_references`` (alias ``drop_rf``, paragraph-semantic only), all flowing
|
||||
through the existing per-document ``chunk_options`` channel. Float/enum chunk
|
||||
parameters are still rejected with a friendly error; adding them is a matter of
|
||||
registering new :class:`ParamSpec` entries and an extra ``kind`` branch in
|
||||
:func:`parse_chunk_params`.
|
||||
|
||||
This module is import-cheap and has no dependency on
|
||||
:mod:`lightrag.parser.routing` so it can be reused without import cycles.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import re
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
from lightrag.constants import (
|
||||
PARSER_ENGINE_DOCLING,
|
||||
PARSER_ENGINE_MINERU,
|
||||
PROCESS_OPTION_CHUNK_FIXED,
|
||||
PROCESS_OPTION_CHUNK_PARAGRAH,
|
||||
PROCESS_OPTION_CHUNK_RECURSIVE,
|
||||
PROCESS_OPTION_CHUNK_VECTOR,
|
||||
)
|
||||
|
||||
# Characters a parameter value may never contain. ``]`` can never appear
|
||||
# anyway (``_PARSER_HINT_RE`` terminates the bracket on ``]``) but is rejected
|
||||
# explicitly so the error is friendly rather than a silent truncation.
|
||||
_VALUE_FORBIDDEN = frozenset(",()]")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Parenthesis-aware scanners (shared by routing's rule/engine/options splits)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def split_top_level(text: str, separators: str) -> list[str]:
|
||||
"""Split ``text`` on any char in ``separators`` at parenthesis depth 0.
|
||||
|
||||
Characters inside ``(...)`` are protected, so a comma used to separate
|
||||
parameters never splits the surrounding rule / options string. Always
|
||||
returns at least one element (the whole string when no separator is hit).
|
||||
"""
|
||||
parts: list[str] = []
|
||||
depth = 0
|
||||
start = 0
|
||||
for i, ch in enumerate(text):
|
||||
if ch == "(":
|
||||
depth += 1
|
||||
elif ch == ")":
|
||||
if depth > 0:
|
||||
depth -= 1
|
||||
elif depth == 0 and ch in separators:
|
||||
parts.append(text[start:i])
|
||||
start = i + 1
|
||||
parts.append(text[start:])
|
||||
return parts
|
||||
|
||||
|
||||
def take_paren_block(text: str, i: int) -> tuple[str | None, int]:
|
||||
"""Read a balanced ``(...)`` block starting at ``text[i]``.
|
||||
|
||||
Returns ``(inner, index_just_after_closing_paren)`` when ``text[i]`` opens
|
||||
a balanced block, otherwise ``(None, i)`` — used for both "no block here"
|
||||
(``text[i] != '('``) and "unterminated block". Callers that need to flag
|
||||
an unterminated block compare the returned index / ``None`` accordingly.
|
||||
"""
|
||||
if i >= len(text) or text[i] != "(":
|
||||
return None, i
|
||||
depth = 0
|
||||
for j in range(i, len(text)):
|
||||
if text[j] == "(":
|
||||
depth += 1
|
||||
elif text[j] == ")":
|
||||
depth -= 1
|
||||
if depth == 0:
|
||||
return text[i + 1 : j], j + 1
|
||||
return None, i # unterminated
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Parameter schema registry
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ParamSpec:
|
||||
"""Declares one tunable hint parameter for a set of chunk selectors.
|
||||
|
||||
``kind`` is ``"int"`` or ``"bool"`` today; the field exists so future
|
||||
types (``float``/``enum``/``range_list``) slot in without a structural
|
||||
change. ``targets`` is the set of chunk selector chars the parameter is
|
||||
valid for (e.g. overlap is invalid for ``V``, ``drop_references`` only
|
||||
applies to ``P``).
|
||||
"""
|
||||
|
||||
canonical: str
|
||||
aliases: frozenset[str]
|
||||
kind: str
|
||||
targets: frozenset[str]
|
||||
is_list: bool = False
|
||||
min_value: int | None = None
|
||||
names: frozenset[str] = field(init=False)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
object.__setattr__(self, "names", frozenset({self.canonical}) | self.aliases)
|
||||
|
||||
|
||||
_ALL_CHUNK_SELECTORS = frozenset(
|
||||
{
|
||||
PROCESS_OPTION_CHUNK_FIXED,
|
||||
PROCESS_OPTION_CHUNK_RECURSIVE,
|
||||
PROCESS_OPTION_CHUNK_VECTOR,
|
||||
PROCESS_OPTION_CHUNK_PARAGRAH,
|
||||
}
|
||||
)
|
||||
|
||||
# Phase 1 registry — keep this list small; expand it (and only it) to grow
|
||||
# hint-parameter coverage.
|
||||
_CHUNK_PARAM_SPECS: tuple[ParamSpec, ...] = (
|
||||
ParamSpec(
|
||||
canonical="chunk_token_size",
|
||||
aliases=frozenset({"chunk_ts"}),
|
||||
kind="int",
|
||||
targets=_ALL_CHUNK_SELECTORS,
|
||||
min_value=1,
|
||||
),
|
||||
ParamSpec(
|
||||
canonical="chunk_overlap_token_size",
|
||||
aliases=frozenset({"chunk_ol"}),
|
||||
kind="int",
|
||||
# Semantic-vector (V) chunking has no overlap concept.
|
||||
targets=frozenset(
|
||||
{
|
||||
PROCESS_OPTION_CHUNK_FIXED,
|
||||
PROCESS_OPTION_CHUNK_RECURSIVE,
|
||||
PROCESS_OPTION_CHUNK_PARAGRAH,
|
||||
}
|
||||
),
|
||||
min_value=0,
|
||||
),
|
||||
# Paragraph-semantic only: drop the trailing reference section before
|
||||
# chunking. Detection-tuning knobs (tail window / heading prefixes) are
|
||||
# env-only and read live by the chunker, so only the switch is a hint param.
|
||||
ParamSpec(
|
||||
canonical="drop_references",
|
||||
aliases=frozenset({"drop_rf"}),
|
||||
kind="bool",
|
||||
targets=frozenset({PROCESS_OPTION_CHUNK_PARAGRAH}),
|
||||
),
|
||||
)
|
||||
|
||||
_CHUNK_PARAM_BY_NAME: dict[str, ParamSpec] = {}
|
||||
for _spec in _CHUNK_PARAM_SPECS:
|
||||
for _name in _spec.names:
|
||||
_CHUNK_PARAM_BY_NAME[_name] = _spec
|
||||
|
||||
|
||||
def supported_chunk_param_names() -> str:
|
||||
"""Comma-joined canonical names, for friendly error messages."""
|
||||
return ", ".join(sorted(spec.canonical for spec in _CHUNK_PARAM_SPECS))
|
||||
|
||||
|
||||
def _parse_int(value: str) -> int | None:
|
||||
try:
|
||||
return int(value, 10)
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
|
||||
|
||||
def parse_chunk_params(
|
||||
text: str, *, selector: str, label: str
|
||||
) -> tuple[dict[str, Any], list[str]]:
|
||||
"""Parse one chunk-strategy parameter block into a canonical dict.
|
||||
|
||||
``text`` is the raw text inside ``(...)`` (parameter separators only —
|
||||
no surrounding parens). ``selector`` is the chunk char the block is
|
||||
attached to (``F``/``R``/``V``/``P``). Returns ``(canonical_dict,
|
||||
errors)``; ``errors`` is empty iff the block is fully valid. Aliases are
|
||||
normalised to their canonical name in the returned dict.
|
||||
|
||||
A boolean parameter may be written bare as a flag — ``P(drop_rf)`` is
|
||||
shorthand for ``P(drop_rf=true)``. Non-boolean parameters still require
|
||||
the explicit ``key=value`` form.
|
||||
"""
|
||||
result: dict[str, Any] = {}
|
||||
errors: list[str] = []
|
||||
seen: set[str] = set()
|
||||
|
||||
for raw in split_top_level(text, ","):
|
||||
segment = raw.strip()
|
||||
if not segment:
|
||||
errors.append(f"{label}: empty parameter")
|
||||
continue
|
||||
if "=" in segment:
|
||||
key, _, value = segment.partition("=")
|
||||
key = key.strip()
|
||||
value = value.strip()
|
||||
flag_form = False
|
||||
else:
|
||||
# Bare flag form, e.g. ``drop_rf``. Only valid for boolean
|
||||
# parameters, where it is shorthand for ``drop_rf=true``.
|
||||
key = segment
|
||||
value = ""
|
||||
flag_form = True
|
||||
|
||||
spec = _CHUNK_PARAM_BY_NAME.get(key)
|
||||
if spec is None:
|
||||
errors.append(
|
||||
f"{label}: unknown parameter {key!r}; supported parameters: "
|
||||
f"{supported_chunk_param_names()}"
|
||||
)
|
||||
continue
|
||||
if selector not in spec.targets:
|
||||
errors.append(
|
||||
f"{label}: parameter {spec.canonical!r} is not supported for "
|
||||
f"chunk strategy {selector!r}"
|
||||
)
|
||||
continue
|
||||
if flag_form:
|
||||
if spec.kind != "bool":
|
||||
errors.append(
|
||||
f"{label}: parameter {spec.canonical!r} must be written as "
|
||||
"'key=value'; only boolean flags may be written bare"
|
||||
)
|
||||
continue
|
||||
value = "true" # bare boolean flag means True
|
||||
if any(ch in _VALUE_FORBIDDEN for ch in value):
|
||||
errors.append(
|
||||
f"{label}: value for {spec.canonical!r} may not contain any of "
|
||||
"',' '(' ')' ']'"
|
||||
)
|
||||
continue
|
||||
if spec.canonical in seen and not spec.is_list:
|
||||
errors.append(f"{label}: parameter {spec.canonical!r} may not be repeated")
|
||||
continue
|
||||
seen.add(spec.canonical)
|
||||
|
||||
if spec.kind == "int":
|
||||
parsed = _parse_int(value)
|
||||
if parsed is None:
|
||||
errors.append(
|
||||
f"{label}: value for {spec.canonical!r} must be an integer, "
|
||||
f"got {value!r}"
|
||||
)
|
||||
continue
|
||||
if spec.min_value is not None and parsed < spec.min_value:
|
||||
errors.append(
|
||||
f"{label}: value for {spec.canonical!r} must be >= "
|
||||
f"{spec.min_value}, got {parsed}"
|
||||
)
|
||||
continue
|
||||
result[spec.canonical] = parsed
|
||||
elif spec.kind == "bool":
|
||||
parsed_bool = _parse_bool(value)
|
||||
if parsed_bool is None:
|
||||
errors.append(
|
||||
f"{label}: value for {spec.canonical!r} must be a boolean "
|
||||
f"(true/false), got {value!r}"
|
||||
)
|
||||
continue
|
||||
result[spec.canonical] = parsed_bool
|
||||
else: # pragma: no cover - only int/bool kinds registered today
|
||||
errors.append(
|
||||
f"{label}: parameter {spec.canonical!r} has unsupported type "
|
||||
f"{spec.kind!r}"
|
||||
)
|
||||
|
||||
# Cross-field invariant: reject an explicit overlap >= size pair here so
|
||||
# every caller (rule startup validation AND filename-hint validation)
|
||||
# rejects it uniformly, instead of only failing later at enqueue.
|
||||
overlap_error = chunk_param_overlap_error(result)
|
||||
if overlap_error is not None:
|
||||
errors.append(f"{label}: {overlap_error}")
|
||||
|
||||
return result, errors
|
||||
|
||||
|
||||
def chunk_param_overlap_error(params: Mapping[str, Any]) -> str | None:
|
||||
"""Return an error string when a block sets an invalid overlap/size pair.
|
||||
|
||||
Only checks the cross-field invariant when **both** ``chunk_token_size`` and
|
||||
``chunk_overlap_token_size`` are explicitly present in ``params`` (a parsed
|
||||
canonical dict). When only one is given the effective value depends on
|
||||
env / ``addon_params`` and cannot be evaluated here — that case is left to
|
||||
the upload-time ``_validate_effective_chunk_overlap`` on the resolved
|
||||
snapshot. Returns ``None`` when the pair is valid or not fully specified.
|
||||
"""
|
||||
size = params.get("chunk_token_size")
|
||||
overlap = params.get("chunk_overlap_token_size")
|
||||
if size is not None and overlap is not None and overlap >= size:
|
||||
return (
|
||||
f"chunk_overlap_token_size ({overlap}) must be < chunk_token_size ({size})"
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Engine parameters (Phase 2) — per-file params attached to the engine token,
|
||||
# e.g. ``mineru(page_range=1-3,language=en)`` / ``docling(force_ocr=true)``.
|
||||
# Keyed by engine name (unlike chunk params, which are keyed by F/R/V/P).
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_BOOL_TRUE = frozenset({"1", "true", "yes", "on", "t", "y"})
|
||||
_BOOL_FALSE = frozenset({"0", "false", "no", "off", "f", "n"})
|
||||
# A single page-range segment: ``N`` or ``N-M`` (validated for positivity /
|
||||
# ordering separately).
|
||||
_PAGE_SEGMENT_RE = re.compile(r"^\d+(?:-\d+)?$")
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class EngineParamSpec:
|
||||
"""Declares one tunable engine parameter (Phase 2).
|
||||
|
||||
Separate from the chunk :class:`ParamSpec` (which requires a ``targets``
|
||||
set of F/R/V/P selectors that is meaningless for engines). ``kind`` is one
|
||||
of ``"str"`` / ``"enum"`` / ``"bool"``. ``is_list`` marks a repeated-key
|
||||
parameter (``page_range``) whose canonical value is a comma-joined string.
|
||||
``enum_values`` constrains an ``"enum"`` parameter.
|
||||
"""
|
||||
|
||||
canonical: str
|
||||
aliases: frozenset[str]
|
||||
kind: str
|
||||
is_list: bool = False
|
||||
enum_values: frozenset[str] | None = None
|
||||
names: frozenset[str] = field(init=False)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
object.__setattr__(self, "names", frozenset({self.canonical}) | self.aliases)
|
||||
|
||||
|
||||
# Phase 2 registry — keyed by engine name; expand by adding specs (and, for a
|
||||
# new engine, a new dict entry). An engine absent here accepts NO parameters.
|
||||
_ENGINE_PARAM_SPECS: dict[str, tuple[EngineParamSpec, ...]] = {
|
||||
PARSER_ENGINE_MINERU: (
|
||||
EngineParamSpec(
|
||||
canonical="page_range",
|
||||
aliases=frozenset({"pr"}),
|
||||
kind="str",
|
||||
is_list=True,
|
||||
),
|
||||
EngineParamSpec(canonical="language", aliases=frozenset(), kind="str"),
|
||||
EngineParamSpec(
|
||||
canonical="local_parse_method",
|
||||
aliases=frozenset({"local_pm"}),
|
||||
kind="enum",
|
||||
enum_values=frozenset({"auto", "txt", "ocr"}),
|
||||
),
|
||||
),
|
||||
PARSER_ENGINE_DOCLING: (
|
||||
EngineParamSpec(canonical="force_ocr", aliases=frozenset({"ocr"}), kind="bool"),
|
||||
),
|
||||
}
|
||||
|
||||
_ENGINE_PARAM_BY_NAME: dict[str, dict[str, EngineParamSpec]] = {
|
||||
engine: {name: spec for spec in specs for name in spec.names}
|
||||
for engine, specs in _ENGINE_PARAM_SPECS.items()
|
||||
}
|
||||
|
||||
|
||||
def engine_params_supported(engine: str) -> bool:
|
||||
"""Whether ``engine`` declares any tunable hint parameters."""
|
||||
return engine in _ENGINE_PARAM_SPECS
|
||||
|
||||
|
||||
def supported_engine_param_names(engine: str) -> str:
|
||||
"""Comma-joined canonical engine-param names, for error messages."""
|
||||
specs = _ENGINE_PARAM_SPECS.get(engine, ())
|
||||
return ", ".join(sorted(spec.canonical for spec in specs))
|
||||
|
||||
|
||||
def _parse_bool(value: str) -> bool | None:
|
||||
low = value.strip().lower()
|
||||
if low in _BOOL_TRUE:
|
||||
return True
|
||||
if low in _BOOL_FALSE:
|
||||
return False
|
||||
return None
|
||||
|
||||
|
||||
def _mineru_api_mode_is_local() -> bool:
|
||||
"""True when MinerU runs in local mode (the default when unset).
|
||||
|
||||
Read here (rather than importing ``mineru.cache``) to keep this module a
|
||||
dependency-free leaf. Mirrors ``cache._normalize_api_mode``: anything that
|
||||
is not ``official`` is treated as local.
|
||||
"""
|
||||
return (os.getenv("MINERU_API_MODE", "") or "").strip().lower() != "official"
|
||||
|
||||
|
||||
def _validate_page_range_segments(parts: list[str]) -> list[str]:
|
||||
"""Validate page-range segment shape + the MinerU local single-segment rule.
|
||||
|
||||
``official`` mode forwards a multi-segment list verbatim; ``local`` accepts
|
||||
only a single page / range (mirrors ``cache.local_page_bounds``). The
|
||||
download-time ``local_page_bounds`` remains the final backstop.
|
||||
"""
|
||||
errors: list[str] = []
|
||||
for seg in parts:
|
||||
if not _PAGE_SEGMENT_RE.match(seg):
|
||||
errors.append(
|
||||
f"page_range segment {seg!r} must be a page 'N' or range 'N-M'"
|
||||
)
|
||||
continue
|
||||
if "-" in seg:
|
||||
left, _, right = seg.partition("-")
|
||||
if int(left) < 1 or int(right) < 1:
|
||||
errors.append(f"page_range segment {seg!r} must use positive pages")
|
||||
elif int(right) < int(left):
|
||||
errors.append(f"page_range segment {seg!r} must have end >= start")
|
||||
elif int(seg) < 1:
|
||||
errors.append(f"page_range segment {seg!r} must be a positive page")
|
||||
if not errors and len(parts) > 1 and _mineru_api_mode_is_local():
|
||||
errors.append(
|
||||
"page_range with MINERU_API_MODE=local supports only a single page "
|
||||
"or range such as '1-10'; use MINERU_API_MODE=official for a "
|
||||
"multi-segment list"
|
||||
)
|
||||
return errors
|
||||
|
||||
|
||||
def _coerce_engine_value(
|
||||
spec: EngineParamSpec, value: str, *, label: str
|
||||
) -> tuple[Any, str | None]:
|
||||
"""Validate + coerce a single engine-param value to its canonical type.
|
||||
|
||||
Returns ``(coerced, error)``; ``error`` is ``None`` when valid. Shared by
|
||||
the text path (:func:`parse_engine_params`) and the resolved-dict path
|
||||
(:func:`normalize_engine_params`) so both apply identical rules. For the
|
||||
list-type ``page_range`` the ``value`` is the already-joined comma string.
|
||||
"""
|
||||
# ``local_parse_method`` only feeds the local MinerU request + signature;
|
||||
# the official API neither sends it nor folds it into the cache key, so
|
||||
# accepting it under official mode would persist a directive that silently
|
||||
# does nothing. Reject it here (mirrors the page_range mode rule).
|
||||
if spec.canonical == "local_parse_method" and not _mineru_api_mode_is_local():
|
||||
return None, (
|
||||
f"{label}: 'local_parse_method' only applies to "
|
||||
"MINERU_API_MODE=local (the default); the official API ignores it"
|
||||
)
|
||||
if spec.kind == "bool":
|
||||
parsed = _parse_bool(value)
|
||||
if parsed is None:
|
||||
return None, (
|
||||
f"{label}: value for {spec.canonical!r} must be a boolean "
|
||||
f"(true/false), got {value!r}"
|
||||
)
|
||||
return parsed, None
|
||||
if spec.kind == "enum":
|
||||
if spec.enum_values is not None and value not in spec.enum_values:
|
||||
allowed = ", ".join(sorted(spec.enum_values))
|
||||
return None, (
|
||||
f"{label}: value for {spec.canonical!r} must be one of "
|
||||
f"{allowed}, got {value!r}"
|
||||
)
|
||||
return value, None
|
||||
# str (incl. the list-type page_range, whose value is a comma-joined string)
|
||||
if not value:
|
||||
return None, f"{label}: value for {spec.canonical!r} must be non-empty"
|
||||
if spec.is_list and spec.canonical == "page_range":
|
||||
segments = [p.strip() for p in value.split(",")]
|
||||
seg_errors = _validate_page_range_segments(segments)
|
||||
if seg_errors:
|
||||
return None, f"{label}: " + "; ".join(seg_errors)
|
||||
return ",".join(segments), None
|
||||
return value, None
|
||||
|
||||
|
||||
def parse_engine_params(
|
||||
text: str, *, engine: str, label: str
|
||||
) -> tuple[dict[str, Any], list[str]]:
|
||||
"""Parse one engine parameter block into a canonical, coerced dict.
|
||||
|
||||
``text`` is the raw text inside ``(...)`` for an engine token; ``engine`` is
|
||||
the (bare) engine the block is attached to. Returns ``(canonical_dict,
|
||||
errors)``; aliases are normalised to canonical and values are coerced to
|
||||
their declared type (so ``force_ocr`` is a real ``bool``). A list-type
|
||||
``page_range`` collects repeated keys and joins them with ``,``.
|
||||
"""
|
||||
by_name = _ENGINE_PARAM_BY_NAME.get(engine)
|
||||
if by_name is None:
|
||||
if text.strip():
|
||||
return {}, [f"{label}: parser engine {engine!r} does not accept parameters"]
|
||||
return {}, []
|
||||
|
||||
result: dict[str, Any] = {}
|
||||
list_values: dict[str, list[str]] = {}
|
||||
errors: list[str] = []
|
||||
seen: set[str] = set()
|
||||
|
||||
for raw in split_top_level(text, ","):
|
||||
segment = raw.strip()
|
||||
if not segment:
|
||||
errors.append(f"{label}: empty parameter")
|
||||
continue
|
||||
if "=" not in segment:
|
||||
if _PAGE_SEGMENT_RE.match(segment) and (
|
||||
"page_range" in by_name or "pr" in by_name
|
||||
):
|
||||
errors.append(
|
||||
f"{label}: page lists must repeat the key, e.g. "
|
||||
"'page_range=1-3,page_range=5' (a comma only separates "
|
||||
"parameters)"
|
||||
)
|
||||
else:
|
||||
errors.append(
|
||||
f"{label}: parameter {segment!r} must be written as "
|
||||
"'key=value' (flag parameters are not supported yet)"
|
||||
)
|
||||
continue
|
||||
key, _, value = segment.partition("=")
|
||||
key = key.strip()
|
||||
value = value.strip()
|
||||
|
||||
spec = by_name.get(key)
|
||||
if spec is None:
|
||||
errors.append(
|
||||
f"{label}: unknown parameter {key!r} for engine {engine!r}; "
|
||||
f"supported parameters: {supported_engine_param_names(engine)}"
|
||||
)
|
||||
continue
|
||||
if any(ch in _VALUE_FORBIDDEN for ch in value):
|
||||
errors.append(
|
||||
f"{label}: value for {spec.canonical!r} may not contain any of "
|
||||
"',' '(' ')' ']'"
|
||||
)
|
||||
continue
|
||||
if spec.is_list:
|
||||
list_values.setdefault(spec.canonical, []).append(value)
|
||||
continue
|
||||
if spec.canonical in seen:
|
||||
errors.append(f"{label}: parameter {spec.canonical!r} may not be repeated")
|
||||
continue
|
||||
seen.add(spec.canonical)
|
||||
coerced, error = _coerce_engine_value(spec, value, label=label)
|
||||
if error is not None:
|
||||
errors.append(error)
|
||||
continue
|
||||
result[spec.canonical] = coerced
|
||||
|
||||
# Join repeated-key list params, then coerce/validate the joined value.
|
||||
for canonical, values in list_values.items():
|
||||
spec = by_name[canonical]
|
||||
coerced, error = _coerce_engine_value(spec, ",".join(values), label=label)
|
||||
if error is not None:
|
||||
errors.append(error)
|
||||
continue
|
||||
result[canonical] = coerced
|
||||
|
||||
return result, errors
|
||||
|
||||
|
||||
def normalize_engine_params(
|
||||
engine: str, params: Mapping[str, Any]
|
||||
) -> tuple[dict[str, Any], list[str]]:
|
||||
"""Normalise + validate an already-resolved engine-param dict.
|
||||
|
||||
Used by the pipeline layer where direct (SDK/API) callers bypass routing's
|
||||
text parsing. Returns a **coerced** dict (e.g. ``force_ocr`` becomes a real
|
||||
``bool``, ``page_range`` a validated comma-joined string) so what gets
|
||||
persisted is exactly what the engine override seam consumes. Accepts a
|
||||
``page_range`` value as either a list or a comma-joined string.
|
||||
"""
|
||||
if not params:
|
||||
return {}, []
|
||||
by_name = _ENGINE_PARAM_BY_NAME.get(engine)
|
||||
if by_name is None:
|
||||
return {}, [f"parser engine {engine!r} does not accept parameters"]
|
||||
|
||||
result: dict[str, Any] = {}
|
||||
errors: list[str] = []
|
||||
for key, value in params.items():
|
||||
spec = by_name.get(str(key))
|
||||
if spec is None:
|
||||
errors.append(
|
||||
f"unknown parameter {key!r} for engine {engine!r}; supported "
|
||||
f"parameters: {supported_engine_param_names(engine)}"
|
||||
)
|
||||
continue
|
||||
if spec.is_list and isinstance(value, (list, tuple)):
|
||||
value = ",".join(str(v).strip() for v in value)
|
||||
else:
|
||||
value = str(value).strip()
|
||||
coerced, error = _coerce_engine_value(
|
||||
spec, value, label=f"engine parameter {spec.canonical!r}"
|
||||
)
|
||||
if error is not None:
|
||||
errors.append(error)
|
||||
continue
|
||||
result[spec.canonical] = coerced
|
||||
return result, errors
|
||||
|
||||
|
||||
def render_engine_params(
|
||||
engine: str, params: Mapping[str, Any]
|
||||
) -> tuple[str, list[str]]:
|
||||
"""Render a resolved engine-param dict to the canonical inner text.
|
||||
|
||||
Returns ``(inner_text, errors)`` where ``inner_text`` is the
|
||||
``key=value,...`` string that goes inside the ``parse_engine`` parens (e.g.
|
||||
``page_range=1-3,page_range=5,language=en``). Normalises first (so the
|
||||
output always round-trips through :func:`parse_engine_params`); a list-type
|
||||
value is emitted as **repeated keys**, a bool as ``true``/``false``. Keys
|
||||
are sorted for deterministic output.
|
||||
"""
|
||||
normalized, errors = normalize_engine_params(engine, params)
|
||||
if errors:
|
||||
return "", errors
|
||||
by_name = _ENGINE_PARAM_BY_NAME.get(engine, {})
|
||||
parts: list[str] = []
|
||||
for canonical in sorted(normalized):
|
||||
spec = by_name[canonical]
|
||||
value = normalized[canonical]
|
||||
if spec.is_list:
|
||||
parts.extend(f"{canonical}={seg}" for seg in str(value).split(","))
|
||||
elif spec.kind == "bool":
|
||||
parts.append(f"{canonical}={'true' if value else 'false'}")
|
||||
else:
|
||||
parts.append(f"{canonical}={value}")
|
||||
return ",".join(parts), []
|
||||
@@ -0,0 +1,66 @@
|
||||
"""Third-party parser engine discovery (``lightrag.parsers`` entry points).
|
||||
|
||||
A third-party package exposes parser engines by declaring an entry point in
|
||||
the ``lightrag.parsers`` group::
|
||||
|
||||
# pyproject.toml of the third-party package
|
||||
[project.entry-points."lightrag.parsers"]
|
||||
myengine = "my_pkg.lightrag_plugin:register"
|
||||
|
||||
Each entry point must resolve to a **zero-argument callable** that performs
|
||||
its own :func:`lightrag.parser.registry.register_parser` call(s). The
|
||||
callable should be import-cheap: defer the parser implementation import to
|
||||
the ``ParserSpec.impl`` string (the registry already loads it lazily).
|
||||
|
||||
:func:`load_third_party_parsers` is invoked once per process by both
|
||||
entrypoints that drive parsers — the API server (``create_app``, before
|
||||
routing-rule validation so ``LIGHTRAG_PARSER`` may reference third-party
|
||||
engine names) and the debug CLI (``lightrag.parser.cli.main``, before
|
||||
``--engine`` choices are built). Library users embedding LightRAG directly
|
||||
can call it themselves before constructing pipelines.
|
||||
|
||||
See ``docs/ThirdPartyParser-zh.md`` for the full plugin authoring guide.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from importlib.metadata import entry_points
|
||||
|
||||
ENTRY_POINT_GROUP = "lightrag.parsers"
|
||||
|
||||
logger = logging.getLogger("lightrag")
|
||||
|
||||
_loaded = False
|
||||
|
||||
|
||||
def load_third_party_parsers(*, force: bool = False) -> list[str]:
|
||||
"""Discover and run all ``lightrag.parsers`` entry points.
|
||||
|
||||
Idempotent per process (``force=True`` re-runs, for tests). Returns the
|
||||
names of the entry points that loaded successfully. A failing plugin is
|
||||
logged with its origin and skipped — one broken third-party package must
|
||||
not take down server startup or the debug CLI; the built-in engines are
|
||||
registered statically and are never affected.
|
||||
"""
|
||||
global _loaded
|
||||
if _loaded and not force:
|
||||
return []
|
||||
_loaded = True
|
||||
|
||||
loaded: list[str] = []
|
||||
for ep in entry_points(group=ENTRY_POINT_GROUP):
|
||||
try:
|
||||
register = ep.load()
|
||||
register()
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"[parser-plugins] failed to load parser plugin %r (%s): %s",
|
||||
ep.name,
|
||||
ep.value,
|
||||
e,
|
||||
)
|
||||
continue
|
||||
loaded.append(ep.name)
|
||||
logger.info("[parser-plugins] loaded parser plugin %r (%s)", ep.name, ep.value)
|
||||
return loaded
|
||||
@@ -0,0 +1,315 @@
|
||||
"""Central registry for parser engines (mirrors the storage-layer convention).
|
||||
|
||||
Like :data:`lightrag.kg.STORAGES` / ``STORAGE_IMPLEMENTATIONS``, this module
|
||||
holds a module-level literal table of lightweight :class:`ParserSpec`
|
||||
metadata. Loading this module imports **no** parser implementation, so
|
||||
capability queries (suffixes / endpoint / supported engines) never trigger a
|
||||
heavy ``mineru``/``docling`` facade import (which would pull ``httpx`` etc.).
|
||||
The implementation class is imported lazily — only when a document is
|
||||
actually parsed — via :func:`get_parser`.
|
||||
|
||||
Capability data lives here (single source of truth); behaviour lives in the
|
||||
parser classes. ``constants.PARSER_ENGINE_*`` keeps only the bare name
|
||||
strings used as identifiers / registry keys.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib
|
||||
import os
|
||||
from dataclasses import dataclass, field
|
||||
from typing import TYPE_CHECKING, Callable
|
||||
|
||||
from lightrag.constants import (
|
||||
PARSER_ENGINE_DOCLING,
|
||||
PARSER_ENGINE_LEGACY,
|
||||
PARSER_ENGINE_MINERU,
|
||||
PARSER_ENGINE_NATIVE,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from lightrag.parser.base import BaseParser
|
||||
|
||||
# Internal format-handler engine keys (not user-selectable).
|
||||
PARSER_ENGINE_REUSE = "reuse"
|
||||
PARSER_ENGINE_PASSTHROUGH = "passthrough"
|
||||
|
||||
_VALID_MINERU_API_MODES = {"official", "local"}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Endpoint capability closures (env-only; no network). Canonical home —
|
||||
# routing.py delegates here.
|
||||
# ---------------------------------------------------------------------------
|
||||
def _mineru_endpoint_configured() -> bool:
|
||||
mode = os.getenv("MINERU_API_MODE", "local").strip().lower()
|
||||
if mode == "official":
|
||||
return bool(os.getenv("MINERU_API_TOKEN", "").strip())
|
||||
if mode == "local":
|
||||
return bool(os.getenv("MINERU_LOCAL_ENDPOINT", "").strip())
|
||||
return False
|
||||
|
||||
|
||||
def _mineru_endpoint_requirement() -> str | None:
|
||||
mode = os.getenv("MINERU_API_MODE", "local").strip().lower()
|
||||
if mode == "official":
|
||||
return "MINERU_API_TOKEN"
|
||||
if mode == "local":
|
||||
return "MINERU_LOCAL_ENDPOINT"
|
||||
allowed = ", ".join(sorted(_VALID_MINERU_API_MODES))
|
||||
return f"valid MINERU_API_MODE ({allowed})"
|
||||
|
||||
|
||||
def _env_endpoint_configured(env_name: str) -> Callable[[], bool]:
|
||||
return lambda: bool(os.getenv(env_name, "").strip())
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ParserSpec:
|
||||
"""Lightweight, import-cheap metadata for one parser engine.
|
||||
|
||||
Holds everything the pipeline needs to *route* and *gate* a document
|
||||
without importing the parser implementation. ``impl`` is a
|
||||
``"module:Class"`` string imported lazily by :func:`get_parser`.
|
||||
"""
|
||||
|
||||
engine_name: str
|
||||
impl: str
|
||||
suffixes: frozenset[str]
|
||||
user_selectable: bool = True
|
||||
queue_group: str = PARSER_ENGINE_NATIVE
|
||||
# Worker count for this spec's queue_group. The registrant bakes in any
|
||||
# env override at registration time (e.g.
|
||||
# ``concurrency=int(os.getenv("MAX_PARALLEL_PARSE_MYENGINE", "3"))``),
|
||||
# mirroring how the built-in ``max_parallel_parse_*`` LightRAG fields read
|
||||
# their env. ``None`` means this spec does not own its group's concurrency
|
||||
# (built-in groups are sized by the LightRAG instance field instead).
|
||||
concurrency: int | None = None
|
||||
endpoint_configured: Callable[[], bool] = field(default=lambda: True)
|
||||
endpoint_requirement: Callable[[], str | None] = field(default=lambda: None)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Suffix capabilities (single source of truth; replaces
|
||||
# constants.PARSER_ENGINE_SUFFIX_CAPABILITIES).
|
||||
# ---------------------------------------------------------------------------
|
||||
_LEGACY_SUFFIXES = frozenset(
|
||||
{
|
||||
"txt",
|
||||
"md",
|
||||
"mdx",
|
||||
"pdf",
|
||||
"docx",
|
||||
"pptx",
|
||||
"xlsx",
|
||||
"rtf",
|
||||
"odt",
|
||||
"tex",
|
||||
"epub",
|
||||
"html",
|
||||
"htm",
|
||||
"csv",
|
||||
"json",
|
||||
"xml",
|
||||
"yaml",
|
||||
"yml",
|
||||
"log",
|
||||
"conf",
|
||||
"ini",
|
||||
"properties",
|
||||
"sql",
|
||||
"bat",
|
||||
"sh",
|
||||
"c",
|
||||
"h",
|
||||
"cpp",
|
||||
"hpp",
|
||||
"py",
|
||||
"java",
|
||||
"js",
|
||||
"ts",
|
||||
"swift",
|
||||
"go",
|
||||
"rb",
|
||||
"php",
|
||||
"css",
|
||||
"scss",
|
||||
"less",
|
||||
}
|
||||
)
|
||||
_MINERU_SUFFIXES = frozenset(
|
||||
{
|
||||
"pdf",
|
||||
"doc",
|
||||
"docx",
|
||||
"ppt",
|
||||
"pptx",
|
||||
"xls",
|
||||
"xlsx",
|
||||
"png",
|
||||
"jpg",
|
||||
"jpeg",
|
||||
"jp2",
|
||||
"webp",
|
||||
"gif",
|
||||
"bmp",
|
||||
}
|
||||
)
|
||||
_DOCLING_SUFFIXES = frozenset(
|
||||
{
|
||||
"pdf",
|
||||
"docx",
|
||||
"pptx",
|
||||
"xlsx",
|
||||
"md",
|
||||
"html",
|
||||
"xhtml",
|
||||
"png",
|
||||
"jpg",
|
||||
"jpeg",
|
||||
"tiff",
|
||||
"webp",
|
||||
"bmp",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
_REGISTRY: dict[str, ParserSpec] = {
|
||||
PARSER_ENGINE_NATIVE: ParserSpec(
|
||||
engine_name=PARSER_ENGINE_NATIVE,
|
||||
# Single ``native`` engine; the dispatcher picks docx vs markdown by
|
||||
# source suffix (see lightrag.parser.native_dispatch).
|
||||
impl="lightrag.parser.native_dispatch:NativeParser",
|
||||
suffixes=frozenset({"docx", "md", "textpack"}),
|
||||
queue_group=PARSER_ENGINE_NATIVE,
|
||||
# Built-in groups are sized by the LightRAG ``max_parallel_parse_*``
|
||||
# instance field (supports constructor override), so no spec-level
|
||||
# ``concurrency`` here.
|
||||
),
|
||||
PARSER_ENGINE_LEGACY: ParserSpec(
|
||||
engine_name=PARSER_ENGINE_LEGACY,
|
||||
impl="lightrag.parser.legacy.parser:LegacyParser",
|
||||
suffixes=_LEGACY_SUFFIXES,
|
||||
queue_group=PARSER_ENGINE_NATIVE, # shares native pool (local, no network)
|
||||
),
|
||||
PARSER_ENGINE_MINERU: ParserSpec(
|
||||
engine_name=PARSER_ENGINE_MINERU,
|
||||
impl="lightrag.parser.external.mineru.parser:MinerUParser",
|
||||
suffixes=_MINERU_SUFFIXES,
|
||||
queue_group=PARSER_ENGINE_MINERU, # sized by max_parallel_parse_mineru
|
||||
endpoint_configured=_mineru_endpoint_configured,
|
||||
endpoint_requirement=_mineru_endpoint_requirement,
|
||||
),
|
||||
PARSER_ENGINE_DOCLING: ParserSpec(
|
||||
engine_name=PARSER_ENGINE_DOCLING,
|
||||
impl="lightrag.parser.external.docling.parser:DoclingParser",
|
||||
suffixes=_DOCLING_SUFFIXES,
|
||||
queue_group=PARSER_ENGINE_DOCLING, # sized by max_parallel_parse_docling
|
||||
endpoint_configured=_env_endpoint_configured("DOCLING_ENDPOINT"),
|
||||
endpoint_requirement=lambda: "DOCLING_ENDPOINT",
|
||||
),
|
||||
PARSER_ENGINE_REUSE: ParserSpec(
|
||||
engine_name=PARSER_ENGINE_REUSE,
|
||||
impl="lightrag.parser.noop:ReuseParser",
|
||||
suffixes=frozenset(),
|
||||
user_selectable=False,
|
||||
),
|
||||
PARSER_ENGINE_PASSTHROUGH: ParserSpec(
|
||||
engine_name=PARSER_ENGINE_PASSTHROUGH,
|
||||
impl="lightrag.parser.noop:PassthroughParser",
|
||||
suffixes=frozenset(),
|
||||
user_selectable=False,
|
||||
),
|
||||
}
|
||||
|
||||
# (engine_name, impl) -> instance. Keyed on impl so a re-registration with a
|
||||
# different implementation is not served a stale cached instance.
|
||||
_INSTANCE_CACHE: dict[tuple[str, str], "BaseParser"] = {}
|
||||
|
||||
|
||||
def register_parser(spec: ParserSpec) -> None:
|
||||
"""Register (or override) a parser engine spec."""
|
||||
_REGISTRY[spec.engine_name] = spec
|
||||
|
||||
|
||||
def parser_specs_snapshot() -> dict[str, ParserSpec]:
|
||||
"""Return a shallow snapshot of the registry.
|
||||
|
||||
The pipeline takes one snapshot at batch start and threads it through
|
||||
queue construction, routing and the parse workers so a concurrent
|
||||
``register_parser`` cannot change the engine set mid-batch.
|
||||
"""
|
||||
return dict(_REGISTRY)
|
||||
|
||||
|
||||
def _table(specs: dict[str, ParserSpec] | None) -> dict[str, ParserSpec]:
|
||||
return specs if specs is not None else _REGISTRY
|
||||
|
||||
|
||||
def get_parser(engine: str, *, specs: dict[str, ParserSpec] | None = None):
|
||||
"""Return a (cached) parser instance for ``engine`` or ``None``.
|
||||
|
||||
Imports the implementation lazily via ``importlib`` — only here does the
|
||||
heavy engine package (and e.g. ``httpx``) get pulled in.
|
||||
"""
|
||||
spec = _table(specs).get(engine)
|
||||
if spec is None:
|
||||
return None
|
||||
cache_key = (engine, spec.impl)
|
||||
inst = _INSTANCE_CACHE.get(cache_key)
|
||||
if inst is None:
|
||||
module_path, _, cls_name = spec.impl.partition(":")
|
||||
cls = getattr(importlib.import_module(module_path), cls_name)
|
||||
inst = cls()
|
||||
_INSTANCE_CACHE[cache_key] = inst
|
||||
return inst
|
||||
|
||||
|
||||
def supported_parser_engines(
|
||||
specs: dict[str, ParserSpec] | None = None,
|
||||
) -> frozenset[str]:
|
||||
"""User-selectable engine names (replaces SUPPORTED_PARSER_ENGINES)."""
|
||||
return frozenset(
|
||||
name for name, spec in _table(specs).items() if spec.user_selectable
|
||||
)
|
||||
|
||||
|
||||
def available_engine_suffixes(
|
||||
specs: dict[str, ParserSpec] | None = None,
|
||||
) -> frozenset[str]:
|
||||
"""Suffixes (lowercase, no dot) parseable by a *currently usable* engine.
|
||||
|
||||
Union over user-selectable engines whose ``endpoint_configured()`` gate
|
||||
passes. This is the single source for the API upload allowlist and the
|
||||
input-directory scan (``DocumentManager.supported_extensions``): in a
|
||||
default deployment (no external endpoints) it equals the local engines'
|
||||
suffixes (legacy ∪ native); configuring e.g. ``MINERU_LOCAL_ENDPOINT``
|
||||
admits mineru's image/office suffixes; a registered third-party engine's
|
||||
suffixes join automatically (subject to its own endpoint gate).
|
||||
"""
|
||||
out: set[str] = set()
|
||||
for spec in _table(specs).values():
|
||||
if spec.user_selectable and spec.endpoint_configured():
|
||||
out |= spec.suffixes
|
||||
return frozenset(out)
|
||||
|
||||
|
||||
def suffix_capabilities(
|
||||
engine: str, specs: dict[str, ParserSpec] | None = None
|
||||
) -> frozenset[str]:
|
||||
spec = _table(specs).get(engine)
|
||||
return spec.suffixes if spec is not None else frozenset()
|
||||
|
||||
|
||||
def engine_endpoint_configured(
|
||||
engine: str, specs: dict[str, ParserSpec] | None = None
|
||||
) -> bool:
|
||||
spec = _table(specs).get(engine)
|
||||
return spec.endpoint_configured() if spec is not None else True
|
||||
|
||||
|
||||
def engine_endpoint_requirement(
|
||||
engine: str, specs: dict[str, ParserSpec] | None = None
|
||||
) -> str | None:
|
||||
spec = _table(specs).get(engine)
|
||||
return spec.endpoint_requirement() if spec is not None else None
|
||||
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user