Coverage for astrocyte/provider.py: 69%

1"""Astrocyte provider protocols — the SPIs that backends implement. 

2 

3These are abstract contracts (Python Protocols), not concrete implementations. 

4Providers must implement every method in the relevant Protocol to be recognized 

5by Astrocyte's runtime type check (``isinstance(obj, VectorStore)`` etc.). 

6 

7Built-in implementations: 

8- In-memory test doubles: ``astrocyte.testing.in_memory`` 

9- PostgreSQL + pgvector: ``astrocyte_postgres.store.PostgresStore`` 

10- Optional wheels (repo ``adapters-storage-py/``): ``astrocyte-postgres``, ``astrocyte-qdrant``, ``astrocyte-neo4j``, ``astrocyte-elasticsearch`` 

11 

12Each protocol has a SPI_VERSION ClassVar for compatibility checking. 

13All methods are async except capabilities() and transport methods. 

14 

15SPI versioning: Astrocyte checks SPI_VERSION at registration time. 

16- Version 1: base protocol (all methods required) 

17- Version 2+: may add optional methods; Astrocyte adapts calls accordingly. 

18Providers with unrecognized versions are rejected. 

19""" 

20 

21from __future__ import annotations 

22 

23from datetime import datetime 

24from typing import TYPE_CHECKING, ClassVar, Protocol, runtime_checkable 

25 

26if TYPE_CHECKING: 

27 from astrocyte.types import ( 

28 Completion, 

29 Document, 

30 DocumentFilters, 

31 DocumentHit, 

32 EngineCapabilities, 

33 Entity, 

34 EntityCandidateMatch, 

35 EntityLink, 

36 ForgetRequest, 

37 ForgetResult, 

38 GraphHit, 

39 HealthStatus, 

40 HttpClientContext, 

41 LLMCapabilities, 

42 MemoryEntityAssociation, 

43 MemoryLink, 

44 MentalModel, 

45 Message, 

46 PageIndexDocument, 

47 PageIndexFact, 

48 PageIndexFactHit, 

49 PageIndexSection, 

50 PageIndexSectionEntity, 

51 PageIndexSectionLink, 

52 RecallRequest, 

53 RecallResult, 

54 ReflectRequest, 

55 ReflectResult, 

56 RetainRequest, 

57 RetainResult, 

58 SourceChunk, 

59 SourceDocument, 

60 ToolDefinition, 

61 TransportCapabilities, 

62 VectorFilters, 

63 VectorHit, 

64 VectorItem, 

65 WikiPage, 

66 WikiPageHit, 

67 ) 

68 

69 

70# --------------------------------------------------------------------------- 

71# Tier 1: Storage Providers 

72# --------------------------------------------------------------------------- 

73 

74 

75@runtime_checkable 

76class VectorStore(Protocol): 

77 """SPI for vector database adapters. Required for Tier 1. 

78 

79 Implement this protocol to connect Astrocyte to a vector database 

80 (e.g., pgvector, Qdrant, Pinecone). Astrocyte uses ``@runtime_checkable`` 

81 to verify your class satisfies this contract at registration time — 

82 no inheritance required, just implement every method below. 

83 

84 See ``astrocyte_postgres.store.PostgresStore`` for a production example, 

85 or ``astrocyte.testing.in_memory.InMemoryVectorStore`` for a minimal reference. 

86 """ 

87 

88 SPI_VERSION: ClassVar[int] = 1 

89 

90 async def store_vectors(self, items: list[VectorItem]) -> list[str]: 

91 """Persist vectors with metadata. Upsert semantics — existing IDs are overwritten. 

92 

93 Implementations must validate that each item's vector length matches the 

94 configured embedding dimensions. 

95 

96 Returns: 

97 List of stored IDs (same order as ``items``). If ``items`` is empty, 

98 implementations must treat this as a no-op and return an empty list. 

99 

100 Raises: 

101 ValueError: If any vector's length does not match embedding dimensions. 

102 """ 

103 pass 

104 

105 async def search_similar( 

106 self, 

107 query_vector: list[float], 

108 bank_id: str, 

109 limit: int = 10, 

110 filters: VectorFilters | None = None, 

111 ) -> list[VectorHit]: 

112 """Find vectors most similar to ``query_vector`` within a bank. 

113 

114 Returns: 

115 Hits sorted by similarity descending. Each hit's ``score`` must be 

116 in the range [0.0, 1.0] where 1.0 is an exact match. Returns an 

117 empty list if ``bank_id`` does not exist. 

118 

119 Raises: 

120 ValueError: If ``query_vector`` length does not match embedding dimensions. 

121 """ 

122 pass 

123 

124 async def delete(self, ids: list[str], bank_id: str) -> int: 

125 """Delete vectors by ID within a bank. Bank isolation is required — 

126 an ID in ``bank-2`` must not be deleted when ``bank_id`` is ``bank-1``. 

127 

128 Returns: 

129 Count of vectors actually deleted. Returns 0 for empty ``ids`` list. 

130 """ 

131 pass 

132 

133 async def list_vectors( 

134 self, 

135 bank_id: str, 

136 offset: int = 0, 

137 limit: int = 100, 

138 ) -> list[VectorItem]: 

139 """List vectors in a bank with pagination. Used by consolidation 

140 to scan and deduplicate vectors. 

141 

142 Implementations must return vectors in a **stable order** (e.g., by ID) 

143 so that pagination produces consistent, complete results. 

144 

145 Returns: 

146 Up to ``limit`` vectors starting at ``offset``. Empty list if 

147 the bank does not exist or ``offset`` is past the end. 

148 """ 

149 pass 

150 

151 async def health(self) -> HealthStatus: 

152 """Check database connectivity. 

153 

154 Returns: 

155 ``HealthStatus(healthy=True, ...)`` if the database is reachable, 

156 ``HealthStatus(healthy=False, ...)`` with a diagnostic message otherwise. 

157 """ 

158 pass 

159 

160 

161@runtime_checkable 

162class GraphStore(Protocol): 

163 """SPI for graph database adapters. Optional for Tier 1. 

164 

165 Implement this protocol to add entity/relationship-based retrieval 

166 alongside vector search. Astrocyte's multi-strategy retrieval pipeline 

167 will query both VectorStore and GraphStore in parallel when available. 

168 """ 

169 

170 SPI_VERSION: ClassVar[int] = 1 

171 

172 async def store_entities(self, entities: list[Entity], bank_id: str) -> list[str]: 

173 """Persist or update entities (people, concepts, objects) in the graph. 

174 

175 Returns: 

176 List of stored entity IDs. 

177 """ 

178 pass 

179 

180 async def store_links(self, links: list[EntityLink], bank_id: str) -> list[str]: 

181 """Persist relationships between entities. 

182 

183 Returns: 

184 List of stored link IDs. 

185 """ 

186 pass 

187 

188 async def link_memories_to_entities( 

189 self, 

190 associations: list[MemoryEntityAssociation], 

191 bank_id: str, 

192 ) -> None: 

193 """Associate memory IDs (from VectorStore) with entity IDs in the graph. 

194 

195 This enables graph-based recall: "find memories connected to entity X." 

196 """ 

197 pass 

198 

199 async def store_memory_links( 

200 self, 

201 links: list["MemoryLink"], 

202 bank_id: str, 

203 ) -> list[str]: 

204 """Persist directional memory-to-memory links (Hindsight parity). 

205 

206 Used for causal chains (``caused_by``) and the precomputed 

207 semantic kNN graph (``semantic``). Adapters that don't 

208 implement this should silently no-op (return ``[]``) — the 

209 spread / link-expansion path probes via ``getattr`` and 

210 degrades to entity-only when the method is absent. 

211 

212 Returns the list of stored link IDs. 

213 """ 

214 return [] 

215 

216 async def find_memory_links( 

217 self, 

218 seed_memory_ids: list[str], 

219 bank_id: str, 

220 *, 

221 link_types: list[str] | None = None, 

222 limit: int = 200, 

223 ) -> list["MemoryLink"]: 

224 """Return memory links touching any seed (either direction). 

225 

226 Adapter contract: results SHOULD include links where the seed 

227 is the source OR the target — the link-expansion retrieval 

228 treats both directions for ``semantic`` and only the target 

229 direction for ``caused_by`` (effect → cause). Callers filter 

230 as needed. 

231 

232 Adapters that haven't implemented this return ``[]``; the 

233 link-expansion retrieval probes via ``getattr`` and skips the 

234 memory-link signal when the method is absent. 

235 """ 

236 return [] 

237 

238 async def query_neighbors( 

239 self, 

240 entity_ids: list[str], 

241 bank_id: str, 

242 max_depth: int = 2, 

243 limit: int = 20, 

244 ) -> list[GraphHit]: 

245 """Traverse the graph from ``entity_ids`` up to ``max_depth`` hops 

246 and return connected memories. 

247 

248 Returns: 

249 Graph hits sorted by relevance (e.g., shortest path, edge weight). 

250 """ 

251 pass 

252 

253 async def query_entities( 

254 self, 

255 query: str, 

256 bank_id: str, 

257 limit: int = 10, 

258 ) -> list[Entity]: 

259 """Search entities by name or alias (fuzzy or exact match). 

260 

261 Returns: 

262 Matching entities, useful for resolving entity references before 

263 calling ``query_neighbors``. 

264 """ 

265 pass 

266 

267 async def find_entity_candidates( 

268 self, 

269 name: str, 

270 bank_id: str, 

271 threshold: float = 0.8, 

272 limit: int = 5, 

273 ) -> list[Entity]: 

274 """Return entities whose name is similar to *name* above *threshold*. 

275 

276 Used by the entity resolution pipeline (M11) to find candidate 

277 entities that may be aliases of a newly-extracted entity before 

278 calling the LLM confirmation step. 

279 

280 A simple implementation may use case-insensitive substring matching; 

281 production adapters should use vector similarity or full-text search. 

282 

283 Returns: 

284 Candidate entities ordered by similarity descending. 

285 """ 

286 pass 

287 

288 async def find_entity_candidates_scored( 

289 self, 

290 name: str, 

291 bank_id: str, 

292 *, 

293 name_embedding: list[float] | None = None, 

294 trigram_threshold: float = 0.3, 

295 limit: int = 10, 

296 ) -> list[EntityCandidateMatch]: 

297 """Return scored entity candidates for the Hindsight entity-resolution cascade. 

298 

299 Scores each candidate against two cheap signals: 

300 

301 - ``name_similarity`` — trigram similarity between the candidate's 

302 name and ``name`` in ``[0.0, 1.0]``. Production adapters use 

303 ``pg_trgm.similarity()``; the in-memory adapter approximates with 

304 :class:`difflib.SequenceMatcher`. 

305 - ``embedding_similarity`` — cosine similarity between the 

306 candidate's stored embedding and ``name_embedding``, or ``None`` 

307 when either side has no embedding. 

308 

309 The resolver uses these scores to autolink, skip, or escalate to LLM 

310 without further round-trips. 

311 

312 Args: 

313 name: Surface form to resolve. 

314 bank_id: Bank scope. 

315 name_embedding: Optional embedding of ``name`` for the embedding 

316 tier. When ``None``, ``embedding_similarity`` is ``None`` for 

317 every result and the cascade falls back to trigram + LLM. 

318 trigram_threshold: Adapter-side prefilter — candidates with a 

319 trigram similarity strictly below this value are dropped. 

320 Default ``0.3`` matches PostgreSQL's ``pg_trgm.similarity_threshold``. 

321 limit: Maximum candidates returned, ordered by descending score. 

322 

323 Returns: 

324 ``list[EntityCandidateMatch]`` ordered by best score first. 

325 Defaults to an empty list when no candidates clear the threshold. 

326 """ 

327 pass 

328 

329 async def store_entity_link(self, link: EntityLink, bank_id: str) -> str: 

330 """Persist a single typed relationship between two entities. 

331 

332 Unlike ``store_links`` (bulk, co-occurrence), this method is called 

333 by the entity resolution pipeline for confirmed alias links that carry 

334 ``evidence`` and ``confidence``. 

335 

336 Returns: 

337 The stored link ID. 

338 """ 

339 pass 

340 

341 async def health(self) -> HealthStatus: 

342 """Check database connectivity.""" 

343 pass 

344 

345 

346@runtime_checkable 

347class DocumentStore(Protocol): 

348 """SPI for document storage with full-text search. Optional for Tier 1. 

349 

350 Implement this protocol to add BM25/lexical search alongside vector 

351 similarity. Useful for keyword-heavy queries where semantic search 

352 alone may miss exact matches. 

353 """ 

354 

355 SPI_VERSION: ClassVar[int] = 1 

356 

357 async def store_document(self, document: Document, bank_id: str) -> str: 

358 """Persist a source document for full-text indexing. 

359 

360 Returns: 

361 The stored document ID. 

362 """ 

363 pass 

364 

365 async def search_fulltext( 

366 self, 

367 query: str, 

368 bank_id: str, 

369 limit: int = 10, 

370 filters: DocumentFilters | None = None, 

371 ) -> list[DocumentHit]: 

372 """BM25 full-text search over stored content. 

373 

374 Returns: 

375 Hits sorted by BM25 relevance score descending. 

376 """ 

377 pass 

378 

379 async def get_document(self, document_id: str, bank_id: str) -> Document | None: 

380 """Retrieve a single document by ID. 

381 

382 Returns: 

383 The document, or ``None`` if not found in this bank. 

384 """ 

385 pass 

386 

387 async def health(self) -> HealthStatus: 

388 """Check database connectivity.""" 

389 pass 

390 

391 

392# --------------------------------------------------------------------------- 

393# Tier 1: Wiki Store (M8) 

394# --------------------------------------------------------------------------- 

395 

396 

397@runtime_checkable 

398class WikiStore(Protocol): 

399 """SPI for wiki page storage (M8 LLM wiki compile). Optional. 

400 

401 WikiStore persists structured WikiPage metadata. Vector embeddings of 

402 compiled pages are stored separately in the VectorStore with 

403 ``memory_layer="compiled"`` and ``fact_type="wiki"``, so recall tiering 

404 can search them via the standard ``search_similar`` path. 

405 

406 Implement this protocol to enable ``brain.compile()`` persistence. 

407 See ``astrocyte.testing.in_memory.InMemoryWikiStore`` for a reference. 

408 """ 

409 

410 SPI_VERSION: ClassVar[int] = 1 

411 

412 async def upsert_page(self, page: WikiPage, bank_id: str) -> str: 

413 """Create or update a wiki page. Upsert semantics — if a page with 

414 the same ``page_id`` exists in this bank, its revision is incremented 

415 and content replaced. The previous revision is archived (not deleted). 

416 

417 Returns: 

418 The stored ``page_id``. 

419 """ 

420 pass 

421 

422 async def get_page(self, page_id: str, bank_id: str) -> WikiPage | None: 

423 """Retrieve the current revision of a wiki page by ID. 

424 

425 Returns: 

426 The page, or ``None`` if not found in this bank. 

427 """ 

428 pass 

429 

430 async def list_pages( 

431 self, 

432 bank_id: str, 

433 scope: str | None = None, 

434 kind: str | None = None, 

435 ) -> list[WikiPage]: 

436 """List current-revision wiki pages for a bank. 

437 

438 Args: 

439 bank_id: The bank to list pages for. 

440 scope: If set, return only pages whose ``scope`` matches exactly. 

441 kind: If set, return only pages of this kind ("entity", "topic", "concept"). 

442 

443 Returns: 

444 All matching pages, unsorted. 

445 """ 

446 pass 

447 

448 async def delete_page(self, page_id: str, bank_id: str) -> bool: 

449 """Delete a wiki page (current revision and audit log). 

450 

451 Returns: 

452 ``True`` if the page was found and deleted, ``False`` if not found. 

453 """ 

454 pass 

455 

456 async def health(self) -> HealthStatus: 

457 """Check storage connectivity.""" 

458 pass 

459 

460 

461# --------------------------------------------------------------------------- 

462# Section recall (M9): PageIndex tree + section graph store 

463# --------------------------------------------------------------------------- 

464 

465 

466@runtime_checkable 

467class PageIndexStore(Protocol): 

468 """SPI for the section recall store (M9). Optional. 

469 

470 Backs the three-layer recall stack defined in 

471 ``docs/_design/recall.md``. Holds: 

472 

473 - ``PageIndexDocument`` — one row per conversation/document, with the 

474 canonical markdown that the picker slices for synth excerpts. 

475 - ``PageIndexSection`` — tree nodes (the M9 recall primitive). 

476 - ``PageIndexSectionEntity`` — Hindsight's unit_entities pattern at 

477 section grain (PR2 commit A populates). 

478 - ``PageIndexSectionLink`` — Hindsight's memory_links pattern at 

479 section grain (PR2 commit D populates). 

480 

481 PR1 (this milestone) only exercises ``save_document`` / 

482 ``save_sections`` / ``load_skeleton`` / ``load_document`` — the 

483 minimum needed to port the Phase A POC picker to a Postgres-backed 

484 cache. Entity / link methods land empty in PR1 and are populated in 

485 PR2. 

486 

487 See ``astrocyte.testing.in_memory.InMemoryPageIndexStore`` for the 

488 reference fixture; ``astrocyte_postgres.PostgresPageIndexStore`` for 

489 the production backend. 

490 """ 

491 

492 SPI_VERSION: ClassVar[int] = 1 

493 

494 async def save_document(self, doc: PageIndexDocument) -> str: 

495 """Upsert a document. Returns the document id (UUID string). 

496 

497 Upsert keyed on ``(bank_id, source_id)`` — re-running tree-build 

498 for the same conversation replaces the prior row. 

499 """ 

500 

501 async def save_sections( 

502 self, 

503 document_id: str, 

504 sections: list[PageIndexSection], 

505 ) -> int: 

506 """Bulk-replace all sections for a document. Returns the count 

507 written. Existing rows for ``document_id`` are deleted first 

508 (the tree rebuild is treated as atomic — no partial trees).""" 

509 

510 async def load_document( 

511 self, 

512 bank_id: str, 

513 source_id: str, 

514 ) -> PageIndexDocument | None: 

515 """Fetch one document by (bank_id, source_id). Returns ``None`` 

516 when not found (caller decides whether to build it).""" 

517 

518 async def load_skeleton(self, document_id: str) -> list[PageIndexSection]: 

519 """Fetch all sections for a document, ordered by ``line_num``. 

520 

521 Returned sections do NOT carry ``summary_embedding`` (that field 

522 is only used by the semantic strategy in PR2; the picker doesn't 

523 need it). Implementations may project it out for cheaper reads. 

524 """ 

525 

526 async def save_section_embeddings( 

527 self, 

528 document_id: str, 

529 embeddings: list[tuple[int, list[float]]], 

530 ) -> int: 

531 """PR2 commit A: bulk-update ``summary_embedding`` on existing 

532 section rows. Skips rows whose ``line_num`` doesn't already 

533 exist (the tree-build step is the source of truth for which 

534 sections exist). Returns rows updated.""" 

535 

536 async def save_section_entities( 

537 self, 

538 entities: list[PageIndexSectionEntity], 

539 ) -> int: 

540 """Bulk-insert entity-mention rows. Idempotent on the composite 

541 primary key. Returns rows written. PR2 commit A populates.""" 

542 

543 async def save_section_links( 

544 self, 

545 links: list[PageIndexSectionLink], 

546 ) -> int: 

547 """Bulk-insert section-link rows. Idempotent on the composite 

548 primary key. Returns rows written. PR2 commit B/D populates.""" 

549 

550 async def populate_semantic_knn_links( 

551 self, 

552 document_id: str, 

553 *, 

554 top_k: int = 5, 

555 min_similarity: float = 0.5, 

556 ) -> int: 

557 """PR2 D.7.1: Populate ``section_links`` with kNN edges over 

558 ``summary_embedding``. For each section, link to its top_k 

559 most-similar OTHER sections in the same document with 

560 ``link_type='semantic_knn'`` and weight = cosine similarity. 

561 

562 Pure SQL — no LLM call. Idempotent. Returns rows inserted. 

563 

564 Why this exists: D.7's LLM-based link extraction over-emits on 

565 LoCoMo and severely under-emits on LME (5 vs 100 links/doc). 

566 kNN restores dense topical bridging for graph_expand.""" 

567 

568 # ── PR2 commit B: query methods for the 5 parallel strategies ── 

569 # 

570 # These are pure read methods that the section recall orchestrator 

571 # (``astrocyte.pipeline.section_recall``) calls in parallel. Each 

572 # returns a ranked list of ``(document_id, line_num, score)`` 

573 # tuples; the orchestrator fuses them via RRF. 

574 # 

575 # Single-bank scoping for now (PR2). Cross-bank scoping comes when 

576 # multi-bank section recall ships in M10+. 

577 

578 async def search_sections_semantic( 

579 self, 

580 bank_id: str, 

581 query_embedding: list[float], 

582 *, 

583 top_k: int = 20, 

584 session_filter: str | None = None, 

585 ) -> list[tuple[str, int, float]]: 

586 """PR2 commit B / semantic strategy: cosine-similarity over 

587 ``summary_embedding``. Returns ``(document_id, line_num, score)`` 

588 tuples, ordered by similarity desc. Score is in [0, 1] (1 - distance). 

589 

590 M31 Fix 2: ``session_filter`` (opaque str) scopes results to 

591 sections whose ``session_id`` matches. ``None`` (default) 

592 retrieves cross-session. 

593 """ 

594 

595 async def search_sections_keyword( 

596 self, 

597 bank_id: str, 

598 query: str, 

599 *, 

600 top_k: int = 20, 

601 speaker: str | None = None, 

602 document_id: str | None = None, 

603 session_filter: str | None = None, 

604 ) -> list[tuple[str, int, float]]: 

605 """PR2 commit B / keyword strategy: full-text search (``tsvector``) 

606 over section titles + summaries. Returns ``(document_id, 

607 line_num, score)`` tuples. The optional ``speaker`` filter is 

608 for LME assistant-recall (``WHERE speaker = 'assistant'``). 

609 The optional ``document_id`` filter (PR2.6) scopes the search 

610 to a single document — used by ``temporal_arithmetic.find_event_date`` 

611 to avoid bank-wide top-K starvation when 50+ documents share a 

612 bank and the matching one isn't in the unfiltered top-5. 

613 

614 M31 Fix 2: ``session_filter`` (opaque str) scopes to sections 

615 whose ``session_id`` matches.""" 

616 

617 async def search_sections_by_entities( 

618 self, 

619 bank_id: str, 

620 entity_names: list[str], 

621 *, 

622 top_k: int = 20, 

623 session_filter: str | None = None, 

624 ) -> list[tuple[str, int, float]]: 

625 """PR2 commit B / entity strategy: Hindsight's CTE pattern. 

626 Returns ``(document_id, line_num, score)`` where score is the 

627 count of distinct matching entity_names per section. 

628 

629 M31 Fix 2: ``session_filter`` joins to sections to scope by 

630 ``session_id``.""" 

631 

632 async def search_sections_temporal( 

633 self, 

634 bank_id: str, 

635 date_range: tuple[datetime, datetime], 

636 *, 

637 top_k: int = 20, 

638 session_filter: str | None = None, 

639 ) -> list[tuple[str, int, float]]: 

640 """PR2 commit B / temporal strategy: date-range filter on 

641 ``session_date`` (uses the partial btree index from migration 

642 015). Score is uniform (1.0) — temporal is a filter, not a 

643 ranker; ranking happens in fusion. 

644 

645 M31 Fix 2: ``session_filter`` scopes to sections whose 

646 ``session_id`` matches.""" 

647 

648 async def expand_section_links( 

649 self, 

650 seeds: list[tuple[str, int]], 

651 *, 

652 link_types: list[str] | None = None, 

653 top_k: int = 20, 

654 ) -> list[tuple[str, int, float]]: 

655 """PR2 commit B / graph-expand strategy: 1-hop expansion through 

656 ``section_links`` from the given ``(document_id, line_num)`` 

657 seeds. ``link_types`` filters to e.g. ['semantic_knn', 'causal']; 

658 None means all types. Score is link weight.""" 

659 

660 async def expand_sections_by_shared_entities( 

661 self, 

662 bank_id: str, 

663 seeds: list[tuple[str, int]], 

664 *, 

665 top_k: int = 20, 

666 exclude_seeds: bool = True, 

667 ) -> list[tuple[str, int, float]]: 

668 """Fix 3 (conv-run-4) / entity spreading activation: 1-hop 

669 expansion through entity co-occurrence in 

670 ``astrocyte_pi_section_entities``. 

671 

672 For each seed ``(document_id, line_num)``, find every section 

673 in the same bank that shares at least one entity with the seed. 

674 Returns ``(document_id, line_num, score)`` where score is the 

675 count of distinct shared entities. ``exclude_seeds=True`` 

676 filters the seeds themselves out of the result. 

677 

678 Why this is distinct from ``expand_section_links``: links rely 

679 on ``astrocyte_pi_section_links`` being populated (semantic_knn 

680 and the LLM-extracted causal/elaborates edges). In conversation 

681 ingest the link density is sparse, so the existing graph_expand 

682 strategy can't bridge Denver → Red Rocks even when both 

683 sections mention the same entity. Entity spread uses the 

684 already-populated ``astrocyte_pi_section_entities`` table 

685 directly, which gives a denser bridge with the same SQL cost. 

686 """ 

687 

688 # ── M12.1: fact-grain layer (atomic facts alongside sections) ─────── 

689 

690 async def save_facts(self, facts: list["PageIndexFact"]) -> int: 

691 """M12.1: persist atomic facts extracted from sections. 

692 

693 Each fact's ``embedding`` may be ``None`` (caller embeds in a 

694 separate batched pass) or populated. Returns the number of rows 

695 inserted. Implementations should validate that 

696 ``(document_id, line_num)`` references an existing section and 

697 let the FK cascade handle deletion.""" 

698 

699 async def update_fact_embeddings( 

700 self, 

701 embeddings: list[tuple[str, list[float]]], 

702 ) -> int: 

703 """M12.1: write the embedding column for a batch of facts. 

704 

705 ``embeddings`` is ``[(fact_id, embedding_vector), ...]``. 

706 Returns the number of rows updated. Separated from 

707 :meth:`save_facts` so the bench can extract facts and embed 

708 them in two parallel passes — Hindsight's retain shape.""" 

709 

710 async def search_facts_semantic( 

711 self, 

712 bank_id: str, 

713 query_embedding: list[float], 

714 *, 

715 top_k: int = 20, 

716 document_id: str | None = None, 

717 fact_type: str | None = None, 

718 session_filter: str | None = None, 

719 ) -> list["PageIndexFactHit"]: 

720 """M12.1: cosine-similarity search over fact embeddings. 

721 

722 M31 Fix 2: ``session_filter`` scopes results to facts whose 

723 anchoring section has matching ``session_id`` (EXISTS subquery 

724 in the Postgres adapter, _matches_session helper in the 

725 in-memory adapter). Top-level facts (no anchor) are excluded 

726 when ``session_filter`` is set.""" 

727 

728 async def search_facts_by_entity( 

729 self, 

730 bank_id: str, 

731 entity_name: str, 

732 *, 

733 top_k: int = 50, 

734 document_id: str | None = None, 

735 fact_type: str | None = None, 

736 session_filter: str | None = None, 

737 ) -> list["PageIndexFactHit"]: 

738 """M12.1: list facts whose ``entities`` array contains 

739 ``entity_name`` (or a case-insensitive variant). Used by the 

740 agent's counting tool — e.g. for "how many doctors" the agent 

741 queries ``entity_name="role:doctor"`` and counts the rows. 

742 

743 M31 Fix 2: ``session_filter`` mirrors the semantic search 

744 contract. For cross-session link-expansion usage (M27), 

745 ``session_filter`` should be ``None`` (the default) so the 

746 whole point of cross-session traversal is preserved. 

747 

748 M34-3: ``fact_type`` filters to a single fact_type (e.g. 

749 ``'preference'``) so per-fact-type segmented retrieval can 

750 produce independent candidate pools per type.""" 

751 

752 async def search_facts_temporal( 

753 self, 

754 bank_id: str, 

755 date_range: tuple[datetime, datetime], 

756 *, 

757 top_k: int = 50, 

758 document_id: str | None = None, 

759 fact_type: str | None = None, 

760 session_filter: str | None = None, 

761 ) -> list["PageIndexFactHit"]: 

762 """M12.1: list facts whose ``occurred_start`` falls in the 

763 date range. Powers temporal-arithmetic queries that previously 

764 had to resort to keyword search over section summaries. 

765 

766 M31 Fix 2: ``session_filter`` scopes to one session's 

767 temporal facts. 

768 

769 M34-3: ``fact_type`` filters to a single fact_type so the 

770 temporal channel's per-type segment doesn't leak across types.""" 

771 

772 async def search_facts_keyword( 

773 self, 

774 bank_id: str, 

775 query: str, 

776 *, 

777 top_k: int = 20, 

778 document_id: str | None = None, 

779 fact_type: str | None = None, 

780 session_filter: str | None = None, 

781 ) -> list["PageIndexFactHit"]: 

782 """M31c — Hindsight-parity BM25/full-text search over fact 

783 text. Complements ``search_facts_semantic`` (vector cosine): 

784 BM25 catches specific noun/number matches (e.g. "500 Mbps", 

785 "Philips LED", "University of Melbourne") that semantic 

786 embeddings tend to under-weight relative to thematically 

787 similar but topically off content. 

788 

789 Implementations: 

790 

791 - Postgres adapter: ``plainto_tsquery`` + ``ts_rank_cd`` over 

792 ``to_tsvector('english', fact_text)`` (computed on-the-fly; 

793 no schema migration needed). Mirrors the 

794 ``search_sections_keyword`` pattern that's been validated 

795 since M9. 

796 - InMemory adapter: token-overlap scoring (count of distinct 

797 query tokens appearing in fact text), case-insensitive. 

798 

799 Returns ``(document_id, line_num, score)``-bearing 

800 ``PageIndexFactHit`` records, ordered by score desc. Score is 

801 rank-rather-than-similarity (BM25 magnitudes are not directly 

802 comparable to vector cosine — RRF fusion in ``fact_recall`` 

803 handles the scale difference). 

804 

805 Diagnostic that motivated this: SSU bench failures showed the 

806 cross-encoder rerank can rank a specific factoid statement 

807 ("User upgraded to 500 Mbps internet plan") deep in the 

808 candidate pool while surfacing thematically-related but 

809 irrelevant content ("commute", "road trip") at the top. 

810 BM25 catches the literal-keyword match that semantic alone 

811 misses.""" 

812 

813 async def count_facts_matching( 

814 self, 

815 bank_id: str, 

816 document_id: str, 

817 *, 

818 entity_pattern: str | None = None, 

819 fact_type: str | None = None, 

820 ) -> int: 

821 """M12.1: deterministic count of facts matching the given 

822 filters. ``entity_pattern`` is a substring (case-insensitive) 

823 matched against any element of the ``entities`` array. Used 

824 by the agent's counting tool for questions like "how many 

825 doctors did I visit?" — agent calls with 

826 ``entity_pattern="role:doctor"`` and gets back the exact 

827 count.""" 

828 

829 async def save_section_event_dates( 

830 self, 

831 document_id: str, 

832 event_dates: list[tuple[int, datetime, datetime | None]], 

833 ) -> int: 

834 """M11.1: update ``occurred_start`` / ``occurred_end`` on 

835 existing section rows. 

836 

837 ``event_dates`` is a list of ``(line_num, start, end_or_None)`` 

838 tuples. Rows not in the list are left untouched. Returns the 

839 number of rows updated. Used by the bench retain pipeline 

840 after section_event_extraction runs.""" 

841 

842 async def load_sections_with_embeddings( 

843 self, 

844 bank_id: str, 

845 document_id: str, 

846 ) -> list["PageIndexSection"]: 

847 """M10.1: load sections including ``summary_embedding`` (which 

848 :meth:`load_skeleton` deliberately drops as a cost optimisation). 

849 

850 Used by :func:`section_compile.compile_sections_for_document` to 

851 cluster sections via DBSCAN. Returns sections in arbitrary order; 

852 caller sorts as needed.""" 

853 

854 async def save_wiki_page( 

855 self, 

856 *, 

857 page: "WikiPage", 

858 embedding: list[float] | None, 

859 provenance: list[tuple[str, int]], 

860 ) -> str: 

861 """M10.1: persist one observation as a wiki page + its 

862 section-grain provenance + (optional) embedding for the wiki 

863 recall strategy. 

864 

865 Implementations must: 

866 - Insert / upsert the wiki_page row. 

867 - Insert provenance rows ``(wiki_page_id, document_id, line_num)`` 

868 via ``astrocyte_pi_wiki_provenance`` (migration 015). 

869 - Set ``current_embedding`` on the wiki page when the column 

870 exists (migration 018 adds it). Pass ``None`` to skip 

871 embedding (caller does its own deferred indexing). 

872 

873 Returns the page's ``page_id``.""" 

874 

875 async def search_wiki_pages_semantic( 

876 self, 

877 bank_id: str, 

878 query_embedding: list[float], 

879 *, 

880 top_k: int = 5, 

881 document_id: str | None = None, 

882 ) -> list["WikiPageHit"]: 

883 """M10.1: cosine-similarity search over 

884 ``astrocyte_wiki_pages.current_embedding`` (migration 018). 

885 

886 Returns up to ``top_k`` :class:`~astrocyte.types.WikiPageHit` 

887 rows ordered by similarity desc. The optional ``document_id`` 

888 filter scopes to a single document (the bench is per-doc; prod 

889 deployments may want bank-wide).""" 

890 

891 async def count_wiki_pages_for_doc( 

892 self, 

893 bank_id: str, 

894 document_id: str, 

895 ) -> int: 

896 """M10.1: idempotency check for the consolidation pass. 

897 ``compile_sections_for_document`` skips work when this returns 

898 > 0.""" 

899 

900 async def list_wiki_pages_for_doc( 

901 self, 

902 bank_id: str, 

903 document_id: str, 

904 ) -> list["WikiPage"]: 

905 """M12.6: enumerate current-revision wiki pages with 

906 ``scope = 'document:<id>'``. Returns full WikiPage rows with 

907 markdown content populated; used by the revision pass to load 

908 existing pages, send them through an LLM update, and persist 

909 bumped revisions. 

910 

911 Default implementation returns ``[]`` for stores that don't 

912 yet implement wiki enumeration — the revision pass becomes a 

913 no-op on those backends rather than crashing.""" 

914 return [] 

915 

916 async def list_wikis_affected_by_entities( 

917 self, 

918 bank_id: str, 

919 entities: list[str], 

920 *, 

921 min_overlap: int = 1, 

922 limit: int = 8, 

923 ) -> list[tuple["WikiPage", int, list[str]]]: 

924 """M14.2: find existing wiki pages whose provenance sections 

925 share entities with a new source, returning fully-hydrated 

926 WikiPage rows + overlap stats. 

927 

928 Backs Karpathy's update-affected-wikis op. The aggregation 

929 ``HAVING`` cuts the wiki set down to the small subset (~5-15 

930 per source) where the entity-overlap signal is real; returning 

931 full rows means callers don't need a separate ``get_page`` 

932 round-trip per result. 

933 

934 Args: 

935 bank_id: tenant scope. 

936 entities: distinct entity names from the new source. 

937 min_overlap: minimum shared entities to flag a wiki as 

938 affected. ``1`` is the loosest threshold. 

939 limit: cap on returned rows (top-N by overlap count). 

940 

941 Returns: 

942 ``[(WikiPage, overlap_count, shared_entities), ...]`` 

943 sorted descending by ``overlap_count``, ties broken by 

944 ``page_id`` ascending for stability. 

945 

946 Default implementation returns ``[]`` for stores that haven't 

947 yet been migrated — the update pass becomes a no-op on those 

948 backends rather than crashing. 

949 """ 

950 _ = bank_id, entities, min_overlap, limit # noqa: ARG002 

951 return [] 

952 

953 async def list_distinct_entities( 

954 self, 

955 bank_id: str, 

956 document_id: str, 

957 *, 

958 pattern: str | None = None, 

959 limit: int = 50, 

960 ) -> list[tuple[str, int]]: 

961 """PR2.6 / agentic counting tool: list distinct entity names in a 

962 document with their per-section mention counts. 

963 

964 ``pattern`` is a SQL ILIKE substring (case-insensitive) — pass 

965 ``"doctor"`` to match ``"Dr. Smith"``, ``"%kit%"`` to match any 

966 entity containing ``"kit"``. ``None`` returns the top-``limit`` 

967 entities by mention count (useful when the agent doesn't yet 

968 know what to look for). 

969 

970 Returns ``[(entity_name, count), ...]`` ordered by count desc 

971 then name asc. ``count`` is the number of distinct sections in 

972 which the entity appears (the (document_id, line_num, 

973 entity_name) PK in 015 means one section can mention an entity 

974 at most once). 

975 

976 Used by the reflect agent's ``list_entities`` tool to count 

977 across-session mentions deterministically — fixes the LME 

978 multi-session counting holes from the PR2.6 gate where the LLM 

979 couldn't reliably aggregate entity mentions from raw section 

980 text.""" 

981 

982 async def health(self) -> HealthStatus: 

983 """Check storage connectivity.""" 

984 

985 

986# --------------------------------------------------------------------------- 

987# Tier 1: Mental Model Store (first-class — see Hindsight comparison) 

988# --------------------------------------------------------------------------- 

989 

990 

991@runtime_checkable 

992class MentalModelStore(Protocol): 

993 """SPI for first-class mental-model storage. 

994 

995 Mental models are curated, refreshable saved-reflect summaries — the 

996 "Caroline prefers async updates" / "Project X status: blocked on 

997 review" durable artifacts that live longer than any single recall. 

998 

999 Previously these piggybacked on :class:`WikiStore` via 

1000 ``kind="concept"`` + ``metadata["_mental_model"] = True``, which 

1001 overloaded the wiki layer's lifecycle (revisions, lint issues, 

1002 cross_links) for a fundamentally different concept. This SPI breaks 

1003 them out into their own table with their own lifecycle and version 

1004 history. 

1005 

1006 See ``astrocyte.testing.in_memory.InMemoryMentalModelStore`` for a 

1007 reference and ``astrocyte_postgres.PostgresMentalModelStore`` for the 

1008 production-grade implementation. 

1009 """ 

1010 

1011 SPI_VERSION: ClassVar[int] = 1 

1012 

1013 async def upsert(self, model: "MentalModel", bank_id: str) -> int: 

1014 """Create or refresh a mental model. Upsert semantics — if a 

1015 model with the same ``model_id`` exists in this bank, the 

1016 revision is incremented and the new content replaces the 

1017 current; the previous revision is archived (not deleted). 

1018 

1019 Args: 

1020 model: The mental model to store. ``revision`` and 

1021 ``refreshed_at`` are assigned by the store. 

1022 bank_id: Tenant-scoped bank identifier. 

1023 

1024 Returns: 

1025 The new revision number. 

1026 """ 

1027 pass 

1028 

1029 async def get(self, model_id: str, bank_id: str) -> "MentalModel | None": 

1030 """Retrieve the current revision of a mental model. 

1031 

1032 Returns ``None`` if the model doesn't exist or has been deleted. 

1033 """ 

1034 pass 

1035 

1036 async def list( 

1037 self, 

1038 bank_id: str, 

1039 *, 

1040 scope: str | None = None, 

1041 kind: str | None = None, 

1042 ) -> list["MentalModel"]: 

1043 """List current-revision mental models in a bank. 

1044 

1045 Args: 

1046 bank_id: Tenant-scoped bank identifier. 

1047 scope: If set, return only models with matching scope 

1048 (e.g. ``"person:alice"``, or the default ``"bank"``). 

1049 kind: M14.6 — if set, return only models of the given 

1050 sub-type (``"general"`` | ``"preference"``). When 

1051 ``None``, returns all kinds. 

1052 """ 

1053 pass 

1054 

1055 async def delete(self, model_id: str, bank_id: str) -> bool: 

1056 """Soft-delete a mental model. 

1057 

1058 Returns ``True`` if the model was found and deleted, ``False`` if 

1059 it didn't exist or was already deleted. Implementations may keep 

1060 the row for audit (deleted_at timestamp) or hard-delete; either 

1061 is acceptable as long as ``get`` returns ``None`` afterward. 

1062 """ 

1063 pass 

1064 

1065 async def update_via_ops( 

1066 self, 

1067 model_id: str, 

1068 bank_id: str, 

1069 operations_json: list[dict], 

1070 ) -> "tuple[int, dict] | None": 

1071 """M21 — apply structured delta operations to a mental model. 

1072 

1073 Loads the current model, applies the operations against its 

1074 ``structured_doc`` via 

1075 :func:`astrocyte.pipeline.delta_ops.apply_operations`, 

1076 re-renders ``content`` from the new doc, and upserts the 

1077 result. Lazy-migrates legacy rows whose ``structured_doc`` is 

1078 ``None`` by parsing the raw ``content`` markdown on first 

1079 refresh. 

1080 

1081 Args: 

1082 model_id: The model to modify. 

1083 bank_id: Tenant-scoped bank identifier. 

1084 operations_json: List of JSON-shaped op dicts (each 

1085 matching one of the :class:`Operation` variants in 

1086 :mod:`astrocyte.pipeline.delta_ops`). Schema-invalid 

1087 ops are dropped by the operation parser; per-op 

1088 validation drops ops that reference unknown sections. 

1089 

1090 Returns: 

1091 ``(new_revision, applied_delta_summary)`` on success. 

1092 ``None`` when the model doesn't exist. 

1093 The summary is the audit log from 

1094 :class:`~astrocyte.pipeline.delta_ops.AppliedDelta` — 

1095 ``{"applied": [...], "skipped": [...], "changed": bool}`` 

1096 — useful for telemetry / debugging the LLM's output. 

1097 

1098 Conservative-failure contract: invalid ops drop with a logged