From 6e205ccdcbf0a290361192bffb7e8d4c93abfc8e Mon Sep 17 00:00:00 2001
From: Logan Yang <logancyang@gmail.com>
Date: Tue, 24 Dec 2024 18:24:27 -0800
Subject: [PATCH] Fix lost embeddings (#967)

---
 src/search/chunkedStorage.ts | 21 ++++++++++-----------
 src/search/dbOperations.ts   |  1 +
 2 files changed, 11 insertions(+), 11 deletions(-)

diff --git a/src/search/chunkedStorage.ts b/src/search/chunkedStorage.ts
index aa968e18..f1e346f9 100644
--- a/src/search/chunkedStorage.ts
+++ b/src/search/chunkedStorage.ts
@@ -1,8 +1,7 @@
 import { CustomError } from "@/error";
 import { getSettings } from "@/settings/model";
-import { create, load, Orama, save } from "@orama/orama";
+import { create, load, Orama, RawData, save } from "@orama/orama";
 import { App } from "obsidian";
-import { DBOperations } from "./dbOperations";
 
 const CHUNK_PREFIX = "copilot-index-chunk-";
 const LEGACY_INDEX_SUFFIX = ".json";
@@ -106,8 +105,7 @@ export class ChunkedStorage {
 
   async saveDatabase(db: Orama<any>): Promise<void> {
     try {
-      const rawData = await save(db);
-      const documents = await DBOperations.getAllDocuments(db);
+      const rawData: RawData = await save(db);
       const numPartitions = getSettings().numPartitions;
 
       if (numPartitions === 1) {
@@ -117,18 +115,19 @@ export class ChunkedStorage {
           legacyPath,
           JSON.stringify({
             ...rawData,
-            documents: documents,
             schema: db.schema,
           })
         );
         return;
       }
 
+      const rawDocs = (rawData as any).docs?.docs;
+
       if (getSettings().debug) {
-        console.log(`Starting save with ${documents.length} total documents`);
+        console.log(`Starting save with ${rawDocs.length ?? 0} total documents`);
       }
 
-      if (!documents || documents.length === 0) {
+      if (!rawDocs || rawDocs.length === 0) {
         const metadata: ChunkMetadata = {
           numPartitions,
           vectorLength: db.schema.embedding.match(/\d+/)[0],
@@ -147,7 +146,7 @@ export class ChunkedStorage {
         return;
       }
 
-      const partitions = this.distributeDocumentsToPartitions(documents, numPartitions);
+      const partitions = this.distributeDocumentsToPartitions(rawDocs, numPartitions);
 
       const metadata: ChunkMetadata = {
         numPartitions,
@@ -155,7 +154,7 @@ export class ChunkedStorage {
         schema: db.schema,
         lastModified: Date.now(),
         documentPartitions: Object.fromEntries(
-          documents.map((doc) => [doc.id, this.assignDocumentToPartition(doc.id, numPartitions)])
+          rawDocs.map((doc: any) => [doc.id, this.assignDocumentToPartition(doc.id, numPartitions)])
         ),
       };
 
@@ -185,9 +184,9 @@ export class ChunkedStorage {
       }
 
       if (getSettings().debug) {
-        if (savedTotal !== documents.length) {
+        if (savedTotal !== rawDocs.length) {
           console.error(
-            `Document count mismatch during save! Original: ${documents.length}, Saved: ${savedTotal}`
+            `Document count mismatch during save! Original: ${rawDocs.length}, Saved: ${savedTotal}`
           );
         }
       }
diff --git a/src/search/dbOperations.ts b/src/search/dbOperations.ts
index 2be5f28f..6bca9107 100644
--- a/src/search/dbOperations.ts
+++ b/src/search/dbOperations.ts
@@ -447,6 +447,7 @@ export class DBOperations {
     const result = await search(db, {
       term: "",
       limit: 100000,
+      includeVectors: true,
     });
     return result.hits.map((hit) => hit.document);
   }