mirror of
https://github.com/VectifyAI/PageIndex.git
synced 2026-07-15 21:11:05 +02:00
b2756c73d2
Engineering-quality cleanups from the SDK review (no behavior change):
- Return-type discoverability: add pageindex/types.py with TypedDicts
(DocumentInfo, DocumentDetail, PageContent) and annotate Collection /
Backend methods with them; add docstrings to every public Collection
method (including the get_page_content `pages` spec). Exported from the
package. Zero runtime cost — these are plain dicts.
- Backend protocol as a real contract:
* query_stream is an async generator, so the protocol now declares it
as `def ... -> AsyncIterator[QueryEvent]` (not `async def`, which
typed it as a coroutine and never matched the implementations).
* custom-parser support is expressed as a runtime_checkable
SupportsParserRegistration capability protocol; the client uses
isinstance(...) instead of hasattr(...) duck-typing.
- Parser layering: move count_tokens into a leaf module pageindex/tokens.py
so parser/* imports it from there instead of reaching back into
pageindex.index (a reverse dependency). index.utils re-exports it for
backward compatibility.
Adds tests/test_architecture.py enforcing: parser never imports index,
count_tokens is a single shared leaf, the capability protocol works,
both backends satisfy Backend, and the TypedDicts are exported.
Claude-Session: https://claude.ai/code/session_01Kx5DgKbhK1N8autqXH8SmS
88 lines
3.2 KiB
Python
88 lines
3.2 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")
|
|
|
|
with open(path, "r", encoding="utf-8") 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)
|
|
|
|
def _extract_headers(self, lines: list[str]) -> list[dict]:
|
|
header_pattern = r"^(#{1,6})\s+(.+)$"
|
|
code_block_pattern = r"^```"
|
|
headers = []
|
|
in_code_block = False
|
|
|
|
for line_num, line in enumerate(lines, 1):
|
|
stripped = line.strip()
|
|
if re.match(code_block_pattern, stripped):
|
|
in_code_block = not in_code_block
|
|
continue
|
|
if not in_code_block 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
|