feat(repair): preserve embeddings on rebuild — cherry-pick of upstream MemPalace#1367 (embeddings portion only)

jphein · claude · jphein · commit aa84fa40cde7 · 2026-05-13T15:37:20.000-07:00
Cherry-picks the `_extract_drawers` embeddings-preserve fix from upstream PR MemPalace#1367 (@zhapostolski), per the cherry-pick evaluation in #31 and @messelink's review guidance: take the embeddings fix, SKIP the divergence-floor cap (which would re-create _refresh_vector_disabled_flag false-positives on our sync_threshold= 50000 palace per MemPalace#1287's design). What the embeddings-preserve fix does: - _extract_drawers now includes "embeddings" in the col.get() call - Returns a 4-tuple: ids, docs, metas, embeddings (or None when any embedding slot is missing — HNSW unreadable case) - _rebuild_collection_via_temp accepts an all_embeddings= kwarg and passes embeddings through to upsert() when available - rebuild_index threads embeddings from _extract_drawers into _rebuild_collection_via_temp When embeddings ARE available (the common case), rebuild skips ONNX recomputation entirely. On a 150K+ palace this cuts rebuild from hours to minutes. Fallback to recompute when any embedding is None (the bug condition we're trying to repair). Updates two call sites in mempalace/cli.py (the cmd_repair path) and mempalace/repair.py (the rebuild_index path) for the new 4-tuple signature. Adds 3 tests: - test_extract_drawers_preserves_embeddings_when_all_present - test_extract_drawers_falls_back_to_none_when_any_embedding_missing - test_extract_drawers_embeddings_absent_in_batch_signals_recompute Plus updates the existing 5 _extract_drawers tests to unpack the new 4-tuple. All 75 tests pass. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
diff --git a/mempalace/cli.py b/mempalace/cli.py
@@ -1378,8 +1378,13 @@ def cmd_repair(args):
     # Extract all drawers in batches
     print("\n  Extracting drawers...")
     batch_size = 5000
-    all_ids, all_docs, all_metas = _extract_drawers(col, total, batch_size)
-    print(f"  Extracted {len(all_ids)} drawers")
+    all_ids, all_docs, all_metas, all_embeddings = _extract_drawers(col, total, batch_size)
+    emb_status = (
+        "with stored embeddings"
+        if all_embeddings is not None
+        else "without embeddings (will recompute)"
+    )
+    print(f"  Extracted {len(all_ids)} drawers ({emb_status})")
 
     # ── #1208 guard ──────────────────────────────────────────────────
     # Cross-check against the SQLite ground truth before doing anything
@@ -1422,6 +1427,7 @@ def cmd_repair(args):
             batch_size,
             collection_name=collection_name,
             progress=print,
+            all_embeddings=all_embeddings,
         )
     except RebuildCollectionError as e:
         print(f"  Repair failed: {e}")
diff --git a/mempalace/repair.py b/mempalace/repair.py
@@ -130,12 +130,32 @@ def _paginate_ids(col, where=None):
 
 
 def _extract_drawers(col, total: int, batch_size: int):
+    """Extract drawers from the source collection for a rebuild.
+
+    Returns ``(ids, documents, metadatas, embeddings)`` — embeddings is a
+    list aligned with ids when every slot has a vector, or ``None`` when
+    any slot is missing (HNSW segment unreadable; caller must let
+    ChromaDB recompute during upsert).
+
+    Embeddings-preserve fix cherry-picked from MemPalace/mempalace#1367
+    (Zharko Apostolski) — the embeddings portion only, NOT the
+    divergence-floor cap (per @messelink's review on jphein/mempalace#31:
+    the divergence cap would false-positive on our sync_threshold=50000
+    palace). Without this, _extract_drawers omitted embeddings entirely
+    and rebuild had to recompute every vector through ONNX, multi-hour
+    cost on a 150K+ palace.
+    """
     all_ids = []
     all_docs = []
     all_metas = []
+    all_embeddings = []
     offset = 0
     while offset < total:
-        batch = col.get(limit=batch_size, offset=offset, include=["documents", "metadatas"])
+        batch = col.get(
+            limit=batch_size,
+            offset=offset,
+            include=["documents", "metadatas", "embeddings"],
+        )
         if not batch["ids"]:
             break
         all_ids.extend(batch["ids"])
@@ -151,8 +171,22 @@ def _extract_drawers(col, total: int, batch_size: int):
             for m in batch["metadatas"]
         ]
         all_metas.extend(sanitized_metas)
+        # embeddings may be None when the HNSW segment is unreadable (the
+        # exact scenario we are trying to repair). Fall back to None so the
+        # caller can pass embeddings=None and let ChromaDB recompute them.
+        batch_embeddings = batch.get("embeddings")
+        if batch_embeddings is not None and any(e is None for e in batch_embeddings):
+            batch_embeddings = None
+        all_embeddings.extend(
+            batch_embeddings
+            if batch_embeddings is not None
+            else [None] * len(batch["ids"])
+        )
         offset += len(batch["ids"])
-    return all_ids, all_docs, all_metas
+    # If any embedding slot is None the whole list is unusable for direct
+    # upsert — signal that to the caller with None so ChromaDB recomputes.
+    has_all = all_embeddings and all(e is not None for e in all_embeddings)
+    return all_ids, all_docs, all_metas, all_embeddings if has_all else None
 
 
 def _verify_collection_count(col, expected: int, label: str) -> None:
@@ -194,12 +228,18 @@ def _rebuild_collection_via_temp(
     batch_size: int,
     collection_name: Optional[str] = None,
     progress=print,
+    all_embeddings=None,
 ) -> int:
     expected = len(all_ids)
     collection_name = collection_name or _drawers_collection_name()
     temp_name = f"{collection_name}__repair_tmp"
     live_replaced = False
 
+    if all_embeddings is not None:
+        progress("  Re-using stored embeddings (no recomputation needed).")
+    else:
+        progress("  Stored embeddings unavailable — will recompute during rebuild.")
+
     try:
         _delete_collection_if_exists(backend, palace_path, temp_name)
 
@@ -210,7 +250,10 @@ def _rebuild_collection_via_temp(
             batch_ids = all_ids[i : i + batch_size]
             batch_docs = all_docs[i : i + batch_size]
             batch_metas = all_metas[i : i + batch_size]
-            temp_col.upsert(documents=batch_docs, ids=batch_ids, metadatas=batch_metas)
+            kwargs = {"documents": batch_docs, "ids": batch_ids, "metadatas": batch_metas}
+            if all_embeddings is not None:
+                kwargs["embeddings"] = all_embeddings[i : i + batch_size]
+            temp_col.upsert(**kwargs)
             staged += len(batch_ids)
             progress(f"  Staged {staged}/{expected} drawers...")
         _verify_collection_count(temp_col, expected, "temporary rebuild")
@@ -225,7 +268,10 @@ def _rebuild_collection_via_temp(
             batch_ids = all_ids[i : i + batch_size]
             batch_docs = all_docs[i : i + batch_size]
             batch_metas = all_metas[i : i + batch_size]
-            new_col.upsert(documents=batch_docs, ids=batch_ids, metadatas=batch_metas)
+            kwargs = {"documents": batch_docs, "ids": batch_ids, "metadatas": batch_metas}
+            if all_embeddings is not None:
+                kwargs["embeddings"] = all_embeddings[i : i + batch_size]
+            new_col.upsert(**kwargs)
             rebuilt += len(batch_ids)
             progress(f"  Re-filed {rebuilt}/{expected} drawers...")
         _verify_collection_count(new_col, expected, "rebuilt live collection")
@@ -755,8 +801,13 @@ def rebuild_index(
     # Extract all drawers in batches
     progress("\n  Extracting drawers...")
     batch_size = 5000
-    all_ids, all_docs, all_metas = _extract_drawers(col, total, batch_size)
-    progress(f"  Extracted {len(all_ids)} drawers")
+    all_ids, all_docs, all_metas, all_embeddings = _extract_drawers(col, total, batch_size)
+    emb_status = (
+        "with stored embeddings"
+        if all_embeddings is not None
+        else "without embeddings (will recompute)"
+    )
+    progress(f"  Extracted {len(all_ids)} drawers ({emb_status})")
 
     # ── #1208 guard ──────────────────────────────────────────────────
     # Refuse to ``delete_collection`` + rebuild when extraction looks
@@ -793,6 +844,7 @@ def rebuild_index(
             batch_size,
             collection_name=collection_name,
             progress=progress,
+            all_embeddings=all_embeddings,
         )
     except RebuildCollectionError as e:
         progress(f"\n  ERROR during rebuild: {e}")
diff --git a/tests/test_repair.py b/tests/test_repair.py
@@ -87,7 +87,7 @@ def test_extract_drawers_preserves_valid_metadata():
         "documents": ["doc1", "doc2"],
         "metadatas": [{"wing": "a", "room": "1"}, {"wing": "b", "room": "2"}],
     }
-    all_ids, all_docs, all_metas = repair._extract_drawers(col, total=2, batch_size=2)
+    all_ids, all_docs, all_metas, _ = repair._extract_drawers(col, total=2, batch_size=2)
     assert all_ids == ["id1", "id2"]
     assert all_docs == ["doc1", "doc2"]
     assert all_metas == [{"wing": "a", "room": "1"}, {"wing": "b", "room": "2"}]
@@ -106,7 +106,7 @@ def test_extract_drawers_sanitizes_none_metadata():
         "documents": ["doc1", "doc2", "doc3"],
         "metadatas": [{"wing": "a"}, None, {"wing": "c"}],
     }
-    _, _, all_metas = repair._extract_drawers(col, total=3, batch_size=3)
+    _, _, all_metas, _ = repair._extract_drawers(col, total=3, batch_size=3)
     assert all_metas[0] == {"wing": "a"}
     assert all_metas[1] == {"_repaired_empty_meta": True}
     assert all_metas[2] == {"wing": "c"}
@@ -124,7 +124,7 @@ def test_extract_drawers_sanitizes_empty_dict_metadata():
         "documents": ["doc1", "doc2"],
         "metadatas": [{}, {"wing": "b"}],
     }
-    _, _, all_metas = repair._extract_drawers(col, total=2, batch_size=2)
+    _, _, all_metas, _ = repair._extract_drawers(col, total=2, batch_size=2)
     assert all_metas[0] == {"_repaired_empty_meta": True}
     assert all_metas[1] == {"wing": "b"}
 
@@ -142,7 +142,7 @@ def test_extract_drawers_sanitization_preserves_alignment():
         "documents": ["d1", "d2", "d3", "d4"],
         "metadatas": [None, {"k": "v"}, {}, None],
     }
-    all_ids, all_docs, all_metas = repair._extract_drawers(col, total=4, batch_size=4)
+    all_ids, all_docs, all_metas, _ = repair._extract_drawers(col, total=4, batch_size=4)
     assert len(all_ids) == len(all_docs) == len(all_metas) == 4
     assert all_ids == ["id1", "id2", "id3", "id4"]
     assert all_metas[0] == {"_repaired_empty_meta": True}
@@ -159,11 +159,67 @@ def test_extract_drawers_multiple_batches():
         {"ids": ["id3"], "documents": ["d3"], "metadatas": [{}]},
         {"ids": [], "documents": [], "metadatas": []},
     ]
-    all_ids, all_docs, all_metas = repair._extract_drawers(col, total=3, batch_size=2)
+    all_ids, all_docs, all_metas, _ = repair._extract_drawers(col, total=3, batch_size=2)
     assert all_ids == ["id1", "id2", "id3"]
     assert all_metas == [{"a": 1}, {"_repaired_empty_meta": True}, {"_repaired_empty_meta": True}]
 
 
+# ── embeddings preservation (upstream #1367 cherry-pick) ────────────
+
+
+def test_extract_drawers_preserves_embeddings_when_all_present():
+    """When every drawer has an embedding, return them aligned with ids.
+
+    The fast-path of upstream #1367: re-uses stored vectors during repair
+    instead of forcing ChromaDB to recompute via ONNX. Saves hours on
+    large palaces.
+    """
+    col = MagicMock()
+    col.get.return_value = {
+        "ids": ["id1", "id2"],
+        "documents": ["d1", "d2"],
+        "metadatas": [{"a": 1}, {"b": 2}],
+        "embeddings": [[0.1] * 384, [0.2] * 384],
+    }
+    _, _, _, all_embeddings = repair._extract_drawers(col, total=2, batch_size=2)
+    assert all_embeddings is not None
+    assert len(all_embeddings) == 2
+    assert all_embeddings[0] == [0.1] * 384
+
+
+def test_extract_drawers_falls_back_to_none_when_any_embedding_missing():
+    """If any drawer's embedding is None (HNSW unreadable), return None.
+
+    Signals to the rebuild caller that ChromaDB must recompute embeddings
+    rather than try to upsert with partial vectors. The whole-or-nothing
+    semantics mean we never end up with a half-vectored collection.
+    """
+    col = MagicMock()
+    col.get.return_value = {
+        "ids": ["id1", "id2", "id3"],
+        "documents": ["d1", "d2", "d3"],
+        "metadatas": [{}, {}, {}],
+        "embeddings": [[0.1] * 384, None, [0.3] * 384],
+    }
+    _, _, _, all_embeddings = repair._extract_drawers(col, total=3, batch_size=3)
+    assert all_embeddings is None, "any None forces caller to recompute all"
+
+
+def test_extract_drawers_embeddings_absent_in_batch_signals_recompute():
+    """When chromadb returns no `embeddings` key at all, the result is None."""
+    col = MagicMock()
+    col.get.return_value = {
+        "ids": ["id1"],
+        "documents": ["d1"],
+        "metadatas": [{}],
+        # 'embeddings' key absent entirely (chromadb may return None for
+        # this on unreadable HNSW)
+        "embeddings": None,
+    }
+    _, _, _, all_embeddings = repair._extract_drawers(col, total=1, batch_size=1)
+    assert all_embeddings is None
+
+
 # ── scan_palace ───────────────────────────────────────────────────────
 
 
