Coverage for astrocyte/documents/builders/summarizer.py: 98%

1"""Adaptive per-node summarization (PageIndex parity). 

2 

3Walks a ``DocumentTree`` and populates each node's ``summary`` field 

4according to PageIndex's rule: 

5 

6 - If node text < threshold tokens: ``summary.text = node.text`` (no 

7 LLM call; ``kind="raw"``). 

8 - If node text ≥ threshold tokens: call LLM with PageIndex's prompt; 

9 ``kind="llm"``. 

10 - For internal nodes (have children): ``kind="prefix"`` — the summary 

11 is meant to be read as "what this whole subtree is about." 

12 

13Default threshold: 200 tokens (PageIndex default; configurable). 

14 

15Public API: 

16 AdaptiveSummarizer(llm_call, threshold_tokens=200, model="gpt-4o-mini") 

17 .summarize_tree(tree) # in-place: sets node.summary on every node 

18 .summarize_node(node) # single-node helper 

19 

20``llm_call`` is a coroutine ``async def fn(prompt: str) -> str``. Pass a 

21real LLM provider in production; pass a fake in tests. Keeps the 

22summarizer testable without an OpenAI key. 

23""" 

24 

25from __future__ import annotations 

26 

27import asyncio 

28import logging 

29from typing import Awaitable, Callable 

30 

31from astrocyte.documents.types import DocumentTree, NodeSummary, TreeNode 

32from astrocyte.policy.homeostasis import count_tokens 

33 

34logger = logging.getLogger(__name__) 

35 

36 

37# PageIndex parity — the exact prompt from ``generate_node_summary`` in 

38# https://github.com/AstrocyteAI/PageIndex/blob/8cbcc08/pageindex/utils.py 

39# (permalink to the pinned fork the bench installs; byte-identical to 

40# upstream VectifyAI/PageIndex for this file). 

41PAGEINDEX_SUMMARY_PROMPT = """You are given a part of a document, your task is to generate a description of the partial document about what are main points covered in the partial document. 

42 

43 Partial Document Text: {node_text} 

44 

45 Directly return the description, do not include any other text. 

46 """ 

47 

48DEFAULT_THRESHOLD_TOKENS = 200 

49 

50 

51LlmCall = Callable[[str], Awaitable[str]] 

52 

53 

54class AdaptiveSummarizer: 

55 """Per-node summarizer with PageIndex-style adaptive LLM gating. 

56 

57 Cost shape: most LME / LoCoMo session-blocks are < 200 tokens and 

58 skip the LLM entirely. LLM is only called for genuinely long nodes 

59 where the raw text would be too long to use directly downstream. 

60 """ 

61 

62 def __init__( 

63 self, 

64 llm_call: LlmCall, 

65 *, 

66 threshold_tokens: int = DEFAULT_THRESHOLD_TOKENS, 

67 model: str = "gpt-4o-mini", 

68 max_concurrent_llm: int = 8, 

69 ) -> None: 

70 if threshold_tokens < 1: 

71 raise ValueError(f"threshold_tokens must be >= 1, got {threshold_tokens}") 

72 self._llm_call = llm_call 

73 self._threshold = threshold_tokens 

74 self.model = model 

75 self._sem = asyncio.Semaphore(max_concurrent_llm) 

76 

77 # ── single-node helper ──────────────────────────────────────────── 

78 

79 async def summarize_node(self, node: TreeNode) -> NodeSummary: 

80 """Produce a NodeSummary for a single node. 

81 

82 Returns the summary (and sets ``node.summary``). Internal-node 

83 determination is left to the caller because it depends on tree 

84 traversal context — use ``summarize_tree`` for the typical case. 

85 """ 

86 text = node.text or "" 

87 tokens = count_tokens(text) 

88 if tokens < self._threshold: 

89 summary = NodeSummary(text=text, kind="raw", token_count=tokens) 

90 node.summary = summary 

91 return summary 

92 

93 async with self._sem: 

94 try: 

95 prompt = PAGEINDEX_SUMMARY_PROMPT.format(node_text=text) 

96 description = await self._llm_call(prompt) 

97 except Exception as exc: # noqa: BLE001 

98 logger.warning( 

99 "summarize_node: LLM call failed for node=%s title=%r tokens=%d: %s", 

100 node.id, 

101 node.title, 

102 tokens, 

103 exc, 

104 ) 

105 # Failure mode: degrade to raw text. Downstream still has SOMETHING. 

106 summary = NodeSummary(text=text, kind="raw", token_count=tokens) 

107 node.summary = summary 

108 return summary 

109 

110 summary_text = (description or "").strip() or text 

111 summary = NodeSummary( 

112 text=summary_text, 

113 kind="llm", 

114 token_count=count_tokens(summary_text), 

115 ) 

116 node.summary = summary 

117 return summary 

118 

119 # ── tree-level walk ─────────────────────────────────────────────── 

120 

121 async def summarize_tree(self, tree: DocumentTree) -> None: 

122 """Summarize every node in the tree. Mutates in place. 

123 

124 Runs node summarizations concurrently (bounded by 

125 ``max_concurrent_llm``). After completion, each node's 

126 ``summary`` is populated; internal nodes get ``kind="prefix"`` 

127 re-labeled so callers can distinguish "this is summary of a 

128 subtree" from "this is summary of a leaf." 

129 """ 

130 nodes = tree.all_nodes() 

131 if not nodes: 

132 return 

133 

134 # Step 1: summarize every node (concurrent, semaphore-bounded inside) 

135 await asyncio.gather(*(self.summarize_node(n) for n in nodes)) 

136 

137 # Step 2: re-label internal nodes' summary kind. 

138 # PageIndex distinguishes prefix_summary (internal) from summary (leaf). 

139 for n in nodes: 

140 if n.summary is not None and not n.is_leaf(): 

141 # Promote raw→prefix when an internal node had small text 

142 # (so its summary is its own heading-and-intro lines) AND 

143 # promote llm→prefix when it had a generated description. 

144 # Both indicate "this represents the subtree." 

145 n.summary = NodeSummary( 

146 text=n.summary.text, 

147 kind="prefix", 

148 token_count=n.summary.token_count, 

149 ) 

150 

151 # ── introspection ───────────────────────────────────────────────── 

152 

153 @property 

154 def threshold_tokens(self) -> int: 

155 return self._threshold