Coverage for astrocyte/pipeline/retain_stage.py: 78%
560 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-10-29 22:10 +0000
1from __future__ import annotations
3import asyncio
4import hashlib
5import logging
6import uuid
7from dataclasses import dataclass
8from datetime import datetime, timezone
9from typing import TYPE_CHECKING, Any
11from astrocyte.config import ExtractionProfileConfig
12from astrocyte.pipeline.chunking import chunk_text
13from astrocyte.pipeline.embedding import generate_embeddings
14from astrocyte.pipeline.entity_extraction import extract_entities
15from astrocyte.pipeline.extraction import (
16 merged_user_and_builtin_profiles,
17 prepare_retain_input,
18 resolve_retain_chunking,
19)
20from astrocyte.recall.authority import merge_retain_metadata_authority_tier
21from astrocyte.types import (
22 Entity,
23 EntityLink,
24 MemoryEntityAssociation,
25 MemoryLink,
26 RetainRequest,
27 RetainResult,
28 VectorHit,
29 VectorItem,
30)
32if TYPE_CHECKING:
33 pass
36from astrocyte.pipeline._orchestrator_common import (
37 _build_cooccurrence_pairs,
38 _entities_from_metadata,
39)
40from astrocyte.pipeline.provenance import SOURCE_KEY
42_logger = logging.getLogger("astrocyte.mip")
44#: Nearest neighbours fetched per chunk for the store-side dedup check. More
45#: than one so a negation-guarded or observation hit does not hide a genuine
46#: duplicate ranked just below it.
47_STORE_DEDUP_NEIGHBOURS = 3
48#: Store rows that are syntheses, not retained memories (``fact_type`` or
49#: ``memory_layer``); never treated as the original of a new chunk.
50_NOT_RETAINED_LAYERS = frozenset({"observation", "model"})
54def _chunk_metadata(base: dict | None, index: int, count: int) -> dict | None:
55 """A chunk's metadata: the retain's, plus its position when the text was
56 split. ``retained_at`` can't order chunks: they are stored within one
57 tick of a coarse clock (Windows: ~15 ms)."""
58 return base if count < 2 or base is None else {**base, "_chunk_index": index}
60@dataclass(frozen=True)
61class _FactOverlay:
62 """What structured extraction learned about one chunk, to apply to its
63 stored row. ``None`` / empty means "the model said nothing": the row
64 keeps the request's values."""
66 metadata: dict[str, Any]
67 fact_type: str | None
68 occurred_at: datetime | None
71def _apply_fact_overlay(
72 metadata: dict | None,
73 overlay: _FactOverlay | None,
74 *,
75 fact_type: str,
76 occurred_at: datetime | None,
77 profile_fact_type: str | None,
78) -> tuple[dict | None, str, datetime | None]:
79 """Merge one chunk's extraction results into its stored fields.
81 Precedence: a profile that sets ``fact_type`` explicitly is an operator
82 decision and wins; otherwise the model's classification replaces the
83 default. An extracted ``occurred_start`` is event time and replaces the
84 request's timestamp, which is kept as ``_mentioned_at`` so nothing is
85 lost. The ``_fact_*`` keys are prefixed and never overwrite caller keys.
86 """
87 if overlay is None:
88 return metadata, fact_type, occurred_at
89 merged = {**overlay.metadata, **(metadata or {})} if overlay.metadata else metadata
90 if overlay.fact_type is not None and not profile_fact_type:
91 fact_type = overlay.fact_type
92 if overlay.occurred_at is not None:
93 if occurred_at is not None and overlay.occurred_at != occurred_at:
94 merged = {**(merged or {}), "_mentioned_at": occurred_at.isoformat()}
95 occurred_at = overlay.occurred_at
96 return merged, fact_type, occurred_at
99class RetainStageMixin:
100 """Retain pipeline: chunk → extract → embed → persist (+ retain_many).
102 A behavior slice of PipelineOrchestrator, composed into it as a mixin.
103 ``self`` is the full orchestrator at runtime (attributes come from its
104 ``__init__`` / ``apply_config``); shared helpers live on sibling mixins.
105 """
107 async def _attach_entity_name_embeddings(self, entities: list[Entity]) -> None:
108 """Embed each entity's name and attach the vector to ``entity.embedding``.
110 Powers the Hindsight-inspired entity-resolution cascade — at retain
111 time we generate one batched embedding call across all entity names,
112 then attach the resulting vector to each entity. Both
113 ``store_entities`` (for adapters that persist embeddings) and
114 ``EntityResolver.resolve()`` (for the cosine-similarity tier of the
115 cascade) consume the attached vectors.
117 Skips entities whose ``embedding`` is already set (caller-supplied)
118 and skips entities with empty/whitespace names. On failure logs a
119 warning and leaves embeddings unset — the resolver degrades to
120 trigram-only and the cost is correctness for that one batch, not
121 a retain failure.
122 """
123 targets = [e for e in entities if e.embedding is None and e.name and e.name.strip()]
124 if not targets:
125 return
126 try:
127 vectors = await generate_embeddings(
128 [e.name.strip() for e in targets],
129 self.llm_provider,
130 )
131 except Exception as exc:
132 _logger.warning(
133 "entity name embedding failed (resolver will fall back to trigram-only tier): %s",
134 exc,
135 )
136 return
137 for entity, vec in zip(targets, vectors, strict=False):
138 entity.embedding = vec
140 async def _structured_fact_extraction_for_text(
141 self,
142 prepared,
143 request: RetainRequest,
144 *,
145 chunk_strategy: str = "paragraph",
146 chunk_max_size: int | None = None,
147 chunk_overlap: int | None = None,
148 ) -> tuple[
149 list[str] | None,
150 list[Entity] | None,
151 list[tuple[int, str]] | None, # (entity_idx_in_full_list, memory_idx) associations
152 list[tuple[int, int, float]] | None, # (source_memory_idx, target_memory_idx, confidence) caused_by
153 list[_FactOverlay] | None, # per fact, aligned with fact_texts
154 ]:
155 """Run the structured 5-dim fact extraction path.
157 Returns ``(fact_texts, entities, associations, caused_by, overlays)``
158 when the path is enabled and produces facts; all ``None`` otherwise.
159 ``overlays`` carries each fact's metadata, classification, and event
160 time for its stored row; before 2026-10-04 these were computed and
161 then discarded, so extraction paid for metadata nothing stored. Caller falls through to the legacy chunk_text +
162 extract_entities path when None is returned.
164 Indices in ``associations`` and ``caused_by`` reference positions
165 in the returned ``fact_texts``/``entities`` lists. Callers
166 resolve them to memory IDs after store_vectors.
167 """
168 # SFE supersedes the legacy ``prepared.extract_entities`` gate.
169 # The legacy setting controlled which entity-extraction PATH
170 # ran (metadata vs LLM); SFE is a third path that produces
171 # both entities and structured facts in one call. When SFE is
172 # enabled, it fires regardless of the profile's entity
173 # extraction mode — the profile's metadata-entities are
174 # ignored in favor of the richer SFE output.
175 if not self.structured_fact_extraction_enabled or self.llm_provider is None:
176 return None, None, None, None, None
178 try:
179 # Pre-chunk using the SAME strategy the legacy retain path
180 # would use (resolved from the extraction profile), then ask
181 # the LLM to enrich each chunk with metadata WITHOUT
182 # paraphrasing. Stored memory text = original chunk
183 # vocabulary at the granularity the profile specifies
184 # (e.g. dialogue turns for conversations, paragraphs for
185 # documents).
186 #
187 # The legacy "concise" path (``extract_facts``) was removed
188 # in M9 — it caused severe recall_hit_rate degradation
189 # because the LLM-paraphrased ``what`` field lost the
190 # surface vocabulary that question embeddings share. Verbatim
191 # is now the only supported mode; ``validate_astrocyte_config``
192 # rejects ``extraction_mode: concise`` at config load.
193 #
194 # Open-domain regression observed on 2026-05-02 traced to
195 # hardcoded "paragraph" strategy producing one giant chunk
196 # per LoCoMo session (no paragraph breaks in dialogue text).
197 # Profile-driven strategy fixes that.
198 from astrocyte.pipeline.chunking import chunk_text
199 from astrocyte.pipeline.fact_extraction import (
200 extract_facts_verbatim,
201 extract_facts_verbatim_parallel,
202 materialize_facts,
203 )
205 pre_chunk_kwargs: dict[str, int] = {}
206 if chunk_max_size is not None:
207 pre_chunk_kwargs["max_chunk_size"] = chunk_max_size
208 if chunk_overlap is not None:
209 pre_chunk_kwargs["overlap"] = chunk_overlap
210 chunks_local = chunk_text(
211 prepared.text,
212 strategy=chunk_strategy,
213 **pre_chunk_kwargs,
214 )
215 # Phase 3: route to the per-chunk parallel path when opt-in.
216 # Drops cross-chunk causal_relations — pair with
217 # ``causal_links.enabled=false``.
218 if self.structured_fact_extraction_parallel_chunks:
219 facts = await extract_facts_verbatim_parallel(
220 chunks_local,
221 self.llm_provider,
222 event_date=request.occurred_at,
223 max_concurrency=(self.structured_fact_extraction_parallel_chunks_max_concurrency),
224 )
225 else:
226 facts = await extract_facts_verbatim(
227 chunks_local,
228 self.llm_provider,
229 event_date=request.occurred_at,
230 )
231 except Exception as exc:
232 _logger.warning(
233 "structured fact extraction failed (%s); falling back to legacy chunk + entity-extraction path.",
234 exc,
235 )
236 return None, None, None, None, None
237 if not facts:
238 return None, None, None, None, None
240 # Materialize without embeddings here — embeddings are batched
241 # later for cost. We only need the list of fact texts and the
242 # pre-extracted entities + association index map.
243 materialized = materialize_facts(
244 facts,
245 bank_id=request.bank_id,
246 occurred_at=request.occurred_at,
247 )
248 fact_texts = [item.text for item in materialized.vector_items]
249 entities = list(materialized.entities)
251 # Build (memory_idx, entity_idx) pairs from the materialized
252 # association tuples (which reference items by ID). We need
253 # indices because the caller hasn't assigned final memory IDs
254 # yet — those come after dedup + store_vectors.
255 item_id_to_idx = {item.id: i for i, item in enumerate(materialized.vector_items)}
256 ent_id_to_idx = {e.id: i for i, e in enumerate(entities)}
257 associations: list[tuple[int, str]] = []
258 for item_id, ent_id in materialized.memory_entity_associations:
259 mem_idx = item_id_to_idx.get(item_id)
260 ent_idx = ent_id_to_idx.get(ent_id)
261 if mem_idx is None or ent_idx is None:
262 continue
263 associations.append((ent_idx, mem_idx))
265 # caused_by edges: indices both into fact list (= memory list).
266 caused_by: list[tuple[int, int, float]] = []
267 for link in materialized.memory_links:
268 src_idx = item_id_to_idx.get(link.source_memory_id)
269 tgt_idx = item_id_to_idx.get(link.target_memory_id)
270 if src_idx is None or tgt_idx is None:
271 continue
272 caused_by.append((src_idx, tgt_idx, float(link.confidence)))
274 # One VectorItem per fact, in order (materialize_facts), so the two
275 # zip exactly. A default fact_type is not a classification.
276 overlays = [
277 _FactOverlay(
278 metadata={
279 k: v
280 for k, v in (item.metadata or {}).items()
281 if k.startswith("_fact_") and (k != "_fact_type" or fact.fact_type_classified)
282 },
283 fact_type=fact.fact_type if fact.fact_type_classified else None,
284 occurred_at=fact.occurred_start,
285 )
286 for fact, item in zip(facts, materialized.vector_items, strict=True)
287 ]
289 return fact_texts, entities, associations, caused_by, overlays
291 async def _persist_semantic_links(
292 self,
293 bank_id: str,
294 memory_ids: list[str],
295 embeddings: list[list[float]],
296 ) -> None:
297 """Compute & persist semantic-kNN edges for a freshly-stored batch.
299 Hindsight parity (C3a): each new memory gets ``MemoryLink(
300 link_type="semantic", weight=cosine)`` edges to its top-K
301 most-similar existing memories with similarity ≥ threshold.
302 Best-effort — failures log and continue rather than aborting
303 retain.
304 """
305 if not self.semantic_link_graph_enabled or self.graph_store is None or not memory_ids:
306 return
307 try:
308 from astrocyte.pipeline.semantic_link_graph import compute_semantic_links
310 links = await compute_semantic_links(
311 bank_id=bank_id,
312 new_memory_ids=memory_ids,
313 new_embeddings=embeddings,
314 vector_store=self.vector_store,
315 top_k=self.semantic_link_graph_top_k,
316 similarity_threshold=self.semantic_link_graph_threshold,
317 )
318 except Exception as exc:
319 _logger.warning("semantic_link_graph computation failed: %s", exc)
320 return
321 if not links:
322 return
323 try:
324 await self.graph_store.store_memory_links(links, bank_id)
325 except Exception as exc:
326 _logger.warning("storing semantic memory_links failed: %s", exc)
328 async def _process_record_entities_with_retry(
329 self,
330 record: dict[str, Any],
331 *,
332 max_retries: int = 10,
333 base_delay: float = 0.1,
334 ) -> None:
335 """Retry-wrapper around :meth:`_process_record_entities`.
337 AGE's entity ``MERGE`` path takes per-label advisory locks
338 (``pg_advisory_xact_lock(label_oid)``) so concurrent retains
339 racing on the same label name don't produce duplicate label
340 rows. Three or more retains contending on the same label can
341 produce a CYCLIC wait: P1 holds advisory, waits on tx; P2 waits
342 on advisory; P3 holds tx, waits on advisory — Postgres breaks
343 the cycle by aborting one of them with ``DeadlockDetected``.
345 This is timing-sensitive: with HNSW vector inserts the retain
346 latency was high enough that the cyclic pattern was rare in
347 practice. With DiskANN's faster inserts (2026-05-06 default
348 switch) the LoCoMo bench surfaced it within 20 retains.
349 Postgres-style deadlocks are inherently retriable — only the
350 loser's transaction rolls back, so we just try again with
351 exponential backoff.
353 Tuning history:
355 - 2026-05-06 (initial): ``max_retries=3``, ``base_delay=0.2``.
356 Absorbed most deadlocks but the LoCoMo bench surfaced 4-way
357 cyclic waits that the 3 attempts couldn't outlast.
358 - 2026-05-06 (revised): ``max_retries=10``, ``base_delay=0.1``.
359 Retain throughput evidence: with ``concurrency=10`` retains
360 and a per-bank entity overlap rate around 30% (LoCoMo
361 characters appearing in many sessions), deadlocks fire on
362 ~5% of retains. 10 retries with 0.1s × 2^n backoff
363 (capped at 3.6s by the worst case) gives 51.1s max stall —
364 ample headroom for the cyclic-wait probability tail. The
365 bench earlier observed 6 deadlocks absorbed before one
366 exhausted — bumping to 10 attempts more than doubles the
367 probability mass we cover.
368 """
369 last_exc: BaseException | None = None
370 for attempt in range(max_retries):
371 try:
372 await self._process_record_entities(record)
373 return
374 except Exception as exc: # noqa: BLE001 — match psycopg without import dep
375 # Match by class name + message so we don't need a hard
376 # dependency on psycopg in this module (the orchestrator
377 # is provider-agnostic; deadlocks are a Postgres-side
378 # concept that bubbles up via whatever DB driver the
379 # graph_store happens to use).
380 name = type(exc).__name__.lower()
381 msg = str(exc).lower()
382 is_deadlock = "deadlock" in name or "deadlock detected" in msg
383 if not is_deadlock or attempt == max_retries - 1:
384 raise
385 last_exc = exc
386 delay = base_delay * (2**attempt)
387 _logger.warning(
388 "_process_record_entities deadlock on attempt %d/%d (%s); retrying in %.2fs",
389 attempt + 1,
390 max_retries,
391 exc,
392 delay,
393 )
394 await asyncio.sleep(delay)
395 # Defensive — loop body either returns or re-raises; this is
396 # only reachable if max_retries=0 (which we guard against above
397 # via the attempt < max_retries-1 check).
398 if last_exc is not None:
399 raise last_exc
401 async def _process_record_entities(self, record: dict[str, Any]) -> None:
402 """Run all entity-related I/O for a single retain_many record.
404 Encapsulates the embed-names / store-entities / link-memories /
405 co-occurrence-links / entity-resolution cascade so :meth:`retain_many`
406 can dispatch one task per record via :func:`asyncio.gather`. Each
407 record's writes target distinct ``(bank_id, id)`` and
408 ``(bank_id, memory_id, entity_id)`` keys, so concurrent records
409 don't contend on row-level locks.
411 Failures in entity resolution are caught and logged — they must
412 not abort retain because the raw memories are already stored.
413 """
414 if self.graph_store is None:
415 return
416 request: RetainRequest = record["request"]
417 prepared = record["prepared"]
418 memory_ids: list[str] = record["memory_ids"]
419 entities: list[Entity] = record["entities"]
420 if not entities:
421 return
423 if self.entity_resolver is not None:
424 async with self._profiler.time("entity_emb"):
425 await self._attach_entity_name_embeddings(entities)
426 # Path B (Hindsight-style): rewrite tentative IDs to canonical
427 # IDs BEFORE storage so different surface forms that match an
428 # existing canonical never produce duplicate entity rows. The
429 # post-store ``resolve()`` alias pass becomes redundant in this
430 # mode and is skipped below.
431 if self.entity_resolver.canonical_resolution:
432 try:
433 async with self._profiler.time("entity_resolve"):
434 await self.entity_resolver.resolve_canonical_ids_in_place(
435 new_entities=entities,
436 bank_id=request.bank_id,
437 graph_store=self.graph_store,
438 event_date=request.occurred_at,
439 )
440 except Exception as exc:
441 _logger.warning(
442 "canonical resolution failed during retain (falling back to tentative IDs): %s",
443 exc,
444 )
446 async with self._profiler.time("entity_store"):
447 entity_ids = await self.graph_store.store_entities(entities, request.bank_id)
449 # Per-fact entity associations when SFE supplied them
450 # (each memory linked only to entities IT mentions); legacy path
451 # uses the Cartesian product.
452 sfe_associations = record.get("sfe_associations")
453 if sfe_associations is not None:
454 associations = [
455 MemoryEntityAssociation(
456 memory_id=memory_ids[mem_idx],
457 entity_id=entity_ids[ent_idx],
458 )
459 for ent_idx, mem_idx in sfe_associations
460 if 0 <= ent_idx < len(entity_ids) and 0 <= mem_idx < len(memory_ids)
461 ]
462 else:
463 associations = [
464 MemoryEntityAssociation(memory_id=mid, entity_id=eid) for mid in memory_ids for eid in entity_ids
465 ]
466 async with self._profiler.time("entity_link_mem"):
467 await self.graph_store.link_memories_to_entities(associations, request.bank_id)
468 # Same cap as the single-record retain path — see comment there.
469 # Without it the all-pairs Cartesian product dominates retain
470 # wall on entity-dense workloads (LME profile 2026-05-06).
471 if self.entity_cooccurrence_enabled:
472 pairs = _build_cooccurrence_pairs(
473 entity_ids,
474 self.entity_cooccurrence_max_entities,
475 )
476 if pairs:
477 links = [EntityLink(entity_a=a, entity_b=b, link_type="co_occurs") for a, b in pairs]
478 async with self._profiler.time("entity_co_occur"):
479 await self.graph_store.store_links(links, request.bank_id)
481 # MemoryLinks (caused_by). Two paths:
482 # - SFE supplied them inline → resolve indices, store.
483 # - Legacy fact_causal_extraction LLM call per record.
484 sfe_caused_by = record.get("sfe_caused_by")
485 chunks: list[str] = record.get("chunks") or []
486 if sfe_caused_by:
487 memory_links = [
488 MemoryLink(
489 source_memory_id=memory_ids[src_idx],
490 target_memory_id=memory_ids[tgt_idx],
491 link_type="caused_by",
492 confidence=conf,
493 weight=1.0,
494 created_at=datetime.now(timezone.utc),
495 metadata={"bank_id": request.bank_id, "source": "fact_extraction"},
496 )
497 for src_idx, tgt_idx, conf in sfe_caused_by
498 if 0 <= src_idx < len(memory_ids) and 0 <= tgt_idx < len(memory_ids) and src_idx != tgt_idx
499 ]
500 if memory_links:
501 try:
502 await self.graph_store.store_memory_links(
503 memory_links,
504 request.bank_id,
505 )
506 except Exception as exc:
507 _logger.warning("storing memory_links failed: %s", exc)
508 elif (
509 self.causal_links_enabled
510 and len(memory_ids) > 1
511 and len(chunks) == len(memory_ids)
512 and self.llm_provider is not None
513 ):
514 try:
515 from astrocyte.pipeline.fact_causal_extraction import (
516 build_memory_links_from_relations,
517 extract_fact_causal_relations,
518 )
520 relations = await extract_fact_causal_relations(
521 chunks,
522 self.llm_provider,
523 max_pairs_per_fact=self.causal_max_pairs_per_memory,
524 min_confidence=self.causal_min_confidence,
525 )
526 memory_links = build_memory_links_from_relations(
527 relations,
528 memory_ids,
529 bank_id=request.bank_id,
530 )
531 except Exception as exc:
532 _logger.warning("fact-level causal extraction failed: %s", exc)
533 memory_links = []
534 if memory_links:
535 try:
536 await self.graph_store.store_memory_links(
537 memory_links,
538 request.bank_id,
539 )
540 except Exception as exc:
541 _logger.warning("storing memory_links failed: %s", exc)
543 # Legacy two-stage flow: post-store alias resolution. Skipped when
544 # ``canonical_resolution=True`` because IDs are already canonical.
545 if self.entity_resolver is not None and not self.entity_resolver.canonical_resolution:
546 try:
547 await self.entity_resolver.resolve(
548 new_entities=entities,
549 source_text=prepared.text,
550 bank_id=request.bank_id,
551 graph_store=self.graph_store,
552 llm_provider=self.llm_provider,
553 event_date=request.occurred_at,
554 )
555 except Exception as exc:
556 _logger.warning("entity resolution failed during retain: %s", exc)
558 async def _provision_source_provenance(
559 self,
560 request: RetainRequest,
561 prepared_text: str,
562 chunks: list[str],
563 ) -> list[str | None]:
564 """Stamp the (document, chunks) pair into the SourceStore.
566 Returns the per-chunk ``chunk_id`` list, parallel to ``chunks``.
567 On failure (or when the SourceStore is not configured / the flag
568 is off), returns ``[None] * len(chunks)`` so retain continues
569 with anonymous vectors — provenance must never break ingest.
571 The document id is deterministic (SHA-256 of the prepared text,
572 prefixed with ``doc:``) so the same input ingested twice resolves
573 to the same SourceDocument and the dedup probe is hit. Chunk ids
574 are ``{document_id}:{i}`` for the same reason.
575 """
576 n = len(chunks)
577 empty: list[str | None] = [None] * n
578 if self.source_store is None or not self.source_retain_provenance:
579 return empty
581 # Late import: keeps the module import-cost low and avoids a hard
582 # dependency on the M10 types from the orchestrator's surface area.
583 from astrocyte.types import SourceChunk, SourceDocument
585 try:
586 doc_hash = hashlib.sha256(prepared_text.encode("utf-8")).hexdigest()
587 doc_id_default = f"doc:{doc_hash[:16]}"
588 document = SourceDocument(
589 id=doc_id_default,
590 bank_id=request.bank_id,
591 title=None,
592 source_uri=request.source,
593 content_hash=doc_hash,
594 content_type=request.content_type,
595 metadata=None,
596 )
597 # store_document is idempotent on (bank_id, content_hash); the
598 # returned id may differ from doc_id_default if a prior live
599 # row with the same hash already exists.
600 stored_doc_id = await self.source_store.store_document(document) # type: ignore[union-attr]
602 source_chunks: list[SourceChunk] = []
603 for i, chunk_text_str in enumerate(chunks):
604 chunk_hash = hashlib.sha256(chunk_text_str.encode("utf-8")).hexdigest()
605 source_chunks.append(
606 SourceChunk(
607 id=f"{stored_doc_id}:{i}",
608 bank_id=request.bank_id,
609 document_id=stored_doc_id,
610 chunk_index=i,
611 text=chunk_text_str,
612 content_hash=chunk_hash,
613 metadata=None,
614 )
615 )
616 chunk_ids = await self.source_store.store_chunks(source_chunks) # type: ignore[union-attr]
617 if len(chunk_ids) != n:
618 # Should not happen for well-behaved adapters; defend
619 # against shape mismatch by falling back rather than
620 # silently misaligning chunk_ids onto VectorItems.
621 _logger.warning(
622 "source_store.store_chunks returned %d ids, expected %d — falling back",
623 len(chunk_ids),
624 n,
625 )
626 return empty
627 return list(chunk_ids)
628 except Exception as exc:
629 # Defensive: provenance is best-effort. Log and degrade to
630 # anonymous vectors so a SourceStore outage doesn't break ingest.
631 _logger.warning("source-aware retain failed; continuing without provenance: %s", exc)
632 return empty
634 async def _store_entity_graph(
635 self,
636 *,
637 entities: list[Entity],
638 bank_id: str,
639 memory_ids: list[str],
640 chunks: list[str],
641 sfe_associations: list[tuple[int, int]] | None,
642 sfe_caused_by: list[tuple[int, int, float]] | None,
643 source_text: str,
644 occurred_at: datetime | None,
645 ) -> None:
646 """Retain phase 5 — persist entities, associations, and links.
648 No-op without a graph store. Covers: name embeddings + canonical
649 resolution, entity storage, memory↔entity associations (SFE
650 fact-grained or legacy Cartesian), co-occurrence links, causal
651 MemoryLinks (SFE-inline or legacy LLM pass), and the legacy
652 two-stage entity-resolution cascade.
653 """
654 if not (self.graph_store and entities):
655 return
656 # 5a. Attach name embeddings before persistence so the
657 # Hindsight-inspired entity-resolution cascade has the cheap
658 # cosine-similarity tier available. Only fires when an
659 # entity_resolver is wired up — without one, the embedding
660 # would be persisted but never used, wasting API budget.
661 if self.entity_resolver is not None:
662 async with self._profiler.time("entity_emb"):
663 await self._attach_entity_name_embeddings(entities)
664 # Path B (Hindsight): rewrite tentative IDs to canonicals
665 # before storage. Skipped when canonical_resolution=False
666 # to preserve the legacy two-stage (store → resolve aliases)
667 # flow.
668 if self.entity_resolver.canonical_resolution:
669 try:
670 async with self._profiler.time("entity_resolve"):
671 await self.entity_resolver.resolve_canonical_ids_in_place(
672 new_entities=entities,
673 bank_id=bank_id,
674 graph_store=self.graph_store,
675 event_date=occurred_at,
676 )
677 except Exception as exc:
678 _logger.warning(
679 "canonical resolution failed during retain (falling back to tentative IDs): %s",
680 exc,
681 )
683 async with self._profiler.time("entity_store"):
684 entity_ids = await self.graph_store.store_entities(entities, bank_id)
686 # Link memories to entities. Two paths:
687 # - Structured fact extraction provides per-fact-per-entity
688 # associations (each memory linked only to entities IT mentions),
689 # matching Hindsight's fact-grained granularity.
690 # - Legacy path uses the Cartesian product (every memory linked
691 # to every entity in the batch), which works when extraction
692 # was over the whole text rather than per fact.
693 if sfe_associations is not None:
694 associations = [
695 MemoryEntityAssociation(
696 memory_id=memory_ids[mem_idx],
697 entity_id=entity_ids[ent_idx],
698 )
699 for ent_idx, mem_idx in sfe_associations
700 if 0 <= ent_idx < len(entity_ids) and 0 <= mem_idx < len(memory_ids)
701 ]
702 else:
703 associations = [
704 MemoryEntityAssociation(memory_id=mid, entity_id=eid) for mid in memory_ids for eid in entity_ids
705 ]
706 async with self._profiler.time("entity_link_mem"):
707 await self.graph_store.link_memories_to_entities(associations, bank_id)
709 # Create co-occurrence links between entities. The 2026-05-06
710 # LME retain-profile measured the unbounded all-pairs
711 # Cartesian product at 34% of retain wall (52s tail at the
712 # session-100 mark). The cap in
713 # ``_build_cooccurrence_pairs`` bounds the work to O(K²) per
714 # retain regardless of corpus size.
715 if self.entity_cooccurrence_enabled:
716 pairs = _build_cooccurrence_pairs(
717 entity_ids,
718 self.entity_cooccurrence_max_entities,
719 )
720 if pairs:
721 links = [EntityLink(entity_a=a, entity_b=b, link_type="co_occurs") for a, b in pairs]
722 async with self._profiler.time("entity_co_occur"):
723 await self.graph_store.store_links(links, bank_id)
725 # Causal MemoryLinks. Two paths:
726 # - Structured fact extraction supplies them inline (no
727 # extra LLM call); resolve indices to memory IDs and store.
728 # - Legacy path runs a separate fact_causal_extraction LLM call.
729 if sfe_caused_by is not None and sfe_caused_by:
730 memory_links = [
731 MemoryLink(
732 source_memory_id=memory_ids[src_idx],
733 target_memory_id=memory_ids[tgt_idx],
734 link_type="caused_by",
735 confidence=conf,
736 weight=1.0,
737 created_at=datetime.now(timezone.utc),
738 metadata={"bank_id": bank_id, "source": "fact_extraction"},
739 )
740 for src_idx, tgt_idx, conf in sfe_caused_by
741 if 0 <= src_idx < len(memory_ids) and 0 <= tgt_idx < len(memory_ids) and src_idx != tgt_idx
742 ]
743 if memory_links:
744 try:
745 await self.graph_store.store_memory_links(
746 memory_links,
747 bank_id,
748 )
749 except Exception as exc:
750 _logger.warning("storing memory_links failed: %s", exc)
751 elif (
752 self.causal_links_enabled
753 and len(memory_ids) > 1
754 and len(chunks) == len(memory_ids)
755 and self.llm_provider is not None
756 ):
757 # Legacy fact-causal-extraction LLM pass (separate call).
758 try:
759 from astrocyte.pipeline.fact_causal_extraction import (
760 build_memory_links_from_relations,
761 extract_fact_causal_relations,
762 )
764 relations = await extract_fact_causal_relations(
765 chunks,
766 self.llm_provider,
767 max_pairs_per_fact=self.causal_max_pairs_per_memory,
768 min_confidence=self.causal_min_confidence,
769 )
770 memory_links = build_memory_links_from_relations(
771 relations,
772 memory_ids,
773 bank_id=bank_id,
774 )
775 except Exception as exc:
776 _logger.warning("fact-level causal extraction failed: %s", exc)
777 memory_links = []
778 if memory_links:
779 try:
780 await self.graph_store.store_memory_links(
781 memory_links,
782 bank_id,
783 )
784 except Exception as exc:
785 _logger.warning("storing memory_links failed: %s", exc)
787 # 5b. Entity resolution (M11) — Hindsight-inspired tiered cascade.
788 # Skipped when canonical_resolution=True (Path B) because IDs
789 # are already canonical from the pre-store pass above. Otherwise
790 # the legacy two-stage path runs the cascade and writes
791 # alias_of links for matched candidates.
792 if self.entity_resolver is not None and not self.entity_resolver.canonical_resolution:
793 try:
794 await self.entity_resolver.resolve(
795 new_entities=entities,
796 source_text=source_text,
797 bank_id=bank_id,
798 graph_store=self.graph_store,
799 llm_provider=self.llm_provider,
800 event_date=occurred_at,
801 )
802 except Exception as exc:
803 # Resolution failures must never abort a retain — degrade gracefully.
804 _logger.warning("entity resolution failed during retain: %s", exc)
806 def _spawn_observation_consolidation(
807 self,
808 *,
809 chunks: list[str],
810 embeddings: list[list[float]],
811 memory_ids: list[str],
812 bank_id: str,
813 tags: list[str] | None,
814 ) -> None:
815 """Fire-and-forget observation consolidation (retain phase 7).
817 Runs after the retain response is returned so it never adds latency to
818 the caller; failures must never surface — the raw memories are already
819 stored and consolidation is best-effort. The representative vector is
820 the first chunk's embedding, sufficient for the observation similarity
821 search.
822 """
823 if self._observation_consolidator is None or not memory_ids or not chunks:
824 return
826 representative_vec = embeddings[0]
827 consolidator = self._observation_consolidator
828 first_chunk = chunks[0]
829 all_ids = list(memory_ids)
830 vs = self.vector_store
831 llm = self.llm_provider
833 async def _run_consolidation() -> None:
834 try:
835 await consolidator.consolidate(
836 new_memory_text=first_chunk,
837 new_memory_ids=all_ids,
838 bank_id=bank_id,
839 vector_store=vs,
840 llm_provider=llm,
841 query_vector=representative_vec,
842 scope="|".join(sorted(tags)) if tags else None,
843 )
844 except Exception as exc:
845 _logger.warning("Observation consolidation task failed for bank %s: %s", bank_id, exc)
847 self._submit_consolidation(_run_consolidation, bank_id)
849 def _submit_consolidation(self, factory, bank_id: str) -> None:
850 """Run now, defer, or — only past both bounds — shed.
852 Three tiers, each bounded, because the original defect was an
853 *unbounded* backlog starving foreground retain() through the provider's
854 own concurrency limit:
856 1. **Run** while fewer than ``max_pending_consolidations`` are in flight.
857 2. **Defer** into a FIFO of at most ``max_deferred_consolidations``
858 zero-arg factories. Deferred entries hold no provider slot, so they
859 cannot starve foreground work; they start as in-flight tasks finish
860 (``_on_consolidation_done``). This replaces silently dropping the
861 overflow — the n=250 self-eval measured ~1,000+ sheds per 50 items,
862 roughly 30% of the observation layer never built.
863 3. **Shed** only when the deferred queue is also full.
865 What deferral does NOT fix: a *sustained* arrival rate above the
866 consolidation service rate. A finite queue absorbs a burst, but under
867 continuous saturated ingest it fills and tier 3 resumes at the old
868 rate. That case needs faster decisions or more provider concurrency;
869 rising ``consolidations_shed`` alongside a full queue is the signal.
870 """
871 cap = getattr(self, "max_pending_consolidations", 0)
872 if not cap or len(self._background_tasks) < cap:
873 self._start_consolidation(factory)
874 return
876 deferred = getattr(self, "_deferred_consolidations", None)
877 dcap = getattr(self, "max_deferred_consolidations", 0)
878 if deferred is not None and dcap and len(deferred) < dcap and not getattr(self, "_consolidation_closed", False):
879 deferred.append(factory)
880 self.consolidations_deferred = getattr(self, "consolidations_deferred", 0) + 1
881 return
883 shed = getattr(self, "consolidations_shed", 0) + 1
884 self.consolidations_shed = shed
885 # WARNING on a geometric backoff (1st, 10th, 100th, …). DEBUG was the
886 # wrong level: shedding means part of the observation layer is never
887 # built, so recall is quietly thinner than the operator thinks. It was
888 # invisible through an entire benchmark run while DEBUG was off.
889 digits = str(shed)
890 if digits[0] == "1" and set(digits[1:]) <= {"0"}:
891 _logger.warning(
892 "Observation consolidation shed (%d so far): %d in flight (cap %d), "
893 "%d deferred (cap %d). Derived memory is incomplete for bank %s — "
894 "ingest is outpacing consolidation; a full deferred queue means this "
895 "is sustained, not a burst.",
896 shed,
897 len(self._background_tasks),
898 cap,
899 len(deferred) if deferred is not None else 0,
900 dcap,
901 bank_id,
902 )
904 def _start_consolidation(self, factory) -> None:
905 task = asyncio.create_task(factory())
906 self._background_tasks.add(task)
907 task.add_done_callback(self._on_consolidation_done)
909 def _on_consolidation_done(self, task) -> None:
910 self._background_tasks.discard(task)
911 self._drain_deferred()
913 def _drain_deferred(self) -> None:
914 """Promote deferred work into freed in-flight slots, oldest first."""
915 # Defense in depth only. The load-bearing protection is the ORDER in
916 # ``PipelineOrchestrator.shutdown()`` — it closes and clears this queue
917 # before cancelling tasks, so by the time a cancelled task's callback
918 # lands here the queue is already empty (a negative control confirms
919 # removing that ordering spawns tasks mid-shutdown; removing this check
920 # alone does not). Kept so a future reordering fails safe.
921 if getattr(self, "_consolidation_closed", False):
922 return
923 deferred = getattr(self, "_deferred_consolidations", None)
924 if not deferred:
925 return
926 cap = getattr(self, "max_pending_consolidations", 0)
927 while deferred and (not cap or len(self._background_tasks) < cap):
928 self._start_consolidation(deferred.popleft())
930 def _config_dedup_action(self, bank_id: str) -> str:
931 """``signal_quality.dedup.action`` for ``bank_id`` — the bank's own
932 setting if it has one, else the top-level one. Used when no matched MIP
933 rule sets ``dedup.action``."""
934 bank = getattr(self, "dedup_by_bank", {}).get(bank_id)
935 return bank.pipeline_action if bank is not None else getattr(self, "dedup_action", "skip_chunk")
937 async def _find_duplicate_chunks(
938 self,
939 bank_id: str,
940 chunks: list[str],
941 embeddings: list[list[float]],
942 threshold_override: float | None,
943 ) -> list[bool]:
944 """Which chunks are near-duplicates of memories the bank already holds
945 (see :meth:`_find_duplicate_ids`)."""
946 found = await self._find_duplicate_ids(bank_id, chunks, embeddings, threshold_override)
947 return [memory_id is not None for memory_id in found]
949 async def _find_duplicate_ids(
950 self,
951 bank_id: str,
952 chunks: list[str],
953 embeddings: list[list[float]],
954 threshold_override: float | None,
955 ) -> list[str | None]:
956 """For each chunk, the id of a memory the bank already holds that it
957 near-duplicates, or ``None``.
959 The in-process ``DedupDetector`` cache only knows what this process
960 retained, so a restarted daemon, a CLI run or a gateway restart would
961 store a duplicate of a memory written earlier. Chunks the cache does
962 not match are checked against their nearest neighbours in the vector
963 store, under the same threshold and negation guard. Observations and
964 mental models are skipped: a raw memory restating a synthesis is new
965 evidence, and the cache never held them either. A failed lookup counts
966 as "not a duplicate" — a redundant row is cheaper than a lost fact.
968 Costs one ``search_similar`` per chunk the cache missed;
969 ``dedup_consult_store = False`` turns it off. ``dedup_enabled = False``
970 skips the whole check: nothing is a duplicate. A bank listed in
971 ``dedup_by_bank`` uses its own enabled / threshold / consult_store; a
972 MIP ``threshold_override`` still beats the bank's threshold.
973 """
974 bank = getattr(self, "dedup_by_bank", {}).get(bank_id)
975 if bank is not None:
976 enabled, consult_store = bank.enabled, bank.consult_store
977 if threshold_override is None:
978 threshold_override = bank.similarity_threshold
979 else:
980 enabled = getattr(self, "dedup_enabled", True)
981 consult_store = getattr(self, "dedup_consult_store", True)
982 if not enabled:
983 return [None] * len(chunks)
984 dups = [
985 self._dedup.find_duplicate(bank_id, emb, threshold_override=threshold_override, text=chunk)
986 for chunk, emb in zip(chunks, embeddings)
987 ]
988 misses = [i for i, dup in enumerate(dups) if dup is None]
989 if not misses or not consult_store or self.vector_store is None:
990 return dups
992 async def nearest(i: int) -> list[VectorHit]:
993 try:
994 return await self.vector_store.search_similar(embeddings[i], bank_id, limit=_STORE_DEDUP_NEIGHBOURS)
995 except Exception as exc:
996 _logger.warning("retain dedup: store lookup failed in bank %r (%s); keeping the chunk", bank_id, exc)
997 return []
999 # Sequential on purpose: SqliteStore scans the bank per query, and
1000 # concurrent scans thrash — 10 chunks against 50k rows took 12 s
1001 # gathered vs 1.5 s one after another.
1002 async with self._profiler.time("dedup_store"):
1003 neighbours = [await nearest(i) for i in misses]
1004 for i, hits in zip(misses, neighbours):
1005 for hit in hits:
1006 if hit.fact_type in _NOT_RETAINED_LAYERS or hit.memory_layer in _NOT_RETAINED_LAYERS:
1007 continue
1008 if self._dedup.matches([(hit.score, hit.text)], threshold_override=threshold_override, text=chunks[i])[0]:
1009 dups[i] = hit.id
1010 break
1011 return dups
1013 async def retain(self, request: RetainRequest) -> RetainResult:
1014 """Retain pipeline: normalize → chunk → extract entities → embed → store."""
1015 # 0–1. Raw → normalizer → profile metadata/tags (M3 extraction chain)
1016 profile: ExtractionProfileConfig | None = None
1017 name = getattr(request, "extraction_profile", None)
1018 profile_table = merged_user_and_builtin_profiles(self.extraction_profiles)
1019 if name:
1020 profile = profile_table.get(name)
1022 # MIP pipeline overrides (from RoutingDecision) — optional, no behavior
1023 # change when absent. ``chunker`` and ``dedup`` are consumed here;
1024 # ``rerank`` and ``reflect`` apply to recall/reflect (Phase 2).
1025 mip_pipeline = request.mip_pipeline
1026 mip_chunker = mip_pipeline.chunker if mip_pipeline else None
1027 mip_dedup = mip_pipeline.dedup if mip_pipeline else None
1028 dedup_threshold_override = mip_dedup.threshold if mip_dedup else None
1029 dedup_action = (mip_dedup.action if mip_dedup else None) or self._config_dedup_action(request.bank_id)
1031 prepared = prepare_retain_input(
1032 request,
1033 profile,
1034 graph_store_configured=self.graph_store is not None,
1035 )
1036 chunking = resolve_retain_chunking(
1037 prepared.effective_content_type,
1038 profile=profile,
1039 default_strategy=self.chunk_strategy,
1040 default_max_chunk_size=self.max_chunk_size,
1041 mip_chunker=mip_chunker,
1042 )
1043 chunk_kwargs: dict[str, int] = {"max_chunk_size": chunking.max_size}
1044 if chunking.overlap is not None:
1045 chunk_kwargs["overlap"] = chunking.overlap
1047 # Structured 5-dim fact extraction (opt-in). Replaces chunking +
1048 # entity extraction with a single LLM call producing facts.
1049 # Falls back to legacy chunk_text + extract_entities when
1050 # disabled or when extraction returns no facts.
1051 # SFE uses its OWN chunking strategy (not the profile's), because
1052 # SFE has different granularity needs: large chunks give the LLM
1053 # full context for metadata extraction. Default "paragraph" wins
1054 # +2.5pts on LoCoMo over the profile-driven "dialogue" choice.
1055 async with self._profiler.time("sfe"):
1056 (
1057 sfe_fact_texts,
1058 sfe_entities,
1059 sfe_associations,
1060 sfe_caused_by,
1061 sfe_overlays,
1062 ) = await self._structured_fact_extraction_for_text(
1063 prepared,
1064 request,
1065 chunk_strategy=self.structured_fact_extraction_chunk_strategy,
1066 chunk_max_size=(
1067 self.structured_fact_extraction_chunk_max_size
1068 if self.structured_fact_extraction_chunk_max_size is not None
1069 else chunking.max_size
1070 ),
1071 chunk_overlap=chunking.overlap,
1072 )
1073 if sfe_fact_texts is not None:
1074 chunks = list(sfe_fact_texts)
1075 else:
1076 chunks = chunk_text(prepared.text, strategy=chunking.strategy, **chunk_kwargs)
1077 if not chunks:
1078 return RetainResult(stored=False, error="No content after chunking")
1080 # 2. Generate embeddings for all chunks
1081 async with self._profiler.time("embed"):
1082 embeddings = await generate_embeddings(chunks, self.llm_provider)
1084 # 2b. Per-chunk dedup — behavior depends on dedup_action (MIP rule,
1085 # else signal_quality.dedup.action for the bank):
1086 # "skip_chunk" (default): drop duplicate chunks, keep the rest
1087 # "skip": if any chunk is a duplicate, reject the entire retain
1088 # "warn": keep all chunks regardless of duplicates
1089 # "update": not yet implemented; falls back to "skip_chunk"
1090 keep_indices: list[int] = []
1091 any_duplicate = False
1092 duplicates = await self._find_duplicate_chunks(request.bank_id, chunks, embeddings, dedup_threshold_override)
1093 for i, is_dup in enumerate(duplicates):
1094 if is_dup:
1095 any_duplicate = True
1096 if dedup_action == "warn" or not is_dup:
1097 keep_indices.append(i)
1099 if dedup_action == "skip" and any_duplicate:
1100 return RetainResult(stored=False, deduplicated=True, error="Duplicate chunk(s) found; rule action=skip")
1102 if not keep_indices:
1103 return RetainResult(stored=False, deduplicated=True, error="All chunks are near-duplicates")
1105 chunks = [chunks[i] for i in keep_indices]
1106 embeddings = [embeddings[i] for i in keep_indices]
1107 if sfe_overlays is not None:
1108 sfe_overlays = [sfe_overlays[i] for i in keep_indices]
1110 # Remap structured-fact-extraction indices through dedup. The
1111 # association and caused_by lists referenced positions in the
1112 # ORIGINAL chunks; after dedup, indices need remapping to the
1113 # surviving positions (or the records dropped entirely).
1114 if sfe_associations is not None or sfe_caused_by is not None:
1115 old_to_new = {old: new for new, old in enumerate(keep_indices)}
1116 if sfe_associations is not None:
1117 sfe_associations = [
1118 (ent_idx, old_to_new[mem_idx]) for ent_idx, mem_idx in sfe_associations if mem_idx in old_to_new
1119 ]
1120 if sfe_caused_by is not None:
1121 sfe_caused_by = [
1122 (old_to_new[src], old_to_new[tgt], conf)
1123 for src, tgt, conf in sfe_caused_by
1124 if src in old_to_new and tgt in old_to_new
1125 ]
1127 # 3. Extract entities (profile can disable; metadata profiles avoid LLM calls).
1128 # Structured fact extraction provides entities pre-extracted —
1129 # skip the second LLM call when it ran successfully.
1130 if sfe_entities is not None:
1131 entities = list(sfe_entities)
1132 elif profile is not None and profile.entity_extraction == "metadata":
1133 entities = _entities_from_metadata(prepared.metadata)
1134 elif prepared.extract_entities:
1135 entities = await extract_entities(prepared.text, self.llm_provider)
1136 else:
1137 entities = []
1139 profile_tier = profile.authority_tier if profile else None
1140 chunk_metadata = merge_retain_metadata_authority_tier(
1141 prepared.metadata,
1142 bank_id=request.bank_id,
1143 profile_authority_tier=profile_tier,
1144 recall_authority=self.recall_authority,
1145 )
1147 # Persist MIP provenance so recall can warn on rule-version drift (P2).
1148 if request.mip_rule_name or (mip_pipeline and mip_pipeline.version is not None):
1149 chunk_metadata = dict(chunk_metadata or {})
1150 if request.mip_rule_name:
1151 chunk_metadata["_mip.rule"] = request.mip_rule_name
1152 if mip_pipeline and mip_pipeline.version is not None:
1153 chunk_metadata["_mip.pipeline_version"] = int(mip_pipeline.version)
1155 # Stamp ``_created_at`` so (1) temporal retrieval can rank by
1156 # recency, (2) MIP forget.min_age_days can enforce age guards, and
1157 # (3) lifecycle TTL has a consistent retain timestamp. Callers that
1158 # already set ``_created_at`` (imports / replay) win via setdefault.
1159 # Mirrors InMemoryEngineProvider.retain so pipeline- and engine-
1160 # backed deployments share the same observability contract.
1161 chunk_metadata = dict(chunk_metadata or {})
1162 chunk_metadata.setdefault(
1163 "_created_at",
1164 datetime.now(timezone.utc).isoformat(),
1165 )
1166 # Ties this retain's chunks together. ``_created_at`` cannot: on a
1167 # coarse clock (Windows ticks every ~15 ms) consecutive retains share it.
1168 chunk_metadata.setdefault("_retain_id", uuid.uuid4().hex)
1169 # Where the content came from (a URL, a file, an agent): recall hits
1170 # carry it as MemoryHit.source. It reached only the optional source
1171 # store before, so on the default path every hit's source was None.
1172 if request.source:
1173 chunk_metadata.setdefault(SOURCE_KEY, request.source)
1175 # 3b. M10: persist source-document + chunk provenance, get back per-chunk
1176 # ``chunk_id``s that we'll stamp onto each VectorItem so recall can
1177 # later resolve "which document/chunk did this memory come from".
1178 # Returns ``[None] * len(chunks)`` when source_store is not configured
1179 # or the flag is off — fully backward-compatible.
1180 chunk_ids = await self._provision_source_provenance(
1181 request,
1182 prepared.text,
1183 chunks,
1184 )
1186 # 4. Store vectors
1187 memory_ids: list[str] = []
1188 items: list[VectorItem] = []
1189 for index, (chunk, embedding, chunk_id) in enumerate(zip(chunks, embeddings, chunk_ids)):
1190 mem_id = uuid.uuid4().hex[:16]
1191 memory_ids.append(mem_id)
1192 item_meta, item_fact_type, item_occurred = _apply_fact_overlay(
1193 _chunk_metadata(chunk_metadata, index, len(chunks)),
1194 sfe_overlays[index] if sfe_overlays is not None else None,
1195 fact_type=prepared.fact_type,
1196 occurred_at=request.occurred_at,
1197 profile_fact_type=profile.fact_type if profile is not None else None,
1198 )
1199 items.append(
1200 VectorItem(
1201 id=mem_id,
1202 bank_id=request.bank_id,
1203 vector=embedding,
1204 text=chunk,
1205 metadata=item_meta,
1206 tags=prepared.tags,
1207 fact_type=item_fact_type,
1208 occurred_at=item_occurred,
1209 retained_at=datetime.now(timezone.utc), # M9: wall-clock store time
1210 chunk_id=chunk_id, # M10: source-chunk backreference
1211 )
1212 )
1214 async with self._profiler.time("store_vec"):
1215 await self.vector_store.store_vectors(items)
1217 # Hindsight-parity semantic-kNN graph (C3a). Best-effort: links
1218 # the new memories to their nearest existing neighbors so the
1219 # link-expansion retrieval CTE has a precomputed semantic
1220 # signal at recall time. Skipped when disabled.
1221 await self._persist_semantic_links(
1222 request.bank_id,
1223 [item.id for item in items],
1224 [item.vector for item in items],
1225 )
1227 # 4b. Mirror chunks into the document store so keyword (BM25) search
1228 # has content to retrieve from. Previously recall.parallel_retrieve
1229 # ran a keyword strategy when ``document_store`` was configured —
1230 # but nothing was ever stored there. Fixed here: every chunk that
1231 # lands in the vector store also lands in the document store, with
1232 # the same memory_id so RRF fusion can dedupe across strategies.
1233 # Failures don't abort the retain — degrading to semantic-only
1234 # retrieval is preferable to losing the whole memory.
1235 if self.document_store is not None:
1236 from astrocyte.types import Document
1238 for item in items:
1239 try:
1240 await self.document_store.store_document(
1241 Document(
1242 id=item.id,
1243 text=item.text,
1244 metadata=item.metadata,
1245 tags=item.tags,
1246 ),
1247 request.bank_id,
1248 )
1249 except Exception as exc:
1250 _logger.warning(
1251 "document_store.store_document failed for chunk %s: %s "
1252 "(keyword retrieval will miss this chunk)",
1253 item.id,
1254 exc,
1255 )
1257 # 5. Store entities and links (if graph store configured) — see
1258 # _store_entity_graph for the phase body.
1259 await self._store_entity_graph(
1260 entities=entities,
1261 bank_id=request.bank_id,
1262 memory_ids=memory_ids,
1263 chunks=chunks,
1264 sfe_associations=sfe_associations,
1265 sfe_caused_by=sfe_caused_by,
1266 source_text=prepared.text,
1267 occurred_at=request.occurred_at,
1268 )
1270 # 6. Update dedup cache with stored embeddings
1271 for mem_id, emb, chunk in zip(memory_ids, embeddings, chunks):
1272 self._dedup.add(request.bank_id, mem_id, emb, text=chunk)
1274 # 7. Observation consolidation — see _spawn_observation_consolidation.
1275 self._spawn_observation_consolidation(
1276 chunks=chunks,
1277 embeddings=embeddings,
1278 memory_ids=memory_ids,
1279 bank_id=request.bank_id,
1280 tags=request.tags,
1281 )
1283 return RetainResult(
1284 stored=True,
1285 memory_id=memory_ids[0] if memory_ids else None,
1286 )
1288 async def retain_many(self, requests: list[RetainRequest]) -> list[RetainResult]:
1289 """Retain multiple requests while batching embedding generation and vector writes."""
1291 if not requests:
1292 return []
1294 results: list[RetainResult | None] = [None] * len(requests)
1295 records: list[dict[str, Any]] = []
1296 all_chunks: list[str] = []
1297 profile_table = merged_user_and_builtin_profiles(self.extraction_profiles)
1299 for index, request in enumerate(requests):
1300 profile = profile_table.get(request.extraction_profile) if request.extraction_profile else None
1301 mip_pipeline = request.mip_pipeline
1302 mip_chunker = mip_pipeline.chunker if mip_pipeline else None
1303 prepared = prepare_retain_input(
1304 request,
1305 profile,
1306 graph_store_configured=self.graph_store is not None,
1307 )
1308 chunking = resolve_retain_chunking(
1309 prepared.effective_content_type,
1310 profile=profile,
1311 default_strategy=self.chunk_strategy,
1312 default_max_chunk_size=self.max_chunk_size,
1313 mip_chunker=mip_chunker,
1314 )
1315 chunk_kwargs: dict[str, int] = {"max_chunk_size": chunking.max_size}
1316 if chunking.overlap is not None:
1317 chunk_kwargs["overlap"] = chunking.overlap
1319 # Structured 5-dim fact extraction (opt-in) — same branch as
1320 # retain(); produces fact texts + pre-extracted entities +
1321 # caused_by index pairs for this record.
1322 # SFE uses its own chunking strategy (see retain() above).
1323 async with self._profiler.time("sfe"):
1324 (
1325 sfe_fact_texts,
1326 sfe_entities,
1327 sfe_associations,
1328 sfe_caused_by,
1329 sfe_overlays,
1330 ) = await self._structured_fact_extraction_for_text(
1331 prepared,
1332 request,
1333 chunk_strategy=self.structured_fact_extraction_chunk_strategy,
1334 chunk_max_size=(
1335 self.structured_fact_extraction_chunk_max_size
1336 if self.structured_fact_extraction_chunk_max_size is not None
1337 else chunking.max_size
1338 ),
1339 chunk_overlap=chunking.overlap,
1340 )
1341 if sfe_fact_texts is not None:
1342 chunks = list(sfe_fact_texts)
1343 else:
1344 chunks = chunk_text(prepared.text, strategy=chunking.strategy, **chunk_kwargs)
1345 if not chunks:
1346 results[index] = RetainResult(stored=False, error="No content after chunking")
1347 continue
1349 start = len(all_chunks)
1350 all_chunks.extend(chunks)
1351 records.append(
1352 {
1353 "index": index,
1354 "request": request,
1355 "profile": profile,
1356 "prepared": prepared,
1357 "chunks": chunks,
1358 "start": start,
1359 "end": len(all_chunks),
1360 # Carry the SFE artefacts through to the second loop and
1361 # _process_record_entities. None entries mean the record
1362 # took the legacy chunk_text + extract_entities path.
1363 "sfe_entities": sfe_entities,
1364 "sfe_associations": sfe_associations,
1365 "sfe_caused_by": sfe_caused_by,
1366 "sfe_overlays": sfe_overlays,
1367 }
1368 )
1370 if not records:
1371 return [result or RetainResult(stored=False, error="No content after chunking") for result in results]
1373 async with self._profiler.time("embed"):
1374 all_embeddings = await generate_embeddings(all_chunks, self.llm_provider)
1375 stored_records: list[dict[str, Any]] = []
1376 all_items: list[VectorItem] = []
1378 for record in records:
1379 request: RetainRequest = record["request"]
1380 profile: ExtractionProfileConfig | None = record["profile"]
1381 prepared = record["prepared"]
1382 chunks = record["chunks"]
1383 embeddings = all_embeddings[record["start"] : record["end"]]
1384 mip_pipeline = request.mip_pipeline
1385 mip_dedup = mip_pipeline.dedup if mip_pipeline else None
1386 dedup_threshold_override = mip_dedup.threshold if mip_dedup else None
1387 dedup_action = (mip_dedup.action if mip_dedup else None) or self._config_dedup_action(request.bank_id)
1389 keep_indices: list[int] = []
1390 any_duplicate = False
1391 duplicates = await self._find_duplicate_chunks(
1392 request.bank_id, chunks, embeddings, dedup_threshold_override
1393 )
1394 for chunk_index, is_dup in enumerate(duplicates):
1395 if is_dup:
1396 any_duplicate = True
1397 if dedup_action == "warn" or not is_dup:
1398 keep_indices.append(chunk_index)
1400 result_index = record["index"]
1401 if dedup_action == "skip" and any_duplicate:
1402 results[result_index] = RetainResult(
1403 stored=False,
1404 deduplicated=True,
1405 error="Duplicate chunk(s) found; rule action=skip",
1406 )
1407 continue
1408 if not keep_indices:
1409 results[result_index] = RetainResult(
1410 stored=False,
1411 deduplicated=True,
1412 error="All chunks are near-duplicates",
1413 )
1414 continue
1416 chunks = [chunks[i] for i in keep_indices]
1417 embeddings = [embeddings[i] for i in keep_indices]
1418 sfe_overlays = record["sfe_overlays"]
1419 if sfe_overlays is not None:
1420 sfe_overlays = [sfe_overlays[i] for i in keep_indices]
1422 # Remap SFE indices through dedup (mirrors retain() path).
1423 sfe_associations = record["sfe_associations"]
1424 sfe_caused_by = record["sfe_caused_by"]
1425 sfe_entities = record["sfe_entities"]
1426 if sfe_associations is not None or sfe_caused_by is not None:
1427 old_to_new = {old: new for new, old in enumerate(keep_indices)}
1428 if sfe_associations is not None:
1429 sfe_associations = [
1430 (ent_idx, old_to_new[mem_idx]) for ent_idx, mem_idx in sfe_associations if mem_idx in old_to_new
1431 ]
1432 if sfe_caused_by is not None:
1433 sfe_caused_by = [
1434 (old_to_new[src], old_to_new[tgt], conf)
1435 for src, tgt, conf in sfe_caused_by
1436 if src in old_to_new and tgt in old_to_new
1437 ]
1438 # Persist remapped versions for _process_record_entities.
1439 record["sfe_associations"] = sfe_associations
1440 record["sfe_caused_by"] = sfe_caused_by
1442 if sfe_entities is not None:
1443 entities = list(sfe_entities)
1444 elif profile is not None and profile.entity_extraction == "metadata":
1445 entities = _entities_from_metadata(prepared.metadata)
1446 elif prepared.extract_entities:
1447 entities = await extract_entities(prepared.text, self.llm_provider)
1448 else:
1449 entities = []
1451 profile_tier = profile.authority_tier if profile else None
1452 chunk_metadata = merge_retain_metadata_authority_tier(
1453 prepared.metadata,
1454 bank_id=request.bank_id,
1455 profile_authority_tier=profile_tier,
1456 recall_authority=self.recall_authority,
1457 )
1458 chunk_metadata = dict(chunk_metadata or {})
1459 if request.mip_rule_name:
1460 chunk_metadata["_mip.rule"] = request.mip_rule_name
1461 if mip_pipeline and mip_pipeline.version is not None:
1462 chunk_metadata["_mip.pipeline_version"] = int(mip_pipeline.version)
1463 chunk_metadata.setdefault("_created_at", datetime.now(timezone.utc).isoformat())
1464 chunk_metadata.setdefault("_retain_id", uuid.uuid4().hex) # see the single-retain path
1465 if request.source:
1466 chunk_metadata.setdefault(SOURCE_KEY, request.source)
1468 # M10 source-aware retain — same helper as the single-retain path.
1469 chunk_ids = await self._provision_source_provenance(
1470 request,
1471 prepared.text,
1472 chunks,
1473 )
1475 memory_ids: list[str] = []
1476 items: list[VectorItem] = []
1477 for index, (chunk, embedding, chunk_id) in enumerate(zip(chunks, embeddings, chunk_ids, strict=False)):
1478 mem_id = uuid.uuid4().hex[:16]
1479 memory_ids.append(mem_id)
1480 item_meta, item_fact_type, item_occurred = _apply_fact_overlay(
1481 _chunk_metadata(chunk_metadata, index, len(chunks)),
1482 sfe_overlays[index] if sfe_overlays is not None else None,
1483 fact_type=prepared.fact_type,
1484 occurred_at=request.occurred_at,
1485 profile_fact_type=record["profile"].fact_type if record["profile"] is not None else None,
1486 )
1487 items.append(
1488 VectorItem(
1489 id=mem_id,
1490 bank_id=request.bank_id,
1491 vector=embedding,
1492 text=chunk,
1493 metadata=item_meta,
1494 tags=prepared.tags,
1495 fact_type=item_fact_type,
1496 occurred_at=item_occurred,
1497 retained_at=datetime.now(timezone.utc),
1498 chunk_id=chunk_id, # M10: source-chunk backreference
1499 )
1500 )
1502 record.update(
1503 {
1504 "chunks": chunks,
1505 "embeddings": embeddings,
1506 "entities": entities,
1507 "memory_ids": memory_ids,
1508 "items": items,
1509 }
1510 )
1511 stored_records.append(record)
1512 all_items.extend(items)
1514 if all_items:
1515 async with self._profiler.time("store_vec"):
1516 await self.vector_store.store_vectors(all_items)
1518 # Hindsight-parity semantic-kNN graph (C3a) — same call as the
1519 # single-retain path, applied per record so each batch's
1520 # neighbors-search is scoped to its own bank.
1521 for record in stored_records:
1522 await self._persist_semantic_links(
1523 record["request"].bank_id,
1524 record["memory_ids"],
1525 record["embeddings"],
1526 )
1528 if self.document_store is not None:
1529 from astrocyte.types import Document
1531 for record in stored_records:
1532 request = record["request"]
1533 for item in record["items"]:
1534 try:
1535 await self.document_store.store_document(
1536 Document(
1537 id=item.id,
1538 text=item.text,
1539 metadata=item.metadata,
1540 tags=item.tags,
1541 ),
1542 request.bank_id,
1543 )
1544 except Exception as exc:
1545 _logger.warning(
1546 "document_store.store_document failed for chunk %s: %s "
1547 "(keyword retrieval will miss this chunk)",
1548 item.id,
1549 exc,
1550 )
1552 # ── Phase 1: parallel entity processing across all records ──
1553 # Each record's entity work (embed names, store entities, write
1554 # co-occurrence links, run resolution cascade) is independent of
1555 # other records and dominates retain wall-clock when the bank is
1556 # large. ``asyncio.gather`` lets all records' LLM/DB I/O run
1557 # concurrently — for batches of 10 records this gives ~10× speedup
1558 # on the entity-resolution-bound retain phase. The non-I/O steps
1559 # (dedup cache, result assembly) stay sequential below.
1560 if self.graph_store is not None:
1561 entity_records = [r for r in stored_records if r["entities"]]
1562 if entity_records:
1563 await asyncio.gather(*[self._process_record_entities_with_retry(r) for r in entity_records])
1565 for record in stored_records:
1566 request: RetainRequest = record["request"]
1567 memory_ids: list[str] = record["memory_ids"]
1568 embeddings: list[list[float]] = record["embeddings"]
1569 chunks: list[str] = record["chunks"]
1571 for mem_id, embedding, chunk in zip(memory_ids, embeddings, chunks, strict=False):
1572 self._dedup.add(request.bank_id, mem_id, embedding, text=chunk)
1574 self._spawn_observation_consolidation(
1575 chunks=chunks,
1576 embeddings=embeddings,
1577 memory_ids=memory_ids,
1578 bank_id=request.bank_id,
1579 tags=request.tags,
1580 )
1582 results[record["index"]] = RetainResult(
1583 stored=True,
1584 memory_id=memory_ids[0] if memory_ids else None,
1585 )
1587 return [result or RetainResult(stored=False, error="Not retained") for result in results]