diff --git a/hindsight-api-slim/hindsight_api/engine/consolidation/consolidator.py b/hindsight-api-slim/hindsight_api/engine/consolidation/consolidator.py index fc0906d693..49b77efea7 100644 --- a/hindsight-api-slim/hindsight_api/engine/consolidation/consolidator.py +++ b/hindsight-api-slim/hindsight_api/engine/consolidation/consolidator.py @@ -303,7 +303,7 @@ async def _dedup_reconcile_create( live_source_ids = await _filter_live_source_memories(conn, bank_id, create_source_ids) if not live_source_ids: return None - if store.writes_memory_rows_in_sql: + if store.writes_memory_rows_in_sql_for(bank_id): # Oracle-safe: _native_search_vector_update emits the to_tsvector clause only for a # native PG tsvector column, "" otherwise (see #3021 — the raw ::regconfig cast # breaks Oracle). RETURNING-gate on the twin's probe-time text so a concurrent @@ -385,7 +385,7 @@ async def _dedup_reconcile_update( store = get_memories() async with acquire_with_retry(pool) as conn: async with conn.transaction(): - if store.writes_memory_rows_in_sql: + if store.writes_memory_rows_in_sql_for(bank_id): # Snapshot the updated row's sources with a PLAIN read (no FOR UPDATE). Lock order # must be sources-before-observation: _filter_live_source_memories below takes # FOR SHARE on the SOURCE rows first, then the fold UPDATE locks the observation @@ -583,7 +583,7 @@ async def _filter_live_source_memories( if not source_memory_ids: return [] store = get_memories() - if store.writes_memory_rows_in_sql: + if store.writes_memory_rows_in_sql_for(bank_id): 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, @@ -612,7 +612,7 @@ async def _any_live_source_memory( if not source_memory_ids: return False store = get_memories() - if store.writes_memory_rows_in_sql: + if store.writes_memory_rows_in_sql_for(bank_id): found = await conn.fetchval( f"SELECT 1 FROM {fq_table('memory_units')} WHERE id = ANY($1::uuid[]) AND bank_id = $2 LIMIT 1", source_memory_ids, @@ -732,7 +732,7 @@ async def _count_observations_for_scope( Observations with no tags are not counted (the limit does not apply to them). """ store = get_memories() - if store.writes_memory_rows_in_sql: + if store.writes_memory_rows_in_sql_for(bank_id): 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[]", @@ -2245,7 +2245,7 @@ async def _execute_update_action( merged_tags = list(existing_tags | source_tags) t0 = time.time() - if store.writes_memory_rows_in_sql: + if store.writes_memory_rows_in_sql_for(bank_id): updated_rows = await conn.execute_rows_affected( f""" UPDATE {fq_table("memory_units")} @@ -2396,7 +2396,7 @@ async def _execute_delete_action( ) -> None: """Delete a superseded or contradicted observation.""" store = get_memories() - if store.writes_memory_rows_in_sql: + if store.writes_memory_rows_in_sql_for(bank_id): 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), @@ -2782,7 +2782,7 @@ async def _create_observation_directly( source_memory_ids = live_source_memory_ids t0 = time.time() - if store.writes_memory_rows_in_sql: + if store.writes_memory_rows_in_sql_for(bank_id): # Query varies based on text search backend. from ..schema import _is_oracle # noqa: PLC0415 diff --git a/hindsight-api-slim/hindsight_api/engine/maintenance.py b/hindsight-api-slim/hindsight_api/engine/maintenance.py index c2d271a5bf..45cf8939f7 100644 --- a/hindsight-api-slim/hindsight_api/engine/maintenance.py +++ b/hindsight-api-slim/hindsight_api/engine/maintenance.py @@ -138,7 +138,12 @@ def _any_job_enabled() -> bool: @staticmethod def _cross_store_recovery_enabled() -> bool: """True when the memories store keeps memories outside SQL and therefore has - cross-store write-group txns a crashed writer could leave undecided.""" + cross-store write-group txns a crashed writer could leave undecided. + + Deliberately reads the PROCESS-LEVEL class attribute, not the per-bank + ``writes_memory_rows_in_sql_for(bank_id)`` — this only decides whether the recovery LOOP + needs to run at all. A store that routes some banks outside SQL keeps the class attribute + False so the loop runs, then ``recover_pending_txns`` is bank-scoped inside it.""" try: from .memories import get_memories diff --git a/hindsight-api-slim/hindsight_api/engine/memories/base.py b/hindsight-api-slim/hindsight_api/engine/memories/base.py index c64fa033e3..c7b1514f7a 100644 --- a/hindsight-api-slim/hindsight_api/engine/memories/base.py +++ b/hindsight-api-slim/hindsight_api/engine/memories/base.py @@ -394,6 +394,20 @@ def name(self) -> str: #: the inline SQL. Cold, never-searched, key-based — see docs/documents-chunks.md. owns_document_store: bool = False + def writes_memory_rows_in_sql_for(self, bank_id: str) -> bool: + """Per-bank form of :attr:`writes_memory_rows_in_sql`. Defaults to the class attribute, so a + single-store extension needs no override. A store that keeps different banks in different + backends (some in SQL, some not) overrides this to answer PER BANK; every *bank-scoped* call + site consults this instead of the class attribute, so mixed banks each take the correct path. + (The few process-level gates — e.g. "is cross-store txn recovery relevant at all" — keep + reading the class attribute.)""" + return self.writes_memory_rows_in_sql + + def owns_document_store_for(self, bank_id: str) -> bool: + """Per-bank form of :attr:`owns_document_store`. Defaults to the class attribute; a store + that keeps some banks in a separate backend overrides it. See :meth:`writes_memory_rows_in_sql_for`.""" + return self.owns_document_store + # ------------------------------------------------------------------ lifecycle async def initialize(self) -> None: diff --git a/hindsight-api-slim/hindsight_api/engine/memory_engine.py b/hindsight-api-slim/hindsight_api/engine/memory_engine.py index afd392215a..22df1ed570 100644 --- a/hindsight-api-slim/hindsight_api/engine/memory_engine.py +++ b/hindsight-api-slim/hindsight_api/engine/memory_engine.py @@ -6016,7 +6016,7 @@ def to_tuple_format(results): from .memories import get_memories _obs_store = get_memories() - if observation_ids_ordered and not _obs_store.writes_memory_rows_in_sql: + if observation_ids_ordered and not _obs_store.writes_memory_rows_in_sql_for(bank_id): # 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. @@ -6098,7 +6098,7 @@ def to_tuple_format(results): # row, so it selects one fewer column and keeps the asyncpg Records as-is — no # per-chunk ``dict`` allocation for an overlay it never runs. _chunk_store = get_memories() - _owns_docs = _chunk_store.owns_document_store + _owns_docs = _chunk_store.owns_document_store_for(bank_id) if _owns_docs: _chunk_cols = "chunk_id, chunk_text, chunk_index, document_id" else: @@ -6295,7 +6295,7 @@ def _source_fact_dict( # Resolve each observation's sources. This is a recall hot path, so the SQL # store reads only the two columns it needs rather than a full memory row; a # store that owns its rows answers from its own objects via one addressed read. - if store.writes_memory_rows_in_sql: + if store.writes_memory_rows_in_sql_for(bank_id): obs_rows = [ {"id": str(r["id"]), "source_memory_ids": r["source_memory_ids"]} for r in await sf_conn.fetch( @@ -6332,7 +6332,7 @@ def _source_fact_dict( # needed, so the SQL store selects those (bank-scoped) instead of the full # 17-column memory row — the difference is measurable on this hot path. if source_ids_ordered: - if store.writes_memory_rows_in_sql: + if store.writes_memory_rows_in_sql_for(bank_id): source_row_by_id = { str(r["id"]): _source_fact_dict( uid=str(r["id"]), @@ -6705,7 +6705,7 @@ async def get_document( from .memories import get_memories _store = get_memories() - if _store.writes_memory_rows_in_sql: + if _store.writes_memory_rows_in_sql_for(bank_id): # 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( @@ -6767,7 +6767,7 @@ async def get_document( # A store that owns the document store 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: + if _store.owns_document_store_for(bank_id): _rec = await _store.get_document_record( bank_id=bank_id, document_id=document_id, include_text=True ) @@ -6844,7 +6844,7 @@ async def delete_document( from .memories import get_memories _store = get_memories() - if _store.writes_memory_rows_in_sql: + if _store.writes_memory_rows_in_sql_for(bank_id): 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, @@ -6890,7 +6890,7 @@ async def delete_document( # 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: + if deleted and not _store.writes_memory_rows_in_sql_for(bank_id): _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 @@ -6899,7 +6899,7 @@ async def delete_document( # 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: + if _store.owns_document_store_for(bank_id): await _store.delete_document_record(bank_id=bank_id, document_id=document_id, txn=_del_txn) # Invalidate observations referencing these (now-deleted) memories @@ -7012,7 +7012,7 @@ async def update_document( from .memories import MemoryPatch, get_memories _store = get_memories() - if tags is not None and not _store.writes_memory_rows_in_sql: + if tags is not None and not _store.writes_memory_rows_in_sql_for(bank_id): # 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 @@ -7186,7 +7186,7 @@ async def delete_memory_unit( from .memories import get_memories _store = get_memories() - if _store.writes_memory_rows_in_sql: + if _store.writes_memory_rows_in_sql_for(bank_id): row = await conn.fetchrow( f"SELECT bank_id, fact_type FROM {fq_table('memory_units')} WHERE id = $1", str(unit_uuid), @@ -7215,7 +7215,7 @@ 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). - if _store.writes_memory_rows_in_sql: + if _store.writes_memory_rows_in_sql_for(bank_id): deleted = await conn.fetchval( f"DELETE FROM {fq_table('memory_units')} WHERE id = $1 RETURNING id", unit_id ) @@ -7529,7 +7529,7 @@ async def delete_bank( from .memories import get_memories as _get_memories_for_scope _scope_store = _get_memories_for_scope() - if _scope_store.writes_memory_rows_in_sql: + if _scope_store.writes_memory_rows_in_sql_for(bank_id): unit_id_rows = await conn.fetch( f"SELECT id FROM {fq_table('memory_units')} WHERE bank_id = $1 AND fact_type = $2", bank_id, @@ -7679,7 +7679,7 @@ async def delete_bank( from .memories import DeletePredicate, get_memories store = get_memories() - if not store.writes_memory_rows_in_sql: + if not store.writes_memory_rows_in_sql_for(bank_id): if fact_type: await store.delete_where(bank_id, DeletePredicate(fact_types=[fact_type])) else: @@ -7729,7 +7729,7 @@ async def clear_observations( backend = await self._get_backend() async with acquire_with_retry(backend) as conn: async with conn.transaction(): - if store.writes_memory_rows_in_sql: + if store.writes_memory_rows_in_sql_for(bank_id): # Count observations before deletion count = await conn.fetchval( f"SELECT COUNT(*) FROM {fq_table('memory_units')} WHERE bank_id = $1 AND fact_type = 'observation'", @@ -7862,7 +7862,7 @@ async def retry_failed_consolidation( store = get_memories() backend = await self._get_backend() async with acquire_with_retry(backend) as conn: - if store.writes_memory_rows_in_sql: + if store.writes_memory_rows_in_sql_for(bank_id): count = await conn.fetchval( f""" SELECT COUNT(*) FROM {fq_table("memory_units")} @@ -7942,7 +7942,7 @@ async def clear_observations_for_memory( from .memories import get_memories _store = get_memories() - if _store.writes_memory_rows_in_sql: + if _store.writes_memory_rows_in_sql_for(bank_id): await conn.execute( f""" UPDATE {fq_table("memory_units")} @@ -9390,7 +9390,7 @@ async def get_chunk( from .memories import get_memories _store = get_memories() - if _store.owns_document_store: + if _store.owns_document_store_for(chunk["bank_id"]): _t = await _store.get_chunk_text( bank_id=chunk["bank_id"], document_id=chunk["document_id"], @@ -9481,7 +9481,7 @@ async def list_document_chunks( from .memories import get_memories _store = get_memories() - if _store.owns_document_store: + if _store.owns_document_store_for(bank_id): _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)) @@ -11631,7 +11631,7 @@ async def _compute_bank_stats(self, bank_id: str) -> dict[str, Any]: from .memories import get_memories _store = get_memories() - if _store.writes_memory_rows_in_sql: + if _store.writes_memory_rows_in_sql_for(bank_id): consolidation_row = await conn.fetchrow( f""" SELECT diff --git a/hindsight-api-slim/hindsight_api/engine/reflect/tools.py b/hindsight-api-slim/hindsight_api/engine/reflect/tools.py index cf4c7fd094..86a66ecb5c 100644 --- a/hindsight-api-slim/hindsight_api/engine/reflect/tools.py +++ b/hindsight-api-slim/hindsight_api/engine/reflect/tools.py @@ -406,7 +406,7 @@ async def tool_expand( from ..memories import get_memories _store = get_memories() - if _store.writes_memory_rows_in_sql: + if _store.writes_memory_rows_in_sql_for(bank_id): memories = await conn.fetch( f""" SELECT id, text, chunk_id, document_id, fact_type, context 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 e94e272984..304366c2d9 100644 --- a/hindsight-api-slim/hindsight_api/engine/retain/bank_utils.py +++ b/hindsight-api-slim/hindsight_api/engine/retain/bank_utils.py @@ -499,7 +499,7 @@ async def list_banks(pool) -> list: last_write = max(write_times) if write_times else None fact_count = row["fact_count"] - if not _store.writes_memory_rows_in_sql: + if not _store.writes_memory_rows_in_sql_for(row["bank_id"]): fact_count = sum( (await _store.count_memories(conn=conn, fq_table=fq_table, bank_id=row["bank_id"])).values() ) 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 22dcd30b4f..4621462322 100644 --- a/hindsight-api-slim/hindsight_api/engine/retain/chunk_storage.py +++ b/hindsight-api-slim/hindsight_api/engine/retain/chunk_storage.py @@ -146,7 +146,7 @@ async def delete_chunks_by_ids(conn, chunk_ids: list[str], bank_id: str | None = from ..memories import META_CHUNK_ID, DeletePredicate, get_memories _store = get_memories() - if bank_id and not _store.writes_memory_rows_in_sql: + if bank_id and not _store.writes_memory_rows_in_sql_for(bank_id): for _cid in chunk_ids: await _store.delete_where(bank_id, DeletePredicate(metadata_equals={META_CHUNK_ID: _cid}), txn=txn) @@ -239,7 +239,7 @@ async def store_chunks_batch( # 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: + if get_memories().owns_document_store_for(bank_id): store_text = False # Prepare chunk data for batch insert 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 0c97804c30..c013f68289 100644 --- a/hindsight-api-slim/hindsight_api/engine/retain/fact_storage.py +++ b/hindsight-api-slim/hindsight_api/engine/retain/fact_storage.py @@ -380,7 +380,7 @@ async def _upsert_document_row( # 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: + if get_memories().owns_document_store_for(bank_id): original_text = None await conn.execute( f""" @@ -422,7 +422,7 @@ async def update_memory_units_metadata_and_tags( from ..memories import MemoryPatch, get_memories store = get_memories() - if not store.writes_memory_rows_in_sql: + if not store.writes_memory_rows_in_sql_for(bank_id): # 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( diff --git a/hindsight-api-slim/hindsight_api/engine/retain/orchestrator.py b/hindsight-api-slim/hindsight_api/engine/retain/orchestrator.py index 4e3ba2b68f..162d20a86a 100644 --- a/hindsight-api-slim/hindsight_api/engine/retain/orchestrator.py +++ b/hindsight-api-slim/hindsight_api/engine/retain/orchestrator.py @@ -1278,7 +1278,7 @@ async def _store_document_bodies( from ..memories import get_memories store = get_memories() - if not store.owns_document_store: + if not store.owns_document_store_for(bank_id): 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 diff --git a/hindsight-api-slim/tests/test_list_banks_non_sql_store.py b/hindsight-api-slim/tests/test_list_banks_non_sql_store.py new file mode 100644 index 0000000000..bd175af9c0 --- /dev/null +++ b/hindsight-api-slim/tests/test_list_banks_non_sql_store.py @@ -0,0 +1,67 @@ +"""Regression test: list_banks must consult the per-bank capability with the +*current row's* bank id, and source fact_count from the store when that bank +keeps its memories outside SQL. + +Bug (introduced with per-bank store capabilities, #3350): list_banks called +``_store.writes_memory_rows_in_sql_for(bank_id)`` with a bare ``bank_id`` name +that is not in scope inside the per-row loop (the row's id is ``row["bank_id"]``). +Because the argument is evaluated before the call, this raised +``NameError: name 'bank_id' is not defined`` for *every* org on the very first +bank — i.e. GET /banks 500'd outright — regardless of the store's capability. + +This test swaps in a store that reports ``writes_memory_rows_in_sql_for -> False`` +(the non-SQL branch the feature added), and asserts list_banks (a) does not raise, +(b) calls the capability + count_memories with the correct per-bank id, and +(c) surfaces the store's live count as fact_count. + +Runs via: uv run pytest tests/test_list_banks_non_sql_store.py -v +""" + +from __future__ import annotations + +import pytest + +import hindsight_api.engine.memories as memories_mod +from hindsight_api.models import RequestContext + + +class _NonSqlStore: + """A store that keeps memory rows outside SQL: list_banks must count via the store.""" + + def __init__(self): + self.capability_calls: list[str] = [] + self.count_calls: list[str] = [] + + def writes_memory_rows_in_sql_for(self, bank_id: str) -> bool: + self.capability_calls.append(bank_id) + return False + + async def count_memories(self, *, conn, fq_table, bank_id: str) -> dict: + self.count_calls.append(bank_id) + return {"world": 7} + + +@pytest.mark.asyncio +async def test_list_banks_counts_via_store_for_non_sql_bank(memory, monkeypatch): + bank_id = "list_banks_non_sql_bank" + request_context = RequestContext(api_key=None, api_key_id=None, tenant_id=None, internal=False) + + store = _NonSqlStore() + monkeypatch.setattr(memories_mod, "get_memories", lambda: store) + + try: + await memory.get_bank_profile(bank_id, request_context=request_context) + + # Must not raise NameError; must reach the store's non-SQL count path. + banks = await memory.list_banks(request_context=request_context) + + entry = next((b for b in banks if b["bank_id"] == bank_id), None) + assert entry is not None, f"bank {bank_id!r} not present in list_banks output" + + # The capability + count were consulted with the row's real bank id. + assert bank_id in store.capability_calls + assert bank_id in store.count_calls + # fact_count came from the store (sum of the per-type counts), not the empty SQL join. + assert entry["fact_count"] == 7 + finally: + await memory.delete_bank(bank_id, request_context=request_context) diff --git a/hindsight-api-slim/tests/test_memories_extension.py b/hindsight-api-slim/tests/test_memories_extension.py index 91d946963f..19c4446e21 100644 --- a/hindsight-api-slim/tests/test_memories_extension.py +++ b/hindsight-api-slim/tests/test_memories_extension.py @@ -601,6 +601,59 @@ async def test_maintenance_passes_are_optional(restore_default_store): await store.record_unit_entities(conn=None, ops=None, fq_table=None, unit_ids=["u"], entity_ids=["e"]) +# --------------------------------------------------------------------------- +# Per-bank store capabilities. A store may route different banks to different +# backends, so every BANK-SCOPED call site asks per bank — +# writes_memory_rows_in_sql_for(bank_id) / owns_document_store_for(bank_id) — +# rather than reading the process-global class attribute. The class attribute +# stays the single-store default the _for methods fall back to. +# --------------------------------------------------------------------------- + + +def test_per_bank_capability_defaults_to_the_class_attribute(): + """A single-store extension needs no override: the _for methods return the class attr, so + every existing store keeps its exact behaviour for every bank.""" + pg = PostgresMemories({}) + assert (pg.writes_memory_rows_in_sql, pg.owns_document_store) == (True, False) + assert pg.writes_memory_rows_in_sql_for("any-bank") is True + assert pg.owns_document_store_for("any-bank") is False + + mem = InMemoryMemories({}) # owns its rows AND its document store + assert mem.writes_memory_rows_in_sql_for("any-bank") is False + assert mem.owns_document_store_for("any-bank") is True + + +def test_a_store_answers_capabilities_per_bank(): + """The point of the _for methods: a store that keeps some banks in SQL and others in a + separate store answers PER BANK, so mixed banks in one process each take the right path.""" + + class PerBankStore(InMemoryMemories): + name = "per-bank" + # The loop-level class attr stays False so cross-store txn recovery still runs; the + # per-bank answer is what every bank-scoped site consults. + writes_memory_rows_in_sql = False + + def __init__(self, config=None): + super().__init__(config) + self.sql_banks = {"legacy-bank"} + + def writes_memory_rows_in_sql_for(self, bank_id): + return bank_id in self.sql_banks + + def owns_document_store_for(self, bank_id): + return bank_id not in self.sql_banks + + store = PerBankStore({}) + # A SQL-backed bank looks like Postgres (host does inline SQL, keeps documents in SQL)... + assert store.writes_memory_rows_in_sql_for("legacy-bank") is True + assert store.owns_document_store_for("legacy-bank") is False + # ...a store-backed bank owns its rows and its document store. + assert store.writes_memory_rows_in_sql_for("new-bank") is False + assert store.owns_document_store_for("new-bank") is True + # The process-level gate (cross-store recovery loop) still fires off the class attr. + assert store.writes_memory_rows_in_sql is False + + # --------------------------------------------------------------------------- # Interface conformance: the stub must stay a COMPLETE, signature-compatible # implementation of every MemoriesExtension method. This is the guard that keeps