Migrate Zilliz to class (#4729)

migrate zilliz to class

Author: shatfield4

server/utils/helpers/index.js

@@ -107,7 +107,7 @@ function getVectorDbClass(getExactly = null) {
       return Milvus;
     case "zilliz":
       const { Zilliz } = require("../vectorDbProviders/zilliz");
-      return Zilliz;
+      return new Zilliz();
     case "astra":
       const { AstraDB } = require("../vectorDbProviders/astra");
       return new AstraDB();

server/utils/vectorDbProviders/zilliz/index.js

@@ -13,19 +13,23 @@ const { sourceIdentifier } = require("../../chats");
 
 // Zilliz is basically a copy of Milvus DB class with a different constructor
 // to connect to the cloud
-const Zilliz = {
-  name: "Zilliz",
+class Zilliz {
+  constructor() {
+    this.name = "Zilliz";
+  }
+
   // Milvus/Zilliz only allows letters, numbers, and underscores in collection names
   // so we need to enforce that by re-normalizing the names when communicating with
   // the DB.
   // If the first char of the collection is not an underscore or letter the collection name will be invalid.
-  normalize: function (inputString) {
+  normalize(inputString) {
     let normalized = inputString.replace(/[^a-zA-Z0-9_]/g, "_");
     if (new RegExp(/^[a-zA-Z_]/).test(normalized.slice(0, 1)))
       normalized = `anythingllm_${normalized}`;
     return normalized;
-  },
-  connect: async function () {
+  }
+
+  async connect() {
     if (process.env.VECTOR_DB !== "zilliz")
       throw new Error("Zilliz::Invalid ENV settings");
 
@@ -41,12 +45,14 @@ const Zilliz = {
       );
 
     return { client };
-  },
-  heartbeat: async function () {
+  }
+
+  async heartbeat() {
     await this.connect();
     return { heartbeat: Number(new Date()) };
-  },
-  totalVectors: async function () {
+  }
+
+  async totalVectors() {
     const { client } = await this.connect();
     const { collection_names } = await client.listCollections();
     const total = collection_names.reduce(async (acc, collection_name) => {
@@ -56,27 +62,31 @@ const Zilliz = {
       return Number(acc) + Number(statistics?.data?.row_count ?? 0);
     }, 0);
     return total;
-  },
-  namespaceCount: async function (_namespace = null) {
+  }
+
+  async namespaceCount(_namespace = null) {
     const { client } = await this.connect();
     const statistics = await client.getCollectionStatistics({
       collection_name: this.normalize(_namespace),
     });
     return Number(statistics?.data?.row_count ?? 0);
-  },
-  namespace: async function (client, namespace = null) {
+  }
+
+  async namespace(client, namespace = null) {
     if (!namespace) throw new Error("No namespace value provided.");
     const collection = await client
       .getCollectionStatistics({ collection_name: this.normalize(namespace) })
       .catch(() => null);
     return collection;
-  },
-  hasNamespace: async function (namespace = null) {
+  }
+
+  async hasNamespace(namespace = null) {
     if (!namespace) return false;
     const { client } = await this.connect();
     return await this.namespaceExists(client, namespace);
-  },
-  namespaceExists: async function (client, namespace = null) {
+  }
+
+  async namespaceExists(client, namespace = null) {
     if (!namespace) throw new Error("No namespace value provided.");
     const { value } = await client
       .hasCollection({ collection_name: this.normalize(namespace) })
@@ -85,15 +95,17 @@ const Zilliz = {
         return { value: false };
       });
     return value;
-  },
-  deleteVectorsInNamespace: async function (client, namespace = null) {
+  }
+
+  async deleteVectorsInNamespace(client, namespace = null) {
     await client.dropCollection({ collection_name: this.normalize(namespace) });
     return true;
-  },
+  }
+
   // Zilliz requires a dimension aspect for collection creation
   // we pass this in from the first chunk to infer the dimensions like other
   // providers do.
-  getOrCreateCollection: async function (client, namespace, dimensions = null) {
+  async getOrCreateCollection(client, namespace, dimensions = null) {
     const isExists = await this.namespaceExists(client, namespace);
     if (!isExists) {
       if (!dimensions)
@@ -134,8 +146,9 @@ const Zilliz = {
         collection_name: this.normalize(namespace),
       });
     }
-  },
-  addDocumentToNamespace: async function (
+  }
+
+  async addDocumentToNamespace(
     namespace,
     documentData = {},
     fullFilePath = null,
@@ -261,8 +274,9 @@ const Zilliz = {
       console.error("addDocumentToNamespace", e.message);
       return { vectorized: false, error: e.message };
     }
-  },
-  deleteDocumentFromNamespace: async function (namespace, docId) {
+  }
+
+  async deleteDocumentFromNamespace(namespace, docId) {
     const { DocumentVectors } = require("../../../models/vectors");
     const { client } = await this.connect();
     if (!(await this.namespaceExists(client, namespace))) return;
@@ -284,8 +298,9 @@ const Zilliz = {
     // on a later call.
     await client.flushSync({ collection_names: [this.normalize(namespace)] });
     return true;
-  },
-  performSimilaritySearch: async function ({
+  }
+
+  async performSimilaritySearch({
     namespace = null,
     input = "",
     LLMConnector = null,
@@ -323,8 +338,9 @@ const Zilliz = {
       sources: this.curateSources(sources),
       message: false,
     };
-  },
-  similarityResponse: async function ({
+  }
+
+  async similarityResponse({
     client,
     namespace,
     queryVector,
@@ -358,8 +374,9 @@ const Zilliz = {
       result.scores.push(match.score);
     });
     return result;
-  },
-  "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();
@@ -369,8 +386,9 @@ const Zilliz = {
     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)))
@@ -382,8 +400,9 @@ const Zilliz = {
     return {
       message: `Namespace ${namespace} was deleted along with ${vectorCount} vectors.`,
     };
-  },
-  curateSources: function (sources = []) {
+  }
+
+  curateSources(sources = []) {
     const documents = [];
     for (const source of sources) {
       const { metadata = {} } = source;
@@ -396,7 +415,7 @@ const Zilliz = {
     }
 
     return documents;
-  },
-};
+  }
+}
 
 module.exports.Zilliz = Zilliz;