Coverage for astrocyte/provider.py: 70%

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 **Optional methods.** Not part of the protocol (so ``isinstance`` checks 

88 don't require them); Astrocyte probes for them with ``getattr`` and falls 

89 back, or reports the capability as unsupported, when they are absent: 

90 

91 - ``list_recent_vectors(bank_id, limit, filters)`` — newest live vectors 

92 first; recency retrieval uses it instead of scanning ``list_vectors``. 

93 - ``get_by_chunk_ids(chunk_ids, bank_id)`` — sibling chunks of a source 

94 chunk, for chunk expansion at recall. 

95 - ``list_changes(bank_id, *, after, limit) -> list[MemoryChange]`` — the 

96 bank's change feed for team-memory sync: **every change to a synced 

97 row, in (changed_at, id) order**. Live rows carry their current values 

98 (so a later metadata or status change shows up as an upsert of the same 

99 id); forgotten rows are tombstones (``deleted=True``, no text). 

100 ``changed_at`` is the row's last change to any synced field: the store 

101 must set it on **every** write that changes a row (insert: its 

102 ``retained_at``; overwrite, restore or any field update: the time of 

103 the change; forget: the time of the forget), and rows that predate it 

104 backfill as ``max(retained_at, forgotten_at)``. ``after`` is a ``(changed_at, id)`` 

105 position; only entries strictly after it are returned, so paging by the 

106 last entry's position never skips or repeats a row, even when many rows 

107 share one ``changed_at``. Ids compare byte-wise (``COLLATE "C"`` on 

108 Postgres). Requires soft deletes (or remembered tombstones): a store 

109 that erases rows outright cannot report that they were forgotten. 

110 ``Astrocyte.list_changes`` (``GET /v1/banks/{bank_id}/changes``) raises 

111 :class:`~astrocyte.errors.CapabilityNotSupported` without it. 

112 - ``lookup_ids(ids) -> list[MemoryChange]`` — the current state of each id 

113 that exists **in any bank**: a live row (full record) or a tombstone. 

114 Ids never stored (or purged) are absent. Deliberately cross-bank: both 

115 SQL stores key rows on ``id`` alone, so a writer that must not take over 

116 another bank's row has to be able to see it. Callers must never return 

117 another bank's content. 

118 - ``insert_vectors(items) -> list[str]`` — like ``store_vectors`` but 

119 insert-only: an item whose id already exists in any bank, live or 

120 forgotten, is skipped, never overwritten (``ON CONFLICT DO NOTHING``). 

121 Returns the ids actually inserted, so a concurrent writer that got there 

122 first is detected rather than clobbered. 

123 ``Astrocyte.push_records`` (``POST /v1/banks/{bank_id}/sync/push``) 

124 needs both, and raises ``CapabilityNotSupported`` without them. 

125 - ``erase(bank_id, ids) -> int`` — erase already-forgotten memories of 

126 ``ids`` from storage for good (text, embedding, metadata and anything 

127 derived from them), keeping each one's tombstone: ``list_changes`` and 

128 ``lookup_ids`` still report it deleted, and ``insert_vectors`` never 

129 stores the id again. Live rows are untouched, so erase never bypasses 

130 forget. Returns how many were erased. ``Astrocyte.erase`` (a gateway 

131 forget with ``erase: true``, a DSAR erasure) raises 

132 ``CapabilityNotSupported`` without it. 

133 - ``save_legal_hold(hold)``, ``delete_legal_hold(bank_id, hold_id) -> bool``, 

134 ``list_legal_holds(bank_id) -> list[LegalHold]`` — persisted legal 

135 holds, so a hold survives restarts and binds every process and replica 

136 sharing the store. ``save`` replaces a hold with the same 

137 ``(bank_id, hold_id)``. ``Astrocyte.place_legal_hold`` uses them, and 

138 ``forget`` / ``run_lifecycle`` re-read them before checking; without 

139 them, holds are kept in the process only. 

140 """ 

141 

142 SPI_VERSION: ClassVar[int] = 1 

143 

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

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

146 

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

148 configured embedding dimensions. 

149 

150 Returns: 

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

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

153 

154 Raises: 

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

156 """ 

157 pass 

158 

159 async def search_similar( 

160 self, 

161 query_vector: list[float], 

162 bank_id: str, 

163 limit: int = 10, 

164 filters: VectorFilters | None = None, 

165 ) -> list[VectorHit]: 

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

167 

168 Returns: 

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

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

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

172 

173 Raises: 

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

175 """ 

176 pass 

177 

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

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

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

181 

182 Returns: 

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

184 """ 

185 pass 

186 

187 async def list_vectors( 

188 self, 

189 bank_id: str, 

190 offset: int = 0, 

191 limit: int = 100, 

192 ) -> list[VectorItem]: 

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

194 to scan and deduplicate vectors. 

195 

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

197 so that pagination produces consistent, complete results. 

198 

199 Returns: 

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

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

202 """ 

203 pass 

204 

205 async def health(self) -> HealthStatus: 

206 """Check database connectivity. 

207 

208 Returns: 

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

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

211 """ 

212 pass 

213 

214 

215@runtime_checkable 

216class GraphStore(Protocol): 

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

218 

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

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

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

222 """ 

223 

224 SPI_VERSION: ClassVar[int] = 1 

225 

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

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

228 

229 Returns: 

230 List of stored entity IDs. 

231 """ 

232 pass 

233 

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

235 """Persist relationships between entities. 

236 

237 Returns: 

238 List of stored link IDs. 

239 """ 

240 pass 

241 

242 async def link_memories_to_entities( 

243 self, 

244 associations: list[MemoryEntityAssociation], 

245 bank_id: str, 

246 ) -> None: 

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

248 

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

250 """ 

251 pass 

252 

253 async def store_memory_links( 

254 self, 

255 links: list["MemoryLink"], 

256 bank_id: str, 

257 ) -> list[str]: 

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

259 

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

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

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

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

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

265 

266 Returns the list of stored link IDs. 

267 """ 

268 return [] 

269 

270 async def find_memory_links( 

271 self, 

272 seed_memory_ids: list[str], 

273 bank_id: str, 

274 *, 

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

276 limit: int = 200, 

277 ) -> list["MemoryLink"]: 

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

279 

280 Adapter contract: results SHOULD include links where the seed 

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

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

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

284 as needed. 

285 

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

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

288 memory-link signal when the method is absent. 

289 """ 

290 return [] 

291 

292 async def query_neighbors( 

293 self, 

294 entity_ids: list[str], 

295 bank_id: str, 

296 max_depth: int = 2, 

297 limit: int = 20, 

298 ) -> list[GraphHit]: 

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

300 and return connected memories. 

301 

302 Returns: 

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

304 """ 

305 pass 

306 

307 async def query_entities( 

308 self, 

309 query: str, 

310 bank_id: str, 

311 limit: int = 10, 

312 ) -> list[Entity]: 

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

314 

315 Returns: 

316 Matching entities, useful for resolving entity references before 

317 calling ``query_neighbors``. 

318 """ 

319 pass 

320 

321 async def find_entity_candidates( 

322 self, 

323 name: str, 

324 bank_id: str, 

325 threshold: float = 0.8, 

326 limit: int = 5, 

327 ) -> list[Entity]: 

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

329 

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

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

332 calling the LLM confirmation step. 

333 

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

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

336 

337 Returns: 

338 Candidate entities ordered by similarity descending. 

339 """ 

340 pass 

341 

342 async def find_entity_candidates_scored( 

343 self, 

344 name: str, 

345 bank_id: str, 

346 *, 

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

348 trigram_threshold: float = 0.3, 

349 limit: int = 10, 

350 ) -> list[EntityCandidateMatch]: 

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

352 

353 Scores each candidate against two cheap signals: 

354 

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

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

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

358 :class:`difflib.SequenceMatcher`. 

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

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

361 when either side has no embedding. 

362 

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

364 without further round-trips. 

365 

366 Args: 

367 name: Surface form to resolve. 

368 bank_id: Bank scope. 

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

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

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

372 trigram_threshold: Adapter-side prefilter — candidates with a 

373 trigram similarity strictly below this value are dropped. 

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

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

376 

377 Returns: 

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

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

380 """ 

381 pass 

382 

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

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

385 

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

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

388 ``evidence`` and ``confidence``. 

389 

390 Returns: 

391 The stored link ID. 

392 """ 

393 pass 

394 

395 async def health(self) -> HealthStatus: 

396 """Check database connectivity.""" 

397 pass 

398 

399 

400@runtime_checkable 

401class DocumentStore(Protocol): 

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

403 

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

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

406 alone may miss exact matches. 

407 """ 

408 

409 SPI_VERSION: ClassVar[int] = 1 

410 

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

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

413 

414 Returns: 

415 The stored document ID. 

416 """ 

417 pass 

418 

419 async def search_fulltext( 

420 self, 

421 query: str, 

422 bank_id: str, 

423 limit: int = 10, 

424 filters: DocumentFilters | None = None, 

425 ) -> list[DocumentHit]: 

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

427 

428 Returns: 

429 Hits sorted by BM25 relevance score descending. 

430 """ 

431 pass 

432 

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

434 """Retrieve a single document by ID. 

435 

436 Returns: 

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

438 """ 

439 pass 

440 

441 async def health(self) -> HealthStatus: 

442 """Check database connectivity.""" 

443 pass 

444 

445 

446# --------------------------------------------------------------------------- 

447# Tier 1: Wiki Store (M8) 

448# --------------------------------------------------------------------------- 

449 

450 

451@runtime_checkable 

452class WikiStore(Protocol): 

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

454 

455 WikiStore persists structured WikiPage metadata. Vector embeddings of 

456 compiled pages are stored separately in the VectorStore with 

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

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

459 

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

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

462 """ 

463 

464 SPI_VERSION: ClassVar[int] = 1 

465 

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

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

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

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

470 

471 Returns: 

472 The stored ``page_id``. 

473 """ 

474 pass 

475 

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

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

478 

479 Returns: 

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

481 """ 

482 pass 

483 

484 async def list_pages( 

485 self, 

486 bank_id: str, 

487 scope: str | None = None, 

488 kind: str | None = None, 

489 ) -> list[WikiPage]: 

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

491 

492 Args: 

493 bank_id: The bank to list pages for. 

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

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

496 

497 Returns: 

498 All matching pages, unsorted. 

499 """ 

500 pass 

501 

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

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

504 

505 Returns: 

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

507 """ 

508 pass 

509 

510 async def health(self) -> HealthStatus: 

511 """Check storage connectivity.""" 

512 pass 

513 

514 

515# --------------------------------------------------------------------------- 

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

517# --------------------------------------------------------------------------- 

518 

519 

520@runtime_checkable 

521class PageIndexDocumentStore(Protocol): 

522 """Document + section-skeleton persistence (save/load). 

523 

524 A narrow role slice of the section-recall SPI. Consumers should depend on 

525 the narrowest role they use; :class:`PageIndexStore` is the union of all 

526 roles, retained for implementers and existing annotations. 

527 """ 

528 

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

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

531 

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

533 for the same conversation replaces the prior row. 

534 """ 

535 

536 async def save_sections( 

537 self, 

538 document_id: str, 

539 sections: list[PageIndexSection], 

540 ) -> int: 

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

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

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

544 

545 async def load_document( 

546 self, 

547 bank_id: str, 

548 source_id: str, 

549 ) -> PageIndexDocument | None: 

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

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

552 

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

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

555 

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

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

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

559 """ 

560 

561 async def load_sections_with_embeddings( 

562 self, 

563 bank_id: str, 

564 document_id: str, 

565 ) -> list["PageIndexSection"]: 

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

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

568 

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

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

571 caller sorts as needed.""" 

572 

573 

574class PageIndexSectionWriter(Protocol): 

575 """Section enrichment writes (embeddings, entities, links, event dates). 

576 

577 A narrow role slice of the section-recall SPI. Consumers should depend on 

578 the narrowest role they use; :class:`PageIndexStore` is the union of all 

579 roles, retained for implementers and existing annotations. 

580 """ 

581 

582 async def save_section_embeddings( 

583 self, 

584 document_id: str, 

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

586 ) -> int: 

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

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

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

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

591 

592 async def save_section_entities( 

593 self, 

594 entities: list[PageIndexSectionEntity], 

595 ) -> int: 

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

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

598 

599 async def save_section_links( 

600 self, 

601 links: list[PageIndexSectionLink], 

602 ) -> int: 

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

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

605 

606 async def populate_semantic_knn_links( 

607 self, 

608 document_id: str, 

609 *, 

610 top_k: int = 5, 

611 min_similarity: float = 0.5, 

612 ) -> int: 

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

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

615 most-similar OTHER sections in the same document with 

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

617 

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

619 

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

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

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

623 

624 async def save_section_event_dates( 

625 self, 

626 document_id: str, 

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

628 ) -> int: 

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

630 existing section rows. 

631 

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

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

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

635 after section_event_extraction runs.""" 

636 

637 

638class PageIndexSectionSearch(Protocol): 

639 """Section-grain retrieval: parallel search + graph-expand strategies. 

640 

641 A narrow role slice of the section-recall SPI. Consumers should depend on 

642 the narrowest role they use; :class:`PageIndexStore` is the union of all 

643 roles, retained for implementers and existing annotations. 

644 """ 

645 

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

647 # 

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

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

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

651 # tuples; the orchestrator fuses them via RRF. 

652 # 

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

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

655 

656 async def search_sections_semantic( 

657 self, 

658 bank_id: str, 

659 query_embedding: list[float], 

660 *, 

661 top_k: int = 20, 

662 session_filter: str | None = None, 

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

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

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

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

667 

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

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

670 retrieves cross-session. 

671 """ 

672 

673 async def search_sections_keyword( 

674 self, 

675 bank_id: str, 

676 query: str, 

677 *, 

678 top_k: int = 20, 

679 speaker: str | None = None, 

680 document_id: str | None = None, 

681 session_filter: str | None = None, 

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

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

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

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

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

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

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

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

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

691 

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

693 whose ``session_id`` matches.""" 

694 

695 async def search_sections_by_entities( 

696 self, 

697 bank_id: str, 

698 entity_names: list[str], 

699 *, 

700 top_k: int = 20, 

701 session_filter: str | None = None, 

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

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

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

705 count of distinct matching entity_names per section. 

706 

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

708 ``session_id``.""" 

709 

710 async def search_sections_temporal( 

711 self, 

712 bank_id: str, 

713 date_range: tuple[datetime, datetime], 

714 *, 

715 top_k: int = 20, 

716 session_filter: str | None = None, 

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

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

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

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

721 ranker; ranking happens in fusion. 

722 

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

724 ``session_id`` matches.""" 

725 

726 async def expand_section_links( 

727 self, 

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

729 *, 

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

731 top_k: int = 20, 

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

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

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

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

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

737 

738 async def expand_sections_by_shared_entities( 

739 self, 

740 bank_id: str, 

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

742 *, 

743 top_k: int = 20, 

744 exclude_seeds: bool = True, 

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

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

747 expansion through entity co-occurrence in 

748 ``astrocyte_pi_section_entities``. 

749 

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

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

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

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

754 filters the seeds themselves out of the result. 

755 

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

757 on ``astrocyte_pi_section_links`` being populated (semantic_knn 

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

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

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

761 sections mention the same entity. Entity spread uses the 

762 already-populated ``astrocyte_pi_section_entities`` table 

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

764 """ 

765 

766 

767class PageIndexFactStore(Protocol): 

768 """M12.1 atomic fact-grain layer (persist + search facts). 

769 

770 A narrow role slice of the section-recall SPI. Consumers should depend on 

771 the narrowest role they use; :class:`PageIndexStore` is the union of all 

772 roles, retained for implementers and existing annotations. 

773 """ 

774 

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

776 

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

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

779 

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

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

782 inserted. Implementations should validate that 

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

784 let the FK cascade handle deletion.""" 

785 

786 async def update_fact_embeddings( 

787 self, 

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

789 ) -> int: 

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

791 

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

793 Returns the number of rows updated. Separated from 

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

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

796 

797 async def search_facts_semantic( 

798 self, 

799 bank_id: str, 

800 query_embedding: list[float], 

801 *, 

802 top_k: int = 20, 

803 document_id: str | None = None, 

804 fact_type: str | None = None, 

805 session_filter: str | None = None, 

806 ) -> list["PageIndexFactHit"]: 

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

808 

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

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

811 in the Postgres adapter, _matches_session helper in the 

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

813 when ``session_filter`` is set.""" 

814 

815 async def search_facts_by_entity( 

816 self, 

817 bank_id: str, 

818 entity_name: str, 

819 *, 

820 top_k: int = 50, 

821 document_id: str | None = None, 

822 fact_type: str | None = None, 

823 session_filter: str | None = None, 

824 ) -> list["PageIndexFactHit"]: 

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

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

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

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

829 

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

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

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

833 whole point of cross-session traversal is preserved. 

834 

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

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

837 produce independent candidate pools per type.""" 

838 

839 async def search_facts_temporal( 

840 self, 

841 bank_id: str, 

842 date_range: tuple[datetime, datetime], 

843 *, 

844 top_k: int = 50, 

845 document_id: str | None = None, 

846 fact_type: str | None = None, 

847 session_filter: str | None = None, 

848 ) -> list["PageIndexFactHit"]: 

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

850 date range. Powers temporal-arithmetic queries that previously 

851 had to resort to keyword search over section summaries. 

852 

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

854 temporal facts. 

855 

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

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

858 

859 async def search_facts_keyword( 

860 self, 

861 bank_id: str, 

862 query: str, 

863 *, 

864 top_k: int = 20, 

865 document_id: str | None = None, 

866 fact_type: str | None = None, 

867 session_filter: str | None = None, 

868 ) -> list["PageIndexFactHit"]: 

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

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

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

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

873 embeddings tend to under-weight relative to thematically 

874 similar but topically off content. 

875 

876 Implementations: 

877 

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

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

880 no schema migration needed). Mirrors the 

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

882 since M9. 

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

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

885 

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

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

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

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

890 handles the scale difference). 

891 

892 Diagnostic that motivated this: SSU bench failures showed the 

893 cross-encoder rerank can rank a specific factoid statement 

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

895 candidate pool while surfacing thematically-related but 

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

897 BM25 catches the literal-keyword match that semantic alone 

898 misses.""" 

899 

900 async def count_facts_matching( 

901 self, 

902 bank_id: str, 

903 document_id: str, 

904 *, 

905 entity_pattern: str | None = None, 

906 fact_type: str | None = None, 

907 ) -> int: 

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

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

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

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

912 doctors did I visit?" — agent calls with 

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

914 count.""" 

915 

916 

917class PageIndexWikiStore(Protocol): 

918 """Compiled wiki-page persistence and lookup. 

919 

920 A narrow role slice of the section-recall SPI. Consumers should depend on 

921 the narrowest role they use; :class:`PageIndexStore` is the union of all 

922 roles, retained for implementers and existing annotations. 

923 """ 

924 

925 async def save_wiki_page( 

926 self, 

927 *, 

928 page: "WikiPage", 

929 embedding: list[float] | None, 

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

931 ) -> str: 

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

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

934 recall strategy. 

935 

936 Implementations must: 

937 - Insert / upsert the wiki_page row. 

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

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

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

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

942 embedding (caller does its own deferred indexing). 

943 

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

945 

946 async def search_wiki_pages_semantic( 

947 self, 

948 bank_id: str, 

949 query_embedding: list[float], 

950 *, 

951 top_k: int = 5, 

952 document_id: str | None = None, 

953 ) -> list["WikiPageHit"]: 

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

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

956 

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

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

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

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

961 

962 async def count_wiki_pages_for_doc( 

963 self, 

964 bank_id: str, 

965 document_id: str, 

966 ) -> int: 

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

968 ``compile_sections_for_document`` skips work when this returns 

969 > 0.""" 

970 

971 async def list_wiki_pages_for_doc( 

972 self, 

973 bank_id: str, 

974 document_id: str, 

975 ) -> list["WikiPage"]: 

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

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

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

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

980 bumped revisions. 

981 

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

983 yet implement wiki enumeration — the revision pass becomes a 

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

985 return [] 

986 

987 async def list_wikis_affected_by_entities( 

988 self, 

989 bank_id: str, 

990 entities: list[str], 

991 *, 

992 min_overlap: int = 1, 

993 limit: int = 8, 

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

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

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

997 WikiPage rows + overlap stats. 

998 

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

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

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

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

1003 round-trip per result. 

1004 

1005 Args: 

1006 bank_id: tenant scope. 

1007 entities: distinct entity names from the new source. 

1008 min_overlap: minimum shared entities to flag a wiki as 

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

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

1011 

1012 Returns: 

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

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

1015 ``page_id`` ascending for stability. 

1016 

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

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

1019 backends rather than crashing. 

1020 """ 

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

1022 return [] 

1023 

1024 async def list_distinct_entities( 

1025 self, 

1026 bank_id: str, 

1027 document_id: str, 

1028 *, 

1029 pattern: str | None = None, 

1030 limit: int = 50, 

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

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

1033 document with their per-section mention counts. 

1034 

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

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

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

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

1039 know what to look for). 

1040 

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

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

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

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

1045 at most once). 

1046 

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

1048 across-session mentions deterministically — fixes the LME 

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

1050 couldn't reliably aggregate entity mentions from raw section 

1051 text.""" 

1052 

1053 

1054class PageIndexLifecycle(Protocol): 

1055 """SPI versioning + health. 

1056 

1057 A narrow role slice of the section-recall SPI. Consumers should depend on 

1058 the narrowest role they use; :class:`PageIndexStore` is the union of all 

1059 roles, retained for implementers and existing annotations. 

1060 """ 

1061 

1062 SPI_VERSION: ClassVar[int] = 1 

1063 

1064 async def health(self) -> HealthStatus: 

1065 """Check storage connectivity.""" 

1066 

1067 

1068# --------------------------------------------------------------------------- 

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

1070# --------------------------------------------------------------------------- 

1071 

1072 

1073@runtime_checkable 

1074 

1075 

1076class PageIndexStore( 

1077 PageIndexDocumentStore, 

1078 PageIndexSectionWriter, 

1079 PageIndexSectionSearch, 

1080 PageIndexFactStore, 

1081 PageIndexWikiStore, 

1082 PageIndexLifecycle, 

1083 Protocol, 

1084): 

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

1086 

1087 Backs the three-layer recall stack defined in 

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

1089 

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

1091 canonical markdown that the picker slices for synth excerpts. 

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

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

1094 section grain (PR2 commit A populates). 

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

1096 section grain (PR2 commit D populates). 

1097 

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

1099 ``save_sections`` / ``load_skeleton`` / ``load_document`` — the 

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

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

1102 PR2. 

1103 

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

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

1106 the production backend. 

1107 """ 

1108 

1109 

1110class MentalModelStore(Protocol):