From 737a0bd3ae94966df073675893946c6abe1d691e Mon Sep 17 00:00:00 2001
From: Cohee <18619528+Cohee1207@users.noreply.github.com>
Date: Fri, 23 Feb 2024 22:37:00 +0200
Subject: [PATCH] Fix purge extras and mistral vectors

---
 public/scripts/extensions/vectors/index.js | 17 +++++++++++++----
 src/endpoints/vectors.js                   |  6 ++++--
 2 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/public/scripts/extensions/vectors/index.js b/public/scripts/extensions/vectors/index.js
index 5821a7a30..e318d6518 100644
--- a/public/scripts/extensions/vectors/index.js
+++ b/public/scripts/extensions/vectors/index.js
@@ -530,10 +530,15 @@ async function queryCollection(collectionId, searchText, topK) {
     return results;
 }
 
+/**
+ * Purges the vector index for a collection.
+ * @param {string} collectionId Collection ID to purge
+ * @returns <Promise<boolean>> True if deleted, false if not
+ */
 async function purgeVectorIndex(collectionId) {
     try {
         if (!settings.enabled_chats) {
-            return;
+            return true;
         }
 
         const response = await fetch('/api/vector/purge', {
@@ -549,9 +554,10 @@ async function purgeVectorIndex(collectionId) {
         }
 
         console.log(`Vectors: Purged vector index for collection ${collectionId}`);
-
+        return true;
     } catch (error) {
         console.error('Vectors: Failed to purge', error);
+        return false;
     }
 }
 
@@ -566,8 +572,11 @@ async function onPurgeClick() {
         toastr.info('No chat selected', 'Purge aborted');
         return;
     }
-    await purgeVectorIndex(chatId);
-    toastr.success('Vector index purged', 'Purge successful');
+    if (await purgeVectorIndex(chatId)) {
+        toastr.success('Vector index purged', 'Purge successful');
+    } else {
+        toastr.error('Failed to purge vector index', 'Purge failed');
+    }
 }
 
 async function onViewStatsClick() {
diff --git a/src/endpoints/vectors.js b/src/endpoints/vectors.js
index e50cc0bcf..b2ff9e2d0 100644
--- a/src/endpoints/vectors.js
+++ b/src/endpoints/vectors.js
@@ -4,6 +4,9 @@ const express = require('express');
 const sanitize = require('sanitize-filename');
 const { jsonParser } = require('../express-common');
 
+// Don't forget to add new sources to the SOURCES array
+const SOURCES = ['transformers', 'mistral', 'openai', 'extras', 'palm'];
+
 /**
  * Gets the vector for the given text from the given source.
  * @param {string} source - The source of the vector
@@ -261,8 +264,7 @@ router.post('/purge', jsonParser, async (req, res) => {
 
         const collectionId = String(req.body.collectionId);
 
-        const sources = ['transformers', 'openai', 'palm'];
-        for (const source of sources) {
+        for (const source of SOURCES) {
             const index = await getIndex(collectionId, source, false);
 
             const exists = await index.isIndexCreated();