Coverage for astrocyte/pipeline/section_compile.py: 42%
183 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-10-29 22:10 +0000
1"""M10.1 — section-aware consolidation engine.
3Hindsight-aligned: aggregation happens at retain time so query-time
4agents can find pre-aggregated observations via wiki recall, instead
5of trying to count distinct items across many section excerpts (which
6the LLM is unreliable at).
8Algorithm per document:
101. Load sections with ``summary_embedding`` populated.
112. DBSCAN-cluster by cosine similarity (eps=0.30, min_samples=2 by
12 default) — same approach as the M8 ``CompileEngine`` over raw
13 memory_units, just at section grain.
143. For each cluster, single LLM call synthesizes a short observation:
15 a 5-10 word title and a 1-3 sentence aggregated claim that cites
16 specific dates / counts / entities. This is the primitive that
17 answers "how many doctors did I visit?" without making the agent
18 count across sessions at query time.
194. Save each observation as a :class:`~astrocyte.types.WikiPage`
20 (kind="topic"), with provenance back to the source sections via
21 ``astrocyte_pi_wiki_provenance`` (migration 015) and an embedding
22 on ``astrocyte_wiki_pages.current_embedding`` (migration 018).
24Differs from M8's :class:`~astrocyte.pipeline.compile.CompileEngine`:
25M8 reads from ``VectorStore`` (memory_units); this module reads from
26``PageIndexStore`` (sections). Both produce ``WikiPage`` rows; both
27use the same ``_dbscan`` clustering. The split exists because the
28PageIndex POC bypasses ``VectorStore`` entirely.
30See:
31- ``docs/_design/recall.md`` §12 (PR2.6 close-out, M10 plan)
32- ``adr/adr-006-three-layer-recall-stack.md`` — Layer 1 wiki tier
33- ``hindsight/hindsight-api-slim/hindsight_api/engine/consolidation/consolidator.py``
34 — Hindsight's analogous flow over raw memories
35"""
37from __future__ import annotations
39import json
40import logging
41import re
42from datetime import datetime, timezone
43from typing import TYPE_CHECKING
45from astrocyte.pipeline._json_tolerant import tolerant_json_loads_or_raise
46from astrocyte.pipeline.compile import _dbscan # reuse pure-Python DBSCAN
47from astrocyte.types import Message, WikiPage
49if TYPE_CHECKING:
50 from astrocyte.provider import LLMProvider, PageIndexStore
51 from astrocyte.types import PageIndexSection
53_logger = logging.getLogger("astrocyte.pipeline.section_compile")
56# ── Synthesis prompt ────────────────────────────────────────────────
59_OBSERVATION_PROMPT = """\
60You are an observation synthesizer. You will be given several sections \
61from a conversation that all discuss the same topic, presented in \
62CHRONOLOGICAL ORDER (oldest first). Your job: write ONE concise \
63observation that aggregates the user's experience or stable facts \
64across these sections, reflecting the LATEST stable state.
66Pay attention to LATER sections that may supersede earlier ones. If \
67the user changes their mind, switches providers, or adds new items \
68over time, the observation should reflect the LATEST state. Capture \
69the chronological progression when it matters (e.g. "previously did \
70X, now does Y") rather than averaging across time.
72Output a JSON object with EXACTLY two fields:
73- "title": a 5-10 word summary that uses CATEGORY nouns the reader \
74might query (e.g. "Doctors visited", "Model kits worked on", \
75"Camping trips taken"). Title should mention the category (doctor / \
76kit / trip / restaurant / etc.) so the wiki search can match it.
77- "content": 1-3 sentences. MANDATORY: lead with an explicit COUNT \
78followed by a colon-separated enumeration. The reader will use your \
79content as the answer to questions like "how many X?" / "total Y?" \
80/ "list of Z?". DO NOT generalize or theme — enumerate the actual \
81distinct items. When listing items added over time, the order is \
82chronological (earliest first).
84REQUIRED CONTENT SHAPE:
85"User <verb>ed N distinct <category>s: <item1> (<context>), \
86<item2> (<context>), ..., <itemN> (<context>)."
88Examples of good observations:
89{{"title": "Doctors visited", "content": "User visited 3 distinct \
90doctors: a primary care physician (Dr. Patel, May), an ENT specialist \
91(Dr. Lee, June), and a dermatologist (referral, July)."}}
93{{"title": "Model kits worked on", "content": "User worked on 5 model \
94kits: a Revell F-15 Eagle, a B-29 bomber, a 1/72 Spitfire, a Tiger I \
95tank, and a Sherman tank."}}
97{{"title": "Camping trips taken", "content": "User went on 4 camping \
98trips totaling 8 days: Yosemite (3 days, April), Lake Tahoe (2 days, \
99June), Olympic NP (2 days, August), Grand Canyon (1 day, October)."}}
101{{"title": "Korean restaurants visited", "content": "User visited 3 \
102distinct Korean restaurants: Yeon Su (Chelsea), Cote (Flatiron), and \
103Hanjip (LA, during travel)."}}
105ANTI-PATTERNS to avoid:
106- "User is managing health" (no count, no enumeration) — WRONG
107- "User explored various trips" (vague, no specifics) — WRONG
108- "User had multiple doctor visits" (says "multiple" without counting) — WRONG
110When the cluster has only 1 distinct fact, still lead with "1": \
111"User has 1 ongoing project: <description>."
113OUTPUT MUST BE VALID JSON. No prose around it.
115Sections in this cluster:
116{sections}
117"""
120# ── Cluster helpers ─────────────────────────────────────────────────
123def _slugify(s: str) -> str:
124 """Convert a title to a URL-safe slug. e.g. 'Doctors visited' -> 'doctors-visited'."""
125 s = s.lower().strip()
126 s = re.sub(r"[^a-z0-9\s-]", "", s)
127 s = re.sub(r"\s+", "-", s)
128 return s[:60] or "obs"
131def _format_section_for_prompt(section: PageIndexSection) -> str:
132 date = section.session_date.strftime("%Y-%m-%d") if section.session_date is not None else "no-date"
133 summary = (section.summary or section.title or "").strip()
134 return f"[line={section.line_num}, date={date}]\n{summary}"
137async def _synthesize_observation(
138 *,
139 provider: LLMProvider,
140 model: str | None,
141 sections: list[PageIndexSection],
142) -> tuple[str, str] | None:
143 """LLM call — returns (title, content) or None on parse failure.
145 M12.6: sections are sorted by ``session_date`` (with ``line_num``
146 fallback) so the synth prompt sees them in chronological order.
147 The Karpathy "latest state wins" guidance in the prompt requires
148 this ordering to be meaningful.
149 """
150 if not sections:
151 return None
152 ordered = sorted(
153 sections,
154 key=lambda s: (
155 s.session_date or datetime.min.replace(tzinfo=timezone.utc),
156 s.line_num,
157 ),
158 )
159 rendered = "\n\n".join(_format_section_for_prompt(s) for s in ordered)
160 msg = _OBSERVATION_PROMPT.format(sections=rendered)
161 try:
162 completion = await provider.complete(
163 [Message(role="user", content=msg)],
164 model=model,
165 max_tokens=300,
166 temperature=0.0,
167 response_format={"type": "json_object"},
168 )
169 except Exception as exc: # noqa: BLE001
170 _logger.warning(
171 "section_compile.synthesize: LLM call failed (%s)",
172 exc,
173 )
174 return None
175 try:
176 data = tolerant_json_loads_or_raise(completion.text)
177 except json.JSONDecodeError as exc:
178 _logger.warning(
179 "section_compile.synthesize: JSON parse failed (%s) text=%r",
180 exc,
181 completion.text[:200],
182 )
183 return None
184 title = str(data.get("title", "")).strip()
185 content = str(data.get("content", "")).strip()
186 if not title or not content:
187 return None
188 return title, content
191# ── Revision pass (Karpathy incremental update, M12.6) ─────────────
194_REVISION_PROMPT = """\
195You are revising a knowledge-base entry to reflect the LATEST stable \
196state across its source sections.
198Current entry (titled "{title}", revision {revision}):
199{content}
201Source sections in CHRONOLOGICAL ORDER (oldest first):
202{sections}
204Examine the sections in time order. If a LATER section supersedes, \
205contradicts, or extends an earlier claim, the entry should reflect \
206the latest state. If the entry already does, no revision is needed.
208Common revision triggers:
209- Entry says "User visited 3 doctors" but a later section adds a 4th → bump count
210- Entry says "User prefers brand X" but later sections say switched to Y → "User now prefers Y (previously X)"
211- Entry uses past tense for an ongoing thing → use ongoing tense
212- Entry omits items that appear in later sections → add them
213- Entry includes superseded items as if current → mark as previous
215If revising: the new content follows the SAME shape as initial compile \
216(lead with explicit COUNT and colon-separated enumeration; cite dates / \
217entities when present).
219Respond as JSON only:
220{{
221 "verdict": "OK" or "REVISE",
222 "revised_title": "<5-10 words; only when REVISE>",
223 "revised_content": "<1-3 sentences; only when REVISE>"
224}}
225"""
228async def _revise_observation(
229 *,
230 provider: LLMProvider,
231 model: str | None,
232 page: WikiPage,
233 sections: list[PageIndexSection],
234) -> tuple[str, str] | None:
235 """LLM call — returns (new_title, new_content) if revised, ``None``
236 if the page already reflects the latest state.
238 Sections must be sorted chronologically by the caller. Resilient
239 to LLM / parse failures (returns ``None``).
240 """
241 if not sections or not page.content.strip():
242 return None
243 rendered = "\n\n".join(_format_section_for_prompt(s) for s in sections)
244 msg = _REVISION_PROMPT.format(
245 title=page.title,
246 revision=page.revision,
247 content=page.content,
248 sections=rendered,
249 )
250 try:
251 completion = await provider.complete(
252 [Message(role="user", content=msg)],
253 model=model,
254 max_tokens=400,
255 temperature=0.0,
256 response_format={"type": "json_object"},
257 )
258 except Exception as exc: # noqa: BLE001
259 _logger.warning(
260 "section_compile.revise: LLM call failed for page=%s (%s)",
261 page.page_id,
262 exc,
263 )
264 return None
265 try:
266 data = tolerant_json_loads_or_raise(completion.text)
267 except json.JSONDecodeError as exc:
268 _logger.warning(
269 "section_compile.revise: JSON parse failed for page=%s (%s) text=%r",
270 page.page_id,
271 exc,
272 completion.text[:200],
273 )
274 return None
275 verdict = str(data.get("verdict", "")).strip().upper()
276 if verdict != "REVISE":
277 return None
278 new_title = str(data.get("revised_title", "")).strip()
279 new_content = str(data.get("revised_content", "")).strip()
280 if not new_title or not new_content:
281 return None
282 if new_title == page.title and new_content == page.content:
283 return None
284 return new_title, new_content
287async def revise_wikis_for_document(
288 *,
289 store: PageIndexStore,
290 bank_id: str,
291 document_id: str,
292 provider: LLMProvider,
293 model: str | None = None,
294 embedding_model: str | None = None,
295) -> int:
296 """Karpathy-style revision pass — for each existing wiki of this
297 document, re-check it against its provenance sections in
298 chronological order and persist any revisions.
300 Returns: number of pages actually revised.
302 Idempotent: if no page needs revision, the call is a no-op (just
303 one LLM call per page). Subsequent calls on already-revised pages
304 should converge to no-op as the wiki stabilises.
305 """
306 try:
307 pages = await store.list_wiki_pages_for_doc(bank_id, document_id)
308 except Exception as exc: # noqa: BLE001
309 _logger.warning(
310 "revise_wikis: list_wiki_pages_for_doc failed for doc=%s (%s)",
311 document_id,
312 exc,
313 )
314 return 0
315 if not pages:
316 return 0
318 try:
319 all_sections = await store.load_sections_with_embeddings(bank_id, document_id)
320 except Exception as exc: # noqa: BLE001
321 _logger.warning(
322 "revise_wikis: load_sections failed for doc=%s (%s)",
323 document_id,
324 exc,
325 )
326 return 0
327 section_by_line = {s.line_num: s for s in all_sections}
329 revised_count = 0
330 for page in pages:
331 # Resolve provenance source_ids back to sections.
332 provenance_lines: list[int] = []
333 for src in page.source_ids or []:
334 # source_ids are formatted "docId:lineNum"
335 if ":" not in src:
336 continue
337 _, _, line_str = src.rpartition(":")
338 try:
339 provenance_lines.append(int(line_str))
340 except ValueError:
341 continue
342 prov_sections = [section_by_line[ln] for ln in provenance_lines if ln in section_by_line]
343 if not prov_sections:
344 continue
345 # Sort chronologically — session_date first, line_num fallback.
346 prov_sections.sort(
347 key=lambda s: (
348 s.session_date or datetime.min.replace(tzinfo=timezone.utc),
349 s.line_num,
350 ),
351 )
352 result = await _revise_observation(
353 provider=provider,
354 model=model,
355 page=page,
356 sections=prov_sections,
357 )
358 if result is None:
359 continue
360 new_title, new_content = result
362 # Save as a new revision. ``save_wiki_page`` upserts the page row
363 # and inserts a new revision row keyed by (page_uuid, revision_number).
364 revised_page = WikiPage(
365 page_id=page.page_id,
366 bank_id=page.bank_id,
367 kind=page.kind,
368 title=new_title,
369 content=new_content,
370 scope=page.scope,
371 source_ids=list(page.source_ids or []),
372 cross_links=list(page.cross_links or []),
373 revision=page.revision + 1,
374 revised_at=datetime.now(tz=timezone.utc),
375 tags=page.tags,
376 metadata=page.metadata,
377 )
378 try:
379 new_embedding = (
380 await provider.embed(
381 [f"{new_title}\n\n{new_content}"],
382 model=embedding_model,
383 )
384 )[0]
385 except Exception as exc: # noqa: BLE001
386 _logger.warning(
387 "revise_wikis: embed failed for page=%s (%s)",
388 page.page_id,
389 exc,
390 )
391 new_embedding = None
393 provenance_pairs = [(s.document_id, s.line_num) for s in prov_sections]
394 try:
395 await store.save_wiki_page(
396 page=revised_page,
397 embedding=new_embedding,
398 provenance=provenance_pairs,
399 )
400 revised_count += 1
401 except Exception as exc: # noqa: BLE001
402 _logger.warning(
403 "revise_wikis.save failed for page=%s (%s)",
404 page.page_id,
405 exc,
406 )
408 if revised_count:
409 _logger.info(
410 "revise_wikis: doc=%s revised %d/%d pages",
411 document_id,
412 revised_count,
413 len(pages),
414 )
415 return revised_count
418# ── Public entry point ─────────────────────────────────────────────
421async def compile_sections_for_document(
422 *,
423 store: PageIndexStore,
424 bank_id: str,
425 document_id: str,
426 provider: LLMProvider,
427 model: str | None = None,
428 embedding_model: str | None = None,
429 eps: float = 0.55,
430 min_samples: int = 2,
431 max_pages_per_doc: int = 10,
432) -> list[str]:
433 """Cluster sections in this document, synthesize an observation per
434 cluster, persist as :class:`WikiPage` rows.
436 Returns the list of newly-created ``page_id``s. When the document
437 already has wiki pages (idempotency check), returns ``[]``.
439 ``eps`` / ``min_samples`` tune DBSCAN: defaults match LME-shaped
440 chat conversations (sessions are heterogeneous; eps=0.30 forces
441 fairly close cosine similarity). Lower ``min_samples`` than M8's
442 default (2 vs 3) because section-grain clusters tend to be smaller
443 than memory-unit clusters.
445 ``max_pages_per_doc`` caps cost: at 5-7 LLM calls per doc × 50 LME
446 docs we'd burn ~$2-3 per gate. The cap keeps it bounded if a
447 pathological doc clusters into many small groups.
448 """
449 # Idempotency: skip if pages already exist for this doc.
450 try:
451 existing = await store.count_wiki_pages_for_doc(bank_id, document_id)
452 except Exception as exc: # noqa: BLE001
453 _logger.warning(
454 "section_compile.count failed for doc=%s (%s)",
455 document_id,
456 exc,
457 )
458 existing = 0
459 if existing > 0:
460 _logger.debug(
461 "section_compile: doc=%s already has %d pages — skip initial compile, run revision pass only",
462 document_id,
463 existing,
464 )
465 # M12.6: even when initial compile is skipped, run the Karpathy
466 # revision pass against the existing pages. Lets a code change
467 # (e.g. new prompt) take effect against cached banks without a
468 # full re-extraction.
469 try:
470 await revise_wikis_for_document(
471 store=store,
472 bank_id=bank_id,
473 document_id=document_id,
474 provider=provider,
475 model=model,
476 embedding_model=embedding_model,
477 )
478 except Exception as exc: # noqa: BLE001
479 _logger.warning(
480 "revise_wikis: cache-hit revision pass failed for doc=%s (%s)",
481 document_id,
482 exc,
483 )
484 return []
486 # Load sections with embeddings populated.
487 sections = await store.load_sections_with_embeddings(bank_id, document_id)
488 if not sections:
489 return []
490 sections_with_emb = [s for s in sections if s.summary_embedding]
491 if len(sections_with_emb) < min_samples:
492 # Too few embedded sections to cluster meaningfully.
493 return []
495 # DBSCAN over (line_num_str, embedding) pairs.
496 items = [(str(s.line_num), s.summary_embedding) for s in sections_with_emb]
497 labels = _dbscan(items, eps=eps, min_samples=min_samples)
498 by_cluster: dict[int, list[PageIndexSection]] = {}
499 section_by_line = {s.line_num: s for s in sections_with_emb}
500 for line_str, cluster_id in labels.items():
501 if cluster_id == -1:
502 continue
503 by_cluster.setdefault(cluster_id, []).append(
504 section_by_line[int(line_str)],
505 )
507 if not by_cluster:
508 _logger.debug(
509 "section_compile: doc=%s yielded no clusters (eps=%s min=%s)",
510 document_id,
511 eps,
512 min_samples,
513 )
514 return []
516 # Cap cluster count to control cost. Prefer larger clusters (more
517 # signal per LLM call).
518 clusters = sorted(by_cluster.values(), key=len, reverse=True)[:max_pages_per_doc]
520 created_page_ids: list[str] = []
521 now = datetime.now(tz=timezone.utc)
522 pending_pages: list[WikiPage] = []
523 pending_provenance: list[tuple[str, list[tuple[str, int]]]] = []
524 pending_summaries: list[str] = [] # used for embedding pass
526 for cluster in clusters:
527 synth = await _synthesize_observation(
528 provider=provider,
529 model=model,
530 sections=cluster,
531 )
532 if synth is None:
533 continue
534 title, content = synth
535 page_id = f"obs:{document_id[:8]}:{_slugify(title)}"
536 source_ids = [f"{s.document_id}:{s.line_num}" for s in cluster]
537 page = WikiPage(
538 page_id=page_id,
539 bank_id=bank_id,
540 kind="topic",
541 title=title,
542 content=content,
543 scope=f"document:{document_id}",
544 source_ids=source_ids,
545 cross_links=[],
546 revision=1,
547 revised_at=now,
548 tags=None,
549 metadata=None,
550 )
551 pending_pages.append(page)
552 provenance_pairs = [(s.document_id, s.line_num) for s in cluster]
553 pending_provenance.append((page_id, provenance_pairs))
554 # Embed the title + content together so the wiki search
555 # strategy matches both the topic label and the aggregated
556 # facts. Mirrors how Hindsight indexes observations.
557 pending_summaries.append(f"{title}\n\n{content}")
559 if not pending_pages:
560 return []
562 # Single batched embedding call for all observations in this doc.
563 try:
564 embeddings = await provider.embed(
565 pending_summaries,
566 model=embedding_model,
567 )
568 except Exception as exc: # noqa: BLE001
569 _logger.warning(
570 "section_compile: embedding batch failed for doc=%s (%s)",
571 document_id,
572 exc,
573 )
574 embeddings = [None] * len(pending_pages)
576 for page, embedding, (page_id, provenance_pairs) in zip(
577 pending_pages,
578 embeddings,
579 pending_provenance,
580 ):
581 try:
582 await store.save_wiki_page(
583 page=page,
584 embedding=embedding,
585 provenance=provenance_pairs,
586 )
587 except Exception as exc: # noqa: BLE001
588 _logger.warning(
589 "section_compile.save_wiki_page failed page_id=%s (%s)",
590 page_id,
591 exc,
592 )
593 continue
594 created_page_ids.append(page_id)
596 _logger.info(
597 "section_compile: doc=%s created %d pages from %d clusters",
598 document_id,
599 len(created_page_ids),
600 len(clusters),
601 )
603 # M12.6: Karpathy-style revision pass over the freshly compiled
604 # pages. Catches "latest state" patterns that single-shot synthesis
605 # missed (e.g. counts that grow over time, providers that change).
606 # Idempotent — pages already at latest state are unchanged.
607 try:
608 await revise_wikis_for_document(
609 store=store,
610 bank_id=bank_id,
611 document_id=document_id,
612 provider=provider,
613 model=model,
614 embedding_model=embedding_model,
615 )
616 except Exception as exc: # noqa: BLE001
617 _logger.warning(
618 "revise_wikis: post-initial revision pass failed for doc=%s (%s)",
619 document_id,
620 exc,
621 )
623 return created_page_ids