Merge pull request #64 from Mintplex-Labs/reset-connection

enable vector db wiping

backend/endpoints/v1/tools/index.js

@@ -8,6 +8,9 @@ const {
 const {
   organizationMigrationJob,
 } = require("../../../utils/jobs/organizationMigrationJob");
+const {
+  organizationResetJob,
+} = require("../../../utils/jobs/organizationResetJob");
 
 process.env.NODE_ENV === "development"
   ? require("dotenv").config({ path: `.env.${process.env.NODE_ENV}` })
@@ -88,6 +91,54 @@ function toolEndpoints(app) {
       }
     }
   );
+
+  app.post(
+    "/v1/tools/org/:orgSlug/reset",
+    [validSessionForUser],
+    async function (request, response) {
+      try {
+        const { orgSlug } = request.params;
+        const user = await userFromSession(request);
+        if (!user || user.role !== "admin") {
+          response.sendStatus(403).end();
+          return;
+        }
+
+        const organization = await Organization.getWithOwner(user.id, {
+          slug: orgSlug,
+        });
+        if (!organization) {
+          response
+            .status(200)
+            .json({ success: false, message: "No org found." });
+          return;
+        }
+
+        const pendingJobs = await Queue.where({
+          status: "pending",
+          organization_id: organization.id,
+        });
+        if (pendingJobs.length > 0) {
+          response
+            .status(200)
+            .json({
+              success: false,
+              message:
+                "There are pending jobs for this organization - you cannot reset it at this time.",
+            });
+          return;
+        }
+
+        await organizationResetJob(organization, user);
+        response
+          .status(200)
+          .json({ success: true, message: "Migration job queued." });
+      } catch (e) {
+        console.log(e.message, e);
+        response.sendStatus(500).end();
+      }
+    }
+  );
 }
 
 module.exports = { toolEndpoints };

backend/utils/jobs/organizationResetJob/index.js

@@ -0,0 +1,25 @@
+const { Queue } = require("../../../models/queue");
+
+async function organizationResetJob(organization, user) {
+  const taskName = `organization/reset`;
+  const jobData = { organization };
+  const { job, error } = await Queue.create(
+    taskName,
+    jobData,
+    user.id,
+    organization.id
+  );
+  if (!!error) return { job, error };
+  await Queue.sendJob({
+    name: taskName,
+    data: {
+      jobId: job.id,
+      ...jobData,
+    },
+  });
+  return { job, error: null };
+}
+
+module.exports = {
+  organizationResetJob,
+};

frontend/src/App.tsx

@@ -27,6 +27,7 @@ const SystemSettingsView = lazy(() => import('./pages/SystemSettings'));
 const MigrateConnectionView = lazy(
   () => import('./pages/Tools/MigrateConnection')
 );
+const ResetConnectionView = lazy(() => import('./pages/Tools/ResetConnection'));
 
 function App() {
   return (
@@ -83,6 +84,10 @@ function App() {
             path="/dashboard/:slug/tools/db-migration"
             element={<PrivateRoute Component={MigrateConnectionView} />}
           />
+          <Route
+            path="/dashboard/:slug/tools/db-reset"
+            element={<PrivateRoute Component={ResetConnectionView} />}
+          />
 
           <Route path="/auth/sign-up" element={<SignUp />} />
           <Route path="/auth/sign-in" element={<SignIn />} />

frontend/src/models/tools.ts

@@ -19,6 +19,21 @@ const Tools = {
         return { success: false, message: e.message };
       });
   },
+  resetOrg: async (
+    slug: string
+  ): Promise<{ success: boolean; message: null | string }> => {
+    return fetch(`${API_BASE}/v1/tools/org/${slug}/reset`, {
+      method: 'POST',
+      cache: 'no-cache',
+      headers: baseHeaders(),
+    })
+      .then((res) => res.json())
+      .then((res) => res)
+      .catch((e) => {
+        console.error(e);
+        return { success: false, message: e.message };
+      });
+  },
 };
 
 export default Tools;

frontend/src/pages/Tools/ResetConnection/index.tsx

@@ -0,0 +1,170 @@
+import { FullScreenLoader } from '../../../components/Preloader';
+import useUser from '../../../hooks/useUser';
+import { useState, useEffect } from 'react';
+import DefaultLayout from '../../../layout/DefaultLayout';
+import User from '../../../models/user';
+import paths from '../../../utils/paths';
+import AppLayout from '../../../layout/AppLayout';
+import { useParams } from 'react-router-dom';
+import Organization from '../../../models/organization';
+import { Info, Loader } from 'react-feather';
+import { nFormatter } from '../../../utils/numbers';
+import pluralize from 'pluralize';
+import Tools from '../../../models/tools';
+import showToast from '../../../utils/toast';
+
+export default function ResetConnectionView() {
+  const { user } = useUser();
+  const { slug } = useParams();
+  const [loading, setLoading] = useState<boolean>(true);
+  const [organizations, setOrganizations] = useState<object[]>([]);
+  const [organization, setOrganization] = useState<object | null>(null);
+  const [totalVectorCount, setTotalVectorCount] = useState<number>(0);
+
+  useEffect(() => {
+    async function userOrgs() {
+      if (!slug) return false;
+
+      const orgs = await User.organizations();
+      if (orgs.length === 0) {
+        window.location.replace(paths.onboarding.orgName());
+        return false;
+      }
+      for (const org of orgs) {
+        const connector = await Organization.connector(org.slug);
+        org.connector = connector || {};
+      }
+
+      const focusedOrg =
+        orgs?.find((org: any) => org.slug === slug) || orgs?.[0];
+      const { remoteCount } = (
+        await Organization.stats(focusedOrg.slug, 'vectorCounts')
+      ).value;
+
+      setOrganizations(orgs);
+      setOrganization(focusedOrg);
+      setTotalVectorCount(remoteCount);
+      setLoading(false);
+    }
+    userOrgs();
+  }, [user.uid, window.location.pathname]);
+
+  if (loading || organizations.length === 0) {
+    return (
+      <DefaultLayout>
+        <FullScreenLoader />
+      </DefaultLayout>
+    );
+  }
+
+  if (!organization) return null;
+  return (
+    <AppLayout
+      headerEntity={organization}
+      headerProp="uuid"
+      organizations={organizations}
+      organization={organization}
+      workspaces={[]}
+    >
+      <div className="flex-1 rounded-sm bg-white px-6 pb-6">
+        <div className="flex items-start justify-between">
+          <div className="mb-6">
+            <h4 className="text-3xl font-semibold text-black">
+              Reset your vector database
+            </h4>
+            <p className="mt-2 text-gray-600">
+              This tool will wipe all existing data in this vector database
+              connection.
+              <br />
+              You should only use this tool if you are absolutely sure you want
+              to lose this data.
+            </p>
+          </div>
+        </div>
+        <div className="flex w-1/2 flex-col gap-y-2">
+          <SubmitResetJob
+            organization={organization}
+            totalVectorCount={totalVectorCount}
+          />
+        </div>
+      </div>
+    </AppLayout>
+  );
+}
+
+function SubmitResetJob({
+  organization,
+  totalVectorCount,
+}: {
+  organization: any;
+  totalVectorCount: number;
+}) {
+  const [status, setStatus] = useState<
+    'idle' | 'loading' | 'success' | 'failed'
+  >('idle');
+
+  async function submitReset() {
+    setStatus('loading');
+    const { success, message } = await Tools.resetOrg(organization.slug);
+
+    if (success) {
+      setStatus('success');
+      showToast('Reset job has started as background job.', 'success');
+      return;
+    }
+
+    setStatus('failed');
+    showToast(message as string, 'error');
+  }
+
+  if (status === 'success') {
+    return (
+      <div className="mt-2 flex flex-col gap-y-1">
+        <a
+          href={paths.jobs(organization)}
+          className="full flex items-center justify-center gap-x-1 rounded-lg bg-green-100 px-4 py-2 font-semibold text-green-800 hover:bg-green-200"
+        >
+          View job in background jobs &rarr;
+        </a>
+      </div>
+    );
+  }
+
+  if (status === 'loading') {
+    return (
+      <div className="mt-2 flex flex-col gap-y-1">
+        <button
+          disabled={true}
+          className="full flex items-center justify-center gap-x-1 rounded-lg bg-blue-600 px-4 py-2 font-semibold text-white hover:bg-blue-700"
+        >
+          <Loader size={14} className="animate-spin" />
+          Deleting {nFormatter(totalVectorCount)}{' '}
+          {pluralize('vector', totalVectorCount)} from{' '}
+          {organization.connector.type} vector database...
+        </button>
+        <p className="mx-auto w-4/5 text-center text-xs text-gray-500">
+          This task will run in the background. You can view its progress in
+          Background Jobs.
+        </p>
+      </div>
+    );
+  }
+
+  return (
+    <div className="mt-2 flex flex-col gap-y-1">
+      <button
+        onClick={submitReset}
+        className="full rounded-lg bg-blue-600 px-4 py-2 font-semibold text-white hover:bg-blue-700"
+      >
+        Delete {nFormatter(totalVectorCount)}{' '}
+        {pluralize('vector', totalVectorCount)} from{' '}
+        {organization.connector.type} vector database
+      </button>
+      <p className="mx-auto w-4/5 text-center text-xs text-gray-500">
+        Depending on the number of vectors this process can take a few moments.
+        Please do not edit or update this vector database while the process is
+        running.
+      </p>
+    </div>
+  );
+}

frontend/src/pages/Tools/ToolsList/index.tsx

@@ -1,13 +1,9 @@
 import { useState } from 'react';
-import { ChevronDown, Info } from 'react-feather';
-import useUser from '../../../hooks/useUser';
+import { Info } from 'react-feather';
 import { APP_NAME } from '../../../utils/constants';
-import moment from 'moment';
 import paths from '../../../utils/paths';
 
 export default function ToolsList({ organization }: { organization: any }) {
-  const { user } = useUser();
-
   return (
     <div className="col-span-12 flex-1 rounded-sm bg-white pb-6 xl:col-span-4">
       <div className="flex items-start justify-between">
@@ -23,6 +19,12 @@ export default function ToolsList({ organization }: { organization: any }) {
       </div>
 
       <div className="px-6">
+        <ToolItem
+          title="Reset vector database"
+          description="Wipe out all existing data in your vector database in one click."
+          available={true}
+          linkTo={paths.tools.resetTool(organization)}
+        />
         <ToolItem
           title="Migrate vector database to another provider"
           description="Take all of your vectors to another connected vector database provider."

frontend/src/utils/paths.ts

@@ -40,6 +40,9 @@ const paths = {
     migrationTool: function ({ slug }: { slug: string }) {
       return `/dashboard/${slug}/tools/db-migration`;
     },
+    resetTool: function ({ slug }: { slug: string }) {
+      return `/dashboard/${slug}/tools/db-reset`;
+    },
   },
   dashboard: function () {
     return '/dashboard';

workers/functions/resetOrganization/index.js

@@ -0,0 +1,149 @@
+const {
+  OrganizationConnection,
+} = require('../../../backend/models/organizationConnection');
+const {
+  OrganizationWorkspace,
+} = require('../../../backend/models/organizationWorkspace');
+const { Queue } = require('../../../backend/models/queue');
+const {
+  selectConnector,
+} = require('../../../backend/utils/vectordatabases/providers');
+
+const { InngestClient } = require('../../utils/inngest');
+
+const resetOrganization = InngestClient.createFunction(
+  { name: 'Reset vector database connected for organization' },
+  { event: 'organization/reset' },
+  async ({ event, step: _step, logger }) => {
+    var result = {};
+    const { organization, jobId } = event.data;
+    const connector = await OrganizationConnection.get({
+      organization_id: Number(organization.id),
+    });
+    const vectorDb = selectConnector(connector);
+
+    if (vectorDb.name === 'pinecone') {
+      try {
+        const namespaces = await vectorDb.namespaces();
+        const { pineconeIndex } = await vectorDb.connect();
+        for (const collection of namespaces) {
+          await pineconeIndex.delete1({
+            namespace: collection.name,
+            deleteAll: true,
+          });
+        }
+
+        await removeWorkspaces(organization);
+        result = {
+          message: `All namespaces deleted from Pinecone.`,
+        };
+        await Queue.updateJob(jobId, Queue.status.complete, result);
+        return { result };
+      } catch (e) {
+        const result = {
+          message: `Job failed with error`,
+          error: e.message,
+          details: e,
+        };
+        await Queue.updateJob(jobId, Queue.status.failed, result);
+        return { result };
+      }
+    }
+
+    if (vectorDb.name === 'chroma') {
+      try {
+        // cannot use .reset as it may not be enabled.
+        const namespaces = await vectorDb.namespaces();
+        const { client } = await vectorDb.connect();
+        for (const collectionInfo of namespaces) {
+          await client.deleteCollection({ name: collectionInfo.name });
+        }
+
+        await removeWorkspaces(organization);
+        result = {
+          message: `All namespaces deleted from Chroma.`,
+        };
+        await Queue.updateJob(jobId, Queue.status.complete, result);
+        return { result };
+      } catch (e) {
+        const result = {
+          message: `Job failed with error`,
+          error: e.message,
+          details: e,
+        };
+        await Queue.updateJob(jobId, Queue.status.failed, result);
+        return { result };
+      }
+    }
+
+    if (vectorDb.name === 'qdrant') {
+      try {
+        const namespaces = await vectorDb.namespaces();
+        const { client } = await vectorDb.connect();
+        for (const collectionInfo of namespaces) {
+          await client.deleteCollection(collectionInfo.name);
+        }
+
+        await removeWorkspaces(organization);
+        result = {
+          message: `All namespaces deleted from Chroma.`,
+        };
+        await Queue.updateJob(jobId, Queue.status.complete, result);
+        return { result };
+      } catch (e) {
+        const result = {
+          message: `Job failed with error`,
+          error: e.message,
+          details: e,
+        };
+        await Queue.updateJob(jobId, Queue.status.failed, result);
+        return { result };
+      }
+    }
+
+    if (vectorDb.name === 'weaviate') {
+      try {
+        const { client } = await weaviateClient.connect();
+        const collections = await vectorDb.collections();
+        for (const collectionInfo of collections) {
+          await client.schema
+            .classDeleter()
+            .withClassName(collectionInfo.name)
+            .do();
+        }
+
+        await removeWorkspaces(organization);
+        result = {
+          message: `All namespaces deleted from Chroma.`,
+        };
+        await Queue.updateJob(jobId, Queue.status.complete, result);
+        return { result };
+      } catch (e) {
+        const result = {
+          message: `Job failed with error`,
+          error: e.message,
+          details: e,
+        };
+        await Queue.updateJob(jobId, Queue.status.failed, result);
+        return { result };
+      }
+    }
+
+    result = { message: `Nothing to do.` };
+    await Queue.updateJob(jobId, Queue.status.complete, result);
+    return { result };
+  }
+);
+
+async function removeWorkspaces(organization) {
+  const workspaces = await OrganizationWorkspace.where({
+    organization_id: Number(organization.id),
+  });
+  OrganizationWorkspace.delete({
+    id: { in: workspaces.map((ws) => ws.id) },
+  });
+}
+
+module.exports = {
+  resetOrganization,
+};

workers/index.js

@@ -34,6 +34,7 @@ const { cloneWeaviateWorkspace } = require("./functions/cloneWeaviateWorkspace")
 const { cloneWeaviateDocument } = require("./functions/cloneWeaviateDocument");
 const { addWeaviateDocuments } = require("./functions/addWeaviateDocuments");
 const { migrateOrganization } = require("./functions/migrateOrganization");
+const { resetOrganization } = require("./functions/resetOrganization");
 const app = express();
 
 app.use(cors({ origin: true }));
@@ -98,6 +99,7 @@ app.use(
     newWorkspaceCreated,
     workspaceDeleted,
     migrateOrganization,
+    resetOrganization,
   ], { landingPage: true })
 );