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