langchain 0.0.179 → 0.0.181

package/dist/vectorstores/cassandra.cjs

@@ -17,19 +17,19 @@ class CassandraStore extends base_js_1.VectorStore {
         return "cassandra";
     }
     constructor(embeddings, args) {
-        const argsWithDefaults = {
-            indices: [],
-            maxConcurrency: 25,
-            batchSize: 1,
-            ...args,
-        };
-        super(embeddings, argsWithDefaults);
+        super(embeddings, args);
         Object.defineProperty(this, "client", {
             enumerable: true,
             configurable: true,
             writable: true,
             value: void 0
         });
+        Object.defineProperty(this, "vectorType", {
+            enumerable: true,
+            configurable: true,
+            writable: true,
+            value: void 0
+        });
         Object.defineProperty(this, "dimensions", {
             enumerable: true,
             configurable: true,
@@ -54,6 +54,18 @@ class CassandraStore extends base_js_1.VectorStore {
             writable: true,
             value: void 0
         });
+        Object.defineProperty(this, "withClause", {
+            enumerable: true,
+            configurable: true,
+            writable: true,
+            value: void 0
+        });
+        Object.defineProperty(this, "selectColumns", {
+            enumerable: true,
+            configurable: true,
+            writable: true,
+            value: void 0
+        });
         Object.defineProperty(this, "table", {
             enumerable: true,
             configurable: true,
@@ -84,19 +96,27 @@ class CassandraStore extends base_js_1.VectorStore {
             writable: true,
             value: void 0
         });
-        this.asyncCaller = new async_caller_js_1.AsyncCaller(argsWithDefaults ?? {});
+        const { indices = [], maxConcurrency = 25, withClause = "", batchSize = 1, vectorType = "cosine", dimensions, keyspace, table, primaryKey, metadataColumns, } = args;
+        const argsWithDefaults = {
+            ...args,
+            indices,
+            maxConcurrency,
+            withClause,
+            batchSize,
+            vectorType,
+        };
+        this.asyncCaller = new async_caller_js_1.AsyncCaller(argsWithDefaults);
         this.client = new cassandra_driver_1.Client(argsWithDefaults);
-        this.dimensions = argsWithDefaults.dimensions;
-        this.keyspace = argsWithDefaults.keyspace;
-        this.table = argsWithDefaults.table;
-        this.primaryKey = argsWithDefaults.primaryKey;
-        this.metadataColumns = argsWithDefaults.metadataColumns;
-        this.indices = argsWithDefaults.indices;
-        this.batchSize = argsWithDefaults.batchSize;
-        if (this.batchSize < 1) {
-            console.warn("batchSize must be greater than or equal to 1, defaulting to 1");
-            this.batchSize = 1;
-        }
+        // Assign properties
+        this.vectorType = vectorType;
+        this.dimensions = dimensions;
+        this.keyspace = keyspace;
+        this.table = table;
+        this.primaryKey = Array.isArray(primaryKey) ? primaryKey : [primaryKey];
+        this.metadataColumns = metadataColumns;
+        this.withClause = withClause.trim().replace(/^with\s*/i, "");
+        this.indices = indices;
+        this.batchSize = batchSize >= 1 ? batchSize : 1;
     }
     /**
      * Method to save vectors to the Cassandra database.
@@ -132,16 +152,45 @@ class CassandraStore extends base_js_1.VectorStore {
         if (!this.isInitialized) {
             await this.initialize();
         }
-        const queryStr = this.buildSearchQuery(query, k, filter);
-        const queryResultSet = await this.client.execute(queryStr);
-        return queryResultSet?.rows.map((row, index) => {
+        // Ensure we have an array of Filter from the public interface
+        const filters = this.asFilters(filter);
+        const queryStr = this.buildSearchQuery(filters);
+        // Search query will be of format:
+        //   SELECT ..., text, similarity_x(?) AS similarity_score
+        //     FROM ...
+        //   <WHERE ...>
+        //    ORDER BY vector ANN OF ?
+        //    LIMIT ?
+        // If any filter values are specified, they will be in the WHERE clause as
+        //   filter.name filter.operator ?
+        // queryParams is a list of bind variables sent with the prepared statement
+        const queryParams = [];
+        const vectorAsFloat32Array = new Float32Array(query);
+        queryParams.push(vectorAsFloat32Array);
+        if (filters) {
+            const values = filters.map(({ value }) => value);
+            queryParams.push(...values);
+        }
+        queryParams.push(vectorAsFloat32Array);
+        queryParams.push(k);
+        const queryResultSet = await this.client.execute(queryStr, queryParams, {
+            prepare: true,
+        });
+        return queryResultSet?.rows.map((row) => {
             const textContent = row.text;
-            const sanitizedRow = Object.assign(row, {});
-            delete sanitizedRow.vector;
+            const sanitizedRow = { ...row };
             delete sanitizedRow.text;
+            delete sanitizedRow.similarity_score;
+            // A null value in Cassandra evaluates to a deleted column
+            // as this is treated as a tombstone record for the cell.
+            Object.keys(sanitizedRow).forEach((key) => {
+                if (sanitizedRow[key] === null) {
+                    delete sanitizedRow[key];
+                }
+            });
             return [
                 new document_js_1.Document({ pageContent: textContent, metadata: sanitizedRow }),
-                index,
+                row.similarity_score,
             ];
         });
     }
@@ -194,41 +243,142 @@ class CassandraStore extends base_js_1.VectorStore {
      * @returns Promise that resolves when the database has been initialized.
      */
     async initialize() {
-        await this.client.execute(`CREATE TABLE IF NOT EXISTS ${this.keyspace}.${this.table} (
-      ${this.primaryKey.name} ${this.primaryKey.type} PRIMARY KEY,
-      text TEXT,
+        let cql = "";
+        cql = `CREATE TABLE IF NOT EXISTS ${this.keyspace}.${this.table} (
+      ${this.primaryKey.map((col) => `${col.name} ${col.type}`).join(", ")}
+      , text TEXT
       ${this.metadataColumns.length > 0
-            ? this.metadataColumns.map((col) => `${col.name} ${col.type},`)
+            ? ", " +
+                this.metadataColumns
+                    .map((col) => `${col.name} ${col.type}`)
+                    .join(", ")
             : ""}
-      vector VECTOR<FLOAT, ${this.dimensions}>
-    );`);
-        await this.client
-            .execute(`CREATE CUSTOM INDEX IF NOT EXISTS idx_vector_${this.table}
-  ON ${this.keyspace}.${this.table}(vector) USING 'StorageAttachedIndex';`);
+      , vector VECTOR<FLOAT, ${this.dimensions}>
+      , ${this.buildPrimaryKey(this.primaryKey)}
+    ) ${this.withClause ? `WITH ${this.withClause}` : ""};`;
+        await this.client.execute(cql);
+        this.selectColumns = `${this.primaryKey
+            .map((col) => `${col.name}`)
+            .join(", ")}
+                          ${this.metadataColumns.length > 0
+            ? ", " +
+                this.metadataColumns
+                    .map((col) => `${col.name}`)
+                    .join(", ")
+            : ""}`;
+        cql = `CREATE CUSTOM INDEX IF NOT EXISTS idx_vector_${this.table}
+           ON ${this.keyspace}.${this.table}(vector) USING 'StorageAttachedIndex' WITH OPTIONS = {'similarity_function': '${this.vectorType.toUpperCase()}'};`;
+        await this.client.execute(cql);
         for await (const { name, value } of this.indices) {
-            await this.client
-                .execute(`CREATE CUSTOM INDEX IF NOT EXISTS idx_${this.table}_${name}
-  ON ${this.keyspace}.${this.table} ${value} USING 'StorageAttachedIndex';`);
+            cql = `CREATE CUSTOM INDEX IF NOT EXISTS idx_${this.table}_${name}
+             ON ${this.keyspace}.${this.table} ${value} USING 'StorageAttachedIndex';`;
+            await this.client.execute(cql);
         }
         this.isInitialized = true;
     }
-    buildWhereClause(filter) {
-        const whereClause = Object.entries(filter)
-            .map(([key, value]) => `${key} = '${value}'`)
-            .join(" AND ");
-        return `WHERE ${whereClause}`;
+    /**
+     * Method to build the PRIMARY KEY clause for CREATE TABLE.
+     * @param columns: list of Column to include in the key
+     * @returns The clause, including PRIMARY KEY
+     */
+    buildPrimaryKey(columns) {
+        // Partition columns may be specified with optional attribute col.partition
+        const partitionColumns = columns
+            .filter((col) => col.partition)
+            .map((col) => col.name)
+            .join(", ");
+        // All columns not part of the partition key are clustering columns
+        const clusteringColumns = columns
+            .filter((col) => !col.partition)
+            .map((col) => col.name)
+            .join(", ");
+        let primaryKey = "";
+        // If partition columns are specified, they are included in a () wrapper
+        // If not, the clustering columns are used, and the first clustering column
+        // is the partition key per normal Cassandra behaviour.
+        if (partitionColumns) {
+            primaryKey = `PRIMARY KEY ((${partitionColumns}), ${clusteringColumns})`;
+        }
+        else {
+            primaryKey = `PRIMARY KEY (${clusteringColumns})`;
+        }
+        return primaryKey;
+    }
+    /**
+     * Type guard to check if an object is a Filter.
+     * @param obj: the object to check
+     * @returns boolean indicating if the object is a Filter
+     */
+    isFilter(obj) {
+        return (typeof obj === "object" && obj !== null && "name" in obj && "value" in obj);
+    }
+    /**
+     * Helper to convert Record<string,unknown> to a Filter[]
+     * @param record: a key-value Record collection
+     * @returns Record as a Filter[]
+     */
+    convertToFilters(record) {
+        return Object.entries(record).map(([name, value]) => ({
+            name,
+            value,
+            operator: "=",
+        }));
+    }
+    /**
+     * Input santisation method for filters, as FilterType is not required to be
+     * Filter[], but we want to use Filter[] internally.
+     * @param record: the proposed filter
+     * @returns A Filter[], which may be empty
+     */
+    asFilters(record) {
+        if (!record) {
+            return [];
+        }
+        // If record is already an array
+        if (Array.isArray(record)) {
+            return record.flatMap((item) => {
+                // Check if item is a Filter before passing it to convertToFilters
+                if (this.isFilter(item)) {
+                    return [item];
+                }
+                else {
+                    // Here item is treated as Record<string, unknown>
+                    return this.convertToFilters(item);
+                }
+            });
+        }
+        // If record is a single Filter object, return it in an array
+        if (this.isFilter(record)) {
+            return [record];
+        }
+        // If record is a Record<string, unknown>, convert it to an array of Filter
+        return this.convertToFilters(record);
+    }
+    /**
+     * Method to build the WHERE clause of a CQL query, using bind variable ?
+     * @param filters list of filters to include in the WHERE clause
+     * @returns The WHERE clause
+     */
+    buildWhereClause(filters) {
+        if (!filters || filters.length === 0) {
+            return "";
+        }
+        const whereConditions = filters.map(({ name, operator = "=" }) => `${name} ${operator} ?`);
+        return `WHERE ${whereConditions.join(" AND ")}`;
     }
     /**
      * Method to build an CQL query for searching for similar vectors in the
      * Cassandra database.
      * @param query The query vector.
      * @param k The number of similar vectors to return.
-     * @param filter
+     * @param filters
      * @returns The CQL query string.
      */
-    buildSearchQuery(query, k = 1, filter = undefined) {
-        const whereClause = filter ? this.buildWhereClause(filter) : "";
-        return `SELECT * FROM ${this.keyspace}.${this.table} ${whereClause} ORDER BY vector ANN OF [${query}] LIMIT ${k}`;
+    buildSearchQuery(filters) {
+        const whereClause = filters ? this.buildWhereClause(filters) : "";
+        const cqlQuery = `SELECT ${this.selectColumns}, text, similarity_${this.vectorType}(vector, ?) AS similarity_score
+                        FROM ${this.keyspace}.${this.table} ${whereClause} ORDER BY vector ANN OF ? LIMIT ?`;
+        return cqlQuery;
     }
     /**
      * Method for inserting vectors and documents into the Cassandra database in a batch.
@@ -256,7 +406,7 @@ class CassandraStore extends base_js_1.VectorStore {
             const metadataInsert = metadataColNames.length > 0 ? ", " + metadataColNames.join(", ") : "";
             // Construct the query string and parameters
             const query = {
-                query: `INSERT INTO ${this.keyspace}.${this.table} (vector, text${metadataInsert}) 
+                query: `INSERT INTO ${this.keyspace}.${this.table} (vector, text${metadataInsert})
                 VALUES (?, ?${", ?".repeat(metadataColNames.length)})`,
                 params: [preparedVector, document.pageContent, ...metadataVals],
             };

package/dist/vectorstores/cassandra.d.ts

@@ -6,17 +6,27 @@ import { Document } from "../document.js";
 export interface Column {
     type: string;
     name: string;
+    partition?: boolean;
 }
 export interface Index {
     name: string;
     value: string;
 }
+export interface Filter {
+    name: string;
+    value: unknown;
+    operator?: string;
+}
+export type WhereClause = Filter[] | Filter | Record<string, unknown>;
+export type SupportedVectorTypes = "cosine" | "dot_product" | "euclidean";
 export interface CassandraLibArgs extends DseClientOptions, AsyncCallerParams {
     table: string;
     keyspace: string;
+    vectorType?: SupportedVectorTypes;
     dimensions: number;
-    primaryKey: Column;
+    primaryKey: Column | Column[];
     metadataColumns: Column[];
+    withClause?: string;
     indices?: Index[];
     batchSize?: number;
 }
@@ -27,12 +37,15 @@ export interface CassandraLibArgs extends DseClientOptions, AsyncCallerParams {
  * texts or documents.
  */
 export declare class CassandraStore extends VectorStore {
-    FilterType: Record<string, any>;
+    FilterType: WhereClause;
     private client;
+    private readonly vectorType;
     private readonly dimensions;
     private readonly keyspace;
     private primaryKey;
     private metadataColumns;
+    private withClause;
+    private selectColumns;
     private readonly table;
     private indices;
     private isInitialized;
@@ -60,7 +73,7 @@ export declare class CassandraStore extends VectorStore {
      * @param filter
      * @returns Promise that resolves with an array of tuples, each containing a Document and a score.
      */
-    similaritySearchVectorWithScore(query: number[], k: number, filter?: this["FilterType"]): Promise<[Document, number][]>;
+    similaritySearchVectorWithScore(query: number[], k: number, filter?: WhereClause): Promise<[Document, number][]>;
     /**
      * Static method to create an instance of CassandraStore from texts.
      * @param texts The texts to use.
@@ -91,13 +104,43 @@ export declare class CassandraStore extends VectorStore {
      * @returns Promise that resolves when the database has been initialized.
      */
     private initialize;
+    /**
+     * Method to build the PRIMARY KEY clause for CREATE TABLE.
+     * @param columns: list of Column to include in the key
+     * @returns The clause, including PRIMARY KEY
+     */
+    private buildPrimaryKey;
+    /**
+     * Type guard to check if an object is a Filter.
+     * @param obj: the object to check
+     * @returns boolean indicating if the object is a Filter
+     */
+    private isFilter;
+    /**
+     * Helper to convert Record<string,unknown> to a Filter[]
+     * @param record: a key-value Record collection
+     * @returns Record as a Filter[]
+     */
+    private convertToFilters;
+    /**
+     * Input santisation method for filters, as FilterType is not required to be
+     * Filter[], but we want to use Filter[] internally.
+     * @param record: the proposed filter
+     * @returns A Filter[], which may be empty
+     */
+    private asFilters;
+    /**
+     * Method to build the WHERE clause of a CQL query, using bind variable ?
+     * @param filters list of filters to include in the WHERE clause
+     * @returns The WHERE clause
+     */
     private buildWhereClause;
     /**
      * Method to build an CQL query for searching for similar vectors in the
      * Cassandra database.
      * @param query The query vector.
      * @param k The number of similar vectors to return.
-     * @param filter
+     * @param filters
      * @returns The CQL query string.
      */
     private buildSearchQuery;