chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -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"]
|
||||
Reference in New Issue
Block a user