feat: add fallback to metadataDictToNode

examples/astradb/README.md

@@ -14,18 +14,27 @@ Here are two sample scripts which work well with the sample data in the Astra Po
 
 - `ASTRA_DB_APPLICATION_TOKEN`: The generated app token for your Astra database
 - `ASTRA_DB_ENDPOINT`: The API endpoint for your Astra database
+- `ASTRA_DB_NAMESPACE`: (Optional) The namespace where your collection is stored defaults to `default_keyspace`
 - `OPENAI_API_KEY`: Your OpenAI key
 
 2. `cd` Into the `examples` directory
 3. run `npm i`
 
-## Load the data
+## Example load and query
+
+Loads and queries a simple vectorstore with some documents about Astra DB
+
+run `ts-node astradb/example`
+
+## Movie Reviews Example
+
+### Load the data
 
 This sample loads the same dataset of movie reviews as the Astra Portal sample dataset. (Feel free to load the data in your the Astra Data Explorer to compare)
 
 run `ts-node astradb/load`
 
-## Use RAG to Query the data
+### Use RAG to Query the data
 
 Check out your data in the Astra Data Explorer and change the sample query as you see fit.
 

examples/astradb/example.ts

@@ -0,0 +1,58 @@
+import {
+  AstraDBVectorStore,
+  Document,
+  storageContextFromDefaults,
+  VectorStoreIndex,
+} from "llamaindex";
+
+const collectionName = "test_collection";
+
+async function main() {
+  try {
+    const docs = [
+      new Document({
+        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",
+        },
+      }),
+    ];
+
+    const astraVS = new AstraDBVectorStore();
+    await astraVS.create(collectionName, {
+      vector: { dimension: 1536, metric: "cosine" },
+    });
+    await astraVS.connect(collectionName);
+
+    const ctx = await storageContextFromDefaults({ vectorStore: astraVS });
+    const index = await VectorStoreIndex.fromDocuments(docs, {
+      storageContext: ctx,
+    });
+
+    const queryEngine = index.asQueryEngine();
+    const response = await queryEngine.query({
+      query: "Describe AstraDB.",
+    });
+
+    console.log(response.toString());
+  } catch (e) {
+    console.error(e);
+  }
+}
+
+main();

examples/astradb/load.ts

@@ -10,9 +10,9 @@ const collectionName = "movie_reviews";
 async function main() {
   try {
     const reader = new PapaCSVReader(false);
-    const docs = await reader.loadData("../data/movie_reviews.csv");
+    const docs = await reader.loadData("./data/movie_reviews.csv");
 
-    const astraVS = new AstraDBVectorStore();
+    const astraVS = new AstraDBVectorStore({ contentKey: "reviewtext" });
     await astraVS.create(collectionName, {
       vector: { dimension: 1536, metric: "cosine" },
     });

examples/astradb/query.ts

@@ -8,7 +8,7 @@ const collectionName = "movie_reviews";
 
 async function main() {
   try {
-    const astraVS = new AstraDBVectorStore();
+    const astraVS = new AstraDBVectorStore({ contentKey: "reviewtext" });
     await astraVS.connect(collectionName);
 
     const ctx = serviceContextFromDefaults();
@@ -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);

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;
 
@@ -12,7 +13,7 @@ export class AstraDBVectorStore implements VectorStore {
 
   astraDBClient: AstraDB;
   idKey: string;
-  contentKey: string | undefined; // if undefined the entirety of the node aside from the id and embedding will be stored as content
+  contentKey: string;
   metadataKey: string;
 
   private collection: Collection | undefined;
@@ -22,6 +23,7 @@ export class AstraDBVectorStore implements VectorStore {
       params?: {
         token: string;
         endpoint: string;
+        namespace: string;
       };
     },
   ) {
@@ -40,11 +42,15 @@ export class AstraDBVectorStore implements VectorStore {
       if (!endpoint) {
         throw new Error("Must specify ASTRA_DB_ENDPOINT via env variable.");
       }
-      this.astraDBClient = new AstraDB(token, endpoint);
+      const namespace =
+        init?.params?.namespace ??
+        process.env.ASTRA_DB_NAMESPACE ??
+        "default_keyspace";
+      this.astraDBClient = new AstraDB(token, endpoint, namespace);
     }
 
     this.idKey = init?.idKey ?? "_id";
-    this.contentKey = init?.contentKey;
+    this.contentKey = init?.contentKey ?? "content";
     this.metadataKey = init?.metadataKey ?? "metadata";
   }
 
@@ -102,12 +108,20 @@ export class AstraDBVectorStore implements VectorStore {
     if (!nodes || nodes.length === 0) {
       return [];
     }
+
     const dataToInsert = nodes.map((node) => {
+      const metadata = nodeToMetadata(
+        node,
+        true,
+        this.contentKey,
+        this.flatMetadata,
+      );
+
       return {
-        _id: node.id_,
         $vector: node.getEmbedding(),
-        content: node.getContent(MetadataMode.ALL),
-        metadata: node.metadata,
+        [this.idKey]: node.id_,
+        [this.contentKey]: node.getContent(MetadataMode.NONE),
+        [this.metadataKey]: metadata,
       };
     });
 
@@ -122,11 +136,10 @@ export class AstraDBVectorStore implements VectorStore {
 
     for (const batch of batchData) {
       console.debug(`Inserting batch of size ${batch.length}`);
-
-      const result = await collection.insertMany(batch);
+      await collection.insertMany(batch);
     }
 
-    return dataToInsert.map((node) => node._id);
+    return dataToInsert.map((node) => node?.[this.idKey] as string);
   }
 
   /**
@@ -185,27 +198,24 @@ export class AstraDBVectorStore implements VectorStore {
     const similarities: number[] = [];
 
     await cursor.forEach(async (row: Record<string, any>) => {
-      const id = row[this.idKey];
-      const embedding = row.$vector;
-      const similarity = row.$similarity;
-      const metadata = row[this.metadataKey];
+      const {
+        $vector: embedding,
+        $similarity: similarity,
+        [this.idKey]: id,
+        [this.contentKey]: content,
+        [this.metadataKey]: metadata = {},
+        ...rest
+      } = row;
 
-      // Remove fields from content
-      delete row[this.idKey];
-      delete row.$similarity;
-      delete row.$vector;
-      delete row[this.metadataKey];
-
-      const content = this.contentKey
-        ? row[this.contentKey]
-        : JSON.stringify(row);
-
-      const node = new Document({
-        id_: id,
-        text: content,
-        metadata: metadata ?? {},
-        embedding: embedding,
+      const node = metadataDictToNode(metadata, {
+        fallback: {
+          id,
+          text: content,
+          metadata,
+          ...rest,
+        },
       });
+      node.setContent(content);
 
       ids.push(id);
       similarities.push(similarity);

packages/core/src/storage/vectorStore/utils.ts

@@ -36,7 +36,16 @@ export function nodeToMetadata(
   return metadata;
 }
 
-export function metadataDictToNode(metadata: Metadata): BaseNode {
+type MetadataDictToNodeOptions = {
+  // If the metadata doesn't contain node content, use this object as a fallback, for usage see
+  // AstraDBVectorStore.ts
+  fallback: Record<string, any>;
+};
+
+export function metadataDictToNode(
+  metadata: Metadata,
+  options?: MetadataDictToNodeOptions,
+): BaseNode {
   const {
     _node_content: nodeContent,
     _node_type: nodeType,
@@ -45,11 +54,17 @@ export function metadataDictToNode(metadata: Metadata): BaseNode {
     ref_doc_id,
     ...rest
   } = metadata;
+  let nodeObj;
   if (!nodeContent) {
-    throw new Error("Node content not found in metadata.");
+    if (options?.fallback) {
+      nodeObj = options?.fallback;
+    } else {
+      throw new Error("Node content not found in metadata.");
+    }
+  } else {
+    nodeObj = JSON.parse(nodeContent);
+    nodeObj.metadata = rest;
   }
-  const nodeObj = JSON.parse(nodeContent);
-  nodeObj.metadata = rest;
 
   // Note: we're using the name of the class stored in `_node_type`
   // and not the type attribute to reconstruct

packages/eslint-config-custom/index.js

@@ -14,6 +14,7 @@ module.exports = {
 
           "ASTRA_DB_APPLICATION_TOKEN",
           "ASTRA_DB_ENDPOINT",
+          "ASTRA_DB_NAMESPACE",
 
           "AZURE_OPENAI_KEY",
           "AZURE_OPENAI_ENDPOINT",