fix persistence bug

.changeset/short-boats-confess.md

+---
+"llamaindex": patch
+---
+
+Fix persistence bug (thanks @HenryHengZJ)

apps/simple/persist.ts

-import fs from "fs/promises";
-import {
-  Document,
-  VectorStoreIndex,
-  storageContextFromDefaults,
-} from "llamaindex";
-
-async function main() {
-  // Load essay from abramov.txt in Node
-  const essay = await fs.readFile(
-    "node_modules/llamaindex/examples/abramov.txt",
-    "utf-8"
-  );
-
-  // Create Document object with essay
-  const document = new Document({ text: essay });
-
-  // Split text and create embeddings. Store them in a VectorStoreIndex with persistence
-  const storageContext = await storageContextFromDefaults({
-    persistDir: "./storage",
-  });
-  const index = await VectorStoreIndex.fromDocuments([document], {
-    storageContext,
-  });
-
-  // Query the index
-  const queryEngine = index.asQueryEngine();
-  const response = await queryEngine.query(
-    "What did the author do in college?"
-  );
-
-  // Output response
-  console.log(response.toString());
-}
-
-main().catch(console.error);

apps/simple/storageContext.ts

-import { Document, VectorStoreIndex, storageContextFromDefaults } from "llamaindex";
+import {
+  Document,
+  VectorStoreIndex,
+  storageContextFromDefaults,
+} from "llamaindex";
 import essay from "./essay";
 
-
 async function main() {
   // Create Document object with essay
   const document = new Document({ text: essay });
 
   // Split text and create embeddings. Store them in a VectorStoreIndex
   // persist the vector store automatically with the storage context
-  const storageContext = await storageContextFromDefaults({ persistDir: "./storage" });
-  const index = await VectorStoreIndex.fromDocuments([document], { storageContext });
+  const storageContext = await storageContextFromDefaults({
+    persistDir: "./storage",
+  });
+  const index = await VectorStoreIndex.fromDocuments([document], {
+    storageContext,
+  });
 
   // Query the index
   const queryEngine = index.asQueryEngine();
@@ -21,9 +28,14 @@ async function main() {
   console.log(response.toString());
 
   // load the index
-  const loadedIndex = await VectorStoreIndex.init({ storageContext });
-  const laodedQueryEngine = loadedIndex.asQueryEngine();
-  const loadedResponse = await laodedQueryEngine.query(
+  const secondStorageContext = await storageContextFromDefaults({
+    persistDir: "./storage",
+  });
+  const loadedIndex = await VectorStoreIndex.init({
+    storageContext: secondStorageContext,
+  });
+  const loadedQueryEngine = loadedIndex.asQueryEngine();
+  const loadedResponse = await loadedQueryEngine.query(
     "What did the author do growing up?"
   );
   console.log(loadedResponse.toString());

examples/storageContext.ts

-import { Document, VectorStoreIndex, storageContextFromDefaults } from "llamaindex";
+import {
+  Document,
+  VectorStoreIndex,
+  storageContextFromDefaults,
+} from "llamaindex";
 import essay from "./essay";
 
-
 async function main() {
   // Create Document object with essay
   const document = new Document({ text: essay });
 
   // Split text and create embeddings. Store them in a VectorStoreIndex
   // persist the vector store automatically with the storage context
-  const storageContext = await storageContextFromDefaults({ persistDir: "./storage" });
-  const index = await VectorStoreIndex.fromDocuments([document], { storageContext });
+  const storageContext = await storageContextFromDefaults({
+    persistDir: "./storage",
+  });
+  const index = await VectorStoreIndex.fromDocuments([document], {
+    storageContext,
+  });
 
   // Query the index
   const queryEngine = index.asQueryEngine();
@@ -21,9 +28,14 @@ async function main() {
   console.log(response.toString());
 
   // load the index
-  const loadedIndex = await VectorStoreIndex.init({ storageContext });
-  const laodedQueryEngine = loadedIndex.asQueryEngine();
-  const loadedResponse = await laodedQueryEngine.query(
+  const secondStorageContext = await storageContextFromDefaults({
+    persistDir: "./storage",
+  });
+  const loadedIndex = await VectorStoreIndex.init({
+    storageContext: secondStorageContext,
+  });
+  const loadedQueryEngine = loadedIndex.asQueryEngine();
+  const loadedResponse = await loadedQueryEngine.query(
     "What did the author do growing up?"
   );
   console.log(loadedResponse.toString());

packages/core/src/Node.ts

@@ -128,6 +128,14 @@ export abstract class BaseNode {
       hash: this.hash,
     };
   }
+
+  /**
+   * Used with built in JSON.stringify
+   * @returns
+   */
+  toJSON(): Record<string, any> {
+    return { ...this, type: this.getType() };
+  }
 }
 
 /**
@@ -232,6 +240,23 @@ export class Document extends TextNode {
   }
 }
 
+export function jsonToNode(json: any) {
+  if (!json.type) {
+    throw new Error("Node type not found");
+  }
+
+  switch (json.type) {
+    case ObjectType.TEXT:
+      return new TextNode(json);
+    case ObjectType.INDEX:
+      return new IndexNode(json);
+    case ObjectType.DOCUMENT:
+      return new Document(json);
+    default:
+      throw new Error(`Invalid node type: ${json.type}`);
+  }
+}
+
 // export class ImageDocument extends Document {
 //   image?: string;
 // }

packages/core/src/indices/BaseIndex.ts

-import { Document, BaseNode } from "../Node";
+import { Document, BaseNode, jsonToNode } from "../Node";
 import { v4 as uuidv4 } from "uuid";
 import { BaseRetriever } from "../Retriever";
 import { ServiceContext } from "../ServiceContext";
@@ -74,7 +74,12 @@ export function jsonToIndexStruct(json: any): IndexStruct {
     return indexList;
   } else if (json.type === IndexStructType.SIMPLE_DICT) {
     const indexDict = new IndexDict(json.indexId, json.summary);
-    indexDict.nodesDict = json.nodesDict;
+    indexDict.nodesDict = Object.entries(json.nodesDict).reduce<
+      Record<string, BaseNode>
+    >((acc, [key, value]) => {
+      acc[key] = jsonToNode(value);
+      return acc;
+    }, {});
     return indexDict;
   } else {
     throw new Error(`Unknown index struct type: ${json.type}`);

packages/core/src/storage/constants.ts

 export const DEFAULT_COLLECTION = "data";
 export const DEFAULT_PERSIST_DIR = "./storage";
 export const DEFAULT_INDEX_STORE_PERSIST_FILENAME = "index_store.json";
-export const DEFAULT_DOC_STORE_PERSIST_FILENAME = "docstore.json";
+export const DEFAULT_DOC_STORE_PERSIST_FILENAME = "doc_store.json";
 export const DEFAULT_VECTOR_STORE_PERSIST_FILENAME = "vector_store.json";
 export const DEFAULT_GRAPH_STORE_PERSIST_FILENAME = "graph_store.json";
 export const DEFAULT_NAMESPACE = "docstore";