PR comments on metadata handling

mfortman11 · Jan 31, 2024 · 733f74a · 733f74a
1 parent 0c843b0
commit 733f74a
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 14 deletions.
diff --git a/examples/astradb/example.ts b/examples/astradb/example.ts
@@ -14,18 +14,21 @@ async function main() {
         text: "AstraDB is built on Apache Cassandra",
         metadata: {
           id: 123,
+          foo: "bar",
         },
       }),
       new Document({
         text: "AstraDB is a NoSQL DB",
         metadata: {
           id: 456,
+          foo: "baz",
         },
       }),
       new Document({
         text: "AstraDB supports vector search",
         metadata: {
           id: 789,
+          foo: "qux",
         },
       }),
     ];

diff --git a/examples/astradb/query.ts b/examples/astradb/query.ts
@@ -19,7 +19,7 @@ async function main() {
     const queryEngine = await index.asQueryEngine({ retriever });
 
     const results = await queryEngine.query({
-      query: "What is the best reviewed movie?",
+      query: 'How was "La Sapienza" reviewed?',
     });
 
     console.log(results.response);

diff --git a/packages/core/src/storage/vectorStore/AstraDBVectorStore.ts b/packages/core/src/storage/vectorStore/AstraDBVectorStore.ts
@@ -1,8 +1,9 @@
 import { AstraDB } from "@datastax/astra-db-ts";
 import { Collection } from "@datastax/astra-db-ts/dist/collections";
 import { CreateCollectionOptions } from "@datastax/astra-db-ts/dist/collections/options";
-import { BaseNode, Document, MetadataMode } from "../../Node";
+import { BaseNode, MetadataMode } from "../../Node";
 import { VectorStore, VectorStoreQuery, VectorStoreQueryResult } from "./types";
+import { metadataDictToNode, nodeToMetadata } from "./utils";
 
 const MAX_INSERT_BATCH_SIZE = 20;
 
@@ -13,6 +14,7 @@ export class AstraDBVectorStore implements VectorStore {
   astraDBClient: AstraDB;
   idKey: string;
   contentKey: string;
+  metadataKey: string;
 
   private collection: Collection | undefined;
 
@@ -49,6 +51,7 @@ export class AstraDBVectorStore implements VectorStore {
 
     this.idKey = init?.idKey ?? "_id";
     this.contentKey = init?.contentKey ?? "text";
+    this.metadataKey = init?.metadataKey ?? "metadata";
   }
 
   /**
@@ -107,11 +110,18 @@ export class AstraDBVectorStore implements VectorStore {
     }
 
     const dataToInsert = nodes.map((node) => {
+      const metadata = nodeToMetadata(
+        node,
+        true,
+        this.contentKey,
+        this.flatMetadata,
+      );
+
       return {
-        $vector: node.embedding,
+        $vector: node.getEmbedding(),
         [this.idKey]: node.id_,
         [this.contentKey]: node.getContent(MetadataMode.NONE),
-        ...node.metadata,
+        [this.metadataKey]: metadata,
       };
     });
 
@@ -129,7 +139,7 @@ export class AstraDBVectorStore implements VectorStore {
       await collection.insertMany(batch);
     }
 
-    return dataToInsert.map((node) => node._id);
+    return dataToInsert.map((node) => node?.[this.idKey] as string);
   }
 
   /**
@@ -189,19 +199,25 @@ export class AstraDBVectorStore implements VectorStore {
 
     await cursor.forEach(async (row: Record<string, any>) => {
       const {
-        [this.idKey]: id,
-        [this.contentKey]: content,
         $vector: embedding,
         $similarity: similarity,
-        ...metadata
+        [this.idKey]: id,
+        [this.contentKey]: content,
+        [this.metadataKey]: metadata = {},
+        ...rest
       } = row;
 
-      const node = new Document({
-        id_: id,
-        text: content,
-        metadata,
-        embedding,
-      });
+      if (!metadata?.["_node_content"]) {
+        metadata["_node_content"] = JSON.stringify({
+          id,
+          text: content,
+          metadata,
+          ...rest,
+        });
+      }
+
+      const node = metadataDictToNode(metadata);
+      node.setContent(content);
 
       ids.push(id);
       similarities.push(similarity);