feature: Integrate Astra as vectorDBProvider (#648)

* feature: Integrate Astra as vectorDBProvider feature: Integrate Astra as vectorDBProvider * Update .env.example * Add env.example to docker example file Update spellcheck fo Astra Update Astra key for vector selection Update order of AstraDB options Resize Astra logo image to 330x330 Update methods of Astra to take in latest vectorDB params like TopN and more Update Astra interface to support default methods and avoid crash errors from 404 collections Update Astra interface to comply to max chunk insertion limitations Update Astra interface to dynamically set dimensionality from chunk 0 size on creation * reset workspaces --------- Co-authored-by: timothycarambat <[email protected]>
Mintplex-Labs · Jan 26, 2024 · 5614e2e · 5614e2e
1 parent 21653b0
commit 5614e2e
Show file tree

Hide file tree

Showing 16 changed files with 536 additions and 6 deletions.
diff --git a/.vscode/settings.json b/.vscode/settings.json
@@ -1,5 +1,6 @@
 {
   "cSpell.words": [
+    "Astra",
     "Dockerized",
     "Langchain",
     "Milvus",

diff --git a/docker/.env.example b/docker/.env.example
@@ -103,6 +103,11 @@ GID='1000'
 # ZILLIZ_ENDPOINT="https://sample.api.gcp-us-west1.zillizcloud.com"
 # ZILLIZ_API_TOKEN=api-token-here
 
+# Enable all below if you are using vector database: Astra DB.
+# VECTOR_DB="astra"
+# ASTRA_DB_APPLICATION_TOKEN=
+# ASTRA_DB_ENDPOINT=
+
 # CLOUD DEPLOYMENT VARIRABLES ONLY
 # AUTH_TOKEN="hunter2" # This is the password to your application if remote hosting.
 

diff --git a/frontend/src/components/Modals/MangeWorkspace/Settings/index.jsx b/frontend/src/components/Modals/MangeWorkspace/Settings/index.jsx
@@ -44,6 +44,7 @@ export default function WorkspaceSettings({ active, workspace, settings }) {
   const formEl = useRef(null);
   const [saving, setSaving] = useState(false);
   const [hasChanges, setHasChanges] = useState(false);
+  const [deleting, setDeleting] = useState(false);
   const defaults = recommendedSettings(settings?.LLMProvider);
 
   const handleUpdate = async (e) => {
@@ -72,7 +73,15 @@ export default function WorkspaceSettings({ active, workspace, settings }) {
       )
     )
       return false;
-    await Workspace.delete(workspace.slug);
+
+    setDeleting(true);
+    const success = await Workspace.delete(workspace.slug);
+    if (!success) {
+      showToast("Workspace could not be deleted!", "error", { clear: true });
+      setDeleting(false);
+      return;
+    }
+
     workspace.slug === slug
       ? (window.location = paths.home())
       : window.location.reload();
@@ -310,7 +319,11 @@ export default function WorkspaceSettings({ active, workspace, settings }) {
         </div>
       </div>
       <div className="flex items-center justify-between p-2 md:p-6 space-x-2 border-t rounded-b border-gray-600">
-        <DeleteWorkspace workspace={workspace} onClick={deleteWorkspace} />
+        <DeleteWorkspace
+          deleting={deleting}
+          workspace={workspace}
+          onClick={deleteWorkspace}
+        />
         {hasChanges && (
           <button
             type="submit"
@@ -324,7 +337,7 @@ export default function WorkspaceSettings({ active, workspace, settings }) {
   );
 }
 
-function DeleteWorkspace({ workspace, onClick }) {
+function DeleteWorkspace({ deleting, workspace, onClick }) {
   const [canDelete, setCanDelete] = useState(false);
   useEffect(() => {
     async function fetchKeys() {
@@ -337,11 +350,12 @@ function DeleteWorkspace({ workspace, onClick }) {
   if (!canDelete) return null;
   return (
     <button
+      disabled={deleting}
       onClick={onClick}
       type="button"
-      className="transition-all duration-300 border border-transparent rounded-lg whitespace-nowrap text-sm px-5 py-2.5 focus:z-10 bg-transparent text-white hover:text-white hover:bg-red-600"
+      className="transition-all duration-300 border border-transparent rounded-lg whitespace-nowrap text-sm px-5 py-2.5 focus:z-10 bg-transparent text-white hover:text-white hover:bg-red-600 disabled:bg-red-600 disabled:text-red-200 disabled:animate-pulse"
     >
-      Delete Workspace
+      {deleting ? "Deleting Workspace..." : "Delete Workspace"}
     </button>
   );
 }

diff --git a/frontend/src/components/VectorDBSelection/AstraDBOptions/index.jsx b/frontend/src/components/VectorDBSelection/AstraDBOptions/index.jsx
@@ -0,0 +1,41 @@
+export default function AstraDBOptions({ 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">
+            Astra DB Endpoint
+          </label>
+          <input
+            type="url"
+            name="AstraDBEndpoint"
+            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="Astra DB API endpoint"
+            defaultValue={settings?.AstraDBEndpoint}
+            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">
+            Astra DB Application Token
+          </label>
+          <input
+            type="password"
+            name="AstraDBApplicationToken"
+            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="AstraCS:..."
+            defaultValue={
+              settings?.AstraDBApplicationToken ? "*".repeat(20) : ""
+            }
+            required={true}
+            autoComplete="off"
+            spellCheck={false}
+          />
+        </div>
+      </div>
+    </div>
+  );
+}
diff --git a/frontend/src/media/vectordbs/astraDB.png b/frontend/src/media/vectordbs/astraDB.png
diff --git a/frontend/src/pages/GeneralSettings/VectorDatabase/index.jsx b/frontend/src/pages/GeneralSettings/VectorDatabase/index.jsx
@@ -10,6 +10,7 @@ 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 AstraDBLogo from "@/media/vectordbs/astraDB.png";
 import PreLoader from "@/components/Preloader";
 import ChangeWarningModal from "@/components/ChangeWarning";
 import { MagnifyingGlass } from "@phosphor-icons/react";
@@ -23,6 +24,7 @@ import MilvusDBOptions from "@/components/VectorDBSelection/MilvusDBOptions";
 import ZillizCloudOptions from "@/components/VectorDBSelection/ZillizCloudOptions";
 import { useModal } from "@/hooks/useModal";
 import ModalWrapper from "@/components/ModalWrapper";
+import AstraDBOptions from "@/components/VectorDBSelection/AstraDBOptions";
 
 export default function GeneralVectorDatabase() {
   const [saving, setSaving] = useState(false);
@@ -100,6 +102,13 @@ export default function GeneralVectorDatabase() {
       options: <MilvusDBOptions settings={settings} />,
       description: "Open-source, highly scalable, and blazing fast.",
     },
+    {
+      name: "AstraDB",
+      value: "astra",
+      logo: AstraDBLogo,
+      options: <AstraDBOptions settings={settings} />,
+      description: "Vector Search for Real-world GenAI.",
+    },
   ];
 
   const updateVectorChoice = (selection) => {

diff --git a/frontend/src/pages/OnboardingFlow/Steps/DataHandling/index.jsx b/frontend/src/pages/OnboardingFlow/Steps/DataHandling/index.jsx
@@ -11,6 +11,7 @@ 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 AstraDBLogo from "@/media/vectordbs/astraDB.png";
 import ChromaLogo from "@/media/vectordbs/chroma.png";
 import PineconeLogo from "@/media/vectordbs/pinecone.png";
 import LanceDbLogo from "@/media/vectordbs/lancedb.png";
@@ -147,6 +148,13 @@ const VECTOR_DB_PRIVACY = {
     ],
     logo: ZillizLogo,
   },
+  astra: {
+    name: "AstraDB",
+    description: [
+      "Your vectors and document text are stored on your cloud AstraDB database.",
+    ],
+    logo: AstraDBLogo,
+  },
   lancedb: {
     name: "LanceDB",
     description: [

diff --git a/frontend/src/pages/OnboardingFlow/Steps/VectorDatabaseConnection/index.jsx b/frontend/src/pages/OnboardingFlow/Steps/VectorDatabaseConnection/index.jsx
@@ -7,6 +7,7 @@ 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 AstraDBLogo from "@/media/vectordbs/astraDB.png";
 import System from "@/models/system";
 import paths from "@/utils/paths";
 import PineconeDBOptions from "@/components/VectorDBSelection/PineconeDBOptions";
@@ -16,6 +17,7 @@ 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 AstraDBOptions from "@/components/VectorDBSelection/AstraDBOptions";
 import showToast from "@/utils/toast";
 import { useNavigate } from "react-router-dom";
 import VectorDBItem from "@/components/VectorDBSelection/VectorDBItem";
@@ -100,6 +102,13 @@ export default function VectorDatabaseConnection({
       options: <MilvusOptions settings={settings} />,
       description: "Open-source, highly scalable, and blazing fast.",
     },
+    {
+      name: "AstraDB",
+      value: "astra",
+      logo: AstraDBLogo,
+      options: <AstraDBOptions settings={settings} />,
+      description: "Vector Search for Real-world GenAI.",
+    },
   ];
 
   function handleForward() {

diff --git a/server/.env.example b/server/.env.example
@@ -76,6 +76,11 @@ JWT_SECRET="my-random-string-for-seeding" # Please generate random string at lea
 # PINECONE_API_KEY=
 # PINECONE_INDEX=
 
+# Enable all below if you are using vector database: Astra DB.
+# VECTOR_DB="astra"
+# ASTRA_DB_APPLICATION_TOKEN=
+# ASTRA_DB_ENDPOINT=
+
 # Enable all below if you are using vector database: LanceDB.
 VECTOR_DB="lancedb"
 

diff --git a/server/models/systemSettings.js b/server/models/systemSettings.js
@@ -68,6 +68,12 @@ const SystemSettings = {
             ZillizApiToken: process.env.ZILLIZ_API_TOKEN,
           }
         : {}),
+      ...(vectorDB === "astra"
+        ? {
+            AstraDBApplicationToken: process?.env?.ASTRA_DB_APPLICATION_TOKEN,
+            AstraDBEndpoint: process?.env?.ASTRA_DB_ENDPOINT,
+          }
+        : {}),
       LLMProvider: llmProvider,
       ...(llmProvider === "openai"
         ? {

diff --git a/server/package.json b/server/package.json
@@ -22,6 +22,7 @@
   "dependencies": {
     "@anthropic-ai/sdk": "^0.8.1",
     "@azure/openai": "1.0.0-beta.10",
+    "@datastax/astra-db-ts": "^0.1.3",
     "@google/generative-ai": "^0.1.3",
     "@googleapis/youtube": "^9.0.0",
     "@pinecone-database/pinecone": "^2.0.1",

diff --git a/server/utils/helpers/index.js b/server/utils/helpers/index.js
@@ -22,6 +22,9 @@ function getVectorDbClass() {
     case "zilliz":
       const { Zilliz } = require("../vectorDbProviders/zilliz");
       return Zilliz;
+    case "astra":
+      const { AstraDB } = require("../vectorDbProviders/astra");
+      return AstraDB;
     default:
       throw new Error("ENV: No VECTOR_DB value found in environment!");
   }

diff --git a/server/utils/helpers/updateENV.js b/server/utils/helpers/updateENV.js
@@ -204,6 +204,17 @@ const KEY_MAPPING = {
     checks: [isNotEmpty],
   },
 
+  // Astra DB Options
+
+  AstraDBApplicationToken: {
+    envKey: "ASTRA_DB_APPLICATION_TOKEN",
+    checks: [isNotEmpty],
+  },
+  AstraDBEndpoint: {
+    envKey: "ASTRA_DB_ENDPOINT",
+    checks: [isNotEmpty],
+  },
+
   // Together Ai Options
   TogetherAiApiKey: {
     envKey: "TOGETHER_AI_API_KEY",
@@ -322,6 +333,7 @@ function supportedVectorDB(input = "") {
     "qdrant",
     "milvus",
     "zilliz",
+    "astra",
   ];
   return supported.includes(input)
     ? null

diff --git a/server/utils/vectorDbProviders/astra/ASTRA_SETUP.md b/server/utils/vectorDbProviders/astra/ASTRA_SETUP.md
@@ -0,0 +1,22 @@
+# How to setup Astra Vector Database for AnythingLLM
+
+[Official Astra DB Docs](https://docs.datastax.com/en/astra/astra-db-vector/get-started/quickstart.html) for reference.
+
+### How to get started
+
+**Requirements**
+
+- Astra Vector Database with active status.
+
+**Instructions**
+
+- [Create an Astra account or sign in to an existing Astra account](astra.datastax.com)
+- Create an Astra Serverless(Vector) Database.
+- Make sure DB is in active state.
+- Get `API ENDPOINT`and `Application Token` from Overview screen
+
+```
+VECTOR_DB="astra"
+ASTRA_DB_ENDPOINT=Astra DB API endpoint
+ASTRA_DB_APPLICATION_TOKEN=AstraCS:..
+```