diff --git a/packages/llamaindex/src/storage/docStore/KVDocumentStore.ts b/packages/llamaindex/src/storage/docStore/KVDocumentStore.ts index 223969b61c..510cbac68c 100644 --- 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> { @@ -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 { diff --git a/packages/llamaindex/src/storage/docStore/PostgresDocumentStore.ts b/packages/llamaindex/src/storage/docStore/PostgresDocumentStore.ts index 2a221d9533..72f20a3029 100644 --- 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 index 1735b81eea..b70ef14287 100644 --- 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) {