diff --git a/hindsight-api-slim/hindsight_api/alembic/versions/a1c9e7f3b2d8_observation_history_drop_memory_units_fk.py b/hindsight-api-slim/hindsight_api/alembic/versions/a1c9e7f3b2d8_observation_history_drop_memory_units_fk.py new file mode 100644 index 0000000000..a1334465f2 --- /dev/null +++ b/hindsight-api-slim/hindsight_api/alembic/versions/a1c9e7f3b2d8_observation_history_drop_memory_units_fk.py @@ -0,0 +1,67 @@ +"""Drop observation_history's FK to memory_units. + +The history table records one snapshot per observation change, keyed by +``(bank_id, observation_id)``. Its foreign key to ``memory_units`` existed only to +cascade-delete history when the observation row went away. + +That assumes every observation *is* a ``memory_units`` row, which is true only +while Postgres is the memories store. When another store owns the memories the +observation lives there and Postgres holds no row for it, so every history insert +raises a foreign-key violation — swallowed by the writer as "a race with parallel +consolidation" and logged at warning level. The audit trail goes silently empty. + +Dropping the constraint lets history be recorded wherever the observation is +stored. The cleanup the cascade used to do is now explicit, in the paths that +delete observations (``_execute_delete_action``, ``clear_observations``, +``delete_bank``). Rows orphaned by a path that misses — a document delete +cascading through ``memory_units``, for instance — are invisible to readers, +which always filter by ``(bank_id, observation_id)``, and are reclaimed when the +bank is deleted. + +Oracle builds this schema through its own DDL runner and never had the +constraint, so the Oracle slot is a deliberate no-op. + +Revision ID: a1c9e7f3b2d8 +Revises: c7d1e9a4b3f2 +""" + +from collections.abc import Sequence + +from alembic import op + +from hindsight_api.alembic._dialect import run_for_dialect + +revision: str = "a1c9e7f3b2d8" +down_revision: str | Sequence[str] | None = "c7d1e9a4b3f2" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + +_CONSTRAINT = "observation_history_observation_id_fkey" + + +def _pg_upgrade() -> None: + op.execute(f"ALTER TABLE observation_history DROP CONSTRAINT IF EXISTS {_CONSTRAINT}") + + +def _pg_downgrade() -> None: + # Re-adding the FK requires every row to reference a live memory_unit, so + # clear any history whose observation is not a Postgres row first — those are + # exactly the rows this migration made possible. + op.execute( + "DELETE FROM observation_history h " + "WHERE NOT EXISTS (SELECT 1 FROM memory_units m WHERE m.id = h.observation_id)" + ) + op.execute( + f"ALTER TABLE observation_history ADD CONSTRAINT {_CONSTRAINT} " + "FOREIGN KEY (observation_id) REFERENCES memory_units(id) ON DELETE CASCADE" + ) + + +def upgrade() -> None: + # Oracle never had the constraint (its schema is built by a separate DDL + # runner), so only Postgres has anything to drop. + run_for_dialect(pg=_pg_upgrade, oracle=None) + + +def downgrade() -> None: + run_for_dialect(pg=_pg_downgrade, oracle=None) diff --git a/hindsight-api-slim/hindsight_api/engine/consolidation/consolidator.py b/hindsight-api-slim/hindsight_api/engine/consolidation/consolidator.py index a4fed11cd4..098ec3f3ad 100644 --- a/hindsight-api-slim/hindsight_api/engine/consolidation/consolidator.py +++ b/hindsight-api-slim/hindsight_api/engine/consolidation/consolidator.py @@ -42,6 +42,7 @@ trace_context_of, ) from ..llm_wrapper import sanitize_llm_output +from ..memories import FactRecord, get_memories from ..memory_engine import Budget, fq_table from ..retain import embedding_utils from .prompts import ( @@ -246,6 +247,7 @@ async def _dedup_reconcile_create( create_text: str, create_source_ids: list[uuid.UUID], tags: list[str] | None, + txn=None, ) -> str | None: """Semantic dedup for a single CREATE (create-time, focused 1-by-1). @@ -259,27 +261,33 @@ async def _dedup_reconcile_create( if not outcome.should_merge or outcome.best_id is None: return None - # Fold the new source facts into the twin and persist the merged text. We keep the twin's - # existing embedding: the merged text is >= threshold similar, so the stored vector stays - # representative and we avoid a re-embed + a dialect-specific vector UPDATE. - search_vector_clause = ( - f",\n search_vector = to_tsvector('{config.text_search_extension_native_language}'::regconfig, COALESCE($1, ''))" - if config.text_search_extension == "native" - else "" - ) - await conn.execute( - f""" - UPDATE {fq_table("memory_units")} - SET text = $1, - source_memory_ids = (SELECT array_agg(DISTINCT e) FROM unnest(source_memory_ids || $2::uuid[]) e), - proof_count = (SELECT count(DISTINCT e) FROM unnest(source_memory_ids || $2::uuid[]) e), - updated_at = now(){search_vector_clause} - WHERE id = $3::uuid - """, - outcome.merged_text, - create_source_ids, - uuid.UUID(outcome.best_id), - ) + # Fold the new source facts into the twin and persist the merged text. The SQL path keeps the + # twin's existing embedding (the merged text is >= threshold similar, so it stays + # representative and avoids a re-embed + a dialect-specific vector UPDATE). + store = get_memories() + if store.writes_memory_rows_in_sql: + search_vector_clause = ( + f",\n search_vector = to_tsvector('{config.text_search_extension_native_language}'::regconfig, COALESCE($1, ''))" + if config.text_search_extension == "native" + else "" + ) + await conn.execute( + f""" + UPDATE {fq_table("memory_units")} + SET text = $1, + source_memory_ids = (SELECT array_agg(DISTINCT e) FROM unnest(source_memory_ids || $2::uuid[]) e), + proof_count = (SELECT count(DISTINCT e) FROM unnest(source_memory_ids || $2::uuid[]) e), + updated_at = now(){search_vector_clause} + WHERE id = $3::uuid + """, + outcome.merged_text, + create_source_ids, + uuid.UUID(outcome.best_id), + ) + else: + await _reconcile_merge_via_store( + store, conn, memory_engine, bank_id, outcome.best_id, outcome.merged_text, create_source_ids, txn=txn + ) return outcome.best_id @@ -293,6 +301,7 @@ async def _dedup_reconcile_update( updated_text: str, updated_emb_str: str | None, tags: list[str] | None, + txn=None, ) -> None: """Semantic dedup for an UPDATE (after the observation was rewritten + re-embedded). @@ -322,30 +331,38 @@ async def _dedup_reconcile_update( # the create path) then delete the now-redundant updated row. The all_strict/any tag match # guarantees twin and updated share scope, so dropping the updated row's tags loses no # visibility. Temporal fields follow the surviving twin (minimal scope; matches create). - search_vector_clause = ( - f",\n search_vector = to_tsvector('{config.text_search_extension_native_language}'::regconfig, COALESCE($1, ''))" - if config.text_search_extension == "native" - else "" - ) - await conn.execute( - f""" - UPDATE {fq_table("memory_units")} t - SET text = $1, - source_memory_ids = ( - SELECT array_agg(DISTINCT e) FROM unnest(t.source_memory_ids || u.source_memory_ids) e - ), - proof_count = ( - SELECT count(DISTINCT e) FROM unnest(t.source_memory_ids || u.source_memory_ids) e - ), - updated_at = now(){search_vector_clause} - FROM {fq_table("memory_units")} u - WHERE t.id = $2::uuid AND u.id = $3::uuid - """, - outcome.merged_text, - uuid.UUID(outcome.best_id), - uuid.UUID(updated_id), - ) - await _execute_delete_action(conn, bank_id, updated_id) + store = get_memories() + if store.writes_memory_rows_in_sql: + search_vector_clause = ( + f",\n search_vector = to_tsvector('{config.text_search_extension_native_language}'::regconfig, COALESCE($1, ''))" + if config.text_search_extension == "native" + else "" + ) + await conn.execute( + f""" + UPDATE {fq_table("memory_units")} t + SET text = $1, + source_memory_ids = ( + SELECT array_agg(DISTINCT e) FROM unnest(t.source_memory_ids || u.source_memory_ids) e + ), + proof_count = ( + SELECT count(DISTINCT e) FROM unnest(t.source_memory_ids || u.source_memory_ids) e + ), + updated_at = now(){search_vector_clause} + FROM {fq_table("memory_units")} u + WHERE t.id = $2::uuid AND u.id = $3::uuid + """, + outcome.merged_text, + uuid.UUID(outcome.best_id), + uuid.UUID(updated_id), + ) + else: + updated_obs = await store.get_memories(conn=conn, fq_table=fq_table, bank_id=bank_id, unit_ids=[updated_id]) + updated_sources = list(updated_obs[0].source_memory_ids or []) if updated_obs else [] + await _reconcile_merge_via_store( + store, conn, memory_engine, bank_id, outcome.best_id, outcome.merged_text, updated_sources, txn=txn + ) + await _execute_delete_action(conn, bank_id, updated_id, txn=txn) logger.info( "[CONSOLIDATION] dedup-merged updated observation %s into %s (cosine>=%.2f)", updated_id[:8], @@ -369,13 +386,21 @@ class _BatchDeltas: def _parse_observation_scopes(memory: dict[str, Any]) -> Any: - """Parse the per-memory ``observation_scopes`` column from a DB row. + """Parse the per-memory ``observation_scopes`` value. - asyncpg may return JSONB as a raw JSON string depending on driver settings; - accept both that and a pre-parsed value. + The value arrives already decoded when read through the memories store (its + reader coerces the JSONB column) or as raw JSON text from a driver without a + JSONB codec. A scalar mode such as ``"per_tag"`` decodes to a bare string that + is not itself valid JSON, so a blind ``json.loads`` would raise on it — try to + parse, but treat an unparseable string as an already-decoded scalar. """ raw = memory.get("observation_scopes") - return json.loads(raw) if isinstance(raw, str) else raw + if not isinstance(raw, str): + return raw + try: + return json.loads(raw) + except (json.JSONDecodeError, ValueError): + return raw def _resolve_obs_tags_list(memory: dict[str, Any]) -> list[list[str]] | None: @@ -471,18 +496,15 @@ async def _filter_live_source_memories( """ if not source_memory_ids: return [] - rows = await conn.fetch( - f""" - SELECT id - FROM {fq_table("memory_units")} - WHERE id = ANY($1::uuid[]) AND bank_id = $2 - FOR SHARE - """, - source_memory_ids, - bank_id, + # Which sources still exist, asked of the store (a non-Postgres store keeps them elsewhere). The + # FOR SHARE lock the Postgres path used is belt-and-suspenders: the orphan race is actually + # closed by the delete path running its stale-observation sweep *after* deleting the source, + # so an existence check is sufficient here. + present = await get_memories().get_memories( + conn=conn, fq_table=fq_table, bank_id=bank_id, unit_ids=[str(mid) for mid in source_memory_ids] ) - live = {row["id"] for row in rows} - return [mid for mid in source_memory_ids if mid in live] + live = {str(m.unit_id) for m in present} + return [mid for mid in source_memory_ids if str(mid) in live] class _CreateAction(BaseModel): @@ -591,12 +613,33 @@ async def _count_observations_for_scope( Returns the count of observations whose tags contain all specified tags. Observations with no tags are not counted (the limit does not apply to them). """ - return await conn.fetchval( - f"SELECT COUNT(*) FROM {fq_table('memory_units')} " - f"WHERE bank_id = $1 AND fact_type = 'observation' AND tags @> $2::varchar[]", - bank_id, - tags, - ) + store = get_memories() + if store.writes_memory_rows_in_sql: + return await conn.fetchval( + f"SELECT COUNT(*) FROM {fq_table('memory_units')} " + f"WHERE bank_id = $1 AND fact_type = 'observation' AND tags @> $2::varchar[]", + bank_id, + tags, + ) + # A store that keeps observations outside Postgres: count them through it (tag containment). + total = 0 + page_token = "" + for _ in range(100): + page = await store.scan_memories( + conn=conn, + fq_table=fq_table, + bank_id=bank_id, + fact_types=["observation"], + tags=tags or None, + tags_match="all", + limit=500, + page_token=page_token, + ) + total += len(page.memories) + page_token = page.next_page_token + if not page_token: + break + return total @dataclass(frozen=True) @@ -767,6 +810,120 @@ def flush(self) -> None: logger.info(log_output) +def _as_dt(v: "datetime | str | None") -> "datetime | None": + """Coerce an ISO string to a datetime. Recall results can carry timestamps as strings while + the store's addressed reads hand back datetimes, so normalise before comparing.""" + return datetime.fromisoformat(v) if isinstance(v, str) else v + + +def _merge_min(a: "datetime | str | None", b: "datetime | str | None") -> "datetime | None": + """SQL ``LEAST(a, COALESCE(b, a))`` in Python: the earlier of two times, ignoring None.""" + a, b = _as_dt(a), _as_dt(b) + return a if b is None else b if a is None else min(a, b) + + +def _merge_max(a: "datetime | str | None", b: "datetime | str | None") -> "datetime | None": + """SQL ``GREATEST(a, COALESCE(b, a))`` in Python: the later of two times, ignoring None.""" + a, b = _as_dt(a), _as_dt(b) + return a if b is None else b if a is None else max(a, b) + + +async def _reconcile_merge_via_store( + store, + conn, + memory_engine: "MemoryEngine", + bank_id: str, + observation_id: str, + merged_text: str, + add_source_ids: list, + txn=None, +) -> None: + """Dedup merge for a store that owns its rows: fold the extra source facts and the merged text + into the twin observation and re-upsert it, preserving its other fields. Re-embeds the merged + text because ``get_memories`` does not return the stored vector (the SQL path reuses it in + place instead).""" + current = await store.get_memories(conn=conn, fq_table=fq_table, bank_id=bank_id, unit_ids=[observation_id]) + cur = current[0] if current else None + if cur is None: + return + merged_sources = list(dict.fromkeys([*(cur.source_memory_ids or []), *(str(s) for s in add_source_ids)])) + embeddings = await embedding_utils.generate_embeddings_batch(memory_engine.embeddings, [merged_text]) + await store.upsert_observation( + conn=conn, + bank_id=bank_id, + txn=txn, + record=FactRecord( + unit_id=observation_id, + text=merged_text, + embedding=str(embeddings[0]) if embeddings else None, + fact_type="observation", + tags=list(cur.tags or []), + proof_count=len(merged_sources), + source_memory_ids=merged_sources, + event_date=cur.event_date, + occurred_start=cur.occurred_start, + occurred_end=cur.occurred_end, + mentioned_at=cur.mentioned_at, + created_at=cur.created_at, + ), + ) + + +async def _fetch_unconsolidated_rows( + conn, + bank_id: str, + fact_types: list[str], + limit: int, + observation_scopes: list[list[str]] | None, +) -> list[dict[str, Any]]: + """Unconsolidated candidate facts, read through the memories store. + + The store owns the memories, so this must ask it rather than query ``memory_units`` + directly — otherwise a store that keeps its rows elsewhere yields nothing and + consolidation silently produces no observations. Returns the same row-dict shape the + consolidation loop consumes. Mirrors the job's scope filter: with scopes, OR each + "tags ⊇ scope" and merge oldest-first; without, one unscoped read. + """ + store = get_memories() + scopes: list[list[str] | None] = list(observation_scopes) if observation_scopes else [None] + by_id: dict[str, Any] = {} + for scope in scopes: + for m in await store.find_unconsolidated( + conn=conn, fq_table=fq_table, bank_id=bank_id, fact_types=fact_types, limit=limit, scope_tags=scope + ): + by_id.setdefault(m.unit_id, m) + ordered = sorted(by_id.values(), key=lambda m: (m.created_at is None, m.created_at))[:limit] + return [ + { + "id": uuid.UUID(m.unit_id), + "text": m.text, + "fact_type": m.fact_type, + "occurred_start": m.occurred_start, + "occurred_end": m.occurred_end, + "event_date": m.event_date, + "tags": list(m.tags or []), + "mentioned_at": m.mentioned_at, + "observation_scopes": m.observation_scopes, + } + for m in ordered + ] + + +#: Cap on the store-side count of unconsolidated facts. Used only for the "is there work?" +#: gate and progress reporting, so a floor at this size is harmless on a huge backlog. +_COUNT_LIMIT = 100_000 + + +async def _count_unconsolidated_rows( + conn, + bank_id: str, + fact_types: list[str], + observation_scopes: list[list[str]] | None, +) -> int: + """Count of unconsolidated candidate facts, read through the store (bounded by ``_COUNT_LIMIT``).""" + return len(await _fetch_unconsolidated_rows(conn, bank_id, fact_types, _COUNT_LIMIT, observation_scopes)) + + async def run_consolidation_job( memory_engine: "MemoryEngine", bank_id: str, @@ -856,31 +1013,8 @@ async def _run_consolidation_job( perf.record_timing("fetch_bank", time.time() - t0) - # Build optional scope filter clause. When observation_scopes is provided, - # only process memories whose tags contain all tags in at least one scope. - scope_clause = "" - scope_params: list[Any] = [bank_id] - if observation_scopes: - or_parts: list[str] = [] - for scope_tags in observation_scopes: - idx = len(scope_params) + 1 - or_parts.append(f"tags @> ${idx}::varchar[]") - scope_params.append(scope_tags) - scope_clause = " AND (" + " OR ".join(or_parts) + ")" - - # Count total unconsolidated memories for progress logging - total_count = await conn.fetchval( - f""" - SELECT COUNT(*) - FROM {fq_table("memory_units")} - WHERE bank_id = $1 - AND consolidated_at IS NULL - AND consolidation_failed_at IS NULL - AND fact_type IN ('experience', 'world') - {scope_clause} - """, - *scope_params, - ) + # Count total unconsolidated memories for progress logging — through the store. + total_count = await _count_unconsolidated_rows(conn, bank_id, ["experience", "world"], observation_scopes) if total_count == 0: logger.debug(f"No new memories to consolidate for bank {bank_id}") @@ -905,19 +1039,7 @@ async def _count_unconsolidated() -> int: it. When that happens we re-count to report a real total (processed + remaining) instead of pinning the bar at 100%.""" async with acquire_with_retry(pool) as count_conn: - pending = await count_conn.fetchval( - f""" - SELECT COUNT(*) - FROM {fq_table("memory_units")} - WHERE bank_id = $1 - AND consolidated_at IS NULL - AND consolidation_failed_at IS NULL - AND fact_type IN ('experience', 'world') - {scope_clause} - """, - *scope_params, - ) - return pending or 0 + return await _count_unconsolidated_rows(count_conn, bank_id, ["experience", "world"], observation_scopes) async def _progress_total(processed: int) -> int: # Cheap path: while we're still within the start-of-job estimate it's exact, so @@ -965,26 +1087,12 @@ async def _progress_total(processed: int) -> int: min(max_memories_per_batch, int(round_remaining)) if round_limit_enabled else max_memories_per_batch ) - # Fetch next batch of unconsolidated memories + # Fetch next batch of unconsolidated memories — through the store, so a store that + # keeps its rows outside Postgres is read too. async with acquire_with_retry(pool) as conn: t0 = time.time() - # scope_params[0] is bank_id; append fetch_limit after scope params - fetch_params = list(scope_params) + [fetch_limit] - limit_idx = len(fetch_params) - memories = await conn.fetch( - f""" - SELECT id, text, fact_type, occurred_start, occurred_end, event_date, tags, mentioned_at, - observation_scopes - FROM {fq_table("memory_units")} - WHERE bank_id = $1 - AND consolidated_at IS NULL - AND consolidation_failed_at IS NULL - AND fact_type IN ('experience', 'world') - {scope_clause} - ORDER BY created_at ASC - LIMIT ${limit_idx} - """, - *fetch_params, + memories = await _fetch_unconsolidated_rows( + conn, bank_id, ["experience", "world"], fetch_limit, observation_scopes ) perf.record_timing("fetch_memories", time.time() - t0) @@ -1045,6 +1153,14 @@ async def _process_one_llm_batch(llm_batch_local: list[dict[str, Any]], batch_nu succeeded_ids: list[Any] = [] failed_ids: list[Any] = [] + # One cross-store write-group per LLM batch: MINT the txn up front (no Postgres held) + # and tag every observation upsert/delete + the mark_consolidated stamps with it, so + # they are durable-but-invisible in memlake while this batch runs its LLM work. The + # witness row + decide happen in ONE short transaction at the end (below) — we must not + # hold a Postgres transaction across the LLM calls in the sub-batch loop. + _txn_provider = get_memories() + _batch_txn = await _txn_provider.mint_txn(bank_id=bank_id, mutating=True) + pending: list[list[dict[str, Any]]] = [llm_batch_local] while pending: sub_batch = pending.pop(0) @@ -1067,6 +1183,7 @@ async def _process_one_llm_batch(llm_batch_local: list[dict[str, Any]], batch_nu perf=batch_perf, config=config, obs_tags_override=obs_tags, + txn=_batch_txn, ) sub_deleted += pass_deleted sub_llm_failed = sub_llm_failed or pass_failed @@ -1103,6 +1220,7 @@ async def _process_one_llm_batch(llm_batch_local: list[dict[str, Any]], batch_nu request_context=request_context, perf=batch_perf, config=config, + txn=_batch_txn, ) all_deleted += sub_deleted @@ -1125,17 +1243,38 @@ async def _process_one_llm_batch(llm_batch_local: list[dict[str, Any]], batch_nu succeeded_ids.extend(m["id"] for m in sub_batch) all_results.extend(sub_results) + # Mark through the store so the flag lands wherever the source facts live — tagged + # with this batch's txn, so the marks become visible together with the observations + # above. Then record the witness row and commit in this ONE short transaction (no LLM + # work inside it): its commit is the batch's fate, and `decide` publishes the group. async with acquire_with_retry(pool) as conn: + store = get_memories() + now = datetime.now(timezone.utc) if succeeded_ids: - await conn.executemany( - f"UPDATE {fq_table('memory_units')} SET consolidated_at = NOW() WHERE id = $1", - [(mem_id,) for mem_id in succeeded_ids], + await store.mark_consolidated( + conn=conn, + fq_table=fq_table, + bank_id=bank_id, + unit_ids=[str(mem_id) for mem_id in succeeded_ids], + when=now, + failed=False, + txn=_batch_txn, ) if failed_ids: - await conn.executemany( - f"UPDATE {fq_table('memory_units')} SET consolidation_failed_at = NOW() WHERE id = $1", - [(mem_id,) for mem_id in failed_ids], + await store.mark_consolidated( + conn=conn, + fq_table=fq_table, + bank_id=bank_id, + unit_ids=[str(mem_id) for mem_id in failed_ids], + when=now, + failed=True, + txn=_batch_txn, ) + async with conn.transaction(): + await _txn_provider.write_txn_witness(_batch_txn, conn=conn, fq_table=fq_table) + # Postgres committed the witness: publish the batch's write-group. On a crash before + # here the writes stay invisible and the recovery sweep resolves them (spec §5). + await _txn_provider.decide_txn(_batch_txn, commit=True) cancelled_local = False if operation_id and not await memory_engine._check_op_alive(operation_id): @@ -1510,6 +1649,7 @@ async def _process_memory_batch( perf: ConsolidationPerfLog | None = None, config: Any = None, obs_tags_override: list[str] | None = None, + txn=None, ) -> tuple[list[dict[str, Any]], int, bool]: """ Process a batch of memories in a single LLM call. @@ -1642,7 +1782,7 @@ async def _process_memory_batch( f"Batch consolidation: rejected delete — observation {delete.observation_id} not in unioned recall" ) continue - await _execute_delete_action(conn=conn, bank_id=bank_id, observation_id=delete.observation_id) + await _execute_delete_action(conn=conn, bank_id=bank_id, observation_id=delete.observation_id, txn=txn) deleted_count += 1 for update in llm_result.updates: @@ -1670,6 +1810,7 @@ async def _process_memory_batch( source_occurred_end=agg.occurred_end, source_mentioned_at=agg.mentioned_at, perf=perf, + txn=txn, ) for m in source_mems: per_memory_updated.add(str(m["id"])) @@ -1687,6 +1828,7 @@ async def _process_memory_batch( update.text, updated_emb_str, agg.tags, + txn=txn, ) # Deterministic dedup guard: map the observations the LLM was SHOWN by their @@ -1726,7 +1868,15 @@ async def _process_memory_batch( # near-identical observation (LLM-adjudicated, 1-by-1) instead of inserting a dup. if dedup_enabled: merged_into = await _dedup_reconcile_create( - conn, memory_engine, bank_id, config, dedup_llm_config, create.text, create_source_ids, agg.tags + conn, + memory_engine, + bank_id, + config, + dedup_llm_config, + create.text, + create_source_ids, + agg.tags, + txn=txn, ) if merged_into is not None: logger.info( @@ -1750,6 +1900,7 @@ async def _process_memory_batch( occurred_end=agg.occurred_end, mentioned_at=agg.mentioned_at, perf=perf, + txn=txn, ) for m in source_mems: per_memory_created.add(str(m["id"])) @@ -1860,6 +2011,7 @@ async def _execute_update_action( source_occurred_end: datetime | None = None, source_mentioned_at: datetime | None = None, perf: ConsolidationPerfLog | None = None, + txn=None, ) -> str | None: """ Update an existing observation. @@ -1917,30 +2069,57 @@ async def _execute_update_action( ) t0 = time.time() - await conn.execute( - f""" - UPDATE {fq_table("memory_units")} - SET text = $1, - embedding = $2::vector, - source_memory_ids = $3, - proof_count = $4, - tags = $9, - updated_at = now(), - occurred_start = LEAST(occurred_start, COALESCE($6, occurred_start)), - occurred_end = GREATEST(occurred_end, COALESCE($7, occurred_end)), - mentioned_at = GREATEST(mentioned_at, COALESCE($8, mentioned_at)){search_vector_clause} - WHERE id = $5 - """, - new_text, - embedding_str, - source_ids, - len(source_ids), - uuid.UUID(observation_id), - source_occurred_start, - source_occurred_end, - source_mentioned_at, - merged_tags, - ) + store = get_memories() + if store.writes_memory_rows_in_sql: + await conn.execute( + f""" + UPDATE {fq_table("memory_units")} + SET text = $1, + embedding = $2::vector, + source_memory_ids = $3, + proof_count = $4, + tags = $9, + updated_at = now(), + occurred_start = LEAST(occurred_start, COALESCE($6, occurred_start)), + occurred_end = GREATEST(occurred_end, COALESCE($7, occurred_end)), + mentioned_at = GREATEST(mentioned_at, COALESCE($8, mentioned_at)){search_vector_clause} + WHERE id = $5 + """, + new_text, + embedding_str, + source_ids, + len(source_ids), + uuid.UUID(observation_id), + source_occurred_start, + source_occurred_end, + source_mentioned_at, + merged_tags, + ) + else: + # Upsert overwrites the whole observation, so start from its current state (fetched from + # the store) and apply the same merge the SQL does — LEAST/GREATEST on the times — while + # preserving fields the update never touches (event_date, created_at). + current = await store.get_memories(conn=conn, fq_table=fq_table, bank_id=bank_id, unit_ids=[observation_id]) + cur = current[0] if current else None + await store.upsert_observation( + conn=conn, + bank_id=bank_id, + txn=txn, + record=FactRecord( + unit_id=observation_id, + text=new_text, + embedding=embedding_str, + fact_type="observation", + tags=merged_tags, + proof_count=len(source_ids), + source_memory_ids=[str(s) for s in source_ids], + event_date=cur.event_date if cur else None, + occurred_start=_merge_min(model.occurred_start, source_occurred_start), + occurred_end=_merge_max(model.occurred_end, source_occurred_end), + mentioned_at=_merge_max(model.mentioned_at, source_mentioned_at), + created_at=cur.created_at if cur else None, + ), + ) # Record the pre-update snapshot in the dedicated observation_history table # (one row per change), then trim to the configured cap. History lived in a @@ -1989,6 +2168,7 @@ async def _execute_create_action( occurred_end: datetime | None = None, mentioned_at: datetime | None = None, perf: ConsolidationPerfLog | None = None, + txn=None, ) -> None: """ Create a new observation from one or more source memories. @@ -2008,6 +2188,7 @@ async def _execute_create_action( occurred_end=occurred_end, mentioned_at=mentioned_at, perf=perf, + txn=txn, ) # Map the new observation onto the consolidation trace as a produced memory. new_id = created.get("observation_id") @@ -2020,13 +2201,18 @@ async def _execute_delete_action( conn: "Connection", bank_id: str, observation_id: str, + txn=None, ) -> None: """Delete a superseded or contradicted observation.""" - await conn.execute( - f"DELETE FROM {fq_table('memory_units')} WHERE id = $1 AND bank_id = $2 AND fact_type = 'observation'", - uuid.UUID(observation_id), - bank_id, - ) + store = get_memories() + if store.writes_memory_rows_in_sql: + await conn.execute( + f"DELETE FROM {fq_table('memory_units')} WHERE id = $1 AND bank_id = $2 AND fact_type = 'observation'", + uuid.UUID(observation_id), + bank_id, + ) + else: + await store.delete_facts(bank_id, [observation_id], txn=txn) logger.debug(f"Deleted observation {observation_id}") @@ -2373,6 +2559,7 @@ async def _create_observation_directly( occurred_end: datetime | None = None, mentioned_at: datetime | None = None, perf: ConsolidationPerfLog | None = None, + txn=None, ) -> dict[str, Any]: """Create an observation from one or more source memories with pre-processed text.""" live_source_memory_ids = await _filter_live_source_memories(conn, bank_id, source_memory_ids) @@ -2399,70 +2586,96 @@ async def _create_observation_directly( t0 = time.time() observation_id = uuid.uuid4() - # Query varies based on text search backend - config = get_config() - if config.text_search_extension == "vchord": - # VectorChord: manually tokenize and insert search_vector - query = f""" - INSERT INTO {fq_table("memory_units")} ( - id, bank_id, text, fact_type, embedding, proof_count, source_memory_ids, - tags, event_date, occurred_start, occurred_end, mentioned_at, search_vector - ) - VALUES ($1, $2, $3, 'observation', $4::vector, 1, $5, $6, $7, $8, $9, $10, - tokenize($3, 'llmlingua2')::bm25_catalog.bm25vector) - RETURNING id - """ - elif config.text_search_extension == "native": - # Native: search_vector is populated with to_tsvector() using the - # configured native language dictionary, matching the batch insert - # path in ops_postgresql.insert_facts_batch. - query = f""" - INSERT INTO {fq_table("memory_units")} ( - id, bank_id, text, fact_type, embedding, proof_count, source_memory_ids, - tags, event_date, occurred_start, occurred_end, mentioned_at, search_vector - ) - VALUES ($1, $2, $3, 'observation', $4::vector, 1, $5, $6, $7, $8, $9, $10, - to_tsvector('{config.text_search_extension_native_language}'::regconfig, COALESCE($3, ''))) - RETURNING id - """ - else: # pg_textsearch, pgroonga, pg_search: indexes operate on base text columns directly - query = f""" - INSERT INTO {fq_table("memory_units")} ( - id, bank_id, text, fact_type, embedding, proof_count, source_memory_ids, - tags, event_date, occurred_start, occurred_end, mentioned_at - ) - VALUES ($1, $2, $3, 'observation', $4::vector, 1, $5, $6, $7, $8, $9, $10) - RETURNING id - """ + # Write the observation. A SQL store keeps it as a `memory_units` row (inline below, with the + # search_vector the configured backend needs); a store that owns its rows takes it through + # upsert_observation as a normal Observation-type memory carrying all of its own state. + store = get_memories() + if store.writes_memory_rows_in_sql: + config = get_config() + if config.text_search_extension == "vchord": + # VectorChord: manually tokenize and insert search_vector + query = f""" + INSERT INTO {fq_table("memory_units")} ( + id, bank_id, text, fact_type, embedding, proof_count, source_memory_ids, + tags, event_date, occurred_start, occurred_end, mentioned_at, search_vector + ) + VALUES ($1, $2, $3, 'observation', $4::vector, 1, $5, $6, $7, $8, $9, $10, + tokenize($3, 'llmlingua2')::bm25_catalog.bm25vector) + RETURNING id + """ + elif config.text_search_extension == "native": + # Native: search_vector is populated with to_tsvector() using the + # configured native language dictionary, matching the batch insert + # path in ops_postgresql.insert_facts_batch. + query = f""" + INSERT INTO {fq_table("memory_units")} ( + id, bank_id, text, fact_type, embedding, proof_count, source_memory_ids, + tags, event_date, occurred_start, occurred_end, mentioned_at, search_vector + ) + VALUES ($1, $2, $3, 'observation', $4::vector, 1, $5, $6, $7, $8, $9, $10, + to_tsvector('{config.text_search_extension_native_language}'::regconfig, COALESCE($3, ''))) + RETURNING id + """ + else: # pg_textsearch, pgroonga, pg_search: indexes operate on base text columns directly + query = f""" + INSERT INTO {fq_table("memory_units")} ( + id, bank_id, text, fact_type, embedding, proof_count, source_memory_ids, + tags, event_date, occurred_start, occurred_end, mentioned_at + ) + VALUES ($1, $2, $3, 'observation', $4::vector, 1, $5, $6, $7, $8, $9, $10) + RETURNING id + """ - row = await conn.fetchrow( - query, - observation_id, - bank_id, - observation_text, - embedding_str, - source_memory_ids, - obs_tags, - obs_event_date, - obs_occurred_start, - obs_occurred_end, - obs_mentioned_at, - ) + row = await conn.fetchrow( + query, + observation_id, + bank_id, + observation_text, + embedding_str, + source_memory_ids, + obs_tags, + obs_event_date, + obs_occurred_start, + obs_occurred_end, + obs_mentioned_at, + ) + created_id = row["id"] - # Populate observation_sources junction table (Oracle only — PG uses native array ops). - if memory_engine._backend.ops.uses_observation_sources_table and source_memory_ids: - await conn.executemany( - f""" - INSERT INTO {fq_table("observation_sources")} (observation_id, source_id) - VALUES ($1, $2) - ON CONFLICT (observation_id, source_id) DO NOTHING - """, - [(observation_id, sid) for sid in dict.fromkeys(source_memory_ids)], + # Populate observation_sources junction table (Oracle only — PG uses native array ops). + if memory_engine._backend.ops.uses_observation_sources_table and source_memory_ids: + await conn.executemany( + f""" + INSERT INTO {fq_table("observation_sources")} (observation_id, source_id) + VALUES ($1, $2) + ON CONFLICT (observation_id, source_id) DO NOTHING + """, + [(observation_id, sid) for sid in dict.fromkeys(source_memory_ids)], + ) + else: + await store.upsert_observation( + conn=conn, + bank_id=bank_id, + txn=txn, + record=FactRecord( + unit_id=str(observation_id), + text=observation_text, + embedding=embedding_str, + fact_type="observation", + tags=list(obs_tags), + proof_count=1, + source_memory_ids=[str(s) for s in source_memory_ids], + event_date=obs_event_date, + occurred_start=obs_occurred_start, + occurred_end=obs_occurred_end, + mentioned_at=obs_mentioned_at, + created_at=now, + ), ) + created_id = observation_id if perf: perf.record_timing("db_write", time.time() - t0) logger.debug(f"Created observation {observation_id} from {len(source_memory_ids)} memories (tags: {obs_tags})") - return {"action": "created", "observation_id": str(row["id"]), "tags": obs_tags} + return {"action": "created", "observation_id": str(created_id), "tags": obs_tags} diff --git a/hindsight-api-slim/hindsight_api/engine/db/__init__.py b/hindsight-api-slim/hindsight_api/engine/db/__init__.py index 66c4fd6b90..3e9fde9fd8 100644 --- a/hindsight-api-slim/hindsight_api/engine/db/__init__.py +++ b/hindsight-api-slim/hindsight_api/engine/db/__init__.py @@ -67,16 +67,28 @@ def create_database_backend(backend_type: str) -> DatabaseBackend: return _get_backend_class(backend_type)() +_OPS_CACHE: dict[str, DataAccessOps] = {} + + def create_data_access_ops(backend_type: str) -> DataAccessOps: - """Factory: create a DataAccessOps by backend name. + """Factory: the DataAccessOps for a backend name. + + Returns a per-dialect SINGLETON: ``DataAccessOps`` is stateless (it only builds and runs SQL), + so one shared instance per dialect is correct — and it means the database backend and the + memories store hold the *same* ops object, so a test that patches a method on it (e.g. + ``enqueue_graph_maintenance``) observes every caller regardless of which layer issued it. Args: backend_type: One of "postgresql" or "oracle". Returns: - A DataAccessOps instance. + The shared DataAccessOps instance for that backend. Raises: ValueError: If backend_type is not recognized. """ - return _get_ops_class(backend_type)() + ops = _OPS_CACHE.get(backend_type) + if ops is None: + ops = _get_ops_class(backend_type)() + _OPS_CACHE[backend_type] = ops + return ops diff --git a/hindsight-api-slim/hindsight_api/engine/entity_resolver.py b/hindsight-api-slim/hindsight_api/engine/entity_resolver.py index 0adf28c46f..062fb59efe 100644 --- a/hindsight-api-slim/hindsight_api/engine/entity_resolver.py +++ b/hindsight-api-slim/hindsight_api/engine/entity_resolver.py @@ -864,6 +864,7 @@ async def link_units_to_entities_batch( self, unit_entity_pairs: list[tuple[str, str]] | list[tuple[str, str, datetime | None]], conn=None, + bank_id: str | None = None, ): """ Link multiple memory units to entities in batch (MUCH faster than sequential). @@ -891,22 +892,32 @@ async def link_units_to_entities_batch( if conn is None: async with acquire_with_retry(self.pool) as conn: - return await self._link_units_to_entities_batch_impl(conn, normalized) + return await self._link_units_to_entities_batch_impl(conn, normalized, bank_id) else: - return await self._link_units_to_entities_batch_impl(conn, normalized) + return await self._link_units_to_entities_batch_impl(conn, normalized, bank_id) - async def _link_units_to_entities_batch_impl(self, conn, unit_entity_pairs: list[tuple[str, str, datetime | None]]): + async def _link_units_to_entities_batch_impl( + self, conn, unit_entity_pairs: list[tuple[str, str, datetime | None]], bank_id: str | None = None + ): # Sorted bulk insert to prevent deadlocks from inconsistent lock ordering # across concurrent transactions on the unit_entities unique index. sorted_pairs = sorted(unit_entity_pairs, key=lambda t: (t[0], t[1])) unit_ids = [p[0] for p in sorted_pairs] entity_ids = [p[1] for p in sorted_pairs] - await self._ops.bulk_insert_unit_entities( - conn, - fq_table("unit_entities"), - unit_ids, - entity_ids, + # The unit→entity posting belongs to whoever stores the memory, so the + # memories store records it. Co-occurrence below is separate and unaffected: + # it references only `entities`, which stays in Postgres either way, and is + # read by the entity-graph endpoint and by resolution's disambiguation signal. + from .memories import get_memories + + await get_memories().record_unit_entities( + conn=conn, + ops=self._ops, + fq_table=fq_table, + bank_id=bank_id, + unit_ids=unit_ids, + entity_ids=entity_ids, ) # Build maps keyed by unit_id: diff --git a/hindsight-api-slim/hindsight_api/engine/graph_maintenance.py b/hindsight-api-slim/hindsight_api/engine/graph_maintenance.py index d925812afc..314c3b9072 100644 --- a/hindsight-api-slim/hindsight_api/engine/graph_maintenance.py +++ b/hindsight-api-slim/hindsight_api/engine/graph_maintenance.py @@ -5,25 +5,25 @@ 1. **Relink top-up.** Drain ``graph_maintenance_queue`` (units whose outgoing temporal/semantic links lost a neighbour to a delete). For each, count current outgoing links per type; if below cap, run the - same probes retain uses (:func:`fetch_temporal_neighbors`, - :func:`compute_semantic_links_ann`) and insert the missing links. - ``bulk_insert_links`` has ``ON CONFLICT DO NOTHING`` on the uniqueness - key, so we can re-probe freely and the DB de-dupes. + same probes retain uses and insert the missing links. 2. **Orphan entity prune.** Delete ``entities`` rows in the bank that no - longer have any ``unit_entities`` references. FK ON DELETE CASCADE on + longer have any live memory references. FK ON DELETE CASCADE on ``entity_cooccurrences`` then removes any cooccurrence row pointing at the pruned entities. 3. **Stale cooccurrence prune.** Defensive sweep for cooccurrence rows - where both endpoints still exist but no current memory_unit references + where both endpoints still exist but no current memory references both of them — the cooccurrence was real at the time it was recorded, but every unit that witnessed it has since been deleted. -All three passes run on every invocation. The queue is the only source -of work for pass 1; passes 2 and 3 are bank-wide sweeps backed by indexes -on ``entities(bank_id)`` and ``unit_entities(entity_id)``, so they're -cheap when there's nothing to do. +Each pass is work the *memories store* owns, because each is a query over +`memory_links`, `unit_entities` and `entities` — the slice the store carves +out. This module orchestrates them (drain the queue, wrap the sweep in a +deadlock-retry) and asks the store to do the part that touches storage. A store +whose links travel inside its memories has no `memory_links` to dangle and no +join table to sweep, so its relink and cooccurrence passes are no-ops and the +job simply prunes the orphan `entities` rows, which stay in Postgres regardless. The worker dedupes on bank: a second job for the same bank is dropped while one is pending. Once processing starts, a new job becomes the @@ -35,19 +35,15 @@ import logging import time -import uuid as uuid_module from dataclasses import dataclass -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING -from ..config import get_config from ..models import RequestContext from .db.base import DatabaseConnection -from .retain.link_utils import ( - MAX_TEMPORAL_LINKS_PER_UNIT, - _bulk_insert_links, - _normalize_datetime, - compute_semantic_links_ann, -) + +# Re-exported for callers and tests that import the link caps from here; the caps +# themselves live with the retain-time link builders the relink pass mirrors. +from .retain.link_utils import MAX_TEMPORAL_LINKS_PER_UNIT # noqa: F401 from .schema import fq_table if TYPE_CHECKING: @@ -59,14 +55,11 @@ # time. If you change one, change the other — otherwise victims would either # never reach the cap (probe returns less than the cap) or stay perpetually # under it (cap is higher than retain creates). +# +# Kept here as well as in the store's Postgres relink pass because the relink +# tests import it from this module; the two must not drift. MAX_SEMANTIC_LINKS_PER_UNIT = 50 -# Worker fetches this many rows per relink-loop iteration. Bounds -# per-iteration probe/insert latency so a 10k-row backlog doesn't hold a -# worker slot for minutes. Chosen so the typical iteration runs in well -# under 1s. -_DRAIN_BATCH_SIZE = 50 - # Retry budget for the idempotent Pass 2/3 entity/cooccurrence sweep. Higher # than db_utils' default (3) because the sweep has no client waiting on it and # is safe to rerun, so we'd rather spend a longer jittered-backoff tail than @@ -104,7 +97,7 @@ async def enqueue_relink_victims( conn: DatabaseConnection, bank_id: str, affected_unit_ids: list[str], - ops: Any, + ops: Any = None, include_affected_units: bool = False, ) -> int: """Enqueue surviving units whose outgoing temporal/semantic links pointed at @@ -121,6 +114,11 @@ async def enqueue_relink_victims( — the drain skips queue rows with no live unit — so callers should only set it when the unit survives the transaction. + Delegated to the memories store: finding the victims is a `memory_links` + query, and a store whose links are inline has none, so it returns 0 and the + relink pass has nothing to do. ``ops`` is accepted for callers that still pass + it and ignored by a store that resolves what it needs from ``conn``. + Args: conn: Database connection inside the active transaction. bank_id: Bank owning the affected units. @@ -129,55 +127,23 @@ async def enqueue_relink_victims( ops: ``DataAccessOps`` instance, supplies the dialect-specific bulk-insert path. include_affected_units: Also enqueue ``affected_unit_ids`` themselves, - for callers that leave them live. One combined insert (rather than a - second call) keeps the queue's sorted lock ordering intact: two - transactions editing mutually linked units would otherwise take the - ``(bank_id, unit_id)`` keys in opposite orders and deadlock. + for callers that leave them live. Returns: - Number of distinct units passed to the queue insert. + Number of distinct victim units enqueued (0 for a store with no links). """ if not affected_unit_ids: return 0 - affected_uuids = [uuid_module.UUID(uid) if isinstance(uid, str) else uid for uid in affected_unit_ids] - affected_str_set = {str(uid) for uid in affected_uuids} - - # Find units (other than the affected ones) that have an outgoing - # temporal/semantic link pointing at an affected unit. Entity links are - # intentionally excluded — they're scheduled for removal and would only - # add noise to the recompute job. - victim_rows = await conn.fetch( - f""" - SELECT DISTINCT from_unit_id - FROM {fq_table("memory_links")} - WHERE to_unit_id = ANY($1::uuid[]) - AND bank_id = $2 - AND link_type IN ('temporal', 'semantic') - """, - affected_uuids, - bank_id, - ) - - relink_ids = {row["from_unit_id"] for row in victim_rows if str(row["from_unit_id"]) not in affected_str_set} - if include_affected_units: - relink_ids.update(affected_uuids) - - if not relink_ids: - return 0 + from .memories import get_memories - await ops.enqueue_graph_maintenance( - conn, - fq_table("graph_maintenance_queue"), - bank_id, - list(relink_ids), - ) - - logger.debug( - f"[GRAPH_MAINT] Enqueued {len(relink_ids)} units for relinking in " - f"bank={bank_id} ({len(affected_unit_ids)} units affected)" + return await get_memories().enqueue_relink_victims( + conn=conn, + fq_table=fq_table, + bank_id=bank_id, + affected_unit_ids=affected_unit_ids, + include_affected_units=include_affected_units, ) - return len(relink_ids) async def run_graph_maintenance_job( @@ -193,88 +159,52 @@ async def run_graph_maintenance_job( Per-pass counters from :class:`JobResult`. """ del request_context # accepted for symmetry with other run_*_job helpers + from ..config import get_config + from .memories import get_memories + backend = await memory_engine._get_backend() - ops = backend.ops + store = get_memories() + config = get_config() result = JobResult() job_start = time.time() - semantic_link_min_similarity = get_config().semantic_link_min_similarity # --- Pass 1: relink --- - # Per-iteration loop: claim → top up → commit. We rely on submit-time - # dedup to keep at most one job per bank running, so no need for - # SKIP LOCKED. - iterations = 0 - while True: - from .memory_engine import acquire_with_retry - - async with acquire_with_retry(backend) as conn: - async with conn.transaction(): - unit_ids = await ops.claim_graph_maintenance_batch( - conn, - fq_table("graph_maintenance_queue"), - bank_id, - _DRAIN_BATCH_SIZE, - ) - if not unit_ids: - break - - result.relink_links_added += await _relink_batch( - conn, - bank_id, - unit_ids, - ops, - backend, - semantic_link_min_similarity, - ) - - result.relink_units_processed += len(unit_ids) - iterations += 1 - - if iterations > 10000: - # Defensive guard against runaway loops — at 50 units/iter that's - # 500k targets, far beyond any realistic single-bank backlog. - logger.error( - f"[GRAPH_MAINT] bank={bank_id} hit iteration cap ({iterations}); aborting relink ({result.as_dict()})" - ) - break + # The store owns the whole drain loop: it is a claim → top-up → commit over + # its own link table, so how it batches and re-probes is its business. A + # store with no links returns an empty dict and this is a no-op. + relink = await store.relink_pass(backend=backend, fq_table=fq_table, bank_id=bank_id, config=config) + result.relink_units_processed = relink.get("relink_units_processed", 0) + result.relink_links_added = relink.get("relink_links_added", 0) # --- Pass 2 & 3: entity / cooccurrence sweeps --- # Bank-wide single-statement deletes. Cheap when there's nothing to do. # # Unlike Pass 1's queue claim, these DELETEs aren't protected by any - # consistent lock-ordering guarantee: prune_stale_cooccurrences scans - # entity_cooccurrences via a join/NOT EXISTS plan, while retain's - # concurrent cooccurrence upserts (entity_resolver._flush_pending) lock - # the same rows in sorted (entity_id_1, entity_id_2) order. When a sweep - # and a concurrent upsert touch overlapping rows in opposite orders, - # Postgres detects a genuine circular wait and aborts one side with - # DeadlockDetectedError. Both prunes are idempotent bank-wide sweeps — - # rerunning only deletes what's still stale — so retrying the whole - # transaction on deadlock is safe. + # consistent lock-ordering guarantee: the stale-cooccurrence prune scans + # entity_cooccurrences via a join/NOT EXISTS plan, while retain's concurrent + # cooccurrence upserts (entity_resolver._flush_pending) lock the same rows in + # sorted (entity_id_1, entity_id_2) order. When a sweep and a concurrent + # upsert touch overlapping rows in opposite orders, Postgres detects a + # genuine circular wait and aborts one side with DeadlockDetectedError. Both + # prunes are idempotent bank-wide sweeps — rerunning only deletes what's + # still stale — so retrying the whole transaction on deadlock is safe. + # + # The prunes themselves are the store's: the orphan-`entities` sweep applies + # to every store (that registry stays in Postgres), while the cooccurrence + # sweep is a no-op for a store that never wrote `unit_entities`. from .db_utils import retry_with_backoff from .memory_engine import acquire_with_retry async def _run_sweep() -> _SweepCounts: async with acquire_with_retry(backend) as conn: async with conn.transaction(): - orphan_pruned = await ops.prune_orphan_entities( - conn, - fq_table("entities"), - fq_table("unit_entities"), - bank_id, - ) + orphan_pruned = await store.prune_orphan_entities(conn=conn, fq_table=fq_table, bank_id=bank_id) # The orphan prune above cascades cooccurrences via FK. The # explicit cooccurrence pass below catches the *stale-count* - # case: both entities still exist but no current unit - # witnesses them together. - stale_pruned = await ops.prune_stale_cooccurrences( - conn, - fq_table("entity_cooccurrences"), - fq_table("unit_entities"), - fq_table("entities"), - bank_id, - ) + # case: both entities still exist but no current unit witnesses + # them together. + stale_pruned = await store.prune_stale_cooccurrences(conn=conn, fq_table=fq_table, bank_id=bank_id) return _SweepCounts(orphan_entities_pruned=orphan_pruned, stale_cooccurrences_pruned=stale_pruned) # A larger retry budget than the default (3): this is idempotent background @@ -292,130 +222,3 @@ async def _run_sweep() -> _SweepCounts: f"[GRAPH_MAINT] bank={bank_id} done: {result.as_dict()}, elapsed={elapsed:.2f}s, operation_id={operation_id}" ) return result.as_dict() - - -async def _relink_batch( - conn: DatabaseConnection, - bank_id: str, - victim_ids: list[str], - ops: Any, - backend: Any, - semantic_link_min_similarity: float, -) -> int: - """Top up temporal/semantic links for a batch of victim units. Returns rows inserted.""" - # Load each victim's metadata. Victims whose units were deleted between - # enqueue and now silently drop out — exactly the no-op behaviour we want - # for stale queue rows. - victim_uuids = [uuid_module.UUID(vid) for vid in victim_ids] - victim_rows = await conn.fetch( - f""" - SELECT id::text AS id, event_date, fact_type, embedding::text AS embedding - FROM {fq_table("memory_units")} - WHERE id = ANY($1::uuid[]) - AND bank_id = $2 - AND fact_type IN ('experience', 'world') - """, - victim_uuids, - bank_id, - ) - - if not victim_rows: - return 0 - - alive_uuids = [uuid_module.UUID(row["id"]) for row in victim_rows] - - # Count current outgoing temporal/semantic links per victim so we only - # probe for the ones genuinely below cap. Saves the bulk of the work when - # most victims still have plenty of links. - count_rows = await conn.fetch( - f""" - SELECT from_unit_id, link_type, COUNT(*) AS cnt - FROM {fq_table("memory_links")} - WHERE from_unit_id = ANY($1::uuid[]) - AND bank_id = $2 - AND link_type IN ('temporal', 'semantic') - GROUP BY from_unit_id, link_type - """, - alive_uuids, - bank_id, - ) - counts: dict[tuple[str, str], int] = {} - for row in count_rows: - counts[(str(row["from_unit_id"]), row["link_type"])] = int(row["cnt"]) - - # --- Temporal top-up --- - temporal_needs = [r for r in victim_rows if counts.get((r["id"], "temporal"), 0) < MAX_TEMPORAL_LINKS_PER_UNIT] - new_links: list[tuple] = [] - - if temporal_needs: - lateral_unit_ids = [uuid_module.UUID(r["id"]) for r in temporal_needs if r["event_date"] is not None] - lateral_event_dates = [ - _normalize_datetime(r["event_date"]) for r in temporal_needs if r["event_date"] is not None - ] - lateral_fact_types = [r["fact_type"] for r in temporal_needs if r["event_date"] is not None] - - if lateral_unit_ids: - rows = await ops.fetch_temporal_neighbors( - conn, - fq_table("memory_units"), - bank_id, - lateral_unit_ids, - lateral_event_dates, - lateral_fact_types, - MAX_TEMPORAL_LINKS_PER_UNIT, - ) - for row in rows: - time_diff_h = float(row["time_diff_hours"]) - # Mirror the 24h window enforced at retain time. The bidirectional - # index scan returns the K closest neighbours regardless of - # window, so we filter here. - if time_diff_h > 24: - continue - weight = max(0.3, 1.0 - (time_diff_h / 24)) - new_links.append((row["from_id"], str(row["id"]), "temporal", weight, None)) - - # --- Semantic top-up --- - # ANN must run on its own connection: it opens a nested transaction with - # SET LOCAL hnsw.ef_search + CREATE TEMP TABLE ON COMMIT DROP, and nesting - # that inside our current write transaction would commit our writes early. - semantic_needs = [ - r - for r in victim_rows - if counts.get((r["id"], "semantic"), 0) < MAX_SEMANTIC_LINKS_PER_UNIT and r["embedding"] is not None - ] - if semantic_needs: - from .memory_engine import acquire_with_retry - - seed_ids = [r["id"] for r in semantic_needs] - seed_embs = [r["embedding"] for r in semantic_needs] - seed_ftypes = [r["fact_type"] for r in semantic_needs] - async with acquire_with_retry(backend) as ann_conn: - try: - ann_links = await compute_semantic_links_ann( - ann_conn, - bank_id, - seed_ids, - seed_embs, - fact_types=seed_ftypes, - threshold=semantic_link_min_similarity, - ) - # Strip self-links (rare but possible because the ANN probe - # has no exclude list — see the comment in compute_semantic_links_ann). - ann_links = [lnk for lnk in ann_links if lnk[0] != lnk[1]] - new_links.extend(ann_links) - except Exception as e: - # ANN uses PG-specific HNSW syntax; on dialects/configs where - # it isn't available we still want the temporal top-up to land. - logger.warning(f"[GRAPH_MAINT] Semantic top-up failed for bank={bank_id}: {type(e).__name__}: {e}") - - if not new_links: - return 0 - - await _bulk_insert_links( - conn, - new_links, - bank_id=bank_id, - skip_exists_check=False, - ops=ops, - ) - return len(new_links) diff --git a/hindsight-api-slim/hindsight_api/engine/maintenance.py b/hindsight-api-slim/hindsight_api/engine/maintenance.py index 99ade76606..4a33ae5fc9 100644 --- a/hindsight-api-slim/hindsight_api/engine/maintenance.py +++ b/hindsight-api-slim/hindsight_api/engine/maintenance.py @@ -54,6 +54,13 @@ # sets the drain rate for a backlog. Kept at one-per-tick (the value it used while # it rode the worker's poll loop) so throughput is unchanged by the move. _OPERATION_CLEANUP_INTERVAL_SECONDS = 60 +# Cross-store txn recovery (memlake store only): a backstop for a writer that crashed between +# its memlake writes and the decide. The happy path decides inline after commit, so this rarely +# finds work; five minutes bounds how long a crashed txn stalls its namespace's fold. +_TXN_RECOVERY_INTERVAL_SECONDS = 300 +# A pending txn is left alone for this long from first sighting before the sweep aborts an +# unwitnessed one — the writer may still be mid-flight (PendingTxn carries no timestamp). +_TXN_RECOVERY_GRACE_SECONDS = 300 class MaintenanceLoop: @@ -65,6 +72,9 @@ def __init__(self, engine: "MemoryEngine") -> None: self._stop = asyncio.Event() # Monotonic timestamps of the last run per job, keyed by job name. self._last_run: dict[str, float] = {} + # Cross-store txn recovery: first-sighting time per pending txn_id, so an unwitnessed + # txn gets a grace period before the sweep aborts it. Persists across ticks. + self._txn_first_seen: dict[str, float] = {} # ── lifecycle ────────────────────────────────────────────────────────── @@ -109,7 +119,25 @@ def _any_job_enabled() -> bool: llm_on = cfg.llm_trace_enabled and cfg.llm_trace_retention_days > 0 mm_refresh_on = cfg.mental_model_refresh_tick_seconds > 0 op_cleanup_on = cfg.operation_retention_days > 0 - return reconcile_on or audit_on or llm_on or mm_refresh_on or op_cleanup_on + return ( + reconcile_on + or audit_on + or llm_on + or mm_refresh_on + or op_cleanup_on + or MaintenanceLoop._memlake_recovery_enabled() + ) + + @staticmethod + def _memlake_recovery_enabled() -> bool: + """True when the memories store keeps memories outside SQL (memlake) and therefore has + cross-store write-group txns a crashed writer could leave undecided.""" + try: + from .memories import get_memories + + return not get_memories().writes_memory_rows_in_sql + except Exception: + return False # ── loop ─────────────────────────────────────────────────────────────── @@ -145,6 +173,8 @@ async def _tick(self) -> None: await self._run_timed("scheduled mental model refresh", self._run_scheduled_mm_refresh()) if cfg.operation_retention_days > 0 and self._is_due("operation_cleanup", _OPERATION_CLEANUP_INTERVAL_SECONDS): await self._run_timed("operation cleanup", self._run_operation_cleanup(cfg)) + if self._memlake_recovery_enabled() and self._is_due("txn_recovery", _TXN_RECOVERY_INTERVAL_SECONDS): + await self._run_timed("memlake txn recovery", self._run_txn_recovery()) async def _run_timed(self, name: str, coro: Coroutine[Any, Any, None]) -> None: """Run a maintenance job and emit one timing line for it. @@ -259,6 +289,41 @@ async def _run_operation_cleanup(self, cfg: HindsightConfig) -> None: if pruned: logger.info(f"Operation cleanup: pruned {pruned} operation(s) total") + # ── memlake cross-store txn recovery ───────────────────────────────────── + + async def _run_txn_recovery(self) -> None: + """Resolve write-group txns a crashed writer left undecided, for the memlake store. + + For each bank, the store lists its namespace's pending txns and decides each against the + Postgres witness table (present ⇒ commit, absent past the grace ⇒ abort — never on + assumption), then reaps expired witness rows. A no-op for the SQL stores. Best-effort: a + failure here only delays a stalled fold until the next tick. + """ + from .memories import get_memories + + store = get_memories() + if store.writes_memory_rows_in_sql: + return + backend = self._engine._backend + try: + async with acquire_with_retry(backend, max_retries=1) as conn: + bank_ids = [r[0] for r in await conn.fetch(f"SELECT bank_id FROM {fq_table('banks')}")] + if not bank_ids: + return + decided = await store.recover_pending_txns( + conn=conn, + fq_table=fq_table, + bank_ids=bank_ids, + first_seen=self._txn_first_seen, + now=time.monotonic(), + grace_seconds=_TXN_RECOVERY_GRACE_SECONDS, + ) + except Exception as e: + logger.warning(f"Memlake txn recovery failed: {e}") + return + if decided: + logger.info(f"Memlake txn recovery: decided {decided} undecided txn(s)") + # ── consolidation reconcile ────────────────────────────────────────────── async def _run_reconcile(self) -> None: diff --git a/hindsight-api-slim/hindsight_api/engine/memories/__init__.py b/hindsight-api-slim/hindsight_api/engine/memories/__init__.py new file mode 100644 index 0000000000..876e5e1467 --- /dev/null +++ b/hindsight-api-slim/hindsight_api/engine/memories/__init__.py @@ -0,0 +1,90 @@ +"""The memories store: which one is installed, and how the engine reaches it. + +Resolved through the ordinary extension loader — ``HINDSIGHT_API_MEMORIES_EXTENSION`` +names a ``module:Class``, and ``HINDSIGHT_API_MEMORIES_*`` becomes its config — so +this behaves like every other extension point. Unset (the normal case) means +:class:`~hindsight_api.engine.memories.postgres.PostgresMemories`: rows in +`memory_units`, links in `memory_links` / `unit_entities`, retrieval as SQL. +""" + +from __future__ import annotations + +import logging + +from .base import ( + FACT_TYPE_TO_MEMORY_TYPE, + MEMORY_TYPE_TO_FACT_TYPE, + META_CHUNK_ID, + CausalEdgeRecord, + DeletePredicate, + FactRecord, + MemoriesExtension, + MemoryPatch, + ScanPage, + StoredMemory, + build_fact_records, + build_text_signals, + source_key, +) + +logger = logging.getLogger(__name__) + +_memories: MemoriesExtension | None = None + + +def create_memories(context=None) -> MemoriesExtension: + """Build the configured memories store, or the Postgres default.""" + from ...extensions.loader import load_extension + + loaded = load_extension("MEMORIES", MemoriesExtension, context=context) + if loaded is not None: + logger.info("[memories] store=%s (memory rows do not go to postgres)", loaded.name) + return loaded + + from .postgres import PostgresMemories + + return PostgresMemories({}) + + +def get_memories() -> MemoriesExtension: + """The process-wide memories store, built on first use. + + Retrieval and the retain pipeline reach it through call chains that do not + carry the engine, so it is resolved here rather than threaded through every + signature. + """ + global _memories + if _memories is None: + _memories = create_memories() + return _memories + + +def set_memories(memories: MemoriesExtension | None) -> None: + """Override the store (tests, and engine startup after initialize()).""" + global _memories + _memories = memories + # The graph arm's retriever is chosen from the store and then cached, so it + # has to be re-resolved whenever the store changes. + from ..search.retrieval import set_default_graph_retriever + + set_default_graph_retriever(None) + + +__all__ = [ + "FACT_TYPE_TO_MEMORY_TYPE", + "MEMORY_TYPE_TO_FACT_TYPE", + "META_CHUNK_ID", + "CausalEdgeRecord", + "DeletePredicate", + "FactRecord", + "MemoriesExtension", + "MemoryPatch", + "ScanPage", + "StoredMemory", + "build_fact_records", + "build_text_signals", + "create_memories", + "get_memories", + "set_memories", + "source_key", +] diff --git a/hindsight-api-slim/hindsight_api/engine/memories/base.py b/hindsight-api-slim/hindsight_api/engine/memories/base.py new file mode 100644 index 0000000000..be60eab9c1 --- /dev/null +++ b/hindsight-api-slim/hindsight_api/engine/memories/base.py @@ -0,0 +1,1138 @@ +"""Extension interface for the *memories* slice of storage. + +`memory_units` and the link tables around it (`memory_links`, `unit_entities`) +are the one part of the schema that is a search index as much as a table: every +recall arm — semantic, BM25, graph, temporal — is a query over them. This module +carves that slice out from behind the raw SQL so a different engine can own it, +without touching how documents, chunks, banks, operations or the entity registry +are stored. + +The default :class:`~hindsight_api.engine.memories.postgres.PostgresMemories` +keeps everything exactly where it has always been: rows in `memory_units`, links +in `memory_links` and `unit_entities`, retrieval as SQL. It is what runs unless +an extension is configured, and it is the implementation the test suite +exercises. + +An alternative implementation is loaded like any other Hindsight extension:: + + HINDSIGHT_API_MEMORIES_EXTENSION=mypackage.memories:MyMemories + HINDSIGHT_API_MEMORIES_SOME_SETTING=value + +Such an implementation is the **sole store** for memories: no memory- or +link-shaped row reaches Postgres at all. Unit ids are minted by +:meth:`MemoriesExtension.allocate_unit_ids` rather than by an INSERT's RETURNING +clause, facts carry their entity ids and causal edges inline instead of becoming +join rows, and recall results come back fully populated with no Postgres +hydration. Everything else — documents, chunks, banks, the `entities` registry — +stays in Postgres either way. + +Every operation the engine needs is a method here, so no call site branches on +which implementation is installed. Where the two differ, they differ by what the +method does: the Postgres implementation writes join rows and reprocesses links; +one that owns the store no-ops those passes and does its own thing. +""" + +from __future__ import annotations + +import json +import uuid +from abc import ABC, abstractmethod +from dataclasses import dataclass, field +from datetime import datetime, timezone +from typing import TYPE_CHECKING, Any + +from ...extensions.base import Extension + +if TYPE_CHECKING: # pragma: no cover - typing only + from ..search.retrieval import GraphRetriever, SemanticBm25Result + + +class MemoryTxn: + """Opaque token for a cross-store write-group transaction, threaded from + :meth:`MemoriesExtension.begin_txn` through the write calls to + :meth:`MemoriesExtension.decide_txn`. + + A store that keeps memories in the same database as the caller's transaction has nothing + to coordinate and returns ``None`` from ``begin_txn`` — the write methods then receive + ``txn=None`` and behave exactly as before. A store that writes memories to a *separate* + system (memlake) subclasses this to carry whatever it needs to defer the writes' + visibility until the caller's transaction is known to have committed. The base is + deliberately empty: only the store that minted a handle interprets it.""" + + +# Hindsight's fact_type strings <-> a numeric memory type. An implementation that +# indexes per type (results are never fused across types) lines up exactly with +# how recall queries one arm set per fact_type. +FACT_TYPE_TO_MEMORY_TYPE: dict[str, int] = { + "world": 1, + "experience": 2, + "observation": 3, +} +MEMORY_TYPE_TO_FACT_TYPE: dict[int, str] = {v: k for k, v in FACT_TYPE_TO_MEMORY_TYPE.items()} + +# Keys used in an implementation's opaque metadata bag for the `memory_units` +# columns it has no first-class model of. These round-trip verbatim: they are +# stored without interpretation and returned on every hit, which is what lets +# recall rebuild a full result row without touching Postgres. +# +# Nothing here is queryable — an implementation cannot filter or sort on these. A +# column that retrieval must *filter* on has to be modelled properly instead. +META_CONTEXT = "context" +META_DOCUMENT_ID = "document_id" +META_CHUNK_ID = "chunk_id" +META_METADATA_JSON = "metadata_json" +META_OBSERVATION_SCOPES = "observation_scopes" +META_TEXT_SIGNALS = "text_signals" +META_CREATED_AT = "created_at" +META_UPDATED_AT = "updated_at" +# Observation bookkeeping. `source_memory_ids` is a JSON list: an implementation +# with no edge relation carries an observation's sources denormalised. +META_SOURCE_MEMORY_IDS = "source_memory_ids" +META_CONSOLIDATED_AT = "consolidated_at" +# A *positive* flag mirroring META_CONSOLIDATED_AT, because a metadata predicate +# can only match equality — there is no "key is absent". Consolidation's candidate +# query is "not yet consolidated", so it needs a value to match on: every memory is +# written with "0" and flipped to "1" once folded into an observation. +META_CONSOLIDATED_FLAG = "consolidated" +CONSOLIDATED_NO = "0" +CONSOLIDATED_YES = "1" + +#: Prefix for the per-source metadata key an observation carries, one per source. +#: The forward list (:data:`META_SOURCE_MEMORY_IDS`) reads an observation's +#: sources; these read the other direction — "observations built on this fact" — +#: as an equality predicate rather than a corpus walk. +META_SOURCE_KEY_PREFIX = "src:" + + +def source_key(unit_id: str) -> str: + """The metadata key marking an observation as built on ``unit_id``.""" + return f"{META_SOURCE_KEY_PREFIX}{unit_id}" + + +@dataclass +class CausalEdgeRecord: + """A causal edge, resolved to the target's unit id.""" + + target_unit_id: str + relation_type: str # "caused_by" for retain; legacy types on transfer import + weight: float = 1.0 + + +@dataclass +class StoredMemory: + """A memory read by address rather than by ranking. + + What comes back from a get-by-id or a scan: no arm scores, because nothing + ranked it. Shaped like a `memory_units` row so the callers that render one + (the curation UI, export) need no second shape. + """ + + unit_id: str + text: str + fact_type: str + context: str | None = None + document_id: str | None = None + chunk_id: str | None = None + tags: list[str] = field(default_factory=list) + metadata: dict | None = None + proof_count: int = 1 + event_date: datetime | None = None + occurred_start: datetime | None = None + occurred_end: datetime | None = None + mentioned_at: datetime | None = None + created_at: datetime | None = None + # Which observation scopes a memory is routed to. Consolidation reads it off + # its candidates to decide which observation each one belongs in, so it has + # to survive the round trip through the store. + observation_scopes: list | None = None + entity_ids: list[str] = field(default_factory=list) + source_memory_ids: list[str] = field(default_factory=list) + consolidated_at: datetime | None = None + # Derived kNN edges `(target_unit_id, weight)`, populated only when the read + # asked for them — the ranking path never does. + semantic_edges: list[tuple[str, float]] = field(default_factory=list) + + +@dataclass +class MemoryPatch: + """A partial update to one memory. Unset fields are left alone. + + ``proof_count_delta`` is relative; everything else is an absolute set. + ``metadata`` merges into the existing bag rather than replacing it. + """ + + unit_id: str + text: str | None = None + # Either a float list or the pgvector literal '[0.1,0.2,...]' — Hindsight + # carries embeddings in both forms depending on the call site. + embedding: list[float] | str | None = None + tags: list[str] | None = None + event_date: datetime | None = None + occurred_start: datetime | None = None + occurred_end: datetime | None = None + mentioned_at: datetime | None = None + metadata: dict[str, str] | None = None + proof_count_delta: int = 0 + + +@dataclass +class DeletePredicate: + """Which memories a predicate-delete removes: type AND metadata AND tags. + + An empty predicate is refused unless ``delete_all`` — a stray empty filter + must not be able to wipe a bank. + """ + + fact_types: list[str] | None = None + metadata_equals: dict[str, str] | None = None + tags: list[str] | None = None + tags_match: str = "any" + delete_all: bool = False + + def is_empty(self) -> bool: + # A fact_type restriction is a real constraint, so a predicate carrying only + # ``fact_types`` is NOT empty — it scopes the delete to those types (e.g. clearing + # just a bank's observations), and must not be refused as a stray empty filter. + return not self.metadata_equals and not self.tags and not self.fact_types + + +@dataclass +class ScanPage: + """One page of a scan, plus the cursor for the next. + + ``next_page_token`` is empty when the walk is exhausted. It is a *position*, + not a snapshot: concurrent writes can shift later pages, so a scan is + eventually-complete browsing rather than a consistent iterator. + """ + + memories: list[StoredMemory] = field(default_factory=list) + next_page_token: str = "" + + +@dataclass +class FactRecord: + """One memory unit, as an implementation that owns the store needs to see it. + + There is no row behind this — it is the *whole* record — so it carries every + column recall returns, plus the edges that would otherwise have become + `memory_links` and `unit_entities` rows. + """ + + unit_id: str # UUID string + text: str + # A float list, or the pgvector literal '[0.1,...]' — Hindsight produces both. + embedding: list[float] | str + fact_type: str + tags: list[str] = field(default_factory=list) + proof_count: int = 1 + context: str | None = None + document_id: str | None = None + chunk_id: str | None = None + metadata: dict | None = None + observation_scopes: list | str | None = None + # Entity names + spelled-out date tokens Hindsight folds into its BM25 document. + text_signals: str | None = None + event_date: datetime | None = None + occurred_start: datetime | None = None + occurred_end: datetime | None = None + mentioned_at: datetime | None = None + created_at: datetime | None = None + # What would have become `unit_entities` rows: the entity registry stays in + # Postgres, but the unit→entity posting travels with the memory. + entity_ids: list[str] = field(default_factory=list) + # What would have become causal `memory_links` rows. + causal_edges: list[CausalEdgeRecord] = field(default_factory=list) + # Observations only: the facts this observation was consolidated from. + source_memory_ids: list[str] = field(default_factory=list) + # When this memory was folded into an observation (sources only). + consolidated_at: datetime | None = None + + def metadata_bag(self) -> dict[str, str]: + """Render the non-modelled columns as an opaque str→str bag.""" + bag: dict[str, str] = {} + if self.context: + bag[META_CONTEXT] = self.context + if self.document_id: + bag[META_DOCUMENT_ID] = self.document_id + if self.chunk_id: + bag[META_CHUNK_ID] = self.chunk_id + if self.metadata: + bag[META_METADATA_JSON] = json.dumps(self.metadata) + if self.observation_scopes is not None: + bag[META_OBSERVATION_SCOPES] = json.dumps(self.observation_scopes) + if self.text_signals: + bag[META_TEXT_SIGNALS] = self.text_signals + if self.created_at is not None: + bag[META_CREATED_AT] = self.created_at.isoformat() + # Hindsight filters recall's created_after/created_before window on + # updated_at. A freshly written fact has updated_at == created_at. + stamp = self.created_at + if stamp is not None: + bag[META_UPDATED_AT] = stamp.isoformat() + if self.source_memory_ids: + # Forward direction: the list, for reading an observation's sources back. + bag[META_SOURCE_MEMORY_IDS] = json.dumps(self.source_memory_ids) + # Backward direction: one key per source, so "observations built on + # this fact" is an equality predicate rather than a corpus walk. + for source_id in self.source_memory_ids: + bag[source_key(source_id)] = "1" + if self.consolidated_at is not None: + bag[META_CONSOLIDATED_AT] = self.consolidated_at.isoformat() + # Observations are not themselves consolidated, so only sources carry the flag. + if self.fact_type != "observation": + bag[META_CONSOLIDATED_FLAG] = CONSOLIDATED_YES if self.consolidated_at else CONSOLIDATED_NO + return bag + + +def build_text_signals(fact) -> str | None: + """Entity names + spelled-out dates — the enrichment Hindsight folds into BM25. + + Mirrors the signal construction the `memory_units` INSERT performs, so an + implementation that owns the store produces the same searchable document the + SQL path does. + """ + parts: list[str] = [] + if fact.entities: + parts.extend(e.name for e in fact.entities) + stamps = [fact.occurred_start] + if fact.occurred_end and fact.occurred_end != fact.occurred_start: + stamps.append(fact.occurred_end) + for stamp in stamps: + if stamp is None: + continue + try: + parts.append(stamp.strftime("%B %d %Y").lstrip("0").replace(" 0", " ")) + except (ValueError, AttributeError): + pass + return " ".join(parts) if parts else None + + +def build_fact_records( + unit_ids: list[str], + facts: list, + document_id: str | None = None, + unit_entity_ids: dict[str, list[str]] | None = None, +) -> list[FactRecord]: + """Turn the retain pipeline's facts into records, edges resolved. + + ``unit_entity_ids`` is the unit→entity posting that would otherwise become + `unit_entities` rows; causal relations become the memory's causal edges. Both + travel with the memory, which is why a store that owns them writes once rather + than inserting and then linking. + + Only called by implementations that own the store — the Postgres one already + wrote all of this and never builds a record. + """ + now = datetime.now(timezone.utc) + records: list[FactRecord] = [] + for index, (unit_id, fact) in enumerate(zip(unit_ids, facts)): + entity_ids = (unit_entity_ids or {}).get(str(unit_id)) + if entity_ids is None: + entity_ids = [str(e.entity_id) for e in (fact.entities or []) if e.entity_id is not None] + + causal_edges = [] + for relation in fact.causal_relations or []: + target = relation.target_fact_index + # Targets are indices into this batch; a stale index would otherwise + # produce an edge pointing at the wrong memory. + if not isinstance(target, int) or not 0 <= target < len(unit_ids) or target == index: + continue + causal_edges.append( + CausalEdgeRecord(target_unit_id=str(unit_ids[target]), relation_type=relation.relation_type) + ) + + records.append( + FactRecord( + unit_id=str(unit_id), + text=fact.fact_text, + embedding=fact.embedding, + fact_type=fact.fact_type, + tags=fact.tags or [], + context=fact.context, + document_id=fact.document_id or document_id, + chunk_id=fact.chunk_id, + metadata=fact.metadata, + observation_scopes=fact.observation_scopes, + text_signals=build_text_signals(fact), + event_date=fact.occurred_start if fact.occurred_start is not None else fact.mentioned_at, + occurred_start=fact.occurred_start, + occurred_end=fact.occurred_end, + mentioned_at=fact.mentioned_at, + created_at=now, + entity_ids=entity_ids, + causal_edges=causal_edges, + ) + ) + return records + + +class MemoriesExtension(Extension, ABC): + """Storage + retrieval for memory units and their links, behind one interface. + + Loaded with the ``MEMORIES`` prefix; see the module docstring. Subclasses get + ``self.config`` (the ``HINDSIGHT_API_MEMORIES_*`` environment) and + ``self.context`` from :class:`~hindsight_api.extensions.base.Extension`. + + Methods are grouped by what calls them: the retain write path, the recall + arms, addressed reads for curation/export, and the maintenance passes. The + Postgres implementation is the reference for what each one must mean. + """ + + #: Name for logs and the startup banner. + name: str = "postgres" + + #: Whether memories live as rows in the SQL ``memory_units`` table. True for the SQL stores + #: (Postgres/Oracle), whose ``upsert_observation`` / ``delete_facts`` are no-ops because the + #: consolidator writes those rows inline. A store that keeps memories elsewhere sets this + #: False so the consolidator skips the inline SQL and routes the write through the store — + #: then all of an observation's state lives wherever the store keeps it, not in Postgres. + writes_memory_rows_in_sql: bool = True + + #: Whether this store owns the document/chunk BODIES — a document's extracted text, its chunk + #: texts, and its original uploaded file. Default False: Postgres keeps ``documents.original_text`` + #: / ``chunks.chunk_text`` and the file goes through ``file_storage``. A store that sets this True + #: (memlake) owns a dedicated document store, so the retain and read paths route document/chunk + #: bodies through the ``put_document`` / ``get_document_record`` / ``get_chunk_text`` / + #: ``list_chunk_texts`` / ``count_chunks`` / ``document_content_hash`` methods below instead of + #: the inline SQL. Cold, never-searched, key-based — see docs/documents-chunks.md. + owns_document_store: bool = False + + # ------------------------------------------------------------------ lifecycle + + async def initialize(self) -> None: + """Open connections/channels. Called once during engine startup. + + Separate from :meth:`Extension.on_startup` because the memories store has + to be live before the engine finishes booting, not alongside the HTTP app. + """ + + async def shutdown(self) -> None: + """Release resources. Called during engine shutdown.""" + + async def ensure_namespace(self, bank_id: str) -> None: + """Ensure per-bank storage exists. Idempotent.""" + + def allocate_unit_ids(self, count: int) -> list[str]: + """Mint unit ids for a batch about to be written. + + The Postgres path never calls this — its ids come back from the INSERT's + RETURNING clause — so this is what an implementation that owns the store + uses to name memories before writing them. + """ + return [str(uuid.uuid4()) for _ in range(count)] + + # ------------------------------------------------------------------ writes + + async def begin_txn(self, *, conn, fq_table, bank_id: str, mutating: bool) -> "MemoryTxn | None": + """Open a cross-store write-group transaction around a unit of work, or ``None``. + + Called INSIDE the caller's database transaction, before the writes that belong to it. + The returned handle is threaded (as ``txn=``) into every write of the unit and finally + into :meth:`decide_txn` once the caller's transaction has settled. + + Default is ``None``: a store whose memories live in the caller's own database needs no + cross-store coordination — its writes are already covered by that transaction, and the + ``txn`` kwarg is ignored everywhere. A store that writes to a *separate* system returns + a handle so those writes can be held invisible until the transaction is known to have + committed. ``mutating`` distinguishes a unit that only creates new memories (safe to + write plainly and compensate on abort) from one that changes or removes existing ones + (whose previous value only deferred visibility can preserve).""" + return None + + async def decide_txn(self, txn: "MemoryTxn | None", *, commit: bool) -> None: + """Resolve a handle from :meth:`begin_txn` after its transaction settled. + + ``commit=True`` once the caller's transaction has COMMITTED, ``commit=False`` if it + aborted. A no-op for ``None``. For a separate-store implementation this is where the + held writes are made visible (commit) or discarded/compensated (abort).""" + return None + + async def mint_txn(self, *, bank_id: str, mutating: bool) -> "MemoryTxn | None": + """Mint a write-group handle WITHOUT opening a database transaction — the split form of + :meth:`begin_txn` for a unit of work (consolidation) that runs slow work between its + writes and must not hold a transaction across it. Tag the writes with the handle, then + :meth:`write_txn_witness` + commit in one short transaction at the end, then + :meth:`decide_txn`. Default ``None`` (no cross-store coordination).""" + return None + + async def write_txn_witness(self, txn: "MemoryTxn | None", *, conn, fq_table) -> None: + """Record a :meth:`mint_txn` handle's commit witness in the caller's transaction, just + before it commits. No-op for ``None``.""" + return None + + async def recover_pending_txns( + self, + *, + conn, + fq_table, + bank_ids: list[str], + first_seen: dict[str, float], + now: float, + grace_seconds: float = 300.0, + witness_ttl_seconds: float = 3600.0, + ) -> int: + """Backstop for a crashed writer: resolve each bank's undecided write-group txns against + the witness table. Only a store that keeps memory rows outside SQL has cross-store txns to + recover; the default (Postgres) has none, and the maintenance loop skips it. Returns the + number of txns decided.""" + return 0 + + @abstractmethod + async def insert_facts( + self, + *, + conn, + ops, + bank_id: str, + facts: list, + document_id: str | None = None, + defer_index: bool = False, + txn: "MemoryTxn | None" = None, + ) -> list[str]: + """Store a batch of extracted facts and return their unit ids, in order. + + ``defer_index`` asks for ids *without* the write, because the retain + orchestrator can only supply entity ids and causal edges after Phase-1 + placeholders have been remapped onto real unit ids; it then calls + :meth:`index_facts` with the complete picture. An implementation whose + write is the row insert itself ignores the flag. + + ``conn`` and ``ops`` are the live Postgres connection and dialect ops, + used only by an implementation that keeps its rows there. + """ + + async def index_facts( + self, + bank_id: str, + unit_ids: list[str], + facts: list, + document_id: str | None = None, + unit_entity_ids: dict[str, list[str]] | None = None, + ) -> None: + """Index facts whose ids came from a deferred :meth:`insert_facts`. + + A no-op by default: for Postgres the row *is* the index entry, so there is + nothing left to do, and nothing is built. :func:`build_fact_records` turns + the arguments into records for implementations that need them. + """ + + @abstractmethod + async def delete_facts(self, bank_id: str, unit_ids: list[str], *, txn: "MemoryTxn | None" = None) -> None: + """Remove units. Safe to call for ids that were never written.""" + + async def delete_where(self, bank_id: str, predicate: DeletePredicate, txn=None) -> int: + """Remove every memory matching ``predicate``. Returns the count when known. + + May be implemented lazily (recording the delete and materializing it + later), in which case the returned count is 0 rather than a scan. + """ + raise NotImplementedError + + @abstractmethod + async def delete_document( + self, *, conn, fq_table, bank_id: str, document_id: str, txn: "MemoryTxn | None" = None + ) -> None: + """Remove every memory belonging to ``document_id``. + + Called when a document is replaced, so it races the replacement's writes: + an implementation must remove only what was written *before* this call, + never the facts arriving moments later. + """ + + # ------------------------------------------------------ document/chunk bodies + # + # Only relevant when :attr:`owns_document_store` is True: a store that keeps document/chunk + # BODIES (extracted text, chunk texts, original file) in its own dedicated store rather than in + # ``documents.original_text`` / ``chunks.chunk_text`` / ``file_storage``. The retain and read + # paths branch on ``owns_document_store`` and call these instead of the inline SQL. All bodies + # are cold and never-searched; the document is passed whole (text + ordered chunk texts + file) + # so the store can pack and dedup it — see docs/documents-chunks.md. + + async def put_document( + self, + *, + bank_id: str, + document_id: str, + content_hash: str, + original_text: "str | None", + chunk_texts: list[str], + tags: "list[str] | None" = None, + metadata: "dict | None" = None, + file_bytes: "bytes | None" = None, + file_content_type: str = "", + file_original_name: str = "", + txn: "MemoryTxn | None" = None, + ) -> None: + """Store (or replace) a document's bodies: its extracted text, its ordered chunk texts, and + optionally the original uploaded file. Idempotent by content — re-ingest re-uploads only + what changed. Under a ``txn`` the record commits atomically with the retain's facts.""" + raise NotImplementedError + + async def document_content_hash(self, *, bank_id: str, document_id: str) -> "str | None": + """The stored document's content hash, for the idempotent-skip check; ``None`` if absent.""" + raise NotImplementedError + + async def get_document_record(self, *, bank_id: str, document_id: str, include_text: bool = False) -> "dict | None": + """A document's metadata (and, if asked, its extracted ``original_text``), or ``None``.""" + raise NotImplementedError + + async def get_chunk_text(self, *, bank_id: str, document_id: str, chunk_index: int) -> "str | None": + """One chunk's text by position, or ``None`` if the document/index does not exist.""" + raise NotImplementedError + + async def list_chunk_texts(self, *, bank_id: str, document_id: str) -> "list[str] | None": + """Every chunk's text in order, or ``None`` if the document does not exist.""" + raise NotImplementedError + + async def count_chunks(self, *, bank_id: str, document_id: str) -> int: + """How many chunks a document has (0 if it does not exist).""" + raise NotImplementedError + + async def delete_document_record(self, *, bank_id: str, document_id: str, txn: "MemoryTxn | None" = None) -> None: + """Delete a document's RECORD and bodies from the document store — an EXPLICIT document + deletion, distinct from :meth:`delete_document` (which drops only the document's facts on + re-ingest and must not touch the record, since the replacement's ``put_document`` overwrites + it). No-op for a store that does not own the document store.""" + raise NotImplementedError + + async def delete_namespace(self, bank_id: str) -> None: + """Drop a bank's entire storage. Irreversible. + + A no-op for Postgres, where deleting the bank cascades to its rows. + """ + + async def delete_observations(self, *, conn, fq_table, bank_id: str, txn=None) -> None: + """Remove every observation in a bank, leaving the facts behind it.""" + raise NotImplementedError + + async def update_memories(self, bank_id: str, patches: list[MemoryPatch], txn=None) -> None: + """Apply partial updates. Only the fields set on each patch change.""" + raise NotImplementedError + + # ------------------------------------------------------------------ recall arms + + @abstractmethod + async def search( + self, + *, + conn, + bank_id: str, + fact_types: list[str], + query_embedding: str, + query_text: str, + limit: int, + tags: list[str] | None = None, + tags_match: str = "any", + tag_groups: list | None = None, + created_after: datetime | None = None, + created_before: datetime | None = None, + min_semantic: float | None = None, + min_keyword: float | None = None, + graph_seed_min_similarity: float | None = None, + ) -> "dict[str, SemanticBm25Result]": + """Run the semantic + BM25 arms. + + Returns ``{fact_type: SemanticBm25Result(semantic, bm25, graph_seeds)}`` of + ``RetrievalResult`` — the contract ``retrieve_semantic_bm25_combined`` has. + ``graph_seed_min_similarity`` restricts which semantic hits seed the graph + arm (Postgres populates ``graph_seeds``; a store with its own graph arm + leaves it ``None``). + """ + + @abstractmethod + async def temporal_search( + self, + *, + conn, + bank_id: str, + fact_types: list[str], + query_embedding: str, + start_date: datetime, + end_date: datetime, + limit: int, + semantic_threshold: float = 0.1, + tags: list[str] | None = None, + tags_match: str = "any", + tag_groups: list | None = None, + created_after: datetime | None = None, + created_before: datetime | None = None, + ) -> dict[str, list]: + """Run the temporal arm over ``[start_date, end_date]``. + + Returns ``{fact_type: [RetrievalResult]}``: entry points whose effective + time — ``COALESCE(occurred_start, mentioned_at, occurred_end)`` — falls in + the window, spread one hop and scored by proximity to it. + """ + + def graph_retriever(self) -> "GraphRetriever | None": + """The retriever backing the graph arm, or ``None`` to use the configured one. + + ``None`` means the links are in Postgres and ``config.graph_retriever`` + chooses among the SQL retrievers, as it always has. An implementation that + owns the links returns its own, because the SQL retrievers would walk + tables it never wrote to. + """ + return None + + # ------------------------------------------------------------------ addressed reads + # + # Not retrieval: these serve the curation UI, export, consolidation and stats. + # Every one has a `memory_units` query behind it in the Postgres implementation. + + @abstractmethod + async def get_memories(self, *, conn, fq_table, bank_id: str, unit_ids: list[str]) -> list[StoredMemory]: + """Fetch memories by id. Missing or deleted ids are simply absent.""" + + @abstractmethod + async def scan_memories( + self, + *, + conn, + fq_table, + bank_id: str, + fact_types: list[str] | None = None, + limit: int = 100, + page_token: str = "", + tags: list[str] | None = None, + tags_match: str = "any", + tag_groups: list | None = None, + document_id: str | None = None, + metadata_equals: dict[str, str] | None = None, + skip: int = 0, + include_edges: bool = False, + ) -> ScanPage: + """Page through stored memories. + + A full walk by construction — cost grows with the corpus — so this is for + browsing and export, never for retrieval. + + ``document_id`` is its own filter rather than an entry in + ``metadata_equals`` because it is not metadata everywhere: Postgres has a + real column for it, and a store that keeps it in an opaque bag must still + be asked the same question. + + ``tags_match`` selects a flat tag mode; ``tag_groups`` is the compound form + (a list of AND/OR/NOT trees, AND-ed together) for conditions a flat filter + cannot express, the same shape ``search`` takes. Both are AND-ed with + ``metadata_equals``; a scan walks every member, so they filter what a page + returns rather than what it reads. + """ + + @abstractmethod + async def count_memories(self, *, conn, fq_table, bank_id: str) -> dict[str, int]: + """Live memory count per fact_type.""" + + @abstractmethod + async def list_tags(self, *, conn, fq_table, bank_id: str) -> dict[str, int]: + """Distinct tags in a bank and how many live memories carry each.""" + + @abstractmethod + async def find_unconsolidated( + self, + *, + conn, + fq_table, + bank_id: str, + fact_types: list[str], + limit: int, + scope_tags: list[str] | None = None, + ) -> list[StoredMemory]: + """Memories not yet folded into an observation, oldest first. + + ``scope_tags`` restricts to memories carrying *every* one of them, the + same containment the SQL ``tags @> scope`` expresses. + """ + + async def find_failed_consolidation(self, *, conn, fq_table, bank_id: str) -> list[StoredMemory]: + """Source memories the consolidator marked as permanently failed, for retry to requeue. + + Gated like ``find_unconsolidated``: a SQL store keeps the failure marker in a column and + answers the retry inline, so this default is empty; a store that keeps memories outside + SQL overrides it. Returns experience/world memories only (observations are never + consolidated) — the caller clears them with ``mark_consolidated(when=None)``. + """ + return [] + + @abstractmethod + async def mark_consolidated( + self, + *, + conn, + fq_table, + bank_id: str, + unit_ids: list[str], + when: datetime | None, + failed: bool = False, + txn: "MemoryTxn | None" = None, + ) -> None: + """Stamp (or clear, with ``when=None``) the consolidated marker on sources. + + ``failed`` stamps the failure marker instead, so a memory the LLM could + not consolidate is not retried forever. + """ + + @abstractmethod + async def entity_memory_counts( + self, *, conn, fq_table, bank_id: str, entity_ids: list[str] | None = None + ) -> dict[str, int]: + """Live memory count per entity id. + + Entities with no live memories are absent, so an id passed in and not + returned is an orphan. + """ + + @abstractmethod + async def entities_for_units(self, *, conn, fq_table, bank_id: str, unit_ids: list[str]) -> dict[str, list[str]]: + """The entity ids each unit carries, keyed by unit id.""" + + @abstractmethod + async def entity_map_for_units( + self, *, conn, fq_table, bank_id: str, unit_ids: list[str] + ) -> dict[str, list[dict[str, str]]]: + """``{unit_id: [{entity_id, canonical_name}]}`` — the named form recall renders. + + Like :meth:`entities_for_units` but carrying each entity's label, because + recall shows the name on the fact. An observation with no direct postings + inherits its source memories' entities, so a hit reads the same either way. + """ + + @abstractmethod + async def any_memory_updated_since( + self, + *, + conn, + fq_table, + bank_id: str, + since: datetime, + fact_types: list[str] | None = None, + tags: list[str] | None = None, + tags_match: str = "any", + tag_groups: list | None = None, + ) -> bool: + """Whether any memory in the given scope was written after ``since``. + + Backs the mental-model staleness check, so it must be cheap: a bounded + existence test, never a count. The scope is the mental model's — its flat + tags or compound ``tag_groups``, plus an optional ``fact_types`` filter — + so the same scope that gates a refresh decides whether one is due. + """ + + # ------------------------------------------------------------------ count surfaces + # + # The stats/admin views that aggregate memories by a key: consolidation + # freshness, per-document counts, ingestion over time, observation scopes. For + # Postgres each is one GROUP BY; a store without a queryable index over these + # keys answers them by walking, so cost is O(matching) — acceptable for + # admin/stats surfaces, and the reason these are their own methods rather than + # uses of `count_memories`. + + async def consolidation_freshness(self, *, conn, fq_table, bank_id: str) -> dict[str, Any]: + """``{"last_consolidated_at", "pending", "failed"}`` for a bank. + + ``pending`` / ``failed`` count the world/experience facts not yet folded + into an observation, and those the LLM gave up on. Backs + ``get_bank_freshness``, which reflect() calls often, so keep it cheap. + """ + raise NotImplementedError + + async def document_memory_counts(self, *, conn, fq_table, bank_id: str, document_ids: list[str]) -> dict[str, int]: + """Live memory count per document id, for the documents named. Absent = 0.""" + raise NotImplementedError + + async def link_counts(self, *, conn, fq_table, bank_id: str) -> dict[str, int]: + """``{link_type: count}`` of live links in a bank, for the stats page's link total. + + Keyed by link type (the caller sums the values); an absent type is zero. A store + must answer from its own link representation — Postgres counts ``memory_links`` rows + plus entity-derived edges; a store that keeps links inside the memory counts those — + so the stats page never disagrees with the graph view about whether links exist. + """ + raise NotImplementedError + + async def memories_timeseries( + self, *, conn, fq_table, bank_id: str, time_field: str, trunc: str, since: datetime + ) -> list[dict[str, Any]]: + """``[{"bucket": datetime, "fact_type": str, "count": int}]`` since ``since``. + + Memories bucketed by ``time_field`` truncated to ``trunc`` (minute / hour / + day) on UTC boundaries, broken down by fact_type — the caller fills the + empty buckets. ``time_field`` is one of created_at / mentioned_at / + occurred_start (the event-time fields fall back to created_at per memory). + """ + raise NotImplementedError + + async def observation_scope_counts(self, *, conn, fq_table, bank_id: str) -> list[dict[str, Any]]: + """``[{"tags": list[str], "count": int}]`` — observations grouped by scope. + + A scope is the sorted set of tags an observation was consolidated with; + ``[]`` is the global (untagged) scope. Most-populous first. + """ + raise NotImplementedError + + # ------------------------------------------------------------------ curation reads + # + # These back the curation UI and the bank/entity views. They page and filter, + # which is why they are their own methods rather than uses of `scan_memories`: + # a scan walks the corpus, and these must not. + + @abstractmethod + async def list_memory_units( + self, + *, + conn, + ops, + fq_table, + bank_id: str, + fact_type: str | None = None, + search_query: str | None = None, + consolidation_state: str | None = None, + state: str | None = None, + document_id: str | None = None, + entity_id: str | None = None, + tags: list[str] | None = None, + tags_match: str = "any", + created_before: "datetime | None" = None, + limit: int = 100, + offset: int = 0, + ) -> dict[str, Any]: + """One page of the curation list: ``{"items": [...], "total": int}``. + + ``total`` is the count matching the filters, not the page size, because + the UI pages on it. + """ + + @abstractmethod + async def get_memory_unit(self, *, conn, ops, fq_table, bank_id: str, unit_id: str) -> dict[str, Any] | None: + """One memory rendered for the curation detail view, or ``None``.""" + + # ------------------------------------------------------------------ curation archive + # + # Invalidation is *structural*, not a flag: a memory the curator rejects is + # moved out of every recall surface into an archive it can be restored from, + # so recall / consolidation / graph never need a "valid?" predicate. The two + # implementations realize the archive differently — Postgres moves the row to + # a sibling table, a store that owns its memories moves it to a sibling + # namespace — but the lifecycle is the same, so it lives behind these methods. + + @abstractmethod + async def get_archived_memory(self, *, conn, fq_table, bank_id: str, unit_id: str) -> StoredMemory | None: + """An *invalidated* memory read from the archive, or ``None``. + + Only invalidated memories are in the archive, so a live or missing id + returns ``None`` — which is how a caller tells "invalidated" from "live" + without a state column. + """ + + @abstractmethod + async def invalidate_memory( + self, *, conn, fq_table, bank_id: str, unit_id: str, reason: str | None, txn=None + ) -> bool: + """Move a live memory into the archive, out of every recall surface. + + Returns ``True`` if it was live and is now archived, ``False`` if there was + no live memory with that id. The memory stays retrievable via + :meth:`get_archived_memory` and restorable via :meth:`restore_memory`; + ``reason`` is recorded alongside it. + """ + + @abstractmethod + async def set_invalidation_reason(self, *, conn, fq_table, bank_id: str, unit_id: str, reason: str | None) -> None: + """Update the recorded reason on a memory that is already archived.""" + + @abstractmethod + async def restore_memory(self, *, conn, fq_table, bank_id: str, unit_id: str, txn=None) -> StoredMemory | None: + """Move an archived memory back to the live set, restoring its entity postings. + + Returns the restored memory (so the caller can recompute its embedding — + the archive need not keep one), or ``None`` if it was not archived. + """ + + @abstractmethod + async def set_memory_embedding(self, *, conn, fq_table, bank_id: str, unit_id: str, embedding, txn=None) -> None: + """Write a memory's embedding, recomputed by the caller. + + Its own method because the general :meth:`update_memories` is a no-op for + the store whose write is the row itself — reverting or editing a memory has + to put a freshly computed vector back on it, so this is a real write for + both. ``embedding`` is a float list or the pgvector literal. + """ + + async def clear_unit_entities(self, *, conn, fq_table, bank_id: str, unit_id: str) -> None: + """Drop a unit's entity postings, ahead of an edit re-resolving them. + + A no-op for a store that keeps entity ids on the memory itself — the edit's + rewrite replaces the whole set, so there is nothing to clear first. + """ + + async def apply_edit( + self, + *, + conn, + fq_table, + bank_id: str, + unit_id: str, + text: str, + context: str | None, + fact_type: str, + occurred_start, + occurred_end, + event_date, + mentioned_at, + entity_ids: list[str] | None, + txn=None, + ) -> None: + """Apply a curation field edit to a live memory. + + Writes the new text / context / fact_type / occurred window, resets the + consolidation markers (the memory re-consolidates) and stamps the edit + time, and drops the memory's derived links (they are recomputed). The + embedding is *not* written here — the caller re-embeds from the new fields + and calls :meth:`set_memory_embedding` after. + + ``entity_ids`` is the resolved entity set the memory should now carry; a + store that keeps them on the memory writes them here, one that keeps them + in a join table has already re-linked them and ignores this. ``None`` means + the entity set was not part of this edit. + """ + raise NotImplementedError + + @abstractmethod + async def list_entities( + self, + *, + conn, + fq_table, + bank_id: str, + search: str | None = None, + limit: int = 100, + offset: int = 0, + ) -> dict[str, Any]: + """Entities in a bank with their live memory counts, paged.""" + + @abstractmethod + async def graph_units( + self, + *, + conn, + fq_table, + bank_id: str, + fact_type: str | None = None, + search_query: str | None = None, + document_id: str | None = None, + chunk_id: str | None = None, + tags: list[str] | None = None, + tags_match: str = "all_strict", + limit: int = 1000, + ) -> dict[str, Any]: + """Memory nodes for the graph view, plus the total matching count. + + Returns ``{"units": [...], "total": int}``: the page of nodes (newest + first, capped at ``limit``) and how many match the filters. ``document_id`` + / ``chunk_id`` also match an observation whose sources carry them. + """ + + @abstractmethod + async def graph_entity_rows(self, *, conn, fq_table, bank_id: str, unit_ids: list[str]) -> list[dict[str, Any]]: + """``(unit_id, entity_id, canonical_name)`` rows for the graph view's entity edges.""" + + @abstractmethod + async def graph_direct_links(self, *, conn, fq_table, bank_id: str, unit_ids: list[str]) -> list[dict[str, Any]]: + """Memory-to-memory edges among ``unit_ids`` for the graph view.""" + + # ------------------------------------------------------------------ observations + + async def upsert_observation( + self, *, conn, bank_id: str, record: FactRecord, txn: "MemoryTxn | None" = None + ) -> None: + """Write an observation, replacing any earlier one with the same id.""" + raise NotImplementedError + + @abstractmethod + async def observations_for_sources( + self, *, conn, ops, fq_table, bank_id: str, unit_ids: list[str] + ) -> list[StoredMemory]: + """Observations consolidated from any of ``unit_ids``.""" + + @abstractmethod + async def delete_stale_observations(self, *, conn, ops, fq_table, bank_id: str, fact_ids: list) -> int: + """Delete observations built on ``fact_ids`` and requeue surviving sources. + + Returns how many observations were removed. Called whenever facts are + replaced or deleted, so an observation never outlives the facts it + summarises; sources that survive go back in the consolidation queue. + """ + + # ------------------------------------------------------------------ maintenance + # + # The graph-maintenance job orchestrates these; each pass asks the store to do + # the part it owns. A store whose links are inline has nothing to relink and no + # join table to sweep, so those passes are no-ops for it. + + async def record_unit_entities( + self, *, conn, ops, fq_table, bank_id: str | None = None, unit_ids: list[Any], entity_ids: list[Any] + ) -> None: + """Record the unit→entity postings for a batch of memories. + + ``unit_ids`` and ``entity_ids`` are parallel: a unit that mentions three + entities appears three times. The `entities` registry itself stays in + Postgres regardless; this is the join from a memory to the entities it + mentions. ``bank_id`` is passed because a store that keeps the posting on + the memory (rather than in a global join table) needs to know which + namespace the units live in — the Postgres join is keyed by global unit id + and ignores it. + """ + + async def enqueue_relink_victims( + self, *, conn, fq_table, bank_id: str, affected_unit_ids: list, include_affected_units: bool = False + ) -> int: + """Queue memories that lost a link when ``affected_unit_ids`` changed. + + Zero for a store with no link table to dangle: nothing can point at a + deleted memory if the pointers travel inside the memories themselves. + ``include_affected_units`` (also enqueue the affected units themselves, for + edits that leave them live) is honoured only by a store with a link table. + """ + return 0 + + async def relink_pass(self, *, backend, fq_table, bank_id: str, config) -> dict: + """Top up links for queued victims. ``{}`` when there is nothing to relink.""" + return {} + + async def prune_orphan_entities(self, *, conn, fq_table, bank_id: str) -> int: + """Delete `entities` rows no live memory references. Returns the count.""" + return 0 + + async def prune_stale_cooccurrences(self, *, conn, fq_table, bank_id: str) -> int: + """Delete co-occurrence rows whose witnessing memories are all gone.""" + return 0 + + +__all__ = [ + "CONSOLIDATED_NO", + "CONSOLIDATED_YES", + "FACT_TYPE_TO_MEMORY_TYPE", + "MEMORY_TYPE_TO_FACT_TYPE", + "META_CHUNK_ID", + "META_CONSOLIDATED_AT", + "META_CONSOLIDATED_FLAG", + "META_CONTEXT", + "META_CREATED_AT", + "META_DOCUMENT_ID", + "META_METADATA_JSON", + "META_OBSERVATION_SCOPES", + "META_SOURCE_KEY_PREFIX", + "META_SOURCE_MEMORY_IDS", + "META_TEXT_SIGNALS", + "META_UPDATED_AT", + "CausalEdgeRecord", + "DeletePredicate", + "FactRecord", + "MemoriesExtension", + "MemoryPatch", + "MemoryTxn", + "ScanPage", + "StoredMemory", + "build_fact_records", + "build_text_signals", + "source_key", +] diff --git a/hindsight-api-slim/hindsight_api/engine/memories/pg/__init__.py b/hindsight-api-slim/hindsight_api/engine/memories/pg/__init__.py new file mode 100644 index 0000000000..8a5227935a --- /dev/null +++ b/hindsight-api-slim/hindsight_api/engine/memories/pg/__init__.py @@ -0,0 +1,20 @@ +"""The Postgres memories implementation, split by what calls it. + +:class:`~hindsight_api.engine.memories.postgres.PostgresMemories` is a thin class +over these modules; the queries live here, grouped by concern rather than piled +behind one object: + +* :mod:`counts` — the stats/admin aggregates (freshness, per-doc, timeseries, scopes) +* :mod:`curation` — the memory/entity list and detail views +* :mod:`graph` — the graph view, entity postings, and the maintenance passes +* :mod:`reads` — addressed reads: get, scan, count, tags, consolidation state +* :mod:`writes` — inserts, deletes, and observation invalidation + +Every function here takes the live connection and Hindsight's ``fq_table`` +resolver rather than reaching for globals, so each is callable from a +transaction the caller already owns. +""" + +from __future__ import annotations + +__all__ = ["counts", "curation", "graph", "reads", "writes"] diff --git a/hindsight-api-slim/hindsight_api/engine/memories/pg/counts.py b/hindsight-api-slim/hindsight_api/engine/memories/pg/counts.py new file mode 100644 index 0000000000..3f909d0191 --- /dev/null +++ b/hindsight-api-slim/hindsight_api/engine/memories/pg/counts.py @@ -0,0 +1,151 @@ +"""The count/aggregate surfaces: consolidation freshness, per-document counts, +ingestion over time, observation scopes. + +Each is one ``GROUP BY`` (or filtered ``COUNT``) over `memory_units`. They back +the stats and admin views, not retrieval, so they are grouped here away from the +addressed reads. The SQL is lifted verbatim from the engine methods that used to +carry it; only the connection and ``fq_table`` resolver are now parameters. +""" + +from __future__ import annotations + +from collections.abc import Callable +from datetime import datetime +from typing import Any + + +async def consolidation_freshness(*, conn, fq_table: Callable[[str], str], bank_id: str) -> dict[str, Any]: + """Last consolidation time plus the pending / failed fact counts, in one scan. + + All three come from a single pass so keeping ``failed`` — part of the + published contract — costs nothing over reflect()'s ``pending`` read. + """ + row = await conn.fetchrow( + f""" + SELECT + MAX(consolidated_at) AS last_consolidated_at, + COUNT(*) FILTER (WHERE consolidated_at IS NULL AND fact_type IN ('experience', 'world')) AS pending, + COUNT(*) FILTER (WHERE consolidation_failed_at IS NOT NULL AND fact_type IN ('experience', 'world')) AS failed + FROM {fq_table("memory_units")} + WHERE bank_id = $1 + """, + bank_id, + ) + if row is None: + return {"last_consolidated_at": None, "pending": 0, "failed": 0} + return { + "last_consolidated_at": row["last_consolidated_at"], + "pending": row["pending"] or 0, + "failed": row["failed"] or 0, + } + + +async def link_counts(*, conn, fq_table: Callable[[str], str], bank_id: str) -> dict[str, int]: + """``{link_type: count}`` of live links in a bank. + + Non-entity links (temporal / semantic / caused_by) are a single ``GROUP BY`` over + ``memory_links``. Entity links are no longer stored there — they are derived on demand + from ``unit_entities``, replicating the historical writer cap of ``MAX_LINKS_PER_ENTITY`` + bidirectional edges per shared entity — so they are aggregated to one ``entity`` scalar. + """ + max_links_per_entity = 10 + non_entity_link_rows = await conn.fetch( + f""" + SELECT link_type, COUNT(*) as count + FROM {fq_table("memory_links")} + WHERE bank_id = $1 + GROUP BY link_type + """, + bank_id, + ) + entity_total_row = await conn.fetchrow( + f""" + WITH per_entity AS ( + SELECT ue.entity_id, COUNT(*) AS n + FROM {fq_table("unit_entities")} ue + JOIN {fq_table("memory_units")} mu ON mu.id = ue.unit_id + WHERE mu.bank_id = $1 + GROUP BY ue.entity_id + ) + SELECT COALESCE(SUM(LEAST(n - 1, $2)), 0)::bigint AS count + FROM per_entity + """, + bank_id, + max_links_per_entity, + ) + entity_link_total = int(entity_total_row["count"] or 0) if entity_total_row else 0 + counts: dict[str, int] = {row["link_type"]: row["count"] for row in non_entity_link_rows} + if entity_link_total > 0: + counts["entity"] = entity_link_total + return counts + + +async def document_memory_counts( + *, conn, fq_table: Callable[[str], str], bank_id: str, document_ids: list[str] +) -> dict[str, int]: + """Live memory count per document id, for the ids given.""" + if not document_ids: + return {} + rows = await conn.fetch( + f""" + SELECT document_id, COUNT(*) AS unit_count + FROM {fq_table("memory_units")} + WHERE bank_id = $1 AND document_id = ANY($2::text[]) + GROUP BY document_id + """, + bank_id, + list(document_ids), + ) + return {row["document_id"]: row["unit_count"] for row in rows} + + +async def memories_timeseries( + *, conn, fq_table: Callable[[str], str], bank_id: str, time_field: str, trunc: str, since: datetime +) -> list[dict[str, Any]]: + """Memories bucketed by ``time_field`` (truncated to ``trunc``) and fact_type. + + ``time_field`` is whitelisted by the caller before it reaches here — it is + interpolated into SQL. Event-time fields fall back to ``created_at`` per row so + rows without an event timestamp still appear. + """ + bucket_expr = time_field if time_field == "created_at" else f"COALESCE({time_field}, created_at)" + rows = await conn.fetch( + f""" + SELECT date_trunc('{trunc}', {bucket_expr} AT TIME ZONE 'UTC') AS bucket, + fact_type, COUNT(*) AS count + FROM {fq_table("memory_units")} + WHERE bank_id = $1 AND {bucket_expr} >= $2 + GROUP BY bucket, fact_type + ORDER BY bucket + """, + bank_id, + since, + ) + return [{"bucket": r["bucket"], "fact_type": r["fact_type"], "count": r["count"]} for r in rows] + + +async def observation_scope_counts(*, conn, fq_table: Callable[[str], str], bank_id: str) -> list[dict[str, Any]]: + """Observations grouped by scope (their sorted tag set), most-populous first.""" + rows = await conn.fetch( + f""" + SELECT scope, COUNT(*) AS count + FROM ( + SELECT COALESCE(ARRAY(SELECT unnest(tags) ORDER BY 1), '{{}}'::text[]) AS scope + FROM {fq_table("memory_units")} + WHERE bank_id = $1 AND fact_type = 'observation' + ) s + GROUP BY scope + ORDER BY count DESC, scope + """, + bank_id, + ) + return [{"tags": list(r["scope"]), "count": r["count"]} for r in rows] + + +__all__ = [ + "consolidation_freshness", + "document_memory_counts", + "link_counts", + "memories_timeseries", + "observation_scope_counts", +] diff --git a/hindsight-api-slim/hindsight_api/engine/memories/pg/curation.py b/hindsight-api-slim/hindsight_api/engine/memories/pg/curation.py new file mode 100644 index 0000000000..474da2eb4d --- /dev/null +++ b/hindsight-api-slim/hindsight_api/engine/memories/pg/curation.py @@ -0,0 +1,506 @@ +"""Curation reads: the memory list, the memory detail view, and the entity list. + +These back the curation UI — the table of memories a bank holds, the detail panel +for one of them, and the entity roster beside it. They are paged and filtered +rather than ranked: nothing here scores anything, and nothing walks the corpus. + +Two things separate them from the addressed reads in :mod:`reads`. They render +*view* dicts (ISO strings, joined entity names, a ``state`` discriminator) rather +than :class:`~hindsight_api.engine.memories.base.StoredMemory`, because the HTTP +layer serialises what comes back verbatim. And they read the archive as well as +the live table: curation moves an invalidated fact to `invalidated_memory_units`, +so "show me the invalidated ones" is a different table, not a different predicate. + +Authentication, operation validation and audit stay with the engine methods that +call these — only the queries and their row rendering live here. +""" + +from __future__ import annotations + +import json +from typing import Any + +from ...search.tags import build_tags_where_clause + + +def _entity_rows_for_units_sql(*, ops, fq_table, unit_ids_placeholder: int) -> str: + """SQL SELECT producing ``(unit_id, entity_id, canonical_name)`` rows for + the given unit IDs. + + Direct rows come from ``unit_entities``. Observations rarely carry + direct rows there; their entity association lives transitively through + their source memories (``source_memory_ids`` on PG, the + ``observation_sources`` junction on Oracle). When an observation has + no direct entity rows the SELECT inherits its source memories' + entities, so the result is the same set callers would get from + ``get_memory_unit``. + + ``unit_ids_placeholder`` is the 1-based parameter index that holds the + ``uuid[]`` of unit IDs. The placeholder is referenced twice — both + sides of the UNION need it — so callers should not reuse the slot. + """ + ue = fq_table("unit_entities") + ents = fq_table("entities") + mu = fq_table("memory_units") + p = unit_ids_placeholder + + direct = ( + f"SELECT ue.unit_id, e.id AS entity_id, e.canonical_name " + f"FROM {ue} ue " + f"JOIN {ents} e ON e.id = ue.entity_id " + f"WHERE ue.unit_id = ANY(${p}::uuid[])" + ) + + if ops.uses_observation_sources_table: + os_t = fq_table("observation_sources") + inherited = ( + f"SELECT os.observation_id AS unit_id, e.id AS entity_id, e.canonical_name " + f"FROM {os_t} os " + f"JOIN {ue} src_ue ON src_ue.unit_id = os.source_id " + f"JOIN {ents} e ON e.id = src_ue.entity_id " + f"WHERE os.observation_id = ANY(${p}::uuid[]) " + f"AND NOT EXISTS (SELECT 1 FROM {ue} d WHERE d.unit_id = os.observation_id)" + ) + else: + inherited = ( + f"SELECT obs.id AS unit_id, e.id AS entity_id, e.canonical_name " + f"FROM {mu} obs " + f"CROSS JOIN LATERAL unnest(obs.source_memory_ids) AS src_id " + f"JOIN {ue} src_ue ON src_ue.unit_id = src_id " + f"JOIN {ents} e ON e.id = src_ue.entity_id " + f"WHERE obs.id = ANY(${p}::uuid[]) " + f"AND obs.fact_type = 'observation' " + f"AND obs.source_memory_ids IS NOT NULL " + f"AND NOT EXISTS (SELECT 1 FROM {ue} d WHERE d.unit_id = obs.id)" + ) + + return f"({direct}) UNION ({inherited})" + + +async def list_memory_units( + *, + conn, + ops, + fq_table, + bank_id: str, + fact_type: str | None = None, + search_query: str | None = None, + consolidation_state: str | None = None, + state: str | None = None, + document_id: str | None = None, + entity_id: str | None = None, + tags: list[str] | None = None, + tags_match: str = "any", + created_before: datetime | None = None, + limit: int = 100, + offset: int = 0, +) -> dict[str, Any]: + """ + List memory units for table view with optional full-text search. + + Args: + conn: Open database connection (the caller owns the transaction). + ops: Dialect ops. Unused by this query; part of the interface signature. + fq_table: Table-name resolver. + bank_id: Filter by bank ID + fact_type: Filter by fact type (world, experience) + search_query: Full-text search query (searches text and context fields) + document_id: Optional filter to a single source document. + tags: Optional list of tag names to filter by. When omitted, no tag + filtering is applied (except tags_match='exact', which then selects + the untagged/global scope). + tags_match: How to combine tags (same modes as recall): 'any' (OR, + default) or 'all' (AND) both also include untagged units; + 'any_strict'/'all_strict' exclude untagged units; 'exact' matches + units whose tag set equals the given tags exactly. + state: Optional curation-state filter ('valid' or 'invalidated'). + Invalidated facts live in a separate archive table; 'invalidated' + reads that archive. Omitted/('valid') lists live facts. + consolidation_state: Optional filter on consolidation state. One of + 'failed' (consolidation permanently failed and awaiting recovery), + 'pending' (not yet consolidated, no failure), or + 'done' (successfully consolidated). Only applies to source memory + types (world/experience). + limit: Maximum number of results to return + offset: Offset for pagination + + Returns: + Dict with items (list of memory units) and total count + """ + if state is not None and state not in ("valid", "invalidated"): + raise ValueError(f"Invalid state '{state}': expected 'valid' or 'invalidated'.") + if entity_id is not None: + import uuid as _uuid + + try: + _uuid.UUID(entity_id) + except ValueError: + raise ValueError(f"Invalid entity_id: '{entity_id}' is not a valid UUID") from None + # Invalidated facts live in a separate archive table; pick the source + # accordingly. Default (state is None) lists live facts. + is_archived = state == "invalidated" + source_table = fq_table("invalidated_memory_units") if is_archived else fq_table("memory_units") + + # Build query conditions + query_conditions = [] + query_params = [] + param_count = 0 + + if bank_id: + param_count += 1 + query_conditions.append(f"bank_id = ${param_count}") + query_params.append(bank_id) + + if fact_type: + param_count += 1 + query_conditions.append(f"fact_type = ${param_count}") + query_params.append(fact_type) + + if document_id: + param_count += 1 + query_conditions.append(f"document_id = ${param_count}") + query_params.append(document_id) + + if entity_id: + # Reverse lookup via the stored entity links. Entity links reference live memory units, so + # this yields nothing against the invalidated archive (documented on the method). + param_count += 1 + query_conditions.append( + f"id IN (SELECT unit_id FROM {fq_table('unit_entities')} WHERE entity_id = ${param_count}::uuid)" + ) + query_params.append(entity_id) + + if search_query: + # Full-text search on text and context fields using ILIKE + param_count += 1 + query_conditions.append(f"(text ILIKE ${param_count} OR context ILIKE ${param_count})") + query_params.append(f"%{search_query}%") + + if consolidation_state: + # Named apart from `state`, which the engine method used to shadow here; + # `is_archived` was already resolved above, so behaviour is unchanged. + wanted = consolidation_state.lower() + if wanted == "failed": + query_conditions.append("consolidation_failed_at IS NOT NULL AND fact_type IN ('experience', 'world')") + elif wanted == "pending": + query_conditions.append( + "consolidated_at IS NULL AND consolidation_failed_at IS NULL AND fact_type IN ('experience', 'world')" + ) + elif wanted == "done": + query_conditions.append("consolidated_at IS NOT NULL AND fact_type IN ('experience', 'world')") + else: + raise ValueError( + f"Invalid consolidation_state '{consolidation_state}': expected 'failed', 'pending', or 'done'." + ) + + if tags: + tags_clause, tags_params, next_param = build_tags_where_clause(tags, param_count + 1, "", tags_match) + if tags_clause: + query_conditions.append(tags_clause.removeprefix("AND ")) + query_params.extend(tags_params) + param_count = next_param - 1 + elif tags_match == "exact": + # Exact match with no tags is the "global" scope: rows that carry no + # tags at all. (Other match modes treat empty tags as "no filter".) + query_conditions.append("(tags IS NULL OR tags = '{}')") + + if created_before is not None: + param_count += 1 + query_conditions.append(f"created_at < ${param_count}") + query_params.append(created_before) + + where_clause = "WHERE " + " AND ".join(query_conditions) if query_conditions else "" + + # Get total count + count_query = f""" + SELECT COUNT(*) as total + FROM {source_table} + {where_clause} + """ + count_result = await conn.fetchrow(count_query, *query_params) + total = count_result["total"] + + # Get units with limit and offset + param_count += 1 + limit_param = f"${param_count}" + query_params.append(limit) + + param_count += 1 + offset_param = f"${param_count}" + query_params.append(offset) + + # The archive carries invalidation bookkeeping; the live table doesn't. + curation_cols = ( + "invalidation_reason, invalidated_at" + if is_archived + else "NULL::text AS invalidation_reason, NULL::timestamptz AS invalidated_at" + ) + units = await conn.fetch( + f""" + SELECT id, text, event_date, context, fact_type, document_id, + mentioned_at, occurred_start, occurred_end, chunk_id, proof_count, + tags, metadata, consolidated_at, consolidation_failed_at, edited_at, {curation_cols} + FROM {source_table} + {where_clause} + ORDER BY mentioned_at DESC NULLS LAST, created_at DESC + LIMIT {limit_param} OFFSET {offset_param} + """, + *query_params, + ) + + # Get entity information for these units + if units: + unit_ids = [row["id"] for row in units] + unit_entities = await conn.fetch( + f""" + SELECT ue.unit_id, e.canonical_name + FROM {fq_table("unit_entities")} ue + JOIN {fq_table("entities")} e ON ue.entity_id = e.id + WHERE ue.unit_id = ANY($1::uuid[]) + ORDER BY ue.unit_id + """, + unit_ids, + ) + else: + unit_entities = [] + + # Build entity mapping + entity_map: dict[Any, list[str]] = {} + for row in unit_entities: + unit_id = row["unit_id"] + entity_name = row["canonical_name"] + if unit_id not in entity_map: + entity_map[unit_id] = [] + entity_map[unit_id].append(entity_name) + + # Build result items + items = [] + for row in units: + unit_id = row["id"] + entities = entity_map.get(unit_id, []) + + items.append( + { + "id": str(unit_id), + "text": row["text"], + "context": row["context"] if row["context"] else "", + "date": row["event_date"].isoformat() if row["event_date"] else "", + "fact_type": row["fact_type"], + "document_id": row["document_id"], + "mentioned_at": row["mentioned_at"].isoformat() if row["mentioned_at"] else None, + "occurred_start": row["occurred_start"].isoformat() if row["occurred_start"] else None, + "occurred_end": row["occurred_end"].isoformat() if row["occurred_end"] else None, + "entities": ", ".join(entities) if entities else "", + "chunk_id": row["chunk_id"] if row["chunk_id"] else None, + "proof_count": row["proof_count"] if row["proof_count"] is not None else 1, + "tags": list(row["tags"]) if row["tags"] else [], + "metadata": conn.parse_json(row["metadata"]) if row["metadata"] is not None else {}, + "consolidated_at": row["consolidated_at"].isoformat() if row["consolidated_at"] else None, + "consolidation_failed_at": ( + row["consolidation_failed_at"].isoformat() if row["consolidation_failed_at"] else None + ), + "state": "invalidated" if is_archived else "valid", + "invalidation_reason": row["invalidation_reason"], + "invalidated_at": row["invalidated_at"].isoformat() if row["invalidated_at"] else None, + "edited_at": row["edited_at"].isoformat() if row["edited_at"] else None, + } + ) + + return {"items": items, "total": total, "limit": limit, "offset": offset} + + +async def get_memory_unit(*, conn, ops, fq_table, bank_id: str, unit_id: str) -> dict[str, Any] | None: + """ + Get a single memory unit by ID. + + Args: + conn: Open database connection (the caller owns the transaction). + ops: Dialect ops, for the observation→source entity inheritance shape. + fq_table: Table-name resolver. + bank_id: Bank ID + unit_id: Memory unit ID (the caller validates it is a UUID) + + Returns: + Dict with memory unit data or None if not found + """ + # Get the memory unit (include source_memory_ids for mental models). + # Curation moves invalidated facts to invalidated_memory_units, so fall + # back to the archive (with its invalidation bookkeeping) on a miss. + select_cols = ( + "id, text, context, event_date, occurred_start, occurred_end, " + "mentioned_at, fact_type, document_id, chunk_id, tags, metadata, source_memory_ids, " + "observation_scopes, edited_at" + ) + row = await conn.fetchrow( + f"SELECT {select_cols}, NULL::text AS invalidation_reason, NULL::timestamptz AS invalidated_at " + f"FROM {fq_table('memory_units')} WHERE id = $1 AND bank_id = $2", + unit_id, + bank_id, + ) + unit_state = "valid" + if not row: + row = await conn.fetchrow( + f"SELECT {select_cols}, invalidation_reason, invalidated_at " + f"FROM {fq_table('invalidated_memory_units')} WHERE id = $1 AND bank_id = $2", + unit_id, + bank_id, + ) + unit_state = "invalidated" + + if not row: + return None + + # Get entity information. _entity_rows_for_units_sql handles the + # observation→source_memory_ids inheritance fallback in SQL, so a + # single query covers direct rows and inherited ones. + entities_rows = await conn.fetch( + _entity_rows_for_units_sql(ops=ops, fq_table=fq_table, unit_ids_placeholder=1), + [row["id"]], + ) + entities = [r["canonical_name"] for r in entities_rows] + + result: dict[str, Any] = { + "id": str(row["id"]), + "text": row["text"], + "context": row["context"] if row["context"] else "", + "date": row["event_date"].isoformat() if row["event_date"] else "", + "type": row["fact_type"], + "mentioned_at": row["mentioned_at"].isoformat() if row["mentioned_at"] else None, + "occurred_start": row["occurred_start"].isoformat() if row["occurred_start"] else None, + "occurred_end": row["occurred_end"].isoformat() if row["occurred_end"] else None, + "entities": entities, + "document_id": row["document_id"] if row["document_id"] else None, + "chunk_id": str(row["chunk_id"]) if row["chunk_id"] else None, + "tags": row["tags"] if row["tags"] else [], + "metadata": conn.parse_json(row["metadata"]) if row["metadata"] is not None else {}, + "observation_scopes": ( + conn.parse_json(row["observation_scopes"]) if row["observation_scopes"] is not None else None + ), + "state": unit_state, + "invalidation_reason": row["invalidation_reason"], + "invalidated_at": row["invalidated_at"].isoformat() if row["invalidated_at"] else None, + "edited_at": row["edited_at"].isoformat() if row["edited_at"] else None, + } + + # For observations, include source_memory_ids + # history is deprecated here - use GET /memories/{id}/history instead + if row["fact_type"] == "observation": + result["history"] = [] + + if row["fact_type"] == "observation" and row["source_memory_ids"]: + source_ids = row["source_memory_ids"] + result["source_memory_ids"] = [str(sid) for sid in source_ids] + + # Fetch source memories + source_rows = await conn.fetch( + f""" + SELECT id, text, fact_type, context, occurred_start, mentioned_at + FROM {fq_table("memory_units")} + WHERE id = ANY($1::uuid[]) + ORDER BY mentioned_at DESC NULLS LAST + """, + source_ids, + ) + result["source_memories"] = [ + { + "id": str(r["id"]), + "text": r["text"], + "type": r["fact_type"], + "context": r["context"], + "occurred_start": r["occurred_start"].isoformat() if r["occurred_start"] else None, + "mentioned_at": r["mentioned_at"].isoformat() if r["mentioned_at"] else None, + } + for r in source_rows + ] + + return result + + +async def list_entities( + *, + conn, + fq_table, + bank_id: str, + search: str | None = None, + limit: int = 100, + offset: int = 0, +) -> dict[str, Any]: + """ + List all entities for a bank with pagination. + + Args: + conn: Open database connection (the caller owns the transaction). + fq_table: Table-name resolver. + bank_id: bank IDentifier + search: Optional case-insensitive substring match on canonical_name. + limit: Maximum number of entities to return + offset: Offset for pagination + + Returns: + Dict with items, total, limit, offset + """ + conditions = ["bank_id = $1"] + params: list[Any] = [bank_id] + if search: + # Substring match, same ILIKE shape entity lookup uses elsewhere. Applied + # to the count too, so the UI pages over the filtered set. + params.append(f"%{search}%") + conditions.append(f"canonical_name ILIKE ${len(params)}") + where_clause = " AND ".join(conditions) + + # Get total count + total_row = await conn.fetchrow( + f""" + SELECT COUNT(*) as total + FROM {fq_table("entities")} + WHERE {where_clause} + """, + *params, + ) + total = total_row["total"] if total_row else 0 + + # Get paginated entities + rows = await conn.fetch( + f""" + SELECT id, canonical_name, mention_count, first_seen, last_seen, metadata + FROM {fq_table("entities")} + WHERE {where_clause} + ORDER BY mention_count DESC, last_seen DESC, id ASC + LIMIT ${len(params) + 1} OFFSET ${len(params) + 2} + """, + *params, + limit, + offset, + ) + + entities = [] + for row in rows: + # Handle metadata - may be dict, JSON string, or None + metadata = row["metadata"] + if metadata is None: + metadata = {} + elif isinstance(metadata, str): + try: + metadata = json.loads(metadata) + except json.JSONDecodeError: + metadata = {} + + entities.append( + { + "id": str(row["id"]), + "canonical_name": row["canonical_name"], + "mention_count": row["mention_count"], + "first_seen": row["first_seen"].isoformat() if row["first_seen"] else None, + "last_seen": row["last_seen"].isoformat() if row["last_seen"] else None, + "metadata": metadata, + } + ) + return { + "items": entities, + "total": total, + "limit": limit, + "offset": offset, + } + + +__all__ = ["get_memory_unit", "list_entities", "list_memory_units"] diff --git a/hindsight-api-slim/hindsight_api/engine/memories/pg/graph.py b/hindsight-api-slim/hindsight_api/engine/memories/pg/graph.py new file mode 100644 index 0000000000..085218dad8 --- /dev/null +++ b/hindsight-api-slim/hindsight_api/engine/memories/pg/graph.py @@ -0,0 +1,780 @@ +"""Graph-shaped reads and the link-maintenance passes, in SQL. + +Everything here is a query over the *joins* around `memory_units` rather than +over the memories themselves: `unit_entities` (which entities a memory mentions) +and `memory_links` (memory-to-memory temporal/semantic/causal edges). + +Two groups of callers: + +* **The graph view.** :func:`graph_units`, :func:`graph_entity_rows` and + :func:`graph_direct_links` return raw rows; the engine still owns the + filtering, the observation inheritance, the derived entity edges, the + colouring and the response assembly. These functions answer only "which + memories", "which entity postings" and "which stored edges". +* **The graph-maintenance job.** :func:`enqueue_relink_victims` runs inside the + delete transaction; :func:`relink_pass`, :func:`prune_orphan_entities` and + :func:`prune_stale_cooccurrences` are the three reconciliation passes the job + drives. The job keeps the orchestration (pass ordering, the deadlock retry + around the sweeps, the timing log); each function here does the pass's work. + +:func:`entity_memory_counts` and :func:`entities_for_units` are the two entity +postings reads that are not part of the graph view but read the same join table. + +A store whose links travel inside the memory has nothing to relink and no join +table to sweep, which is why these are methods on the interface at all: it +answers them with zeroes rather than with SQL. +""" + +from __future__ import annotations + +import logging +import uuid as uuid_module +from collections.abc import Callable +from typing import Any + +from ....config import get_config +from ...db.base import DatabaseConnection +from ...retain.link_utils import ( + MAX_TEMPORAL_LINKS_PER_UNIT, + _bulk_insert_links, + _normalize_datetime, + compute_semantic_links_ann, +) + +logger = logging.getLogger(__name__) + +# Mirrors the ``top_k`` default in ``compute_semantic_links_ann`` at retain +# time. If you change one, change the other — otherwise victims would either +# never reach the cap (probe returns less than the cap) or stay perpetually +# under it (cap is higher than retain creates). +MAX_SEMANTIC_LINKS_PER_UNIT = 50 + +# Worker fetches this many rows per relink-loop iteration. Bounds +# per-iteration probe/insert latency so a 10k-row backlog doesn't hold a +# worker slot for minutes. Chosen so the typical iteration runs in well +# under 1s. +_DRAIN_BATCH_SIZE = 50 + +# Defensive guard against runaway relink loops — at _DRAIN_BATCH_SIZE units per +# iteration that's 500k targets, far beyond any realistic single-bank backlog. +_RELINK_ITERATION_CAP = 10000 + +# Cap at 10k edges — the UI can't usefully render more, and uncapped queries +# on highly-connected graphs (e.g. 1000 nodes with 500k+ edges) are too slow. +_GRAPH_MAX_EDGES = 10000 + +# Columns the graph view renders: nodes take id/text/date/context/entities, +# the table rows take the rest, and `source_memory_ids` is what lets the caller +# inherit an observation's links and entities from the facts behind it. +_GRAPH_UNIT_COLUMNS = ( + "id, text, event_date, context, occurred_start, occurred_end, mentioned_at, " + "document_id, chunk_id, fact_type, tags, created_at, proof_count, source_memory_ids" +) + + +# DataAccessOps is stateless and cached per dialect, so resolving it from the +# connection on each call is a dict lookup — cheap enough that functions taking a +# bare `conn` (no backend, no ops passed) can sniff the dialect here rather than +# threading ops through every signature. +def _ops_for(conn: DatabaseConnection) -> Any: + """The ``DataAccessOps`` matching the connection's SQL dialect. + + This is the SQL memories store, and SQL means Postgres *or* Oracle — the two + speak different dialects (Oracle inherits entity links through the + ``observation_sources`` junction, Postgres through ``source_memory_ids`` + arrays), so the ops have to follow the connection rather than assume Postgres. + ``create_data_access_ops`` caches per dialect, so this is a dict lookup after + the first call and returns the same instance the backend uses. + """ + from ...db import create_data_access_ops + + return create_data_access_ops(getattr(conn, "backend_type", "postgresql")) + + +def _as_uuids(unit_ids: list) -> list: + """Coerce a mixed list of uuid strings / UUIDs to UUIDs for a ``uuid[]`` bind.""" + return [uuid_module.UUID(uid) if isinstance(uid, str) else uid for uid in unit_ids] + + +# ---------------------------------------------------------------- graph view + + +def _observations_via_source_match( + fq_table: Callable[[str], str], + ops: Any, + source_column: str, + source_placeholder: int, + bank_placeholder: int | None, +) -> str: + """A predicate matching observations whose *sources* satisfy `` = $n``. + + Observations carry no `document_id` / `chunk_id` of their own; the link to a + source row lives in `source_memory_ids` (native array) or the + `observation_sources` junction, depending on the dialect. + """ + if ops.uses_observation_sources_table: + bank_clause = f" AND src.bank_id = ${bank_placeholder}" if bank_placeholder else "" + return ( + f"id IN (SELECT os.observation_id " + f"FROM {fq_table('observation_sources')} os " + f"JOIN {fq_table('memory_units')} src ON src.id = os.source_id " + f"WHERE src.{source_column} = ${source_placeholder}{bank_clause})" + ) + bank_clause = f" AND bank_id = ${bank_placeholder}" if bank_placeholder else "" + return ( + f"source_memory_ids && (SELECT array_agg(id) " + f"FROM {fq_table('memory_units')} " + f"WHERE {source_column} = ${source_placeholder}{bank_clause})" + ) + + +async def graph_units( + *, + conn: DatabaseConnection, + fq_table: Callable[[str], str], + bank_id: str | None = None, + fact_type: str | None = None, + search_query: str | None = None, + document_id: str | None = None, + chunk_id: str | None = None, + tags: list[str] | None = None, + tags_match: str = "all_strict", + limit: int = 1000, +) -> dict[str, Any]: + """Memory nodes for the graph view, plus the total matching count. + + Returns ``{"units": [...], "total": int}``: ``units`` is the page (newest + first, capped at ``limit``); ``total`` is how many match the filters, which + the UI shows alongside the page. ``document_id`` / ``chunk_id`` also match an + observation whose *sources* carry them, since observations have neither of + their own. + """ + from ...search.tags import build_tags_where_clause_simple + + ops = _ops_for(conn) + conditions: list[str] = [] + params: list[Any] = [] + + bank_placeholder: int | None = None + if bank_id: + params.append(bank_id) + bank_placeholder = len(params) + conditions.append(f"bank_id = ${bank_placeholder}") + + if fact_type: + params.append(fact_type) + conditions.append(f"fact_type = ${len(params)}") + + if document_id: + params.append(document_id) + obs = _observations_via_source_match(fq_table, ops, "document_id", len(params), bank_placeholder) + conditions.append(f"(document_id = ${len(params)} OR (fact_type = 'observation' AND {obs}))") + + if chunk_id: + params.append(chunk_id) + obs = _observations_via_source_match(fq_table, ops, "chunk_id", len(params), bank_placeholder) + conditions.append(f"(chunk_id = ${len(params)} OR (fact_type = 'observation' AND {obs}))") + + if search_query: + params.append(f"%{search_query}%") + conditions.append(f"(text ILIKE ${len(params)} OR context ILIKE ${len(params)})") + + if tags: + tag_clause = build_tags_where_clause_simple(tags, len(params) + 1, match=tags_match) + if tag_clause: + conditions.append(tag_clause.removeprefix("AND ")) + params.append(tags) + elif tags_match == "exact": + # Exact match with no tags is the "global" scope: rows carrying no tags at + # all. (Other modes treat empty tags as "no filter".) + conditions.append("(tags IS NULL OR tags = '{}')") + + where_clause = "WHERE " + " AND ".join(conditions) if conditions else "" + + total_row = await conn.fetchrow( + f"SELECT COUNT(*) AS total FROM {fq_table('memory_units')} {where_clause}", + *params, + ) + total = total_row["total"] if total_row else 0 + + params.append(limit) + rows = await conn.fetch( + f""" + SELECT {_GRAPH_UNIT_COLUMNS} + FROM {fq_table("memory_units")} + {where_clause} + ORDER BY mentioned_at DESC NULLS LAST, event_date DESC + LIMIT ${len(params)} + """, + *params, + ) + return {"units": [dict(row) for row in rows], "total": total} + + +async def graph_entity_rows( + *, + conn: DatabaseConnection, + fq_table: Callable[[str], str], + bank_id: str, + unit_ids: list[str], +) -> list[dict[str, Any]]: + """``(unit_id, entity_id, canonical_name)`` rows for the graph view's entity edges. + + Direct `unit_entities` postings only. An observation's entities are inherited + from its source memories by the caller, which is why the ids it passes here + are the visible units *plus* their source memories. + + Scoped by unit id rather than by bank: the ids already came from a + bank-scoped :func:`graph_units`, and `unit_entities` carries no bank column. + """ + if not unit_ids: + return [] + + rows = await conn.fetch( + f""" + SELECT ue.unit_id, e.id AS entity_id, e.canonical_name + FROM {fq_table("unit_entities")} ue + JOIN {fq_table("entities")} e ON ue.entity_id = e.id + WHERE ue.unit_id = ANY($1::uuid[]) + ORDER BY ue.unit_id + """, + _as_uuids(unit_ids), + ) + return [dict(row) for row in rows] + + +async def graph_direct_links( + *, + conn: DatabaseConnection, + fq_table: Callable[[str], str], + bank_id: str, + unit_ids: list[str], +) -> list[dict[str, Any]]: + """Memory-to-memory edges with *both* endpoints in ``unit_ids``. + + Entity edges are derived by the caller from `unit_entities` so we don't + materialize them in `memory_links` anymore (dropped in migration + e9b2c7d1f3a4) — no link_type filter is needed. ``entity_name`` is selected as + NULL so the row shape matches the derived edges the caller mixes these with. + + Pass the visible units *and* the source memories they inherit from: the + caller copies a source memory's links onto the observations built on it. + """ + if not unit_ids: + return [] + + rows = await conn.fetch( + f""" + SELECT ml.from_unit_id, + ml.to_unit_id, + ml.link_type, + ml.weight, + NULL::text AS entity_name + FROM {fq_table("memory_links")} ml + WHERE ml.from_unit_id = ANY($1::uuid[]) + AND ml.to_unit_id = ANY($1::uuid[]) + ORDER BY ml.weight DESC NULLS LAST + LIMIT $2 + """, + _as_uuids(unit_ids), + _GRAPH_MAX_EDGES, + ) + return [dict(row) for row in rows] + + +# ------------------------------------------------------------ entity postings + + +async def entity_memory_counts( + *, + conn: DatabaseConnection, + fq_table: Callable[[str], str], + bank_id: str, + entity_ids: list[str] | None = None, +) -> dict[str, int]: + """Live memory count per entity id, for the entities in ``bank_id``. + + The GROUP BY is what makes this an orphan test: an entity with no surviving + `unit_entities` row produces no group, so it is simply absent from the + result rather than present with a zero. + + Scoped through ``memory_units.bank_id`` — `unit_entities` has no bank column, + and joining is what keeps the count to *live* memories (deleted units take + their postings with them via ON DELETE CASCADE). + """ + params: list[Any] = [bank_id] + entity_filter = "" + if entity_ids is not None: + if not entity_ids: + return {} + params.append(_as_uuids(entity_ids)) + entity_filter = f"AND ue.entity_id = ANY(${len(params)}::uuid[])" + + rows = await conn.fetch( + f""" + SELECT ue.entity_id, COUNT(*) AS n + FROM {fq_table("unit_entities")} ue + JOIN {fq_table("memory_units")} mu ON mu.id = ue.unit_id + WHERE mu.bank_id = $1 + {entity_filter} + GROUP BY ue.entity_id + """, + *params, + ) + return {str(row["entity_id"]): int(row["n"]) for row in rows} + + +def _entity_rows_for_units_sql( + fq_table: Callable[[str], str], + ops: Any, + unit_ids_placeholder: int, +) -> str: + """SQL SELECT producing ``(unit_id, entity_id, canonical_name)`` rows for + the given unit IDs. + + Direct rows come from ``unit_entities``. Observations rarely carry + direct rows there; their entity association lives transitively through + their source memories (``source_memory_ids`` on PG, the + ``observation_sources`` junction on Oracle). When an observation has + no direct entity rows the SELECT inherits its source memories' + entities, so the result is the same set callers would get from + ``get_memory_unit``. + + ``unit_ids_placeholder`` is the 1-based parameter index that holds the + ``uuid[]`` of unit IDs. The placeholder is referenced twice — both + sides of the UNION need it — so callers should not reuse the slot. + """ + ue = fq_table("unit_entities") + ents = fq_table("entities") + mu = fq_table("memory_units") + p = unit_ids_placeholder + + direct = ( + f"SELECT ue.unit_id, e.id AS entity_id, e.canonical_name " + f"FROM {ue} ue " + f"JOIN {ents} e ON e.id = ue.entity_id " + f"WHERE ue.unit_id = ANY(${p}::uuid[])" + ) + + if ops.uses_observation_sources_table: + os_t = fq_table("observation_sources") + inherited = ( + f"SELECT os.observation_id AS unit_id, e.id AS entity_id, e.canonical_name " + f"FROM {os_t} os " + f"JOIN {ue} src_ue ON src_ue.unit_id = os.source_id " + f"JOIN {ents} e ON e.id = src_ue.entity_id " + f"WHERE os.observation_id = ANY(${p}::uuid[]) " + f"AND NOT EXISTS (SELECT 1 FROM {ue} d WHERE d.unit_id = os.observation_id)" + ) + else: + inherited = ( + f"SELECT obs.id AS unit_id, e.id AS entity_id, e.canonical_name " + f"FROM {mu} obs " + f"CROSS JOIN LATERAL unnest(obs.source_memory_ids) AS src_id " + f"JOIN {ue} src_ue ON src_ue.unit_id = src_id " + f"JOIN {ents} e ON e.id = src_ue.entity_id " + f"WHERE obs.id = ANY(${p}::uuid[]) " + f"AND obs.fact_type = 'observation' " + f"AND obs.source_memory_ids IS NOT NULL " + f"AND NOT EXISTS (SELECT 1 FROM {ue} d WHERE d.unit_id = obs.id)" + ) + + return f"({direct}) UNION ({inherited})" + + +async def entities_for_units( + *, + conn: DatabaseConnection, + fq_table: Callable[[str], str], + bank_id: str, + unit_ids: list[str], +) -> dict[str, list[str]]: + """The entity ids each unit carries, keyed by unit id. + + Observations inherit their source memories' entities when they carry no + direct postings of their own — see :func:`_entity_rows_for_units_sql`. Units + with no entities are absent rather than mapped to an empty list. + """ + if not unit_ids: + return {} + + rows = await conn.fetch( + _entity_rows_for_units_sql(fq_table, _ops_for(conn), unit_ids_placeholder=1), + _as_uuids(unit_ids), + ) + + # UNION already de-duplicates whole rows, but a unit can reach the same + # entity through more than one source memory, so dedupe per unit while + # preserving the order the rows arrived in. + by_unit: dict[str, list[str]] = {} + for row in rows: + unit_key = str(row["unit_id"]) + entity_id = str(row["entity_id"]) + bucket = by_unit.setdefault(unit_key, []) + if entity_id not in bucket: + bucket.append(entity_id) + return by_unit + + +async def entity_map_for_units( + *, + conn: DatabaseConnection, + fq_table: Callable[[str], str], + bank_id: str, + unit_ids: list[str], +) -> dict[str, list[dict[str, str]]]: + """``{unit_id: [{entity_id, canonical_name}]}`` — the recall/curation shape. + + The named twin of :func:`entities_for_units`: recall renders the entity name + on each fact, so it needs the label, not just the id. Observation-via-source + inheritance and the per-unit dedupe are identical. + """ + if not unit_ids: + return {} + + rows = await conn.fetch( + _entity_rows_for_units_sql(fq_table, _ops_for(conn), unit_ids_placeholder=1), + _as_uuids(unit_ids), + ) + by_unit: dict[str, list[dict[str, str]]] = {} + for row in rows: + unit_key = str(row["unit_id"]) + entity_id = str(row["entity_id"]) + bucket = by_unit.setdefault(unit_key, []) + if not any(existing["entity_id"] == entity_id for existing in bucket): + bucket.append({"entity_id": entity_id, "canonical_name": row["canonical_name"]}) + return by_unit + + +# --------------------------------------------------------------- maintenance + + +async def enqueue_relink_victims( + *, + conn: DatabaseConnection, + fq_table: Callable[[str], str], + bank_id: str, + affected_unit_ids: list, + include_affected_units: bool = False, +) -> int: + """Enqueue surviving units whose outgoing temporal/semantic links pointed at + ``affected_unit_ids`` for later link top-up. + + Must run inside the same transaction that drops those links, *before* the + delete (or cascade) fires — once the rows are gone, the join that finds the + victims returns nothing. + + Args: + conn: Database connection inside the active transaction. + fq_table: Schema-qualifying table-name resolver. + bank_id: Bank owning the affected units. + affected_unit_ids: Memory_unit IDs whose incident temporal/semantic links + are about to be (or are being) removed. + include_affected_units: Also enqueue ``affected_unit_ids`` themselves — for + an edit that deletes a unit's links but leaves the unit live, so its own + outgoing adjacency is rebuilt too. One combined insert keeps the queue's + sorted lock ordering intact. + + Returns: + Number of distinct victim units enqueued (after dedup against rows + already in the queue). + """ + if not affected_unit_ids: + return 0 + + ops = _ops_for(conn) + affected_uuids = _as_uuids(affected_unit_ids) + affected_str_set = {str(uid) for uid in affected_uuids} + + # Find units (other than the affected ones) that have an outgoing + # temporal/semantic link pointing at an affected unit. Entity links are + # intentionally excluded — they're scheduled for removal and would only + # add noise to the recompute job. + victim_rows = await conn.fetch( + f""" + SELECT DISTINCT from_unit_id + FROM {fq_table("memory_links")} + WHERE to_unit_id = ANY($1::uuid[]) + AND bank_id = $2 + AND link_type IN ('temporal', 'semantic') + """, + affected_uuids, + bank_id, + ) + + victim_ids = {row["from_unit_id"] for row in victim_rows if str(row["from_unit_id"]) not in affected_str_set} + if include_affected_units: + victim_ids.update(affected_uuids) + + if not victim_ids: + return 0 + + await ops.enqueue_graph_maintenance( + conn, + fq_table("graph_maintenance_queue"), + bank_id, + list(victim_ids), + ) + + logger.debug( + f"[GRAPH_MAINT] Enqueued {len(victim_ids)} relink victims in " + f"bank={bank_id} ({len(affected_unit_ids)} units affected)" + ) + return len(victim_ids) + + +async def relink_pass( + *, + backend: Any, + fq_table: Callable[[str], str], + bank_id: str, + config: Any, +) -> dict: + """Drain ``graph_maintenance_queue`` for ``bank_id``, topping up lost links. + + Per-iteration loop: claim → top up → commit. We rely on submit-time + dedup to keep at most one job per bank running, so no need for + SKIP LOCKED. + + Takes ``backend`` rather than a connection because the loop spans several + transactions — one per claimed batch, plus a separate connection for the ANN + probe — so it has to acquire its own. + + ``config`` is the caller's resolved configuration. The Postgres pass takes + its caps from retain's link_utils (so relink and retain agree on what "full" + means) and never reads it; it is accepted so a store that *does* tune its + relinking gets it. + + Returns: + ``{"relink_units_processed": int, "relink_links_added": int}``. + """ + del config # accepted for symmetry with stores that tune their own relinking + ops = backend.ops + + units_processed = 0 + links_added = 0 + iterations = 0 + while True: + from ...memory_engine import acquire_with_retry + + async with acquire_with_retry(backend) as conn: + async with conn.transaction(): + unit_ids = await ops.claim_graph_maintenance_batch( + conn, + fq_table("graph_maintenance_queue"), + bank_id, + _DRAIN_BATCH_SIZE, + ) + if not unit_ids: + break + + links_added += await _relink_batch(conn, fq_table, bank_id, unit_ids, ops, backend) + + units_processed += len(unit_ids) + iterations += 1 + + if iterations > _RELINK_ITERATION_CAP: + # Defensive guard against runaway loops — at 50 units/iter that's + # 500k targets, far beyond any realistic single-bank backlog. + logger.error( + f"[GRAPH_MAINT] bank={bank_id} hit iteration cap ({iterations}); aborting relink " + f"(units_processed={units_processed}, links_added={links_added})" + ) + break + + return {"relink_units_processed": units_processed, "relink_links_added": links_added} + + +async def _relink_batch( + conn: DatabaseConnection, + fq_table: Callable[[str], str], + bank_id: str, + victim_ids: list[str], + ops: Any, + backend: Any, +) -> int: + """Top up temporal/semantic links for a batch of victim units. Returns rows inserted.""" + # Load each victim's metadata. Victims whose units were deleted between + # enqueue and now silently drop out — exactly the no-op behaviour we want + # for stale queue rows. + victim_uuids = [uuid_module.UUID(vid) for vid in victim_ids] + victim_rows = await conn.fetch( + f""" + SELECT id::text AS id, event_date, fact_type, embedding::text AS embedding + FROM {fq_table("memory_units")} + WHERE id = ANY($1::uuid[]) + AND bank_id = $2 + AND fact_type IN ('experience', 'world') + """, + victim_uuids, + bank_id, + ) + + if not victim_rows: + return 0 + + alive_uuids = [uuid_module.UUID(row["id"]) for row in victim_rows] + + # Count current outgoing temporal/semantic links per victim so we only + # probe for the ones genuinely below cap. Saves the bulk of the work when + # most victims still have plenty of links. + count_rows = await conn.fetch( + f""" + SELECT from_unit_id, link_type, COUNT(*) AS cnt + FROM {fq_table("memory_links")} + WHERE from_unit_id = ANY($1::uuid[]) + AND bank_id = $2 + AND link_type IN ('temporal', 'semantic') + GROUP BY from_unit_id, link_type + """, + alive_uuids, + bank_id, + ) + counts: dict[tuple[str, str], int] = {} + for row in count_rows: + counts[(str(row["from_unit_id"]), row["link_type"])] = int(row["cnt"]) + + # --- Temporal top-up --- + temporal_needs = [r for r in victim_rows if counts.get((r["id"], "temporal"), 0) < MAX_TEMPORAL_LINKS_PER_UNIT] + new_links: list[tuple] = [] + + if temporal_needs: + lateral_unit_ids = [uuid_module.UUID(r["id"]) for r in temporal_needs if r["event_date"] is not None] + lateral_event_dates = [ + _normalize_datetime(r["event_date"]) for r in temporal_needs if r["event_date"] is not None + ] + lateral_fact_types = [r["fact_type"] for r in temporal_needs if r["event_date"] is not None] + + if lateral_unit_ids: + rows = await ops.fetch_temporal_neighbors( + conn, + fq_table("memory_units"), + bank_id, + lateral_unit_ids, + lateral_event_dates, + lateral_fact_types, + MAX_TEMPORAL_LINKS_PER_UNIT, + ) + for row in rows: + time_diff_h = float(row["time_diff_hours"]) + # Mirror the 24h window enforced at retain time. The bidirectional + # index scan returns the K closest neighbours regardless of + # window, so we filter here. + if time_diff_h > 24: + continue + weight = max(0.3, 1.0 - (time_diff_h / 24)) + new_links.append((row["from_id"], str(row["id"]), "temporal", weight, None)) + + # --- Semantic top-up --- + # ANN must run on its own connection: it opens a nested transaction with + # SET LOCAL hnsw.ef_search + CREATE TEMP TABLE ON COMMIT DROP, and nesting + # that inside our current write transaction would commit our writes early. + semantic_needs = [ + r + for r in victim_rows + if counts.get((r["id"], "semantic"), 0) < MAX_SEMANTIC_LINKS_PER_UNIT and r["embedding"] is not None + ] + if semantic_needs: + from ...memory_engine import acquire_with_retry + + seed_ids = [r["id"] for r in semantic_needs] + seed_embs = [r["embedding"] for r in semantic_needs] + seed_ftypes = [r["fact_type"] for r in semantic_needs] + async with acquire_with_retry(backend) as ann_conn: + try: + ann_links = await compute_semantic_links_ann( + ann_conn, + bank_id, + seed_ids, + seed_embs, + fact_types=seed_ftypes, + threshold=get_config().semantic_link_min_similarity, + ) + # Strip self-links (rare but possible because the ANN probe + # has no exclude list — see the comment in compute_semantic_links_ann). + ann_links = [lnk for lnk in ann_links if lnk[0] != lnk[1]] + new_links.extend(ann_links) + except Exception as e: + # ANN uses PG-specific HNSW syntax; on dialects/configs where + # it isn't available we still want the temporal top-up to land. + logger.warning(f"[GRAPH_MAINT] Semantic top-up failed for bank={bank_id}: {type(e).__name__}: {e}") + + if not new_links: + return 0 + + await _bulk_insert_links( + conn, + new_links, + bank_id=bank_id, + skip_exists_check=False, + ops=ops, + ) + return len(new_links) + + +async def prune_orphan_entities( + *, + conn: DatabaseConnection, + fq_table: Callable[[str], str], + bank_id: str, +) -> int: + """Delete ``entities`` rows in the bank with no remaining ``unit_entities`` + references. Returns the number pruned. + + FK ON DELETE CASCADE on ``entity_cooccurrences`` then removes any + cooccurrence row pointing at the pruned entities — which is why this runs + before :func:`prune_stale_cooccurrences` rather than after. + + A bank-wide single-statement delete, cheap when there's nothing to do. It is + idempotent (rerunning only deletes what is still orphaned), so the caller is + free to retry the whole transaction on deadlock. + """ + ops = _ops_for(conn) + return await ops.prune_orphan_entities( + conn, + fq_table("entities"), + fq_table("unit_entities"), + bank_id, + ) + + +async def prune_stale_cooccurrences( + *, + conn: DatabaseConnection, + fq_table: Callable[[str], str], + bank_id: str, +) -> int: + """Delete cooccurrence rows no current memory witnesses. Returns the count. + + Defensive sweep for rows where both endpoints still exist but no current + memory_unit references both of them — the cooccurrence was real at the time + it was recorded, but every unit that witnessed it has since been deleted. + :func:`prune_orphan_entities` cascades the *missing-entity* case via FK; this + pass catches the *stale-count* case it cannot see. + + Like the orphan prune, a bank-wide idempotent sweep backed by indexes, so + it's cheap when there's nothing to do and safe for the caller to retry. + """ + ops = _ops_for(conn) + return await ops.prune_stale_cooccurrences( + conn, + fq_table("entity_cooccurrences"), + fq_table("unit_entities"), + fq_table("entities"), + bank_id, + ) + + +__all__ = [ + "MAX_SEMANTIC_LINKS_PER_UNIT", + "enqueue_relink_victims", + "entities_for_units", + "entity_map_for_units", + "entity_memory_counts", + "graph_direct_links", + "graph_entity_rows", + "graph_units", + "prune_orphan_entities", + "prune_stale_cooccurrences", + "relink_pass", +] diff --git a/hindsight-api-slim/hindsight_api/engine/memories/pg/reads.py b/hindsight-api-slim/hindsight_api/engine/memories/pg/reads.py new file mode 100644 index 0000000000..caadf37df6 --- /dev/null +++ b/hindsight-api-slim/hindsight_api/engine/memories/pg/reads.py @@ -0,0 +1,451 @@ +"""Addressed reads over `memory_units`: get, scan, count, tags, consolidation state. + +Not retrieval — nothing here ranks. These are the queries behind the curation +detail view, export, the bank-stats panel and the consolidation queue, lifted out +of the call sites that used to issue them inline (``memory_engine``, +``transfer/export``, ``consolidation/consolidator``) so +:class:`~hindsight_api.engine.memories.postgres.PostgresMemories` can delegate +rather than embed SQL. + +Every function takes the live connection and Hindsight's ``fq_table`` resolver, so +each one runs inside whatever transaction the caller already holds; none of them +acquires a connection of its own. + +**Cursor semantics.** ``scan_memories``'s ``page_token`` is opaque to callers, and +for Postgres it is simply a *numeric offset rendered as a decimal string* against +the scan's fixed ``ORDER BY created_at, id``. An empty token means "start at the +beginning", and an empty token comes back once the walk is exhausted (i.e. the +final short page). An offset cursor is a position rather than a snapshot — exactly +the guarantee :class:`~hindsight_api.engine.memories.base.ScanPage` documents: +rows written or deleted mid-walk can shift later pages, so a scan is +eventually-complete browsing rather than a consistent iterator. ``skip`` is applied +*on top of* the decoded cursor, so a caller that pages with both should pass +``skip`` only on the first call — the returned token already accounts for it. +""" + +from __future__ import annotations + +import json +import uuid +from collections.abc import Callable +from datetime import datetime +from typing import Any + +from ...search.tags import ( + build_tag_groups_where_clause, + build_tags_where_clause, + build_tags_where_clause_simple, +) +from ..base import ScanPage, StoredMemory + +# The `memory_units` projection every read here shares. Superset of the by-id +# SELECT the recall source-facts path used (text/fact_type/context/timestamps/ +# document_id/chunk_id/tags/metadata), plus the observation bookkeeping columns +# `StoredMemory` carries: source_memory_ids and consolidated_at. +_MEMORY_COLUMNS = """ + id, text, fact_type, context, document_id, chunk_id, tags, metadata, + proof_count, event_date, occurred_start, occurred_end, mentioned_at, + created_at, source_memory_ids, consolidated_at, observation_scopes +""" + +# The scan's order. Fixed (created_at, id) like the export loader's, because an +# offset cursor is only meaningful against a total order. +_SCAN_ORDER = "ORDER BY created_at, id" + + +def _as_json(value: Any) -> Any: + """Coerce an asyncpg JSONB column (str or already-decoded) to a Python object. + + Connections differ in whether a JSONB codec is registered, so the column + arrives either as text or as the decoded object. + """ + if value is None: + return None + if isinstance(value, str): + try: + return json.loads(value) + except json.JSONDecodeError: + # A valid scalar such as `"combined"` arrives already decoded on + # connections that do register a decoder. + return value + return value + + +def _as_uuids(unit_ids: list[Any]) -> list[uuid.UUID]: + """Unit ids as UUIDs, dropping anything unparseable. + + A malformed id is treated the same way a deleted one is — simply absent from + the result — rather than failing the whole read. + """ + out: list[uuid.UUID] = [] + for unit_id in unit_ids or []: + if isinstance(unit_id, uuid.UUID): + out.append(unit_id) + continue + try: + out.append(uuid.UUID(str(unit_id))) + except (ValueError, AttributeError, TypeError): + continue + return out + + +def _column(row: Any, name: str, default: Any = None) -> Any: + """One column of an asyncpg Record, tolerating a narrower projection.""" + try: + return row[name] + except (KeyError, IndexError): + return default + + +def _stored_from_row(row: Any) -> StoredMemory: + """Map a `memory_units` row onto :class:`StoredMemory`. + + Shared by every read in this module so the row → dataclass mapping exists + once. ``entity_ids`` stays empty: the unit→entity posting lives in + `unit_entities` and is served by ``entities_for_units``, not by a join here. + """ + source_ids = _column(row, "source_memory_ids") or [] + return StoredMemory( + unit_id=str(row["id"]), + text=row["text"], + fact_type=row["fact_type"], + context=_column(row, "context"), + document_id=_column(row, "document_id"), + chunk_id=str(_column(row, "chunk_id")) if _column(row, "chunk_id") else None, + tags=list(_column(row, "tags") or []), + metadata=_as_json(_column(row, "metadata")), + proof_count=_column(row, "proof_count") or 1, + event_date=_column(row, "event_date"), + occurred_start=_column(row, "occurred_start"), + occurred_end=_column(row, "occurred_end"), + mentioned_at=_column(row, "mentioned_at"), + created_at=_column(row, "created_at"), + source_memory_ids=[str(sid) for sid in source_ids], + consolidated_at=_column(row, "consolidated_at"), + # Consolidation routes a candidate by its scopes, so this has to survive + # the trip through the store rather than being re-queried per memory. + observation_scopes=_as_json(_column(row, "observation_scopes")), + ) + + +def _decode_page_token(page_token: str) -> int: + """Decode the offset cursor. Empty, malformed or negative all mean "start".""" + if not page_token: + return 0 + try: + offset = int(page_token) + except (TypeError, ValueError): + return 0 + return offset if offset > 0 else 0 + + +async def get_memories( + *, conn, fq_table: Callable[[str], str], bank_id: str, unit_ids: list[str] +) -> list[StoredMemory]: + """Fetch memories by id. Missing or deleted ids are simply absent.""" + ids = _as_uuids(unit_ids) + if not ids: + return [] + rows = await conn.fetch( + f""" + SELECT {_MEMORY_COLUMNS} + FROM {fq_table("memory_units")} + WHERE bank_id = $1 AND id = ANY($2::uuid[]) + """, + bank_id, + ids, + ) + return [_stored_from_row(row) for row in rows] + + +async def _semantic_edges( + *, conn, fq_table: Callable[[str], str], bank_id: str, unit_ids: list[uuid.UUID] +) -> dict[str, list[tuple[str, float]]]: + """Derived kNN edges for ``unit_ids``, keyed by unit id. + + Walked in both directions, like the graph arm's semantic expansion: a + `memory_links` row is written once, so a unit's neighbourhood is the union of + the edges leaving it and those arriving at it. + """ + if not unit_ids: + return {} + rows = await conn.fetch( + f""" + SELECT from_unit_id AS unit_id, to_unit_id AS target_id, weight + FROM {fq_table("memory_links")} + WHERE bank_id = $1 AND link_type = 'semantic' AND from_unit_id = ANY($2::uuid[]) + UNION ALL + SELECT to_unit_id AS unit_id, from_unit_id AS target_id, weight + FROM {fq_table("memory_links")} + WHERE bank_id = $1 AND link_type = 'semantic' AND to_unit_id = ANY($2::uuid[]) + """, + bank_id, + unit_ids, + ) + edges: dict[str, list[tuple[str, float]]] = {} + for row in rows: + edges.setdefault(str(row["unit_id"]), []).append((str(row["target_id"]), float(row["weight"] or 0.0))) + return edges + + +async def scan_memories( + *, + conn, + fq_table: Callable[[str], str], + bank_id: str, + fact_types: list[str] | None = None, + limit: int = 100, + page_token: str = "", + tags: list[str] | None = None, + tags_match: str = "any", + tag_groups: list | None = None, + document_id: str | None = None, + metadata_equals: dict[str, str] | None = None, + skip: int = 0, + include_edges: bool = False, +) -> ScanPage: + """Page through stored memories. A full walk — for browsing and export only. + + See the module docstring for the ``page_token`` (offset) cursor semantics. + """ + if limit is None or limit <= 0: + return ScanPage() + + where: list[str] = ["bank_id = $1"] + params: list[Any] = [bank_id] + + if fact_types: + params.append(list(fact_types)) + where.append(f"fact_type = ANY(${len(params)})") + + if document_id is not None: + # A real column here, which is why it is not folded into + # `metadata_equals`: only a store without the column keeps it in the bag. + params.append(document_id) + where.append(f"document_id = ${len(params)}") + + if metadata_equals: + # str→str equality across every key, which is exactly JSONB containment. + params.append(json.dumps(metadata_equals)) + where.append(f"metadata @> ${len(params)}::jsonb") + + # The tags clause owns its own `AND` prefix and, per the helper's contract, + # only consumes a bind param when `tags` is non-empty (match="exact" with no + # tags is the untagged/global scope and needs none). + tags_clause = build_tags_where_clause_simple(tags, len(params) + 1, match=tags_match) + if tags: + params.append(list(tags)) + + # Compound tag groups (AND/OR/NOT trees), AND-ed on. Also owns its `AND` prefix and appends + # one bind param per leaf; empty/absent groups yield no clause and no params. + groups_clause, group_params, _ = build_tag_groups_where_clause(tag_groups, param_offset=len(params) + 1) + params.extend(group_params) + + offset = _decode_page_token(page_token) + max(int(skip or 0), 0) + params.append(limit) + limit_idx = len(params) + params.append(offset) + offset_idx = len(params) + + rows = await conn.fetch( + f""" + SELECT {_MEMORY_COLUMNS} + FROM {fq_table("memory_units")} + WHERE {" AND ".join(where)} {tags_clause} {groups_clause} + {_SCAN_ORDER} + LIMIT ${limit_idx} OFFSET ${offset_idx} + """, + *params, + ) + + memories = [_stored_from_row(row) for row in rows] + if include_edges and memories: + edges = await _semantic_edges( + conn=conn, fq_table=fq_table, bank_id=bank_id, unit_ids=_as_uuids([m.unit_id for m in memories]) + ) + for memory in memories: + memory.semantic_edges = edges.get(memory.unit_id, []) + + # A short page means the walk is exhausted, so the cursor goes empty. + next_token = str(offset + len(rows)) if len(rows) == limit else "" + return ScanPage(memories=memories, next_page_token=next_token) + + +async def count_memories(*, conn, fq_table: Callable[[str], str], bank_id: str) -> dict[str, int]: + """Live memory count per fact_type. The bank-stats node counts.""" + rows = await conn.fetch( + f""" + SELECT fact_type, COUNT(*) as count + FROM {fq_table("memory_units")} + WHERE bank_id = $1 + GROUP BY fact_type + """, + bank_id, + ) + return {row["fact_type"]: int(row["count"]) for row in rows} + + +async def list_tags(*, conn, fq_table: Callable[[str], str], bank_id: str) -> dict[str, int]: + """Distinct tags in a bank and how many live memories carry each. + + The engine's generic tag histogram builds its fragments from + ``ops.build_tag_listing_parts`` so it can serve any table on any dialect. This + module's signature carries no ``ops``, and reaching one off the connection + would be a back door into the dialect layer for a query that is already + Postgres-specific by virtue of living under ``pg/`` — so the Postgres + fragments (``unnest(tags)`` + the non-empty guard) are inlined verbatim here. + Unpaged on purpose: the interface returns the whole histogram. + """ + rows = await conn.fetch( + f""" + SELECT tag, COUNT(*) as count + FROM {fq_table("memory_units")}, unnest(tags) AS tag + WHERE bank_id = $1 AND tags IS NOT NULL AND tags != '{{}}' + GROUP BY tag + ORDER BY count DESC, tag ASC + """, + bank_id, + ) + return {row["tag"]: int(row["count"]) for row in rows} + + +async def find_unconsolidated( + *, + conn, + fq_table: Callable[[str], str], + bank_id: str, + fact_types: list[str], + limit: int, + scope_tags: list[str] | None = None, +) -> list[StoredMemory]: + """Memories not yet folded into an observation, oldest first. + + The consolidator's candidate query: never consolidated, never *failed* to + consolidate (a memory the LLM could not handle must not be retried forever), + ordered by ``created_at`` so the queue drains in arrival order. ``scope_tags`` + is the same ``tags @> scope`` containment the job's scope filter uses — the + job ORs several scopes together; one scope is passed here. + """ + where = [ + "bank_id = $1", + "consolidated_at IS NULL", + "consolidation_failed_at IS NULL", + ] + params: list[Any] = [bank_id] + if fact_types: + params.append(list(fact_types)) + where.append(f"fact_type = ANY(${len(params)})") + if scope_tags: + params.append(list(scope_tags)) + where.append(f"tags @> ${len(params)}::varchar[]") + params.append(limit) + + rows = await conn.fetch( + f""" + SELECT {_MEMORY_COLUMNS} + FROM {fq_table("memory_units")} + WHERE {" AND ".join(where)} + ORDER BY created_at ASC + LIMIT ${len(params)} + """, + *params, + ) + return [_stored_from_row(row) for row in rows] + + +async def mark_consolidated( + *, + conn, + fq_table: Callable[[str], str], + bank_id: str, + unit_ids: list[str], + when: datetime | None, + failed: bool = False, +) -> None: + """Stamp (or clear, with ``when=None``) the consolidated marker on sources. + + ``failed`` writes ``consolidation_failed_at`` instead of ``consolidated_at``, + which is what keeps a memory the LLM could not consolidate out of the queue. + + ``when=None`` clears the column rather than stamping it — that is how a source + is requeued once the observation built on it is deleted. The clear keeps the + ``fact_type IN ('experience', 'world')`` guard the requeue sites carry: + observations are never themselves consolidated, so nothing about them should + be reset by a requeue. + + ``updated_at`` is deliberately left alone, matching the consolidator's own + statements: consolidation bookkeeping is not an edit to the memory, and + bumping it would make every consolidation pass look like a write to the + staleness check below. + """ + ids = _as_uuids(unit_ids) + if not ids: + return + column = "consolidation_failed_at" if failed else "consolidated_at" + guard = "" if when is not None else " AND fact_type IN ('experience', 'world')" + await conn.execute( + f""" + UPDATE {fq_table("memory_units")} + SET {column} = $1 + WHERE bank_id = $2 AND id = ANY($3::uuid[]){guard} + """, + when, + bank_id, + ids, + ) + + +async def any_memory_updated_since( + *, + conn, + fq_table: Callable[[str], str], + bank_id: str, + since: datetime, + fact_types: list[str] | None = None, + tags: list[str] | None = None, + tags_match: str = "any", + tag_groups: list | None = None, +) -> bool: + """Whether any memory in ``bank_id``'s scope was written after ``since``. + + Backs the mental-model staleness check, so it is a bounded existence test — + ``LIMIT 1``, never a COUNT: the answer is "is there one", and the planner can + stop at the first hit. The scope is the mental model's: its flat tags (or the + compound ``tag_groups``) plus an optional ``fact_types`` restriction. This is + where the staleness query's WHERE lives, so the same scope that gates a + refresh decides whether one is due. + """ + params: list[Any] = [bank_id, since] + where = ["bank_id = $1", "updated_at > $2"] + + tag_clause, tag_params, next_param = build_tags_where_clause(tags, param_offset=len(params) + 1, match=tags_match) + if tag_clause: + where.append(tag_clause.removeprefix("AND ")) + params.extend(tag_params) + + group_clause, group_params, _ = build_tag_groups_where_clause(tag_groups, param_offset=next_param) + if group_clause: + where.append(group_clause.removeprefix("AND ")) + params.extend(group_params) + # Untagged, no tag_groups → no tag constraint, matching any memory in the bank. + + if fact_types: + params.append(list(fact_types)) + where.append(f"fact_type = ANY(${len(params)}::text[])") + + row = await conn.fetchval( + f"SELECT 1 FROM {fq_table('memory_units')} WHERE {' AND '.join(where)} LIMIT 1", + *params, + ) + return row is not None + + +__all__ = [ + "any_memory_updated_since", + "count_memories", + "find_unconsolidated", + "get_memories", + "list_tags", + "mark_consolidated", + "scan_memories", +] diff --git a/hindsight-api-slim/hindsight_api/engine/memories/pg/writes.py b/hindsight-api-slim/hindsight_api/engine/memories/pg/writes.py new file mode 100644 index 0000000000..d42def0dda --- /dev/null +++ b/hindsight-api-slim/hindsight_api/engine/memories/pg/writes.py @@ -0,0 +1,575 @@ +"""Writes against `memory_units`: the fact insert, the deletes, and observation invalidation. + +Everything here mutates the memories slice and nothing else. The document row, +the chunks, the entity registry and the link tables stay with their own callers — +what lands in this module is only the statements that touch `memory_units` (and, +on backends that keep one, the `observation_sources` junction that hangs off it). + +Each function takes the live connection and Hindsight's ``fq_table`` resolver, so +it runs inside whatever transaction the caller already holds; ``ops`` is the +dialect ops object, which is what lets the same code serve the PG (native array) +and Oracle (junction table) shapes of the observation→source relation. +""" + +from __future__ import annotations + +import json +import logging +import uuid +from collections.abc import Callable +from typing import TYPE_CHECKING + +from ....config import get_config +from ..base import StoredMemory + +if TYPE_CHECKING: # pragma: no cover - typing only + from ...retain.types import ProcessedFact + +logger = logging.getLogger(__name__) + + +async def insert_facts( + *, + conn, + ops, + bank_id: str, + facts: list[ProcessedFact], + document_id: str | None = None, +) -> list[str]: + """Insert facts into the database in batch. + + Args: + conn: Database connection + bank_id: Bank identifier + facts: List of ProcessedFact objects to insert + document_id: Optional document ID to associate with facts + + Returns: + List of unit IDs (UUIDs as strings) for the inserted facts, in the same + order as ``facts``. + """ + if not facts: + return [] + + # Imported here: `retain` reaches back into the engine for `fq_table`, so a + # module-level import would close the cycle once the engine imports this store. + from ...retain.fact_extraction import _sanitize_text + + # Prepare data for batch insert + fact_texts = [] + embeddings = [] + event_dates = [] + occurred_starts = [] + occurred_ends = [] + mentioned_ats = [] + contexts = [] + fact_types = [] + metadata_jsons = [] + chunk_ids = [] + document_ids = [] + tags_list = [] + observation_scopes_list = [] + text_signals_list = [] + + for fact in facts: + fact_texts.append(_sanitize_text(fact.fact_text)) + # Convert embedding to string for asyncpg vector type + embeddings.append(str(fact.embedding)) + # event_date: Use occurred_start if available, otherwise use mentioned_at + # This maintains backward compatibility while handling None occurred_start + event_dates.append(fact.occurred_start if fact.occurred_start is not None else fact.mentioned_at) + occurred_starts.append(fact.occurred_start) + occurred_ends.append(fact.occurred_end) + mentioned_ats.append(fact.mentioned_at) + contexts.append(_sanitize_text(fact.context)) + fact_types.append(fact.fact_type) + metadata_jsons.append(json.dumps(fact.metadata)) + chunk_ids.append(fact.chunk_id) + # Use per-fact document_id if available, otherwise fallback to batch-level document_id + document_ids.append(fact.document_id if fact.document_id else document_id) + # Convert tags to JSON string for proper batch insertion (PostgreSQL unnest doesn't handle 2D arrays well) + tags_list.append(json.dumps(fact.tags if fact.tags else [])) + # observation_scopes: stored as JSONB (string or 2D array), None if not provided + observation_scopes_list.append( + json.dumps(fact.observation_scopes) if fact.observation_scopes is not None else None + ) + # Build text_signals: entity names + date tokens for enriched BM25 indexing + signal_parts = [] + if fact.entities: + signal_parts.extend(e.name for e in fact.entities) + if fact.occurred_start: + try: + signal_parts.append(fact.occurred_start.strftime("%B %d %Y").lstrip("0").replace(" 0", " ")) + except (ValueError, AttributeError): + pass + if fact.occurred_end and fact.occurred_end != fact.occurred_start: + try: + signal_parts.append(fact.occurred_end.strftime("%B %d %Y").lstrip("0").replace(" 0", " ")) + except (ValueError, AttributeError): + pass + text_signals_list.append(" ".join(signal_parts) if signal_parts else None) + + # Batch insert all facts — delegates to DataAccessOps which handles + # unnest (PG) vs row-by-row (Oracle) transparently. + config = get_config() + + return await ops.insert_facts_batch( + conn, + bank_id, + fact_texts, + embeddings, + event_dates, + occurred_starts, + occurred_ends, + mentioned_ats, + contexts, + fact_types, + metadata_jsons, + chunk_ids, + document_ids, + tags_list, + observation_scopes_list, + text_signals_list, + text_search_extension=config.text_search_extension, + ) + + +async def delete_document(*, conn, fq_table: Callable[[str], str], bank_id: str, document_id: str) -> None: + """Delete every memory unit belonging to ``document_id``. + + Explicitly delete memory_units by document_id BEFORE deleting the + document row. The CASCADE from documents→chunks→memory_units only + catches units that have a non-NULL chunk_id FK. Units with chunk_id=NULL + (e.g. from partial writes or edge cases) would survive the cascade. + This explicit delete ensures complete cleanup. + + Called when a document is replaced, so it races the replacement's writes: it + must remove only what was written *before* this call, never the facts + arriving moments later — which the ``document_id``/``bank_id`` predicate + gives for free inside the caller's transaction. + """ + await conn.execute( + f"DELETE FROM {fq_table('memory_units')} WHERE document_id = $1 AND bank_id = $2", + document_id, + bank_id, + ) + + +async def delete_observations(*, conn, fq_table: Callable[[str], str], bank_id: str) -> None: + """Delete all observations in a bank, leaving the facts behind them. + + Only the observation rows: requeuing the surviving sources (clearing + ``consolidated_at``) and resetting the bank's consolidation timestamp belong + to the caller, which owns the bank row. + """ + await conn.execute( + f"DELETE FROM {fq_table('memory_units')} WHERE bank_id = $1 AND fact_type = 'observation'", + bank_id, + ) + + +async def observations_for_sources( + *, + conn, + ops, + fq_table: Callable[[str], str], + bank_id: str, + unit_ids: list[str | uuid.UUID], +) -> list[StoredMemory]: + """Observations consolidated from any of ``unit_ids``. + + Only ``unit_id`` and ``source_memory_ids`` are populated — the caller uses + them to delete the observations and to work out which sources survive, and + the rest of the row is about to be deleted anyway. + """ + if not unit_ids: + return [] + + fact_uuids = [uuid.UUID(str(fid)) if not isinstance(fid, uuid.UUID) else fid for fid in unit_ids] + + if ops is not None and not ops.uses_observation_sources_table: + # PG: use native array overlap operator + rows = await conn.fetch( + f""" + SELECT id, source_memory_ids + FROM {fq_table("memory_units")} + WHERE bank_id = $1 + AND fact_type = 'observation' + AND source_memory_ids && $2::uuid[] + """, + bank_id, + fact_uuids, + ) + else: + # Oracle / default: use observation_sources junction table + rows = await conn.fetch( + f""" + SELECT mu.id, mu.source_memory_ids + FROM {fq_table("memory_units")} mu + WHERE mu.bank_id = $1 + AND mu.fact_type = 'observation' + AND EXISTS ( + SELECT 1 FROM {fq_table("observation_sources")} os + WHERE os.observation_id = mu.id + AND os.source_id = ANY($2::uuid[]) + ) + """, + bank_id, + fact_uuids, + ) + + return [ + StoredMemory( + unit_id=str(row["id"]), + text="", + fact_type="observation", + source_memory_ids=[str(src_id) for src_id in (row["source_memory_ids"] or [])], + ) + for row in rows + ] + + +async def delete_stale_observations( + *, + conn, + ops, + fq_table: Callable[[str], str], + bank_id: str, + fact_ids: list[str | uuid.UUID], +) -> int: + """Delete observations whose source memories are about to be removed. + + Mirrors the cleanup performed by ``MemoryEngine.delete_document`` so that + every code path that removes ``memory_units`` also removes the + observations derived from them. Without this, ingesting a fresh version + of a document via the retain pipeline (which does a full-replace + ``DELETE FROM documents`` cascade) used to leave orphan observations + pointing at memory IDs that no longer existed. + + For each observation referencing any of ``fact_ids``: + 1. Delete the observation row (its text is stale once even one source + memory disappears). + 2. Reset ``consolidated_at = NULL`` on the surviving source memories so + they get re-consolidated under fresh observations on the next run. + + Must be called within an active transaction, before the source memories + are deleted. + + Returns the number of observations deleted. + """ + if not fact_ids: + return 0 + + fact_uuids = [uuid.UUID(str(fid)) if not isinstance(fid, uuid.UUID) else fid for fid in fact_ids] + + affected_obs = await observations_for_sources( + conn=conn, ops=ops, fq_table=fq_table, bank_id=bank_id, unit_ids=fact_uuids + ) + if not affected_obs: + return 0 + + deleted_set = {str(uid) for uid in fact_uuids} + obs_ids = [uuid.UUID(obs.unit_id) for obs in affected_obs] + seen_remaining: set[str] = set() + remaining_source_ids: list[uuid.UUID] = [] + for obs in affected_obs: + for src_str in obs.source_memory_ids: + if src_str not in deleted_set and src_str not in seen_remaining: + remaining_source_ids.append(uuid.UUID(src_str)) + seen_remaining.add(src_str) + + await conn.execute( + f"DELETE FROM {fq_table('memory_units')} WHERE id = ANY($1::uuid[])", + obs_ids, + ) + + if remaining_source_ids: + await conn.execute( + f""" + UPDATE {fq_table("memory_units")} + SET consolidated_at = NULL + WHERE id = ANY($1::uuid[]) + AND fact_type IN ('experience', 'world') + """, + remaining_source_ids, + ) + + logger.info( + f"[OBSERVATIONS] Deleted {len(obs_ids)} observations, reset {len(remaining_source_ids)} " + f"source memories for re-consolidation in bank {bank_id}" + ) + return len(obs_ids) + + +# --------------------------------------------------------------------- curation archive +# +# Invalidation moves a rejected memory between two tables rather than flagging it, +# so recall / consolidation / graph never carry a "valid?" predicate: live facts +# live in `memory_units`, invalidated ones in `invalidated_memory_units`. The +# archive is cold storage — no index, so it drops the `embedding` and +# `search_vector` columns, which are recomputed on the way back. + +# The two recall-surface columns the archive omits. Both follow server config +# (embedding dimension, search backend), so keeping them out of the INSERT…SELECT +# round-trip makes a model or text-backend switch structurally unable to trip a +# type/dimension mismatch (#2209, #2503); each is recomputed on revert. +_ARCHIVE_OMITTED = ('"embedding"', '"search_vector"') + + +async def _memory_unit_columns(conn, fq_table: Callable[[str], str]) -> str: + """The quoted, ordinal column list of `memory_units`. + + Read from the catalog rather than hardcoded so a schema migration cannot make + the archive round-trip drift from the live table (the archive is created via + ``LIKE memory_units``, so the lists line up). + """ + rows = await conn.fetch( + f"SELECT a.attname FROM pg_attribute a " + f"WHERE a.attrelid = '{fq_table('memory_units')}'::regclass " + f"AND a.attnum > 0 AND NOT a.attisdropped ORDER BY a.attnum" + ) + return ", ".join(f'"{r["attname"]}"' for r in rows) + + +async def _archive_columns(conn, fq_table: Callable[[str], str]) -> str: + """`_memory_unit_columns` minus the two the archive does not carry.""" + collist = await _memory_unit_columns(conn, fq_table) + return ", ".join(c for c in (s.strip() for s in collist.split(",")) if c not in _ARCHIVE_OMITTED) + + +_ARCHIVE_SELECT = ( + "id, text, fact_type, context, occurred_start, occurred_end, mentioned_at, " + "document_id, chunk_id, tags, metadata, proof_count, event_date, created_at, " + "consolidated_at, entity_ids" +) + + +def _archived_stored(row: Any) -> StoredMemory: + """Map an `invalidated_memory_units` row onto :class:`StoredMemory`.""" + return StoredMemory( + unit_id=str(row["id"]), + text=row["text"], + fact_type=row["fact_type"], + context=row["context"], + document_id=row["document_id"], + chunk_id=str(row["chunk_id"]) if row["chunk_id"] else None, + tags=list(row["tags"] or []), + metadata=row["metadata"] if isinstance(row["metadata"], dict) else None, + proof_count=row["proof_count"] or 1, + event_date=row["event_date"], + occurred_start=row["occurred_start"], + occurred_end=row["occurred_end"], + mentioned_at=row["mentioned_at"], + created_at=row["created_at"], + consolidated_at=row["consolidated_at"], + entity_ids=[str(e) for e in (row["entity_ids"] or [])], + ) + + +async def get_archived_memory(*, conn, fq_table, bank_id: str, unit_id: str) -> StoredMemory | None: + row = await conn.fetchrow( + f"SELECT {_ARCHIVE_SELECT} FROM {fq_table('invalidated_memory_units')} WHERE id = $1 AND bank_id = $2", + str(unit_id), + bank_id, + ) + return _archived_stored(row) if row else None + + +async def invalidate_memory(*, conn, fq_table, bank_id: str, unit_id: str, reason: str | None) -> bool: + mu = fq_table("memory_units") + arch = fq_table("invalidated_memory_units") + ue = fq_table("unit_entities") + arch_cols = await _archive_columns(conn, fq_table) + + # Snapshot the entity ids before the delete cascade takes `unit_entities`, so + # revert can restore the postings the move is about to drop. + entity_ids = [ + r["entity_id"] for r in await conn.fetch(f"SELECT entity_id FROM {ue} WHERE unit_id = $1", str(unit_id)) + ] + # Causal edges are retain-time extraction output the FK cascade would destroy for good — + # unlike temporal/semantic links they can't be recomputed, so snapshot their descriptors onto + # the archive row and revert rematerializes them (#2864). + from ...retain.link_utils import snapshot_causal_links + + causal_links = await snapshot_causal_links(conn, bank_id, str(unit_id)) + inserted = await conn.fetchval( + f"INSERT INTO {arch} ({arch_cols}, invalidation_reason, invalidated_at, entity_ids, causal_links) " + f"SELECT {arch_cols}, $2, now(), $3::uuid[], $5::jsonb FROM {mu} WHERE id = $1 AND bank_id = $4 " + f"RETURNING id", + str(unit_id), + reason, + entity_ids, + bank_id, + json.dumps([descriptor.as_json_dict() for descriptor in causal_links]), + ) + if inserted is None: + return False + # The cascade prunes `unit_entities` and `memory_links` with the row. + await conn.execute(f"DELETE FROM {mu} WHERE id = $1 AND bank_id = $2", str(unit_id), bank_id) + return True + + +async def set_invalidation_reason(*, conn, fq_table, bank_id: str, unit_id: str, reason: str | None) -> None: + await conn.execute( + f"UPDATE {fq_table('invalidated_memory_units')} SET invalidation_reason = $3 WHERE id = $1 AND bank_id = $2", + str(unit_id), + bank_id, + reason, + ) + + +async def restore_memory(*, conn, fq_table, bank_id: str, unit_id: str) -> StoredMemory | None: + mu = fq_table("memory_units") + arch = fq_table("invalidated_memory_units") + ue = fq_table("unit_entities") + ent = fq_table("entities") + arch_cols = await _archive_columns(conn, fq_table) + + arch_row = await conn.fetchrow( + f"SELECT {_ARCHIVE_SELECT} FROM {arch} WHERE id = $1 AND bank_id = $2", str(unit_id), bank_id + ) + if arch_row is None: + return None + + # Move the row back. The archive omits embedding/search_vector, so both default + # to NULL here; search_vector is rebuilt now, the embedding by the caller. + await conn.execute( + f"INSERT INTO {mu} ({arch_cols}) SELECT {arch_cols} FROM {arch} WHERE id = $1 AND bank_id = $2", + str(unit_id), + bank_id, + ) + # Rebuild search_vector with the *current* backend, so a backend change while + # the fact sat archived cannot leave a stale/wrong-type vector (#2503). None + # means the backend indexes base columns directly and leaves it empty. + from ...db.ops_postgresql import pg_search_vector_expr + + sv_expr = pg_search_vector_expr(get_config()) + if sv_expr is not None: + await conn.execute( + f"UPDATE {mu} SET search_vector = {sv_expr} WHERE id = $1 AND bank_id = $2", str(unit_id), bank_id + ) + # Re-consolidate from scratch; links are rebuilt by graph maintenance. + await conn.execute( + f"UPDATE {mu} SET consolidated_at = NULL, consolidation_failed_at = NULL, updated_at = now() " + f"WHERE id = $1 AND bank_id = $2", + str(unit_id), + bank_id, + ) + # Restore the entity postings for entities that still exist — some may have + # been swept as orphans while the memory was archived. + if arch_row["entity_ids"]: + await conn.execute( + f"INSERT INTO {ue} (unit_id, entity_id) " + f"SELECT $1, eid FROM unnest($2::uuid[]) AS eid " + f"WHERE EXISTS (SELECT 1 FROM {ent} e WHERE e.id = eid AND e.bank_id = $3) " + f"ON CONFLICT DO NOTHING", + str(unit_id), + arch_row["entity_ids"], + bank_id, + ) + # Rematerialize the causal edges parked at invalidation (#2864). Edges whose peer is still + # archived or permanently deleted are skipped — the peer keeps its own copy and recreates the + # edge when it reverts, so the restore is order-independent and idempotent. + from ...retain.link_utils import rematerialize_causal_links + from .graph import _ops_for + + causal_json = await conn.fetchval( + f"SELECT causal_links FROM {arch} WHERE id = $1 AND bank_id = $2", str(unit_id), bank_id + ) + if causal_json: + await rematerialize_causal_links(conn, bank_id, conn.parse_json(causal_json) or [], ops=_ops_for(conn)) + # Invalidation cascaded away this unit's derived outgoing links; queue it so graph maintenance + # rebuilds them (the drain only touches queued units — it never scans for missing adjacency). + await _ops_for(conn).enqueue_graph_maintenance( + conn, fq_table("graph_maintenance_queue"), bank_id, [uuid.UUID(str(unit_id))] + ) + await conn.execute(f"DELETE FROM {arch} WHERE id = $1 AND bank_id = $2", str(unit_id), bank_id) + return _archived_stored(arch_row) + + +async def set_memory_embedding(*, conn, fq_table, bank_id: str, unit_id: str, embedding) -> None: + await conn.execute( + f"UPDATE {fq_table('memory_units')} SET embedding = $3::vector WHERE id = $1 AND bank_id = $2", + str(unit_id), + bank_id, + embedding, + ) + + +async def clear_unit_entities(*, conn, fq_table, bank_id: str, unit_id: str) -> None: + await conn.execute(f"DELETE FROM {fq_table('unit_entities')} WHERE unit_id = $1", str(unit_id)) + + +async def apply_edit( + *, + conn, + fq_table, + bank_id: str, + unit_id: str, + text: str, + context: str | None, + fact_type: str, + occurred_start, + occurred_end, + event_date, + mentioned_at, + entity_ids: list[str] | None, +) -> None: + # `entity_ids` and `mentioned_at` are unused here: the entity postings are + # re-linked into `unit_entities` by the caller, and an edit does not move the + # mention time. Both are on the signature for a store that carries entities on + # the memory and rebuilds it wholesale. + from ...causal_links import CAUSAL_LINK_TYPES + from ...db.ops_postgresql import pg_search_vector_expr + + mu = fq_table("memory_units") + ml = fq_table("memory_links") + # The caller enqueues the relink victims (and the edited unit itself, via + # ``include_affected_units``) before invoking this — one combined queue insert keeps the + # graph-maintenance queue's lock ordering intact. + # Keep the stored text-search vector in sync with the edited text/context. + # Reference the bind parameters, not the columns: PostgreSQL evaluates the + # UPDATE's RHS before the sibling SET assignments land, so a column reference + # would see the pre-edit values. + sv_expr = pg_search_vector_expr(get_config(), text_col="$3", context_col="$4") + sv_clause = f", search_vector = {sv_expr}" if sv_expr else "" + await conn.execute( + f""" + UPDATE {mu} + SET text = $3, context = $4, fact_type = $5, occurred_start = $6, occurred_end = $7, + event_date = $8, consolidated_at = NULL, consolidation_failed_at = NULL, + edited_at = now(), updated_at = now(){sv_clause} + WHERE id = $1 AND bank_id = $2 + """, + str(unit_id), + bank_id, + text, + context, + fact_type, + occurred_start, + occurred_end, + event_date, + ) + # Drop only the DERIVED links — graph maintenance recomputes temporal/semantic. Causal edges + # are retain-time extraction output that nothing recreates, so an edit preserves them (#2864). + await conn.execute( + f"DELETE FROM {ml} WHERE (from_unit_id = $1 OR to_unit_id = $1) AND NOT (link_type = ANY($2::text[]))", + str(unit_id), + list(CAUSAL_LINK_TYPES), + ) + + +__all__ = [ + "apply_edit", + "clear_unit_entities", + "delete_document", + "delete_observations", + "delete_stale_observations", + "get_archived_memory", + "insert_facts", + "invalidate_memory", + "observations_for_sources", + "restore_memory", + "set_invalidation_reason", + "set_memory_embedding", +] diff --git a/hindsight-api-slim/hindsight_api/engine/memories/postgres.py b/hindsight-api-slim/hindsight_api/engine/memories/postgres.py new file mode 100644 index 0000000000..ac675259f8 --- /dev/null +++ b/hindsight-api-slim/hindsight_api/engine/memories/postgres.py @@ -0,0 +1,484 @@ +"""The default memories store: Postgres holds the memories and the links. + +This is the behaviour Hindsight has always had, stated as an implementation of +:class:`~hindsight_api.engine.memories.base.MemoriesExtension` rather than as the +absence of one. Rows go in `memory_units`, the joins around it are `memory_links` +and `unit_entities`, and every read is SQL — writing a row *is* indexing it, so +:meth:`index_facts` has nothing left to do. + +The class is deliberately thin. Each method delegates to a plain function in +:mod:`hindsight_api.engine.memories.pg`, split by what calls it — curation, +graph, reads, writes — so a change to one area is a change to one file, and the +SQL is grouped by concern rather than piled behind a class. The two retrieval +arms delegate further out still, to the query functions that already own them in +:mod:`hindsight_api.engine.search.retrieval`. + +Keeping this as an explicit store (rather than an ``if store is None`` branch at +each call site) means the default path is the one the whole test suite exercises, +and a second implementation cannot change it by accident. +""" + +from __future__ import annotations + +from datetime import datetime +from typing import Any + +from .base import DeletePredicate, MemoriesExtension, MemoryPatch, ScanPage, StoredMemory +from .pg import counts, curation, graph, reads, writes + + +class PostgresMemories(MemoriesExtension): + """Memories in `memory_units`, links in `memory_links` / `unit_entities`.""" + + name = "postgres" + + # ------------------------------------------------------------------ writes + + async def insert_facts( + self, + *, + conn, + ops, + bank_id: str, + facts: list, + document_id: str | None = None, + defer_index: bool = False, + txn=None, + ) -> list[str]: + # `txn` is ignored: Postgres memories live in the caller's own transaction, so the + # write is already atomic with it — there is no separate store to hold invisible. + # `defer_index` is meaningless here: the INSERT that returns the ids is + # also what indexes the facts, so there is nothing to defer. + return await writes.insert_facts(conn=conn, ops=ops, bank_id=bank_id, facts=facts, document_id=document_id) + + async def delete_facts(self, bank_id: str, unit_ids: list[str], *, txn=None) -> None: + """No-op: the caller's `memory_units` DELETE (or its FK cascade) removed them.""" + + async def delete_where(self, bank_id: str, predicate: DeletePredicate, txn=None) -> int: + """No-op: predicate deletes are issued as SQL by the caller that owns the transaction.""" + return 0 + + async def delete_document(self, *, conn, fq_table, bank_id: str, document_id: str, txn=None) -> None: + # `txn` ignored: Postgres memories are covered by the caller's own transaction. + await writes.delete_document(conn=conn, fq_table=fq_table, bank_id=bank_id, document_id=document_id) + + async def delete_namespace(self, bank_id: str) -> None: + """No-op: deleting the bank cascades to its memories.""" + + async def delete_observations(self, *, conn, fq_table, bank_id: str, txn=None) -> None: + await writes.delete_observations(conn=conn, fq_table=fq_table, bank_id=bank_id) + + async def update_memories(self, bank_id: str, patches: list[MemoryPatch], txn=None) -> None: + """No-op: the caller's UPDATE already wrote the row it holds open.""" + + # ------------------------------------------------------------------ recall arms + + async def search( + self, + *, + conn, + bank_id: str, + fact_types: list[str], + query_embedding: str, + query_text: str, + limit: int, + tags: list[str] | None = None, + tags_match: str = "any", + tag_groups: list | None = None, + created_after: datetime | None = None, + created_before: datetime | None = None, + min_semantic: float | None = None, + min_keyword: float | None = None, + graph_seed_min_similarity: float | None = None, + ) -> "dict[str, SemanticBm25Result]": + # Imported here: retrieval imports this package, so a module-level import + # would close the cycle. + from ..search.retrieval import retrieve_semantic_bm25_combined_sql + + return await retrieve_semantic_bm25_combined_sql( + conn, + query_embedding, + query_text, + bank_id, + fact_types, + limit, + tags=tags, + tags_match=tags_match, + tag_groups=tag_groups, + created_after=created_after, + created_before=created_before, + min_semantic=min_semantic, + min_keyword=min_keyword, + graph_seed_min_similarity=graph_seed_min_similarity, + ) + + async def temporal_search( + self, + *, + conn, + bank_id: str, + fact_types: list[str], + query_embedding: str, + start_date: datetime, + end_date: datetime, + limit: int, + semantic_threshold: float = 0.1, + tags: list[str] | None = None, + tags_match: str = "any", + tag_groups: list | None = None, + created_after: datetime | None = None, + created_before: datetime | None = None, + ) -> dict[str, list]: + from ..search.retrieval import retrieve_temporal_combined_sql + + return await retrieve_temporal_combined_sql( + conn, + query_embedding, + bank_id, + fact_types, + start_date, + end_date, + limit, + semantic_threshold=semantic_threshold, + tags=tags, + tags_match=tags_match, + tag_groups=tag_groups, + created_after=created_after, + created_before=created_before, + ) + + # ------------------------------------------------------------------ addressed reads + + async def get_memories(self, *, conn, fq_table, bank_id: str, unit_ids: list[str]) -> list[StoredMemory]: + return await reads.get_memories(conn=conn, fq_table=fq_table, bank_id=bank_id, unit_ids=unit_ids) + + async def scan_memories( + self, + *, + conn, + fq_table, + bank_id: str, + fact_types: list[str] | None = None, + limit: int = 100, + page_token: str = "", + tags: list[str] | None = None, + tags_match: str = "any", + tag_groups: list | None = None, + document_id: str | None = None, + metadata_equals: dict[str, str] | None = None, + skip: int = 0, + include_edges: bool = False, + ) -> ScanPage: + return await reads.scan_memories( + conn=conn, + fq_table=fq_table, + bank_id=bank_id, + fact_types=fact_types, + limit=limit, + page_token=page_token, + tags=tags, + tags_match=tags_match, + tag_groups=tag_groups, + document_id=document_id, + metadata_equals=metadata_equals, + skip=skip, + include_edges=include_edges, + ) + + async def count_memories(self, *, conn, fq_table, bank_id: str) -> dict[str, int]: + return await reads.count_memories(conn=conn, fq_table=fq_table, bank_id=bank_id) + + async def list_tags(self, *, conn, fq_table, bank_id: str) -> dict[str, int]: + return await reads.list_tags(conn=conn, fq_table=fq_table, bank_id=bank_id) + + async def find_unconsolidated( + self, + *, + conn, + fq_table, + bank_id: str, + fact_types: list[str], + limit: int, + scope_tags: list[str] | None = None, + ) -> list[StoredMemory]: + return await reads.find_unconsolidated( + conn=conn, + fq_table=fq_table, + bank_id=bank_id, + fact_types=fact_types, + limit=limit, + scope_tags=scope_tags, + ) + + async def mark_consolidated( + self, + *, + conn, + fq_table, + bank_id: str, + unit_ids: list[str], + when: datetime | None, + failed: bool = False, + txn=None, + ) -> None: + await reads.mark_consolidated( + conn=conn, fq_table=fq_table, bank_id=bank_id, unit_ids=unit_ids, when=when, failed=failed + ) + + async def any_memory_updated_since( + self, + *, + conn, + fq_table, + bank_id: str, + since: datetime, + fact_types: list[str] | None = None, + tags: list[str] | None = None, + tags_match: str = "any", + tag_groups: list | None = None, + ) -> bool: + return await reads.any_memory_updated_since( + conn=conn, + fq_table=fq_table, + bank_id=bank_id, + since=since, + fact_types=fact_types, + tags=tags, + tags_match=tags_match, + tag_groups=tag_groups, + ) + + # -- count surfaces -- + + async def consolidation_freshness(self, *, conn, fq_table, bank_id: str) -> dict[str, Any]: + return await counts.consolidation_freshness(conn=conn, fq_table=fq_table, bank_id=bank_id) + + async def document_memory_counts(self, *, conn, fq_table, bank_id: str, document_ids: list[str]) -> dict[str, int]: + return await counts.document_memory_counts( + conn=conn, fq_table=fq_table, bank_id=bank_id, document_ids=document_ids + ) + + async def link_counts(self, *, conn, fq_table, bank_id: str) -> dict[str, int]: + return await counts.link_counts(conn=conn, fq_table=fq_table, bank_id=bank_id) + + async def memories_timeseries( + self, *, conn, fq_table, bank_id: str, time_field: str, trunc: str, since: datetime + ) -> list[dict[str, Any]]: + return await counts.memories_timeseries( + conn=conn, fq_table=fq_table, bank_id=bank_id, time_field=time_field, trunc=trunc, since=since + ) + + async def observation_scope_counts(self, *, conn, fq_table, bank_id: str) -> list[dict[str, Any]]: + return await counts.observation_scope_counts(conn=conn, fq_table=fq_table, bank_id=bank_id) + + # ------------------------------------------------------------------ observations + + async def upsert_observation(self, *, conn, bank_id: str, record, txn=None) -> None: + """No-op: the observation was written as a `memory_units` row by the caller.""" + + async def observations_for_sources( + self, *, conn, ops, fq_table, bank_id: str, unit_ids: list[str] + ) -> list[StoredMemory]: + return await writes.observations_for_sources( + conn=conn, ops=ops, fq_table=fq_table, bank_id=bank_id, unit_ids=unit_ids + ) + + async def delete_stale_observations(self, *, conn, ops, fq_table, bank_id: str, fact_ids: list) -> int: + return await writes.delete_stale_observations( + conn=conn, ops=ops, fq_table=fq_table, bank_id=bank_id, fact_ids=fact_ids + ) + + # ------------------------------------------------------------------ curation reads + + async def list_memory_units( + self, + *, + conn, + ops, + fq_table, + bank_id: str, + fact_type: str | None = None, + search_query: str | None = None, + consolidation_state: str | None = None, + state: str | None = None, + document_id: str | None = None, + entity_id: str | None = None, + tags: list[str] | None = None, + tags_match: str = "any", + created_before: datetime | None = None, + limit: int = 100, + offset: int = 0, + ) -> dict[str, Any]: + return await curation.list_memory_units( + conn=conn, + ops=ops, + fq_table=fq_table, + bank_id=bank_id, + fact_type=fact_type, + search_query=search_query, + consolidation_state=consolidation_state, + state=state, + document_id=document_id, + entity_id=entity_id, + tags=tags, + tags_match=tags_match, + created_before=created_before, + limit=limit, + offset=offset, + ) + + async def get_memory_unit(self, *, conn, ops, fq_table, bank_id: str, unit_id: str) -> dict[str, Any] | None: + return await curation.get_memory_unit(conn=conn, ops=ops, fq_table=fq_table, bank_id=bank_id, unit_id=unit_id) + + # -- curation archive -- + + async def get_archived_memory(self, *, conn, fq_table, bank_id: str, unit_id: str) -> StoredMemory | None: + return await writes.get_archived_memory(conn=conn, fq_table=fq_table, bank_id=bank_id, unit_id=unit_id) + + async def invalidate_memory( + self, *, conn, fq_table, bank_id: str, unit_id: str, reason: str | None, txn=None + ) -> bool: + return await writes.invalidate_memory( + conn=conn, fq_table=fq_table, bank_id=bank_id, unit_id=unit_id, reason=reason + ) + + async def set_invalidation_reason(self, *, conn, fq_table, bank_id: str, unit_id: str, reason: str | None) -> None: + await writes.set_invalidation_reason( + conn=conn, fq_table=fq_table, bank_id=bank_id, unit_id=unit_id, reason=reason + ) + + async def restore_memory(self, *, conn, fq_table, bank_id: str, unit_id: str, txn=None) -> StoredMemory | None: + return await writes.restore_memory(conn=conn, fq_table=fq_table, bank_id=bank_id, unit_id=unit_id) + + async def set_memory_embedding(self, *, conn, fq_table, bank_id: str, unit_id: str, embedding, txn=None) -> None: + await writes.set_memory_embedding( + conn=conn, fq_table=fq_table, bank_id=bank_id, unit_id=unit_id, embedding=embedding + ) + + async def clear_unit_entities(self, *, conn, fq_table, bank_id: str, unit_id: str) -> None: + await writes.clear_unit_entities(conn=conn, fq_table=fq_table, bank_id=bank_id, unit_id=unit_id) + + async def apply_edit( + self, + *, + conn, + fq_table, + bank_id: str, + unit_id: str, + text: str, + context: str | None, + fact_type: str, + occurred_start, + occurred_end, + event_date, + mentioned_at, + entity_ids: list[str] | None, + txn=None, + ) -> None: + await writes.apply_edit( + conn=conn, + fq_table=fq_table, + bank_id=bank_id, + unit_id=unit_id, + text=text, + context=context, + fact_type=fact_type, + occurred_start=occurred_start, + occurred_end=occurred_end, + event_date=event_date, + mentioned_at=mentioned_at, + entity_ids=entity_ids, + ) + + async def list_entities( + self, + *, + conn, + fq_table, + bank_id: str, + search: str | None = None, + limit: int = 100, + offset: int = 0, + ) -> dict[str, Any]: + return await curation.list_entities( + conn=conn, fq_table=fq_table, bank_id=bank_id, search=search, limit=limit, offset=offset + ) + + # ------------------------------------------------------------------ graph + + async def graph_units( + self, + *, + conn, + fq_table, + bank_id: str, + fact_type: str | None = None, + search_query: str | None = None, + document_id: str | None = None, + chunk_id: str | None = None, + tags: list[str] | None = None, + tags_match: str = "all_strict", + limit: int = 1000, + ) -> dict[str, Any]: + return await graph.graph_units( + conn=conn, + fq_table=fq_table, + bank_id=bank_id, + fact_type=fact_type, + search_query=search_query, + document_id=document_id, + chunk_id=chunk_id, + tags=tags, + tags_match=tags_match, + limit=limit, + ) + + async def graph_entity_rows(self, *, conn, fq_table, bank_id: str, unit_ids: list[str]) -> list[dict[str, Any]]: + return await graph.graph_entity_rows(conn=conn, fq_table=fq_table, bank_id=bank_id, unit_ids=unit_ids) + + async def graph_direct_links(self, *, conn, fq_table, bank_id: str, unit_ids: list[str]) -> list[dict[str, Any]]: + return await graph.graph_direct_links(conn=conn, fq_table=fq_table, bank_id=bank_id, unit_ids=unit_ids) + + async def entity_memory_counts( + self, *, conn, fq_table, bank_id: str, entity_ids: list[str] | None = None + ) -> dict[str, int]: + return await graph.entity_memory_counts(conn=conn, fq_table=fq_table, bank_id=bank_id, entity_ids=entity_ids) + + async def entities_for_units(self, *, conn, fq_table, bank_id: str, unit_ids: list[str]) -> dict[str, list[str]]: + return await graph.entities_for_units(conn=conn, fq_table=fq_table, bank_id=bank_id, unit_ids=unit_ids) + + async def entity_map_for_units( + self, *, conn, fq_table, bank_id: str, unit_ids: list[str] + ) -> dict[str, list[dict[str, str]]]: + return await graph.entity_map_for_units(conn=conn, fq_table=fq_table, bank_id=bank_id, unit_ids=unit_ids) + + # ------------------------------------------------------------------ maintenance + + async def record_unit_entities( + self, *, conn, ops, fq_table, bank_id: str | None = None, unit_ids: list[Any], entity_ids: list[Any] + ) -> None: + # The join is keyed by global unit id, so bank_id is not needed here. + await ops.bulk_insert_unit_entities(conn, fq_table("unit_entities"), unit_ids, entity_ids) + + async def enqueue_relink_victims( + self, *, conn, fq_table, bank_id: str, affected_unit_ids: list, include_affected_units: bool = False + ) -> int: + return await graph.enqueue_relink_victims( + conn=conn, + fq_table=fq_table, + bank_id=bank_id, + affected_unit_ids=affected_unit_ids, + include_affected_units=include_affected_units, + ) + + async def relink_pass(self, *, backend, fq_table, bank_id: str, config) -> dict: + return await graph.relink_pass(backend=backend, fq_table=fq_table, bank_id=bank_id, config=config) + + async def prune_orphan_entities(self, *, conn, fq_table, bank_id: str) -> int: + return await graph.prune_orphan_entities(conn=conn, fq_table=fq_table, bank_id=bank_id) + + async def prune_stale_cooccurrences(self, *, conn, fq_table, bank_id: str) -> int: + return await graph.prune_stale_cooccurrences(conn=conn, fq_table=fq_table, bank_id=bank_id) + + +__all__ = ["PostgresMemories"] diff --git a/hindsight-api-slim/hindsight_api/engine/memory_engine.py b/hindsight-api-slim/hindsight_api/engine/memory_engine.py index 0d6beefada..fa46201ce8 100644 --- a/hindsight-api-slim/hindsight_api/engine/memory_engine.py +++ b/hindsight-api-slim/hindsight_api/engine/memory_engine.py @@ -2912,6 +2912,15 @@ async def init_query_analyzer(): # Query analyzer load is sync and CPU-bound await loop.run_in_executor(None, self.query_analyzer.load) + async def init_memories(): + """Bring up the memories store's own resources (connection pool, + client, …) once at startup. The default Postgres store treats this as + a no-op; a store that owns an external service builds its client here + so the first request does not race an uninitialized handle.""" + from .memories import get_memories + + await get_memories().initialize() + async def verify_llm(): """Verify LLM connections are working for all unique configs. @@ -2996,6 +3005,7 @@ async def verify_llm(): init_embeddings(), init_query_analyzer(), init_cross_encoder(), + init_memories(), ] # Only verify LLM if not skipping @@ -3380,6 +3390,15 @@ async def close(self): # Shutdown task backend await self._task_backend.shutdown() + # Release the memories store's own resources (client/pool). No-op for the + # default Postgres store; symmetric with init_memories() at startup. + try: + from .memories import get_memories + + await get_memories().shutdown() + except Exception as e: + logger.warning(f"Error shutting down memories store: {e}") + # Close HTTP client used for webhook delivery if self._http_client is not None: await self._http_client.aclose() @@ -5201,15 +5220,21 @@ def to_tuple_format(results): if observation_ids: dedup_start = time.time() superseded_ids: set[str] = set() + from .memories import get_memories + async with acquire_with_retry(backend) as dedup_conn: - obs_rows = await dedup_conn.fetch( - f""" - SELECT source_memory_ids - FROM {fq_table("memory_units")} - WHERE id = ANY($1::uuid[]) AND fact_type = 'observation' - """, - observation_ids, - ) + # The observation carries its sources; the store resolves + # them all in one addressed read. + obs_rows = [ + {"source_memory_ids": m.source_memory_ids} + for m in await get_memories().get_memories( + conn=dedup_conn, + fq_table=fq_table, + bank_id=bank_id, + unit_ids=[str(o) for o in observation_ids], + ) + if m.fact_type == "observation" + ] if tracer: tracer.add_phase_metric( "prefer_observations_dedup", @@ -5263,7 +5288,35 @@ def to_tuple_format(results): # Resolve source chunk_ids for all observations in a single query, # ordered by observation rank so per-observation results stay grouped correctly. obs_chunk_ids: dict[str, list[str]] = {} - if observation_ids_ordered: + from .memories import get_memories + + _obs_store = get_memories() + if observation_ids_ordered and not _obs_store.writes_memory_rows_in_sql: + # A store that keeps memories outside SQL: fetch each observation, then its + # source memories, for their chunk_ids — the join the SQL branch does, walked + # in observation-rank order so per-observation grouping is preserved. + obs_units = await _obs_store.get_memories( + conn=None, + fq_table=fq_table, + bank_id=bank_id, + unit_ids=[str(o) for o in observation_ids_ordered], + ) + by_obs = {u.unit_id: u for u in obs_units} + src_ids = [sid for u in obs_units for sid in u.source_memory_ids] + srcs = await _obs_store.get_memories( + conn=None, fq_table=fq_table, bank_id=bank_id, unit_ids=list(dict.fromkeys(src_ids)) + ) + src_chunk = {s.unit_id: s.chunk_id for s in srcs} + for _obs_uuid in observation_ids_ordered: + _obs = by_obs.get(str(_obs_uuid)) + if not _obs: + continue + for _sid in _obs.source_memory_ids: + _cid = src_chunk.get(_sid) + if _cid and _cid not in seen_chunk_ids: + obs_chunk_ids.setdefault(str(_obs_uuid), []).append(_cid) + seen_chunk_ids.add(_cid) + elif observation_ids_ordered: async with acquire_with_retry(backend) as obs_conn: if self._backend.ops.uses_observation_sources_table: obs_source_rows = await obs_conn.fetch( @@ -5448,16 +5501,22 @@ def to_tuple_format(results): if include_source_facts: observation_ids = [uuid.UUID(sr.id) for sr in top_scored if sr.retrieval.fact_type == "observation"] if observation_ids: + from .memories import get_memories + + store = get_memories() async with acquire_with_retry(backend) as sf_conn: - # Fetch source_memory_ids for all observation results - obs_rows = await sf_conn.fetch( - f""" - SELECT id, source_memory_ids - FROM {fq_table("memory_units")} - WHERE id = ANY($1::uuid[]) AND fact_type = 'observation' - """, - observation_ids, - ) + # Each observation carries its sources; one addressed read + # resolves them. + obs_rows = [ + {"id": m.unit_id, "source_memory_ids": m.source_memory_ids} + for m in await store.get_memories( + conn=sf_conn, + fq_table=fq_table, + bank_id=bank_id, + unit_ids=[str(o) for o in observation_ids], + ) + if m.fact_type == "observation" + ] # Collect unique source IDs in order of first appearance seen_source_ids: set[str] = set() @@ -5473,18 +5532,26 @@ def to_tuple_format(results): # Fetch source fact content up to token budget if source_ids_ordered: - import uuid as uuid_module - - source_rows = await sf_conn.fetch( - f""" - SELECT id, text, fact_type, context, occurred_start, occurred_end, - mentioned_at, document_id, chunk_id, tags, metadata - FROM {fq_table("memory_units")} - WHERE id = ANY($1::uuid[]) - """, - [uuid_module.UUID(sid) for sid in source_ids_ordered], - ) - source_row_by_id = {str(r["id"]): r for r in source_rows} + # The source facts, as the store holds them — same + # columns, shaped as dicts so the rendering below is shared. + source_row_by_id = { + m.unit_id: { + "id": m.unit_id, + "text": m.text, + "fact_type": m.fact_type, + "context": m.context, + "occurred_start": m.occurred_start, + "occurred_end": m.occurred_end, + "mentioned_at": m.mentioned_at, + "document_id": m.document_id, + "chunk_id": m.chunk_id, + "tags": list(m.tags or []), + "metadata": m.metadata, + } + for m in await store.get_memories( + conn=sf_conn, fq_table=fq_table, bank_id=bank_id, unit_ids=source_ids_ordered + ) + } encoding = _get_tiktoken_encoding() source_facts_dict = {} @@ -5548,22 +5615,21 @@ def _make_source_fact(sid: str, r: Any) -> MemoryFact: entity_build_start = time.time() # Get entities for each fact if include_entities is requested. - # _entity_rows_for_units_sql resolves both direct unit_entities rows - # and observation-via-source-memory inheritance in a single query. + # The store resolves both a memory's direct entity postings and an + # observation's inherited-from-sources entities in one call. fact_entity_map = {} # unit_id -> list of {entity_id, canonical_name} if include_entities and top_scored: - unit_ids = [uuid.UUID(sr.id) for sr in top_scored] + unit_ids = [sr.id for sr in top_scored] if unit_ids: + from .memories import get_memories + async with acquire_with_retry(backend) as entity_conn: - entity_rows = await entity_conn.fetch( - self._entity_rows_for_units_sql(unit_ids_placeholder=1), - unit_ids, + # The memory carries its own entity ids; the store resolves + # them to names (observations inherit their sources'), the + # `entities` registry staying in postgres. + fact_entity_map = await get_memories().entity_map_for_units( + conn=entity_conn, fq_table=fq_table, bank_id=bank_id, unit_ids=unit_ids ) - for row in entity_rows: - unit_id = str(row["unit_id"]) - fact_entity_map.setdefault(unit_id, []).append( - {"entity_id": str(row["entity_id"]), "canonical_name": row["canonical_name"]} - ) # Convert results to MemoryFact objects # Build per-result scores (final/reranker/semantic/text) keyed by id. @@ -5715,59 +5781,6 @@ def _make_source_fact(sid: str, r: Any) -> MemoryFact: logger.error("\n" + "\n".join(log_buffer), exc_info=True) raise RuntimeError(f"Failed to search memories ({type(e).__name__}): {e!r}") from e - def _entity_rows_for_units_sql(self, unit_ids_placeholder: int) -> str: - """SQL SELECT producing ``(unit_id, entity_id, canonical_name)`` rows for - the given unit IDs. - - Direct rows come from ``unit_entities``. Observations rarely carry - direct rows there; their entity association lives transitively through - their source memories (``source_memory_ids`` on PG, the - ``observation_sources`` junction on Oracle). When an observation has - no direct entity rows the SELECT inherits its source memories' - entities, so the result is the same set callers would get from - ``get_memory_unit``. - - ``unit_ids_placeholder`` is the 1-based parameter index that holds the - ``uuid[]`` of unit IDs. The placeholder is referenced twice — both - sides of the UNION need it — so callers should not reuse the slot. - """ - ue = fq_table("unit_entities") - ents = fq_table("entities") - mu = fq_table("memory_units") - p = unit_ids_placeholder - - direct = ( - f"SELECT ue.unit_id, e.id AS entity_id, e.canonical_name " - f"FROM {ue} ue " - f"JOIN {ents} e ON e.id = ue.entity_id " - f"WHERE ue.unit_id = ANY(${p}::uuid[])" - ) - - if self._backend.ops.uses_observation_sources_table: - os_t = fq_table("observation_sources") - inherited = ( - f"SELECT os.observation_id AS unit_id, e.id AS entity_id, e.canonical_name " - f"FROM {os_t} os " - f"JOIN {ue} src_ue ON src_ue.unit_id = os.source_id " - f"JOIN {ents} e ON e.id = src_ue.entity_id " - f"WHERE os.observation_id = ANY(${p}::uuid[]) " - f"AND NOT EXISTS (SELECT 1 FROM {ue} d WHERE d.unit_id = os.observation_id)" - ) - else: - inherited = ( - f"SELECT obs.id AS unit_id, e.id AS entity_id, e.canonical_name " - f"FROM {mu} obs " - f"CROSS JOIN LATERAL unnest(obs.source_memory_ids) AS src_id " - f"JOIN {ue} src_ue ON src_ue.unit_id = src_id " - f"JOIN {ents} e ON e.id = src_ue.entity_id " - f"WHERE obs.id = ANY(${p}::uuid[]) " - f"AND obs.fact_type = 'observation' " - f"AND obs.source_memory_ids IS NOT NULL " - f"AND NOT EXISTS (SELECT 1 FROM {ue} d WHERE d.unit_id = obs.id)" - ) - - return f"({direct}) UNION ({inherited})" - def _filter_by_token_budget( self, results: list[dict[str, Any]], max_tokens: int ) -> tuple[list[dict[str, Any]], int]: @@ -5867,31 +5880,77 @@ async def get_document( f"WHERE bank_id = $2 AND fact_type = 'observation' AND {obs_match})" ) - # Use a subquery for counts to avoid GROUP BY on CLOB columns - # (Oracle cannot use CLOB types as comparison keys in GROUP BY). - doc = await conn.fetchrow( - f""" - SELECT d.id, d.bank_id, d.original_text, d.content_hash, - d.created_at, d.updated_at, d.tags, d.retain_params, - COALESCE(stats.unit_count, 0) as unit_count, - COALESCE(stats.world_count, 0) as world_count, - COALESCE(stats.experience_count, 0) as experience_count, - COALESCE({observation_count_sql}, 0) as observation_count - FROM {fq_table("documents")} d - LEFT JOIN ( - SELECT mu.document_id, mu.bank_id, - COUNT(mu.id) as unit_count, - COUNT(CASE WHEN mu.fact_type = 'world' THEN 1 END) as world_count, - COUNT(CASE WHEN mu.fact_type = 'experience' THEN 1 END) as experience_count - FROM {fq_table("memory_units")} mu - WHERE mu.document_id = $1 AND mu.bank_id = $2 - GROUP BY mu.document_id, mu.bank_id - ) stats ON stats.document_id = d.id AND stats.bank_id = d.bank_id - WHERE d.id = $1 AND d.bank_id = $2 - """, - document_id, - bank_id, - ) + from .memories import get_memories + + _store = get_memories() + if _store.writes_memory_rows_in_sql: + # Use a subquery for counts to avoid GROUP BY on CLOB columns + # (Oracle cannot use CLOB types as comparison keys in GROUP BY). + doc = await conn.fetchrow( + f""" + SELECT d.id, d.bank_id, d.original_text, d.content_hash, + d.created_at, d.updated_at, d.tags, d.retain_params, + COALESCE(stats.unit_count, 0) as unit_count, + COALESCE(stats.world_count, 0) as world_count, + COALESCE(stats.experience_count, 0) as experience_count, + COALESCE({observation_count_sql}, 0) as observation_count + FROM {fq_table("documents")} d + LEFT JOIN ( + SELECT mu.document_id, mu.bank_id, + COUNT(mu.id) as unit_count, + COUNT(CASE WHEN mu.fact_type = 'world' THEN 1 END) as world_count, + COUNT(CASE WHEN mu.fact_type = 'experience' THEN 1 END) as experience_count + FROM {fq_table("memory_units")} mu + WHERE mu.document_id = $1 AND mu.bank_id = $2 + GROUP BY mu.document_id, mu.bank_id + ) stats ON stats.document_id = d.id AND stats.bank_id = d.bank_id + WHERE d.id = $1 AND d.bank_id = $2 + """, + document_id, + bank_id, + ) + else: + # A store that keeps memories outside SQL: the documents row is still SQL, but its + # per-fact-type counts come from the store (scan the document's memories; count the + # observations built on them via observations_for_sources). + _drow = await conn.fetchrow( + f""" + SELECT d.id, d.bank_id, d.original_text, d.content_hash, + d.created_at, d.updated_at, d.tags, d.retain_params + FROM {fq_table("documents")} d + WHERE d.id = $1 AND d.bank_id = $2 + """, + document_id, + bank_id, + ) + if _drow is None: + doc = None + else: + doc = dict(_drow) + _page = await _store.scan_memories( + conn=conn, fq_table=fq_table, bank_id=bank_id, document_id=document_id, limit=1_000_000 + ) + doc["unit_count"] = len(_page.memories) + doc["world_count"] = sum(1 for m in _page.memories if m.fact_type == "world") + doc["experience_count"] = sum(1 for m in _page.memories if m.fact_type == "experience") + _sids = [m.unit_id for m in _page.memories if m.fact_type in ("experience", "world")] + _obs = ( + await _store.observations_for_sources( + conn=conn, ops=self._backend.ops, fq_table=fq_table, bank_id=bank_id, unit_ids=_sids + ) + if _sids + else [] + ) + doc["observation_count"] = len(_obs) + # A store that owns the document store (memlake) keeps the extracted text in + # its own store, not in documents.original_text (which is NULL here). Overlay + # it from the store so get_document still returns the body. + if _store.owns_document_store: + _rec = await _store.get_document_record( + bank_id=bank_id, document_id=document_id, include_text=True + ) + if _rec is not None: + doc["original_text"] = _rec.get("original_text") if not doc: return None @@ -5954,17 +6013,38 @@ async def delete_document( await self._validate_operation(self._operation_validator.validate_bank_write(ctx)) backend = await self._get_backend() invalidated_obs = 0 + _del_txn = None async with acquire_with_retry(backend) as conn: async with conn.transaction(): - # Get memory unit IDs before deletion (for observation cleanup) - unit_rows = await conn.fetch( - f"SELECT id FROM {fq_table('memory_units')} WHERE document_id = $1 AND fact_type IN ('experience', 'world')", - document_id, - ) - unit_ids = [str(row["id"]) for row in unit_rows] - units_count = await conn.fetchval( - f"SELECT COUNT(*) FROM {fq_table('memory_units')} WHERE document_id = $1", document_id - ) + # Get memory unit IDs before deletion (for observation cleanup). A store that + # keeps memories outside SQL answers by document through the store — memory_units + # is empty for it, so the SQL below would find nothing to clean up. + from .memories import get_memories + + _store = get_memories() + if _store.writes_memory_rows_in_sql: + unit_rows = await conn.fetch( + f"SELECT id FROM {fq_table('memory_units')} WHERE document_id = $1 AND fact_type IN ('experience', 'world')", + document_id, + ) + unit_ids = [str(row["id"]) for row in unit_rows] + units_count = await conn.fetchval( + f"SELECT COUNT(*) FROM {fq_table('memory_units')} WHERE document_id = $1", document_id + ) + else: + src_page = await _store.scan_memories( + conn=conn, + fq_table=fq_table, + bank_id=bank_id, + document_id=document_id, + fact_types=["experience", "world"], + limit=1_000_000, + ) + unit_ids = [m.unit_id for m in src_page.memories] + _doc_counts = await _store.document_memory_counts( + conn=conn, fq_table=fq_table, bank_id=bank_id, document_ids=[document_id] + ) + units_count = _doc_counts.get(document_id, 0) # Capture relink victims BEFORE the cascade — once the source # rows are gone, the join finding them returns nothing. @@ -5984,6 +6064,22 @@ async def delete_document( bank_id, ) + # For a store that keeps memories outside SQL, deleting the documents row does not + # cascade to its memories (they are not SQL rows) — drop them through the store, + # tagged with a write-group so the store tombstone commits atomically with the + # Postgres document delete (a rolled-back delete must not orphan the memories). + if deleted and not _store.writes_memory_rows_in_sql: + _del_txn = await _store.begin_txn(conn=conn, fq_table=fq_table, bank_id=bank_id, mutating=True) + await _store.delete_document( + conn=conn, fq_table=fq_table, bank_id=bank_id, document_id=document_id, txn=_del_txn + ) + # A store that owns the document store also drops the document RECORD (its + # extracted text + chunk bodies; the orphan sweep reclaims the blobs), under the + # same write-group so it commits atomically with the Postgres document delete. + # This is the EXPLICIT deletion — distinct from the re-ingest facts-delete above. + if _store.owns_document_store: + await _store.delete_document_record(bank_id=bank_id, document_id=document_id, txn=_del_txn) + # Invalidate observations referencing these (now-deleted) memories if unit_ids: invalidated_obs = await self._delete_stale_observations_for_memories(conn, bank_id, unit_ids) @@ -5993,6 +6089,11 @@ async def delete_document( "memory_units_deleted": units_count if deleted else 0, } + # Postgres committed the delete: publish the store's tombstone write-group (no-op if + # nothing was deleted or the store keeps memories in SQL). + if _del_txn is not None: + await _store.decide_txn(_del_txn, commit=True) + # Drop any cached stats for this bank — deleting the document changed # the document count and (via cascade) the memory-unit/link counts # get_bank_stats reports, which the TTL would otherwise serve at @@ -6084,6 +6185,31 @@ async def update_document( return False if tags is not None: + from .memories import MemoryPatch, get_memories + + _store = get_memories() + if tags is not None and not _store.writes_memory_rows_in_sql: + # A store that keeps memories outside SQL: retag the document's memories, then + # invalidate the observations built on them and requeue their sources so the + # next consolidation rebuilds them under the new tags (the cascade the SQL + # branch does by hand — delete_stale_observations requeues surviving co-sources). + _doc_page = await _store.scan_memories( + conn=conn, fq_table=fq_table, bank_id=bank_id, document_id=document_id, limit=1_000_000 + ) + _doc_units = _doc_page.memories + if _doc_units: + await _store.update_memories( + bank_id, [MemoryPatch(unit_id=m.unit_id, tags=list(tags)) for m in _doc_units] + ) + _src_ids = [m.unit_id for m in _doc_units if m.fact_type in ("experience", "world")] + if _src_ids: + invalidated_obs = await _store.delete_stale_observations( + conn=conn, ops=self._backend.ops, fq_table=fq_table, bank_id=bank_id, fact_ids=_src_ids + ) + await _store.mark_consolidated( + conn=conn, fq_table=fq_table, bank_id=bank_id, unit_ids=_src_ids, when=None + ) + elif tags is not None: unit_rows = await conn.fetch( f"SELECT id FROM {fq_table('memory_units')} WHERE document_id = $1 AND fact_type IN ('experience', 'world')", document_id, @@ -6194,6 +6320,7 @@ async def delete_memory_unit( self, unit_id: str, *, + bank_id: str | None = None, request_context: "RequestContext", ) -> dict[str, Any]: """ @@ -6226,15 +6353,31 @@ async def delete_memory_unit( invalidated_obs = 0 bank_id_for_consolidation: str | None = None bank_id_for_graph_maintenance: str | None = None + _del_txn = None async with acquire_with_retry(backend) as conn: async with conn.transaction(): - # Get bank_id and fact_type before deletion - row = await conn.fetchrow( - f"SELECT bank_id, fact_type FROM {fq_table('memory_units')} WHERE id = $1", - str(unit_uuid), - ) - bank_id = row["bank_id"] if row else None - fact_type = row["fact_type"] if row else None + # Get bank_id and fact_type before deletion. A SQL store discovers the bank from + # the row itself; a store that keeps memories outside SQL is partitioned by bank, + # so the caller must say which one — hence the optional `bank_id` argument. + from .memories import get_memories + + _store = get_memories() + if _store.writes_memory_rows_in_sql: + row = await conn.fetchrow( + f"SELECT bank_id, fact_type FROM {fq_table('memory_units')} WHERE id = $1", + str(unit_uuid), + ) + bank_id = row["bank_id"] if row else None + fact_type = row["fact_type"] if row else None + else: + _found = ( + await _store.get_memories(conn=conn, fq_table=fq_table, bank_id=bank_id, unit_ids=[unit_id]) + if bank_id + else [] + ) + fact_type = _found[0].fact_type if _found else None + if not _found: + bank_id = None # Capture relink victims BEFORE the cascade — once the row is # gone, the join finding them returns nothing. @@ -6248,9 +6391,16 @@ async def delete_memory_unit( # observations inserted concurrently by consolidation (otherwise a # racing insert committed between the sweep and the delete would # leave an orphan referencing this just-deleted source memory). - deleted = await conn.fetchval( - f"DELETE FROM {fq_table('memory_units')} WHERE id = $1 RETURNING id", unit_id - ) + if _store.writes_memory_rows_in_sql: + deleted = await conn.fetchval( + f"DELETE FROM {fq_table('memory_units')} WHERE id = $1 RETURNING id", unit_id + ) + else: + deleted = unit_id if fact_type is not None else None + if deleted: + # Tag the store tombstone so it commits atomically with this transaction. + _del_txn = await _store.begin_txn(conn=conn, fq_table=fq_table, bank_id=bank_id, mutating=True) + await _store.delete_facts(bank_id, [unit_id], txn=_del_txn) # Invalidate observations referencing this (now-deleted) source memory if bank_id and fact_type in ("experience", "world"): @@ -6272,6 +6422,11 @@ async def delete_memory_unit( else "Memory unit not found", } + # Postgres committed: publish the store's tombstone write-group (no-op if nothing was + # deleted or the store keeps memories in SQL). + if _del_txn is not None: + await _store.decide_txn(_del_txn, commit=True) + # Drop any cached stats for this bank — the deleted unit (and its # cascaded links/entities) changed the counts get_bank_stats reports, # which the TTL would otherwise serve at pre-delete values for up to a @@ -6541,12 +6696,29 @@ async def delete_bank( # observations inserted concurrently by consolidation. unit_ids: list[str] = [] if fact_type in ("experience", "world"): - unit_id_rows = await conn.fetch( - f"SELECT id FROM {fq_table('memory_units')} WHERE bank_id = $1 AND fact_type = $2", - bank_id, - fact_type, - ) - unit_ids = [str(row["id"]) for row in unit_id_rows] + # These ids drive the stale-observation sweep below, so they must come + # from wherever the memories live: reading memory_units for a store that + # keeps them elsewhere yields nothing, and the sweep would silently skip, + # leaving observations behind that outlive the sources they summarise. + from .memories import get_memories as _get_memories_for_scope + + _scope_store = _get_memories_for_scope() + if _scope_store.writes_memory_rows_in_sql: + unit_id_rows = await conn.fetch( + f"SELECT id FROM {fq_table('memory_units')} WHERE bank_id = $1 AND fact_type = $2", + bank_id, + fact_type, + ) + unit_ids = [str(row["id"]) for row in unit_id_rows] + else: + _scope_page = await _scope_store.scan_memories( + conn=conn, + fq_table=fq_table, + bank_id=bank_id, + fact_types=[fact_type], + limit=1_000_000, + ) + unit_ids = [m.unit_id for m in _scope_page.memories] # Delete only memories of a specific fact type units_count = await conn.fetchval( @@ -6652,6 +6824,18 @@ async def delete_bank( lambda: bank_utils.drop_bank_vector_indexes(conn, bank_internal_id, ops=self._backend.ops) ) + # A store that keeps memories outside SQL leaves memory_units empty, so every DELETE + # above was a no-op on its data — it must be told to drop the bank's memories too, or + # they are orphaned. Runs after the transaction: it is an external-store call, not SQL. + from .memories import DeletePredicate, get_memories + + store = get_memories() + if not store.writes_memory_rows_in_sql: + if fact_type: + await store.delete_where(bank_id, DeletePredicate(fact_types=[fact_type])) + else: + await store.delete_namespace(bank_id) + # Drop any cached stats for this bank — counts have changed and the # TTL would otherwise serve pre-delete values for up to a minute. await self._bank_stats_cache.invalidate(get_current_schema(), bank_id) @@ -6690,28 +6874,52 @@ async def clear_observations( bank_id=bank_id, operation=BankWriteOperation.CLEAR_OBSERVATIONS, request_context=request_context ) await self._validate_operation(self._operation_validator.validate_bank_write(ctx)) + from .memories import get_memories + + store = get_memories() backend = await self._get_backend() async with acquire_with_retry(backend) as conn: async with conn.transaction(): - # Count observations before deletion - count = await conn.fetchval( - f"SELECT COUNT(*) FROM {fq_table('memory_units')} WHERE bank_id = $1 AND fact_type = 'observation'", - bank_id, - ) + if store.writes_memory_rows_in_sql: + # Count observations before deletion + count = await conn.fetchval( + f"SELECT COUNT(*) FROM {fq_table('memory_units')} WHERE bank_id = $1 AND fact_type = 'observation'", + bank_id, + ) - # Delete all observations - await conn.execute( - f"DELETE FROM {fq_table('memory_units')} WHERE bank_id = $1 AND fact_type = 'observation'", - bank_id, - ) + # Delete all observations + await conn.execute( + f"DELETE FROM {fq_table('memory_units')} WHERE bank_id = $1 AND fact_type = 'observation'", + bank_id, + ) - # Reset consolidated_at on source memories so they get re-consolidated - await conn.execute( - f"UPDATE {fq_table('memory_units')} SET consolidated_at = NULL WHERE bank_id = $1 AND fact_type IN ('experience', 'world')", - bank_id, - ) + # Reset consolidated_at on source memories so they get re-consolidated + await conn.execute( + f"UPDATE {fq_table('memory_units')} SET consolidated_at = NULL WHERE bank_id = $1 AND fact_type IN ('experience', 'world')", + bank_id, + ) + else: + # A store that keeps memories outside SQL: count + delete the observations + # through the store, then requeue every source (clear its consolidated marker, + # mark_consolidated(when=None)) so the next pass re-consolidates them. + count = (await store.count_memories(conn=conn, fq_table=fq_table, bank_id=bank_id)).get( + "observation", 0 + ) + await store.delete_observations(conn=conn, fq_table=fq_table, bank_id=bank_id) + src_page = await store.scan_memories( + conn=conn, + fq_table=fq_table, + bank_id=bank_id, + fact_types=["experience", "world"], + limit=1_000_000, + ) + src_ids = [m.unit_id for m in src_page.memories] + if src_ids: + await store.mark_consolidated( + conn=conn, fq_table=fq_table, bank_id=bank_id, unit_ids=src_ids, when=None + ) - # Reset consolidation timestamp + # Reset consolidation timestamp (Postgres banks bookkeeping, for every store) await conn.execute( f"UPDATE {fq_table('banks')} SET last_consolidated_at = NULL WHERE bank_id = $1", bank_id, @@ -6755,21 +6963,11 @@ async def list_observation_scopes( ) await self._validate_operation(self._operation_validator.validate_bank_read(ctx)) backend = await self._get_backend() + from .memories import get_memories + async with acquire_with_retry(backend) as conn: - rows = await conn.fetch( - f""" - SELECT scope, COUNT(*) AS count - FROM ( - SELECT COALESCE(ARRAY(SELECT unnest(tags) ORDER BY 1), '{{}}'::text[]) AS scope - FROM {fq_table("memory_units")} - WHERE bank_id = $1 AND fact_type = 'observation' - ) s - GROUP BY scope - ORDER BY count DESC, scope - """, - bank_id, - ) - return {"scopes": [{"tags": list(r["scope"]), "count": r["count"]} for r in rows]} + scopes = await get_memories().observation_scope_counts(conn=conn, fq_table=fq_table, bank_id=bank_id) + return {"scopes": scopes} async def retry_failed_consolidation( self, @@ -6802,27 +7000,41 @@ async def retry_failed_consolidation( request_context=request_context, ) await self._validate_operation(self._operation_validator.validate_bank_write(ctx)) + from .memories import get_memories + + store = get_memories() backend = await self._get_backend() async with acquire_with_retry(backend) as conn: - count = await conn.fetchval( - f""" - SELECT COUNT(*) FROM {fq_table("memory_units")} - WHERE bank_id = $1 - AND consolidation_failed_at IS NOT NULL - AND fact_type IN ('experience', 'world') - """, - bank_id, - ) - await conn.execute( - f""" - UPDATE {fq_table("memory_units")} - SET consolidation_failed_at = NULL, consolidated_at = NULL - WHERE bank_id = $1 - AND consolidation_failed_at IS NOT NULL - AND fact_type IN ('experience', 'world') - """, - bank_id, - ) + if store.writes_memory_rows_in_sql: + count = await conn.fetchval( + f""" + SELECT COUNT(*) FROM {fq_table("memory_units")} + WHERE bank_id = $1 + AND consolidation_failed_at IS NOT NULL + AND fact_type IN ('experience', 'world') + """, + bank_id, + ) + await conn.execute( + f""" + UPDATE {fq_table("memory_units")} + SET consolidation_failed_at = NULL, consolidated_at = NULL + WHERE bank_id = $1 + AND consolidation_failed_at IS NOT NULL + AND fact_type IN ('experience', 'world') + """, + bank_id, + ) + else: + # A store that keeps the failure marker on the memory: find the failed sources and + # requeue them. mark_consolidated(when=None) clears BOTH the failed and consolidated + # markers and returns the memory to the not-yet-consolidated state. + failed = await store.find_failed_consolidation(conn=conn, fq_table=fq_table, bank_id=bank_id) + count = len(failed) + if failed: + await store.mark_consolidated( + conn=conn, fq_table=fq_table, bank_id=bank_id, unit_ids=[m.unit_id for m in failed], when=None + ) return {"retried_count": count or 0} async def clear_observations_for_memory( @@ -6870,17 +7082,25 @@ async def clear_observations_for_memory( # Also reset this memory's own consolidated_at so it gets re-consolidated # (the memory was a source for the deleted observations, so it needs new ones) if deleted_count > 0: - await conn.execute( - f""" - UPDATE {fq_table("memory_units")} - SET consolidated_at = NULL - WHERE id = $1 - AND bank_id = $2 - AND fact_type IN ('experience', 'world') - """, - uuid_module.UUID(memory_id), - bank_id, - ) + from .memories import get_memories + + _store = get_memories() + if _store.writes_memory_rows_in_sql: + await conn.execute( + f""" + UPDATE {fq_table("memory_units")} + SET consolidated_at = NULL + WHERE id = $1 + AND bank_id = $2 + AND fact_type IN ('experience', 'world') + """, + uuid_module.UUID(memory_id), + bank_id, + ) + else: + await _store.mark_consolidated( + conn=conn, fq_table=fq_table, bank_id=bank_id, unit_ids=[memory_id], when=None + ) if deleted_count > 0: config = await self._config_resolver.resolve_full_config(bank_id, request_context) @@ -6919,21 +7139,6 @@ async def _reembed_memory_text( embeddings = await embedding_processing.generate_embeddings_batch(self.embeddings, augmented) return str(embeddings[0]) if embeddings else None - async def _memory_unit_columns(self, conn) -> str: - """Comma-joined, quoted ordinal column list of ``memory_units``. - - Used to move a row verbatim between ``memory_units`` and the curation - archive (``invalidated_memory_units``) without hardcoding the - migration-evolving column set — the archive is created via - ``LIKE memory_units`` so the lists line up. - """ - rows = await conn.fetch( - f"SELECT a.attname FROM pg_attribute a " - f"WHERE a.attrelid = '{fq_table('memory_units')}'::regclass " - f"AND a.attnum > 0 AND NOT a.attisdropped ORDER BY a.attnum" - ) - return ", ".join(f'"{r["attname"]}"' for r in rows) - @_bind_bank_id() async def update_memory_unit( self, @@ -7031,9 +7236,8 @@ def _parse_edit_date(value: str | None) -> datetime | None: await self._validate_operation(self._operation_validator.validate_bank_write(ctx)) backend = await self._get_backend() - from .causal_links import CAUSAL_LINK_TYPES from .graph_maintenance import enqueue_relink_victims - from .retain.link_utils import rematerialize_causal_links, resolve_entities_only, snapshot_causal_links + from .retain.link_utils import resolve_entities_only # Resolve the bank's entity-label taxonomy once when re-resolving entities, # so corrected entities are matched with the same rules retain uses. @@ -7042,56 +7246,45 @@ def _parse_edit_date(value: str | None) -> datetime | None: edit_config = await self._config_resolver.resolve_full_config(bank_id, request_context) entity_labels = getattr(edit_config, "entity_labels", None) - mu = fq_table("memory_units") - arch = fq_table("invalidated_memory_units") - ue = fq_table("unit_entities") - ml = fq_table("memory_links") - ent = fq_table("entities") - need_consolidation = False need_graph = False found = False async with acquire_with_retry(backend) as conn: async with conn.transaction(): - live = await conn.fetchrow( - f"SELECT text, context, fact_type, event_date, occurred_start, occurred_end, mentioned_at " - f"FROM {mu} WHERE id = $1 AND bank_id = $2", - str(memory_uuid), - bank_id, + from .memories import get_memories + + store = get_memories() + # The store decides existence and drives the state changes, so + # invalidate/revert work whichever store owns the memory. `live` + # is the live record (used for the edit path's fields too); + # `archived` is its counterpart in the curation archive. + live_batch = await store.get_memories( + conn=conn, fq_table=fq_table, bank_id=bank_id, unit_ids=[str(memory_uuid)] ) - archived = None - if not live: - archived = await conn.fetchrow( - f"SELECT fact_type FROM {arch} WHERE id = $1 AND bank_id = $2", - str(memory_uuid), - bank_id, + live = live_batch[0] if live_batch else None + archived = ( + None + if live + else await store.get_archived_memory( + conn=conn, fq_table=fq_table, bank_id=bank_id, unit_id=str(memory_uuid) ) + ) record = live or archived if record is None: return None found = True - current_fact_type = record["fact_type"] + # One cross-store write-group for this curation edit/invalidate/revert: the + # store's writes below (apply_edit + re-embed, or the archive move) are tagged so + # they commit together with this Postgres transaction; decided after it commits. + _curation_txn = await store.begin_txn(conn=conn, fq_table=fq_table, bank_id=bank_id, mutating=True) + current_fact_type = record.fact_type if current_fact_type not in ("experience", "world"): raise ValueError( f"Memory '{memory_id}' is a {current_fact_type}; only world/experience facts can be " "curated. Observations are derived and regenerate from their sources." ) - collist = await self._memory_unit_columns(conn) - # The archive is cold storage, never a recall surface and carries no index, - # so the schema gives it neither the `embedding` (dropped in d4f6a8c2e1b3) - # nor the `search_vector` column (dropped in e7c3a9f1b2d5). Both are - # recall-surface columns whose type/shape follows server - # config, so the move in/out is over every memory_units column EXCEPT those - # two; on revert each is recomputed from the unit's text/dates/entities below. - # This makes a model switch (which re-dimensions memory_units) structurally - # unable to trip a vector-dimension mismatch (#2209), and a text-search backend - # switch unable to trip a search_vector type mismatch (#2503), on the - # INSERT … SELECT round-trip. - _archive_omitted = ('"embedding"', '"search_vector"') - arch_cols = ", ".join(c for c in (s.strip() for s in collist.split(",")) if c not in _archive_omitted) - # --- Edit fields (live rows only): text / context / dates / fact_type / entities --- doing_edit = any( v is not None for v in (text, context, occurred_start, occurred_end, new_fact_type) @@ -7099,24 +7292,25 @@ def _parse_edit_date(value: str | None) -> datetime | None: if doing_edit: if not live: raise ValueError("Cannot edit an invalidated memory; revert it to 'valid' first.") - new_text = text if text is not None else live["text"] - new_context = (context or None) if context is not None else live["context"] - new_fact = new_fact_type if new_fact_type is not None else live["fact_type"] + new_text = text if text is not None else live.text + new_context = (context or None) if context is not None else live.context + new_fact = new_fact_type if new_fact_type is not None else live.fact_type new_occ_start = ( - _parse_edit_date(occurred_start) if occurred_start is not None else live["occurred_start"] + _parse_edit_date(occurred_start) if occurred_start is not None else live.occurred_start ) - new_occ_end = _parse_edit_date(occurred_end) if occurred_end is not None else live["occurred_end"] + new_occ_end = _parse_edit_date(occurred_end) if occurred_end is not None else live.occurred_end # event_date (NOT NULL, legacy single date + used by temporal links) # tracks the occurred start when it's set. - new_event_date = new_occ_start or live["event_date"] + new_event_date = new_occ_start or live.event_date # Rebuild the unit's entity set FIRST, so the re-embed below picks # up the corrected canonical names. Reuses retain's resolver # (find-or-create + cooccurrence) rather than touching entities # directly. Orphaned entities + stale cooccurrence are swept by # the graph-maintenance run this edit submits. + edit_entity_ids = None if new_entities is not None: - entity_date = new_occ_start or live["mentioned_at"] + entity_date = new_occ_start or live.mentioned_at entity_resolution = await resolve_entities_only( self.entity_resolver, conn, @@ -7128,8 +7322,13 @@ def _parse_edit_date(value: str | None) -> datetime | None: [[{"text": name, "type": "CONCEPT"} for name in new_entities]], entity_labels=entity_labels, ) - await conn.execute(f"DELETE FROM {ue} WHERE unit_id = $1", str(memory_uuid)) + # Clear the old postings before re-linking — a no-op for a + # store that carries entity ids on the memory. + await store.clear_unit_entities( + conn=conn, fq_table=fq_table, bank_id=bank_id, unit_id=str(memory_uuid) + ) resolved_for_unit = entity_resolution.unit_to_entity_ids.get(str(memory_uuid), []) + edit_entity_ids = [str(eid) for eid in resolved_for_unit] if resolved_for_unit: # Same prune race as retain (#2662): a found (not newly # created) parent can be deleted by graph maintenance @@ -7141,238 +7340,122 @@ def _parse_edit_date(value: str | None) -> datetime | None: await self.entity_resolver.link_units_to_entities_batch( [(str(memory_uuid), eid, entity_date) for eid in resolved_for_unit], conn=conn, + bank_id=bank_id, ) - ent_rows = await conn.fetch( - f"SELECT e.canonical_name FROM {ue} ue JOIN {ent} e ON ue.entity_id = e.id " - f"WHERE ue.unit_id = $1", - str(memory_uuid), + # Capture relink victims before this memory's links change, then + # apply the field edit + new entity set through the store: it + # resets consolidation, stamps the edit, and drops the derived + # links. The embedding is written separately, after the re-embed. + # The edit leaves the unit live but drops its derived links, so it needs its own + # outgoing adjacency rebuilt too — one combined insert with its victims (#2864). + await enqueue_relink_victims( + conn, bank_id, [memory_id], ops=backend.ops, include_affected_units=True ) - new_emb = await self._reembed_memory_text( + await store.apply_edit( + conn=conn, + fq_table=fq_table, + bank_id=bank_id, + unit_id=str(memory_uuid), text=new_text, + context=new_context, + fact_type=new_fact, occurred_start=new_occ_start, occurred_end=new_occ_end, - mentioned_at=live["mentioned_at"], - entities=[r["canonical_name"] for r in ent_rows], - ) - # Keep the stored text-search vector in sync with curated - # text/context edits. Use the incoming parameters here: - # PostgreSQL evaluates UPDATE RHS expressions before the - # sibling SET assignments take effect, so column references - # would see the pre-edit text/context. - from .db.ops_postgresql import pg_search_vector_expr - - sv_expr = pg_search_vector_expr(get_config(), text_col="$3", context_col="$4") - search_vector_clause = ( - f",\n search_vector = {sv_expr}" if sv_expr else "" - ) - # The DELETE below drops this unit's incident DERIVED edges (causal - # ones are preserved, #2864), so the unit itself needs its outgoing - # temporal/semantic adjacency rebuilt — not just the neighbours that - # lost an edge to it. Skip that when this same call also invalidates - # the unit (the block below archives it, and the drain no-ops on a - # queue row with no live unit). - # - # Victims and the edited unit go in ONE insert: the queue insert - # sorts its ids, so a single call keeps the (bank_id, unit_id) lock - # order global. Two separate inserts can deadlock when concurrently - # edited units point at each other. - await enqueue_relink_victims( - conn, - bank_id, - [memory_id], - ops=backend.ops, - include_affected_units=state != "invalidated", + event_date=new_event_date, + mentioned_at=live.mentioned_at, + entity_ids=edit_entity_ids, + txn=_curation_txn, ) - await conn.execute( - f""" - UPDATE {mu} - SET text = $3, context = $4, fact_type = $5, occurred_start = $6, - occurred_end = $7, event_date = $8, embedding = $9::vector, - consolidated_at = NULL, consolidation_failed_at = NULL, - edited_at = now(), updated_at = now(){search_vector_clause} - WHERE id = $1 AND bank_id = $2 - """, - str(memory_uuid), - bank_id, - new_text, - new_context, - new_fact, - new_occ_start, - new_occ_end, - new_event_date, - new_emb, + + # Re-embed from the now-current fields + entity names (through the + # store, so a memory carrying its entities inline resolves too). + emap = await store.entity_map_for_units( + conn=conn, fq_table=fq_table, bank_id=bank_id, unit_ids=[str(memory_uuid)] ) - # Drop only the DERIVED edges — graph maintenance (submitted - # below) recomputes temporal/semantic from the edited dates - # and embedding. Causal edges are retain-time extraction - # output that nothing recreates, so an edit preserves them - # rather than silently destroying them (#2864); a corrected - # fact keeps the causality the extractor asserted for it. - await conn.execute( - f"DELETE FROM {ml} WHERE (from_unit_id = $1 OR to_unit_id = $1) " - f"AND NOT (link_type = ANY($2::text[]))", - str(memory_uuid), - list(CAUSAL_LINK_TYPES), + names = [e["canonical_name"] for e in emap.get(str(memory_uuid), [])] + new_emb = await self._reembed_memory_text( + text=new_text, + occurred_start=new_occ_start, + occurred_end=new_occ_end, + mentioned_at=live.mentioned_at, + entities=names, ) + if new_emb is not None: + await store.set_memory_embedding( + conn=conn, + fq_table=fq_table, + bank_id=bank_id, + unit_id=str(memory_uuid), + embedding=new_emb, + txn=_curation_txn, + ) await self._delete_stale_observations_for_memories(conn, bank_id, [memory_id]) need_consolidation = True need_graph = True # --- Invalidate: move live → archive --- if state == "invalidated" and live: - entity_ids = [ - r["entity_id"] - for r in await conn.fetch(f"SELECT entity_id FROM {ue} WHERE unit_id = $1", str(memory_uuid)) - ] - # Capture relink victims BEFORE the row (and its links) disappear. + # Capture relink victims before the row (and its links) go. await enqueue_relink_victims(conn, bank_id, [memory_id], ops=backend.ops) - # Same for the causal edges: temporal/semantic are recomputed - # by graph maintenance on revert, but causal edges are - # retain-time extraction output the cascade would destroy for - # good. Park their descriptors on the archive row (#2864). - causal_links = await snapshot_causal_links(conn, bank_id, str(memory_uuid)) - await conn.execute( - f"INSERT INTO {arch} ({arch_cols}, invalidation_reason, invalidated_at, entity_ids, " - f"causal_links) " - f"SELECT {arch_cols}, $2, now(), $3::uuid[], $5::jsonb FROM {mu} WHERE id = $1 AND bank_id = $4", - str(memory_uuid), - reason, - entity_ids, - bank_id, - json.dumps([descriptor.as_json_dict() for descriptor in causal_links]), + await store.invalidate_memory( + conn=conn, + fq_table=fq_table, + bank_id=bank_id, + unit_id=str(memory_uuid), + reason=reason, + txn=_curation_txn, ) - # Cascade prunes unit_entities + memory_links; sweep runs after - # the delete so it also catches a racing observation insert. - await conn.execute(f"DELETE FROM {mu} WHERE id = $1 AND bank_id = $2", str(memory_uuid), bank_id) + # Sweep after the move, so a racing observation insert is caught too. await self._delete_stale_observations_for_memories(conn, bank_id, [memory_id]) need_consolidation = True need_graph = True elif state == "invalidated" and archived and reason is not None: # Already archived — just update the recorded reason. - await conn.execute( - f"UPDATE {arch} SET invalidation_reason = $3 WHERE id = $1 AND bank_id = $2", - str(memory_uuid), - bank_id, - reason, + await store.set_invalidation_reason( + conn=conn, fq_table=fq_table, bank_id=bank_id, unit_id=str(memory_uuid), reason=reason ) # --- Revert: move archive → live --- elif state == "valid" and archived: - arch_row = await conn.fetchrow( - f"SELECT entity_ids, causal_links FROM {arch} WHERE id = $1 AND bank_id = $2", - str(memory_uuid), - bank_id, - ) - # The archive keeps neither embedding nor search_vector (see arch_cols - # above), so both default to NULL on the way back and are recomputed here: - # the embedding below once entities are restored, the search_vector now - # from the row's own text/context/text_signals. - await conn.execute( - f"INSERT INTO {mu} ({arch_cols}) SELECT {arch_cols} FROM {arch} WHERE id = $1 AND bank_id = $2", - str(memory_uuid), - bank_id, - ) - # Rebuild search_vector using the *current* text-search backend, so the - # reverted unit is keyword-searchable again (more correct than carrying a - # verbatim copy that could be stale/wrong-type if the backend changed while - # the fact sat archived). None = pgroonga/pg_textsearch/pg_search, which - # index base columns directly and leave search_vector empty (#2503). - from .db.ops_postgresql import pg_search_vector_expr - - sv_expr = pg_search_vector_expr(get_config()) - if sv_expr is not None: - await conn.execute( - f"UPDATE {mu} SET search_vector = {sv_expr} WHERE id = $1 AND bank_id = $2", - str(memory_uuid), - bank_id, - ) - # Re-consolidate from scratch; links are rebuilt by graph maintenance. - await conn.execute( - f"UPDATE {mu} SET consolidated_at = NULL, consolidation_failed_at = NULL, updated_at = now() " - f"WHERE id = $1 AND bank_id = $2", - str(memory_uuid), - bank_id, - ) - # Restore entity associations for entities that still exist (some may - # have been pruned as orphans after the original move). - if arch_row and arch_row["entity_ids"]: - await conn.execute( - f"INSERT INTO {ue} (unit_id, entity_id) " - f"SELECT $1, eid FROM unnest($2::uuid[]) AS eid " - f"WHERE EXISTS (SELECT 1 FROM {ent} e WHERE e.id = eid AND e.bank_id = $3) " - f"ON CONFLICT DO NOTHING", - str(memory_uuid), - arch_row["entity_ids"], - bank_id, - ) - # Rematerialize the causal edges parked at invalidation time. - # Edges whose peer is still archived (or was permanently - # deleted) are skipped — the peer keeps its own copy of the - # descriptor and recreates the edge when it reverts, so the - # restore is order-independent and idempotent. - if arch_row and arch_row["causal_links"]: - await rematerialize_causal_links( - conn, - bank_id, - conn.parse_json(arch_row["causal_links"]) or [], - ops=backend.ops, - ) - # Recompute the embedding (the archive doesn't keep one) so the reverted - # unit is searchable again, using the now-current model's dimension and the - # restored entity set — mirroring how an edit re-embeds. - reverted = await conn.fetchrow( - f"SELECT text, occurred_start, occurred_end, mentioned_at FROM {mu} " - f"WHERE id = $1 AND bank_id = $2", - str(memory_uuid), - bank_id, + restored = await store.restore_memory( + conn=conn, fq_table=fq_table, bank_id=bank_id, unit_id=str(memory_uuid), txn=_curation_txn ) - if reverted: - ent_rows = await conn.fetch( - f"SELECT e.canonical_name FROM {ue} ue JOIN {ent} e ON ue.entity_id = e.id " - f"WHERE ue.unit_id = $1", - str(memory_uuid), + if restored is not None: + # Recompute the embedding — the archive need not keep one — from the + # restored fields and current entity names, with the now-current model + # (the same re-embed an edit does). Names come through the store, so a + # memory whose entities ride on it rather than on a join table resolves too. + emap = await store.entity_map_for_units( + conn=conn, fq_table=fq_table, bank_id=bank_id, unit_ids=[str(memory_uuid)] ) + names = [e["canonical_name"] for e in emap.get(str(memory_uuid), [])] new_emb = await self._reembed_memory_text( - text=reverted["text"], - occurred_start=reverted["occurred_start"], - occurred_end=reverted["occurred_end"], - mentioned_at=reverted["mentioned_at"], - entities=[r["canonical_name"] for r in ent_rows], + text=restored.text, + occurred_start=restored.occurred_start, + occurred_end=restored.occurred_end, + mentioned_at=restored.mentioned_at, + entities=names, ) if new_emb is not None: - await conn.execute( - f"UPDATE {mu} SET embedding = $3::vector WHERE id = $1 AND bank_id = $2", - str(memory_uuid), - bank_id, - new_emb, + await store.set_memory_embedding( + conn=conn, + fq_table=fq_table, + bank_id=bank_id, + unit_id=str(memory_uuid), + embedding=new_emb, + txn=_curation_txn, ) - # Invalidation cascaded away every link incident to this unit. The - # causal ones came back from the archive snapshot above (#2864); the - # derived ones are graph maintenance's job, and it only rebuilds units - # present in the queue — it never scans memory_units for missing - # adjacency. Without this enqueue the submission below short-circuits - # on an empty queue (no_work) and the reverted fact stays off the - # temporal/semantic graph. - # Enqueued last, so the row is fully searchable (text, entities, - # search_vector, embedding) before the drain reads it, and atomic - # with the archive→live move: a rollback takes the work item too. - # Scope: this rebuilds the reverted unit's OUTGOING links. Units - # that pointed at it were relinked elsewhere at invalidation time - # and are not re-queued here. - await backend.ops.enqueue_graph_maintenance( - conn, - fq_table("graph_maintenance_queue"), - bank_id, - [memory_uuid], - ) - await conn.execute(f"DELETE FROM {arch} WHERE id = $1 AND bank_id = $2", str(memory_uuid), bank_id) need_consolidation = True need_graph = True if not found: return None + # Postgres committed the curation change: publish the store's write-group. On a crash + # before here the writes stay invisible and the recovery sweep resolves them (spec §5). + await store.decide_txn(_curation_txn, commit=True) + if need_consolidation: config = await self._config_resolver.resolve_full_config(bank_id, request_context) if config.enable_auto_consolidation: @@ -7468,87 +7551,27 @@ async def get_graph_data( bank_id=bank_id, operation=BankReadOperation.GET_GRAPH_DATA, request_context=request_context ) await self._validate_operation(self._operation_validator.validate_bank_read(ctx)) + from .memories import get_memories + + store = get_memories() backend = await self._get_backend() async with acquire_with_retry(backend) as conn: - # Get memory units, optionally filtered by bank_id and fact_type - query_conditions = [] - query_params = [] - param_count = 0 - - bank_id_placeholder: int | None = None - if bank_id: - param_count += 1 - bank_id_placeholder = param_count - query_conditions.append(f"bank_id = ${param_count}") - query_params.append(bank_id) - - if fact_type: - param_count += 1 - query_conditions.append(f"fact_type = ${param_count}") - query_params.append(fact_type) - - if document_id: - param_count += 1 - obs_match = self._observations_via_source_match_sql( - "document_id", source_placeholder=param_count, bank_placeholder=bank_id_placeholder - ) - query_conditions.append( - f"(document_id = ${param_count} OR (fact_type = 'observation' AND {obs_match}))" - ) - query_params.append(document_id) - - if chunk_id: - param_count += 1 - obs_match = self._observations_via_source_match_sql( - "chunk_id", source_placeholder=param_count, bank_placeholder=bank_id_placeholder - ) - query_conditions.append(f"(chunk_id = ${param_count} OR (fact_type = 'observation' AND {obs_match}))") - query_params.append(chunk_id) - - if q: - param_count += 1 - query_conditions.append(f"(text ILIKE ${param_count} OR context ILIKE ${param_count})") - query_params.append(f"%{q}%") - - if tags: - from .search.tags import build_tags_where_clause_simple - - tag_clause = build_tags_where_clause_simple(tags, param_count + 1, match=tags_match) - if tag_clause: - query_conditions.append(tag_clause.removeprefix("AND ")) - param_count += 1 - query_params.append(tags) - elif tags_match == "exact": - # Exact match with no tags is the "global" scope: rows that carry no - # tags at all. (Other match modes treat empty tags as "no filter".) - query_conditions.append("(tags IS NULL OR tags = '{}')") - - where_clause = "WHERE " + " AND ".join(query_conditions) if query_conditions else "" - - # Get total count first - total_count_result = await conn.fetchrow( - f""" - SELECT COUNT(*) as total - FROM {fq_table("memory_units")} - {where_clause} - """, - *query_params, - ) - total_count = total_count_result["total"] if total_count_result else 0 - - # Get units with limit - param_count += 1 - units = await conn.fetch( - f""" - SELECT id, text, event_date, context, occurred_start, occurred_end, mentioned_at, document_id, chunk_id, fact_type, tags, created_at, proof_count, source_memory_ids - FROM {fq_table("memory_units")} - {where_clause} - ORDER BY mentioned_at DESC NULLS LAST, event_date DESC - LIMIT ${param_count} - """, - *query_params, - limit, - ) + # The nodes, and how many match the filters, come from the store — it + # is the one that knows where the memories live and how to page them. + page = await store.graph_units( + conn=conn, + fq_table=fq_table, + bank_id=bank_id, + fact_type=fact_type, + search_query=q, + document_id=document_id, + chunk_id=chunk_id, + tags=tags, + tags_match=tags_match, + limit=limit, + ) + units = page["units"] + total_count = page["total"] # Get links, filtering to only include links between units of the selected agent # Use DISTINCT ON with LEAST/GREATEST to deduplicate bidirectional links @@ -7568,27 +7591,10 @@ async def get_graph_data( # e9b2c7d1f3a4) — no link_type filter is needed. # Cap at 10k edges — the UI can't usefully render more, and uncapped queries # on highly-connected graphs (e.g. 1000 nodes with 500k+ edges) are too slow. - max_edges = 10000 all_relevant_ids = unit_ids + source_memory_ids - if all_relevant_ids: - links = await conn.fetch( - f""" - SELECT ml.from_unit_id, - ml.to_unit_id, - ml.link_type, - ml.weight, - NULL::text AS entity_name - FROM {fq_table("memory_links")} ml - WHERE ml.from_unit_id = ANY($1::uuid[]) - AND ml.to_unit_id = ANY($1::uuid[]) - ORDER BY ml.weight DESC NULLS LAST - LIMIT $2 - """, - all_relevant_ids, - max_edges, - ) - else: - links = [] + links = await store.graph_direct_links( + conn=conn, fq_table=fq_table, bank_id=bank_id, unit_ids=[str(u) for u in all_relevant_ids] + ) # Copy links from source memories to observations # Observations inherit links from their source memories via source_memory_ids @@ -7650,19 +7656,9 @@ async def get_graph_data( # Fetch entities for visible units AND their source memories # (so observations can inherit entities from source memories) entity_lookup_ids = unit_ids + source_memory_ids - if entity_lookup_ids: - unit_entities = await conn.fetch( - f""" - SELECT ue.unit_id, e.canonical_name - FROM {fq_table("unit_entities")} ue - JOIN {fq_table("entities")} e ON ue.entity_id = e.id - WHERE ue.unit_id = ANY($1::uuid[]) - ORDER BY ue.unit_id - """, - entity_lookup_ids, - ) - else: - unit_entities = [] + unit_entities = await store.graph_entity_rows( + conn=conn, fq_table=fq_table, bank_id=bank_id, unit_ids=[str(u) for u in entity_lookup_ids] + ) # Build entity mapping entity_map = {} @@ -7998,189 +7994,31 @@ async def list_memory_units( bank_id=bank_id, operation=BankReadOperation.LIST_MEMORY_UNITS, request_context=request_context ) await self._validate_operation(self._operation_validator.validate_bank_read(ctx)) - if state is not None and state not in ("valid", "invalidated"): - raise ValueError(f"Invalid state '{state}': expected 'valid' or 'invalidated'.") - if entity_id is not None: - try: - uuid.UUID(entity_id) - except ValueError: - raise ValueError(f"Invalid entity_id: '{entity_id}' is not a valid UUID") - # Invalidated facts live in a separate archive table; pick the source - # accordingly. Default (state is None) lists live facts. - is_archived = state == "invalidated" - source_table = fq_table("invalidated_memory_units") if is_archived else fq_table("memory_units") + from .memories import get_memories + backend = await self._get_backend() async with acquire_with_retry(backend) as conn: - # Build query conditions - query_conditions = [] - query_params = [] - param_count = 0 - - if bank_id: - param_count += 1 - query_conditions.append(f"bank_id = ${param_count}") - query_params.append(bank_id) - - if fact_type: - param_count += 1 - query_conditions.append(f"fact_type = ${param_count}") - query_params.append(fact_type) - - if document_id: - param_count += 1 - query_conditions.append(f"document_id = ${param_count}") - query_params.append(document_id) - - if entity_id: - # Reverse lookup via the stored entity links. Entity links only - # reference live memory units, so this yields nothing against the - # invalidated archive (documented on the method). The - # idx_unit_entities_entity_unit index covers this subquery. - param_count += 1 - query_conditions.append( - f"id IN (SELECT unit_id FROM {fq_table('unit_entities')} WHERE entity_id = ${param_count}::uuid)" - ) - query_params.append(entity_id) - - if search_query: - # Full-text search on text and context fields using ILIKE - param_count += 1 - query_conditions.append(f"(text ILIKE ${param_count} OR context ILIKE ${param_count})") - query_params.append(f"%{search_query}%") - - if consolidation_state: - state = consolidation_state.lower() - if state == "failed": - query_conditions.append( - "consolidation_failed_at IS NOT NULL AND fact_type IN ('experience', 'world')" - ) - elif state == "pending": - query_conditions.append( - "consolidated_at IS NULL AND consolidation_failed_at IS NULL " - "AND fact_type IN ('experience', 'world')" - ) - elif state == "done": - query_conditions.append("consolidated_at IS NOT NULL AND fact_type IN ('experience', 'world')") - else: - raise ValueError( - f"Invalid consolidation_state '{consolidation_state}': expected 'failed', 'pending', or 'done'." - ) - - if tags: - tags_clause, tags_params, next_param = build_tags_where_clause(tags, param_count + 1, "", tags_match) - if tags_clause: - query_conditions.append(tags_clause.removeprefix("AND ")) - query_params.extend(tags_params) - param_count = next_param - 1 - elif tags_match == "exact": - # Exact match with no tags is the "global" scope: rows that carry no - # tags at all. (Other match modes treat empty tags as "no filter".) - query_conditions.append("(tags IS NULL OR tags = '{}')") - - if created_before is not None: - param_count += 1 - query_conditions.append(f"created_at < ${param_count}") - query_params.append(created_before) - - where_clause = "WHERE " + " AND ".join(query_conditions) if query_conditions else "" - - # Get total count - count_query = f""" - SELECT COUNT(*) as total - FROM {source_table} - {where_clause} - """ - count_result = await conn.fetchrow(count_query, *query_params) - total = count_result["total"] - - # Get units with limit and offset - param_count += 1 - limit_param = f"${param_count}" - query_params.append(limit) - - param_count += 1 - offset_param = f"${param_count}" - query_params.append(offset) - - # The archive carries invalidation bookkeeping; the live table doesn't. - curation_cols = ( - "invalidation_reason, invalidated_at" - if is_archived - else "NULL::text AS invalidation_reason, NULL::timestamptz AS invalidated_at" - ) - units = await conn.fetch( - f""" - SELECT id, text, event_date, context, fact_type, document_id, - mentioned_at, occurred_start, occurred_end, chunk_id, proof_count, - tags, metadata, consolidated_at, consolidation_failed_at, edited_at, {curation_cols} - FROM {source_table} - {where_clause} - ORDER BY mentioned_at DESC NULLS LAST, created_at DESC - LIMIT {limit_param} OFFSET {offset_param} - """, - *query_params, + # The memories store owns the list — same page shape wherever the + # memories live. `state` still selects the live vs invalidated view; + # the store validates and resolves it. + return await get_memories().list_memory_units( + conn=conn, + ops=self._backend.ops, + fq_table=fq_table, + bank_id=bank_id, + fact_type=fact_type, + search_query=search_query, + consolidation_state=consolidation_state, + state=state, + document_id=document_id, + entity_id=entity_id, + tags=tags, + tags_match=tags_match, + created_before=created_before, + limit=limit, + offset=offset, ) - # Get entity information for these units - if units: - unit_ids = [row["id"] for row in units] - unit_entities = await conn.fetch( - f""" - SELECT ue.unit_id, e.canonical_name - FROM {fq_table("unit_entities")} ue - JOIN {fq_table("entities")} e ON ue.entity_id = e.id - WHERE ue.unit_id = ANY($1::uuid[]) - ORDER BY ue.unit_id - """, - unit_ids, - ) - else: - unit_entities = [] - - # Build entity mapping - entity_map = {} - for row in unit_entities: - unit_id = row["unit_id"] - entity_name = row["canonical_name"] - if unit_id not in entity_map: - entity_map[unit_id] = [] - entity_map[unit_id].append(entity_name) - - # Build result items - items = [] - for row in units: - unit_id = row["id"] - entities = entity_map.get(unit_id, []) - - items.append( - { - "id": str(unit_id), - "text": row["text"], - "context": row["context"] if row["context"] else "", - "date": row["event_date"].isoformat() if row["event_date"] else "", - "fact_type": row["fact_type"], - "document_id": row["document_id"], - "mentioned_at": row["mentioned_at"].isoformat() if row["mentioned_at"] else None, - "occurred_start": row["occurred_start"].isoformat() if row["occurred_start"] else None, - "occurred_end": row["occurred_end"].isoformat() if row["occurred_end"] else None, - "entities": ", ".join(entities) if entities else "", - "chunk_id": row["chunk_id"] if row["chunk_id"] else None, - "proof_count": row["proof_count"] if row["proof_count"] is not None else 1, - "tags": list(row["tags"]) if row["tags"] else [], - "metadata": conn.parse_json(row["metadata"]) if row["metadata"] is not None else {}, - "consolidated_at": row["consolidated_at"].isoformat() if row["consolidated_at"] else None, - "consolidation_failed_at": ( - row["consolidation_failed_at"].isoformat() if row["consolidation_failed_at"] else None - ), - "state": "invalidated" if is_archived else "valid", - "invalidation_reason": row["invalidation_reason"], - "invalidated_at": row["invalidated_at"].isoformat() if row["invalidated_at"] else None, - "edited_at": row["edited_at"].isoformat() if row["edited_at"] else None, - } - ) - - return {"items": items, "total": total, "limit": limit, "offset": offset} - async def get_memory_unit( self, bank_id: str, @@ -8213,99 +8051,161 @@ async def get_memory_unit( bank_id=bank_id, operation=BankReadOperation.GET_MEMORY_UNIT, request_context=request_context ) await self._validate_operation(self._operation_validator.validate_bank_read(ctx)) + from .memories import get_memories + backend = await self._get_backend() async with acquire_with_retry(backend) as conn: - # Get the memory unit (include source_memory_ids for mental models). - # Curation moves invalidated facts to invalidated_memory_units, so fall - # back to the archive (with its invalidation bookkeeping) on a miss. - select_cols = ( - "id, text, context, event_date, occurred_start, occurred_end, " - "mentioned_at, fact_type, document_id, chunk_id, tags, metadata, source_memory_ids, " - "observation_scopes, edited_at" + # The store renders the detail view — including the observation + # history and source facts it folds in — for a normalized id. + return await get_memories().get_memory_unit( + conn=conn, + ops=self._backend.ops, + fq_table=fq_table, + bank_id=bank_id, + unit_id=str(memory_uuid), ) - row = await conn.fetchrow( - f"SELECT {select_cols}, NULL::text AS invalidation_reason, NULL::timestamptz AS invalidated_at " - f"FROM {fq_table('memory_units')} WHERE id = $1 AND bank_id = $2", - str(memory_uuid), - bank_id, + + async def list_documents( + self, + bank_id: str, + *, + search_query: str | None = None, + tags: list[str] | None = None, + tags_match: "TagsMatch" = "any_strict", + limit: int = 100, + offset: int = 0, + request_context: "RequestContext", + ): + """ + List documents with optional search and pagination. + + Args: + bank_id: bank ID (required) + search_query: Search in document ID + tags: Filter by tags + tags_match: How to match tags (any, all, any_strict, all_strict) + limit: Maximum number of results + offset: Offset for pagination + request_context: Request context for authentication. + + Returns: + Dict with items (list of documents without original_text) and total count + """ + await self._authenticate_tenant(request_context) + if self._operation_validator: + from hindsight_api.extensions import BankReadContext, BankReadOperation + + ctx = BankReadContext( + bank_id=bank_id, operation=BankReadOperation.LIST_DOCUMENTS, request_context=request_context ) - unit_state = "valid" - if not row: - row = await conn.fetchrow( - f"SELECT {select_cols}, invalidation_reason, invalidated_at " - f"FROM {fq_table('invalidated_memory_units')} WHERE id = $1 AND bank_id = $2", - str(memory_uuid), - bank_id, - ) - unit_state = "invalidated" + await self._validate_operation(self._operation_validator.validate_bank_read(ctx)) + backend = await self._get_backend() + async with acquire_with_retry(backend) as conn: + # Build query conditions + query_conditions = [] + query_params = [] + param_count = 0 - if not row: - return None + param_count += 1 + query_conditions.append(f"bank_id = ${param_count}") + query_params.append(bank_id) - # Get entity information. _entity_rows_for_units_sql handles the - # observation→source_memory_ids inheritance fallback in SQL, so a - # single query covers direct rows and inherited ones. - entities_rows = await conn.fetch( - self._entity_rows_for_units_sql(unit_ids_placeholder=1), - [row["id"]], - ) - entities = [r["canonical_name"] for r in entities_rows] - - result = { - "id": str(row["id"]), - "text": row["text"], - "context": row["context"] if row["context"] else "", - "date": row["event_date"].isoformat() if row["event_date"] else "", - "type": row["fact_type"], - "mentioned_at": row["mentioned_at"].isoformat() if row["mentioned_at"] else None, - "occurred_start": row["occurred_start"].isoformat() if row["occurred_start"] else None, - "occurred_end": row["occurred_end"].isoformat() if row["occurred_end"] else None, - "entities": entities, - "document_id": row["document_id"] if row["document_id"] else None, - "chunk_id": str(row["chunk_id"]) if row["chunk_id"] else None, - "tags": row["tags"] if row["tags"] else [], - "metadata": conn.parse_json(row["metadata"]) if row["metadata"] is not None else {}, - "observation_scopes": ( - conn.parse_json(row["observation_scopes"]) if row["observation_scopes"] is not None else None - ), - "state": unit_state, - "invalidation_reason": row["invalidation_reason"], - "invalidated_at": row["invalidated_at"].isoformat() if row["invalidated_at"] else None, - "edited_at": row["edited_at"].isoformat() if row["edited_at"] else None, - } + if search_query: + # Search in document ID + param_count += 1 + query_conditions.append(f"id ILIKE ${param_count}") + query_params.append(f"%{search_query}%") + + tags_clause, tags_params, next_param = build_tags_where_clause( + tags, param_offset=param_count + 1, match=tags_match + ) + query_params.extend(tags_params) + param_count = next_param - 1 # next_param is next available; convert to last used + + where_clause = "WHERE " + " AND ".join(query_conditions) if query_conditions else "" + if tags_clause: + # tags_clause starts with "AND", append after WHERE conditions + where_clause = where_clause + " " + tags_clause if where_clause else "WHERE " + tags_clause[4:].lstrip() - # For observations, include source_memory_ids - # history is deprecated here - use GET /memories/{id}/history instead - if row["fact_type"] == "observation": - result["history"] = [] + # Get total count + count_query = f""" + SELECT COUNT(*) as total + FROM {fq_table("documents")} + {where_clause} + """ + count_result = await conn.fetchrow(count_query, *query_params) + total = count_result["total"] - if row["fact_type"] == "observation" and row["source_memory_ids"]: - source_ids = row["source_memory_ids"] - result["source_memory_ids"] = [str(sid) for sid in source_ids] + # Get documents with limit and offset (without original_text for performance) + param_count += 1 + limit_param = f"${param_count}" + query_params.append(limit) - # Fetch source memories - source_rows = await conn.fetch( - f""" - SELECT id, text, fact_type, context, occurred_start, mentioned_at - FROM {fq_table("memory_units")} - WHERE id = ANY($1::uuid[]) - ORDER BY mentioned_at DESC NULLS LAST - """, - source_ids, + param_count += 1 + offset_param = f"${param_count}" + query_params.append(offset) + + documents = await conn.fetch( + f""" + SELECT + id, + bank_id, + content_hash, + created_at, + updated_at, + LENGTH(original_text) as text_length, + retain_params, + tags + FROM {fq_table("documents")} + {where_clause} + ORDER BY created_at DESC + LIMIT {limit_param} OFFSET {offset_param} + """, + *query_params, + ) + + # Memory count per document — through the store, so a store that keeps + # its memories elsewhere answers it too (this page reports 0 otherwise). + from .memories import get_memories + + doc_ids = [row["id"] for row in documents] + per_doc = ( + await get_memories().document_memory_counts( + conn=conn, fq_table=fq_table, bank_id=bank_id, document_ids=doc_ids ) - result["source_memories"] = [ + if doc_ids + else {} + ) + count_map = {(doc_id, bank_id): count for doc_id, count in per_doc.items()} + + # Build result items + items = [] + for row in documents: + doc_id = row["id"] + bank_id_val = row["bank_id"] + unit_count = count_map.get((doc_id, bank_id_val), 0) + + retain_params_val = conn.parse_json(row["retain_params"]) + + # document_metadata is sourced from retain_params.metadata + document_metadata = retain_params_val.get("metadata") if retain_params_val else None + + items.append( { - "id": str(r["id"]), - "text": r["text"], - "type": r["fact_type"], - "context": r["context"], - "occurred_start": r["occurred_start"].isoformat() if r["occurred_start"] else None, - "mentioned_at": r["mentioned_at"].isoformat() if r["mentioned_at"] else None, + "id": doc_id, + "bank_id": bank_id_val, + "content_hash": row["content_hash"], + "created_at": row["created_at"].isoformat() if row["created_at"] else "", + "updated_at": row["updated_at"].isoformat() if row["updated_at"] else "", + "text_length": row["text_length"] or 0, + "memory_unit_count": unit_count, + "retain_params": retain_params_val or None, + "document_metadata": document_metadata or None, + "tags": row["tags"] if row["tags"] else [], } - for r in source_rows - ] + ) - return result + return {"items": items, "total": total, "limit": limit, "offset": offset} async def get_observation_history( self, @@ -8433,164 +8333,6 @@ def _as_list(v: Any) -> list: enriched.reverse() return enriched - async def list_documents( - self, - bank_id: str, - *, - search_query: str | None = None, - tags: list[str] | None = None, - tags_match: "TagsMatch" = "any_strict", - limit: int = 100, - offset: int = 0, - request_context: "RequestContext", - ): - """ - List documents with optional search and pagination. - - Args: - bank_id: bank ID (required) - search_query: Search in document ID - tags: Filter by tags - tags_match: How to match tags (any, all, any_strict, all_strict) - limit: Maximum number of results - offset: Offset for pagination - request_context: Request context for authentication. - - Returns: - Dict with items (list of documents without original_text) and total count - """ - await self._authenticate_tenant(request_context) - if self._operation_validator: - from hindsight_api.extensions import BankReadContext, BankReadOperation - - ctx = BankReadContext( - bank_id=bank_id, operation=BankReadOperation.LIST_DOCUMENTS, request_context=request_context - ) - await self._validate_operation(self._operation_validator.validate_bank_read(ctx)) - backend = await self._get_backend() - async with acquire_with_retry(backend) as conn: - # Build query conditions - query_conditions = [] - query_params = [] - param_count = 0 - - param_count += 1 - query_conditions.append(f"bank_id = ${param_count}") - query_params.append(bank_id) - - if search_query: - # Search in document ID - param_count += 1 - query_conditions.append(f"id ILIKE ${param_count}") - query_params.append(f"%{search_query}%") - - tags_clause, tags_params, next_param = build_tags_where_clause( - tags, param_offset=param_count + 1, match=tags_match - ) - query_params.extend(tags_params) - param_count = next_param - 1 # next_param is next available; convert to last used - - where_clause = "WHERE " + " AND ".join(query_conditions) if query_conditions else "" - if tags_clause: - # tags_clause starts with "AND", append after WHERE conditions - where_clause = where_clause + " " + tags_clause if where_clause else "WHERE " + tags_clause[4:].lstrip() - - # Get total count - count_query = f""" - SELECT COUNT(*) as total - FROM {fq_table("documents")} - {where_clause} - """ - count_result = await conn.fetchrow(count_query, *query_params) - total = count_result["total"] - - # Get documents with limit and offset (without original_text for performance) - param_count += 1 - limit_param = f"${param_count}" - query_params.append(limit) - - param_count += 1 - offset_param = f"${param_count}" - query_params.append(offset) - - documents = await conn.fetch( - f""" - SELECT - id, - bank_id, - content_hash, - created_at, - updated_at, - LENGTH(original_text) as text_length, - retain_params, - tags - FROM {fq_table("documents")} - {where_clause} - ORDER BY created_at DESC - LIMIT {limit_param} OFFSET {offset_param} - """, - *query_params, - ) - - # Get memory unit count for each document - if documents: - doc_ids = [(row["id"], row["bank_id"]) for row in documents] - - # Create placeholders for the query - placeholders = [] - params_for_count = [] - for i, (doc_id, bank_id_val) in enumerate(doc_ids): - idx_doc = i * 2 + 1 - idx_agent = i * 2 + 2 - placeholders.append(f"(document_id = ${idx_doc} AND bank_id = ${idx_agent})") - params_for_count.extend([doc_id, bank_id_val]) - - where_clause_count = " OR ".join(placeholders) - - unit_counts = await conn.fetch( - f""" - SELECT document_id, bank_id, COUNT(*) as unit_count - FROM {fq_table("memory_units")} - WHERE {where_clause_count} - GROUP BY document_id, bank_id - """, - *params_for_count, - ) - else: - unit_counts = [] - - # Build count mapping - count_map = {(row["document_id"], row["bank_id"]): row["unit_count"] for row in unit_counts} - - # Build result items - items = [] - for row in documents: - doc_id = row["id"] - bank_id_val = row["bank_id"] - unit_count = count_map.get((doc_id, bank_id_val), 0) - - retain_params_val = conn.parse_json(row["retain_params"]) - - # document_metadata is sourced from retain_params.metadata - document_metadata = retain_params_val.get("metadata") if retain_params_val else None - - items.append( - { - "id": doc_id, - "bank_id": bank_id_val, - "content_hash": row["content_hash"], - "created_at": row["created_at"].isoformat() if row["created_at"] else "", - "updated_at": row["updated_at"].isoformat() if row["updated_at"] else "", - "text_length": row["text_length"] or 0, - "memory_unit_count": unit_count, - "retain_params": retain_params_val or None, - "document_metadata": document_metadata or None, - "tags": row["tags"] if row["tags"] else [], - } - ) - - return {"items": items, "total": total, "limit": limit, "offset": offset} - async def get_chunk( self, chunk_id: str, @@ -8636,12 +8378,27 @@ async def get_chunk( ) await self._validate_operation(self._operation_validator.validate_bank_read(ctx)) + # A store that owns the document store (memlake) keeps chunk_text there, not in the SQL + # chunks row (which is empty). Overlay it from the store. + chunk_text = chunk["chunk_text"] + from .memories import get_memories + + _store = get_memories() + if _store.owns_document_store: + _t = await _store.get_chunk_text( + bank_id=chunk["bank_id"], + document_id=chunk["document_id"], + chunk_index=chunk["chunk_index"], + ) + if _t is not None: + chunk_text = _t + return { "chunk_id": chunk["chunk_id"], "document_id": chunk["document_id"], "bank_id": chunk["bank_id"], "chunk_index": chunk["chunk_index"], - "chunk_text": chunk["chunk_text"], + "chunk_text": chunk_text, "created_at": chunk["created_at"].isoformat() if chunk["created_at"] else "", } @@ -8711,13 +8468,25 @@ async def list_document_chunks( offset, ) + # A store that owns the document store (memlake) keeps chunk_text there, not in the SQL + # chunks rows (which are empty). Fetch the document's chunk texts once (ordered by + # index) and overlay each row by its chunk_index. + _texts_by_index: dict[int, str] = {} + from .memories import get_memories + + _store = get_memories() + if _store.owns_document_store: + _texts = await _store.list_chunk_texts(bank_id=bank_id, document_id=document_id) + if _texts is not None: + _texts_by_index = dict(enumerate(_texts)) + items = [ { "chunk_id": row["chunk_id"], "document_id": row["document_id"], "bank_id": row["bank_id"], "chunk_index": row["chunk_index"], - "chunk_text": row["chunk_text"], + "chunk_text": _texts_by_index.get(row["chunk_index"], row["chunk_text"]), "created_at": row["created_at"].isoformat() if row["created_at"] else "", } for row in chunks @@ -10428,64 +10197,19 @@ async def list_entities( bank_id=bank_id, operation=BankReadOperation.LIST_ENTITIES, request_context=request_context ) await self._validate_operation(self._operation_validator.validate_bank_read(ctx)) + from .memories import get_memories + backend = await self._get_backend() async with acquire_with_retry(backend) as conn: - # Get total count - total_row = await conn.fetchrow( - f""" - SELECT COUNT(*) as total - FROM {fq_table("entities")} - WHERE bank_id = $1 - """, - bank_id, - ) - total = total_row["total"] if total_row else 0 - - # Get paginated entities - rows = await conn.fetch( - f""" - SELECT id, canonical_name, mention_count, first_seen, last_seen, metadata - FROM {fq_table("entities")} - WHERE bank_id = $1 - ORDER BY mention_count DESC, last_seen DESC, id ASC - LIMIT $2 OFFSET $3 - """, - bank_id, - limit, - offset, + # The store owns the entity list and its live memory counts. + return await get_memories().list_entities( + conn=conn, + fq_table=fq_table, + bank_id=bank_id, + limit=limit, + offset=offset, ) - entities = [] - for row in rows: - # Handle metadata - may be dict, JSON string, or None - metadata = row["metadata"] - if metadata is None: - metadata = {} - elif isinstance(metadata, str): - import json - - try: - metadata = json.loads(metadata) - except json.JSONDecodeError: - metadata = {} - - entities.append( - { - "id": str(row["id"]), - "canonical_name": row["canonical_name"], - "mention_count": row["mention_count"], - "first_seen": row["first_seen"].isoformat() if row["first_seen"] else None, - "last_seen": row["last_seen"].isoformat() if row["last_seen"] else None, - "metadata": metadata, - } - ) - return { - "items": entities, - "total": total, - "limit": limit, - "offset": offset, - } - async def get_entity_graph( self, bank_id: str, @@ -10626,13 +10350,25 @@ async def list_tags( bank_id=bank_id, operation=BankReadOperation.LIST_TAGS, request_context=request_context ) await self._validate_operation(self._operation_validator.validate_bank_read(ctx)) - return await self._list_tags_from_table( - table="memory_units", - bank_id=bank_id, - pattern=pattern, - limit=limit, - offset=offset, - ) + # Tags live with the memories, so the store produces the histogram; the + # wildcard filter, ordering (count desc, tag asc) and paging are applied + # here because the set is bounded by distinct tags, not the corpus. + from .memories import get_memories + + backend = await self._get_backend() + async with acquire_with_retry(backend) as conn: + histogram = await get_memories().list_tags(conn=conn, fq_table=fq_table, bank_id=bank_id) + items = [{"tag": tag, "count": count} for tag, count in histogram.items()] + if pattern: + # '*' is the wildcard, matched case-insensitively against the whole + # tag — the same anchored semantics `... ILIKE 'user:%'` had. + import re as _re + + regex = _re.compile("^" + ".*".join(_re.escape(part) for part in pattern.split("*")) + "$", _re.IGNORECASE) + items = [item for item in items if regex.match(str(item["tag"]))] + items.sort(key=lambda item: (-item["count"], item["tag"])) + total = len(items) + return {"items": items[offset : offset + limit], "total": total, "limit": limit, "offset": offset} async def list_mental_model_tags( self, @@ -10805,72 +10541,22 @@ async def get_bank_stats( ) async def _compute_bank_stats(self, bank_id: str) -> dict[str, Any]: + from .memories import get_memories + + store = get_memories() backend = await self._get_backend() async with acquire_with_retry(backend) as conn: - # Get node counts by fact_type - node_stats = await conn.fetch( - f""" - SELECT fact_type, COUNT(*) as count - FROM {fq_table("memory_units")} - WHERE bank_id = $1 - GROUP BY fact_type - """, - bank_id, - ) + # Node counts per fact_type come from the store — one metadata read + # for a store that keeps a live count, the same GROUP BY for Postgres. + node_counts = await store.count_memories(conn=conn, fq_table=fq_table, bank_id=bank_id) - # Non-entity link counts — no join, group by link_type only. With a - # (bank_id, link_type) index this is an index-only scan; without one - # it is at worst a single seq scan over memory_links rather than the - # multi-second hash join through memory_units that this used to be. - # The previous shape produced a (fact_type, link_type) matrix; only - # the hindsight-cli `bank stats` renderer still consumes the - # per-fact-type slice, and it tolerates empty maps (the section - # prints with no rows). Response keys are kept populated below for - # schema stability so existing SDK deserializers don't break. - # No link_type filter: entity edges are no longer stored in - # memory_links (dropped in migration e9b2c7d1f3a4 — derived on demand - # from unit_entities), so only temporal/semantic/caused_by rows exist - # here. Omitting the predicate lets the (bank_id, link_type) index - # serve this bank-scoped GROUP BY as an index-only scan. - non_entity_link_rows = await conn.fetch( - f""" - SELECT link_type, COUNT(*) as count - FROM {fq_table("memory_links")} - WHERE bank_id = $1 - GROUP BY link_type - """, - bank_id, - ) - - # Entity links are derived from unit_entities (no longer stored in - # memory_links). Replicate the historical writer cap: each unit - # linked bidirectionally to up to MAX_LINKS_PER_ENTITY others sharing - # each entity. Aggregated to a single scalar — the per-fact-type - # slice doubled the join cost and only fed link_counts_by_fact_type - # / link_breakdown, which the UI ignores and the CLI renders into - # sections that degrade gracefully when empty. - max_links_per_entity = 10 - entity_total_row = await conn.fetchrow( - f""" - WITH per_entity AS ( - SELECT ue.entity_id, COUNT(*) AS n - FROM {fq_table("unit_entities")} ue - JOIN {fq_table("memory_units")} mu ON mu.id = ue.unit_id - WHERE mu.bank_id = $1 - GROUP BY ue.entity_id - ) - SELECT COALESCE(SUM(LEAST(n - 1, $2)), 0)::bigint AS count - FROM per_entity - """, - bank_id, - max_links_per_entity, - ) - entity_link_total = int(entity_total_row["count"] or 0) if entity_total_row else 0 - - link_counts: dict[str, int] = {row["link_type"]: row["count"] for row in non_entity_link_rows} - if entity_link_total > 0: - link_counts["entity"] = entity_link_total + # Link counts come from the store, like node_counts — a store keeps its links in + # its own shape (Postgres in memory_links + unit_entities; another store may keep + # them inside the memory), so the stats page's link total must be asked of the + # store rather than read straight from Postgres tables a non-Postgres store leaves + # empty. Keyed by link type; the response sums the values below. + link_counts = await store.link_counts(conn=conn, fq_table=fq_table, bank_id=bank_id) ops_stats = await conn.fetch( f""" @@ -10885,19 +10571,27 @@ async def _compute_bank_stats(self, bank_id: str) -> dict[str, Any]: f"SELECT COUNT(*) as count FROM {fq_table('documents')} WHERE bank_id = $1", bank_id, ) - consolidation_row = await conn.fetchrow( - f""" - SELECT - MAX(consolidated_at) as last_consolidated_at, - COUNT(*) FILTER (WHERE consolidated_at IS NULL AND fact_type IN ('experience', 'world')) as pending, - COUNT(*) FILTER (WHERE consolidation_failed_at IS NOT NULL AND fact_type IN ('experience', 'world')) as failed - FROM {fq_table("memory_units")} - WHERE bank_id = $1 - """, - bank_id, - ) + # Consolidation freshness (last-consolidated, pending, failed) lives on the memories, + # so a store that keeps them outside SQL must answer this — the memory_units query + # returns 0/None for it. Same {last_consolidated_at, pending, failed} shape either way. + from .memories import get_memories + + _store = get_memories() + if _store.writes_memory_rows_in_sql: + consolidation_row = await conn.fetchrow( + f""" + SELECT + MAX(consolidated_at) as last_consolidated_at, + COUNT(*) FILTER (WHERE consolidated_at IS NULL AND fact_type IN ('experience', 'world')) as pending, + COUNT(*) FILTER (WHERE consolidation_failed_at IS NOT NULL AND fact_type IN ('experience', 'world')) as failed + FROM {fq_table("memory_units")} + WHERE bank_id = $1 + """, + bank_id, + ) + else: + consolidation_row = await _store.consolidation_freshness(conn=conn, fq_table=fq_table, bank_id=bank_id) - node_counts = {row["fact_type"]: row["count"] for row in node_stats} ops_by_status = {row["status"]: row["count"] for row in ops_stats} last_consolidated_at = consolidation_row["last_consolidated_at"] if consolidation_row else None @@ -10951,26 +10645,16 @@ async def get_bank_freshness( # contract (see interface.get_bank_freshness) so the returned shape stays # a strict subset of get_bank_stats. All three come from one scan, so # keeping `failed` costs nothing extra. + from .memories import get_memories + async with acquire_with_retry(backend) as conn: - row = await conn.fetchrow( - f""" - SELECT - MAX(consolidated_at) as last_consolidated_at, - COUNT(*) FILTER (WHERE consolidated_at IS NULL AND fact_type IN ('experience', 'world')) as pending, - COUNT(*) FILTER (WHERE consolidation_failed_at IS NOT NULL AND fact_type IN ('experience', 'world')) as failed - FROM {fq_table("memory_units")} - WHERE bank_id = $1 - """, - bank_id, - ) + fresh = await get_memories().consolidation_freshness(conn=conn, fq_table=fq_table, bank_id=bank_id) - if row is None: - return {"last_consolidated_at": None, "pending_consolidation": 0, "failed_consolidation": 0} - last = row["last_consolidated_at"] + last = fresh["last_consolidated_at"] return { "last_consolidated_at": last.isoformat() if last else None, - "pending_consolidation": row["pending"] or 0, - "failed_consolidation": row["failed"] or 0, + "pending_consolidation": fresh["pending"], + "failed_consolidation": fresh["failed"], } async def _probe_llm(self, llm: Any) -> _LlmProbeOutcome: @@ -11079,22 +10763,16 @@ async def get_memories_timeseries( _ALLOWED_TIME_FIELDS = ("created_at", "mentioned_at", "occurred_start") if time_field not in _ALLOWED_TIME_FIELDS: time_field = "created_at" - # COALESCE onto created_at for event-time fields so null rows don't vanish. - bucket_expr = time_field if time_field == "created_at" else f"COALESCE({time_field}, created_at)" + from .memories import get_memories + + # The window: everything since one full period back. Computed here so the + # store gets a concrete `since` rather than a dialect interval string. + since = datetime.now(timezone.utc) - cfg.step * cfg.count backend = await self._get_backend() async with acquire_with_retry(backend) as conn: - rows = await conn.fetch( - f""" - SELECT date_trunc('{cfg.trunc}', {bucket_expr} AT TIME ZONE 'UTC') AS bucket, - fact_type, COUNT(*) AS count - FROM {fq_table("memory_units")} - WHERE bank_id = $1 - AND {bucket_expr} >= now() - interval '{cfg.interval}' - GROUP BY bucket, fact_type - ORDER BY bucket - """, - bank_id, + rows = await get_memories().memories_timeseries( + conn=conn, fq_table=fq_table, bank_id=bank_id, time_field=time_field, trunc=cfg.trunc, since=since ) # Build the canonical bucket list anchored on the most recent UTC boundary. @@ -12361,15 +12039,21 @@ def _get(key: str) -> Any: fact_types: list[str] = list(trigger.get("fact_types") or []) tag_filtering = _resolve_refresh_tag_filtering(mm_tags, trigger) - scope_filter = self._build_mm_scope_filter(bank_id, tag_filtering, fact_types) - params = [*scope_filter.params, last_refreshed_at] - where = [*scope_filter.where, f"updated_at > ${len(params)}"] + # The scoped existence check belongs to the store: it is a query over the + # memories, and the mental model's scope (tags, tag_groups, fact_types) is + # exactly what decides whether one of them changed since the last refresh. + from .memories import get_memories - row = await conn.fetchrow( - f"SELECT 1 FROM {fq_table('memory_units')} WHERE {' AND '.join(where)} LIMIT 1", - *params, + return await get_memories().any_memory_updated_since( + conn=conn, + fq_table=fq_table, + bank_id=bank_id, + since=last_refreshed_at, + fact_types=fact_types, + tags=tag_filtering.tags, + tags_match=tag_filtering.tags_match, + tag_groups=tag_filtering.tag_groups, ) - return row is not None def _row_to_mental_model(self, row, *, detail: str = "full") -> dict[str, Any]: """Convert a database row to a mental model dict. diff --git a/hindsight-api-slim/hindsight_api/engine/reflect/tools.py b/hindsight-api-slim/hindsight_api/engine/reflect/tools.py index 69e1cda420..d02cc3b176 100644 --- a/hindsight-api-slim/hindsight_api/engine/reflect/tools.py +++ b/hindsight-api-slim/hindsight_api/engine/reflect/tools.py @@ -343,16 +343,37 @@ async def tool_expand( valid_uuids = list(uuid_by_id.values()) - # Batch fetch all memory units - memories = await conn.fetch( - f""" - SELECT id, text, chunk_id, document_id, fact_type, context - FROM {fq_table("memory_units")} - WHERE id = ANY($1) AND bank_id = $2 - """, - valid_uuids, - bank_id, - ) + # Batch fetch all memory units. A store that keeps memories outside SQL answers by id + # through the store; normalize its records to the same UUID-keyed dict shape the SQL rows + # have so the result-building below stays store-agnostic. + from ..memories import get_memories + + _store = get_memories() + if _store.writes_memory_rows_in_sql: + memories = await conn.fetch( + f""" + SELECT id, text, chunk_id, document_id, fact_type, context + FROM {fq_table("memory_units")} + WHERE id = ANY($1) AND bank_id = $2 + """, + valid_uuids, + bank_id, + ) + else: + stored = await _store.get_memories( + conn=conn, fq_table=fq_table, bank_id=bank_id, unit_ids=[str(u) for u in valid_uuids] + ) + memories = [ + { + "id": uuid.UUID(s.unit_id), + "text": s.text, + "chunk_id": s.chunk_id, + "document_id": s.document_id, + "fact_type": s.fact_type, + "context": s.context, + } + for s in stored + ] memory_map = {row["id"]: row for row in memories} # Collect chunk_ids and document_ids for batch fetching diff --git a/hindsight-api-slim/hindsight_api/engine/retain/bank_utils.py b/hindsight-api-slim/hindsight_api/engine/retain/bank_utils.py index e4e4e587c1..3447bd8acb 100644 --- a/hindsight-api-slim/hindsight_api/engine/retain/bank_utils.py +++ b/hindsight-api-slim/hindsight_api/engine/retain/bank_utils.py @@ -450,6 +450,12 @@ async def list_banks(pool) -> list: ) result = [] + # A store that keeps memories outside SQL leaves the memory_units join empty, so its + # per-bank fact_count comes from the store instead (one live count per bank). + from ..memories import get_memories + + _store = get_memories() + for row in rows: disposition_data = row["disposition"] if isinstance(disposition_data, str): @@ -457,6 +463,12 @@ async def list_banks(pool) -> list: last_doc = row["last_document_at"] + fact_count = row["fact_count"] + if not _store.writes_memory_rows_in_sql: + fact_count = sum( + (await _store.count_memories(conn=conn, fq_table=fq_table, bank_id=row["bank_id"])).values() + ) + result.append( { "bank_id": row["bank_id"], @@ -465,7 +477,7 @@ async def list_banks(pool) -> list: "mission": row["mission"] or "", "created_at": row["created_at"].isoformat() if row["created_at"] else None, "updated_at": row["updated_at"].isoformat() if row["updated_at"] else None, - "fact_count": row["fact_count"], + "fact_count": fact_count, "last_document_at": last_doc.isoformat() if last_doc else None, } ) diff --git a/hindsight-api-slim/hindsight_api/engine/retain/chunk_storage.py b/hindsight-api-slim/hindsight_api/engine/retain/chunk_storage.py index 90d7b25b8c..1e7d8d0012 100644 --- a/hindsight-api-slim/hindsight_api/engine/retain/chunk_storage.py +++ b/hindsight-api-slim/hindsight_api/engine/retain/chunk_storage.py @@ -55,16 +55,31 @@ async def load_existing_chunks(conn, bank_id: str, document_id: str) -> list[Exi ] -async def delete_chunks_by_ids(conn, chunk_ids: list[str]) -> None: +async def delete_chunks_by_ids(conn, chunk_ids: list[str], bank_id: str | None = None, txn=None) -> None: """ Delete specific chunks by their IDs. This cascades to memory_units (via FK with CASCADE delete) and their links. + + ``txn`` carries a cross-store write-group handle when this delete is part of a re-ingest: + the store's tombstones must ride the same txn as the replacement writes so they commit + (become visible) together — otherwise an aborted re-ingest could drop the old memories + without landing the new ones. """ if not chunk_ids: return + # The chunks->memory_units FK cascade below does not reach a store that keeps memories + # outside SQL (its memory_units is empty), so drop the memories carrying each deleted + # chunk_id through the store — otherwise a delta re-ingest leaves the old ones as duplicates. + from ..memories import META_CHUNK_ID, DeletePredicate, get_memories + + _store = get_memories() + if bank_id and not _store.writes_memory_rows_in_sql: + for _cid in chunk_ids: + await _store.delete_where(bank_id, DeletePredicate(metadata_equals={META_CHUNK_ID: _cid}), txn=txn) + # PostgreSQL's FK cascade deletes child memory_links in executor-chosen # order. Concurrent chunk deletes for the same bank can then lock overlapping # memory_links in opposite orders and deadlock. Delete links explicitly in a @@ -148,6 +163,13 @@ async def store_chunks_batch( # Fallback to the raw global default (not get_config(), which guards # bank-configurable fields); the retain path always passes the resolved value. store_text = store_document_text if store_document_text is not None else _get_raw_config().store_document_text + # A store that owns a dedicated document store (memlake) keeps the chunk TEXT there, so the + # SQL chunks row carries only its metadata (chunk_id, index, content_hash) with empty text — + # same shape as store_document_text=False, and idempotency is unaffected (content_hash stays). + from ..memories import get_memories + + if get_memories().owns_document_store: + store_text = False # Prepare chunk data for batch insert chunk_ids = [] diff --git a/hindsight-api-slim/hindsight_api/engine/retain/fact_storage.py b/hindsight-api-slim/hindsight_api/engine/retain/fact_storage.py index ef9c3b8052..812a7c79ba 100644 --- a/hindsight-api-slim/hindsight_api/engine/retain/fact_storage.py +++ b/hindsight-api-slim/hindsight_api/engine/retain/fact_storage.py @@ -9,7 +9,7 @@ import uuid from datetime import datetime -from ...config import _get_raw_config, get_config +from ...config import _get_raw_config from ..memory_engine import fq_table from .bank_utils import DEFAULT_DISPOSITION, create_bank_vector_indexes from .fact_extraction import _sanitize_text @@ -17,6 +17,11 @@ logger = logging.getLogger(__name__) +#: Page size for walking a replaced document's outgoing memories. Large enough +#: that one page covers any ordinary document, small enough that a pathological +#: one does not arrive as a single result set. +_OUTGOING_PAGE = 500 + async def get_document_content( conn, @@ -36,16 +41,27 @@ async def get_document_content( async def insert_facts_batch( - conn, bank_id: str, facts: list[ProcessedFact], document_id: str | None = None, ops=None + conn, + bank_id: str, + facts: list[ProcessedFact], + document_id: str | None = None, + ops=None, + defer_index: bool = False, + txn=None, ) -> list[str]: """ - Insert facts into the database in batch. + Store facts and return their unit ids, in order. Args: conn: Database connection bank_id: Bank identifier facts: List of ProcessedFact objects to insert document_id: Optional document ID to associate with facts + defer_index: Ask for ids without the write. The retain orchestrator needs + this because it can only supply entity ids and causal edges after + Phase-1 placeholders have been remapped onto real unit ids; it then + calls `index_facts` with the complete picture. The Postgres store, + whose write *is* the insert that mints the ids, ignores it. Returns: List of unit IDs (UUIDs as strings) for the inserted facts @@ -53,85 +69,37 @@ async def insert_facts_batch( if not facts: return [] - # Prepare data for batch insert - fact_texts = [] - embeddings = [] - event_dates = [] - occurred_starts = [] - occurred_ends = [] - mentioned_ats = [] - contexts = [] - fact_types = [] - metadata_jsons = [] - chunk_ids = [] - document_ids = [] - tags_list = [] - observation_scopes_list = [] - text_signals_list = [] - - for fact in facts: - fact_texts.append(_sanitize_text(fact.fact_text)) - # Convert embedding to string for asyncpg vector type - embeddings.append(str(fact.embedding)) - # event_date: Use occurred_start if available, otherwise use mentioned_at - # This maintains backward compatibility while handling None occurred_start - event_dates.append(fact.occurred_start if fact.occurred_start is not None else fact.mentioned_at) - occurred_starts.append(fact.occurred_start) - occurred_ends.append(fact.occurred_end) - mentioned_ats.append(fact.mentioned_at) - contexts.append(_sanitize_text(fact.context)) - fact_types.append(fact.fact_type) - metadata_jsons.append(json.dumps(fact.metadata)) - chunk_ids.append(fact.chunk_id) - # Use per-fact document_id if available, otherwise fallback to batch-level document_id - document_ids.append(fact.document_id if fact.document_id else document_id) - # Convert tags to JSON string for proper batch insertion (PostgreSQL unnest doesn't handle 2D arrays well) - tags_list.append(json.dumps(fact.tags if fact.tags else [])) - # observation_scopes: stored as JSONB (string or 2D array), None if not provided - observation_scopes_list.append( - json.dumps(fact.observation_scopes) if fact.observation_scopes is not None else None - ) - # Build text_signals: entity names + date tokens for enriched BM25 indexing - signal_parts = [] - if fact.entities: - signal_parts.extend(e.name for e in fact.entities) - if fact.occurred_start: - try: - signal_parts.append(fact.occurred_start.strftime("%B %d %Y").lstrip("0").replace(" 0", " ")) - except (ValueError, AttributeError): - pass - if fact.occurred_end and fact.occurred_end != fact.occurred_start: - try: - signal_parts.append(fact.occurred_end.strftime("%B %d %Y").lstrip("0").replace(" 0", " ")) - except (ValueError, AttributeError): - pass - text_signals_list.append(" ".join(signal_parts) if signal_parts else None) - - # Batch insert all facts — delegates to DataAccessOps which handles - # unnest (PG) vs row-by-row (Oracle) transparently. - config = get_config() - - return await ops.insert_facts_batch( - conn, - bank_id, - fact_texts, - embeddings, - event_dates, - occurred_starts, - occurred_ends, - mentioned_ats, - contexts, - fact_types, - metadata_jsons, - chunk_ids, - document_ids, - tags_list, - observation_scopes_list, - text_signals_list, - text_search_extension=config.text_search_extension, + from ..memories import get_memories + + return await get_memories().insert_facts( + conn=conn, + ops=ops, + bank_id=bank_id, + facts=facts, + document_id=document_id, + defer_index=defer_index, + txn=txn, ) +async def index_facts( + bank_id: str, + unit_ids: list[str], + facts: list[ProcessedFact], + document_id: str | None = None, + unit_entity_ids: dict[str, list[str]] | None = None, +) -> None: + """Complete a deferred `insert_facts_batch`, now that the edges are known. + + ``unit_entity_ids`` is the unit→entity posting and each fact's causal + relations are its edges; both travel with the memory for a store that owns + them. A no-op for the Postgres store, which wrote all of it already. + """ + from ..memories import get_memories + + await get_memories().index_facts(bank_id, unit_ids, facts, document_id, unit_entity_ids) + + async def ensure_bank_exists(conn, bank_id: str, ops=None) -> None: """ Ensure bank exists in the database. @@ -172,94 +140,36 @@ async def delete_stale_observations_for_memories( """Delete observations whose source memories are about to be removed. Mirrors the cleanup performed by ``MemoryEngine.delete_document`` so that - every code path that removes ``memory_units`` also removes the - observations derived from them. Without this, ingesting a fresh version - of a document via the retain pipeline (which does a full-replace - ``DELETE FROM documents`` cascade) used to leave orphan observations - pointing at memory IDs that no longer existed. + every code path that removes memories also removes the observations derived + from them. Without this, ingesting a fresh version of a document via the + retain pipeline (which does a full-replace ``DELETE FROM documents`` + cascade) used to leave orphan observations pointing at memory IDs that no + longer existed. For each observation referencing any of ``fact_ids``: - 1. Delete the observation row (its text is stale once even one source - memory disappears). - 2. Reset ``consolidated_at = NULL`` on the surviving source memories so - they get re-consolidated under fresh observations on the next run. + 1. Delete the observation (its text is stale once even one source memory + disappears). + 2. Reset the consolidated marker on the surviving source memories so they + get re-consolidated under fresh observations on the next run. - Must be called within an active transaction, before the source memories - are deleted. + Must be called within an active transaction, before the source memories are + deleted. - Returns the number of observations deleted. + Returns: + Number of observations deleted. """ if not fact_ids: return 0 - fact_uuids = [uuid.UUID(str(fid)) if not isinstance(fid, uuid.UUID) else fid for fid in fact_ids] - - if ops is not None and not ops.uses_observation_sources_table: - # PG: use native array overlap operator - affected_obs = await conn.fetch( - f""" - SELECT id, source_memory_ids - FROM {fq_table("memory_units")} - WHERE bank_id = $1 - AND fact_type = 'observation' - AND source_memory_ids && $2::uuid[] - """, - bank_id, - fact_uuids, - ) - else: - # Oracle / default: use observation_sources junction table - affected_obs = await conn.fetch( - f""" - SELECT mu.id, mu.source_memory_ids - FROM {fq_table("memory_units")} mu - WHERE mu.bank_id = $1 - AND mu.fact_type = 'observation' - AND EXISTS ( - SELECT 1 FROM {fq_table("observation_sources")} os - WHERE os.observation_id = mu.id - AND os.source_id = ANY($2::uuid[]) - ) - """, - bank_id, - fact_uuids, - ) - - if not affected_obs: - return 0 - - deleted_set = {str(uid) for uid in fact_uuids} - obs_ids = [obs["id"] for obs in affected_obs] - seen_remaining: set[str] = set() - remaining_source_ids: list[uuid.UUID] = [] - for obs in affected_obs: - for src_id in obs["source_memory_ids"] or []: - src_str = str(src_id) - if src_str not in deleted_set and src_str not in seen_remaining: - remaining_source_ids.append(src_id) - seen_remaining.add(src_str) - - await conn.execute( - f"DELETE FROM {fq_table('memory_units')} WHERE id = ANY($1::uuid[])", - obs_ids, - ) - - if remaining_source_ids: - await conn.execute( - f""" - UPDATE {fq_table("memory_units")} - SET consolidated_at = NULL - WHERE id = ANY($1::uuid[]) - AND fact_type IN ('experience', 'world') - """, - remaining_source_ids, - ) + from ..memories import get_memories - logger.info( - f"[OBSERVATIONS] Deleted {len(obs_ids)} observations, reset {len(remaining_source_ids)} " - f"source memories for re-consolidation in bank {bank_id}" + return await get_memories().delete_stale_observations( + conn=conn, + ops=ops, + fq_table=fq_table, + bank_id=bank_id, + fact_ids=fact_ids, ) - return len(obs_ids) async def handle_document_tracking( @@ -272,6 +182,7 @@ async def handle_document_tracking( document_tags: list[str] | None = None, ops=None, store_document_text: bool | None = None, + txn=None, ) -> None: """ Handle document tracking in the database (full-replace mode). @@ -308,14 +219,29 @@ async def handle_document_tracking( # frozen). Same cleanup the explicit ``delete_document`` API performs. preserved_created_at = None if is_first_batch: - existing_unit_rows = await conn.fetch( - f""" - SELECT id FROM {fq_table("memory_units")} - WHERE document_id = $1 AND fact_type IN ('experience', 'world') - """, - document_id, - ) - existing_unit_ids = [row["id"] for row in existing_unit_rows] + from ..memories import get_memories + + store = get_memories() + # Which memories the outgoing version left behind. Asked of the store + # rather than queried here, because it is the store that knows where they + # are. Paged to exhaustion: every one of them is about to be deleted, and + # a document whose facts overflow one page must not keep half of them. + existing_unit_ids: list[str] = [] + page_token = "" + while True: + page = await store.scan_memories( + conn=conn, + fq_table=fq_table, + bank_id=bank_id, + fact_types=["experience", "world"], + document_id=document_id, + limit=_OUTGOING_PAGE, + page_token=page_token, + ) + existing_unit_ids.extend(m.unit_id for m in page.memories) + page_token = page.next_page_token + if not page_token: + break if existing_unit_ids: invalidated = await delete_stale_observations_for_memories(conn, bank_id, existing_unit_ids, ops=ops) if invalidated: @@ -332,16 +258,13 @@ async def handle_document_tracking( from ..graph_maintenance import enqueue_relink_victims await enqueue_relink_victims(conn, bank_id, [str(uid) for uid in existing_unit_ids], ops=ops) + # Explicitly delete memory_units by document_id BEFORE deleting the # document row. The CASCADE from documents→chunks→memory_units only # catches units that have a non-NULL chunk_id FK. Units with chunk_id=NULL # (e.g. from partial writes or edge cases) would survive the cascade. # This explicit delete ensures complete cleanup. - await conn.execute( - f"DELETE FROM {fq_table('memory_units')} WHERE document_id = $1 AND bank_id = $2", - document_id, - bank_id, - ) + await store.delete_document(conn=conn, fq_table=fq_table, bank_id=bank_id, document_id=document_id, txn=txn) # Capture created_at before deletion so re-ingestion preserves it. preserved_created_at = await conn.fetchval( f"DELETE FROM {fq_table('documents')} WHERE id = $1 AND bank_id = $2 RETURNING created_at", @@ -423,6 +346,13 @@ async def _upsert_document_row( # bank-configurable fields); the retain path always passes the resolved value. store_text = store_document_text if store_document_text is not None else _get_raw_config().store_document_text original_text = combined_content if store_text else None + # A store that owns a dedicated document store (memlake) keeps the extracted text there, so the + # SQL documents row holds only its metadata (id, content_hash, tags) with original_text NULL — + # the bulky body is written to the store up front (orchestrator._store_document_bodies). + from ..memories import get_memories + + if get_memories().owns_document_store: + original_text = None await conn.execute( f""" INSERT INTO {fq_table("documents")} (id, bank_id, original_text, content_hash, retain_params, tags, created_at, updated_at) @@ -458,6 +388,20 @@ async def update_memory_units_tags( Returns: Number of memory units updated. """ + from ..memories import MemoryPatch, get_memories + + store = get_memories() + if not store.writes_memory_rows_in_sql: + # A store that keeps memories outside SQL: page the document's memories and patch each + # one's tags through the store — the UPDATE below is a no-op on its empty memory_units. + page = await store.scan_memories( + conn=conn, fq_table=fq_table, bank_id=bank_id, document_id=document_id, limit=1_000_000 + ) + patches = [MemoryPatch(unit_id=m.unit_id, tags=list(tags or [])) for m in page.memories] + if patches: + await store.update_memories(bank_id, patches) + return len(patches) + result = await conn.execute( f""" UPDATE {fq_table("memory_units")} diff --git a/hindsight-api-slim/hindsight_api/engine/retain/orchestrator.py b/hindsight-api-slim/hindsight_api/engine/retain/orchestrator.py index f1ee5d18c3..73f1129f90 100644 --- a/hindsight-api-slim/hindsight_api/engine/retain/orchestrator.py +++ b/hindsight-api-slim/hindsight_api/engine/retain/orchestrator.py @@ -464,6 +464,7 @@ async def _insert_facts_and_links( skip_semantic_links: bool = False, outbox_callback=None, ops=None, + txn=None, ) -> list[list[str]]: """ Phase 2 of the retain pipeline: insert facts and retrieval-critical links. @@ -476,7 +477,7 @@ async def _insert_facts_and_links( memory_links here. """ set_stage("retain.phase2.insert_facts") - unit_ids = await fact_storage.insert_facts_batch(conn, bank_id, processed_facts, ops=ops) + unit_ids = await fact_storage.insert_facts_batch(conn, bank_id, processed_facts, ops=ops, txn=txn) step_start = time.time() log_buffer.append(f" Insert facts: {len(unit_ids)} units in {time.time() - step_start:.3f}s") @@ -501,7 +502,7 @@ async def _insert_facts_and_links( # closing the window where prune_orphan_entities could have deleted one # between Phase-1 resolution and this insert (#2662). await entity_resolver.reassert_entities_batch(bank_id, resolved_entities, conn=conn) - await entity_resolver.link_units_to_entities_batch(unit_entity_pairs, conn=conn) + await entity_resolver.link_units_to_entities_batch(unit_entity_pairs, conn=conn, bank_id=bank_id) log_buffer.append(f" Insert unit_entities: {len(unit_entity_pairs)} pairs in {time.time() - step_start:.3f}s") # Create temporal links @@ -1220,6 +1221,56 @@ async def _process_ann_chunk(chunk_idx: int) -> None: log_buffer.append(f"[streaming] Final ANN: {total_links} total semantic links") +# --------------------------------------------------------------------------- +# Document bodies → the store's dedicated document store (memlake) +# --------------------------------------------------------------------------- + + +async def _store_document_bodies( + *, + bank_id: str, + document_id: str, + combined_content: str, + chunk_texts: list[str], + merged_tags: list[str] | None, + config: Any, + content_hash: str | None = None, +) -> None: + """Route a document's bulky bodies — its extracted text and ordered chunk texts — to the + store's dedicated document store, when the store owns one (memlake). No-op for Postgres. + + Content-addressed and idempotent, so this is safe to call up front, before the facts commit: + a re-ingest re-uploads only the bodies whose hash changed, and a retain that later rolls back + leaves only orphan bodies the store's sweep reclaims (they are referenced by no committed + record). The SQL ``documents``/``chunks`` rows still carry the small metadata (id, + content_hash, chunk_index, tags) — only their bulky text columns are left empty (see + ``fact_storage._upsert_document_row`` / ``chunk_storage.store_chunks_batch``). Cold, + never-searched, key-based — see docs/documents-chunks.md. + """ + from ..memories import get_memories + + store = get_memories() + if not store.owns_document_store: + return + # The record's content_hash must equal what the SQL documents row stores, so a read is + # consistent whichever it comes from: sanitize + sha256 the same combined_content. The + # streaming path already has it (passes it in); delta computes it here. + if content_hash is None: + _sanitized = fact_extraction._sanitize_text(combined_content) or "" + content_hash = hashlib.sha256(_sanitized.encode()).hexdigest() + await store.put_document( + bank_id=bank_id, + document_id=document_id, + content_hash=content_hash, + # Honour store_document_text: when a deployment opts out of keeping the full text, only the + # chunk texts (needed for citation) go to the store, not the whole document body. + original_text=combined_content if getattr(config, "store_document_text", True) else None, + chunk_texts=list(chunk_texts), + tags=list(merged_tags or []), + metadata={}, + ) + + # --------------------------------------------------------------------------- # Streaming chunk batching # --------------------------------------------------------------------------- @@ -1340,6 +1391,23 @@ async def _streaming_retain_batch( # to detect when a concurrent request has taken over the document. # See _run_mini_batch_db_work() for the implementation. retain_params, merged_tags = _build_retain_params(contents_dicts, document_tags) + + # Route the document's bulky bodies (extracted text + ordered chunk texts) to the store's + # dedicated document store (memlake) when it owns one — up front, before the streaming batches + # write facts. Idempotent and content-addressed, so this is safe here and dedups a re-ingest; + # a no-op for a Postgres store (which keeps the text in its own columns below). ``all_pre_chunks`` + # is the full ordered chunk-text list; ``combined_content`` is the full document text (both are + # released as the batches stream, so the write happens now while they are still resident). + await _store_document_bodies( + bank_id=bank_id, + document_id=effective_doc_id, + content_hash=new_content_hash, + combined_content=combined_content, + chunk_texts=all_pre_chunks, + merged_tags=merged_tags, + config=config, + ) + # Track whether document tracking has been done (by the first batch) doc_tracking_done = [False] # Track whether the transactional-outbox callback has already fired inside a @@ -1576,6 +1644,10 @@ async def _process_db_batch( # (cascade-delete + insert doc row) to establish ownership and prevent # concurrent requests from interleaving. Later batches can safely skip. if not doc_tracking_done[0]: + from ..memories import get_memories + + _edge_provider = get_memories() + _edge_txn = None async with acquire_with_retry(pool) as conn: async with conn.transaction(): await conn.execute( @@ -1602,6 +1674,11 @@ async def _process_db_batch( store_document_text=getattr(config, "store_document_text", True), ) else: + # A 0-fact re-ingest still deletes the outgoing memories — tag that + # tombstone with a write-group so it commits atomically with the doc row. + _edge_txn = await _edge_provider.begin_txn( + conn=conn, fq_table=fq_table, bank_id=bank_id, mutating=True + ) await fact_storage.handle_document_tracking( conn, bank_id, @@ -1612,6 +1689,7 @@ async def _process_db_batch( merged_tags, ops=pool.ops, store_document_text=getattr(config, "store_document_text", True), + txn=_edge_txn, ) doc_tracking_done[0] = True # Memory: combined_content has been persisted; release @@ -1619,6 +1697,8 @@ async def _process_db_batch( # a multi-MB string. Nothing reads it after tracking. combined_content = "" log_buffer.append(f"[streaming] Document {effective_doc_id} tracked (0 facts in first batch)") + if _edge_txn is not None: + await _edge_provider.decide_txn(_edge_txn, commit=True) log_buffer.append( f"[streaming] Consumer batch {consumer_batch_idx + 1}: " f"0 facts extracted from {len(batch)} chunks, skipping" @@ -1693,6 +1773,18 @@ async def _run_mini_batch_db_work() -> None: bank_id, ) + # Open the cross-store write-group txn INSIDE this batch's transaction, + # before the first-batch replace deletes any outgoing memories: the delete + # and this batch's writes must ride the same txn so they commit together. + # Streaming is per-batch atomic (each batch its own PG txn), so each batch + # is its own write-group — matching the existing transactional granularity. + from ..memories import get_memories + + _provider = get_memories() + _memlake_txn = await _provider.begin_txn( + conn=conn, fq_table=fq_table, bank_id=bank_id, mutating=True + ) + if not doc_tracking_done[0]: # --- First batch: document tracking (atomic with chunk write) --- if is_recovery: @@ -1720,6 +1812,7 @@ async def _run_mini_batch_db_work() -> None: merged_tags, ops=pool.ops, store_document_text=getattr(config, "store_document_text", True), + txn=_memlake_txn, ) log_buffer.append(f"[streaming] Document {effective_doc_id} tracked (full content)") doc_tracking_done[0] = True @@ -1783,8 +1876,13 @@ async def _run_mini_batch_db_work() -> None: skip_semantic_links=True, outbox_callback=outbox_callback if is_last else None, ops=pool.ops, + txn=_memlake_txn, ) + # Postgres committed this batch: publish its write-group. If it had aborted, + # this is skipped and the recovery sweep resolves the undecided txn (spec §5). + await _provider.decide_txn(_memlake_txn, commit=True) + logger.info(f"[streaming] Phase 2 (write txn): {time.time() - p2_start:.3f}s") # The write TXN above committed the transactional-outbox row in the @@ -1913,6 +2011,10 @@ async def _run_mini_batch_db_work() -> None: # never created by the first batch TXN. Create it now so the document # is tracked regardless of extraction results. if not doc_tracking_done[0] and not pipeline_aborted[0]: + from ..memories import get_memories + + _edge_provider = get_memories() + _edge_txn = None async with acquire_with_retry(pool) as conn: async with conn.transaction(): await conn.execute( @@ -1938,6 +2040,11 @@ async def _run_mini_batch_db_work() -> None: store_document_text=getattr(config, "store_document_text", True), ) else: + # A no-facts re-ingest still deletes the outgoing memories — tag that + # tombstone with a write-group so it commits atomically with the doc row. + _edge_txn = await _edge_provider.begin_txn( + conn=conn, fq_table=fq_table, bank_id=bank_id, mutating=True + ) await fact_storage.handle_document_tracking( conn, bank_id, @@ -1948,12 +2055,15 @@ async def _run_mini_batch_db_work() -> None: merged_tags, ops=pool.ops, store_document_text=getattr(config, "store_document_text", True), + txn=_edge_txn, ) doc_tracking_done[0] = True # Memory: combined_content has been persisted and won't be # read again — release the per-document text now. combined_content = "" log_buffer.append(f"[streaming] Document {effective_doc_id} tracked (no facts extracted)") + if _edge_txn is not None: + await _edge_provider.decide_txn(_edge_txn, commit=True) # Transactional-outbox fallback. The in-TXN fire only runs on a final # facts-bearing batch (is_last=True). When the committed-chunk count lands @@ -2401,8 +2511,29 @@ async def _run_delta_db_work() -> None: retain_params, merged_tags, ) + # Re-store the document's bodies in the store's document store (memlake) with the + # FULL new chunk set — put_document dedups by content hash, so unchanged chunks and + # text re-upload nothing; only what the delta changed moves. A no-op for Postgres. + await _store_document_bodies( + bank_id=bank_id, + document_id=effective_doc_id, + combined_content=combined_content, + chunk_texts=[new_chunks_with_contents[i] for i in sorted(new_chunks_with_contents)], + merged_tags=merged_tags, + config=config, + ) log_buffer.append(f" Document metadata update in {time.time() - step_start:.3f}s") + # Open the cross-store write-group txn INSIDE this transaction, BEFORE the + # tombstones below: a re-ingest deletes the old memories and writes new ones, + # and both must ride the same txn so they become visible together — an aborted + # re-ingest must not drop the old without landing the new. The witness row is + # the commit proof the recovery sweep consults. + from ..memories import get_memories + + _provider = get_memories() + _memlake_txn = await _provider.begin_txn(conn=conn, fq_table=fq_table, bank_id=bank_id, mutating=True) + # Delete changed and removed chunks (cascades to memory_units and links) step_start = time.time() chunks_to_delete = [ @@ -2410,7 +2541,7 @@ async def _run_delta_db_work() -> None: for idx in changed_indices + removed_indices if idx in existing_by_index ] - await chunk_storage.delete_chunks_by_ids(conn, chunks_to_delete) + await chunk_storage.delete_chunks_by_ids(conn, chunks_to_delete, bank_id, txn=_memlake_txn) log_buffer.append( f" Deleted {len(chunks_to_delete)} chunks " f"({len(changed_indices)} changed + {len(removed_indices)} removed) " @@ -2480,8 +2611,14 @@ async def _run_delta_db_work() -> None: semantic_ann_links=phase1.semantic_ann_links, outbox_callback=outbox_callback, ops=pool.ops, + txn=_memlake_txn, ) + # Postgres has committed: publish the write-group so its writes become visible. + # If the transaction had aborted instead, this line is skipped and the recovery + # sweep resolves the undecided txn against the (absent) witness row (spec §5). + await _provider.decide_txn(_memlake_txn, commit=True) + total_time = time.time() - start_time log_buffer.append(f"{'=' * 60}") log_buffer.append( diff --git a/hindsight-api-slim/hindsight_api/engine/search/retrieval.py b/hindsight-api-slim/hindsight_api/engine/search/retrieval.py index b2901d51aa..59935ad4aa 100644 --- a/hindsight-api-slim/hindsight_api/engine/search/retrieval.py +++ b/hindsight-api-slim/hindsight_api/engine/search/retrieval.py @@ -81,9 +81,23 @@ class SemanticBm25Result: def get_default_graph_retriever() -> GraphRetriever: - """Get or create the default graph retriever based on config.""" + """Get or create the default graph retriever. + + The memories store gets first refusal: the SQL retrievers walk `memory_links` + and `unit_entities`, so a store that keeps its links elsewhere has to supply + its own or the graph arm would silently return nothing. A store whose links + are in Postgres returns None and ``config.graph_retriever`` decides, as ever. + """ global _default_graph_retriever if _default_graph_retriever is None: + from ..memories import get_memories + + from_store = get_memories().graph_retriever() + if from_store is not None: + _default_graph_retriever = from_store + logger.info("Using the memories store's graph retriever") + return _default_graph_retriever + config = get_config() retriever_type = config.graph_retriever.lower() if retriever_type == "link_expansion": @@ -95,8 +109,12 @@ def get_default_graph_retriever() -> GraphRetriever: return _default_graph_retriever -def set_default_graph_retriever(retriever: GraphRetriever) -> None: - """Set the default graph retriever (for configuration/testing).""" +def set_default_graph_retriever(retriever: GraphRetriever | None) -> None: + """Set the default graph retriever (for configuration/testing). + + ``None`` clears the cache so the next call re-resolves it — used when the + memories store changes, since the retriever is chosen from it. + """ global _default_graph_retriever _default_graph_retriever = retriever @@ -116,6 +134,47 @@ async def retrieve_semantic_bm25_combined( min_semantic: float | None = None, min_keyword: float | None = None, graph_seed_min_similarity: float | None = None, +) -> dict[str, SemanticBm25Result]: + """Combined semantic + BM25 retrieval, run by the configured memories store. + + With the default Postgres store this calls straight through to + :func:`retrieve_semantic_bm25_combined_sql` below — same query, same results. + """ + from ..memories import get_memories + + return await get_memories().search( + conn=conn, + bank_id=bank_id, + fact_types=fact_types, + query_embedding=query_emb_str, + query_text=query_text, + limit=limit, + tags=tags, + tags_match=tags_match, + tag_groups=tag_groups, + created_after=created_after, + created_before=created_before, + min_semantic=min_semantic, + min_keyword=min_keyword, + graph_seed_min_similarity=graph_seed_min_similarity, + ) + + +async def retrieve_semantic_bm25_combined_sql( + conn, + query_emb_str: str, + query_text: str, + bank_id: str, + fact_types: list[str], + limit: int, + tags: list[str] | None = None, + tags_match: TagsMatch = "any", + tag_groups: list[TagGroup] | None = None, + created_after: datetime | None = None, + created_before: datetime | None = None, + min_semantic: float | None = None, + min_keyword: float | None = None, + graph_seed_min_similarity: float | None = None, ) -> dict[str, SemanticBm25Result]: """ Combined semantic + BM25 retrieval for multiple fact types in a single query. @@ -428,6 +487,45 @@ async def retrieve_temporal_combined( tag_groups: list[TagGroup] | None = None, created_after: datetime | None = None, created_before: datetime | None = None, +) -> dict[str, list[RetrievalResult]]: + """Temporal retrieval, run by the configured memories store. + + The timestamps live with the memories, so whoever holds them runs the arm. + With the default Postgres store this is :func:`retrieve_temporal_combined_sql`. + """ + from ..memories import get_memories + + return await get_memories().temporal_search( + conn=conn, + bank_id=bank_id, + fact_types=fact_types, + query_embedding=query_emb_str, + start_date=start_date, + end_date=end_date, + limit=budget, + semantic_threshold=semantic_threshold, + tags=tags, + tags_match=tags_match, + tag_groups=tag_groups, + created_after=created_after, + created_before=created_before, + ) + + +async def retrieve_temporal_combined_sql( + conn, + query_emb_str: str, + bank_id: str, + fact_types: list[str], + start_date: datetime, + end_date: datetime, + budget: int, + semantic_threshold: float = 0.1, + tags: list[str] | None = None, + tags_match: TagsMatch = "any", + tag_groups: list[TagGroup] | None = None, + created_after: datetime | None = None, + created_before: datetime | None = None, ) -> dict[str, list[RetrievalResult]]: """ Temporal retrieval for multiple fact types in a single query. diff --git a/hindsight-api-slim/hindsight_api/engine/storage/__init__.py b/hindsight-api-slim/hindsight_api/engine/storage/__init__.py index 5fde54e020..214efc601b 100644 --- a/hindsight-api-slim/hindsight_api/engine/storage/__init__.py +++ b/hindsight-api-slim/hindsight_api/engine/storage/__init__.py @@ -75,5 +75,21 @@ def create_file_storage( account_name=config.file_storage_azure_account_name, account_key=config.file_storage_azure_account_key, ) + elif storage_type == "memlake": + # Files go to memlake's document store (each file a single-body document). The target and + # namespace prefix default to the memories store's, so a memlake deployment stores facts, + # documents and files in one place; override with HINDSIGHT_API_FILE_STORAGE_MEMLAKE_TARGET. + import os + + from .memlake import MemlakeFileStorage + + target = ( + os.environ.get("HINDSIGHT_API_FILE_STORAGE_MEMLAKE_TARGET") + or os.environ.get("HINDSIGHT_API_MEMORIES_TARGET") + or "localhost:50051" + ).strip() + prefix = os.environ.get("HINDSIGHT_API_MEMORIES_NAMESPACE_PREFIX", "") + targets = [t.strip() for t in target.split(",")] if "," in target else target + return MemlakeFileStorage(target=targets, namespace_prefix=prefix) else: - raise ValueError(f"Unknown storage type: {storage_type}. Supported: 'native', 's3', 'gcs', 'azure'.") + raise ValueError(f"Unknown storage type: {storage_type}. Supported: 'native', 's3', 'gcs', 'azure', 'memlake'.") diff --git a/hindsight-api-slim/hindsight_api/engine/storage/memlake.py b/hindsight-api-slim/hindsight_api/engine/storage/memlake.py new file mode 100644 index 0000000000..0ac5fb786d --- /dev/null +++ b/hindsight-api-slim/hindsight_api/engine/storage/memlake.py @@ -0,0 +1,102 @@ +"""memlake-backed file storage. + +Stores each uploaded file as a single-body document in memlake's dedicated document store — the +same content-addressed blob machinery that holds documents' extracted text and chunks (see +docs/documents-chunks.md). ``file_storage`` stays the seam the rest of Hindsight uses; this is +simply a backend for it, selected with ``HINDSIGHT_API_FILE_STORAGE_TYPE=memlake``. + +Each storage key maps to one memlake "document" whose only body is the file: ``store`` uploads it +over a presigned PUT (dedup skips an unchanged re-upload), ``retrieve`` / ``get_download_url`` mint +a presigned GET, ``delete`` drops the record (its body is reclaimed by memlake's orphan sweep), and +``exists`` is a record lookup. The file lives in the file's OWN bank namespace (parsed from the +``banks/{bank_id}/…`` key), so dropping the bank drops its files too. Bodies move client↔S3 +directly; the gRPC server only ever carries small metadata. +""" + +from __future__ import annotations + +import asyncio +import hashlib +from functools import partial + +import memlake_client as mc # type: ignore[unresolved-import] + +from .base import FileStorage + + +def _sha256(data: bytes) -> str: + return hashlib.sha256(data).hexdigest() + + +class MemlakeFileStorage(FileStorage): + def __init__(self, target: "str | list[str]", namespace_prefix: str = "") -> None: + self._client = mc.MemlakeClient(target) + self._prefix = namespace_prefix + self._ensured: set[str] = set() + + def _resolve(self, key: str) -> tuple[str, str]: + """(namespace, document_id) for a storage key. The key is ``banks/{bank_id}/files/…``, so + the file lands in that bank's namespace; the document id is a key-safe hash of the full key + (prefixed so it never collides with a real document's record).""" + parts = key.split("/") + bank_id = parts[1] if len(parts) >= 2 and parts[0] == "banks" else "_files" + namespace = f"{self._prefix}{bank_id}" + document_id = "_file_" + _sha256(key.encode()) + return namespace, document_id + + async def _ensure(self, namespace: str) -> None: + if namespace in self._ensured: + return + await asyncio.to_thread(partial(self._client.create_namespace, namespace)) + self._ensured.add(namespace) + + async def store(self, file_data: bytes, key: str, metadata: dict[str, str] | None = None) -> str: + namespace, document_id = self._resolve(key) + await self._ensure(namespace) + meta = metadata or {} + file_hash = _sha256(file_data) + doc = mc.document( + document_id, + content_hash=file_hash, + file_hash=file_hash, + file_bytes=len(file_data), + file_content_type=meta.get("content_type", "application/octet-stream"), + file_original_name=meta.get("original_name", key.rsplit("/", 1)[-1]), + metadata={"storage_key": key}, + ) + resp = await asyncio.to_thread(partial(self._client.put_documents, namespace, [doc])) + # Upload the body only if memlake asked for it (missing — an identical re-store dedups). + if resp.uploads and resp.uploads[0].file.url: + await asyncio.to_thread(self._client.upload_blob, resp.uploads[0].file.url, file_data) + return key + + async def retrieve(self, key: str) -> bytes: + namespace, document_id = self._resolve(key) + resp = await asyncio.to_thread(partial(self._client.get_document, namespace, document_id, include_file=True)) + if not resp.found or not resp.file_url: + raise FileNotFoundError(key) + return await asyncio.to_thread(self._client.download_blob, resp.file_url) + + async def delete(self, key: str) -> None: + namespace, document_id = self._resolve(key) + await asyncio.to_thread(partial(self._client.delete_document, namespace, document_id)) + + async def exists(self, key: str) -> bool: + namespace, document_id = self._resolve(key) + resp = await asyncio.to_thread(partial(self._client.get_document, namespace, document_id)) + return bool(resp.found) + + async def get_download_url(self, key: str, expires_in: int = 3600) -> str: + namespace, document_id = self._resolve(key) + resp = await asyncio.to_thread( + partial( + self._client.get_document, + namespace, + document_id, + include_file=True, + url_ttl_seconds=expires_in, + ) + ) + if not resp.found or not resp.file_url: + raise FileNotFoundError(key) + return resp.file_url diff --git a/hindsight-api-slim/tests/test_bank_stats_cache_invalidation.py b/hindsight-api-slim/tests/test_bank_stats_cache_invalidation.py index 67384519b8..b748149bf8 100644 --- a/hindsight-api-slim/tests/test_bank_stats_cache_invalidation.py +++ b/hindsight-api-slim/tests/test_bank_stats_cache_invalidation.py @@ -13,49 +13,103 @@ """ import uuid +from datetime import datetime, timezone +from types import SimpleNamespace import pytest from hindsight_api import RequestContext from hindsight_api.engine.bank_stats_cache import BankStatsCache -from hindsight_api.engine.memory_engine import MemoryEngine +from hindsight_api.engine.memories import FactRecord, get_memories +from hindsight_api.engine.memory_engine import MemoryEngine, fq_table # A TTL long enough that, absent invalidation, the warmed cache would still be # served on the post-mutation read within the same test. _PINNED_TTL_SECONDS = 300.0 -async def _insert_memory(conn, bank_id: str, text: str, fact_type: str = "experience") -> uuid.UUID: - """Insert a memory unit directly, bypassing the LLM retain pipeline.""" - mem_id = uuid.uuid4() - await conn.execute( - """ - INSERT INTO memory_units (id, bank_id, text, fact_type, event_date, created_at, updated_at, consolidated_at) - VALUES ($1, $2, $3, $4, NOW(), NOW(), NOW(), NOW()) - """, - mem_id, - bank_id, - text, - fact_type, +async def _insert_memory( + memory: MemoryEngine, + conn, + bank_id: str, + text: str, + fact_type: str = "experience", + document_id: str | None = None, +) -> uuid.UUID: + """Seed one memory through the store, bypassing the LLM retain pipeline. + + Goes through ``insert_facts`` rather than an ``INSERT INTO memory_units`` so the fixture + seeds wherever memories actually live: a store that keeps them outside SQL never sees a + raw row, which would leave the bank empty and every assertion below reading zero. + """ + store = get_memories() + fact = SimpleNamespace( + fact_text=text, + embedding=memory.embeddings.encode([text])[0], + fact_type=fact_type, + tags=[], + context=None, + document_id=document_id, + chunk_id=None, + metadata=None, + observation_scopes=None, + entities=[], + causal_relations=[], + occurred_start=None, + occurred_end=None, + mentioned_at=None, ) - return mem_id + unit_ids = await store.insert_facts( + conn=conn, ops=memory._backend.ops, bank_id=bank_id, facts=[fact], document_id=document_id + ) + # The raw insert this replaces stamped consolidated_at, so these fixtures are not a + # consolidation backlog. Keep that: an unconsolidated source would change the stats read below. + await store.mark_consolidated( + conn=conn, fq_table=fq_table, bank_id=bank_id, unit_ids=unit_ids, when=datetime.now(timezone.utc) + ) + return uuid.UUID(unit_ids[0]) + +async def _insert_observation( + memory: MemoryEngine, conn, bank_id: str, text: str, source_memory_ids: list[uuid.UUID] +) -> uuid.UUID: + """Seed one observation, wherever the configured store keeps observations. -async def _insert_observation(conn, bank_id: str, text: str, source_memory_ids: list[uuid.UUID]) -> uuid.UUID: - """Insert an observation unit directly.""" + Gated because the two stores split this differently: the SQL store's ``upsert_observation`` + is a no-op (the consolidator writes that row inline), so SQL is seeded with the insert it + would have written; a store that owns its observations is seeded through the store. + """ + store = get_memories() obs_id = uuid.uuid4() - await conn.execute( - """ - INSERT INTO memory_units ( - id, bank_id, text, fact_type, event_date, source_memory_ids, proof_count, created_at, updated_at - ) VALUES ($1, $2, $3, 'observation', NOW(), $4, $5, NOW(), NOW()) - """, - obs_id, - bank_id, - text, - source_memory_ids, - len(source_memory_ids), - ) + if store.writes_memory_rows_in_sql: + await conn.execute( + """ + INSERT INTO memory_units ( + id, bank_id, text, fact_type, event_date, source_memory_ids, proof_count, created_at, updated_at + ) VALUES ($1, $2, $3, 'observation', NOW(), $4, $5, NOW(), NOW()) + """, + obs_id, + bank_id, + text, + source_memory_ids, + len(source_memory_ids), + ) + else: + now = datetime.now(timezone.utc) + await store.upsert_observation( + conn=conn, + bank_id=bank_id, + record=FactRecord( + unit_id=str(obs_id), + text=text, + embedding=memory.embeddings.encode([text])[0], + fact_type="observation", + proof_count=len(source_memory_ids), + source_memory_ids=[str(s) for s in source_memory_ids], + event_date=now, + created_at=now, + ), + ) return obs_id @@ -72,10 +126,6 @@ async def _insert_document(conn, bank_id: str, doc_id: str) -> None: ) -async def _attach_unit_to_doc(conn, unit_id: uuid.UUID, doc_id: str) -> None: - await conn.execute("UPDATE memory_units SET document_id = $1 WHERE id = $2", doc_id, unit_id) - - async def _ensure_bank(memory: MemoryEngine, bank_id: str, request_context: RequestContext) -> None: await memory.get_bank_profile(bank_id=bank_id, request_context=request_context) @@ -95,15 +145,15 @@ async def test_delete_memory_unit_invalidates_stats_cache( pool = await memory._get_pool() async with pool.acquire() as conn: - m1 = await _insert_memory(conn, bank_id, "Alice loves hiking.") - await _insert_memory(conn, bank_id, "Bob enjoys cycling.") + m1 = await _insert_memory(memory, conn, bank_id, "Alice loves hiking.") + await _insert_memory(memory, conn, bank_id, "Bob enjoys cycling.") _pin_cache(memory) try: before = await memory.get_bank_stats(bank_id, request_context=request_context) assert before["node_counts"].get("experience") == 2 - await memory.delete_memory_unit(str(m1), request_context=request_context) + await memory.delete_memory_unit(str(m1), bank_id=bank_id, request_context=request_context) after = await memory.get_bank_stats(bank_id, request_context=request_context) # Without invalidation the long-TTL cache would still report 2. @@ -120,8 +170,7 @@ async def test_delete_document_invalidates_stats_cache(self, memory: MemoryEngin pool = await memory._get_pool() async with pool.acquire() as conn: await _insert_document(conn, bank_id, document_id) - unit_id = await _insert_memory(conn, bank_id, "Alice works at Acme.") - await _attach_unit_to_doc(conn, unit_id, document_id) + await _insert_memory(memory, conn, bank_id, "Alice works at Acme.", document_id=document_id) _pin_cache(memory) try: @@ -147,8 +196,8 @@ async def test_clear_observations_invalidates_stats_cache( pool = await memory._get_pool() async with pool.acquire() as conn: - m1 = await _insert_memory(conn, bank_id, "Alice loves hiking.") - await _insert_observation(conn, bank_id, "Alice enjoys hiking regularly.", [m1]) + m1 = await _insert_memory(memory, conn, bank_id, "Alice loves hiking.") + await _insert_observation(memory, conn, bank_id, "Alice enjoys hiking regularly.", [m1]) _pin_cache(memory) try: diff --git a/hindsight-api-slim/tests/test_graph_maintenance.py b/hindsight-api-slim/tests/test_graph_maintenance.py index b1c7892c95..9bcf70830a 100644 --- a/hindsight-api-slim/tests/test_graph_maintenance.py +++ b/hindsight-api-slim/tests/test_graph_maintenance.py @@ -12,20 +12,14 @@ from __future__ import annotations import uuid -from contextlib import asynccontextmanager from datetime import UTC, datetime, timedelta -from types import SimpleNamespace -from unittest.mock import AsyncMock, patch import pytest -import hindsight_api.engine.graph_maintenance as graph_maintenance_module -import hindsight_api.engine.memory_engine as memory_engine_module from hindsight_api import RequestContext from hindsight_api.engine.graph_maintenance import ( MAX_SEMANTIC_LINKS_PER_UNIT, MAX_TEMPORAL_LINKS_PER_UNIT, - _relink_batch, enqueue_relink_victims, run_graph_maintenance_job, ) @@ -234,86 +228,6 @@ async def test_dedupes_via_on_conflict(self, memory: MemoryEngine, request_conte assert await _queue_unit_ids(conn, bank_id) == [str(survivor)] - @pytest.mark.asyncio - async def test_include_affected_enqueues_self_without_victims( - self, memory: MemoryEngine, request_context: RequestContext - ): - """Outgoing-only unit: nothing points at it, so the victim lookup is empty. - - This is the case that made an edit a silent no-op (#2889) — the helper - returned 0, the queue stayed empty and submission short-circuited on - ``no_work``, so the edited unit's own outgoing links were never rebuilt. - """ - bank_id = f"test-gm-self-live-{uuid.uuid4().hex[:8]}" - await _ensure_bank(memory, bank_id, request_context) - - pool = await memory._get_pool() - async with pool.acquire() as conn: - edited = await _insert_unit(conn, bank_id, "edited") - target = await _insert_unit(conn, bank_id, "target") - # edited → target only; no incoming derived links. - await _insert_link(conn, bank_id, edited, target, "temporal") - - backend = await memory._get_backend() - async with conn.transaction(): - count = await enqueue_relink_victims( - conn, bank_id, [str(edited)], ops=backend.ops, include_affected_units=True - ) - - assert count == 1 - assert await _queue_unit_ids(conn, bank_id) == [str(edited)] - - @pytest.mark.asyncio - async def test_include_affected_combines_self_and_victims_in_one_insert( - self, memory: MemoryEngine, request_context: RequestContext - ): - """Mutually linked units share a single sorted insert. - - Splitting them across two inserts would let concurrent edits take the - ``(bank_id, unit_id)`` keys in opposite orders and deadlock. - """ - bank_id = f"test-gm-self-both-{uuid.uuid4().hex[:8]}" - await _ensure_bank(memory, bank_id, request_context) - - pool = await memory._get_pool() - async with pool.acquire() as conn: - a = await _insert_unit(conn, bank_id, "a") - b = await _insert_unit(conn, bank_id, "b") - await _insert_link(conn, bank_id, a, b, "temporal") - await _insert_link(conn, bank_id, b, a, "temporal") - - backend = await memory._get_backend() - with patch.object( - backend.ops, "enqueue_graph_maintenance", wraps=backend.ops.enqueue_graph_maintenance - ) as enqueue_mock: - async with conn.transaction(): - count = await enqueue_relink_victims( - conn, bank_id, [str(a)], ops=backend.ops, include_affected_units=True - ) - - assert count == 2 - assert enqueue_mock.await_count == 1, "self and victims must share one lock-ordered insert" - assert await _queue_unit_ids(conn, bank_id) == sorted([str(a), str(b)]) - - @pytest.mark.asyncio - async def test_include_affected_is_opt_in(self, memory: MemoryEngine, request_context: RequestContext): - """Delete callers keep the default: a doomed unit must not queue itself.""" - bank_id = f"test-gm-optin-{uuid.uuid4().hex[:8]}" - await _ensure_bank(memory, bank_id, request_context) - - pool = await memory._get_pool() - async with pool.acquire() as conn: - doomed = await _insert_unit(conn, bank_id, "doomed") - target = await _insert_unit(conn, bank_id, "target") - await _insert_link(conn, bank_id, doomed, target, "temporal") - - backend = await memory._get_backend() - async with conn.transaction(): - count = await enqueue_relink_victims(conn, bank_id, [str(doomed)], ops=backend.ops) - - assert count == 0 - assert await _queue_unit_ids(conn, bank_id) == [] - # --------------------------------------------------------------------------- # delete_document hook @@ -353,54 +267,6 @@ async def test_delete_document_enqueues_cross_doc_victims( class TestRelinkPass: - @pytest.mark.asyncio - async def test_semantic_topup_uses_configured_link_threshold(self, monkeypatch): - """Maintenance relinking must use the same construction gate as retain.""" - victim_id = str(uuid.uuid4()) - conn = SimpleNamespace( - fetch=AsyncMock( - side_effect=[ - [ - { - "id": victim_id, - "event_date": None, - "fact_type": "world", - "embedding": "[1.0]", - } - ], - [], - ] - ) - ) - captured_thresholds: list[float] = [] - - @asynccontextmanager - async def fake_acquire_with_retry(_backend): - yield object() - - async def fake_compute_semantic_links_ann(*_args, **kwargs): - captured_thresholds.append(kwargs["threshold"]) - return [] - - monkeypatch.setattr(memory_engine_module, "acquire_with_retry", fake_acquire_with_retry) - monkeypatch.setattr( - graph_maintenance_module, - "compute_semantic_links_ann", - fake_compute_semantic_links_ann, - ) - - added = await _relink_batch( - conn=conn, - bank_id="bank", - victim_ids=[victim_id], - ops=object(), - backend=object(), - semantic_link_min_similarity=0.86, - ) - - assert added == 0 - assert captured_thresholds == [0.86] - @pytest.mark.asyncio async def test_drains_empty_queue_cleanly(self, memory: MemoryEngine, request_context: RequestContext): bank_id = f"test-gm-empty-{uuid.uuid4().hex[:8]}" diff --git a/hindsight-api-slim/tests/test_memories_extension.py b/hindsight-api-slim/tests/test_memories_extension.py new file mode 100644 index 0000000000..1a4c131d76 --- /dev/null +++ b/hindsight-api-slim/tests/test_memories_extension.py @@ -0,0 +1,325 @@ +"""The memories store is an extension point, and the default is Postgres. + +Two things are worth pinning down here, and neither is about SQL: + +1. **The default is unconditional.** With nothing configured the engine gets + :class:`PostgresMemories`, so every other test in this suite is exercising the + real store rather than a seam that happens to fall through to it. +2. **The interface is complete.** A store that implements + :class:`MemoriesExtension` and touches no database at all can be installed and + used. That is the property the whole extraction exists for: if a call site + still reached past the interface for a `memory_units` row, the stub below + would not be able to answer and the test would fail rather than quietly + falling back to SQL. + +The stub is deliberately a dictionary. Anything cleverer would start re-testing +storage instead of the seam. +""" + +from __future__ import annotations + +import os +from datetime import datetime, timezone +from unittest.mock import patch + +import pytest + +from hindsight_api.engine.memories import create_memories, get_memories, set_memories +from hindsight_api.engine.memories.base import MemoriesExtension, ScanPage, StoredMemory +from hindsight_api.engine.memories.postgres import PostgresMemories + + +class InMemoryMemories(MemoriesExtension): + """A complete store that is a dict — no connection, no tables, no SQL. + + Every method that touches storage is answered from ``self.rows``. The + Postgres handles (``conn``, ``ops``, ``fq_table``) are accepted and ignored, + which is exactly what an implementation that owns the store does with them. + """ + + name = "in-memory" + + def __init__(self, config: dict[str, str] | None = None): + super().__init__(config or {}) + self.rows: dict[str, StoredMemory] = {} + # The curation archive is just a second dict — invalidation moves a memory + # from `rows` to `archive`, exactly as it moves between tables/namespaces. + self.archive: dict[str, StoredMemory] = {} + self.invalidation_reason: dict[str, str | None] = {} + self.embeddings: dict[str, object] = {} + # Proof the engine went through the interface rather than around it. + self.calls: list[str] = [] + + # -- writes -------------------------------------------------------------- + + async def insert_facts(self, *, conn, ops, bank_id, facts, document_id=None, defer_index=False): + self.calls.append("insert_facts") + unit_ids = self.allocate_unit_ids(len(facts)) + if not defer_index: + await self.index_facts(bank_id, unit_ids, facts, document_id) + return unit_ids + + async def index_facts(self, bank_id, unit_ids, facts, document_id=None, unit_entity_ids=None): + self.calls.append("index_facts") + for unit_id, fact in zip(unit_ids, facts): + self.rows[unit_id] = StoredMemory( + unit_id=unit_id, + text=fact.fact_text, + fact_type=fact.fact_type, + document_id=document_id, + tags=list(fact.tags or []), + created_at=datetime.now(timezone.utc), + ) + + async def delete_facts(self, bank_id, unit_ids): + self.calls.append("delete_facts") + for unit_id in unit_ids: + self.rows.pop(str(unit_id), None) + + async def delete_document(self, *, conn, fq_table, bank_id, document_id): + self.calls.append("delete_document") + for unit_id in [k for k, v in self.rows.items() if v.document_id == document_id]: + del self.rows[unit_id] + + async def delete_observations(self, *, conn, fq_table, bank_id): + for unit_id in [k for k, v in self.rows.items() if v.fact_type == "observation"]: + del self.rows[unit_id] + + # -- recall arms --------------------------------------------------------- + + async def search(self, *, conn, bank_id, fact_types, query_embedding, query_text, limit, **kwargs): + self.calls.append("search") + return {ft: ([], []) for ft in fact_types} + + async def temporal_search( + self, *, conn, bank_id, fact_types, query_embedding, start_date, end_date, limit, **kwargs + ): + self.calls.append("temporal_search") + return {ft: [] for ft in fact_types} + + # -- addressed reads ----------------------------------------------------- + + async def get_memories(self, *, conn, fq_table, bank_id, unit_ids): + self.calls.append("get_memories") + return [self.rows[str(u)] for u in unit_ids if str(u) in self.rows] + + async def scan_memories(self, *, conn, fq_table, bank_id, limit=100, page_token="", **kwargs): + start = int(page_token or 0) + ordered = list(self.rows.values())[start : start + limit] + nxt = str(start + limit) if start + limit < len(self.rows) else "" + return ScanPage(memories=ordered, next_page_token=nxt) + + async def count_memories(self, *, conn, fq_table, bank_id): + counts: dict[str, int] = {} + for row in self.rows.values(): + counts[row.fact_type] = counts.get(row.fact_type, 0) + 1 + return counts + + async def list_tags(self, *, conn, fq_table, bank_id): + counts: dict[str, int] = {} + for row in self.rows.values(): + for tag in row.tags: + counts[tag] = counts.get(tag, 0) + 1 + return counts + + async def find_unconsolidated(self, *, conn, fq_table, bank_id, fact_types, limit, scope_tags=None): + out = [r for r in self.rows.values() if r.fact_type in fact_types and r.consolidated_at is None] + if scope_tags: + out = [r for r in out if set(scope_tags).issubset(set(r.tags))] + return out[:limit] + + async def mark_consolidated(self, *, conn, fq_table, bank_id, unit_ids, when, failed=False): + for unit_id in unit_ids: + row = self.rows.get(str(unit_id)) + if row is not None: + row.consolidated_at = when + + async def entity_memory_counts(self, *, conn, fq_table, bank_id, entity_ids=None): + counts: dict[str, int] = {} + for row in self.rows.values(): + for entity_id in row.entity_ids: + counts[entity_id] = counts.get(entity_id, 0) + 1 + return counts if entity_ids is None else {k: v for k, v in counts.items() if k in set(entity_ids)} + + async def entities_for_units(self, *, conn, fq_table, bank_id, unit_ids): + return {str(u): list(self.rows[str(u)].entity_ids) for u in unit_ids if str(u) in self.rows} + + async def entity_map_for_units(self, *, conn, fq_table, bank_id, unit_ids): + # No names in the stub — the entity registry is Postgres's, which the + # stub does not stand in for. The shape is what recall consumes. + return { + str(u): [{"entity_id": e, "canonical_name": e} for e in self.rows[str(u)].entity_ids] + for u in unit_ids + if str(u) in self.rows + } + + async def any_memory_updated_since( + self, *, conn, fq_table, bank_id, since, fact_types=None, tags=None, tags_match="any", tag_groups=None + ): + rows = self.rows.values() + if fact_types: + rows = [r for r in rows if r.fact_type in fact_types] + return any(r.created_at is not None and r.created_at > since for r in rows) + + # -- observations -------------------------------------------------------- + + async def observations_for_sources(self, *, conn, ops, fq_table, bank_id, unit_ids): + wanted = {str(u) for u in unit_ids} + return [r for r in self.rows.values() if wanted & set(r.source_memory_ids)] + + async def delete_stale_observations(self, *, conn, ops, fq_table, bank_id, fact_ids): + stale = await self.observations_for_sources( + conn=conn, ops=ops, fq_table=fq_table, bank_id=bank_id, unit_ids=fact_ids + ) + for obs in stale: + self.rows.pop(obs.unit_id, None) + return len(stale) + + # -- curation reads ------------------------------------------------------ + + async def list_memory_units(self, *, conn, ops, fq_table, bank_id, limit=100, offset=0, **kwargs): + ordered = list(self.rows.values()) + return {"items": ordered[offset : offset + limit], "total": len(ordered), "limit": limit, "offset": offset} + + async def get_memory_unit(self, *, conn, ops, fq_table, bank_id, unit_id): + row = self.rows.get(str(unit_id)) + return None if row is None else {"id": row.unit_id, "text": row.text, "fact_type": row.fact_type} + + # -- curation archive: a second dict is the archive namespace ------------ + + async def get_archived_memory(self, *, conn, fq_table, bank_id, unit_id): + return self.archive.get(str(unit_id)) + + async def invalidate_memory(self, *, conn, fq_table, bank_id, unit_id, reason): + row = self.rows.pop(str(unit_id), None) + if row is None: + return False + self.archive[str(unit_id)] = row + self.invalidation_reason[str(unit_id)] = reason + return True + + async def set_invalidation_reason(self, *, conn, fq_table, bank_id, unit_id, reason): + self.invalidation_reason[str(unit_id)] = reason + + async def restore_memory(self, *, conn, fq_table, bank_id, unit_id): + row = self.archive.pop(str(unit_id), None) + if row is None: + return None + self.rows[str(unit_id)] = row + self.invalidation_reason.pop(str(unit_id), None) + return row + + async def set_memory_embedding(self, *, conn, fq_table, bank_id, unit_id, embedding): + self.embeddings[str(unit_id)] = embedding + + async def list_entities(self, *, conn, fq_table, bank_id, search=None, limit=100, offset=0): + return {"items": [], "total": 0, "limit": limit, "offset": offset} + + async def graph_units(self, *, conn, fq_table, bank_id, limit=1000, **kwargs): + rows = list(self.rows.values())[:limit] + return {"units": [{"id": r.unit_id, "fact_type": r.fact_type} for r in rows], "total": len(self.rows)} + + async def graph_entity_rows(self, *, conn, fq_table, bank_id, unit_ids): + return [] + + async def graph_direct_links(self, *, conn, fq_table, bank_id, unit_ids): + return [] + + +@pytest.fixture +def restore_default_store(): + """Put the process-wide store back, whatever a test did to it.""" + yield + set_memories(None) + + +def test_default_store_is_postgres(restore_default_store): + """Nothing configured means the SQL path — which is what every other test runs.""" + set_memories(None) + with patch.dict(os.environ, {}, clear=False): + os.environ.pop("HINDSIGHT_API_MEMORIES_EXTENSION", None) + assert isinstance(create_memories(), PostgresMemories) + assert get_memories().name == "postgres" + + +def test_a_configured_store_replaces_the_default(restore_default_store): + """The ordinary extension env var selects the store, like every other extension.""" + set_memories(None) + spec = f"{InMemoryMemories.__module__}:{InMemoryMemories.__name__}" + with patch.dict(os.environ, {"HINDSIGHT_API_MEMORIES_EXTENSION": spec}): + store = create_memories() + assert isinstance(store, InMemoryMemories) + assert store.name == "in-memory" + + +def test_the_store_receives_its_prefixed_config(restore_default_store): + """`HINDSIGHT_API_MEMORIES_*` reaches the store, stripped and lowercased.""" + set_memories(None) + spec = f"{InMemoryMemories.__module__}:{InMemoryMemories.__name__}" + env = { + "HINDSIGHT_API_MEMORIES_EXTENSION": spec, + "HINDSIGHT_API_MEMORIES_TARGET": "example:50051", + "HINDSIGHT_API_MEMORIES_NPROBE": "16", + } + with patch.dict(os.environ, env): + store = create_memories() + assert store.config["target"] == "example:50051" + assert store.config["nprobe"] == "16" + assert "extension" not in store.config, "the selector must not leak into the store's own config" + + +def test_the_interface_is_implementable_without_a_database(restore_default_store): + """The point of the extraction: a store with no SQL behind it is a valid store. + + Instantiating an ABC with a missing method raises `TypeError` naming it, so a + method added to the interface without a home here fails loudly rather than at + the first call site that needs it. + """ + store = InMemoryMemories({}) + assert isinstance(store, MemoriesExtension) + + +async def test_a_store_that_owns_its_rows_needs_no_postgres(restore_default_store): + """Write, read back, and delete — with `conn` set to something unusable. + + Passing `None` where the Postgres store would expect a connection is the + assertion: any code path that quietly reached for SQL would raise instead of + returning the rows the store holds. + """ + store = InMemoryMemories({}) + set_memories(store) + + class _Fact: + fact_text = "the cat sat on the mat" + fact_type = "world" + tags = ["animals"] + + unit_ids = await store.insert_facts(conn=None, ops=None, bank_id="bank", facts=[_Fact()], document_id="doc-1") + assert len(unit_ids) == 1 + + got = await store.get_memories(conn=None, fq_table=None, bank_id="bank", unit_ids=unit_ids) + assert [m.text for m in got] == ["the cat sat on the mat"] + assert await store.count_memories(conn=None, fq_table=None, bank_id="bank") == {"world": 1} + assert await store.list_tags(conn=None, fq_table=None, bank_id="bank") == {"animals": 1} + + # Deleting the document takes its memories with it, with no cascade to rely on. + await store.delete_document(conn=None, fq_table=None, bank_id="bank", document_id="doc-1") + assert await store.get_memories(conn=None, fq_table=None, bank_id="bank", unit_ids=unit_ids) == [] + assert "insert_facts" in store.calls and "get_memories" in store.calls + + +async def test_maintenance_passes_are_optional(restore_default_store): + """A store with inline links has nothing to relink and no join table to sweep. + + These have safe base implementations precisely so such a store does not have + to write four no-op methods to be complete — and so the maintenance job can + call them unconditionally. + """ + store = InMemoryMemories({}) + assert await store.enqueue_relink_victims(conn=None, fq_table=None, bank_id="b", affected_unit_ids=["x"]) == 0 + assert await store.relink_pass(backend=None, fq_table=None, bank_id="b", config=None) == {} + assert await store.prune_orphan_entities(conn=None, fq_table=None, bank_id="b") == 0 + assert await store.prune_stale_cooccurrences(conn=None, fq_table=None, bank_id="b") == 0 + # And recording entity postings is a no-op rather than an error: the posting + # travels on the memory for a store that owns it. + await store.record_unit_entities(conn=None, ops=None, fq_table=None, unit_ids=["u"], entity_ids=["e"]) diff --git a/hindsight-api-slim/tests/test_observation_invalidation.py b/hindsight-api-slim/tests/test_observation_invalidation.py index f3f4c6b9fe..48c5f25ed2 100644 --- a/hindsight-api-slim/tests/test_observation_invalidation.py +++ b/hindsight-api-slim/tests/test_observation_invalidation.py @@ -10,65 +10,149 @@ """ import uuid +from datetime import datetime, timezone +from types import SimpleNamespace from unittest.mock import AsyncMock, patch import pytest from hindsight_api import RequestContext -from hindsight_api.engine.memory_engine import MemoryEngine +from hindsight_api.engine.memories import FactRecord, get_memories +from hindsight_api.engine.memory_engine import MemoryEngine, fq_table # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- -async def _insert_memory(conn, bank_id: str, text: str, fact_type: str = "experience") -> uuid.UUID: - """Insert a memory unit directly, bypassing LLM retain pipeline.""" - mem_id = uuid.uuid4() - await conn.execute( - """ - INSERT INTO memory_units (id, bank_id, text, fact_type, event_date, created_at, updated_at, consolidated_at) - VALUES ($1, $2, $3, $4, NOW(), NOW(), NOW(), NOW()) - """, - mem_id, - bank_id, - text, - fact_type, +async def _insert_memory( + memory: MemoryEngine, + conn, + bank_id: str, + text: str, + fact_type: str = "experience", + document_id: str | None = None, +) -> uuid.UUID: + """Seed one memory through the store, bypassing the LLM retain pipeline. + + Uses insert_facts rather than an INSERT INTO memory_units so the fixture seeds wherever + memories actually live — a store that keeps them outside SQL never sees a raw row. + """ + store = get_memories() + fact = SimpleNamespace( + fact_text=text, + embedding=memory.embeddings.encode([text])[0], + fact_type=fact_type, + tags=[], + context=None, + document_id=document_id, + chunk_id=None, + metadata=None, + observation_scopes=None, + entities=[], + causal_relations=[], + occurred_start=None, + occurred_end=None, + mentioned_at=None, ) - return mem_id + unit_ids = await store.insert_facts( + conn=conn, ops=memory._backend.ops, bank_id=bank_id, facts=[fact], document_id=document_id + ) + # The raw insert this replaces stamped consolidated_at; keep that so these fixtures are a + # consolidated baseline, not a backlog (several tests assert the reset back to NULL). + await store.mark_consolidated( + conn=conn, fq_table=fq_table, bank_id=bank_id, unit_ids=unit_ids, when=datetime.now(timezone.utc) + ) + return uuid.UUID(unit_ids[0]) + +async def _insert_observation( + memory: MemoryEngine, conn, bank_id: str, text: str, source_memory_ids: list[uuid.UUID] +) -> uuid.UUID: + """Seed one observation wherever the configured store keeps observations. -async def _insert_observation(conn, bank_id: str, text: str, source_memory_ids: list[uuid.UUID]) -> uuid.UUID: - """Insert an observation unit directly.""" + Gated: the SQL store's upsert_observation is a no-op (the consolidator writes that row + inline), so SQL is seeded with the insert it would have written. + """ + store = get_memories() obs_id = uuid.uuid4() - await conn.execute( - """ - INSERT INTO memory_units ( - id, bank_id, text, fact_type, event_date, source_memory_ids, proof_count, created_at, updated_at - ) VALUES ($1, $2, $3, 'observation', NOW(), $4, $5, NOW(), NOW()) - """, - obs_id, - bank_id, - text, - source_memory_ids, - len(source_memory_ids), - ) + if store.writes_memory_rows_in_sql: + await conn.execute( + """ + INSERT INTO memory_units ( + id, bank_id, text, fact_type, event_date, source_memory_ids, proof_count, created_at, updated_at + ) VALUES ($1, $2, $3, 'observation', NOW(), $4, $5, NOW(), NOW()) + """, + obs_id, + bank_id, + text, + source_memory_ids, + len(source_memory_ids), + ) + else: + now = datetime.now(timezone.utc) + await store.upsert_observation( + conn=conn, + bank_id=bank_id, + record=FactRecord( + unit_id=str(obs_id), + text=text, + embedding=memory.embeddings.encode([text])[0], + fact_type="observation", + proof_count=len(source_memory_ids), + source_memory_ids=[str(s) for s in source_memory_ids], + event_date=now, + created_at=now, + ), + ) return obs_id async def _get_observation_ids(conn, bank_id: str) -> list[str]: - rows = await conn.fetch( - "SELECT id FROM memory_units WHERE bank_id = $1 AND fact_type = 'observation'", - bank_id, + """Ids of the bank's observations, read from whichever store holds them.""" + store = get_memories() + if store.writes_memory_rows_in_sql: + rows = await conn.fetch( + "SELECT id FROM memory_units WHERE bank_id = $1 AND fact_type = 'observation'", + bank_id, + ) + return [str(r["id"]) for r in rows] + page = await store.scan_memories( + conn=conn, fq_table=fq_table, bank_id=bank_id, fact_types=["observation"], limit=1_000_000 + ) + return [m.unit_id for m in page.memories] + + +async def _get_consolidated_at(conn, memory_id: uuid.UUID, bank_id: str | None = None): + """A memory's consolidated marker. ``bank_id`` is required for a bank-partitioned store.""" + store = get_memories() + if store.writes_memory_rows_in_sql: + return await conn.fetchval( + "SELECT consolidated_at FROM memory_units WHERE id = $1", + memory_id, + ) + found = ( + await store.get_memories(conn=conn, fq_table=fq_table, bank_id=bank_id, unit_ids=[str(memory_id)]) + if bank_id + else [] ) - return [str(r["id"]) for r in rows] + return found[0].consolidated_at if found else None -async def _get_consolidated_at(conn, memory_id: uuid.UUID): - return await conn.fetchval( - "SELECT consolidated_at FROM memory_units WHERE id = $1", - memory_id, +async def _get_memory(conn, bank_id: str, memory_id): + """One stored memory, read through the store — works for either backend.""" + store = get_memories() + found = await store.get_memories(conn=conn, fq_table=fq_table, bank_id=bank_id, unit_ids=[str(memory_id)]) + return found[0] if found else None + + +async def _count_surviving(conn, bank_id: str, memory_ids: list) -> int: + """How many of these ids still exist, per the store that holds them.""" + store = get_memories() + found = await store.get_memories( + conn=conn, fq_table=fq_table, bank_id=bank_id, unit_ids=[str(m) for m in memory_ids] ) + return len(found) async def _ensure_bank(memory: MemoryEngine, bank_id: str, request_context: RequestContext): @@ -91,11 +175,11 @@ async def test_deleting_source_memory_removes_observation( pool = await memory._get_pool() async with pool.acquire() as conn: - m1 = await _insert_memory(conn, bank_id, "Alice loves hiking.") - m2 = await _insert_memory(conn, bank_id, "Alice goes hiking every weekend.") - obs_id = await _insert_observation(conn, bank_id, "Alice enjoys hiking regularly.", [m1, m2]) + m1 = await _insert_memory(memory, conn, bank_id, "Alice loves hiking.") + m2 = await _insert_memory(memory, conn, bank_id, "Alice goes hiking every weekend.") + obs_id = await _insert_observation(memory, conn, bank_id, "Alice enjoys hiking regularly.", [m1, m2]) - await memory.delete_memory_unit(str(m1), request_context=request_context) + await memory.delete_memory_unit(str(m1), bank_id=bank_id, request_context=request_context) async with pool.acquire() as conn: obs_ids = await _get_observation_ids(conn, bank_id) @@ -113,20 +197,20 @@ async def test_deleting_source_memory_resets_remaining_source_consolidated_at( pool = await memory._get_pool() async with pool.acquire() as conn: - m1 = await _insert_memory(conn, bank_id, "Alice loves hiking.") - m2 = await _insert_memory(conn, bank_id, "Alice goes hiking every weekend.") - await _insert_observation(conn, bank_id, "Alice enjoys hiking regularly.", [m1, m2]) + m1 = await _insert_memory(memory, conn, bank_id, "Alice loves hiking.") + m2 = await _insert_memory(memory, conn, bank_id, "Alice goes hiking every weekend.") + await _insert_observation(memory, conn, bank_id, "Alice enjoys hiking regularly.", [m1, m2]) # Verify m2 starts with consolidated_at set - assert await _get_consolidated_at(conn, m2) is not None + assert await _get_consolidated_at(conn, m2, bank_id) is not None # Patch out consolidation so it doesn't re-set consolidated_at before we can check it with patch.object(memory, "submit_async_consolidation", new=AsyncMock()): - await memory.delete_memory_unit(str(m1), request_context=request_context) + await memory.delete_memory_unit(str(m1), bank_id=bank_id, request_context=request_context) async with pool.acquire() as conn: # m2 should have consolidated_at reset to NULL - consolidated_at = await _get_consolidated_at(conn, m2) + consolidated_at = await _get_consolidated_at(conn, m2, bank_id) assert consolidated_at is None, "Remaining source memory should be reset for re-consolidation" await memory.delete_bank(bank_id, request_context=request_context) @@ -141,19 +225,19 @@ async def test_deleting_non_source_memory_leaves_observations_intact( pool = await memory._get_pool() async with pool.acquire() as conn: - m1 = await _insert_memory(conn, bank_id, "Alice loves hiking.") - m2 = await _insert_memory(conn, bank_id, "Alice goes hiking every weekend.") - unrelated = await _insert_memory(conn, bank_id, "Bob likes cycling.") - obs_id = await _insert_observation(conn, bank_id, "Alice enjoys hiking regularly.", [m1, m2]) + m1 = await _insert_memory(memory, conn, bank_id, "Alice loves hiking.") + m2 = await _insert_memory(memory, conn, bank_id, "Alice goes hiking every weekend.") + unrelated = await _insert_memory(memory, conn, bank_id, "Bob likes cycling.") + obs_id = await _insert_observation(memory, conn, bank_id, "Alice enjoys hiking regularly.", [m1, m2]) - await memory.delete_memory_unit(str(unrelated), request_context=request_context) + await memory.delete_memory_unit(str(unrelated), bank_id=bank_id, request_context=request_context) async with pool.acquire() as conn: obs_ids = await _get_observation_ids(conn, bank_id) assert str(obs_id) in obs_ids, "Observation should remain untouched" # m1 and m2 should still be consolidated - assert await _get_consolidated_at(conn, m1) is not None - assert await _get_consolidated_at(conn, m2) is not None + assert await _get_consolidated_at(conn, m1, bank_id) is not None + assert await _get_consolidated_at(conn, m2, bank_id) is not None await memory.delete_bank(bank_id, request_context=request_context) @@ -167,10 +251,10 @@ async def test_deleting_sole_source_memory_removes_observation_no_remaining_rese pool = await memory._get_pool() async with pool.acquire() as conn: - m1 = await _insert_memory(conn, bank_id, "Alice loves hiking.") - obs_id = await _insert_observation(conn, bank_id, "Alice enjoys hiking.", [m1]) + m1 = await _insert_memory(memory, conn, bank_id, "Alice loves hiking.") + obs_id = await _insert_observation(memory, conn, bank_id, "Alice enjoys hiking.", [m1]) - await memory.delete_memory_unit(str(m1), request_context=request_context) + await memory.delete_memory_unit(str(m1), bank_id=bank_id, request_context=request_context) async with pool.acquire() as conn: obs_ids = await _get_observation_ids(conn, bank_id) @@ -188,15 +272,15 @@ async def test_deleting_observation_type_memory_does_not_trigger_invalidation( pool = await memory._get_pool() async with pool.acquire() as conn: - m1 = await _insert_memory(conn, bank_id, "Alice loves hiking.") - obs_id = await _insert_observation(conn, bank_id, "Alice enjoys hiking.", [m1]) + m1 = await _insert_memory(memory, conn, bank_id, "Alice loves hiking.") + obs_id = await _insert_observation(memory, conn, bank_id, "Alice enjoys hiking.", [m1]) # Delete the observation directly (not the source memory) - await memory.delete_memory_unit(str(obs_id), request_context=request_context) + await memory.delete_memory_unit(str(obs_id), bank_id=bank_id, request_context=request_context) async with pool.acquire() as conn: # Source memory should still be consolidated (not reset) - assert await _get_consolidated_at(conn, m1) is not None + assert await _get_consolidated_at(conn, m1, bank_id) is not None obs_ids = await _get_observation_ids(conn, bank_id) assert str(obs_id) not in obs_ids @@ -228,25 +312,16 @@ async def test_deleting_document_removes_observations(self, memory: MemoryEngine doc_id, bank_id, ) - m1 = uuid.uuid4() - m2 = uuid.uuid4() - for mem_id, text in [(m1, "Alice loves hiking."), (m2, "Alice goes hiking every weekend.")]: - await conn.execute( - """ - INSERT INTO memory_units (id, bank_id, text, fact_type, event_date, document_id, created_at, updated_at, consolidated_at) - VALUES ($1, $2, $3, 'experience', NOW(), $4, NOW(), NOW(), NOW()) - """, - mem_id, - bank_id, - text, - doc_id, - ) + m1 = await _insert_memory(memory, conn, bank_id, "Alice loves hiking.", "experience", document_id=doc_id) + m2 = await _insert_memory( + memory, conn, bank_id, "Alice goes hiking every weekend.", "experience", document_id=doc_id + ) # Standalone memory (not in document) - m3 = await _insert_memory(conn, bank_id, "Alice is an avid outdoor person.") + m3 = await _insert_memory(memory, conn, bank_id, "Alice is an avid outdoor person.") # Observation referencing both doc memories and the standalone memory - obs_id = await _insert_observation(conn, bank_id, "Alice enjoys outdoor activities.", [m1, m2, m3]) + obs_id = await _insert_observation(memory, conn, bank_id, "Alice enjoys outdoor activities.", [m1, m2, m3]) # Patch out consolidation so it doesn't re-set consolidated_at before we can check it with patch.object(memory, "submit_async_consolidation", new=AsyncMock()): @@ -257,7 +332,7 @@ async def test_deleting_document_removes_observations(self, memory: MemoryEngine assert str(obs_id) not in obs_ids, "Observation should have been deleted" # m3 (remaining source) should be reset for re-consolidation - consolidated_at = await _get_consolidated_at(conn, m3) + consolidated_at = await _get_consolidated_at(conn, m3, bank_id) assert consolidated_at is None, "Remaining source memory should be reset" await memory.delete_bank(bank_id, request_context=request_context) @@ -306,22 +381,11 @@ async def test_upsert_document_removes_observations_from_outgoing_memories( doc_id, bank_id, ) - doc_mem_a = uuid.uuid4() - doc_mem_b = uuid.uuid4() - for mem_id, text in [(doc_mem_a, "Old fact A."), (doc_mem_b, "Old fact B.")]: - await conn.execute( - """ - INSERT INTO memory_units (id, bank_id, text, fact_type, event_date, document_id, - created_at, updated_at, consolidated_at) - VALUES ($1, $2, $3, 'experience', NOW(), $4, NOW(), NOW(), NOW()) - """, - mem_id, - bank_id, - text, - doc_id, - ) - standalone_mem = await _insert_memory(conn, bank_id, "Standalone fact C.") + doc_mem_a = await _insert_memory(memory, conn, bank_id, "Old fact A.", "experience", document_id=doc_id) + doc_mem_b = await _insert_memory(memory, conn, bank_id, "Old fact B.", "experience", document_id=doc_id) + standalone_mem = await _insert_memory(memory, conn, bank_id, "Standalone fact C.") obs_id = await _insert_observation( + memory, conn, bank_id, "Aggregated observation joining doc + standalone facts.", @@ -360,14 +424,11 @@ async def test_upsert_document_removes_observations_from_outgoing_memories( # The standalone memory survives (different document_id) and should # be reset for re-consolidation since one of its observations was # invalidated by the upsert. - consolidated_at = await _get_consolidated_at(conn, standalone_mem) + consolidated_at = await _get_consolidated_at(conn, standalone_mem, bank_id) assert consolidated_at is None, "Surviving co-source memory should be reset for re-consolidation" # The two doc-scoped memories are gone via FK cascade. - doc_mem_count = await conn.fetchval( - "SELECT COUNT(*) FROM memory_units WHERE id = ANY($1::uuid[])", - [doc_mem_a, doc_mem_b], - ) + doc_mem_count = await _count_surviving(conn, bank_id, [doc_mem_a, doc_mem_b]) assert doc_mem_count == 0 await memory.delete_bank(bank_id, request_context=request_context) @@ -389,9 +450,9 @@ async def test_clearing_experience_memories_removes_affected_observations( pool = await memory._get_pool() async with pool.acquire() as conn: - exp1 = await _insert_memory(conn, bank_id, "Alice went hiking last week.", "experience") - world1 = await _insert_memory(conn, bank_id, "Alice is a hiker.", "world") - obs_id = await _insert_observation(conn, bank_id, "Alice is a regular hiker.", [exp1, world1]) + exp1 = await _insert_memory(memory, conn, bank_id, "Alice went hiking last week.", "experience") + world1 = await _insert_memory(memory, conn, bank_id, "Alice is a hiker.", "world") + obs_id = await _insert_observation(memory, conn, bank_id, "Alice is a regular hiker.", [exp1, world1]) # Patch out consolidation so it doesn't re-set consolidated_at before we can check it with patch.object(memory, "submit_async_consolidation", new=AsyncMock()): @@ -402,7 +463,7 @@ async def test_clearing_experience_memories_removes_affected_observations( assert str(obs_id) not in obs_ids, "Observation should have been deleted" # world1 (remaining source) should be reset for re-consolidation - consolidated_at = await _get_consolidated_at(conn, world1) + consolidated_at = await _get_consolidated_at(conn, world1, bank_id) assert consolidated_at is None, "World memory should be reset for re-consolidation" await memory.delete_bank(bank_id, request_context=request_context) @@ -417,8 +478,8 @@ async def test_clearing_unrelated_type_leaves_observations_intact( pool = await memory._get_pool() async with pool.acquire() as conn: - world1 = await _insert_memory(conn, bank_id, "Alice is a hiker.", "world") - obs_id = await _insert_observation(conn, bank_id, "Alice is a regular hiker.", [world1]) + world1 = await _insert_memory(memory, conn, bank_id, "Alice is a hiker.", "world") + obs_id = await _insert_observation(memory, conn, bank_id, "Alice is a regular hiker.", [world1]) # Deleting 'experience' type should not affect observations sourced only from 'world' await memory.delete_bank(bank_id, fact_type="experience", request_context=request_context) @@ -446,9 +507,9 @@ async def test_clears_observations_and_resets_all_source_memories( pool = await memory._get_pool() async with pool.acquire() as conn: - m1 = await _insert_memory(conn, bank_id, "Alice loves hiking.") - m2 = await _insert_memory(conn, bank_id, "Alice hikes every weekend.") - obs_id = await _insert_observation(conn, bank_id, "Alice is an avid hiker.", [m1, m2]) + m1 = await _insert_memory(memory, conn, bank_id, "Alice loves hiking.") + m2 = await _insert_memory(memory, conn, bank_id, "Alice hikes every weekend.") + obs_id = await _insert_observation(memory, conn, bank_id, "Alice is an avid hiker.", [m1, m2]) # Patch out consolidation so it doesn't re-set consolidated_at before we can check it with patch.object(memory, "submit_async_consolidation", new=AsyncMock()): @@ -461,8 +522,8 @@ async def test_clears_observations_and_resets_all_source_memories( assert str(obs_id) not in obs_ids, "Observation should be deleted" # Both m1 (target) and m2 (remaining source) should be reset - assert await _get_consolidated_at(conn, m1) is None, "Target memory should be reset" - assert await _get_consolidated_at(conn, m2) is None, "Remaining source should be reset" + assert await _get_consolidated_at(conn, m1, bank_id) is None, "Target memory should be reset" + assert await _get_consolidated_at(conn, m2, bank_id) is None, "Remaining source should be reset" await memory.delete_bank(bank_id, request_context=request_context) @@ -474,7 +535,7 @@ async def test_no_observations_returns_zero(self, memory: MemoryEngine, request_ pool = await memory._get_pool() async with pool.acquire() as conn: - m1 = await _insert_memory(conn, bank_id, "Alice loves hiking.") + m1 = await _insert_memory(memory, conn, bank_id, "Alice loves hiking.") result = await memory.clear_observations_for_memory(bank_id, str(m1), request_context=request_context) @@ -482,7 +543,7 @@ async def test_no_observations_returns_zero(self, memory: MemoryEngine, request_ async with pool.acquire() as conn: # Memory should still be consolidated (no observations were cleared) - assert await _get_consolidated_at(conn, m1) is not None + assert await _get_consolidated_at(conn, m1, bank_id) is not None await memory.delete_bank(bank_id, request_context=request_context) @@ -496,12 +557,12 @@ async def test_only_clears_observations_referencing_target_memory( pool = await memory._get_pool() async with pool.acquire() as conn: - m1 = await _insert_memory(conn, bank_id, "Alice loves hiking.") - m2 = await _insert_memory(conn, bank_id, "Alice hikes every weekend.") - m3 = await _insert_memory(conn, bank_id, "Alice climbed a mountain.") + m1 = await _insert_memory(memory, conn, bank_id, "Alice loves hiking.") + m2 = await _insert_memory(memory, conn, bank_id, "Alice hikes every weekend.") + m3 = await _insert_memory(memory, conn, bank_id, "Alice climbed a mountain.") - obs1_id = await _insert_observation(conn, bank_id, "Alice is an avid hiker.", [m1, m2]) - obs2_id = await _insert_observation(conn, bank_id, "Alice is a mountaineer.", [m3]) + obs1_id = await _insert_observation(memory, conn, bank_id, "Alice is an avid hiker.", [m1, m2]) + obs2_id = await _insert_observation(memory, conn, bank_id, "Alice is a mountaineer.", [m3]) result = await memory.clear_observations_for_memory(bank_id, str(m1), request_context=request_context) @@ -513,7 +574,7 @@ async def test_only_clears_observations_referencing_target_memory( assert str(obs2_id) in obs_ids, "obs2 (does not reference m1) should remain" # m3 should still be consolidated - assert await _get_consolidated_at(conn, m3) is not None + assert await _get_consolidated_at(conn, m3, bank_id) is not None await memory.delete_bank(bank_id, request_context=request_context) @@ -527,11 +588,11 @@ async def test_multiple_observations_for_same_memory_all_cleared( pool = await memory._get_pool() async with pool.acquire() as conn: - m1 = await _insert_memory(conn, bank_id, "Alice loves hiking.") - m2 = await _insert_memory(conn, bank_id, "Alice hikes every weekend.") + m1 = await _insert_memory(memory, conn, bank_id, "Alice loves hiking.") + m2 = await _insert_memory(memory, conn, bank_id, "Alice hikes every weekend.") - obs1_id = await _insert_observation(conn, bank_id, "Alice hikes often.", [m1]) - obs2_id = await _insert_observation(conn, bank_id, "Alice is outdoorsy.", [m1, m2]) + obs1_id = await _insert_observation(memory, conn, bank_id, "Alice hikes often.", [m1]) + obs2_id = await _insert_observation(memory, conn, bank_id, "Alice is outdoorsy.", [m1, m2]) # Patch out consolidation so it doesn't re-set consolidated_at before we can check it with patch.object(memory, "submit_async_consolidation", new=AsyncMock()): @@ -545,8 +606,8 @@ async def test_multiple_observations_for_same_memory_all_cleared( assert str(obs2_id) not in obs_ids # m1 and m2 should both be reset - assert await _get_consolidated_at(conn, m1) is None - assert await _get_consolidated_at(conn, m2) is None + assert await _get_consolidated_at(conn, m1, bank_id) is None + assert await _get_consolidated_at(conn, m2, bank_id) is None await memory.delete_bank(bank_id, request_context=request_context) @@ -557,9 +618,13 @@ async def test_multiple_observations_for_same_memory_all_cleared( async def _insert_document_with_memories( - conn, bank_id: str, doc_id: str, memories: list[tuple[str, str]] + memory: MemoryEngine, conn, bank_id: str, doc_id: str, memories: list[tuple[str, str]] ) -> list[uuid.UUID]: - """Insert a document and attach memory units to it. Returns list of memory UUIDs.""" + """Insert a document and attach memory units to it. Returns list of memory UUIDs. + + The documents row stays SQL (it is Postgres bookkeeping for every store); the memories + go through the store, attached via document_id at insert time. + """ await conn.execute( """ INSERT INTO documents (id, bank_id, original_text, content_hash, created_at, updated_at) @@ -570,19 +635,7 @@ async def _insert_document_with_memories( ) mem_ids = [] for text, fact_type in memories: - mem_id = uuid.uuid4() - await conn.execute( - """ - INSERT INTO memory_units (id, bank_id, text, fact_type, event_date, document_id, created_at, updated_at, consolidated_at) - VALUES ($1, $2, $3, $4, NOW(), $5, NOW(), NOW(), NOW()) - """, - mem_id, - bank_id, - text, - fact_type, - doc_id, - ) - mem_ids.append(mem_id) + mem_ids.append(await _insert_memory(memory, conn, bank_id, text, fact_type, document_id=doc_id)) return mem_ids @@ -596,7 +649,7 @@ async def test_update_tags_returns_updated_document(self, memory: MemoryEngine, pool = await memory._get_pool() async with pool.acquire() as conn: doc_id = f"doc-{uuid.uuid4().hex[:8]}" - await _insert_document_with_memories(conn, bank_id, doc_id, [("Alice loves hiking.", "experience")]) + await _insert_document_with_memories(memory, conn, bank_id, doc_id, [("Alice loves hiking.", "experience")]) result = await memory.update_document(doc_id, bank_id, tags=["new-tag"], request_context=request_context) @@ -628,7 +681,7 @@ async def test_update_tags_propagates_to_memory_units(self, memory: MemoryEngine async with pool.acquire() as conn: doc_id = f"doc-{uuid.uuid4().hex[:8]}" mem_ids = await _insert_document_with_memories( - conn, bank_id, doc_id, [("Alice loves hiking.", "experience"), ("Alice hikes weekly.", "world")] + memory, conn, bank_id, doc_id, [("Alice loves hiking.", "experience"), ("Alice hikes weekly.", "world")] ) with patch.object(memory, "submit_async_consolidation", new=AsyncMock()): @@ -636,8 +689,8 @@ async def test_update_tags_propagates_to_memory_units(self, memory: MemoryEngine async with pool.acquire() as conn: for mem_id in mem_ids: - tags = await conn.fetchval("SELECT tags FROM memory_units WHERE id = $1", mem_id) - assert list(tags) == ["new-tag"], f"Memory unit {mem_id} should have updated tags" + stored_mem = await _get_memory(conn, bank_id, mem_id) + assert list(stored_mem.tags) == ["new-tag"], f"Memory unit {mem_id} should have updated tags" await memory.delete_bank(bank_id, request_context=request_context) @@ -651,9 +704,9 @@ async def test_update_tags_invalidates_observations(self, memory: MemoryEngine, async with pool.acquire() as conn: doc_id = f"doc-{uuid.uuid4().hex[:8]}" mem_ids = await _insert_document_with_memories( - conn, bank_id, doc_id, [("Alice loves hiking.", "experience")] + memory, conn, bank_id, doc_id, [("Alice loves hiking.", "experience")] ) - obs_id = await _insert_observation(conn, bank_id, "Alice is a hiker.", mem_ids) + obs_id = await _insert_observation(memory, conn, bank_id, "Alice is a hiker.", mem_ids) with patch.object(memory, "submit_async_consolidation", new=AsyncMock()): await memory.update_document(doc_id, bank_id, tags=["new-tag"], request_context=request_context) @@ -676,18 +729,18 @@ async def test_update_tags_resets_consolidated_at_on_affected_units( async with pool.acquire() as conn: doc_id = f"doc-{uuid.uuid4().hex[:8]}" mem_ids = await _insert_document_with_memories( - conn, bank_id, doc_id, [("Alice loves hiking.", "experience")] + memory, conn, bank_id, doc_id, [("Alice loves hiking.", "experience")] ) - obs_id = await _insert_observation(conn, bank_id, "Alice is a hiker.", mem_ids) + await _insert_observation(memory, conn, bank_id, "Alice is a hiker.", mem_ids) # Verify memory starts consolidated - assert await _get_consolidated_at(conn, mem_ids[0]) is not None + assert await _get_consolidated_at(conn, mem_ids[0], bank_id) is not None with patch.object(memory, "submit_async_consolidation", new=AsyncMock()): await memory.update_document(doc_id, bank_id, tags=["new-tag"], request_context=request_context) async with pool.acquire() as conn: - consolidated_at = await _get_consolidated_at(conn, mem_ids[0]) + consolidated_at = await _get_consolidated_at(conn, mem_ids[0], bank_id) assert consolidated_at is None, "Memory unit should be reset for re-consolidation" await memory.delete_bank(bank_id, request_context=request_context) @@ -704,9 +757,9 @@ async def test_update_tags_triggers_consolidation_when_observations_invalidated( async with pool.acquire() as conn: doc_id = f"doc-{uuid.uuid4().hex[:8]}" mem_ids = await _insert_document_with_memories( - conn, bank_id, doc_id, [("Alice loves hiking.", "experience")] + memory, conn, bank_id, doc_id, [("Alice loves hiking.", "experience")] ) - await _insert_observation(conn, bank_id, "Alice is a hiker.", mem_ids) + await _insert_observation(memory, conn, bank_id, "Alice is a hiker.", mem_ids) with patch.object(memory, "submit_async_consolidation", new=AsyncMock()) as mock_consolidate: await memory.update_document(doc_id, bank_id, tags=["new-tag"], request_context=request_context) @@ -725,7 +778,7 @@ async def test_update_tags_no_consolidation_when_no_observations( pool = await memory._get_pool() async with pool.acquire() as conn: doc_id = f"doc-{uuid.uuid4().hex[:8]}" - await _insert_document_with_memories(conn, bank_id, doc_id, [("Alice loves hiking.", "experience")]) + await _insert_document_with_memories(memory, conn, bank_id, doc_id, [("Alice loves hiking.", "experience")]) # No observations inserted with patch.object(memory, "submit_async_consolidation", new=AsyncMock()) as mock_consolidate: @@ -746,16 +799,16 @@ async def test_update_tags_resets_co_source_memories_from_other_documents( async with pool.acquire() as conn: doc_id = f"doc-{uuid.uuid4().hex[:8]}" doc_mem_ids = await _insert_document_with_memories( - conn, bank_id, doc_id, [("Alice loves hiking.", "experience")] + memory, conn, bank_id, doc_id, [("Alice loves hiking.", "experience")] ) # Unrelated memory from another document — co-sourced in the same observation - other_mem = await _insert_memory(conn, bank_id, "Alice also rock-climbs.") + other_mem = await _insert_memory(memory, conn, bank_id, "Alice also rock-climbs.") obs_id = await _insert_observation( - conn, bank_id, "Alice loves outdoor activities.", doc_mem_ids + [other_mem] + memory, conn, bank_id, "Alice loves outdoor activities.", doc_mem_ids + [other_mem] ) # Verify other_mem starts consolidated - assert await _get_consolidated_at(conn, other_mem) is not None + assert await _get_consolidated_at(conn, other_mem, bank_id) is not None with patch.object(memory, "submit_async_consolidation", new=AsyncMock()): await memory.update_document(doc_id, bank_id, tags=["new-tag"], request_context=request_context) @@ -765,7 +818,7 @@ async def test_update_tags_resets_co_source_memories_from_other_documents( assert str(obs_id) not in obs_ids, "Observation should have been invalidated" # other_mem (co-source from another document) must also be reset - consolidated_at = await _get_consolidated_at(conn, other_mem) + consolidated_at = await _get_consolidated_at(conn, other_mem, bank_id) assert consolidated_at is None, "Co-source memory from other document should be reset" await memory.delete_bank(bank_id, request_context=request_context) @@ -781,12 +834,10 @@ async def test_update_tags_does_not_affect_unrelated_observations( pool = await memory._get_pool() async with pool.acquire() as conn: doc_id = f"doc-{uuid.uuid4().hex[:8]}" - mem_ids = await _insert_document_with_memories( - conn, bank_id, doc_id, [("Alice loves hiking.", "experience")] - ) + await _insert_document_with_memories(memory, conn, bank_id, doc_id, [("Alice loves hiking.", "experience")]) # Unrelated memory not in the document - unrelated = await _insert_memory(conn, bank_id, "Bob likes cycling.") - unrelated_obs_id = await _insert_observation(conn, bank_id, "Bob is a cyclist.", [unrelated]) + unrelated = await _insert_memory(memory, conn, bank_id, "Bob likes cycling.") + unrelated_obs_id = await _insert_observation(memory, conn, bank_id, "Bob is a cyclist.", [unrelated]) with patch.object(memory, "submit_async_consolidation", new=AsyncMock()): await memory.update_document(doc_id, bank_id, tags=["new-tag"], request_context=request_context) @@ -822,7 +873,7 @@ async def test_create_observation_filters_deleted_source_memories( pool = await memory._get_pool() async with pool.acquire() as conn: - live = await _insert_memory(conn, bank_id, "Alice loves hiking.") + live = await _insert_memory(memory, conn, bank_id, "Alice loves hiking.") dead = uuid.uuid4() # never existed — stands in for a concurrently deleted source result = await _create_observation_directly( @@ -834,10 +885,7 @@ async def test_create_observation_filters_deleted_source_memories( ) assert result["action"] == "created" - stored = await conn.fetchval( - "SELECT source_memory_ids FROM memory_units WHERE id = $1", - uuid.UUID(result["observation_id"]), - ) + stored = (await _get_memory(conn, bank_id, result["observation_id"])).source_memory_ids stored_set = {str(s) for s in stored} assert str(live) in stored_set assert str(dead) not in stored_set, "Deleted source must not appear in stored observation" @@ -883,8 +931,8 @@ async def test_update_observation_skipped_when_all_new_sources_deleted( pool = await memory._get_pool() async with pool.acquire() as conn: - original_source = await _insert_memory(conn, bank_id, "Alice hikes.") - obs_id = await _insert_observation(conn, bank_id, "Alice is a hiker.", [original_source]) + original_source = await _insert_memory(memory, conn, bank_id, "Alice hikes.") + obs_id = await _insert_observation(memory, conn, bank_id, "Alice is a hiker.", [original_source]) original_text = "Alice is a hiker." observation_model = MemoryFact( @@ -905,9 +953,9 @@ async def test_update_observation_skipped_when_all_new_sources_deleted( observations=[observation_model], ) - row = await conn.fetchrow("SELECT text, source_memory_ids FROM memory_units WHERE id = $1", obs_id) - assert row["text"] == original_text, "Observation text must not change" - stored_sources = {str(s) for s in row["source_memory_ids"]} + row = await _get_memory(conn, bank_id, obs_id) + assert row.text == original_text, "Observation text must not change" + stored_sources = {str(s) for s in row.source_memory_ids} assert stored_sources == {str(original_source)}, "Dead sources must not be appended" await memory.delete_bank(bank_id, request_context=request_context) diff --git a/hindsight-dev/benchmarks/common/benchmark_runner.py b/hindsight-dev/benchmarks/common/benchmark_runner.py index e897848150..00632ab5d3 100644 --- a/hindsight-dev/benchmarks/common/benchmark_runner.py +++ b/hindsight-dev/benchmarks/common/benchmark_runner.py @@ -18,6 +18,7 @@ """ import asyncio +import io import json import logging import os @@ -393,6 +394,9 @@ def __init__( self.answer_generator = answer_generator self.answer_evaluator = answer_evaluator self.template_path: Optional[str] = None + # When set, ingestion replays this exported bank archive instead of + # running fact extraction (see import_bank_archive). + self.bank_archive_path: Optional[str] = None self.memory = memory or MemoryEngine( db_url=os.getenv("HINDSIGHT_API_DATABASE_URL", "pg0"), memory_llm_provider=os.getenv("HINDSIGHT_API_LLM_PROVIDER", "groq"), @@ -484,6 +488,52 @@ async def apply_template(self, bank_id: str, manifest_path: str) -> None: request_context=request_context, ) + async def import_bank_archive(self, agent_id: str) -> int: + """Replay an exported archive into ``agent_id`` instead of extracting facts. + + The archive already carries extracted facts, entity names, causal edges + and chunks, so this runs no LLM extraction at all. The one model cost is + re-embedding, which is unavoidable: embeddings are deliberately not + carried across so the target re-embeds with its own model. + + Accepts either export shape. A whole-bank archive + (``hindsight-admin export-bank``) goes through ``import_bank``; a + document archive (the ``/document-transfer`` endpoint) through + ``import_documents``. Both are called directly rather than via the async + import endpoint so the benchmark needs no file storage or worker — it + wants the writes done before it starts asking questions. + """ + import zipfile + from pathlib import Path + + from hindsight_api.engine.transfer.importer import import_bank, import_documents + + archive_bytes = Path(self.bank_archive_path).read_bytes() + with zipfile.ZipFile(io.BytesIO(archive_bytes)) as archive: + is_bank_archive = "banks.json" in archive.namelist() + + memory = self.memory + config = await memory._config_resolver.resolve_full_config(agent_id, RequestContext()) + backend = await memory._get_backend() + common = dict( + backend=backend, + embeddings_model=memory.embeddings, + entity_resolver=memory.entity_resolver, + config=config, + format_date_fn=memory._format_readable_date, + archive_bytes=archive_bytes, + ) + + if is_bank_archive: + # import_bank restores a complete bank and refuses to merge into an + # existing one, so clear the target first (the caller already does + # this for the normal path via delete_bank). + await memory.delete_bank(agent_id, request_context=RequestContext()) + result = await import_bank(**common, target_bank_id=agent_id) + else: + result = await import_documents(**common, bank_id=agent_id, on_conflict="replace") + return result.documents_imported + async def ingest_conversation( self, item: Dict[str, Any], agent_id: str, wait_for_consolidation: bool = False ) -> int: @@ -922,11 +972,19 @@ async def process_single_item( await self.apply_template(agent_id, self.template_path) console.print(" [green]✓[/green] Template applied") - # Ingest conversation step += 1 - console.print(f" [{step}] Ingesting conversation (batch mode)...") - num_sessions = await self.ingest_conversation(item, agent_id, wait_for_consolidation=False) - console.print(f" [green]✓[/green] Ingested {num_sessions} sessions") + if self.bank_archive_path: + # Replay a previously exported bank instead of extracting facts + # again: the archive already holds them, so this re-embeds and + # re-resolves entities without a single LLM extraction call. + console.print(f" [{step}] Importing bank archive (no LLM extraction)...") + num_sessions = await self.import_bank_archive(agent_id) + console.print(f" [green]✓[/green] Imported {num_sessions} documents") + else: + # Ingest conversation + console.print(f" [{step}] Ingesting conversation (batch mode)...") + num_sessions = await self.ingest_conversation(item, agent_id, wait_for_consolidation=False) + console.print(f" [green]✓[/green] Ingested {num_sessions} sessions") else: num_sessions = -1 @@ -981,6 +1039,7 @@ async def run( merge_with_existing: bool = False, # Whether to merge with existing results wait_consolidation: bool = False, # Wait for consolidation to complete before evaluating QA template_path: Optional[str] = None, # Path to a bank template manifest to apply before ingestion + bank_archive: Optional[str] = None, # Exported bank/document ZIP to replay instead of extracting facts ) -> Dict[str, Any]: """ Run the full benchmark evaluation. @@ -1030,6 +1089,14 @@ async def run( if template_path: self.template_path = template_path console.print(f" Bank template: {template_path}") + if bank_archive: + self.bank_archive_path = bank_archive + console.print(f" Bank archive: {bank_archive} (replaying facts, no extraction)") + from hindsight_api.engine.memories import get_memories + + # Named in the header because a benchmark run is only comparable against + # another run of the same store. + console.print(f" Memories store: {get_memories().name}") console.print(" [green]✓[/green] Memory system initialized") # Start a background worker poller when we need to wait for consolidation. diff --git a/hindsight-dev/benchmarks/locomo/locomo_benchmark.py b/hindsight-dev/benchmarks/locomo/locomo_benchmark.py index 32dbd4570b..7aa3efb715 100644 --- a/hindsight-dev/benchmarks/locomo/locomo_benchmark.py +++ b/hindsight-dev/benchmarks/locomo/locomo_benchmark.py @@ -305,6 +305,7 @@ async def run_benchmark( question_index: int = None, wait_consolidation: bool = False, template_path: str = None, + bank_archive: str = None, ): """ Run the LoComo benchmark. @@ -472,6 +473,7 @@ def filtered_get_qa_pairs(item: Dict) -> List[Dict[str, Any]]: merge_with_existing=merge_with_existing, wait_consolidation=wait_consolidation, template_path=template_path, + bank_archive=bank_archive, ) # Display results (final save already happened incrementally) @@ -615,6 +617,17 @@ def generate_markdown_table(results: dict, use_reflect: bool = False): default=None, help="Path to a bank template manifest JSON to apply before ingestion (sets config, mental models, directives)", ) + parser.add_argument( + "--bank-archive", + type=str, + default=None, + help=( + "Path to an exported bank/document ZIP. Replays its already-extracted facts " + "instead of running fact extraction — facts are re-embedded, no LLM extraction. " + "Use it to re-run over a fixed corpus without paying for ingestion again, and to " + "compare two memories stores on identical input." + ), + ) args = parser.parse_args() @@ -636,5 +649,6 @@ def generate_markdown_table(results: dict, use_reflect: bool = False): question_index=args.question_index, wait_consolidation=args.wait_consolidation, template_path=args.template, + bank_archive=args.bank_archive, ) ) diff --git a/hindsight-dev/benchmarks/perf/system_perf.py b/hindsight-dev/benchmarks/perf/system_perf.py index 62ab7ff229..cb38105ba0 100644 --- a/hindsight-dev/benchmarks/perf/system_perf.py +++ b/hindsight-dev/benchmarks/perf/system_perf.py @@ -1268,8 +1268,9 @@ class _GraphMaintTimers: ``run_graph_maintenance_job`` runs them deep inside its own connections and transactions, so the only seam that doesn't perturb the path under test is - wrapping the functions it calls. We patch the symbol the graph_maintenance - module resolves (``compute_semantic_links_ann``) and the bound ops method + wrapping the functions it calls. The relink probe lives in the Postgres + store's maintenance pass now, so we patch the symbol *that* module resolves + (``compute_semantic_links_ann``) and the bound ops method (``fetch_temporal_neighbors``), tallying wall-clock and call counts. """ @@ -1289,8 +1290,8 @@ class _InstrumentedJob: async def _run_graph_maintenance_instrumented(engine: Any, bank_id: str, request_context: Any) -> _InstrumentedJob: """Run the maintenance job with the two relink probes timed.""" - from hindsight_api.engine import graph_maintenance as gm from hindsight_api.engine.graph_maintenance import run_graph_maintenance_job + from hindsight_api.engine.memories.pg import graph as gm timers = _GraphMaintTimers()