Migrate LanceDB to class (#4721)

migrate lancedb to class

Author: shatfield4

server/utils/helpers/index.js

@@ -95,7 +95,7 @@ function getVectorDbClass(getExactly = null) {
       const { ChromaCloud } = require("../vectorDbProviders/chromacloud");
       return ChromaCloud;
     case "lancedb":
-      return LanceDb;
+      return new LanceDb();
     case "weaviate":
       const { Weaviate } = require("../vectorDbProviders/weaviate");
       return Weaviate;
@@ -118,7 +118,7 @@ function getVectorDbClass(getExactly = null) {
       console.error(
         `\x1b[31m[ENV ERROR]\x1b[0m No VECTOR_DB value found in environment! Falling back to LanceDB`
       );
-      return LanceDb;
+      return new LanceDb();
   }
 }
 

server/utils/vectorDbProviders/lance/index.js

@@ -12,32 +12,38 @@ const { NativeEmbeddingReranker } = require("../../EmbeddingRerankers/native");
  * @typedef {import('@lancedb/lancedb').Connection} LanceClient
  */
 
-const LanceDb = {
-  uri: `${
-    !!process.env.STORAGE_DIR ? `${process.env.STORAGE_DIR}/` : "./storage/"
-  }lancedb`,
-  name: "LanceDb",
+class LanceDb {
+  constructor() {
+    this.uri = `${
+      !!process.env.STORAGE_DIR ? `${process.env.STORAGE_DIR}/` : "./storage/"
+    }lancedb`;
+    this.name = "LanceDb";
+  }
 
   /** @returns {Promise<{client: LanceClient}>} */
-  connect: async function () {
+  async connect() {
     const client = await lancedb.connect(this.uri);
     return { client };
-  },
-  distanceToSimilarity: function (distance = null) {
+  }
+
+  distanceToSimilarity(distance = null) {
     if (distance === null || typeof distance !== "number") return 0.0;
     if (distance >= 1.0) return 1;
     if (distance < 0) return 1 - Math.abs(distance);
     return 1 - distance;
-  },
-  heartbeat: async function () {
+  }
+
+  async heartbeat() {
     await this.connect();
     return { heartbeat: Number(new Date()) };
-  },
-  tables: async function () {
+  }
+
+  async tables() {
     const { client } = await this.connect();
     return await client.tableNames();
-  },
-  totalVectors: async function () {
+  }
+
+  async totalVectors() {
     const { client } = await this.connect();
     const tables = await client.tableNames();
     let count = 0;
@@ -46,15 +52,17 @@ const LanceDb = {
       count += await table.countRows();
     }
     return count;
-  },
-  namespaceCount: async function (_namespace = null) {
+  }
+
+  async namespaceCount(_namespace = null) {
     const { client } = await this.connect();
     const exists = await this.namespaceExists(client, _namespace);
     if (!exists) return 0;
 
     const table = await client.openTable(_namespace);
     return (await table.countRows()) || 0;
-  },
+  }
+
   /**
    * Performs a SimilaritySearch + Reranking on a namespace.
    * @param {Object} params - The parameters for the rerankedSimilarityResponse.
@@ -67,7 +75,7 @@ const LanceDb = {
    * @param {string[]} params.filterIdentifiers - The identifiers of the documents to filter out.
    * @returns
    */
-  rerankedSimilarityResponse: async function ({
+  async rerankedSimilarityResponse({
     client,
     namespace,
     query,
@@ -138,7 +146,7 @@ const LanceDb = {
       });
 
     return result;
-  },
+  }
 
   /**
    * Performs a SimilaritySearch on a give LanceDB namespace.
@@ -151,7 +159,7 @@ const LanceDb = {
    * @param {string[]} params.filterIdentifiers
    * @returns
    */
-  similarityResponse: async function ({
+  async similarityResponse({
     client,
     namespace,
     queryVector,
@@ -192,30 +200,32 @@ const LanceDb = {
     });
 
     return result;
-  },
+  }
+
   /**
    *
    * @param {LanceClient} client
    * @param {string} namespace
    * @returns
    */
-  namespace: async function (client, namespace = null) {
+  async namespace(client, namespace = null) {
     if (!namespace) throw new Error("No namespace value provided.");
     const collection = await client.openTable(namespace).catch(() => false);
     if (!collection) return null;
 
     return {
       ...collection,
     };
-  },
+  }
+
   /**
    *
    * @param {LanceClient} client
    * @param {number[]} data
    * @param {string} namespace
    * @returns
    */
-  updateOrCreateCollection: async function (client, data = [], namespace) {
+  async updateOrCreateCollection(client, data = [], namespace) {
     const hasNamespace = await this.hasNamespace(namespace);
     if (hasNamespace) {
       const collection = await client.openTable(namespace);
@@ -225,35 +235,39 @@ const LanceDb = {
 
     await client.createTable(namespace, data);
     return true;
-  },
-  hasNamespace: async function (namespace = null) {
+  }
+
+  async hasNamespace(namespace = null) {
     if (!namespace) return false;
     const { client } = await this.connect();
     const exists = await this.namespaceExists(client, namespace);
     return exists;
-  },
+  }
+
   /**
    *
    * @param {LanceClient} client
    * @param {string} namespace
    * @returns
    */
-  namespaceExists: async function (client, namespace = null) {
+  async namespaceExists(client, namespace = null) {
     if (!namespace) throw new Error("No namespace value provided.");
     const collections = await client.tableNames();
     return collections.includes(namespace);
-  },
+  }
+
   /**
    *
    * @param {LanceClient} client
    * @param {string} namespace
    * @returns
    */
-  deleteVectorsInNamespace: async function (client, namespace = null) {
+  async deleteVectorsInNamespace(client, namespace = null) {
     await client.dropTable(namespace);
     return true;
-  },
-  deleteDocumentFromNamespace: async function (namespace, docId) {
+  }
+
+  async deleteDocumentFromNamespace(namespace, docId) {
     const { client } = await this.connect();
     const exists = await this.namespaceExists(client, namespace);
     if (!exists) {
@@ -272,8 +286,9 @@ const LanceDb = {
     if (vectorIds.length === 0) return;
     await table.delete(`id IN (${vectorIds.map((v) => `'${v}'`).join(",")})`);
     return true;
-  },
-  addDocumentToNamespace: async function (
+  }
+
+  async addDocumentToNamespace(
     namespace,
     documentData = {},
     fullFilePath = null,
@@ -376,8 +391,9 @@ const LanceDb = {
       console.error("addDocumentToNamespace", e.message);
       return { vectorized: false, error: e.message };
     }
-  },
-  performSimilaritySearch: async function ({
+  }
+
+  async performSimilaritySearch({
     namespace = null,
     input = "",
     LLMConnector = null,
@@ -427,8 +443,9 @@ const LanceDb = {
       sources: this.curateSources(sources),
       message: false,
     };
-  },
-  "namespace-stats": async function (reqBody = {}) {
+  }
+
+  async "namespace-stats"(reqBody = {}) {
     const { namespace = null } = reqBody;
     if (!namespace) throw new Error("namespace required");
     const { client } = await this.connect();
@@ -438,8 +455,9 @@ const LanceDb = {
     return stats
       ? stats
       : { message: "No stats were able to be fetched from DB for namespace" };
-  },
-  "delete-namespace": async function (reqBody = {}) {
+  }
+
+  async "delete-namespace"(reqBody = {}) {
     const { namespace = null } = reqBody;
     const { client } = await this.connect();
     if (!(await this.namespaceExists(client, namespace)))
@@ -449,14 +467,16 @@ const LanceDb = {
     return {
       message: `Namespace ${namespace} was deleted.`,
     };
-  },
-  reset: async function () {
+  }
+
+  async reset() {
     const { client } = await this.connect();
     const fs = require("fs");
     fs.rm(`${client.uri}`, { recursive: true }, () => null);
     return { reset: true };
-  },
-  curateSources: function (sources = []) {
+  }
+
+  curateSources(sources = []) {
     const documents = [];
     for (const source of sources) {
       const { text, vector: _v, _distance: _d, ...rest } = source;
@@ -470,7 +490,7 @@ const LanceDb = {
     }
 
     return documents;
-  },
-};
+  }
+}
 
 module.exports.LanceDb = LanceDb;