From 9a4a917c481f68b974ebf555306d0864003cc96a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Boschi?= Date: Tue, 28 Jul 2026 18:01:36 +0200 Subject: [PATCH] feat(memories): pluggable memories store + cross-store txns + documents/chunks Squashes the feat/pluggable-memories-provider work into one commit. - Carve the `memory_units` + link slice out from behind raw SQL into a pluggable MemoriesExtension (engine/memories/), so a different engine (memlake) can own memories, links, retrieval, consolidation and curation while documents, chunks, banks and the entity registry stay in Postgres. The default PostgresMemories keeps everything exactly where it was; every call site routes through the store interface rather than branching on the implementation. - Route recall (semantic+BM25+graph), scan/get, stats/counts, consolidation writes, curation edits, bank/document deletion, entity postings and graph reads through the store. - Cross-store write-group transactions (begin/decide/mint/witness + recovery sweep) so a store that keeps memories elsewhere commits atomically with the Postgres side of a retain/consolidation/curation/delete. - Documents & chunks: when the store owns a dedicated document store (owns_document_store), a document's bulky extracted text + chunk texts move out of Postgres into it (Postgres keeps thin rows: id, content_hash, chunk_index, tags); reads overlay the text from the store; the original file goes through a memlake FileStorage backend. All gated so the Postgres path is unchanged. --- ...bservation_history_drop_memory_units_fk.py | 67 + .../engine/consolidation/consolidator.py | 647 ++++-- .../hindsight_api/engine/db/__init__.py | 18 +- .../hindsight_api/engine/entity_resolver.py | 27 +- .../hindsight_api/engine/graph_maintenance.py | 317 +-- .../hindsight_api/engine/maintenance.py | 67 +- .../hindsight_api/engine/memories/__init__.py | 90 + .../hindsight_api/engine/memories/base.py | 1138 ++++++++++ .../engine/memories/pg/__init__.py | 20 + .../engine/memories/pg/counts.py | 151 ++ .../engine/memories/pg/curation.py | 506 +++++ .../hindsight_api/engine/memories/pg/graph.py | 780 +++++++ .../hindsight_api/engine/memories/pg/reads.py | 451 ++++ .../engine/memories/pg/writes.py | 575 +++++ .../hindsight_api/engine/memories/postgres.py | 484 ++++ .../hindsight_api/engine/memory_engine.py | 1998 +++++++---------- .../hindsight_api/engine/reflect/tools.py | 41 +- .../hindsight_api/engine/retain/bank_utils.py | 14 +- .../engine/retain/chunk_storage.py | 24 +- .../engine/retain/fact_storage.py | 284 +-- .../engine/retain/orchestrator.py | 143 +- .../hindsight_api/engine/search/retrieval.py | 104 +- .../hindsight_api/engine/storage/__init__.py | 18 +- .../hindsight_api/engine/storage/memlake.py | 102 + .../test_bank_stats_cache_invalidation.py | 127 +- .../tests/test_graph_maintenance.py | 134 -- .../tests/test_memories_extension.py | 325 +++ .../tests/test_observation_invalidation.py | 386 ++-- .../benchmarks/common/benchmark_runner.py | 75 +- .../benchmarks/locomo/locomo_benchmark.py | 14 + hindsight-dev/benchmarks/perf/system_perf.py | 7 +- 31 files changed, 6953 insertions(+), 2181 deletions(-) create mode 100644 hindsight-api-slim/hindsight_api/alembic/versions/a1c9e7f3b2d8_observation_history_drop_memory_units_fk.py create mode 100644 hindsight-api-slim/hindsight_api/engine/memories/__init__.py create mode 100644 hindsight-api-slim/hindsight_api/engine/memories/base.py create mode 100644 hindsight-api-slim/hindsight_api/engine/memories/pg/__init__.py create mode 100644 hindsight-api-slim/hindsight_api/engine/memories/pg/counts.py create mode 100644 hindsight-api-slim/hindsight_api/engine/memories/pg/curation.py create mode 100644 hindsight-api-slim/hindsight_api/engine/memories/pg/graph.py create mode 100644 hindsight-api-slim/hindsight_api/engine/memories/pg/reads.py create mode 100644 hindsight-api-slim/hindsight_api/engine/memories/pg/writes.py create mode 100644 hindsight-api-slim/hindsight_api/engine/memories/postgres.py create mode 100644 hindsight-api-slim/hindsight_api/engine/storage/memlake.py create mode 100644 hindsight-api-slim/tests/test_memories_extension.py 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()