PageIndex/pageindex/parser/markdown.py
Ray 9980fc9414 feat: restore status and page/line counts in local get_document
Parsers self-report counts via ParsedDocument.metadata (pymupdf
page_count; markdown len(lines)); the backend merges them generically
into the stored document row at index time. New page_count/line_count
columns are added in place for pre-existing DBs. Local status is always
"completed" (indexing is synchronous), matching the cloud backend's
field. Align the agentic demo's get_document tool docstring (and its
col -> collection naming) accordingly.
2026-07-19 15:25:25 +08:00

105 lines
4.3 KiB
Python

import re
from pathlib import Path
from .protocol import ContentNode, ParsedDocument
from ..tokens import count_tokens
class MarkdownParser:
def supported_extensions(self) -> list[str]:
return [".md", ".markdown"]
def parse(self, file_path: str, **kwargs) -> ParsedDocument:
path = Path(file_path)
model = kwargs.get("model")
# utf-8-sig strips a leading BOM if present (common from Windows
# editors/exporters) and is otherwise identical to plain utf-8. Without
# it, a BOM-prefixed first line fails the header regex below (the BOM
# isn't whitespace, so .strip() doesn't remove it), misclassifying the
# document's first heading as unrecognized preamble text.
with open(path, "r", encoding="utf-8-sig") as f:
content = f.read()
lines = content.split("\n")
headers = self._extract_headers(lines)
nodes = self._build_nodes(headers, lines, model, doc_title=path.stem)
return ParsedDocument(doc_name=path.stem, nodes=nodes,
metadata={"line_count": len(lines)})
def _extract_headers(self, lines: list[str]) -> list[dict]:
header_pattern = r"^(#{1,6})\s+(.+)$"
# CommonMark allows both backtick and tilde fences, and a fence is
# closed only by one of the SAME character. Track which char opened the
# block so a ~~~ line inside a ```-fenced block (or vice versa) is
# treated as content, not a close — otherwise the block appears to end
# early and '#'-prefixed lines inside it get misparsed as headings.
fence_pattern = r"^(`{3,}|~{3,})"
headers = []
open_fence = None # the fence char ('`' or '~') of the open block, or None
for line_num, line in enumerate(lines, 1):
stripped = line.strip()
fence = re.match(fence_pattern, stripped)
if fence:
marker = fence.group(1)[0]
if open_fence is None:
open_fence = marker # open a block
elif open_fence == marker:
open_fence = None # matching char closes it
# a non-matching fence char while a block is open is content
continue
if open_fence is None and stripped:
match = re.match(header_pattern, stripped)
if match:
headers.append({
"title": match.group(2).strip(),
"level": len(match.group(1)),
"line_num": line_num,
})
return headers
def _build_nodes(self, headers: list[dict], lines: list[str], model: str | None,
doc_title: str = "Document") -> list[ContentNode]:
nodes = []
# A file with no headings at all still has content — index it as a
# single node instead of producing zero nodes (which would push an
# empty page list into the LLM pipeline).
if not headers:
text = "\n".join(lines).strip()
if text:
nodes.append(ContentNode(
content=text,
tokens=count_tokens(text, model=model),
title=doc_title,
index=1,
level=1,
))
return nodes
# Content before the first heading (abstract, preamble) would
# otherwise be silently dropped and become unretrievable.
preamble = "\n".join(lines[: headers[0]["line_num"] - 1]).strip()
if preamble:
nodes.append(ContentNode(
content=preamble,
tokens=count_tokens(preamble, model=model),
title=doc_title,
index=1,
level=headers[0]["level"],
))
for i, header in enumerate(headers):
start = header["line_num"] - 1
end = headers[i + 1]["line_num"] - 1 if i + 1 < len(headers) else len(lines)
text = "\n".join(lines[start:end]).strip()
tokens = count_tokens(text, model=model)
nodes.append(ContentNode(
content=text,
tokens=tokens,
title=header["title"],
index=header["line_num"],
level=header["level"],
))
return nodes