fix: use native json in pg doc store

run-llama · Sep 22, 2024 · f43d7fa · f43d7fa
1 parent a56fc0c
commit f43d7fa
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 10 deletions.
diff --git a/packages/llamaindex/src/storage/docStore/KVDocumentStore.ts b/packages/llamaindex/src/storage/docStore/KVDocumentStore.ts
@@ -14,13 +14,19 @@ export class KVDocumentStore extends BaseDocumentStore {
   private nodeCollection: string;
   private refDocCollection: string;
   private metadataCollection: string;
+  private nativeJson: boolean;
 
-  constructor(kvstore: BaseKVStore, namespace: string = DEFAULT_NAMESPACE) {
+  constructor(
+    kvstore: BaseKVStore,
+    namespace: string = DEFAULT_NAMESPACE,
+    nativeJson: boolean = false,
+  ) {
     super();
     this.kvstore = kvstore;
     this.nodeCollection = `${namespace}/data`;
     this.refDocCollection = `${namespace}/ref_doc_info`;
     this.metadataCollection = `${namespace}/metadata`;
+    this.nativeJson = nativeJson;
   }
 
   async docs(): Promise<Record<string, BaseNode>> {
@@ -29,7 +35,7 @@ export class KVDocumentStore extends BaseDocumentStore {
     for (const key in jsonDict) {
       const value = jsonDict[key];
       if (isValidDocJson(value)) {
-        docs[key] = jsonToDoc(value);
+        docs[key] = jsonToDoc(value, this.nativeJson);
       } else {
         console.warn(`Invalid JSON for docId ${key}`);
       }
@@ -52,7 +58,7 @@ export class KVDocumentStore extends BaseDocumentStore {
         );
       }
       const nodeKey = doc.id_;
-      const data = docToJson(doc);
+      const data = docToJson(doc, this.nativeJson);
       await this.kvstore.put(nodeKey, data, this.nodeCollection);
       const metadata: DocMetaData = { docHash: doc.hash };
 
@@ -94,7 +100,7 @@ export class KVDocumentStore extends BaseDocumentStore {
     if (!isValidDocJson(json)) {
       throw new Error(`Invalid JSON for docId ${docId}`);
     }
-    return jsonToDoc(json);
+    return jsonToDoc(json, this.nativeJson);
   }
 
   async getRefDocInfo(refDocId: string): Promise<RefDocInfo | undefined> {

diff --git a/packages/llamaindex/src/storage/docStore/PostgresDocumentStore.ts b/packages/llamaindex/src/storage/docStore/PostgresDocumentStore.ts
@@ -26,6 +26,7 @@ export class PostgresDocumentStore extends KVDocumentStore {
           : {}),
     });
     const namespace = config?.namespace || DEFAULT_NAMESPACE;
-    super(kvStore, namespace);
+    const nativeJson = true;
+    super(kvStore, namespace, nativeJson);
   }
 }
diff --git a/packages/llamaindex/src/storage/docStore/utils.ts b/packages/llamaindex/src/storage/docStore/utils.ts
@@ -6,7 +6,7 @@ const DATA_KEY = "__data__";
 
 type DocJson = {
   [TYPE_KEY]: ObjectType;
-  [DATA_KEY]: string;
+  [DATA_KEY]: string | any;
 };
 
 export function isValidDocJson(docJson: any): docJson is DocJson {
@@ -18,16 +18,18 @@ export function isValidDocJson(docJson: any): docJson is DocJson {
   );
 }
 
-export function docToJson(doc: BaseNode): DocJson {
+export function docToJson(doc: BaseNode, nativeJson?: boolean): DocJson {
   return {
-    [DATA_KEY]: JSON.stringify(doc.toJSON()),
+    [DATA_KEY]: nativeJson ? doc.toJSON() : JSON.stringify(doc.toJSON()),
     [TYPE_KEY]: doc.type,
   };
 }
 
-export function jsonToDoc(docDict: DocJson): BaseNode {
+export function jsonToDoc(docDict: DocJson, nativeJson?: boolean): BaseNode {
   const docType = docDict[TYPE_KEY];
-  const dataDict = JSON.parse(docDict[DATA_KEY]);
+  const dataDict = nativeJson
+    ? docDict[DATA_KEY]
+    : JSON.parse(docDict[DATA_KEY]);
   let doc: BaseNode;
 
   if (docType === ObjectType.DOCUMENT) {