Coverage for astrocyte/documents/builders/md_builder.py: 100%

1r"""Markdown → DocumentTree builder. 

2 

3PageIndex parity (see 

4https://github.com/AstrocyteAI/PageIndex/blob/8cbcc08/pageindex/page_index_md.py 

5— permalink to the pinned fork the bench installs; byte-identical to 

6upstream VectifyAI/PageIndex for this file): 

7 

8- Header regex: ``^(#{1,6})\s+(.+)$`` — all six heading levels become 

9 nodes. 

10- Code-block aware: triple-backtick fences are tracked so headings 

11 inside code samples don't accidentally split sections. 

12- Each node owns the text from its heading line up to (but not 

13 including) the next heading at the same-or-shallower level. 

14- Children attach to the most-recent ancestor with strictly shallower 

15 depth — the standard markdown nesting rule. 

16 

17This module ONLY constructs the tree structure. Per-node summarization 

18is the summarizer's job (Phase 2). The tree returned here has 

19``summary=None`` on every node. 

20 

21Public API: 

22 build_markdown_tree(md_text, document_id) -> DocumentTree 

23""" 

24 

25from __future__ import annotations 

26 

27import re 

28 

29from astrocyte.documents.types import DocumentTree, TreeNode 

30 

31_HEADER_RE = re.compile(r"^(#{1,6})\s+(.+)$") 

32_CODE_FENCE_RE = re.compile(r"^```") 

33 

34 

35def _extract_node_headers(md_text: str) -> list[tuple[int, int, str]]: 

36 """Pass 1 of PageIndex's md→tree: locate every header line. 

37 

38 Returns a list of ``(line_num, depth, title)`` triples in document 

39 order. ``line_num`` is 1-indexed. Code blocks are excluded so 

40 headers inside ```fenced``` regions don't appear. 

41 """ 

42 headers: list[tuple[int, int, str]] = [] 

43 in_code = False 

44 for i, line in enumerate(md_text.splitlines(), start=1): 

45 stripped = line.strip() 

46 if _CODE_FENCE_RE.match(stripped): 

47 in_code = not in_code 

48 continue 

49 if in_code: 

50 continue 

51 if not stripped: 

52 continue 

53 m = _HEADER_RE.match(stripped) 

54 if m: 

55 depth = len(m.group(1)) 

56 title = m.group(2).strip() 

57 headers.append((i, depth, title)) 

58 return headers 

59 

60 

61def _build_node_text( 

62 headers: list[tuple[int, int, str]], 

63 md_lines: list[str], 

64 idx: int, 

65) -> tuple[str, int]: 

66 """Compute the text body for header at ``headers[idx]``. 

67 

68 Returns ``(text, end_line)``. The text spans from the header line 

69 (inclusive, 1-indexed) up to but not including the next header 

70 line of any depth (PageIndex behavior — each leaf "owns" its body 

71 until any next header). ``end_line`` is the exclusive end (next 

72 header's line_num, or total_lines+1 if last). 

73 """ 

74 start_line = headers[idx][0] # 1-indexed 

75 if idx + 1 < len(headers): 

76 end_line = headers[idx + 1][0] 

77 else: 

78 end_line = len(md_lines) + 1 

79 # md_lines is 0-indexed; convert 

80 body = "\n".join(md_lines[start_line - 1 : end_line - 1]).strip() 

81 return body, end_line 

82 

83 

84def build_markdown_tree(md_text: str, document_id: str) -> DocumentTree: 

85 """Build a DocumentTree from markdown text. 

86 

87 Implementation mirrors PageIndex's two-pass approach: 

88 

89 Pass 1: scan lines, locate every header (depth + title + line_num). 

90 Pass 2: stack-based tree construction — each header attaches under 

91 the most-recent ancestor with strictly shallower depth. 

92 

93 Each node's text body is the lines from its header through the next 

94 header of any depth, stripped. Roots are nodes whose parent is None 

95 (i.e., the document started with that depth, no shallower header 

96 preceded it). 

97 

98 Edge cases: 

99 - Markdown with NO headers → returns DocumentTree with one 

100 synthetic root holding the full text. (Most documents have at 

101 least one heading; the bench harness for LME/LoCoMo synthesizes 

102 ``## Session N`` headers, so this fallback rarely fires in 

103 practice but keeps the contract well-defined.) 

104 - Empty input → DocumentTree with empty ``roots`` list. 

105 """ 

106 if not md_text or not md_text.strip(): 

107 return DocumentTree(document_id=document_id, roots=[]) 

108 

109 md_lines = md_text.splitlines() 

110 headers = _extract_node_headers(md_text) 

111 

112 if not headers: 

113 # Headerless input — produce one synthetic root 

114 synthetic = TreeNode.new( 

115 parent_id=None, 

116 depth=1, 

117 title="(untitled document)", 

118 text=md_text.strip(), 

119 line_start=1, 

120 line_end=len(md_lines), 

121 ) 

122 return DocumentTree(document_id=document_id, roots=[synthetic]) 

123 

124 # Build all nodes first (no parent assignments yet) 

125 nodes: list[TreeNode] = [] 

126 for idx, (line_num, depth, title) in enumerate(headers): 

127 text, end_line = _build_node_text(headers, md_lines, idx) 

128 nodes.append( 

129 TreeNode.new( 

130 parent_id=None, # filled in by stack pass below 

131 depth=depth, 

132 title=title, 

133 text=text, 

134 line_start=line_num, 

135 line_end=end_line - 1, 

136 ), 

137 ) 

138 

139 # Stack-based parent assignment. Stack holds (depth, node) pairs of 

140 # candidate ancestors. For each new node, pop until top has strictly 

141 # shallower depth; that's the parent. If stack empties → new root. 

142 roots: list[TreeNode] = [] 

143 stack: list[TreeNode] = [] 

144 for node in nodes: 

145 while stack and stack[-1].depth >= node.depth: 

146 stack.pop() 

147 if stack: 

148 stack[-1].add_child(node) 

149 else: 

150 roots.append(node) 

151 stack.append(node) 

152 

153 return DocumentTree(document_id=document_id, roots=roots)