Updated packages/core/src/indices/vectorStore/VectorStoreIndex.ts

packages/core/src/Document.ts

@@ -0,0 +1,15 @@
+import { Node } from './Node';
+
+export class Document extends Node {
+  private text: string;
+
+  constructor(text: string) {
+    super();
+    this.text = text;
+  }
+
+  getContent(): string {
+    return this.text;
+  }
+}
+

packages/core/src/JsonDocument.ts

@@ -0,0 +1,15 @@
+import { Document } from './Document';
+
+export class JsonDocument extends Document {
+  private jsonData: object;
+
+  constructor(jsonData: object) {
+    super();
+    this.jsonData = jsonData;
+  }
+
+  getContent(): string {
+    return JSON.stringify(this.jsonData);
+  }
+}
+

packages/core/src/JsonNode.ts

@@ -0,0 +1,15 @@
+import { Node } from './Node';
+
+export class JsonNode extends Node {
+  private jsonData: object;
+
+  constructor(jsonData: object) {
+    super();
+    this.jsonData = jsonData;
+  }
+
+  getContent(): string {
+    return JSON.stringify(this.jsonData);
+  }
+}
+

packages/core/src/NodeParser.ts

@@ -1,4 +1,6 @@
 import { Document, NodeRelationship, TextNode } from "./Node";
+import { JsonDocument } from "./JsonDocument";
+import { JsonNode } from "./JsonNode";
 import { SentenceSplitter } from "./TextSplitter";
 import { DEFAULT_CHUNK_OVERLAP, DEFAULT_CHUNK_SIZE } from "./constants";
 
@@ -37,10 +39,18 @@ export function getNodesFromDocument(
   const textSplits = getTextSplitsFromDocument(document, textSplitter);
 
   textSplits.forEach((textSplit) => {
-    const node = new TextNode({
-      text: textSplit,
-      metadata: includeMetadata ? document.metadata : {},
-    });
+    let node;
+    if (document instanceof JsonDocument) {
+      node = new JsonNode({
+        jsonData: textSplit,
+        metadata: includeMetadata ? document.metadata : {},
+      });
+    } else {
+      node = new TextNode({
+        text: textSplit,
+        metadata: includeMetadata ? document.metadata : {},
+      });
+    }
     node.relationships[NodeRelationship.SOURCE] = document.asRelatedNodeInfo();
     nodes.push(node);
   });

packages/core/src/indices/vectorStore/VectorStoreIndex.ts

@@ -3,6 +3,8 @@ import {
   BaseNode,
   MetadataMode,
   NodeWithEmbedding,
+  JsonNode,
+} from "../../Node";
 } from "../../Node";
 import { BaseQueryEngine, RetrieverQueryEngine } from "../../QueryEngine";
 import { VectorIndexRetriever } from "./VectorIndexRetriever";
@@ -132,9 +134,16 @@ export class VectorStoreIndex extends BaseIndex<IndexDict> {
       if (logProgress) {
         console.log(`getting embedding for node ${i}/${nodes.length}`);
       }
-      const embedding = await serviceContext.embedModel.getTextEmbedding(
+    let embedding;
+    if (node instanceof JsonNode) {
+      embedding = await serviceContext.embedModel.getJsonEmbedding(
         node.getContent(MetadataMode.EMBED)
       );
+    } else {
+      embedding = await serviceContext.embedModel.getTextEmbedding(
+        node.getContent(MetadataMode.EMBED)
+      );
+    }
       nodesWithEmbeddings.push({ node, embedding });
     }