VectorDB class migration (#4787)

* Migrate Astra to class (#4722)

migrate astra to class

* Migrate LanceDB to class (#4721)

migrate lancedb to class

* Migrate Pinecone to class (#4726)

migrate pinecone to class

* Migrate Zilliz to class (#4729)

migrate zilliz to class

* Migrate Weaviate to class (#4728)

migrate weaviate to class

* Migrate Qdrant to class (#4727)

migrate qdrant to class

* Migrate Milvus to class (#4725)

migrate milvus to class

* Migrate Chroma to class (#4723)

migrate chroma to class

* Migrate Chroma Cloud to class (#4724)

* migrate chroma to class

* migrate chroma cloud to class

* move limits to class field

---------

Co-authored-by: Timothy Carambat <[email protected]>

* Migrate PGVector to class (#4730)

* migrate pgvector to class

* patch pgvector test

* convert connectionString, tableName, and validateConnection to static methods

* move instance properties to class fields

---------

Co-authored-by: Timothy Carambat <[email protected]>

* Refactor Zilliz Cloud vector DB provider (#4749)

simplify zilliz implementation by using milvus as base class

Co-authored-by: Timothy Carambat <[email protected]>

* VectorDatabase base class (#4738)

create generic VectorDatabase base class

Co-authored-by: Timothy Carambat <[email protected]>

* Extend VectorDatabase base class to all providers (#4755)

extend VectorDatabase base class to all providers

* patch lancedb import

* breakout name and add generic logger

* dev tag build

---------

Co-authored-by: Timothy Carambat <[email protected]>

Author: shatfield4

.github/workflows/dev-build.yaml

@@ -6,7 +6,7 @@ concurrency:
 
 on:
   push:
-    branches: ['4841-aws-bedrock-api-key'] # put your current branch to create a build. Core team only.
+    branches: ['vectordb-class-migration'] # put your current branch to create a build. Core team only.
     paths-ignore:
       - '**.md'
       - 'cloud-deployments/*'

server/__tests__/utils/vectorDbProviders/pgvector/index.test.js

@@ -1,4 +1,6 @@
-const { PGVector } = require("../../../../utils/vectorDbProviders/pgvector");
+const { PGVector: PGVectorClass } = require("../../../../utils/vectorDbProviders/pgvector");
+
+const PGVector = new PGVectorClass();
 
 describe("PGVector.sanitizeForJsonb", () => {
   it("returns null/undefined as-is", () => {

server/utils/helpers/index.js

@@ -86,40 +86,40 @@ function getVectorDbClass(getExactly = null) {
   switch (vectorSelection) {
     case "pinecone":
       const { Pinecone } = require("../vectorDbProviders/pinecone");
-      return Pinecone;
+      return new Pinecone();
     case "chroma":
       const { Chroma } = require("../vectorDbProviders/chroma");
-      return Chroma;
+      return new Chroma();
     case "chromacloud":
       const { ChromaCloud } = require("../vectorDbProviders/chromacloud");
-      return ChromaCloud;
+      return new ChromaCloud();
     case "lancedb":
       const { LanceDb } = require("../vectorDbProviders/lance");
-      return LanceDb;
+      return new LanceDb();
     case "weaviate":
       const { Weaviate } = require("../vectorDbProviders/weaviate");
-      return Weaviate;
+      return new Weaviate();
     case "qdrant":
       const { QDrant } = require("../vectorDbProviders/qdrant");
-      return QDrant;
+      return new QDrant();
     case "milvus":
       const { Milvus } = require("../vectorDbProviders/milvus");
-      return Milvus;
+      return new Milvus();
     case "zilliz":
       const { Zilliz } = require("../vectorDbProviders/zilliz");
-      return Zilliz;
+      return new Zilliz();
     case "astra":
       const { AstraDB } = require("../vectorDbProviders/astra");
-      return AstraDB;
+      return new AstraDB();
     case "pgvector":
       const { PGVector } = require("../vectorDbProviders/pgvector");
-      return PGVector;
+      return new PGVector();
     default:
       console.error(
         `\x1b[31m[ENV ERROR]\x1b[0m No VECTOR_DB value found in environment! Falling back to LanceDB`
       );
       const { LanceDb: DefaultLanceDb } = require("../vectorDbProviders/lance");
-      return DefaultLanceDb;
+      return new DefaultLanceDb();
   }
 }
 

server/utils/vectorDbProviders/astra/index.js

@@ -5,6 +5,7 @@ const { storeVectorResult, cachedVectorInformation } = require("../../files");
 const { v4: uuidv4 } = require("uuid");
 const { toChunks, getEmbeddingEngineSelection } = require("../../helpers");
 const { sourceIdentifier } = require("../../chats");
+const { VectorDatabase } = require("../base");
 
 const sanitizeNamespace = (namespace) => {
   // If namespace already starts with ns_, don't add it again
@@ -22,14 +23,21 @@ const collectionExists = async function (client, namespace) {
       return collections.includes(namespace);
     }
   } catch (error) {
-    console.log("Astra::collectionExists check error", error?.message || error);
+    this.logger("collectionExists check error", error?.message || error);
     return false; // Return false for any error to allow creation attempt
   }
 };
 
-const AstraDB = {
-  name: "AstraDB",
-  connect: async function () {
+class AstraDB extends VectorDatabase {
+  constructor() {
+    super();
+  }
+
+  get name() {
+    return "AstraDB";
+  }
+
+  async connect() {
     if (process.env.VECTOR_DB !== "astra")
       throw new Error("AstraDB::Invalid ENV settings");
 
@@ -38,21 +46,24 @@ const AstraDB = {
       process?.env?.ASTRA_DB_ENDPOINT
     );
     return { client };
-  },
-  heartbeat: async function () {
+  }
+
+  async heartbeat() {
     return { heartbeat: Number(new Date()) };
-  },
+  }
+
   // Astra interface will return a valid collection object even if the collection
   // does not actually exist. So we run a simple check which will always throw
   // when the table truly does not exist. Faster than iterating all collections.
-  isRealCollection: async function (astraCollection = null) {
+  async isRealCollection(astraCollection = null) {
     if (!astraCollection) return false;
     return await astraCollection
       .countDocuments()
       .then(() => true)
       .catch(() => false);
-  },
-  totalVectors: async function () {
+  }
+
+  async totalVectors() {
     const { client } = await this.connect();
     const collectionNames = await this.allNamespaces(client);
     var totalVectors = 0;
@@ -62,13 +73,15 @@ const AstraDB = {
       totalVectors += count ? count : 0;
     }
     return totalVectors;
-  },
-  namespaceCount: async function (_namespace = null) {
+  }
+
+  async namespaceCount(_namespace = null) {
     const { client } = await this.connect();
     const namespace = await this.namespace(client, _namespace);
     return namespace?.vectorCount || 0;
-  },
-  namespace: async function (client, namespace = null) {
+  }
+
+  async namespace(client, namespace = null) {
     if (!namespace) throw new Error("No namespace value provided.");
     const sanitizedNamespace = sanitizeNamespace(namespace);
     const collection = await client
@@ -77,7 +90,7 @@ const AstraDB = {
     if (!(await this.isRealCollection(collection))) return null;
 
     const count = await collection.countDocuments().catch((e) => {
-      console.error("Astra::namespaceExists", e.message);
+      this.logger("namespaceExists", e.message);
       return null;
     });
 
@@ -86,27 +99,31 @@ const AstraDB = {
       ...collection,
       vectorCount: typeof count === "number" ? count : 0,
     };
-  },
-  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 sanitizedNamespace = sanitizeNamespace(namespace);
     const collection = await client.collection(sanitizedNamespace);
     return await this.isRealCollection(collection);
-  },
-  deleteVectorsInNamespace: async function (client, namespace = null) {
+  }
+
+  async deleteVectorsInNamespace(client, namespace = null) {
     const sanitizedNamespace = sanitizeNamespace(namespace);
     await client.dropCollection(sanitizedNamespace);
     return true;
-  },
+  }
+
   // AstraDB 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 sanitizedNamespace = sanitizeNamespace(namespace);
     try {
       const exists = await collectionExists(client, sanitizedNamespace);
@@ -132,14 +149,12 @@ const AstraDB = {
 
       return await client.collection(sanitizedNamespace);
     } catch (error) {
-      console.error(
-        "Astra::getOrCreateCollection error",
-        error?.message || error
-      );
+      this.logger("getOrCreateCollection", error?.message || error);
       throw error;
     }
-  },
-  addDocumentToNamespace: async function (
+  }
+
+  async addDocumentToNamespace(
     namespace,
     documentData = {},
     fullFilePath = null,
@@ -151,7 +166,7 @@ const AstraDB = {
       const { pageContent, docId, ...metadata } = documentData;
       if (!pageContent || pageContent.length == 0) return false;
 
-      console.log("Adding new vectorized document into namespace", namespace);
+      this.logger("Adding new vectorized document into namespace", namespace);
       if (!skipCache) {
         const cacheResult = await cachedVectorInformation(fullFilePath);
         if (cacheResult.exists) {
@@ -210,7 +225,7 @@ const AstraDB = {
       });
       const textChunks = await textSplitter.splitText(pageContent);
 
-      console.log("Snippets created from document:", textChunks.length);
+      this.logger("Snippets created from document:", textChunks.length);
       const documentVectors = [];
       const vectors = [];
       const vectorValues = await EmbedderEngine.embedChunks(textChunks);
@@ -246,7 +261,7 @@ const AstraDB = {
       if (vectors.length > 0) {
         const chunks = [];
 
-        console.log("Inserting vectorized chunks into Astra DB.");
+        this.logger("Inserting vectorized chunks into Astra DB.");
 
         // AstraDB has maximum upsert size of 20 records per-request so we have to use a lower chunk size here
         // in order to do the queries - this takes a lot more time than other providers but there
@@ -266,11 +281,12 @@ const AstraDB = {
       await DocumentVectors.bulkInsert(documentVectors);
       return { vectorized: true, error: null };
     } catch (e) {
-      console.error("addDocumentToNamespace", e.message);
+      this.logger("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();
     namespace = sanitizeNamespace(namespace);
@@ -293,8 +309,9 @@ const AstraDB = {
     const indexes = knownDocuments.map((doc) => doc.id);
     await DocumentVectors.deleteIds(indexes);
     return true;
-  },
-  performSimilaritySearch: async function ({
+  }
+
+  async performSimilaritySearch({
     namespace = null,
     input = "",
     LLMConnector = null,
@@ -336,8 +353,9 @@ const AstraDB = {
       sources: this.curateSources(sources),
       message: false,
     };
-  },
-  similarityResponse: async function ({
+  }
+
+  async similarityResponse({
     client,
     namespace,
     queryVector,
@@ -367,8 +385,8 @@ const AstraDB = {
     responses.forEach((response) => {
       if (response.$similarity < similarityThreshold) return;
       if (filterIdentifiers.includes(sourceIdentifier(response.metadata))) {
-        console.log(
-          "AstraDB: A source was filtered from context as it's parent document is pinned."
+        this.logger(
+          "A source was filtered from context as it's parent document is pinned."
         );
         return;
       }
@@ -380,8 +398,9 @@ const AstraDB = {
       result.scores.push(response.$similarity);
     });
     return result;
-  },
-  allNamespaces: async function (client) {
+  }
+
+  async allNamespaces(client) {
     try {
       let header = new Headers();
       header.append("Token", client?.httpClient?.applicationToken);
@@ -403,11 +422,12 @@ const AstraDB = {
       const collections = resp ? JSON.parse(resp)?.status?.collections : [];
       return collections;
     } catch (e) {
-      console.error("Astra::AllNamespace", e);
+      this.logger("AllNamespace", e);
       return [];
     }
-  },
-  "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();
@@ -417,8 +437,9 @@ const AstraDB = {
     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)))
@@ -431,8 +452,9 @@ const AstraDB = {
         details?.vectorCount || "all"
       } vectors.`,
     };
-  },
-  curateSources: function (sources = []) {
+  }
+
+  curateSources(sources = []) {
     const documents = [];
     for (const source of sources) {
       if (Object.keys(source).length > 0) {
@@ -446,7 +468,7 @@ const AstraDB = {
     }
 
     return documents;
-  },
-};
+  }
+}
 
 module.exports.AstraDB = AstraDB;