feat: Add support for Zilliz Cloud by Milvus (#615)

* feat: Add support for Zilliz Cloud by Milvus

* update placeholder text
update data handling stmt

* update zilliz descriptor

.vscode/settings.json

@@ -2,10 +2,13 @@
   "cSpell.words": [
     "Dockerized",
     "Langchain",
+    "Milvus",
     "Ollama",
     "openai",
     "Qdrant",
-    "Weaviate"
+    "vectordbs",
+    "Weaviate",
+    "Zilliz"
   ],
   "eslint.experimental.useFlatConfig": true
 }

README.md

@@ -89,6 +89,7 @@ Some cool features of AnythingLLM
 - [Weaviate](https://weaviate.io)
 - [QDrant](https://qdrant.tech)
 - [Milvus](https://milvus.io)
+- [Zilliz](https://zilliz.com)
 
 ### Technical Overview
 

docker/.env.example

@@ -99,6 +99,11 @@ GID='1000'
 # MILVUS_USERNAME=
 # MILVUS_PASSWORD=
 
+# Enable all below if you are using vector database: Zilliz Cloud.
+# VECTOR_DB="zilliz"
+# ZILLIZ_ENDPOINT="https://sample.api.gcp-us-west1.zillizcloud.com"
+# ZILLIZ_API_TOKEN=api-token-here
+
 # CLOUD DEPLOYMENT VARIRABLES ONLY
 # AUTH_TOKEN="hunter2" # This is the password to your application if remote hosting.
 

frontend/src/components/VectorDBSelection/ZillizCloudOptions/index.jsx

@@ -0,0 +1,38 @@
+export default function ZillizCloudOptions({ settings }) {
+  return (
+    <div className="w-full flex flex-col gap-y-4">
+      <div className="w-full flex items-center gap-4">
+        <div className="flex flex-col w-60">
+          <label className="text-white text-sm font-semibold block mb-4">
+            Cluster Endpoint
+          </label>
+          <input
+            type="text"
+            name="ZillizEndpoint"
+            className="bg-zinc-900 text-white placeholder-white placeholder-opacity-60 text-sm rounded-lg focus:border-white block w-full p-2.5"
+            placeholder="https://sample.api.gcp-us-west1.zillizcloud.com"
+            defaultValue={settings?.ZillizEndpoint}
+            required={true}
+            autoComplete="off"
+            spellCheck={false}
+          />
+        </div>
+
+        <div className="flex flex-col w-60">
+          <label className="text-white text-sm font-semibold block mb-4">
+            API Token
+          </label>
+          <input
+            type="password"
+            name="ZillizApiToken"
+            className="bg-zinc-900 text-white placeholder-white placeholder-opacity-60 text-sm rounded-lg focus:border-white block w-full p-2.5"
+            placeholder="Zilliz cluster API Token"
+            defaultValue={settings?.ZillizApiToken ? "*".repeat(20) : ""}
+            autoComplete="off"
+            spellCheck={false}
+          />
+        </div>
+      </div>
+    </div>
+  );
+}

frontend/src/pages/GeneralSettings/VectorDatabase/index.jsx

@@ -9,6 +9,7 @@ import LanceDbLogo from "@/media/vectordbs/lancedb.png";
 import WeaviateLogo from "@/media/vectordbs/weaviate.png";
 import QDrantLogo from "@/media/vectordbs/qdrant.png";
 import MilvusLogo from "@/media/vectordbs/milvus.png";
+import ZillizLogo from "@/media/vectordbs/zilliz.png";
 import PreLoader from "@/components/Preloader";
 import ChangeWarningModal from "@/components/ChangeWarning";
 import { MagnifyingGlass } from "@phosphor-icons/react";
@@ -19,6 +20,7 @@ import QDrantDBOptions from "@/components/VectorDBSelection/QDrantDBOptions";
 import WeaviateDBOptions from "@/components/VectorDBSelection/WeaviateDBOptions";
 import VectorDBItem from "@/components/VectorDBSelection/VectorDBItem";
 import MilvusDBOptions from "@/components/VectorDBSelection/MilvusDBOptions";
+import ZillizCloudOptions from "@/components/VectorDBSelection/ZillizCloudOptions";
 
 export default function GeneralVectorDatabase() {
   const [saving, setSaving] = useState(false);
@@ -33,7 +35,6 @@ export default function GeneralVectorDatabase() {
   useEffect(() => {
     async function fetchKeys() {
       const _settings = await System.keys();
-      console.log(_settings);
       setSettings(_settings);
       setSelectedVDB(_settings?.VectorDB || "lancedb");
       setHasEmbeddings(_settings?.HasExistingEmbeddings || false);
@@ -66,6 +67,14 @@ export default function GeneralVectorDatabase() {
       options: <PineconeDBOptions settings={settings} />,
       description: "100% cloud-based vector database for enterprise use cases.",
     },
+    {
+      name: "Zilliz Cloud",
+      value: "zilliz",
+      logo: ZillizLogo,
+      options: <ZillizCloudOptions settings={settings} />,
+      description:
+        "Cloud hosted vector database built for enterprise with SOC 2 compliance.",
+    },
     {
       name: "QDrant",
       value: "qdrant",

frontend/src/pages/OnboardingFlow/Steps/DataHandling/index.jsx

@@ -10,6 +10,7 @@ import TogetherAILogo from "@/media/llmprovider/togetherai.png";
 import LMStudioLogo from "@/media/llmprovider/lmstudio.png";
 import LocalAiLogo from "@/media/llmprovider/localai.png";
 import MistralLogo from "@/media/llmprovider/mistral.jpeg";
+import ZillizLogo from "@/media/vectordbs/zilliz.png";
 import ChromaLogo from "@/media/vectordbs/chroma.png";
 import PineconeLogo from "@/media/vectordbs/pinecone.png";
 import LanceDbLogo from "@/media/vectordbs/lancedb.png";
@@ -139,6 +140,13 @@ const VECTOR_DB_PRIVACY = {
     ],
     logo: MilvusLogo,
   },
+  zilliz: {
+    name: "Zilliz Cloud",
+    description: [
+      "Your vectors and document text are stored on your Zilliz cloud cluster.",
+    ],
+    logo: ZillizLogo,
+  },
   lancedb: {
     name: "LanceDB",
     description: [

frontend/src/pages/OnboardingFlow/Steps/VectorDatabaseConnection/index.jsx

@@ -6,6 +6,7 @@ import LanceDbLogo from "@/media/vectordbs/lancedb.png";
 import WeaviateLogo from "@/media/vectordbs/weaviate.png";
 import QDrantLogo from "@/media/vectordbs/qdrant.png";
 import MilvusLogo from "@/media/vectordbs/milvus.png";
+import ZillizLogo from "@/media/vectordbs/zilliz.png";
 import System from "@/models/system";
 import paths from "@/utils/paths";
 import PineconeDBOptions from "@/components/VectorDBSelection/PineconeDBOptions";
@@ -14,6 +15,7 @@ import QDrantDBOptions from "@/components/VectorDBSelection/QDrantDBOptions";
 import WeaviateDBOptions from "@/components/VectorDBSelection/WeaviateDBOptions";
 import LanceDBOptions from "@/components/VectorDBSelection/LanceDBOptions";
 import MilvusOptions from "@/components/VectorDBSelection/MilvusDBOptions";
+import ZillizCloudOptions from "@/components/VectorDBSelection/ZillizCloudOptions";
 import showToast from "@/utils/toast";
 import { useNavigate } from "react-router-dom";
 import VectorDBItem from "@/components/VectorDBSelection/VectorDBItem";
@@ -68,6 +70,14 @@ export default function VectorDatabaseConnection({
       options: <PineconeDBOptions settings={settings} />,
       description: "100% cloud-based vector database for enterprise use cases.",
     },
+    {
+      name: "Zilliz Cloud",
+      value: "zilliz",
+      logo: ZillizLogo,
+      options: <ZillizCloudOptions settings={settings} />,
+      description:
+        "Cloud hosted vector database built for enterprise with SOC 2 compliance.",
+    },
     {
       name: "QDrant",
       value: "qdrant",

server/.env.example

@@ -96,6 +96,11 @@ VECTOR_DB="lancedb"
 # MILVUS_USERNAME=
 # MILVUS_PASSWORD=
 
+# Enable all below if you are using vector database: Zilliz Cloud.
+# VECTOR_DB="zilliz"
+# ZILLIZ_ENDPOINT="https://sample.api.gcp-us-west1.zillizcloud.com"
+# ZILLIZ_API_TOKEN=api-token-here
+
 # CLOUD DEPLOYMENT VARIRABLES ONLY
 # AUTH_TOKEN="hunter2" # This is the password to your application if remote hosting.
 # STORAGE_DIR= # absolute filesystem path with no trailing slash

server/models/systemSettings.js

@@ -63,6 +63,12 @@ const SystemSettings = {
             MilvusPassword: !!process.env.MILVUS_PASSWORD,
           }
         : {}),
+      ...(vectorDB === "zilliz"
+        ? {
+            ZillizEndpoint: process.env.ZILLIZ_ENDPOINT,
+            ZillizApiToken: process.env.ZILLIZ_API_TOKEN,
+          }
+        : {}),
       LLMProvider: llmProvider,
       ...(llmProvider === "openai"
         ? {

server/utils/helpers/index.js

@@ -19,6 +19,9 @@ function getVectorDbClass() {
     case "milvus":
       const { Milvus } = require("../vectorDbProviders/milvus");
       return Milvus;
+    case "zilliz":
+      const { Zilliz } = require("../vectorDbProviders/zilliz");
+      return Zilliz;
     default:
       throw new Error("ENV: No VECTOR_DB value found in environment!");
   }

server/utils/helpers/updateENV.js

@@ -199,6 +199,16 @@ const KEY_MAPPING = {
     checks: [isNotEmpty],
   },
 
+  // Zilliz Cloud Options
+  ZillizEndpoint: {
+    envKey: "ZILLIZ_ENDPOINT",
+    checks: [isValidURL],
+  },
+  ZillizApiToken: {
+    envKey: "ZILLIZ_API_TOKEN",
+    checks: [isNotEmpty],
+  },
+
   // Together Ai Options
   TogetherAiApiKey: {
     envKey: "TOGETHER_AI_API_KEY",
@@ -316,6 +326,7 @@ function supportedVectorDB(input = "") {
     "weaviate",
     "qdrant",
     "milvus",
+    "zilliz",
   ];
   return supported.includes(input)
     ? null

server/utils/vectorDbProviders/zilliz/index.js

@@ -0,0 +1,365 @@
+const {
+  DataType,
+  MetricType,
+  IndexType,
+  MilvusClient,
+} = require("@zilliz/milvus2-sdk-node");
+const { RecursiveCharacterTextSplitter } = require("langchain/text_splitter");
+const { v4: uuidv4 } = require("uuid");
+const { storeVectorResult, cachedVectorInformation } = require("../../files");
+const {
+  toChunks,
+  getLLMProvider,
+  getEmbeddingEngineSelection,
+} = require("../../helpers");
+
+// Zilliz is basically a copy of Milvus DB class with a different constructor
+// to connect to the cloud
+const Zilliz = {
+  name: "Zilliz",
+  connect: async function () {
+    if (process.env.VECTOR_DB !== "zilliz")
+      throw new Error("Zilliz::Invalid ENV settings");
+
+    const client = new MilvusClient({
+      address: process.env.ZILLIZ_ENDPOINT,
+      token: process.env.ZILLIZ_API_TOKEN,
+    });
+
+    const { isHealthy } = await client.checkHealth();
+    if (!isHealthy)
+      throw new Error(
+        "Zilliz::Invalid Heartbeat received - is the instance online?"
+      );
+
+    return { client };
+  },
+  heartbeat: async function () {
+    await this.connect();
+    return { heartbeat: Number(new Date()) };
+  },
+  totalVectors: async function () {
+    const { client } = await this.connect();
+    const { collection_names } = await client.listCollections();
+    const total = collection_names.reduce(async (acc, collection_name) => {
+      const statistics = await client.getCollectionStatistics({
+        collection_name,
+      });
+      return Number(acc) + Number(statistics?.data?.row_count ?? 0);
+    }, 0);
+    return total;
+  },
+  namespaceCount: async function (_namespace = null) {
+    const { client } = await this.connect();
+    const statistics = await client.getCollectionStatistics({
+      collection_name: _namespace,
+    });
+    return Number(statistics?.data?.row_count ?? 0);
+  },
+  namespace: async function (client, namespace = null) {
+    if (!namespace) throw new Error("No namespace value provided.");
+    const collection = await client
+      .getCollectionStatistics({ collection_name: namespace })
+      .catch(() => null);
+    return collection;
+  },
+  hasNamespace: async function (namespace = null) {
+    if (!namespace) return false;
+    const { client } = await this.connect();
+    return await this.namespaceExists(client, namespace);
+  },
+  namespaceExists: async function (client, namespace = null) {
+    if (!namespace) throw new Error("No namespace value provided.");
+    const { value } = await client
+      .hasCollection({ collection_name: namespace })
+      .catch((e) => {
+        console.error("Zilliz::namespaceExists", e.message);
+        return { value: false };
+      });
+    return value;
+  },
+  deleteVectorsInNamespace: async function (client, namespace = null) {
+    await client.dropCollection({ collection_name: 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) {
+    const isExists = await this.namespaceExists(client, namespace);
+    if (!isExists) {
+      if (!dimensions)
+        throw new Error(
+          `Zilliz:getOrCreateCollection Unable to infer vector dimension from input. Open an issue on Github for support.`
+        );
+
+      await client.createCollection({
+        collection_name: namespace,
+        fields: [
+          {
+            name: "id",
+            description: "id",
+            data_type: DataType.VarChar,
+            max_length: 255,
+            is_primary_key: true,
+          },
+          {
+            name: "vector",
+            description: "vector",
+            data_type: DataType.FloatVector,
+            dim: dimensions,
+          },
+          {
+            name: "metadata",
+            decription: "metadata",
+            data_type: DataType.JSON,
+          },
+        ],
+      });
+      await client.createIndex({
+        collection_name: namespace,
+        field_name: "vector",
+        index_type: IndexType.AUTOINDEX,
+        metric_type: MetricType.COSINE,
+      });
+      await client.loadCollectionSync({
+        collection_name: namespace,
+      });
+    }
+  },
+  addDocumentToNamespace: async function (
+    namespace,
+    documentData = {},
+    fullFilePath = null
+  ) {
+    const { DocumentVectors } = require("../../../models/vectors");
+    try {
+      let vectorDimension = null;
+      const { pageContent, docId, ...metadata } = documentData;
+      if (!pageContent || pageContent.length == 0) return false;
+
+      console.log("Adding new vectorized document into namespace", namespace);
+      const cacheResult = await cachedVectorInformation(fullFilePath);
+      if (cacheResult.exists) {
+        const { client } = await this.connect();
+        const { chunks } = cacheResult;
+        const documentVectors = [];
+        vectorDimension = chunks[0][0].values.length || null;
+
+        await this.getOrCreateCollection(client, namespace, vectorDimension);
+        for (const chunk of chunks) {
+          // Before sending to Pinecone and saving the records to our db
+          // we need to assign the id of each chunk that is stored in the cached file.
+          const newChunks = chunk.map((chunk) => {
+            const id = uuidv4();
+            documentVectors.push({ docId, vectorId: id });
+            return { id, vector: chunk.values, metadata: chunk.metadata };
+          });
+          const insertResult = await client.insert({
+            collection_name: namespace,
+            data: newChunks,
+          });
+
+          if (insertResult?.status.error_code !== "Success") {
+            throw new Error(
+              `Error embedding into Zilliz! Reason:${insertResult?.status.reason}`
+            );
+          }
+        }
+        await DocumentVectors.bulkInsert(documentVectors);
+        await client.flushSync({ collection_names: [namespace] });
+        return true;
+      }
+
+      const textSplitter = new RecursiveCharacterTextSplitter({
+        chunkSize:
+          getEmbeddingEngineSelection()?.embeddingMaxChunkLength || 1_000,
+        chunkOverlap: 20,
+      });
+      const textChunks = await textSplitter.splitText(pageContent);
+
+      console.log("Chunks created from document:", textChunks.length);
+      const LLMConnector = getLLMProvider();
+      const documentVectors = [];
+      const vectors = [];
+      const vectorValues = await LLMConnector.embedChunks(textChunks);
+
+      if (!!vectorValues && vectorValues.length > 0) {
+        for (const [i, vector] of vectorValues.entries()) {
+          if (!vectorDimension) vectorDimension = vector.length;
+          const vectorRecord = {
+            id: uuidv4(),
+            values: vector,
+            // [DO NOT REMOVE]
+            // LangChain will be unable to find your text if you embed manually and dont include the `text` key.
+            metadata: { ...metadata, text: textChunks[i] },
+          };
+
+          vectors.push(vectorRecord);
+          documentVectors.push({ docId, vectorId: vectorRecord.id });
+        }
+      } else {
+        throw new Error(
+          "Could not embed document chunks! This document will not be recorded."
+        );
+      }
+
+      if (vectors.length > 0) {
+        const chunks = [];
+        const { client } = await this.connect();
+        await this.getOrCreateCollection(client, namespace, vectorDimension);
+
+        console.log("Inserting vectorized chunks into Zilliz.");
+        for (const chunk of toChunks(vectors, 100)) {
+          chunks.push(chunk);
+          const insertResult = await client.insert({
+            collection_name: namespace,
+            data: chunk.map((item) => ({
+              id: item.id,
+              vector: item.values,
+              metadata: chunk.metadata,
+            })),
+          });
+
+          if (insertResult?.status.error_code !== "Success") {
+            throw new Error(
+              `Error embedding into Zilliz! Reason:${insertResult?.status.reason}`
+            );
+          }
+        }
+        await storeVectorResult(chunks, fullFilePath);
+        await client.flushSync({ collection_names: [namespace] });
+      }
+
+      await DocumentVectors.bulkInsert(documentVectors);
+      return true;
+    } catch (e) {
+      console.error(e);
+      console.error("addDocumentToNamespace", e.message);
+      return false;
+    }
+  },
+  deleteDocumentFromNamespace: async function (namespace, docId) {
+    const { DocumentVectors } = require("../../../models/vectors");
+    const { client } = await this.connect();
+    if (!(await this.namespaceExists(client, namespace))) return;
+    const knownDocuments = await DocumentVectors.where({ docId });
+    if (knownDocuments.length === 0) return;
+
+    const vectorIds = knownDocuments.map((doc) => doc.vectorId);
+    const queryIn = vectorIds.map((v) => `'${v}'`).join(",");
+    await client.deleteEntities({
+      collection_name: namespace,
+      expr: `id in [${queryIn}]`,
+    });
+
+    const indexes = knownDocuments.map((doc) => doc.id);
+    await DocumentVectors.deleteIds(indexes);
+
+    // Even after flushing Zilliz can take some time to re-calc the count
+    // so all we can hope to do is flushSync so that the count can be correct
+    // on a later call.
+    await client.flushSync({ collection_names: [namespace] });
+    return true;
+  },
+  performSimilaritySearch: async function ({
+    namespace = null,
+    input = "",
+    LLMConnector = null,
+    similarityThreshold = 0.25,
+  }) {
+    if (!namespace || !input || !LLMConnector)
+      throw new Error("Invalid request to performSimilaritySearch.");
+
+    const { client } = await this.connect();
+    if (!(await this.namespaceExists(client, namespace))) {
+      return {
+        contextTexts: [],
+        sources: [],
+        message: "Invalid query - no documents found for workspace!",
+      };
+    }
+
+    const queryVector = await LLMConnector.embedTextInput(input);
+    const { contextTexts, sourceDocuments } = await this.similarityResponse(
+      client,
+      namespace,
+      queryVector,
+      similarityThreshold
+    );
+
+    const sources = sourceDocuments.map((metadata, i) => {
+      return { ...metadata, text: contextTexts[i] };
+    });
+    return {
+      contextTexts,
+      sources: this.curateSources(sources),
+      message: false,
+    };
+  },
+  similarityResponse: async function (
+    client,
+    namespace,
+    queryVector,
+    similarityThreshold = 0.25
+  ) {
+    const result = {
+      contextTexts: [],
+      sourceDocuments: [],
+      scores: [],
+    };
+    const response = await client.search({
+      collection_name: namespace,
+      vectors: queryVector,
+    });
+    response.results.forEach((match) => {
+      if (match.score < similarityThreshold) return;
+      result.contextTexts.push(match.metadata.text);
+      result.sourceDocuments.push(match);
+      result.scores.push(match.score);
+    });
+    return result;
+  },
+  "namespace-stats": async function (reqBody = {}) {
+    const { namespace = null } = reqBody;
+    if (!namespace) throw new Error("namespace required");
+    const { client } = await this.connect();
+    if (!(await this.namespaceExists(client, namespace)))
+      throw new Error("Namespace by that name does not exist.");
+    const stats = await this.namespace(client, namespace);
+    return stats
+      ? stats
+      : { message: "No stats were able to be fetched from DB for namespace" };
+  },
+  "delete-namespace": async function (reqBody = {}) {
+    const { namespace = null } = reqBody;
+    const { client } = await this.connect();
+    if (!(await this.namespaceExists(client, namespace)))
+      throw new Error("Namespace by that name does not exist.");
+
+    const statistics = await this.namespace(client, namespace);
+    await this.deleteVectorsInNamespace(client, namespace);
+    const vectorCount = Number(statistics?.data?.row_count ?? 0);
+    return {
+      message: `Namespace ${namespace} was deleted along with ${vectorCount} vectors.`,
+    };
+  },
+  curateSources: function (sources = []) {
+    const documents = [];
+    for (const source of sources) {
+      const { metadata = {} } = source;
+      if (Object.keys(metadata).length > 0) {
+        documents.push({
+          ...metadata,
+          ...(source.hasOwnProperty("pageContent")
+            ? { text: source.pageContent }
+            : {}),
+        });
+      }
+    }
+
+    return documents;
+  },
+};
+
+module.exports.Zilliz = Zilliz;