diff --git a/examples/astradb/example.ts b/examples/astradb/example.ts index dbf016d171..67b0170bc9 100644 --- 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 index 5720a413a1..23985c0d25 100644 --- 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 index d35ff313dd..f8c304156e 100644 --- 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) => { 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);