Coverage for astrocyte/pipeline/mental_model_compile.py: 92%

1"""M11.2 — mental-model compile pass. 

2 

3Synthesizes stable user preferences / persona / habits / background 

4from a document's sections and persists them as ``MentalModel`` rows 

5in the bank's :class:`~astrocyte.provider.MentalModelStore`. Direct 

6port of Hindsight's mental-model tier (saved-reflect summaries 

7recalled BEFORE observations and raw memories). 

8 

9Different from :mod:`section_compile` (wiki pages): 

10 

11- **Wiki pages** are topic-clustered observations 

12 ("User visited 3 doctors", "User worked on 5 model kits"). 

13 Generated per cluster via DBSCAN. Carry section-grain provenance. 

14- **Mental models** are DURABLE user-profile statements 

15 ("User prefers Sony cameras", "User practices Spanish 3×/week"). 

16 Generated per document via one LLM call across all sections. 

17 Span the whole document — no per-cluster scope. 

18 

19Why both layers: a question like "recommend a hotel in Miami" wants 

20the user-profile fact ("user prefers ocean views") regardless of 

21which topic-cluster that was discussed in. Wiki pages cluster by 

22topic; mental models cluster by user-profile dimension. 

23 

24Generic across LME / LoCoMo / future benches — the prompt asks for 

25PROFILE facts, not bench-specific shapes. 

26 

27See: 

28- ``docs/_design/recall.md`` §14 (M11 plan) 

29- ``hindsight-api-slim/hindsight_api/engine/reflect/observations.py`` 

30 for Hindsight's analogous structure (their ``Observation`` rows on 

31 ``MentalModel``). 

32""" 

33 

34from __future__ import annotations 

35 

36import json 

37import logging 

38import re 

39from datetime import datetime, timezone 

40from typing import TYPE_CHECKING 

41 

42from astrocyte.pipeline._json_tolerant import tolerant_json_loads_or_raise 

43from astrocyte.types import MentalModel, Message 

44 

45if TYPE_CHECKING: 

46 from astrocyte.provider import LLMProvider, MentalModelStore, PageIndexStore 

47 

48_logger = logging.getLogger("astrocyte.pipeline.mental_model_compile") 

49 

50 

51_COMPILE_PROMPT = """\ 

52You are extracting STABLE user profile facts from a conversation \ 

53transcript. The user is the speaker labeled "user" in chat turns. 

54 

55Output a JSON object with one key, ``models``, containing an array \ 

56of mental-model objects. Each model has: 

57- "title": 3-7 word noun phrase naming the profile dimension \ 

58(e.g. "Photography gear preference", "Diet", "Dance interests") 

59- "content": one declarative sentence stating the user's stable \ 

60preference, habit, persona, or background. Use specific entities \ 

61when the user named them. The reader will use this as authoritative \ 

62context for "recommend X" / "what do I like" / "tell me about my X" \ 

63questions. 

64 

65What counts as a mental model (extract these): 

66- Stable preferences: "User prefers Sony cameras over Canon." 

67- Persona / background: "User is a parent of two children, lives in \ 

68Brooklyn, works in tech." 

69- Recurring habits: "User attends weekly hip-hop dance classes at \ 

70Street Beats." 

71- Stable opinions: "User finds yoga more relaxing than running." 

72- Hobbies / projects: "User is building a home maintenance app." 

73- Skills / expertise: "User has 10 years of photography experience." 

74- Values / goals: "User is saving for a Europe trip in 2024." 

75 

76What does NOT count (ignore these): 

77- One-off events ("Yesterday I went to the doctor") 

78- Casual mentions without preference signal ("I had coffee") 

79- Things the assistant said about the user (only user's own claims) 

80- Speculation or one-time questions 

81 

82Rules: 

83- Each model must reflect something the user EXPLICITLY stated or \ 

84strongly implied across multiple sections (not single-mention) 

85- Be specific about entities when known (brand names, places, etc.) 

86- Cap at 12 models total — prefer the most-discussed dimensions 

87- If sections are generic chitchat with no stable profile signals, \ 

88return ``{{"models": []}}`` 

89 

90Examples of good output: 

91{{"models": [ 

92 {{"title": "Photography gear preference", "content": "User shoots \ 

93with Sony A7 III and prefers Sony-compatible accessories like the 24-70mm \ 

94G Master lens."}}, 

95 {{"title": "Movie genre taste", "content": "User strongly prefers \ 

96stand-up comedy specials on Netflix, especially recent ones (Ali Wong, \ 

97John Mulaney)."}}, 

98 {{"title": "Dance practice", "content": "User attends weekly hip-hop \ 

99classes at Street Beats and enjoys contemporary as a secondary style."}}, 

100 {{"title": "Diet", "content": "User follows a mostly-vegetarian \ 

101diet with occasional fish; avoids dairy."}} 

102]}} 

103 

104OUTPUT MUST BE VALID JSON. No prose around it. 

105 

106Section summaries (chronological): 

107{sections} 

108""" 

109 

110 

111def _slugify(s: str) -> str: 

112 s = s.lower().strip() 

113 s = re.sub(r"[^a-z0-9\s-]", "", s) 

114 s = re.sub(r"\s+", "-", s) 

115 return s[:60] or "model" 

116 

117 

118def _format_sections_for_prompt(sections, max_sections: int = 60) -> str: 

119 """Render section summaries chronologically for the compile prompt. 

120 

121 Cap at ``max_sections`` to keep the prompt under ~6K tokens. We 

122 rank by line_num ascending (chronological order through the 

123 document) and take the first ``max_sections``; LME chat-history 

124 documents commonly have 30-50 sections so this rarely trims. 

125 """ 

126 chronological = sorted(sections, key=lambda s: s.line_num) 

127 rendered = [] 

128 for s in chronological[:max_sections]: 

129 summary = (s.summary or s.title or "").strip() 

130 if not summary: 

131 continue 

132 date = s.session_date.strftime("%Y-%m-%d") if getattr(s, "session_date", None) is not None else "no-date" 

133 rendered.append(f"[line={s.line_num} date={date}] {summary}") 

134 return "\n".join(rendered) 

135 

136 

137async def compile_mental_models_for_document( 

138 *, 

139 page_index_store: PageIndexStore, 

140 mental_model_store: MentalModelStore, 

141 bank_id: str, 

142 document_id: str, 

143 provider: LLMProvider, 

144 model: str | None = None, 

145) -> list[str]: 

146 """Extract mental models for one document and persist via 

147 :class:`MentalModelStore`. 

148 

149 Returns the list of newly-upserted ``model_id``s. Idempotent: when 

150 models with the same ``model_id`` already exist for this bank, they 

151 are bumped to a new revision (the store's standard upsert 

152 semantics) rather than duplicated. 

153 

154 Scoping: ``scope = f"document:{document_id}"`` — mirrors the wiki 

155 tier's pattern so the bench's per-question retrieval can filter 

156 cleanly to the right document without cross-contamination across 

157 sibling LME conversations in the same bank. 

158 """ 

159 sections = await page_index_store.load_skeleton(document_id) 

160 if not sections: 

161 return [] 

162 

163 prompt = _COMPILE_PROMPT.format( 

164 sections=_format_sections_for_prompt(sections), 

165 ) 

166 try: 

167 completion = await provider.complete( 

168 [Message(role="user", content=prompt)], 

169 model=model, 

170 max_tokens=900, 

171 temperature=0.0, 

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

173 ) 

174 except Exception as exc: # noqa: BLE001 

175 _logger.warning( 

176 "mental_model_compile: LLM call failed doc=%s: %s", 

177 document_id, 

178 exc, 

179 ) 

180 return [] 

181 try: 

182 data = tolerant_json_loads_or_raise(completion.text) 

183 except json.JSONDecodeError as exc: 

184 _logger.warning( 

185 "mental_model_compile: JSON parse failed doc=%s: %s text=%r", 

186 document_id, 

187 exc, 

188 completion.text[:200], 

189 ) 

190 return [] 

191 raw_models = data.get("models") or [] 

192 if not isinstance(raw_models, list): 

193 return [] 

194 

195 now = datetime.now(tz=timezone.utc) 

196 scope = f"document:{document_id}" 

197 upserted: list[str] = [] 

198 # M40 — Anchor the MM's evidence timestamp in conversation time, not 

199 # wall-clock-now. Each compiled MM summarizes the whole document, so 

200 # the earliest section date is the right "when did this evidence 

201 # start existing" anchor. Single-point (matches the synthetic single 

202 # source_id ``"{document_id}:doc"``); compute_trend will classify 

203 # NEW (recent doc relative to reference_date) or STALE (old doc). 

204 _section_dates = [s.session_date for s in sections if getattr(s, "session_date", None) is not None] 

205 _evidence_anchor: datetime | None = min(_section_dates) if _section_dates else None 

206 for entry in raw_models: 

207 if not isinstance(entry, dict): 

208 continue 

209 title = str(entry.get("title", "")).strip() 

210 content = str(entry.get("content", "")).strip() 

211 if not title or not content: 

212 continue 

213 model_id = f"mm:{document_id[:8]}:{_slugify(title)}" 

214 mm = MentalModel( 

215 model_id=model_id, 

216 bank_id=bank_id, 

217 title=title, 

218 content=content, 

219 scope=scope, 

220 source_ids=[f"{document_id}:doc"], 

221 revision=1, # upsert assigns the real revision number 

222 refreshed_at=now, 

223 source_timestamps=[_evidence_anchor] if _evidence_anchor is not None else None, 

224 ) 

225 try: 

226 await mental_model_store.upsert(mm, bank_id) 

227 except Exception as exc: # noqa: BLE001 

228 _logger.warning( 

229 "mental_model_compile.upsert failed model_id=%s: %s", 

230 model_id, 

231 exc, 

232 ) 

233 continue 

234 upserted.append(model_id) 

235 

236 _logger.info( 

237 "mental_model_compile: doc=%s upserted %d models", 

238 document_id, 

239 len(upserted), 

240 ) 

241 return upserted