Coverage for astrocyte/testing/in_memory.py: 85%
1113 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-10-29 22:10 +0000
1"""In-memory provider implementations for testing.
3These are fully functional providers backed by Python dicts/lists.
4Used by conformance tests and integration tests.
5"""
7from __future__ import annotations
9import math
10import uuid
11from datetime import UTC, datetime
12from typing import ClassVar
14from astrocyte.types import (
15 Completion,
16 Document,
17 DocumentFilters,
18 DocumentHit,
19 EngineCapabilities,
20 Entity,
21 EntityCandidateMatch,
22 EntityLink,
23 ForgetRequest,
24 ForgetResult,
25 GraphHit,
26 HealthStatus,
27 LegalHold,
28 LLMCapabilities,
29 MemoryChange,
30 MemoryEntityAssociation,
31 MemoryHit,
32 MentalModel,
33 Message,
34 PageIndexDocument,
35 PageIndexFact,
36 PageIndexSection,
37 PageIndexSectionEntity,
38 PageIndexSectionLink,
39 RecallRequest,
40 RecallResult,
41 ReflectRequest,
42 ReflectResult,
43 RetainRequest,
44 RetainResult,
45 SourceChunk,
46 SourceDocument,
47 TokenUsage,
48 VectorFilters,
49 VectorHit,
50 VectorItem,
51 WikiPage,
52)
55def _utc(dt: datetime) -> datetime:
56 """Naive datetimes are UTC, as the SQL stores' ``timestamptz`` treats them."""
57 return dt.replace(tzinfo=UTC) if dt.tzinfo is None else dt
60def _cosine_sim(a: list[float], b: list[float]) -> float:
61 dot = sum(x * y for x, y in zip(a, b))
62 na = math.sqrt(sum(x * x for x in a))
63 nb = math.sqrt(sum(x * x for x in b))
64 if na == 0 or nb == 0:
65 return 0.0
66 return dot / (na * nb)
69# ---------------------------------------------------------------------------
70# In-memory Vector Store
71# ---------------------------------------------------------------------------
74class InMemoryVectorStore:
75 """Fully functional in-memory vector store for testing."""
77 SPI_VERSION: ClassVar[int] = 1
79 def __init__(self) -> None:
80 self._vectors: dict[str, VectorItem] = {}
81 # Change-feed bookkeeping (``list_changes``). Deletes stay hard in
82 # ``_vectors``; a forgotten id is remembered here as a tombstone
83 # ``(bank_id, forgotten_at)`` the way the SQL stores keep the row with
84 # ``forgotten_at`` set. ``_retained_at`` is each live row's effective
85 # ``retained_at`` (the item's own, else the time it was stored) and
86 # ``_changed_at`` its last change to any synced field.
87 self._tombstones: dict[str, tuple[str, datetime]] = {}
88 self._legal_holds: dict[tuple[str, str], LegalHold] = {}
89 self._retained_at: dict[str, datetime] = {}
90 self._changed_at: dict[str, datetime] = {}
92 async def store_vectors(self, items: list[VectorItem]) -> list[str]:
93 ids = []
94 now = datetime.now(UTC)
95 for item in items:
96 existed = item.id in self._vectors or item.id in self._tombstones
97 self._vectors[item.id] = item
98 # Upsert semantics match the SQL stores: a re-stored id is live
99 # again (their ``forgotten_at = NULL``), and overwriting an
100 # existing row is a change now, whatever its retained_at.
101 self._tombstones.pop(item.id, None)
102 retained_at = _utc(item.retained_at) if item.retained_at else now
103 self._retained_at[item.id] = retained_at
104 self._changed_at[item.id] = max(retained_at, now) if existed else retained_at
105 ids.append(item.id)
106 return ids
108 async def search_similar(
109 self,
110 query_vector: list[float],
111 bank_id: str,
112 limit: int = 10,
113 filters: VectorFilters | None = None,
114 ) -> list[VectorHit]:
115 results: list[tuple[float, VectorItem]] = []
116 for item in self._vectors.values():
117 if item.bank_id != bank_id:
118 continue
119 if filters:
120 if filters.tags and item.tags:
121 if not set(filters.tags) & set(item.tags):
122 continue
123 if filters.fact_types and item.fact_type:
124 if item.fact_type not in filters.fact_types:
125 continue
126 # M9: time-travel filter — exclude items retained after as_of
127 if filters.as_of is not None and item.retained_at is not None:
128 if item.retained_at > filters.as_of:
129 continue
130 if filters.time_range is not None and item.occurred_at is not None:
131 start, end = filters.time_range
132 if item.occurred_at < start or item.occurred_at > end:
133 continue
134 # M31 Fix 2 — session_id filter. ``metadata['session_id']``
135 # is the canonical location (preserves the same shape
136 # production callers use when calling retain()). When the
137 # filter is set but the item has no session_id metadata,
138 # the item is excluded (consistent with the
139 # PageIndexStore behaviour where top-level facts without
140 # an anchoring section are excluded from session-filtered
141 # results).
142 if filters.session_id is not None:
143 item_session = None
144 if item.metadata:
145 item_session = item.metadata.get("session_id") if isinstance(item.metadata, dict) else None
146 if item_session != filters.session_id:
147 continue
148 sim = _cosine_sim(query_vector, item.vector)
149 results.append((sim, item))
151 results.sort(key=lambda x: x[0], reverse=True)
152 return [
153 VectorHit(
154 id=item.id,
155 text=item.text,
156 score=sim,
157 metadata=item.metadata,
158 tags=item.tags,
159 fact_type=item.fact_type,
160 occurred_at=item.occurred_at,
161 memory_layer=item.memory_layer,
162 retained_at=item.retained_at, # M9
163 chunk_id=item.chunk_id, # M10
164 )
165 for sim, item in results[:limit]
166 ]
168 async def delete(self, ids: list[str], bank_id: str) -> int:
169 count = 0
170 now = datetime.now(UTC)
171 for vid in ids:
172 if vid in self._vectors and self._vectors[vid].bank_id == bank_id:
173 del self._vectors[vid]
174 self._changed_at.pop(vid, None)
175 retained_at = self._retained_at.pop(vid, now)
176 self._tombstones[vid] = (bank_id, max(retained_at, now))
177 count += 1
178 return count
180 async def save_legal_hold(self, hold: LegalHold) -> None:
181 self._legal_holds[(hold.bank_id, hold.hold_id)] = hold
183 async def delete_legal_hold(self, bank_id: str, hold_id: str) -> bool:
184 return self._legal_holds.pop((bank_id, hold_id), None) is not None
186 async def list_legal_holds(self, bank_id: str) -> list[LegalHold]:
187 return sorted((h for (b, _), h in self._legal_holds.items() if b == bank_id), key=lambda h: h.hold_id)
189 async def erase(self, bank_id: str, ids: list[str]) -> int:
190 """Erase forgotten memories, keeping their tombstones. A forget here
191 already keeps nothing but the tombstone, so this only counts them."""
192 return sum(1 for vid in dict.fromkeys(ids) if self._tombstones.get(vid, ("",))[0] == bank_id)
194 async def list_changes(
195 self,
196 bank_id: str,
197 *,
198 after: tuple[datetime, str] | None = None,
199 limit: int = 100,
200 ) -> list[MemoryChange]:
201 """Every change to a synced row: live rows (current values) and
202 tombstones, ordered by ``(changed_at, id)``, strictly after ``after``
203 (see ``VectorStore`` optional methods)."""
204 changes = [self._live_change(vid, item) for vid, item in self._vectors.items() if item.bank_id == bank_id]
205 for vid, (tomb_bank, forgotten_at) in self._tombstones.items():
206 if tomb_bank == bank_id:
207 changes.append(MemoryChange(id=vid, bank_id=bank_id, changed_at=forgotten_at, deleted=True))
208 if after is not None:
209 position = (_utc(after[0]), after[1])
210 changes = [c for c in changes if (c.changed_at, c.id) > position]
211 changes.sort(key=lambda c: (c.changed_at, c.id))
212 return changes[: max(limit, 0)]
214 def _live_change(self, vid: str, item: VectorItem) -> MemoryChange:
215 now = datetime.now(UTC) # only for rows placed in _vectors directly
216 retained_at = self._retained_at.get(vid) or now
217 return MemoryChange(
218 id=vid,
219 bank_id=item.bank_id,
220 changed_at=self._changed_at.get(vid) or retained_at,
221 text=item.text,
222 occurred_at=item.occurred_at,
223 retained_at=retained_at,
224 tags=list(item.tags) if item.tags else None,
225 fact_type=item.fact_type,
226 memory_layer=item.memory_layer,
227 metadata=item.metadata,
228 )
230 async def lookup_ids(self, ids: list[str]) -> list[MemoryChange]:
231 """The current state of each id held in any bank: a live row or a
232 tombstone. Unknown ids are absent (see ``VectorStore`` optional methods)."""
233 found: list[MemoryChange] = []
234 for vid in dict.fromkeys(ids):
235 if vid in self._vectors:
236 found.append(self._live_change(vid, self._vectors[vid]))
237 elif vid in self._tombstones:
238 tomb_bank, forgotten_at = self._tombstones[vid]
239 found.append(MemoryChange(id=vid, bank_id=tomb_bank, changed_at=forgotten_at, deleted=True))
240 return found
242 async def insert_vectors(self, items: list[VectorItem]) -> list[str]:
243 """Store items whose id is new to the store; never overwrite. Returns
244 the ids inserted (see ``VectorStore`` optional methods)."""
245 inserted: list[str] = []
246 for item in items:
247 if item.id in self._vectors or item.id in self._tombstones:
248 continue
249 await self.store_vectors([item])
250 inserted.append(item.id)
251 return inserted
253 async def get_by_chunk_ids(
254 self,
255 chunk_ids: list[str],
256 bank_id: str,
257 ) -> list[VectorHit]:
258 """M10 chunk expansion: return all vectors whose ``chunk_id`` is in the list.
260 Used by the recall pipeline's chunk-expansion stage when
261 ``source_aware_retrieval.chunk_expansion`` is on. Score is set
262 to ``1.0`` (the caller applies the expansion multiplier);
263 ``chunk_id`` round-trips so callers can group results by source.
264 """
265 if not chunk_ids:
266 return []
267 wanted = set(chunk_ids)
268 hits: list[VectorHit] = []
269 for item in self._vectors.values():
270 if item.bank_id != bank_id:
271 continue
272 if item.chunk_id is None or item.chunk_id not in wanted:
273 continue
274 hits.append(
275 VectorHit(
276 id=item.id,
277 text=item.text,
278 score=1.0,
279 metadata=item.metadata,
280 tags=item.tags,
281 fact_type=item.fact_type,
282 occurred_at=item.occurred_at,
283 memory_layer=item.memory_layer,
284 retained_at=item.retained_at,
285 chunk_id=item.chunk_id,
286 )
287 )
288 return hits
290 async def list_vectors(
291 self,
292 bank_id: str,
293 offset: int = 0,
294 limit: int = 100,
295 ) -> list[VectorItem]:
296 bank_items = sorted(
297 (v for v in self._vectors.values() if v.bank_id == bank_id),
298 key=lambda v: v.id,
299 )
300 return bank_items[offset : offset + limit]
302 async def health(self) -> HealthStatus:
303 return HealthStatus(healthy=True, message="in-memory vector store")
306# ---------------------------------------------------------------------------
307# In-memory Graph Store
308# ---------------------------------------------------------------------------
311class InMemoryGraphStore:
312 """Fully functional in-memory graph store for testing."""
314 SPI_VERSION: ClassVar[int] = 1
316 def __init__(self) -> None:
317 self._entities: dict[str, dict[str, Entity]] = {} # bank_id → {entity_id → Entity}
318 self._links: dict[str, list[EntityLink]] = {} # bank_id → links
319 self._memory_entity_map: dict[str, list[MemoryEntityAssociation]] = {} # bank_id → assocs
320 # Memory-to-memory links (Hindsight parity) — caused_by, semantic, etc.
321 self._memory_links: dict[str, list] = {} # bank_id → list[MemoryLink]
323 async def store_entities(self, entities: list[Entity], bank_id: str) -> list[str]:
324 if bank_id not in self._entities:
325 self._entities[bank_id] = {}
326 ids = []
327 for entity in entities:
328 self._entities[bank_id][entity.id] = entity
329 ids.append(entity.id)
330 return ids
332 async def store_links(self, links: list[EntityLink], bank_id: str) -> list[str]:
333 if bank_id not in self._links:
334 self._links[bank_id] = []
335 ids = []
336 for link in links:
337 lid = uuid.uuid4().hex[:12]
338 self._links[bank_id].append(link)
339 ids.append(lid)
340 return ids
342 async def link_memories_to_entities(
343 self,
344 associations: list[MemoryEntityAssociation],
345 bank_id: str,
346 ) -> None:
347 if bank_id not in self._memory_entity_map:
348 self._memory_entity_map[bank_id] = []
349 self._memory_entity_map[bank_id].extend(associations)
351 async def query_neighbors(
352 self,
353 entity_ids: list[str],
354 bank_id: str,
355 max_depth: int = 2,
356 limit: int = 20,
357 ) -> list[GraphHit]:
358 # Find memories linked to these entities within this bank.
359 # Per-memory ``connected_entities`` is the SUBSET of entity_ids
360 # that this memory is associated with — spreading activation uses
361 # this to decide which activated entity surfaced the memory.
362 per_memory: dict[str, set[str]] = {}
363 for assoc in self._memory_entity_map.get(bank_id, []):
364 if assoc.entity_id in entity_ids:
365 per_memory.setdefault(assoc.memory_id, set()).add(assoc.entity_id)
367 return [
368 GraphHit(
369 memory_id=mid,
370 text=f"[graph result for {mid}]",
371 connected_entities=sorted(connected),
372 depth=1,
373 score=0.5,
374 )
375 for mid, connected in list(per_memory.items())[:limit]
376 ]
378 async def query_entities(self, query: str, bank_id: str, limit: int = 10) -> list[Entity]:
379 query_lower = query.lower()
380 bank_entities = self._entities.get(bank_id, {})
381 results = [e for e in bank_entities.values() if query_lower in e.name.lower()]
382 return results[:limit]
384 async def find_entity_candidates(
385 self,
386 name: str,
387 bank_id: str,
388 threshold: float = 0.8,
389 limit: int = 5,
390 ) -> list[Entity]:
391 """Return entities whose name contains *name* as a substring (case-insensitive).
393 The in-memory implementation uses substring overlap as a proxy for
394 similarity; production adapters use vector or edit-distance similarity.
395 The *threshold* parameter is accepted for interface compatibility but
396 ignored — substring match is all-or-nothing.
397 """
398 name_lower = name.lower()
399 bank_entities = self._entities.get(bank_id, {})
400 results = [e for e in bank_entities.values() if name_lower in e.name.lower() or e.name.lower() in name_lower]
401 return results[:limit]
403 async def find_entity_candidates_scored(
404 self,
405 name: str,
406 bank_id: str,
407 *,
408 name_embedding: list[float] | None = None,
409 trigram_threshold: float = 0.3,
410 limit: int = 10,
411 ) -> list[EntityCandidateMatch]:
412 """In-memory equivalent of pg_trgm + cosine candidate scoring.
414 - ``name_similarity`` is computed with :class:`difflib.SequenceMatcher`
415 on lowercased names, which approximates PostgreSQL's pg_trgm
416 ``similarity()`` closely enough for tests.
417 - ``embedding_similarity`` is computed with the module's
418 :func:`_cosine_sim` against the candidate's stored embedding when
419 both sides have one; ``None`` otherwise.
420 - ``co_occurring_names`` are derived from the in-memory link list —
421 for each ``co_occurs`` link involving the candidate, the OTHER
422 entity's lowercased name is collected.
423 - ``last_seen`` mirrors PostgreSQL's ``updated_at`` semantics —
424 the in-memory store doesn't track update timestamps, so this is
425 ``None`` unless the entity has a ``last_seen`` in metadata.
427 Candidates with ``name_similarity < trigram_threshold`` are dropped.
428 Results are ordered by ``max(name_similarity, embedding_similarity or 0)``
429 descending.
430 """
431 from difflib import SequenceMatcher
433 bank_entities = self._entities.get(bank_id, {})
434 if not bank_entities:
435 return []
437 name_norm = (name or "").strip().lower()
438 if not name_norm:
439 return []
441 # Pre-build a map: candidate_id → list of co-occurring entity names.
442 # Walks the bank's links once; the resolver only iterates a
443 # pre-filtered candidate slice afterward.
444 cooccurrence_map: dict[str, list[str]] = {}
445 for link in self._links.get(bank_id, []):
446 if link.link_type != "co_occurs":
447 continue
448 other_a = bank_entities.get(link.entity_b)
449 other_b = bank_entities.get(link.entity_a)
450 if other_a is not None:
451 cooccurrence_map.setdefault(link.entity_a, []).append((other_a.name or "").strip().lower())
452 if other_b is not None:
453 cooccurrence_map.setdefault(link.entity_b, []).append((other_b.name or "").strip().lower())
455 scored: list[EntityCandidateMatch] = []
456 for entity in bank_entities.values():
457 cand_name = (entity.name or "").strip().lower()
458 if not cand_name:
459 continue
460 name_sim = SequenceMatcher(None, name_norm, cand_name).ratio()
461 if name_sim < trigram_threshold:
462 continue
464 emb_sim: float | None = None
465 if name_embedding is not None and entity.embedding is not None:
466 emb_sim = _cosine_sim(name_embedding, entity.embedding)
467 # Clamp to [0, 1] — _cosine_sim can return negative values
468 # for vectors that aren't non-negative.
469 emb_sim = max(0.0, min(1.0, emb_sim))
471 co_names = cooccurrence_map.get(entity.id, [])
472 last_seen = None
473 if entity.metadata and "last_seen" in entity.metadata:
474 # InMemoryGraphStore doesn't auto-stamp last_seen — tests can
475 # populate via metadata to exercise the temporal tier.
476 raw = entity.metadata["last_seen"]
477 if isinstance(raw, str):
478 try:
479 last_seen = datetime.fromisoformat(raw)
480 except ValueError:
481 last_seen = None
483 scored.append(
484 EntityCandidateMatch(
485 entity=entity,
486 name_similarity=name_sim,
487 embedding_similarity=emb_sim,
488 co_occurring_names=co_names,
489 last_seen=last_seen,
490 mention_count=getattr(entity, "mention_count", 1),
491 )
492 )
494 def _sort_key(m: EntityCandidateMatch) -> float:
495 return max(m.name_similarity, m.embedding_similarity or 0.0)
497 scored.sort(key=_sort_key, reverse=True)
498 return scored[:limit]
500 async def store_memory_links(
501 self,
502 links: list, # list[MemoryLink] — kept loose to avoid forward-ref noise
503 bank_id: str,
504 ) -> list[str]:
505 """Persist memory-to-memory links (Hindsight-parity)."""
506 store = self._memory_links.setdefault(bank_id, [])
507 ids: list[str] = []
508 for link in links:
509 lid = uuid.uuid4().hex[:12]
510 store.append(link)
511 ids.append(lid)
512 return ids
514 async def find_memory_links(
515 self,
516 seed_memory_ids: list[str],
517 bank_id: str,
518 *,
519 link_types: list[str] | None = None,
520 limit: int = 200,
521 ) -> list: # list[MemoryLink]
522 """Find links touching any seed memory in either direction."""
523 if not seed_memory_ids:
524 return []
525 seeds = set(seed_memory_ids)
526 type_filter = set(link_types) if link_types else None
527 out: list = []
528 for link in self._memory_links.get(bank_id, []):
529 if type_filter is not None and link.link_type not in type_filter:
530 continue
531 if link.source_memory_id in seeds or link.target_memory_id in seeds:
532 out.append(link)
533 if len(out) >= limit:
534 break
535 return out
537 async def get_entity_ids_for_memories(
538 self,
539 memory_ids: list[str],
540 bank_id: str,
541 ) -> dict[str, list[str]]:
542 """Return ``{memory_id: [entity_id, ...]}`` from the association table.
544 Used by the spreading-activation pipeline to seed entity IDs
545 directly from memory↔entity associations (the most-accurate
546 path; metadata-based extraction is the fallback).
547 """
548 if not memory_ids:
549 return {}
550 target = set(memory_ids)
551 out: dict[str, list[str]] = {}
552 for assoc in self._memory_entity_map.get(bank_id, []):
553 if assoc.memory_id in target:
554 out.setdefault(assoc.memory_id, []).append(assoc.entity_id)
555 return out
557 async def expand_entities_via_links(
558 self,
559 entity_ids: list[str],
560 bank_id: str,
561 *,
562 max_hops: int = 2,
563 link_types: list[str] | None = None,
564 ) -> dict[str, int]:
565 """BFS over entity-to-entity links to ``max_hops`` distance.
567 Returns ``{entity_id: hop_distance}`` where ``hop_distance == 0``
568 for entities in the input set, ``1`` for direct neighbors, etc.
570 Used by the spreading-activation pipeline to walk
571 ``co_occurs`` links from seed entities to their graph
572 neighborhood. Mirrors the AGE adapter's recursive-CTE
573 implementation so behavior is consistent across stores.
574 """
575 if not entity_ids:
576 return {}
577 accepted_types = set(link_types or ["co_occurs"])
578 bank_links = self._links.get(bank_id, [])
580 # Pre-build a neighbor map keyed by entity_id for O(1) lookup
581 # during the BFS frontier expansion.
582 neighbors: dict[str, set[str]] = {}
583 for link in bank_links:
584 if link.link_type not in accepted_types:
585 continue
586 neighbors.setdefault(link.entity_a, set()).add(link.entity_b)
587 neighbors.setdefault(link.entity_b, set()).add(link.entity_a)
589 distances: dict[str, int] = {eid: 0 for eid in entity_ids}
590 frontier = set(entity_ids)
591 for hop in range(1, max(1, max_hops) + 1):
592 next_frontier: set[str] = set()
593 for eid in frontier:
594 for nb in neighbors.get(eid, ()):
595 if nb not in distances:
596 distances[nb] = hop
597 next_frontier.add(nb)
598 if not next_frontier:
599 break
600 frontier = next_frontier
601 return distances
603 async def increment_mention_counts(
604 self,
605 entity_ids: list[str],
606 bank_id: str,
607 ) -> None:
608 """Bump ``mention_count`` for each canonical entity ID by 1.
610 Mirrors the AGE adapter's :meth:`increment_mention_counts` for
611 in-memory tests of the resolver's mention-count flow.
612 """
613 bank_entities = self._entities.get(bank_id, {})
614 for eid in entity_ids:
615 entity = bank_entities.get(eid)
616 if entity is None:
617 continue
618 # Entity is a frozen-ish dataclass — replace via dataclasses.replace
619 # to avoid mutating the original object held elsewhere.
620 from dataclasses import replace as _dc_replace
622 bank_entities[eid] = _dc_replace(
623 entity,
624 mention_count=int(getattr(entity, "mention_count", 1)) + 1,
625 )
627 async def store_entity_link(self, link: EntityLink, bank_id: str) -> str:
628 """Store a single resolved entity link (M11 entity resolution)."""
629 if bank_id not in self._links:
630 self._links[bank_id] = []
631 lid = uuid.uuid4().hex[:12]
632 self._links[bank_id].append(link)
633 return lid
635 async def health(self) -> HealthStatus:
636 return HealthStatus(healthy=True, message="in-memory graph store")
639# ---------------------------------------------------------------------------
640# In-memory Document Store
641# ---------------------------------------------------------------------------
644class InMemoryDocumentStore:
645 """Fully functional in-memory document store for testing."""
647 SPI_VERSION: ClassVar[int] = 1
649 def __init__(self) -> None:
650 self._docs: dict[str, tuple[Document, str]] = {} # id -> (doc, bank_id)
652 async def store_document(self, document: Document, bank_id: str) -> str:
653 self._docs[document.id] = (document, bank_id)
654 return document.id
656 async def search_fulltext(
657 self,
658 query: str,
659 bank_id: str,
660 limit: int = 10,
661 filters: DocumentFilters | None = None,
662 ) -> list[DocumentHit]:
663 query_terms = set(query.lower().split())
664 results: list[tuple[float, Document]] = []
665 for doc, bid in self._docs.values():
666 if bid != bank_id:
667 continue
668 doc_terms = set(doc.text.lower().split())
669 overlap = len(query_terms & doc_terms)
670 if overlap > 0:
671 score = overlap / max(len(query_terms), 1)
672 results.append((score, doc))
674 results.sort(key=lambda x: x[0], reverse=True)
675 return [
676 DocumentHit(document_id=doc.id, text=doc.text, score=score, metadata=doc.metadata)
677 for score, doc in results[:limit]
678 ]
680 async def get_document(self, document_id: str, bank_id: str) -> Document | None:
681 entry = self._docs.get(document_id)
682 if entry and entry[1] == bank_id:
683 return entry[0]
684 return None
686 async def health(self) -> HealthStatus:
687 return HealthStatus(healthy=True, message="in-memory document store")
690# ---------------------------------------------------------------------------
691# In-memory Engine Provider
692# ---------------------------------------------------------------------------
695class InMemoryEngineProvider:
696 """Fully functional in-memory engine provider for testing."""
698 SPI_VERSION: ClassVar[int] = 1
700 def __init__(self, supports_reflect: bool = True, supports_forget: bool = True) -> None:
701 self._memories: dict[str, list[MemoryHit]] = {} # bank_id -> memories
702 self._supports_reflect = supports_reflect
703 self._supports_forget = supports_forget
705 def capabilities(self) -> EngineCapabilities:
706 return EngineCapabilities(
707 supports_reflect=self._supports_reflect,
708 supports_forget=self._supports_forget,
709 supports_semantic_search=True,
710 supports_tags=True,
711 supports_metadata=True,
712 )
714 async def health(self) -> HealthStatus:
715 return HealthStatus(healthy=True, message="in-memory engine")
717 async def retain(self, request: RetainRequest) -> RetainResult:
718 if request.bank_id not in self._memories:
719 self._memories[request.bank_id] = []
720 mem_id = uuid.uuid4().hex[:16]
721 # Stamp _created_at so MIP forget min_age_days can be enforced. Existing
722 # callers that already set _created_at (eg. import/replay) win.
723 meta = dict(request.metadata) if request.metadata else {}
724 meta.setdefault("_created_at", datetime.now(UTC).isoformat())
725 self._memories[request.bank_id].append(
726 MemoryHit(
727 text=request.content,
728 score=1.0,
729 fact_type="world",
730 metadata=meta,
731 tags=request.tags,
732 occurred_at=request.occurred_at,
733 source=request.source,
734 memory_id=mem_id,
735 bank_id=request.bank_id,
736 )
737 )
738 return RetainResult(stored=True, memory_id=mem_id)
740 async def recall(self, request: RecallRequest) -> RecallResult:
741 memories = self._memories.get(request.bank_id, [])
743 # Apply filters before scoring
744 # MIP soft-delete: hide records flagged with `_deleted: true` in metadata.
745 # The flag is set by `soft_delete()` when forget.mode == "soft" so that
746 # records survive on disk for audit/restore but disappear from recall.
747 filtered = [m for m in memories if not (m.metadata and m.metadata.get("_deleted") is True)]
748 if request.tags:
749 tag_set = set(request.tags)
750 filtered = [m for m in filtered if m.tags and tag_set & set(m.tags)]
751 if request.fact_types:
752 ft_set = set(request.fact_types)
753 filtered = [m for m in filtered if m.fact_type in ft_set]
754 if request.time_range:
755 t_start, t_end = request.time_range
756 filtered = [m for m in filtered if m.occurred_at and t_start <= m.occurred_at <= t_end]
758 # Wildcard query returns all memories (used by export)
759 if request.query.strip() == "*":
760 scored = [(1.0, mem) for mem in filtered]
761 else:
762 # Simple keyword matching
763 query_terms = set(request.query.lower().split())
764 scored = []
765 for mem in filtered:
766 mem_terms = set(mem.text.lower().split())
767 overlap = len(query_terms & mem_terms)
768 if overlap > 0:
769 score = overlap / max(len(query_terms), 1)
770 scored.append((score, mem))
772 scored.sort(key=lambda x: x[0], reverse=True)
773 hits = [
774 MemoryHit(
775 text=mem.text,
776 score=score,
777 fact_type=mem.fact_type,
778 metadata=mem.metadata,
779 tags=mem.tags,
780 occurred_at=mem.occurred_at,
781 source=mem.source,
782 memory_id=mem.memory_id,
783 bank_id=request.bank_id,
784 )
785 for score, mem in scored[: request.max_results]
786 ]
787 return RecallResult(hits=hits, total_available=len(scored), truncated=len(scored) > request.max_results)
789 async def reflect(self, request: ReflectRequest) -> ReflectResult:
790 if not self._supports_reflect:
791 raise NotImplementedError("reflect not supported")
792 recall_result = await self.recall(RecallRequest(query=request.query, bank_id=request.bank_id))
793 answer = "Synthesis: " + "; ".join(h.text for h in recall_result.hits[:5])
794 return ReflectResult(answer=answer, sources=recall_result.hits[:5])
796 async def soft_delete(self, bank_id: str, memory_ids: list[str]) -> int:
797 """Mark records as soft-deleted by setting ``_deleted: true`` in metadata.
799 Used by MIP forget when ``forget.mode == "soft"``: records remain in
800 the store (auditable, restorable) but are hidden from recall.
801 Returns the number of records that were transitioned to deleted.
802 """
803 if not memory_ids:
804 return 0
805 ids_set = set(memory_ids)
806 count = 0
807 for m in self._memories.get(bank_id, []):
808 if m.memory_id in ids_set:
809 meta = dict(m.metadata) if m.metadata else {}
810 if meta.get("_deleted") is True:
811 continue # already soft-deleted
812 meta["_deleted"] = True
813 m.metadata = meta
814 count += 1
815 return count
817 async def forget(self, request: ForgetRequest) -> ForgetResult:
818 if not self._supports_forget:
819 raise NotImplementedError("forget not supported")
820 bank_memories = self._memories.get(request.bank_id, [])
822 if request.scope == "all":
823 count = len(bank_memories)
824 self._memories[request.bank_id] = []
825 return ForgetResult(deleted_count=count)
827 if request.memory_ids:
828 ids_set = set(request.memory_ids)
829 before = len(bank_memories)
830 bank_memories = [m for m in bank_memories if m.memory_id not in ids_set]
831 self._memories[request.bank_id] = bank_memories
832 return ForgetResult(deleted_count=before - len(bank_memories))
834 # Tag-based and/or date-based deletion
835 if request.tags or request.before_date:
836 before = len(bank_memories)
837 tag_set = set(request.tags) if request.tags else None
839 def _should_keep(m: MemoryHit) -> bool:
840 if tag_set and m.tags and tag_set & set(m.tags):
841 return False
842 if request.before_date and m.occurred_at and m.occurred_at < request.before_date:
843 return False
844 return True
846 bank_memories = [m for m in bank_memories if _should_keep(m)]
847 self._memories[request.bank_id] = bank_memories
848 return ForgetResult(deleted_count=before - len(bank_memories))
850 return ForgetResult(deleted_count=0)
853# ---------------------------------------------------------------------------
854# In-memory Wiki Store (M8)
855# ---------------------------------------------------------------------------
858class InMemoryWikiStore:
859 """Fully functional in-memory wiki store for testing (M8).
861 Stores current-revision pages and an audit log of past revisions.
862 Vector embeddings of pages are managed separately by the VectorStore
863 (stored with ``memory_layer="compiled"``); this store holds only the
864 structured WikiPage metadata.
865 """
867 SPI_VERSION: ClassVar[int] = 1
869 def __init__(self) -> None:
870 # bank_id → {page_id → WikiPage} (current revisions)
871 self._pages: dict[str, dict[str, WikiPage]] = {}
872 # bank_id → {page_id → list[WikiPage]} (past revisions, newest last)
873 self._history: dict[str, dict[str, list[WikiPage]]] = {}
875 async def upsert_page(self, page: WikiPage, bank_id: str) -> str:
876 if bank_id not in self._pages:
877 self._pages[bank_id] = {}
878 self._history[bank_id] = {}
880 existing = self._pages[bank_id].get(page.page_id)
881 if existing is not None:
882 # Archive current revision before replacing
883 self._history[bank_id].setdefault(page.page_id, []).append(existing)
884 # Increment revision on the incoming page
885 from dataclasses import replace as _replace
887 page = _replace(page, revision=existing.revision + 1)
889 self._pages[bank_id][page.page_id] = page
890 return page.page_id
892 async def get_page(self, page_id: str, bank_id: str) -> WikiPage | None:
893 return self._pages.get(bank_id, {}).get(page_id)
895 async def list_pages(
896 self,
897 bank_id: str,
898 scope: str | None = None,
899 kind: str | None = None,
900 ) -> list[WikiPage]:
901 pages = list(self._pages.get(bank_id, {}).values())
902 if scope is not None:
903 pages = [p for p in pages if p.scope == scope]
904 if kind is not None:
905 pages = [p for p in pages if p.kind == kind]
906 return pages
908 async def delete_page(self, page_id: str, bank_id: str) -> bool:
909 bank_pages = self._pages.get(bank_id, {})
910 if page_id not in bank_pages:
911 return False
912 del bank_pages[page_id]
913 self._history.get(bank_id, {}).pop(page_id, None)
914 return True
916 async def health(self) -> HealthStatus:
917 return HealthStatus(healthy=True, message="in-memory wiki store")
919 def revision_history(self, page_id: str, bank_id: str) -> list[WikiPage]:
920 """Return past revisions for a page (oldest first). Testing helper."""
921 return list(self._history.get(bank_id, {}).get(page_id, []))
924# ---------------------------------------------------------------------------
925# In-Memory PageIndex Store (M9 — section recall; see ADR-006/007)
926# ---------------------------------------------------------------------------
929class InMemoryPageIndexStore:
930 """Fully functional in-memory PageIndex store for testing (M9 section recall).
932 Holds documents (one per conversation), sections (tree nodes), entity
933 mentions per section, and section-to-section links. PR1 only exercises
934 the document/section path; PR2 populates entities and links.
936 See :class:`~astrocyte.provider.PageIndexStore` for the SPI.
937 """
939 SPI_VERSION: ClassVar[int] = 1
941 def __init__(self) -> None:
942 # (bank_id, source_id) → PageIndexDocument
943 self._documents: dict[tuple[str, str], PageIndexDocument] = {}
944 # document_id → list[PageIndexSection], ordered by line_num
945 self._sections: dict[str, list[PageIndexSection]] = {}
946 # document_id → list[PageIndexSectionEntity]
947 self._section_entities: dict[str, list[PageIndexSectionEntity]] = {}
948 # document_id → list[PageIndexSectionLink]
949 self._section_links: dict[str, list[PageIndexSectionLink]] = {}
950 # M10.1: bank_id → list[(WikiPage, embedding)]; provenance:
951 # page_id → list[(document_id, line_num)]
952 self._wiki_pages: dict[str, list[tuple["WikiPage", list[float] | None]]] = {}
953 self._wiki_provenance: dict[str, list[tuple[str, int]]] = {}
954 # M12.1: document_id → list[PageIndexFact]
955 self._facts: dict[str, list[PageIndexFact]] = {}
957 async def save_document(self, doc: PageIndexDocument) -> str:
958 # Upsert keyed on (bank_id, source_id). Behaviour matches the
959 # Postgres adapter:
960 # - If an existing row matches the key, preserve its id so
961 # child rows (sections, links) remain valid.
962 # - If no existing row AND the input id is empty/falsy,
963 # assign a fresh UUID (mirrors gen_random_uuid() default).
964 # - Otherwise honour the caller-supplied id.
965 import uuid as _uuid
966 from dataclasses import replace as _replace
968 key = (doc.bank_id, doc.source_id)
969 existing = self._documents.get(key)
970 if existing is not None:
971 doc = _replace(doc, id=existing.id)
972 elif not doc.id:
973 doc = _replace(doc, id=str(_uuid.uuid4()))
974 self._documents[key] = doc
975 return doc.id
977 async def save_sections(
978 self,
979 document_id: str,
980 sections: list[PageIndexSection],
981 ) -> int:
982 # Atomic replace — drop any prior sections for this document
983 # before writing the new tree. This matches the SQL adapter's
984 # DELETE + INSERT pattern; partial trees are never observable.
985 sorted_sections = sorted(sections, key=lambda s: s.line_num)
986 self._sections[document_id] = sorted_sections
987 # Cascading wipe of dependent rows mirrors the FK ON DELETE CASCADE
988 # in migration 015. PR2 will repopulate these.
989 self._section_entities.pop(document_id, None)
990 self._section_links.pop(document_id, None)
991 return len(sorted_sections)
993 async def load_document(
994 self,
995 bank_id: str,
996 source_id: str,
997 ) -> PageIndexDocument | None:
998 return self._documents.get((bank_id, source_id))
1000 async def load_skeleton(self, document_id: str) -> list[PageIndexSection]:
1001 # Return a shallow copy without summary_embedding to mirror the
1002 # SQL adapter's projection (the picker doesn't need embeddings).
1003 from dataclasses import replace as _replace
1005 return [_replace(s, summary_embedding=None) for s in self._sections.get(document_id, [])]
1007 async def save_section_embeddings(
1008 self,
1009 document_id: str,
1010 embeddings: list[tuple[int, list[float]]],
1011 ) -> int:
1012 # Update existing PageIndexSection rows in place. PR2 commit A
1013 # mirrors the Postgres adapter's UPDATE pattern: skip rows with
1014 # no matching line_num (we don't auto-create sections from
1015 # embeddings — the tree-build step is the source of truth).
1016 from dataclasses import replace as _replace
1018 sections = self._sections.get(document_id, [])
1019 if not sections:
1020 return 0
1021 wanted = dict(embeddings) # line_num → vec
1022 n = 0
1023 new_sections: list[PageIndexSection] = []
1024 for s in sections:
1025 vec = wanted.pop(s.line_num, None)
1026 if vec is not None:
1027 new_sections.append(_replace(s, summary_embedding=list(vec)))
1028 n += 1
1029 else:
1030 new_sections.append(s)
1031 if n > 0:
1032 self._sections[document_id] = new_sections
1033 return n
1035 async def save_section_entities(
1036 self,
1037 entities: list[PageIndexSectionEntity],
1038 ) -> int:
1039 n = 0
1040 for e in entities:
1041 bucket = self._section_entities.setdefault(e.document_id, [])
1042 # Idempotent on the composite key.
1043 if not any(
1044 existing.line_num == e.line_num and existing.entity_name == e.entity_name for existing in bucket
1045 ):
1046 bucket.append(e)
1047 n += 1
1048 return n
1050 async def save_section_links(
1051 self,
1052 links: list[PageIndexSectionLink],
1053 ) -> int:
1054 n = 0
1055 for link in links:
1056 bucket = self._section_links.setdefault(link.from_doc, [])
1057 key = (link.from_line, link.to_doc, link.to_line, link.link_type)
1058 if not any(
1059 (existing.from_line, existing.to_doc, existing.to_line, existing.link_type) == key
1060 for existing in bucket
1061 ):
1062 bucket.append(link)
1063 n += 1
1064 return n
1066 async def populate_semantic_knn_links(
1067 self,
1068 document_id: str,
1069 *,
1070 top_k: int = 5,
1071 min_similarity: float = 0.5,
1072 ) -> int:
1073 """PR2 D.7.1 in-memory implementation: pure cosine similarity
1074 over each pair of sections in the document. Mirrors the SQL
1075 adapter's LATERAL kNN — for each section, link to top_k
1076 most-similar OTHER sections (cosine sim ≥ ``min_similarity``)."""
1077 sections = self._sections.get(document_id, [])
1078 embedded = [s for s in sections if s.summary_embedding]
1079 if len(embedded) < 2:
1080 return 0
1082 new_links: list[PageIndexSectionLink] = []
1083 for s1 in embedded:
1084 scored: list[tuple[float, PageIndexSection]] = []
1085 for s2 in embedded:
1086 if s2.line_num == s1.line_num:
1087 continue
1088 sim = _cosine_sim(s1.summary_embedding, s2.summary_embedding)
1089 if sim >= min_similarity:
1090 scored.append((sim, s2))
1091 scored.sort(key=lambda x: x[0], reverse=True)
1092 for sim, s2 in scored[:top_k]:
1093 new_links.append(
1094 PageIndexSectionLink(
1095 from_doc=document_id,
1096 from_line=s1.line_num,
1097 to_doc=document_id,
1098 to_line=s2.line_num,
1099 link_type="semantic_knn",
1100 weight=float(sim),
1101 )
1102 )
1103 # Reuse save_section_links for idempotent insert semantics.
1104 return await self.save_section_links(new_links)
1106 # ── PR2 commit B: parallel-strategy query methods ─────────────────
1108 def _docs_in_bank(self, bank_id: str) -> set[str]:
1109 return {d.id for (b, _src), d in self._documents.items() if b == bank_id}
1111 def _section_passes_session(self, doc_id: str, line_num: int, session_filter: str | None) -> bool:
1112 """M31 Fix 2 — section-level helper (analogue of _matches_session
1113 for facts). Returns True when ``session_filter`` is None OR
1114 the section at (doc_id, line_num) has matching session_id.
1115 """
1116 if session_filter is None:
1117 return True
1118 for s in self._sections.get(doc_id, []):
1119 if s.line_num == line_num:
1120 return s.session_id == session_filter
1121 return False
1123 async def search_sections_semantic(
1124 self,
1125 bank_id: str,
1126 query_embedding: list[float],
1127 *,
1128 top_k: int = 20,
1129 session_filter: str | None = None,
1130 ) -> list[tuple[str, int, float]]:
1131 if not query_embedding:
1132 return []
1133 doc_ids = self._docs_in_bank(bank_id)
1134 scored: list[tuple[str, int, float]] = []
1135 for doc_id in doc_ids:
1136 for s in self._sections.get(doc_id, []):
1137 if not s.summary_embedding:
1138 continue
1139 if session_filter is not None and s.session_id != session_filter:
1140 continue
1141 # Cosine similarity (Postgres adapter uses 1 - distance).
1142 score = _cosine_sim(query_embedding, s.summary_embedding)
1143 scored.append((doc_id, s.line_num, score))
1144 scored.sort(key=lambda x: x[2], reverse=True)
1145 return scored[:top_k]
1147 async def search_sections_keyword(
1148 self,
1149 bank_id: str,
1150 query: str,
1151 *,
1152 top_k: int = 20,
1153 speaker: str | None = None,
1154 document_id: str | None = None,
1155 session_filter: str | None = None,
1156 ) -> list[tuple[str, int, float]]:
1157 if not query.strip():
1158 return []
1159 # Lightweight in-memory keyword scoring: count distinct query
1160 # tokens that appear in title+summary (case-insensitive).
1161 # Mirrors the spirit of ts_rank_cd's "more matched terms = higher
1162 # score" without re-implementing the full tsvector machinery.
1163 terms = [t for t in query.lower().split() if t]
1164 if not terms:
1165 return []
1166 if document_id is not None:
1167 # Caller wants single-doc scope (PR2.6 temporal-arithmetic
1168 # find_event_date use case); skip the bank fanout entirely.
1169 doc_ids = [document_id] if document_id in self._sections else []
1170 else:
1171 doc_ids = self._docs_in_bank(bank_id)
1172 scored: list[tuple[str, int, float]] = []
1173 for doc_id in doc_ids:
1174 for s in self._sections.get(doc_id, []):
1175 if speaker is not None and s.speaker != speaker:
1176 continue
1177 if session_filter is not None and s.session_id != session_filter:
1178 continue
1179 haystack = (s.title or "").lower() + " " + (s.summary or "").lower()
1180 hits = sum(1 for t in terms if t in haystack)
1181 if hits > 0:
1182 scored.append((doc_id, s.line_num, float(hits)))
1183 scored.sort(key=lambda x: x[2], reverse=True)
1184 return scored[:top_k]
1186 async def search_sections_by_entities(
1187 self,
1188 bank_id: str,
1189 entity_names: list[str],
1190 *,
1191 top_k: int = 20,
1192 session_filter: str | None = None,
1193 ) -> list[tuple[str, int, float]]:
1194 if not entity_names:
1195 return []
1196 # Hindsight's CTE pattern: count distinct matching entities per
1197 # section. Case-insensitive match on entity_name.
1198 wanted = {n.lower() for n in entity_names if n and n.strip()}
1199 if not wanted:
1200 return []
1201 doc_ids = self._docs_in_bank(bank_id)
1202 per_section: dict[tuple[str, int], set[str]] = {}
1203 for doc_id in doc_ids:
1204 for e in self._section_entities.get(doc_id, []):
1205 low = e.entity_name.lower()
1206 if low in wanted:
1207 if not self._section_passes_session(doc_id, e.line_num, session_filter):
1208 continue
1209 per_section.setdefault((doc_id, e.line_num), set()).add(low)
1210 scored = [(doc_id, line_num, float(len(matches))) for (doc_id, line_num), matches in per_section.items()]
1211 scored.sort(key=lambda x: x[2], reverse=True)
1212 return scored[:top_k]
1214 async def search_sections_temporal(
1215 self,
1216 bank_id: str,
1217 date_range,
1218 *,
1219 top_k: int = 20,
1220 session_filter: str | None = None,
1221 ) -> list[tuple[str, int, float]]:
1222 start, end = date_range
1223 doc_ids = self._docs_in_bank(bank_id)
1224 scored: list[tuple[str, int, float]] = []
1225 for doc_id in doc_ids:
1226 for s in self._sections.get(doc_id, []):
1227 if session_filter is not None and s.session_id != session_filter:
1228 continue
1229 sd = s.session_date
1230 if sd is not None and start <= sd <= end:
1231 scored.append((doc_id, s.line_num, 1.0))
1232 scored.sort(key=lambda x: (x[0], x[1])) # stable order
1233 return scored[:top_k]
1235 async def expand_section_links(
1236 self,
1237 seeds: list[tuple[str, int]],
1238 *,
1239 link_types: list[str] | None = None,
1240 top_k: int = 20,
1241 ) -> list[tuple[str, int, float]]:
1242 if not seeds:
1243 return []
1244 seed_set = {(d, ln) for d, ln in seeds}
1245 type_filter = set(link_types) if link_types else None
1246 weights: dict[tuple[str, int], float] = {}
1247 for doc_id, links in self._section_links.items():
1248 for link in links:
1249 if type_filter is not None and link.link_type not in type_filter:
1250 continue
1251 # Outgoing edge from a seed
1252 if (link.from_doc, link.from_line) in seed_set:
1253 key = (link.to_doc, link.to_line)
1254 weights[key] = weights.get(key, 0.0) + link.weight
1255 # Incoming edge to a seed
1256 if (link.to_doc, link.to_line) in seed_set:
1257 key = (link.from_doc, link.from_line)
1258 weights[key] = weights.get(key, 0.0) + link.weight
1259 scored = [(d, ln, w) for (d, ln), w in weights.items()]
1260 scored.sort(key=lambda x: x[2], reverse=True)
1261 return scored[:top_k]
1263 async def expand_sections_by_shared_entities(
1264 self,
1265 bank_id: str,
1266 seeds: list[tuple[str, int]],
1267 *,
1268 top_k: int = 20,
1269 exclude_seeds: bool = True,
1270 ) -> list[tuple[str, int, float]]:
1271 if not seeds:
1272 return []
1273 seed_set = {(d, ln) for d, ln in seeds}
1274 doc_ids = self._docs_in_bank(bank_id)
1275 # Collect seed entity names (case-insensitive set).
1276 seed_entities: set[str] = set()
1277 for doc_id, line_num in seed_set:
1278 for e in self._section_entities.get(doc_id, []):
1279 if e.line_num == line_num:
1280 seed_entities.add(e.entity_name.lower())
1281 if not seed_entities:
1282 return []
1283 # For every section in the bank, count distinct overlap with seed_entities.
1284 per_section: dict[tuple[str, int], set[str]] = {}
1285 for doc_id in doc_ids:
1286 for e in self._section_entities.get(doc_id, []):
1287 low = e.entity_name.lower()
1288 if low not in seed_entities:
1289 continue
1290 key = (doc_id, e.line_num)
1291 if exclude_seeds and key in seed_set:
1292 continue
1293 per_section.setdefault(key, set()).add(low)
1294 scored = [(doc_id, line_num, float(len(matches))) for (doc_id, line_num), matches in per_section.items()]
1295 scored.sort(key=lambda x: x[2], reverse=True)
1296 return scored[:top_k]
1298 async def list_distinct_entities(
1299 self,
1300 bank_id: str,
1301 document_id: str,
1302 *,
1303 pattern: str | None = None,
1304 limit: int = 50,
1305 ) -> list[tuple[str, int]]:
1306 # Filter section_entities to this document; substring-match the
1307 # entity_name when pattern is set (case-insensitive). Mirrors
1308 # the Postgres ILIKE semantics: pattern is a substring, ``%``
1309 # wildcards are NOT auto-injected — caller passes them if
1310 # desired (matches Postgres behaviour).
1311 entries = self._section_entities.get(document_id, [])
1312 if not entries:
1313 return []
1314 if pattern is not None:
1315 needle = pattern.lower().replace("%", "")
1316 entries = [e for e in entries if needle in e.entity_name.lower()]
1317 # Distinct (line_num, entity_name) collapses to one mention per
1318 # section — matches the Postgres PK (document_id, line_num,
1319 # entity_name).
1320 seen: set[tuple[int, str]] = set()
1321 counts: dict[str, int] = {}
1322 for e in entries:
1323 key = (e.line_num, e.entity_name)
1324 if key in seen:
1325 continue
1326 seen.add(key)
1327 counts[e.entity_name] = counts.get(e.entity_name, 0) + 1
1328 # Order by count desc, then name asc for determinism.
1329 ordered = sorted(counts.items(), key=lambda kv: (-kv[1], kv[0]))
1330 return ordered[: max(1, limit)]
1332 # ── M12.1 fact-grain ───────────────────────────────────────────
1334 async def save_facts(self, facts) -> int:
1335 if not facts:
1336 return 0
1337 for f in facts:
1338 self._facts.setdefault(f.document_id, []).append(f)
1339 return len(facts)
1341 async def update_fact_embeddings(self, embeddings) -> int:
1342 from dataclasses import replace as _replace
1344 if not embeddings:
1345 return 0
1346 emb_by_id = dict(embeddings)
1347 updated = 0
1348 for doc_id, bucket in self._facts.items():
1349 for i, f in enumerate(bucket):
1350 if f.id in emb_by_id:
1351 bucket[i] = _replace(f, embedding=emb_by_id[f.id])
1352 updated += 1
1353 return updated
1355 def _matches_session(self, fact: "PageIndexFact", session_filter: str | None) -> bool:
1356 """M31 Fix 2 — InMemory analogue of the Postgres EXISTS clause.
1358 Returns True when the fact's anchoring section's session_id matches
1359 ``session_filter``. Top-level facts (no anchor) and facts whose
1360 section has no session_id both fail the filter — intentional, since
1361 ``session_filter`` semantics is "limit to this session's content".
1362 """
1363 if session_filter is None:
1364 return True
1365 if fact.document_id is None or fact.line_num is None:
1366 return False
1367 for s in self._sections.get(fact.document_id, []):
1368 if s.line_num == fact.line_num:
1369 return s.session_id == session_filter
1370 return False
1372 async def search_facts_semantic(
1373 self,
1374 bank_id: str,
1375 query_embedding: list[float],
1376 *,
1377 top_k: int = 20,
1378 document_id: str | None = None,
1379 fact_type: str | None = None,
1380 session_filter: str | None = None,
1381 ):
1382 from astrocyte.types import PageIndexFactHit
1384 if not query_embedding:
1385 return []
1387 def _cos(a, b):
1388 import math
1390 dot = sum(x * y for x, y in zip(a, b))
1391 na = math.sqrt(sum(x * x for x in a))
1392 nb = math.sqrt(sum(x * x for x in b))
1393 return dot / (na * nb) if na > 0 and nb > 0 else 0.0
1395 scored: list[tuple[float, PageIndexFact]] = []
1396 scope = (
1397 self._facts.get(document_id, []) if document_id else [f for bucket in self._facts.values() for f in bucket]
1398 )
1399 for f in scope:
1400 if f.bank_id != bank_id:
1401 continue
1402 if fact_type is not None and f.fact_type != fact_type:
1403 continue
1404 if not f.embedding:
1405 continue
1406 if not self._matches_session(f, session_filter): # M31 Fix 2
1407 continue
1408 scored.append((_cos(query_embedding, f.embedding), f))
1409 scored.sort(key=lambda kv: kv[0], reverse=True)
1410 return [
1411 PageIndexFactHit(
1412 fact_id=f.id,
1413 document_id=f.document_id,
1414 line_num=f.line_num,
1415 text=f.text,
1416 fact_type=f.fact_type,
1417 speaker=f.speaker,
1418 occurred_start=f.occurred_start,
1419 occurred_end=f.occurred_end,
1420 entities=list(f.entities or []),
1421 confidence_score=getattr(f, "confidence_score", None), # M27 / M28-A
1422 mentioned_at=getattr(f, "mentioned_at", None), # M27
1423 event_date=getattr(f, "event_date", None), # M31 Fix 4
1424 score=float(s),
1425 )
1426 for s, f in scored[: max(1, top_k)]
1427 ]
1429 async def search_facts_by_entity(
1430 self,
1431 bank_id: str,
1432 entity_name: str,
1433 *,
1434 top_k: int = 50,
1435 document_id: str | None = None,
1436 fact_type: str | None = None, # M34-3 — per-fact-type segmentation
1437 session_filter: str | None = None,
1438 ):
1439 from astrocyte.types import PageIndexFactHit
1441 needle = entity_name.casefold()
1442 scope = (
1443 self._facts.get(document_id, []) if document_id else [f for bucket in self._facts.values() for f in bucket]
1444 )
1445 hits: list["PageIndexFactHit"] = []
1446 for f in scope:
1447 if f.bank_id != bank_id:
1448 continue
1449 if fact_type is not None and f.fact_type != fact_type:
1450 continue
1451 if not any(needle in (e or "").casefold() for e in (f.entities or [])):
1452 continue
1453 if not self._matches_session(f, session_filter): # M31 Fix 2
1454 continue
1455 hits.append(
1456 PageIndexFactHit(
1457 fact_id=f.id,
1458 document_id=f.document_id,
1459 line_num=f.line_num,
1460 text=f.text,
1461 fact_type=f.fact_type,
1462 speaker=f.speaker,
1463 occurred_start=f.occurred_start,
1464 occurred_end=f.occurred_end,
1465 entities=list(f.entities or []),
1466 confidence_score=getattr(f, "confidence_score", None),
1467 mentioned_at=getattr(f, "mentioned_at", None),
1468 event_date=getattr(f, "event_date", None), # M31 Fix 4
1469 score=1.0,
1470 )
1471 )
1472 if len(hits) >= top_k:
1473 break
1474 return hits
1476 async def search_facts_keyword(
1477 self,
1478 bank_id: str,
1479 query: str,
1480 *,
1481 top_k: int = 20,
1482 document_id: str | None = None,
1483 fact_type: str | None = None, # M34-3 — per-fact-type segmentation
1484 session_filter: str | None = None,
1485 ):
1486 """M31c — token-overlap analogue of Postgres BM25 over fact_text.
1488 Mirrors the spirit of ``search_sections_keyword``: lower-cases
1489 both query and fact text, splits on whitespace, counts how
1490 many distinct query tokens appear as substrings of the fact
1491 text, and ranks by that count desc. Not BM25 proper but the
1492 same family of "literal-keyword-match wins over generic
1493 thematic match" behaviour that fixes the SSU bench failures.
1494 """
1495 from astrocyte.types import PageIndexFactHit
1497 if not query or not query.strip():
1498 return []
1499 terms = [t for t in query.lower().split() if t and len(t) > 1]
1500 if not terms:
1501 return []
1502 terms = list(dict.fromkeys(terms)) # dedupe, preserve order
1503 scope = (
1504 self._facts.get(document_id, [])
1505 if document_id
1506 else [f for bucket in self._facts.values() for f in bucket]
1507 )
1508 scored: list[tuple[float, "PageIndexFact"]] = []
1509 for f in scope:
1510 if f.bank_id != bank_id:
1511 continue
1512 if fact_type is not None and f.fact_type != fact_type:
1513 continue
1514 if not self._matches_session(f, session_filter): # M31 Fix 2
1515 continue
1516 haystack = (f.text or "").lower()
1517 hits = sum(1 for t in terms if t in haystack)
1518 if hits > 0:
1519 scored.append((float(hits), f))
1520 scored.sort(key=lambda kv: kv[0], reverse=True)
1521 return [
1522 PageIndexFactHit(
1523 fact_id=f.id,
1524 document_id=f.document_id,
1525 line_num=f.line_num,
1526 text=f.text,
1527 fact_type=f.fact_type,
1528 speaker=f.speaker,
1529 occurred_start=f.occurred_start,
1530 occurred_end=f.occurred_end,
1531 entities=list(f.entities or []),
1532 confidence_score=getattr(f, "confidence_score", None),
1533 mentioned_at=getattr(f, "mentioned_at", None),
1534 event_date=getattr(f, "event_date", None),
1535 score=float(s),
1536 )
1537 for s, f in scored[: max(1, top_k)]
1538 ]
1540 async def search_facts_temporal(
1541 self,
1542 bank_id: str,
1543 date_range,
1544 *,
1545 top_k: int = 50,
1546 document_id: str | None = None,
1547 fact_type: str | None = None, # M34-3 — per-fact-type segmentation
1548 session_filter: str | None = None,
1549 ):
1550 from astrocyte.types import PageIndexFactHit
1552 start, end = date_range
1553 scope = (
1554 self._facts.get(document_id, []) if document_id else [f for bucket in self._facts.values() for f in bucket]
1555 )
1557 # M33-3b — Hindsight-parity 4-way OR (mirrors the Postgres SQL):
1558 # a fact qualifies if ANY of {spanning occurred range overlap,
1559 # mentioned_at in range, occurred_start in range, occurred_end
1560 # in range}. Sort DESC by freshest available date so the top_k
1561 # cap retains newest candidates (Hindsight: ORDER BY
1562 # COALESCE(occurred_start, mentioned_at, occurred_end) DESC).
1563 def _in_range(d):
1564 return d is not None and start <= d <= end
1566 def _qualifies(f) -> bool:
1567 os_ = f.occurred_start
1568 oe_ = f.occurred_end
1569 ma_ = getattr(f, "mentioned_at", None)
1570 if os_ is not None and oe_ is not None and os_ <= end and oe_ >= start:
1571 return True
1572 return _in_range(ma_) or _in_range(os_) or _in_range(oe_)
1574 def _sort_key(f):
1575 # COALESCE(occurred_start, mentioned_at, occurred_end).
1576 # Push None to the bottom by pairing with a falsy flag.
1577 d = f.occurred_start or getattr(f, "mentioned_at", None) or f.occurred_end
1578 return (d is None, d)
1580 matched = [
1581 f
1582 for f in scope
1583 if f.bank_id == bank_id
1584 and (fact_type is None or f.fact_type == fact_type) # M34-3
1585 and _qualifies(f)
1586 and self._matches_session(f, session_filter) # M31 Fix 2
1587 ]
1588 matched.sort(key=_sort_key, reverse=True)
1589 hits: list["PageIndexFactHit"] = []
1590 for f in matched[:top_k]:
1591 hits.append(
1592 PageIndexFactHit(
1593 fact_id=f.id,
1594 document_id=f.document_id,
1595 line_num=f.line_num,
1596 text=f.text,
1597 fact_type=f.fact_type,
1598 speaker=f.speaker,
1599 occurred_start=f.occurred_start,
1600 occurred_end=f.occurred_end,
1601 entities=list(f.entities or []),
1602 confidence_score=getattr(f, "confidence_score", None),
1603 mentioned_at=getattr(f, "mentioned_at", None),
1604 event_date=getattr(f, "event_date", None), # M31 Fix 4
1605 score=1.0,
1606 )
1607 )
1608 return hits
1610 async def count_facts_matching(
1611 self,
1612 bank_id: str,
1613 document_id: str,
1614 *,
1615 entity_pattern: str | None = None,
1616 fact_type: str | None = None,
1617 ) -> int:
1618 count = 0
1619 for f in self._facts.get(document_id, []):
1620 if f.bank_id != bank_id:
1621 continue
1622 if fact_type is not None and f.fact_type != fact_type:
1623 continue
1624 if entity_pattern is not None:
1625 needle = entity_pattern.lower().replace("%", "")
1626 if not any(needle in (e or "").lower() for e in (f.entities or [])):
1627 continue
1628 count += 1
1629 return count
1631 async def save_section_event_dates(
1632 self,
1633 document_id: str,
1634 event_dates,
1635 ) -> int:
1636 if not event_dates:
1637 return 0
1638 from dataclasses import replace as _replace
1640 sections = self._sections.get(document_id) or []
1641 if not sections:
1642 return 0
1643 date_by_line = {line_num: (start, end) for line_num, start, end in event_dates}
1644 updated = 0
1645 new_list = []
1646 for s in sections:
1647 if s.line_num in date_by_line:
1648 start, end = date_by_line[s.line_num]
1649 new_list.append(_replace(s, occurred_start=start, occurred_end=end))
1650 updated += 1
1651 else:
1652 new_list.append(s)
1653 self._sections[document_id] = new_list
1654 return updated
1656 # ── M10.1 wiki / consolidation ─────────────────────────────────
1658 async def load_sections_with_embeddings(
1659 self,
1660 bank_id: str,
1661 document_id: str,
1662 ) -> list[PageIndexSection]:
1663 del bank_id # documents are bank-scoped via _documents key
1664 return list(self._sections.get(document_id, []))
1666 async def save_wiki_page(
1667 self,
1668 *,
1669 page, # WikiPage
1670 embedding: list[float] | None,
1671 provenance: list[tuple[str, int]],
1672 ) -> str:
1673 bucket = self._wiki_pages.setdefault(page.bank_id, [])
1674 # Upsert: replace any existing page with the same page_id.
1675 bucket[:] = [(p, e) for (p, e) in bucket if p.page_id != page.page_id]
1676 bucket.append((page, embedding))
1677 self._wiki_provenance[page.page_id] = list(provenance)
1678 return page.page_id
1680 async def search_wiki_pages_semantic(
1681 self,
1682 bank_id: str,
1683 query_embedding: list[float],
1684 *,
1685 top_k: int = 5,
1686 document_id: str | None = None,
1687 ):
1688 from astrocyte.types import WikiPageHit
1690 bucket = self._wiki_pages.get(bank_id, [])
1691 if not bucket or not query_embedding:
1692 return []
1693 scope_filter = f"document:{document_id}" if document_id is not None else None
1695 def _cos(a: list[float], b: list[float]) -> float:
1696 import math
1698 dot = sum(x * y for x, y in zip(a, b))
1699 na = math.sqrt(sum(x * x for x in a))
1700 nb = math.sqrt(sum(x * x for x in b))
1701 if na == 0 or nb == 0:
1702 return 0.0
1703 return dot / (na * nb)
1705 scored: list[tuple[float, WikiPage]] = []
1706 for page, emb in bucket:
1707 if scope_filter is not None and page.scope != scope_filter:
1708 continue
1709 if not emb:
1710 continue
1711 scored.append((_cos(query_embedding, emb), page))
1712 scored.sort(key=lambda kv: kv[0], reverse=True)
1713 out: list[WikiPageHit] = []
1714 for score, page in scored[: max(1, top_k)]:
1715 out.append(
1716 WikiPageHit(
1717 page_id=page.page_id,
1718 title=page.title,
1719 content=page.content,
1720 scope=page.scope,
1721 kind=page.kind,
1722 score=float(score),
1723 source_ids=list(page.source_ids),
1724 bank_id=page.bank_id,
1725 )
1726 )
1727 return out
1729 async def count_wiki_pages_for_doc(
1730 self,
1731 bank_id: str,
1732 document_id: str,
1733 ) -> int:
1734 bucket = self._wiki_pages.get(bank_id, [])
1735 scope = f"document:{document_id}"
1736 return sum(1 for p, _ in bucket if p.scope == scope)
1738 async def list_wiki_pages_for_doc(
1739 self,
1740 bank_id: str,
1741 document_id: str,
1742 ) -> list[WikiPage]:
1743 bucket = self._wiki_pages.get(bank_id, [])
1744 scope = f"document:{document_id}"
1745 return [p for p, _ in bucket if p.scope == scope]
1747 async def list_wikis_affected_by_entities(
1748 self,
1749 bank_id: str,
1750 entities: list[str],
1751 *,
1752 min_overlap: int = 1,
1753 limit: int = 8,
1754 ) -> list[tuple[WikiPage, int, list[str]]]:
1755 """M14.2: scan in-memory wiki provenance for entity overlap.
1757 Mirrors the Postgres JOIN — for each wiki page in the bank,
1758 intersect its provenance sections' entity sets with the input
1759 entities. Returns the fully-hydrated ``WikiPage`` row, the
1760 overlap count, and the shared entity names — sorted descending
1761 by overlap count with page_id ascending as a tie-break.
1762 """
1763 if not entities:
1764 return []
1765 query_set = set(entities)
1766 bucket = self._wiki_pages.get(bank_id, [])
1767 results: list[tuple[WikiPage, int, list[str]]] = []
1768 for page, _embedding in bucket:
1769 prov = self._wiki_provenance.get(page.page_id, [])
1770 if not prov:
1771 continue
1772 wiki_entities: set[str] = set()
1773 for doc_id, line_num in prov:
1774 for ent in self._section_entities.get(doc_id, []):
1775 if ent.line_num == line_num:
1776 wiki_entities.add(ent.entity_name)
1777 shared = sorted(wiki_entities & query_set)
1778 if len(shared) >= min_overlap:
1779 results.append((page, len(shared), shared))
1780 results.sort(key=lambda r: (-r[1], r[0].page_id))
1781 return results[:limit]
1783 async def health(self) -> HealthStatus:
1784 return HealthStatus(healthy=True, message="in-memory pageindex store")
1787# ---------------------------------------------------------------------------
1788# In-Memory Mental Model Store (M9 — first-class, replaces wiki piggyback)
1789# ---------------------------------------------------------------------------
1792class InMemoryMentalModelStore:
1793 """Fully functional in-memory mental-model store for testing.
1795 Mirrors :class:`InMemoryWikiStore`'s shape but for the dedicated
1796 :class:`~astrocyte.provider.MentalModelStore` SPI. Holds the current
1797 revision of each model plus a history list of past revisions.
1798 """
1800 SPI_VERSION: ClassVar[int] = 1
1802 def __init__(self) -> None:
1803 # bank_id → {model_id → MentalModel} (current revisions)
1804 self._models: dict[str, dict[str, "MentalModel"]] = {}
1805 # bank_id → {model_id → list[MentalModel]} (past revisions, oldest first)
1806 self._history: dict[str, dict[str, list["MentalModel"]]] = {}
1808 async def upsert(self, model: "MentalModel", bank_id: str) -> int:
1809 from dataclasses import replace as _replace
1810 from datetime import UTC, datetime
1812 if bank_id not in self._models:
1813 self._models[bank_id] = {}
1814 self._history[bank_id] = {}
1816 existing = self._models[bank_id].get(model.model_id)
1817 new_revision = (existing.revision + 1) if existing is not None else 1
1818 if existing is not None:
1819 # Archive the prior current-revision before replacing.
1820 self._history[bank_id].setdefault(model.model_id, []).append(existing)
1822 # Store always assigns revision + refreshed_at (callers don't need
1823 # to fill these in). M40 — validate source_timestamps alignment;
1824 # drop on mismatch rather than persist a wrong-alignment array
1825 # that would silently corrupt trend computation downstream.
1826 src_ts = model.source_timestamps
1827 if src_ts is not None and len(src_ts) != len(model.source_ids):
1828 src_ts = None
1829 stamped = _replace(
1830 model,
1831 bank_id=bank_id,
1832 revision=new_revision,
1833 refreshed_at=datetime.now(UTC),
1834 source_timestamps=src_ts,
1835 )
1836 self._models[bank_id][model.model_id] = stamped
1837 return new_revision
1839 async def get(self, model_id: str, bank_id: str) -> "MentalModel | None":
1840 return self._models.get(bank_id, {}).get(model_id)
1842 async def list(
1843 self,
1844 bank_id: str,
1845 *,
1846 scope: str | None = None,
1847 kind: str | None = None,
1848 ) -> list["MentalModel"]:
1849 models = list(self._models.get(bank_id, {}).values())
1850 if scope is not None:
1851 models = [m for m in models if m.scope == scope]
1852 if kind is not None:
1853 models = [m for m in models if m.kind == kind]
1854 return models
1856 async def delete(self, model_id: str, bank_id: str) -> bool:
1857 bank_models = self._models.get(bank_id, {})
1858 if model_id not in bank_models:
1859 return False
1860 del bank_models[model_id]
1861 # Keep history for audit; tests can call ``revision_history()`` to
1862 # observe past revisions even after delete.
1863 return True
1865 async def update_via_ops(
1866 self,
1867 model_id: str,
1868 bank_id: str,
1869 operations_json: list[dict],
1870 ) -> "tuple[int, dict] | None":
1871 """M21 — apply structured delta operations and re-upsert.
1873 Lazy-migrates legacy rows by parsing ``content`` on first
1874 refresh. Re-renders ``content`` from the new structured doc so
1875 markdown readers see the updated text without a separate
1876 migration pass.
1877 """
1878 from dataclasses import replace as _replace
1880 from astrocyte.pipeline.delta_ops import (
1881 DeltaOperationList,
1882 apply_operations,
1883 )
1884 from astrocyte.pipeline.structured_doc import (
1885 StructuredDocument,
1886 parse_markdown,
1887 render_document,
1888 )
1890 current = self._models.get(bank_id, {}).get(model_id)
1891 if current is None:
1892 return None
1894 # Resolve the current structured doc — lazy-migrate from
1895 # raw markdown when the field is None (legacy rows).
1896 if current.structured_doc is None:
1897 doc = parse_markdown(current.content or "")
1898 else:
1899 doc = StructuredDocument.model_validate(current.structured_doc)
1901 # Parse the LLM/caller's op list via the Pydantic schema.
1902 # Schema-invalid op shapes (unknown ``op`` discriminator,
1903 # missing required fields) raise ValidationError; per the
1904 # conservative-failure contract, we catch and return a
1905 # zero-ops apply (doc stays as-is, no revision bump).
1906 try:
1907 ops_container = DeltaOperationList.model_validate({"operations": operations_json})
1908 except Exception:
1909 # Schema-level failure: nothing applied. Return current revision.
1910 return (current.revision, {"applied": [], "skipped": [], "changed": False})
1912 applied = apply_operations(doc, ops_container.operations)
1913 if not applied.changed:
1914 return (
1915 current.revision,
1916 {"applied": applied.applied, "skipped": applied.skipped, "changed": False},
1917 )
1919 new_doc_dict = applied.document.model_dump()
1920 new_content = render_document(applied.document)
1921 new_model = _replace(
1922 current,
1923 content=new_content,
1924 structured_doc=new_doc_dict,
1925 )
1926 new_revision = await self.upsert(new_model, bank_id)
1927 return (
1928 new_revision,
1929 {"applied": applied.applied, "skipped": applied.skipped, "changed": True},
1930 )
1932 async def refresh(
1933 self,
1934 model_id: str,
1935 bank_id: str,
1936 new_source_ids: list[str],
1937 ) -> "MentalModel | None":
1938 """M28 — merge new source_ids into existing model and bump revision.
1940 In-memory semantics: the production store will re-run an LLM
1941 compile against the merged source set; the test store simply
1942 acknowledges the new sources (deduped, preserving existing
1943 order then appending novel ids) and bumps the revision via
1944 :meth:`upsert`. That keeps tests focused on the wiring and
1945 contract surface without coupling them to LLM behaviour.
1946 """
1947 from dataclasses import replace as _replace
1949 current = self._models.get(bank_id, {}).get(model_id)
1950 if current is None:
1951 return None
1953 # Order-preserving dedup: keep existing source_ids in original
1954 # order, then append novel ids in input order.
1955 seen: set[str] = set(current.source_ids)
1956 merged = list(current.source_ids)
1957 for sid in new_source_ids:
1958 if sid not in seen:
1959 seen.add(sid)
1960 merged.append(sid)
1962 refreshed = _replace(current, source_ids=merged)
1963 await self.upsert(refreshed, bank_id)
1964 return self._models.get(bank_id, {}).get(model_id)
1966 async def health(self) -> HealthStatus:
1967 return HealthStatus(healthy=True, message="in-memory mental model store")
1969 def revision_history(self, model_id: str, bank_id: str) -> list["MentalModel"]:
1970 """Return past revisions for a model (oldest first). Testing helper."""
1971 return list(self._history.get(bank_id, {}).get(model_id, []))
1974# ---------------------------------------------------------------------------
1975# In-Memory Source Store (M10 — documents + chunks normalisation)
1976# ---------------------------------------------------------------------------
1979class InMemorySourceStore:
1980 """Fully functional in-memory source-document store for testing.
1982 Mirrors the :class:`~astrocyte.provider.SourceStore` SPI: document
1983 create/get/list/delete + bulk-chunk insert with content_hash dedup.
1984 Soft-deletes match the Postgres adapter's lifecycle (deleted_at on
1985 documents; cascade-by-store-eviction on chunks).
1986 """
1988 SPI_VERSION: ClassVar[int] = 1
1990 def __init__(self) -> None:
1991 # bank_id → {document_id → SourceDocument}
1992 self._docs: dict[str, dict[str, "SourceDocument"]] = {}
1993 # bank_id → {chunk_id → SourceChunk}
1994 self._chunks: dict[str, dict[str, "SourceChunk"]] = {}
1995 # bank_id → {(content_hash) → document_id} for fast dedup lookup
1996 self._doc_hash_index: dict[str, dict[str, str]] = {}
1997 # bank_id → {(content_hash) → chunk_id} for fast dedup lookup
1998 self._chunk_hash_index: dict[str, dict[str, str]] = {}
2000 async def store_document(self, document: "SourceDocument") -> str:
2001 from dataclasses import replace as _replace
2002 from datetime import UTC
2003 from datetime import datetime as _dt
2005 bank_id = document.bank_id
2006 self._docs.setdefault(bank_id, {})
2007 self._doc_hash_index.setdefault(bank_id, {})
2009 # Dedup by content_hash when set.
2010 if document.content_hash:
2011 existing_id = self._doc_hash_index[bank_id].get(document.content_hash)
2012 if existing_id is not None and existing_id in self._docs[bank_id]:
2013 return existing_id
2015 # Stamp created_at if not provided (store owns this field).
2016 stamped = _replace(
2017 document,
2018 created_at=document.created_at or _dt.now(UTC),
2019 )
2020 self._docs[bank_id][document.id] = stamped
2021 if document.content_hash:
2022 self._doc_hash_index[bank_id][document.content_hash] = document.id
2023 return document.id
2025 async def get_document(
2026 self,
2027 document_id: str,
2028 bank_id: str,
2029 ) -> "SourceDocument | None":
2030 return self._docs.get(bank_id, {}).get(document_id)
2032 async def find_document_by_hash(
2033 self,
2034 content_hash: str,
2035 bank_id: str,
2036 ) -> "SourceDocument | None":
2037 doc_id = self._doc_hash_index.get(bank_id, {}).get(content_hash)
2038 if doc_id is None:
2039 return None
2040 return self._docs.get(bank_id, {}).get(doc_id)
2042 async def list_documents(
2043 self,
2044 bank_id: str,
2045 *,
2046 limit: int = 100,
2047 ) -> list["SourceDocument"]:
2048 docs = list(self._docs.get(bank_id, {}).values())
2049 # Newest-first by created_at; tolerate Nones via a fallback.
2050 from datetime import UTC
2051 from datetime import datetime as _dt
2053 order = {id(d): i for i, d in enumerate(docs)} # insertion order breaks clock ties
2054 docs.sort(key=lambda d: (d.created_at or _dt.fromtimestamp(0, UTC), order[id(d)]), reverse=True)
2055 return docs[:limit]
2057 async def delete_document(self, document_id: str, bank_id: str) -> bool:
2058 bank_docs = self._docs.get(bank_id, {})
2059 doc = bank_docs.pop(document_id, None)
2060 if doc is None:
2061 return False
2062 # Drop the hash-index entry too so a follow-up re-store doesn't
2063 # silently dedup against the deleted row.
2064 if doc.content_hash:
2065 self._doc_hash_index.get(bank_id, {}).pop(doc.content_hash, None)
2066 # Cascade: drop all chunks of this document.
2067 bank_chunks = self._chunks.get(bank_id, {})
2068 to_drop = [cid for cid, chunk in bank_chunks.items() if chunk.document_id == document_id]
2069 for cid in to_drop:
2070 chunk = bank_chunks.pop(cid)
2071 if chunk.content_hash:
2072 self._chunk_hash_index.get(bank_id, {}).pop(chunk.content_hash, None)
2073 return True
2075 async def store_chunks(self, chunks: list["SourceChunk"]) -> list[str]:
2076 from dataclasses import replace as _replace
2077 from datetime import UTC
2078 from datetime import datetime as _dt
2080 ids: list[str] = []
2081 for chunk in chunks:
2082 bank_id = chunk.bank_id
2083 self._chunks.setdefault(bank_id, {})
2084 self._chunk_hash_index.setdefault(bank_id, {})
2086 # Dedup by content_hash when set.
2087 if chunk.content_hash:
2088 existing = self._chunk_hash_index[bank_id].get(chunk.content_hash)
2089 if existing is not None and existing in self._chunks[bank_id]:
2090 ids.append(existing)
2091 continue
2093 stamped = _replace(
2094 chunk,
2095 created_at=chunk.created_at or _dt.now(UTC),
2096 )
2097 self._chunks[bank_id][chunk.id] = stamped
2098 if chunk.content_hash:
2099 self._chunk_hash_index[bank_id][chunk.content_hash] = chunk.id
2100 ids.append(chunk.id)
2101 return ids
2103 async def get_chunk(self, chunk_id: str, bank_id: str) -> "SourceChunk | None":
2104 return self._chunks.get(bank_id, {}).get(chunk_id)
2106 async def list_chunks(
2107 self,
2108 document_id: str,
2109 bank_id: str,
2110 ) -> list["SourceChunk"]:
2111 chunks = [c for c in self._chunks.get(bank_id, {}).values() if c.document_id == document_id]
2112 chunks.sort(key=lambda c: c.chunk_index)
2113 return chunks
2115 async def find_chunk_by_hash(
2116 self,
2117 content_hash: str,
2118 bank_id: str,
2119 ) -> "SourceChunk | None":
2120 chunk_id = self._chunk_hash_index.get(bank_id, {}).get(content_hash)
2121 if chunk_id is None:
2122 return None
2123 return self._chunks.get(bank_id, {}).get(chunk_id)
2125 async def health(self) -> HealthStatus:
2126 return HealthStatus(healthy=True, message="in-memory source store")
2129# ---------------------------------------------------------------------------
2130# Mock LLM Provider
2131# ---------------------------------------------------------------------------
2134class MockLLMProvider:
2135 """Mock LLM provider with bag-of-words embeddings and extractive synthesis.
2137 Embeddings use term-frequency vectors over a shared vocabulary, so
2138 semantically related texts have high cosine similarity — enabling
2139 meaningful vector retrieval without an API key.
2141 Synthesis extracts the most relevant memory text from the prompt context
2142 rather than returning a static string.
2143 """
2145 SPI_VERSION: ClassVar[int] = 1
2147 _embed_dim: ClassVar[int] = 128
2149 def __init__(
2150 self,
2151 default_response: str = "Mock LLM response",
2152 embedding_dimensions: int | None = None,
2153 ) -> None:
2154 self._default_response = default_response
2155 if embedding_dimensions is not None:
2156 self._embed_dim = int(embedding_dimensions)
2157 self._call_count = 0
2158 #: Last ``complete()`` user message content (for tests asserting prompt structure).
2159 self.last_user_message: str | None = None
2161 def capabilities(self) -> LLMCapabilities:
2162 return LLMCapabilities()
2164 async def complete(
2165 self,
2166 messages: list[Message],
2167 model: str | None = None,
2168 max_tokens: int = 1024,
2169 temperature: float = 0.0,
2170 tools: list = None, # list[ToolDefinition] | None — kept untyped to avoid import cycle
2171 tool_choice: str | None = None,
2172 response_format: dict | None = None,
2173 ) -> Completion:
2174 self._call_count += 1
2175 for m in reversed(messages):
2176 if m.role == "user" and isinstance(m.content, str):
2177 self.last_user_message = m.content
2178 break
2179 all_text = " ".join(m.content for m in messages if isinstance(m.content, str))
2180 # Entity extraction prompt
2181 if "extract named entities" in all_text.lower():
2182 return Completion(
2183 text='[{"name": "Test Entity", "entity_type": "OTHER", "aliases": []}]',
2184 model=model or "mock",
2185 usage=TokenUsage(input_tokens=10, output_tokens=20),
2186 )
2187 # Memory synthesis prompt — extract most relevant memory text
2188 if "<memories>" in all_text and "<query>" in all_text:
2189 answer = _extractive_synthesize(all_text)
2190 return Completion(
2191 text=answer,
2192 model=model or "mock",
2193 usage=TokenUsage(input_tokens=50, output_tokens=30),
2194 )
2195 return Completion(
2196 text=self._default_response,
2197 model=model or "mock",
2198 usage=TokenUsage(input_tokens=10, output_tokens=20),
2199 )
2201 async def embed(self, texts: list[str], model: str | None = None) -> list[list[float]]:
2202 """Generate bag-of-words embeddings with real semantic signal.
2204 Uses term-frequency vectors over a shared vocabulary. Texts that share
2205 words have high cosine similarity, enabling meaningful retrieval.
2206 """
2207 return [self._bow_embed(text) for text in texts]
2209 def _bow_embed(self, text: str) -> list[float]:
2210 """Build a term-frequency vector using the hashing trick.
2212 Each word maps to a bucket via ``hash(word) % dim`` and increments
2213 that dimension. All components are non-negative, guaranteeing
2214 non-negative cosine similarities between any two vectors — required
2215 because ``VectorHit.score`` enforces ``>= 0.0``.
2216 """
2217 import hashlib
2218 from string import punctuation
2220 dim = self._embed_dim
2221 vec = [0.0] * dim
2222 tokens = [t.strip(punctuation).lower() for t in text.split() if t.strip(punctuation)]
2224 for token in tokens:
2225 if not token:
2226 continue
2227 h = int(hashlib.md5(token.encode()).hexdigest(), 16)
2228 bucket = h % dim
2229 vec[bucket] += 1.0
2231 # L2 normalize
2232 norm = math.sqrt(sum(x * x for x in vec))
2233 if norm > 0:
2234 vec = [x / norm for x in vec]
2235 return vec
2238def _normalize_terms(text: str) -> set[str]:
2239 """Lowercase, strip punctuation, drop short words."""
2240 from string import punctuation
2242 return {t.strip(punctuation).lower() for t in text.split() if len(t.strip(punctuation)) > 2}
2245def _extractive_synthesize(prompt_text: str) -> str:
2246 """Extract the most query-relevant memory from the synthesis prompt.
2248 Parses ``<memories>`` and ``<query>`` blocks from the reflect prompt,
2249 scores each memory by word overlap with the query, and returns
2250 the top memories concatenated. This gives the mock provider
2251 meaningful reflect answers without a real LLM.
2252 """
2253 import re
2255 # Extract query
2256 query_match = re.search(r"<query>\s*(.*?)\s*</query>", prompt_text, re.DOTALL)
2257 if not query_match:
2258 return "No query found."
2259 query = query_match.group(1).strip()
2260 query_terms = _normalize_terms(query)
2262 # Extract individual memories
2263 memories_match = re.search(r"<memories>\s*(.*?)\s*</memories>", prompt_text, re.DOTALL)
2264 if not memories_match:
2265 return "No memories found."
2267 # Group continuation lines with their [Memory N] header.
2268 # Lines starting with [Memory N] start a new block; other lines
2269 # are appended to the current block.
2270 raw_lines = memories_match.group(1).strip().split("\n")
2271 blocks: list[str] = []
2272 for line in raw_lines:
2273 line = line.strip()
2274 if not line:
2275 continue
2276 if re.match(r"^\[Memory \d+\]", line):
2277 # Strip [Memory N] prefix for cleaner output
2278 text = re.sub(r"^\[Memory \d+\](\s*\([^)]*\))?(\s*\[[^\]]*\])?\s*:\s*", "", line)
2279 blocks.append(text)
2280 elif blocks:
2281 blocks[-1] += " " + line
2282 else:
2283 blocks.append(line)
2285 # Score each memory block by query term overlap
2286 scored: list[tuple[float, str]] = []
2287 for text in blocks:
2288 block_terms = _normalize_terms(text)
2289 overlap = len(query_terms & block_terms)
2290 if overlap > 0:
2291 scored.append((overlap, text))
2293 if not scored:
2294 return "I don't have relevant information to answer this question."
2296 scored.sort(key=lambda x: x[0], reverse=True)
2297 # Return top 3 most relevant memories
2298 top = [text for _, text in scored[:3]]
2299 return " ".join(top)