From 9d90335a12a719ee99c03fe2f0ce4fe3f13bd125 Mon Sep 17 00:00:00 2001
From: Mike Fortman <michael.fortman@datastax.com>
Date: Fri, 1 Dec 2023 15:26:24 -0600
Subject: [PATCH] add function descriptions

---
 examples/astradb/load.ts                      |  2 +-
 .../storage/vectorStore/AstraDBVectorStore.ts | 58 +++++++++++++------
 2 files changed, 42 insertions(+), 18 deletions(-)

diff --git a/examples/astradb/load.ts b/examples/astradb/load.ts
index 4bd875a75..ece5d8678 100644
--- a/examples/astradb/load.ts
+++ b/examples/astradb/load.ts
@@ -12,7 +12,7 @@ async function main() {
 
     const astraVS = new AstraDBVectorStore();
     await astraVS.create(collectionName, {
-      vector: { dimension: 1536, metric: "cosine" },
+      vector: { size: 1536, function: "cosine" },
     });
     await astraVS.connect(collectionName);
 
diff --git a/packages/core/src/storage/vectorStore/AstraDBVectorStore.ts b/packages/core/src/storage/vectorStore/AstraDBVectorStore.ts
index b9fc482a4..04cb8c26f 100644
--- a/packages/core/src/storage/vectorStore/AstraDBVectorStore.ts
+++ b/packages/core/src/storage/vectorStore/AstraDBVectorStore.ts
@@ -1,12 +1,7 @@
 import { AstraDB } from "@datastax/astra-db-ts";
 import { Collection } from "@datastax/astra-db-ts/dist/collections";
 import { BaseNode, MetadataMode } from "../../Node";
-import {
-  VectorStore,
-  VectorStoreQuery,
-  VectorStoreQueryMode,
-  VectorStoreQueryResult,
-} from "./types";
+import { VectorStore, VectorStoreQuery, VectorStoreQueryResult } from "./types";
 import { metadataDictToNode, nodeToMetadata } from "./utils";
 
 const MAX_INSERT_BATCH_SIZE = 20;
@@ -38,13 +33,29 @@ export class AstraDBVectorStore implements VectorStore {
     this.astraDBClient = new AstraDB(token, dbId, region, keyspace);
   }
 
-  async create(collection: string, options?: any): Promise<void> {
+  /**
+   * Create a new collection in your Astra DB vector database.
+   * You must still use connect() to connect to the collection.
+   *
+   * @TODO align options type with the JSON API's expected format
+   * @param collection your new colletion's name
+   * @param options: CreateCollectionOptions used to set the number of vector dimensions and similarity metric
+   * @returns Promise that resolves if the creation did not throw an error.
+   */
+  async create(collection: string, options: any): Promise<void> {
     await this.astraDBClient.createCollection(collection, options);
     console.debug("Created Astra DB collection");
 
     return;
   }
 
+  /**
+   * Connect to an existing collection in your Astra DB vector database.
+   * You must call this before adding, deleting, or querying.
+   *
+   * @param collection your existing colletion's name
+   * @returns Promise that resolves if the connection did not throw an error.
+   */
   async connect(collection: string): Promise<void> {
     this.collection = await this.astraDBClient.collection(collection);
     console.debug("Connected to Astra DB collection");
@@ -52,10 +63,19 @@ export class AstraDBVectorStore implements VectorStore {
     return;
   }
 
+  /**
+   * Get an instance of your Astra DB client.
+   * @returns the AstraDB client
+   */
   client(): AstraDB {
     return this.astraDBClient;
   }
 
+  /**
+   * Add your document(s) to your Astra DB collection.
+   *
+   * @returns and array of node ids which were added
+   */
   async add(nodes: BaseNode[]): Promise<string[]> {
     if (!this.collection) {
       throw new Error("Must connect to collection before adding.");
@@ -94,6 +114,13 @@ export class AstraDBVectorStore implements VectorStore {
     return dataToInsert.map((node) => node._id);
   }
 
+  /**
+   * Delete a document from your Astra DB collection.
+   *
+   * @param refDocId the id of the document to delete
+   * @param deleteOptions: any DeleteOneOptions to pass to the delete query
+   * @returns Promise that resolves if the delete query did not throw an error.
+   */
   async delete(refDocId: string, deleteOptions?: any): Promise<void> {
     if (!this.collection) {
       throw new Error("Must connect to collection before deleting.");
@@ -102,7 +129,7 @@ export class AstraDBVectorStore implements VectorStore {
 
     console.debug(`Deleting row with id ${refDocId}`);
 
-    await this.collection.deleteOne(
+    await collection.deleteOne(
       {
         _id: refDocId,
       },
@@ -110,6 +137,12 @@ export class AstraDBVectorStore implements VectorStore {
     );
   }
 
+  /**
+   * Query documents from your Astra DB collection to get the closest match to your embedding.
+   *
+   * @param query: VectorStoreQuery
+   * @param options: Not used
+   */
   async query(
     query: VectorStoreQuery,
     options?: any,
@@ -119,15 +152,6 @@ export class AstraDBVectorStore implements VectorStore {
     }
     const collection = this.collection;
 
-    const availableQueryModes = [
-      VectorStoreQueryMode.DEFAULT,
-      VectorStoreQueryMode.MMR,
-    ];
-
-    if (!availableQueryModes.includes(query.mode)) {
-      throw new Error("Query mode must be one of: " + availableQueryModes);
-    }
-
     const filters: Record<string, any> = {};
     query.filters?.filters?.forEach((f) => {
       filters[f.key] = f.value;
-- 
GitLab