wehub-resource-sync
392 lines
15 KiB
Python
392 lines
15 KiB
Python
"""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
|