Coverage for astrocyte/pipeline/section_link

1"""PR2 D.7: per-section causal / supersedes / elaborates link extraction.

3Why this exists: the ``graph_expand`` strategy (PR2 commit B) reads

4``section_links`` to do 1-hop expansion from semantic / entity seeds.

5Without this module, ``section_links`` is empty and graph_expand

6returns no hits — which is why **LME multi-session sat at 11% (1/9)**

7across PR2-D's iterations. PR2-D.7 populates the table at retain.

9Three link types we extract:

11| link_type | Meaning | Weight semantics |

12|---------------|--------------------------------------------------|------------------|

13| **causal** | from_section caused (or enabled) to_section | Higher = more direct |

14| **supersedes** | to_section corrects / updates from_section | Time-ordered (newer wins) — LME knowledge-update primitive |

15| **elaborates** | to_section extends / adds detail to from_section | Topical continuity |

17Weights default to 0.7 unless the LLM emits one. We don't tune weights

18at PR2-D.7; the graph_expand strategy sums per-edge weights, so a flat

190.7 makes RRF behaviour predictable.

21## Why per-section, not per-pair

23Pair-wise extraction (look at every (i, j) section pair) is O(N²) at

24retain and dominates LME cost. Per-section is O(N): for each section,

25the LLM looks at the section's content + a window of prior section

26summaries and emits any links it sees. Resulting graph is sparse but

27covers the dominant retrieval cases.

29## Cost shape

31One LLM call per section. For LoCoMo ~30 sections × 10 convs = 300

32calls (~$0.05 first-time, $0 on cache). For LME ~50 sections × 50

33samples = 2500 calls (~$0.40 first-time, $0 on cache).

35Failures degrade silently — link extraction is enrichment, not

36correctness. The picker still works without it; graph_expand just

37returns empty.

39See:

40- docs/_design/recall.md §5 (retain pipeline)

41- docs/_design/adr/adr-006-three-layer-recall-stack.md (graph layer)

42- ``astrocyte_pi_section_links`` schema in 015_tier2_recall.sql

43"""

45from __future__ import annotations

47import logging

48from typing import TYPE_CHECKING

50if TYPE_CHECKING:

51 from astrocyte.provider import LLMProvider

52 from astrocyte.types import PageIndexSection

54from astrocyte.types import Message, PageIndexSectionLink

56logger = logging.getLogger("astrocyte.pipeline.section_link_extraction")

59_VALID_LINK_TYPES = frozenset({"causal", "supersedes", "elaborates"})

61#: Cap how many prior-section summaries we show the LLM per call. Bounds

62#: prompt size so retain wallclock stays bounded; LME's 50-section docs

63#: would otherwise blow past gpt-4o-mini's context window.

64_PRIOR_WINDOW = 8

66#: Cap on emitted links per section. Pathological extractions (every

67#: prior section linked) inflate the index without lifting recall —

68#: graph_expand only needs the strongest few.

69_MAX_LINKS_PER_SECTION = 5

72_EXTRACT_PROMPT = """You are analysing a long conversation transcript. The transcript has been chunked into sessions. For the CURRENT session, identify links to PRIOR sessions in this conversation.

74Three link types:

76- "causal": the current session was caused or enabled by the prior session (e.g., prior session sets up an event that the current session reports the outcome of).

77- "supersedes": the current session corrects or updates information from the prior session (e.g., user previously said X, now says actually Y).

78- "elaborates": the current session adds detail or follow-up to a topic the prior session introduced.

80CURRENT session (line {current_line}):

81{current_text}

83PRIOR sessions (you can link to any of these):

84{prior_summary_block}

86Return ONLY a JSON object: {{"links": [{{"to_line": <int>, "type": "causal|supersedes|elaborates", "weight": <0.0-1.0>, "reason": "<one short clause>"}}, ...]}}.

88Emit at most {max_links} links. Return {{"links": []}} when no clear link exists.

90Constraint: ``to_line`` MUST be one of the prior-session line numbers shown above.

92Output (JSON only):

93"""

96async def extract_links_for_section(

97 provider: "LLMProvider",

98 document_id: str,

99 current: "PageIndexSection",

100 current_text: str,

101 prior_sections: list["PageIndexSection"],

102 *,

103 model: str | None = None,

104) -> list[PageIndexSectionLink]:

105 """One LLM call → up to ``_MAX_LINKS_PER_SECTION`` ``PageIndexSectionLink`` rows.

106

107 ``prior_sections`` is a window of the most-recent prior sections

108 (caller passes them; usually the previous N session-grain

109 sections in the same document). The LLM sees their summaries and

110 line_nums; it can emit links to any subset.

111

112 Returns an empty list on parse failure / no links found / no

113 valid prior window.

114 """

115 if not current_text.strip() or not prior_sections:

116 return []

117

118 # Build the prior-summary block for the prompt.

119 prior_lines: list[str] = []

120 valid_to_lines: set[int] = set()

121 for p in prior_sections[-_PRIOR_WINDOW:]:

122 title = (p.title or "").strip()

123 summary = (p.summary or "").strip()

124 if not title and not summary:

125 continue

126 prior_lines.append(f" - line {p.line_num}: {title}\n summary: {summary[:240]}")

127 valid_to_lines.add(p.line_num)

128 if not prior_lines:

129 return []

130

131 msg = _EXTRACT_PROMPT.format(

132 current_line=current.line_num,

133 current_text=current_text[:4000], # cap; matches entity-extract sizing

134 prior_summary_block="\n".join(prior_lines),

135 max_links=_MAX_LINKS_PER_SECTION,

136 )

137 completion = await provider.complete(

138 messages=[Message(role="user", content=msg)],

139 model=model,

140 max_tokens=400,

141 temperature=0.0,

142 response_format={"type": "json_object"},

143 )

144

145 # Tolerant parse (fence-stripping + outermost-brace slicing) — providers

146 # whose json_object mode is emulated rather than native (e.g. claude_cli)

147 # occasionally wrap the object in ```json fences or prose. Bare

148 # json.loads was safe only under OpenAI's native constrained decoding.

149 from astrocyte.pipeline._json_tolerant import tolerant_json_loads # noqa: PLC0415

150

151 parsed = tolerant_json_loads(completion.text)

152 if not isinstance(parsed, dict):

153 logger.warning(

154 "section_link_extraction: JSON parse failed for doc=%s line=%d",

155 document_id,

156 current.line_num,

157 )

158 return []

159 raw = parsed.get("links") or []

160

161 out: list[PageIndexSectionLink] = []

162 seen_to: set[tuple[int, str]] = set() # (to_line, type) dedupe

163 for entry in raw:

164 if not isinstance(entry, dict):

165 continue

166 to_line = entry.get("to_line")

167 link_type = entry.get("type")

168 if not isinstance(to_line, int) or to_line not in valid_to_lines:

169 continue

170 if link_type not in _VALID_LINK_TYPES:

171 continue

172 # Self-loops are nonsensical; skip.

173 if to_line == current.line_num:

174 continue

175 key = (to_line, link_type)

176 if key in seen_to:

177 continue

178 seen_to.add(key)

179 weight_raw = entry.get("weight", 0.7)

180 try:

181 weight = float(weight_raw)

182 except (TypeError, ValueError):

183 weight = 0.7

184 # Clamp to [0, 1] to match the schema's expected range.

185 weight = max(0.0, min(1.0, weight))

186 out.append(

187 PageIndexSectionLink(

188 from_doc=document_id,

189 from_line=current.line_num,

190 to_doc=document_id,

191 to_line=to_line,

192 link_type=link_type,

193 weight=weight,

194 )

195 )

196 if len(out) >= _MAX_LINKS_PER_SECTION:

197 break

198 return out

Coverage for astrocyte/pipeline/section_link_extraction.py: 0%

58 statements