From dcc2917c10e825eb6da8db14d0af655a6047226c Mon Sep 17 00:00:00 2001 From: zhangchi47 Date: Mon, 10 Aug 2026 17:32:18 +0800 Subject: [PATCH 1/2] fix(transfer): rebuild mental model search state Re-embed restored mental models with the target provider. Rebuild PostgreSQL lexical projections when required. --- .../hindsight_api/engine/memory_engine.py | 5 +- .../hindsight_api/engine/transfer/importer.py | 96 ++++++++++++++++++- .../tests/test_document_transfer.py | 71 +++++++++++++- 3 files changed, 166 insertions(+), 6 deletions(-) diff --git a/hindsight-api-slim/hindsight_api/engine/memory_engine.py b/hindsight-api-slim/hindsight_api/engine/memory_engine.py index 567d58dbe..e59c40d57 100644 --- a/hindsight-api-slim/hindsight_api/engine/memory_engine.py +++ b/hindsight-api-slim/hindsight_api/engine/memory_engine.py @@ -4419,8 +4419,9 @@ async def import_bank_async( """Restore a whole bank from an :func:`transfer.export_bank` archive. Re-embeds facts with this instance's embedding model and rebuilds links and - indexes; restores bank config, mental models, directives and webhooks as - exported (no consolidation/webhooks — a migration restores exact state). The + indexes; restores bank config, mental model logical fields, directives and + webhooks as exported while rebuilding target-derived mental model search + state (no consolidation/webhooks — a migration restores exact state). The target bank must not already exist (import restores a whole bank, not a merge). """ from .transfer import import_bank diff --git a/hindsight-api-slim/hindsight_api/engine/transfer/importer.py b/hindsight-api-slim/hindsight_api/engine/transfer/importer.py index a20784bb4..1df5d0b18 100644 --- a/hindsight-api-slim/hindsight_api/engine/transfer/importer.py +++ b/hindsight-api-slim/hindsight_api/engine/transfer/importer.py @@ -21,7 +21,15 @@ from ..causal_links import CANONICAL_CAUSAL_LINK_TYPE, LEGACY_CAUSAL_LINK_TYPES from ..db_utils import acquire_with_retry -from ..retain import bank_utils, chunk_storage, embedding_processing, fact_storage, link_utils, orchestrator +from ..retain import ( + bank_utils, + chunk_storage, + embedding_processing, + embedding_utils, + fact_storage, + link_utils, + orchestrator, +) from ..retain.types import ( CausalRelation, ChunkMetadata, @@ -349,6 +357,67 @@ async def _restore_rows( return inserted +async def _rebuild_mental_model_search_state( + conn: Any, + rows: list[dict], + *, + bank_id: str, + embedding_values: list[str | None], + config: Any, +) -> None: + """Rebuild target-derived search fields for carried mental models. + + Whole-bank archives deliberately omit ``embedding`` and ``search_vector``. + Mental models are restored as rows rather than through ``create_mental_model``, + so the target-side embedding and text-search lifecycle must be applied here. + Native generated search columns are populated by the INSERT; regular vector + columns (notably VectorChord, and native columns after the configurable-language + migration) need an explicit expression using the target configuration. + """ + if not rows: + return + + search_vector_expr: str | None = None + if config.database_backend == "postgresql": + from ..memory_engine import get_current_schema + + schema = get_current_schema() + column = await conn.fetchrow( + """ + SELECT is_generated + FROM information_schema.columns + WHERE table_schema = $1 AND table_name = 'mental_models' AND column_name = 'search_vector' + """, + schema, + ) + if not (column and column["is_generated"] == "ALWAYS"): + from ..db.ops_postgresql import pg_search_vector_expr + + search_vector_expr = pg_search_vector_expr( + config, + text_col="name", + context_col="''", + signals_col="content", + ) + + for row, embedding in zip(rows, embedding_values): + if search_vector_expr is None: + await conn.execute( + f"UPDATE {fq_table('mental_models')} SET embedding = $1 WHERE bank_id = $2 AND id = $3", + embedding, + bank_id, + row["id"], + ) + else: + await conn.execute( + f"UPDATE {fq_table('mental_models')} SET embedding = $1, search_vector = {search_vector_expr} " + "WHERE bank_id = $2 AND id = $3", + embedding, + bank_id, + row["id"], + ) + + async def import_bank( *, backend: Any, @@ -372,7 +441,8 @@ async def import_bank( target id is present, this raises — delete it first or pass ``target_bank_id`` for a fresh id. A migration restores *exact* state, so unlike the document import it fires no retain webhooks and triggers no consolidation/graph - maintenance: observations and mental models are restored as exported. + maintenance: observations and mental model logical fields are restored as + exported, while target-derived mental model search state is rebuilt. Takes ``resolve_config`` rather than a resolved config because the only correct moment to resolve one is *inside* this function, after the archive's bank row @@ -451,13 +521,33 @@ async def import_bank( facts_imported=doc_result.facts_imported, observations_imported=doc_result.observations_imported, ) + mental_model_rows = parsed.bank_rows.get("mental_models", []) + # A slow embedder must not pin a pooled connection while the target-side + # derived search fields are rebuilt below. + embeddings = ( + await embedding_utils.generate_embeddings_batch( + embeddings_model, + [f"{row['name']} {row['content']}" for row in mental_model_rows], + ) + if mental_model_rows + else [] + ) + mental_model_embedding_values = [str(value) if value else None for value in embeddings] + async with acquire_with_retry(backend) as conn: result.mental_models_imported = await _restore_rows( conn, "mental_models", - parsed.bank_rows.get("mental_models", []), + mental_model_rows, bank_rows_json_encoding=bank_rows_json_encoding, ) + await _rebuild_mental_model_search_state( + conn, + mental_model_rows, + bank_id=bank_id, + embedding_values=mental_model_embedding_values, + config=config, + ) # Restored after mental_models so the (mental_model_id, bank_id) FK resolves. result.mental_model_history_imported = await _restore_rows( conn, diff --git a/hindsight-api-slim/tests/test_document_transfer.py b/hindsight-api-slim/tests/test_document_transfer.py index b4bddfe21..8e646bfd2 100644 --- a/hindsight-api-slim/tests/test_document_transfer.py +++ b/hindsight-api-slim/tests/test_document_transfer.py @@ -10,6 +10,8 @@ import uuid import zipfile from datetime import datetime, timezone +from types import SimpleNamespace +from unittest.mock import AsyncMock import httpx import pytest @@ -20,7 +22,7 @@ from hindsight_api.engine.db_utils import acquire_with_retry from hindsight_api.engine.schema import fq_table from hindsight_api.engine.transfer import import_documents -from hindsight_api.engine.transfer.importer import parse_archive +from hindsight_api.engine.transfer.importer import _rebuild_mental_model_search_state, parse_archive from hindsight_api.engine.transfer.schema import ( SCHEMA_VERSION, TransferCausalRelation, @@ -65,6 +67,59 @@ async def on_retain_complete(self, result: RetainResult) -> None: self.results.append(result) +@pytest.mark.asyncio +async def test_rebuild_mental_model_search_state_rebuilds_regular_search_vector(): + """Imported mental models must use the target text backend.""" + conn = AsyncMock() + conn.fetchrow.return_value = {"is_generated": "NEVER"} + config = SimpleNamespace( + database_backend="postgresql", + text_search_extension="vchord", + text_search_extension_native_language="english", + ) + rows = [ + {"id": "mm-1", "name": "One", "content": "First"}, + {"id": "mm-2", "name": "Two", "content": "Second"}, + ] + + await _rebuild_mental_model_search_state( + conn, + rows, + bank_id="bank-1", + embedding_values=["[0.1, 0.2]", "[0.3, 0.4]"], + config=config, + ) + + assert conn.execute.await_count == 2 + first_query = conn.execute.await_args_list[0].args[0] + assert "embedding = $1" in first_query + assert "search_vector = tokenize(" in first_query + + +@pytest.mark.asyncio +async def test_rebuild_mental_model_search_state_skips_generated_search_vector(): + """PostgreSQL generates native search vectors when the column is generated.""" + conn = AsyncMock() + conn.fetchrow.return_value = {"is_generated": "ALWAYS"} + config = SimpleNamespace( + database_backend="postgresql", + text_search_extension="native", + text_search_extension_native_language="english", + ) + + await _rebuild_mental_model_search_state( + conn, + [{"id": "mm-1", "name": "One", "content": "First"}], + bank_id="bank-1", + embedding_values=["[0.1, 0.2]"], + config=config, + ) + + query = conn.execute.await_args.args[0] + assert "embedding = $1" in query + assert "search_vector =" not in query + + @pytest_asyncio.fixture async def api_client(memory): """Async HTTP client over the FastAPI app backed by the mock-LLM engine.""" @@ -389,6 +444,14 @@ async def _bank_content_snapshot(memory, bank_id): f"WHERE bank_id = $1 AND fact_type != 'observation' AND embedding IS NULL", bank_id, ) + null_mm_emb = await conn.fetchval( + f"SELECT count(*) FROM {fq_table('mental_models')} WHERE bank_id = $1 AND embedding IS NULL", + bank_id, + ) + null_mm_search = await conn.fetchval( + f"SELECT count(*) FROM {fq_table('mental_models')} WHERE bank_id = $1 AND search_vector IS NULL", + bank_id, + ) return { "bank": (bank["name"], _as_json(bank["disposition"]), bank["mission"], _as_json(bank["config"])), "documents": sorted( @@ -404,6 +467,8 @@ async def _bank_content_snapshot(memory, bank_id): (m["subtype"], m["name"], m["description"], tuple(sorted(m["tags"] or []))) for m in mms ), "null_embeddings": null_emb, + "null_mm_embeddings": null_mm_emb, + "null_mm_search_vectors": null_mm_search, } @@ -601,6 +666,8 @@ async def test_bank_export_import_exact_roundtrip(memory, request_context): assert before["facts"] and before["entities"] and before["links"] assert before["webhooks"] and before["directives"] and before["mental_models"] assert before["bank"][0] == "My Bank" + assert before["null_mm_embeddings"] == 0 + assert before["null_mm_search_vectors"] == 0 from hindsight_api.engine.transfer import export_bank @@ -626,6 +693,8 @@ async def test_bank_export_import_exact_roundtrip(memory, request_context): assert after_semantic > 0, "semantic links should be regenerated on import" # Facts were re-embedded on import (no NULL vectors). assert after["null_embeddings"] == 0 + assert after["null_mm_embeddings"] == 0 + assert after["null_mm_search_vectors"] == 0 finally: await memory.delete_bank(bank, request_context=request_context) From b28003bc8b9b042718cefd81eaed49d3e7fa014d Mon Sep 17 00:00:00 2001 From: zhangchi47 Date: Mon, 10 Aug 2026 17:57:20 +0800 Subject: [PATCH 2/2] fix(transfer): preserve knowledge pages in bank archives Export knowledge-page rows and restore parent nodes before children. --- .../hindsight_api/engine/transfer/export.py | 19 +++---- .../hindsight_api/engine/transfer/importer.py | 29 ++++++++++- .../tests/test_document_transfer.py | 51 +++++++++++++++++-- 3 files changed, 82 insertions(+), 17 deletions(-) diff --git a/hindsight-api-slim/hindsight_api/engine/transfer/export.py b/hindsight-api-slim/hindsight_api/engine/transfer/export.py index 5e551fad2..0737bd006 100644 --- a/hindsight-api-slim/hindsight_api/engine/transfer/export.py +++ b/hindsight-api-slim/hindsight_api/engine/transfer/export.py @@ -70,7 +70,7 @@ ) # Carried verbatim as JSON rows (bank config + synthesized state). Embedding-bearing # rows have their vector stripped (see _DERIVED_COLUMNS) and are re-embedded on import. -_BANK_ROW_TABLES = ("banks", "mental_models", "directives", "webhooks") +_BANK_ROW_TABLES = ("banks", "mental_models", "knowledge_pages", "directives", "webhooks") # Bank-scoped child-history carried verbatim. Unlike observations, mental models # keep their (id, bank_id) across export/import, so their refresh history can be # re-attached. The surrogate ``id`` is dropped on dump so the target reassigns it @@ -88,12 +88,6 @@ # to fresh ids, so carrying them would only produce dangling associations. # Revert anything worth keeping on the source before migrating. "invalidated_memory_units", - # Knowledge-base folder/page tree (client-managed metadata over the carried - # mental models). Not carried yet: its self-referential parent_id FK needs a - # parents-first (topological) restore order, which the generic per-row - # _restore_rows doesn't provide — a follow-up. The mental models themselves - # ARE carried, so the target can recreate the tree. - "knowledge_pages", } ) # Derived columns dropped from carried rows so the target regenerates them with @@ -294,11 +288,12 @@ async def export_bank( Produces a superset of the documents archive: the logical document/fact/observation export (replayed and re-embedded on import) plus - the bank's config, mental models, directives and webhooks as JSON rows. With - ``include_history`` the operational tails (audit_log, llm_requests) are also - carried. Intended for migrating a bank to a new instance configured with a - different embedding model / vector / text-search backend — every vector is - regenerated on the target, so nothing here is encoder-specific. + the bank's config, mental models, knowledge pages, directives and webhooks as + JSON rows. With ``include_history`` the operational tails (audit_log, + llm_requests) are also carried. Intended for migrating a bank to a new + instance configured with a different embedding model / vector / text-search + backend — every vector is regenerated on the target, so nothing here is + encoder-specific. ``conn`` is a live connection scoped to the bank's schema (the admin CLI sets ``_current_schema`` and passes its raw connection; the engine acquires one diff --git a/hindsight-api-slim/hindsight_api/engine/transfer/importer.py b/hindsight-api-slim/hindsight_api/engine/transfer/importer.py index 1df5d0b18..2a7fd54a8 100644 --- a/hindsight-api-slim/hindsight_api/engine/transfer/importer.py +++ b/hindsight-api-slim/hindsight_api/engine/transfer/importer.py @@ -238,7 +238,7 @@ async def import_documents( # Bank-level config/state tables restored verbatim from a whole-bank archive. # Order matters for foreign keys: banks (parent) is restored before any child. -_BANK_CHILD_TABLES = ("mental_models", "directives", "webhooks") +_BANK_CHILD_TABLES = ("mental_models", "knowledge_pages", "directives", "webhooks") # Child-history carried verbatim; restored after its parent (mental_models) so the # foreign key resolves. Surrogate ids were dropped on export (the target reassigns # them), so these restore via fresh IDENTITY values. @@ -264,12 +264,31 @@ class ParsedBankArchive: """The bank-level sections of a whole-bank archive (documents read separately).""" manifest: TransferManifest - # table name -> list of verbatim row dicts (banks, mental_models, directives, webhooks) + # table name -> list of verbatim row dicts (banks, mental_models, knowledge_pages, directives, webhooks) bank_rows: dict[str, list[dict]] = field(default_factory=dict) # table name -> rows (audit_log, llm_requests), present only with --include-history history_rows: dict[str, list[dict]] = field(default_factory=dict) +def _order_knowledge_pages(rows: list[dict]) -> list[dict]: + """Order knowledge nodes so every self-referential parent is inserted first.""" + pending = list(rows) + ordered: list[dict] = [] + restored_ids: set[str] = set() + while pending: + next_pending: list[dict] = [] + for row in pending: + if row["parent_id"] is None or row["parent_id"] in restored_ids: + ordered.append(row) + restored_ids.add(row["id"]) + else: + next_pending.append(row) + if len(next_pending) == len(pending): + raise ValueError("Knowledge-page archive contains an unresolved parent or cycle") + pending = next_pending + return ordered + + def parse_bank_archive(archive_bytes: bytes) -> ParsedBankArchive: """Parse the bank-level sections of a whole-bank archive (``archive_type='bank'``).""" with zipfile.ZipFile(io.BytesIO(archive_bytes), "r") as zf: @@ -548,6 +567,12 @@ async def import_bank( embedding_values=mental_model_embedding_values, config=config, ) + await _restore_rows( + conn, + "knowledge_pages", + _order_knowledge_pages(parsed.bank_rows.get("knowledge_pages", [])), + bank_rows_json_encoding=bank_rows_json_encoding, + ) # Restored after mental_models so the (mental_model_id, bank_id) FK resolves. result.mental_model_history_imported = await _restore_rows( conn, diff --git a/hindsight-api-slim/tests/test_document_transfer.py b/hindsight-api-slim/tests/test_document_transfer.py index 8e646bfd2..966c84709 100644 --- a/hindsight-api-slim/tests/test_document_transfer.py +++ b/hindsight-api-slim/tests/test_document_transfer.py @@ -22,7 +22,11 @@ from hindsight_api.engine.db_utils import acquire_with_retry from hindsight_api.engine.schema import fq_table from hindsight_api.engine.transfer import import_documents -from hindsight_api.engine.transfer.importer import _rebuild_mental_model_search_state, parse_archive +from hindsight_api.engine.transfer.importer import ( + _order_knowledge_pages, + _rebuild_mental_model_search_state, + parse_archive, +) from hindsight_api.engine.transfer.schema import ( SCHEMA_VERSION, TransferCausalRelation, @@ -120,6 +124,17 @@ async def test_rebuild_mental_model_search_state_skips_generated_search_vector() assert "search_vector =" not in query +def test_order_knowledge_pages_restores_parents_first(): + """Knowledge-page archive order must not violate the self-referential FK.""" + rows = [ + {"id": "page", "parent_id": "nested"}, + {"id": "nested", "parent_id": "root"}, + {"id": "root", "parent_id": None}, + ] + + assert [row["id"] for row in _order_knowledge_pages(rows)] == ["root", "nested", "page"] + + @pytest_asyncio.fixture async def api_client(memory): """Async HTTP client over the FastAPI app backed by the mock-LLM engine.""" @@ -377,7 +392,7 @@ async def test_export_bank_contents(memory, request_context): assert manifest.bank_rows_json_encoding == "serialized" assert manifest.document_count == 1 assert manifest.webhook_count == 1 - assert "mental_models.json" in names and "directives.json" in names + assert "mental_models.json" in names and "knowledge_pages.json" in names and "directives.json" in names assert "mental_model_history.json" in names assert any(d.endswith(".json") and d.startswith("documents/") for d in names) # No history files unless requested. @@ -439,6 +454,11 @@ async def _bank_content_snapshot(memory, bank_id): mms = await conn.fetch( f"SELECT subtype, name, description, tags FROM {fq_table('mental_models')} WHERE bank_id = $1", bank_id ) + knowledge_pages = await conn.fetch( + f"SELECT id, parent_id, kind, name, mental_model_id, sort_order, managed " + f"FROM {fq_table('knowledge_pages')} WHERE bank_id = $1", + bank_id, + ) null_emb = await conn.fetchval( f"SELECT count(*) FROM {fq_table('memory_units')} " f"WHERE bank_id = $1 AND fact_type != 'observation' AND embedding IS NULL", @@ -466,6 +486,18 @@ async def _bank_content_snapshot(memory, bank_id): "mental_models": sorted( (m["subtype"], m["name"], m["description"], tuple(sorted(m["tags"] or []))) for m in mms ), + "knowledge_pages": sorted( + ( + p["id"], + p["parent_id"], + p["kind"], + p["name"], + p["mental_model_id"], + p["sort_order"], + p["managed"], + ) + for p in knowledge_pages + ), "null_embeddings": null_emb, "null_mm_embeddings": null_mm_emb, "null_mm_search_vectors": null_mm_search, @@ -660,11 +692,24 @@ async def test_bank_export_import_exact_roundtrip(memory, request_context): tags=["people"], request_context=request_context, ) + folder = await memory.create_knowledge_folder(bank, "Docs", request_context=request_context) + nested = await memory.create_knowledge_folder( + bank, "Nested", parent_id=folder["id"], request_context=request_context + ) + await memory.create_knowledge_page( + bank, + "Guide", + "How should the guide be written?", + "The guide content.", + parent_id=nested["id"], + request_context=request_context, + ) before = await _bank_content_snapshot(memory, bank) # Sanity: the source genuinely has rich content in every section we carry. assert before["facts"] and before["entities"] and before["links"] assert before["webhooks"] and before["directives"] and before["mental_models"] + assert len(before["knowledge_pages"]) == 3 assert before["bank"][0] == "My Bank" assert before["null_mm_embeddings"] == 0 assert before["null_mm_search_vectors"] == 0 @@ -679,7 +724,7 @@ async def test_bank_export_import_exact_roundtrip(memory, request_context): assert result.bank_id == bank assert result.webhooks_imported == 1 assert result.directives_imported == 1 - assert result.mental_models_imported == 1 + assert result.mental_models_imported == 2 after = await _bank_content_snapshot(memory, bank) # Semantic links are an ANN-approximate retrieval index regenerated from the