From 56fa17caf27949818d34c367605371ceac803c74 Mon Sep 17 00:00:00 2001
From: Sean Hatfield <seanhatfield5@gmail.com>
Date: Thu, 18 Jan 2024 12:34:20 -0800
Subject: [PATCH] create configurable topN per workspace (#616)

* create configurable topN per workspace

* Update TopN UI text
Fix fallbacks for all providers
Add SQLite CHECK to TOPN value

* merge with master
Update zilliz provider for variable TopN

---------

Co-authored-by: timothycarambat <rambat1010@gmail.com>
---
 .../Modals/MangeWorkspace/Settings/index.jsx  | 35 +++++++++++++++++++
 server/models/workspace.js                    |  1 +
 .../20240118201333_init/migration.sql         |  2 ++
 server/prisma/schema.prisma                   |  1 +
 server/utils/chats/index.js                   |  1 +
 server/utils/chats/stream.js                  |  1 +
 .../utils/vectorDbProviders/chroma/index.js   |  9 +++--
 server/utils/vectorDbProviders/lance/index.js |  9 +++--
 .../utils/vectorDbProviders/milvus/index.js   |  8 +++--
 .../utils/vectorDbProviders/pinecone/index.js |  9 +++--
 .../utils/vectorDbProviders/qdrant/index.js   |  9 +++--
 .../utils/vectorDbProviders/weaviate/index.js |  9 +++--
 .../utils/vectorDbProviders/zilliz/index.js   |  8 +++--
 13 files changed, 83 insertions(+), 19 deletions(-)
 create mode 100644 server/prisma/migrations/20240118201333_init/migration.sql

diff --git a/frontend/src/components/Modals/MangeWorkspace/Settings/index.jsx b/frontend/src/components/Modals/MangeWorkspace/Settings/index.jsx
index da0e7b9f..b288dc6c 100644
--- a/frontend/src/components/Modals/MangeWorkspace/Settings/index.jsx
+++ b/frontend/src/components/Modals/MangeWorkspace/Settings/index.jsx
@@ -21,6 +21,9 @@ function castToType(key, value) {
     similarityThreshold: {
       cast: (value) => parseFloat(value),
     },
+    topN: {
+      cast: (value) => Number(value),
+    },
   };
 
   if (!definitions.hasOwnProperty(key)) return value;
@@ -236,6 +239,38 @@ export default function WorkspaceSettings({ active, workspace, settings }) {
                   autoComplete="off"
                   onChange={() => setHasChanges(true)}
                 />
+
+                <div className="mt-4">
+                  <div className="flex flex-col">
+                    <label
+                      htmlFor="name"
+                      className="block text-sm font-medium text-white"
+                    >
+                      Max Context Snippets
+                    </label>
+                    <p className="text-white text-opacity-60 text-xs font-medium py-1.5">
+                      This setting controls the maximum amount of context
+                      snippets the will be sent to the LLM for per chat or
+                      query.
+                      <br />
+                      <i>Recommended: 4</i>
+                    </p>
+                  </div>
+                  <input
+                    name="topN"
+                    type="number"
+                    min={1}
+                    max={12}
+                    step={1}
+                    onWheel={(e) => e.target.blur()}
+                    defaultValue={workspace?.topN ?? 4}
+                    className="bg-zinc-900 text-white text-sm rounded-lg focus:ring-blue-500 focus:border-blue-500 block w-full p-2.5"
+                    placeholder="4"
+                    required={true}
+                    autoComplete="off"
+                    onChange={() => setHasChanges(true)}
+                  />
+                </div>
                 <div className="mt-4">
                   <div className="flex flex-col">
                     <label
diff --git a/server/models/workspace.js b/server/models/workspace.js
index 6de8053e..9169d193 100644
--- a/server/models/workspace.js
+++ b/server/models/workspace.js
@@ -15,6 +15,7 @@ const Workspace = {
     "openAiPrompt",
     "similarityThreshold",
     "chatModel",
+    "topN",
   ],
 
   new: async function (name = null, creatorId = null) {
diff --git a/server/prisma/migrations/20240118201333_init/migration.sql b/server/prisma/migrations/20240118201333_init/migration.sql
new file mode 100644
index 00000000..aaf47f7a
--- /dev/null
+++ b/server/prisma/migrations/20240118201333_init/migration.sql
@@ -0,0 +1,2 @@
+-- AlterTable
+ALTER TABLE "workspaces" ADD COLUMN "topN" INTEGER DEFAULT 4 CHECK ("topN" > 0);
diff --git a/server/prisma/schema.prisma b/server/prisma/schema.prisma
index 2f632a46..15ddbe23 100644
--- a/server/prisma/schema.prisma
+++ b/server/prisma/schema.prisma
@@ -94,6 +94,7 @@ model workspaces {
   openAiPrompt        String?
   similarityThreshold Float?                @default(0.25)
   chatModel           String?
+  topN                Int?                  @default(4)
   workspace_users     workspace_users[]
   documents           workspace_documents[]
 }
diff --git a/server/utils/chats/index.js b/server/utils/chats/index.js
index 764c7795..102189cb 100644
--- a/server/utils/chats/index.js
+++ b/server/utils/chats/index.js
@@ -129,6 +129,7 @@ async function chatWithWorkspace(
     input: message,
     LLMConnector,
     similarityThreshold: workspace?.similarityThreshold,
+    topN: workspace?.topN,
   });
 
   // Failed similarity search.
diff --git a/server/utils/chats/stream.js b/server/utils/chats/stream.js
index cff565ed..29386531 100644
--- a/server/utils/chats/stream.js
+++ b/server/utils/chats/stream.js
@@ -92,6 +92,7 @@ async function streamChatWithWorkspace(
     input: message,
     LLMConnector,
     similarityThreshold: workspace?.similarityThreshold,
+    topN: workspace?.topN,
   });
 
   // Failed similarity search.
diff --git a/server/utils/vectorDbProviders/chroma/index.js b/server/utils/vectorDbProviders/chroma/index.js
index 28af39e6..23f173dd 100644
--- a/server/utils/vectorDbProviders/chroma/index.js
+++ b/server/utils/vectorDbProviders/chroma/index.js
@@ -67,7 +67,8 @@ const Chroma = {
     client,
     namespace,
     queryVector,
-    similarityThreshold = 0.25
+    similarityThreshold = 0.25,
+    topN = 4
   ) {
     const collection = await client.getCollection({ name: namespace });
     const result = {
@@ -78,7 +79,7 @@ const Chroma = {
 
     const response = await collection.query({
       queryEmbeddings: queryVector,
-      nResults: 4,
+      nResults: topN,
     });
     response.ids[0].forEach((_, i) => {
       if (
@@ -271,6 +272,7 @@ const Chroma = {
     input = "",
     LLMConnector = null,
     similarityThreshold = 0.25,
+    topN = 4,
   }) {
     if (!namespace || !input || !LLMConnector)
       throw new Error("Invalid request to performSimilaritySearch.");
@@ -289,7 +291,8 @@ const Chroma = {
       client,
       namespace,
       queryVector,
-      similarityThreshold
+      similarityThreshold,
+      topN
     );
 
     const sources = sourceDocuments.map((metadata, i) => {
diff --git a/server/utils/vectorDbProviders/lance/index.js b/server/utils/vectorDbProviders/lance/index.js
index 8f243cf9..67705c00 100644
--- a/server/utils/vectorDbProviders/lance/index.js
+++ b/server/utils/vectorDbProviders/lance/index.js
@@ -62,7 +62,8 @@ const LanceDb = {
     client,
     namespace,
     queryVector,
-    similarityThreshold = 0.25
+    similarityThreshold = 0.25,
+    topN = 4
   ) {
     const collection = await client.openTable(namespace);
     const result = {
@@ -74,7 +75,7 @@ const LanceDb = {
     const response = await collection
       .search(queryVector)
       .metricType("cosine")
-      .limit(5)
+      .limit(topN)
       .execute();
 
     response.forEach((item) => {
@@ -240,6 +241,7 @@ const LanceDb = {
     input = "",
     LLMConnector = null,
     similarityThreshold = 0.25,
+    topN = 4,
   }) {
     if (!namespace || !input || !LLMConnector)
       throw new Error("Invalid request to performSimilaritySearch.");
@@ -258,7 +260,8 @@ const LanceDb = {
       client,
       namespace,
       queryVector,
-      similarityThreshold
+      similarityThreshold,
+      topN
     );
 
     const sources = sourceDocuments.map((metadata, i) => {
diff --git a/server/utils/vectorDbProviders/milvus/index.js b/server/utils/vectorDbProviders/milvus/index.js
index 79a13241..d1eec9f7 100644
--- a/server/utils/vectorDbProviders/milvus/index.js
+++ b/server/utils/vectorDbProviders/milvus/index.js
@@ -265,6 +265,7 @@ const Milvus = {
     input = "",
     LLMConnector = null,
     similarityThreshold = 0.25,
+    topN = 4,
   }) {
     if (!namespace || !input || !LLMConnector)
       throw new Error("Invalid request to performSimilaritySearch.");
@@ -283,7 +284,8 @@ const Milvus = {
       client,
       namespace,
       queryVector,
-      similarityThreshold
+      similarityThreshold,
+      topN
     );
 
     const sources = sourceDocuments.map((metadata, i) => {
@@ -299,7 +301,8 @@ const Milvus = {
     client,
     namespace,
     queryVector,
-    similarityThreshold = 0.25
+    similarityThreshold = 0.25,
+    topN = 4
   ) {
     const result = {
       contextTexts: [],
@@ -309,6 +312,7 @@ const Milvus = {
     const response = await client.search({
       collection_name: namespace,
       vectors: queryVector,
+      limit: topN,
     });
     response.results.forEach((match) => {
       if (match.score < similarityThreshold) return;
diff --git a/server/utils/vectorDbProviders/pinecone/index.js b/server/utils/vectorDbProviders/pinecone/index.js
index 594a9aaf..260c0a25 100644
--- a/server/utils/vectorDbProviders/pinecone/index.js
+++ b/server/utils/vectorDbProviders/pinecone/index.js
@@ -44,7 +44,8 @@ const Pinecone = {
     index,
     namespace,
     queryVector,
-    similarityThreshold = 0.25
+    similarityThreshold = 0.25,
+    topN = 4
   ) {
     const result = {
       contextTexts: [],
@@ -55,7 +56,7 @@ const Pinecone = {
       queryRequest: {
         namespace,
         vector: queryVector,
-        topK: 4,
+        topK: topN,
         includeMetadata: true,
       },
     });
@@ -237,6 +238,7 @@ const Pinecone = {
     input = "",
     LLMConnector = null,
     similarityThreshold = 0.25,
+    topN = 4,
   }) {
     if (!namespace || !input || !LLMConnector)
       throw new Error("Invalid request to performSimilaritySearch.");
@@ -252,7 +254,8 @@ const Pinecone = {
       pineconeIndex,
       namespace,
       queryVector,
-      similarityThreshold
+      similarityThreshold,
+      topN
     );
 
     const sources = sourceDocuments.map((metadata, i) => {
diff --git a/server/utils/vectorDbProviders/qdrant/index.js b/server/utils/vectorDbProviders/qdrant/index.js
index 70c069e8..e7e00fe6 100644
--- a/server/utils/vectorDbProviders/qdrant/index.js
+++ b/server/utils/vectorDbProviders/qdrant/index.js
@@ -53,7 +53,8 @@ const QDrant = {
     _client,
     namespace,
     queryVector,
-    similarityThreshold = 0.25
+    similarityThreshold = 0.25,
+    topN = 4
   ) {
     const { client } = await this.connect();
     const result = {
@@ -64,7 +65,7 @@ const QDrant = {
 
     const responses = await client.search(namespace, {
       vector: queryVector,
-      limit: 4,
+      limit: topN,
       with_payload: true,
     });
 
@@ -301,6 +302,7 @@ const QDrant = {
     input = "",
     LLMConnector = null,
     similarityThreshold = 0.25,
+    topN = 4,
   }) {
     if (!namespace || !input || !LLMConnector)
       throw new Error("Invalid request to performSimilaritySearch.");
@@ -319,7 +321,8 @@ const QDrant = {
       client,
       namespace,
       queryVector,
-      similarityThreshold
+      similarityThreshold,
+      topN
     );
 
     const sources = sourceDocuments.map((metadata, i) => {
diff --git a/server/utils/vectorDbProviders/weaviate/index.js b/server/utils/vectorDbProviders/weaviate/index.js
index ac89315a..13668303 100644
--- a/server/utils/vectorDbProviders/weaviate/index.js
+++ b/server/utils/vectorDbProviders/weaviate/index.js
@@ -80,7 +80,8 @@ const Weaviate = {
     client,
     namespace,
     queryVector,
-    similarityThreshold = 0.25
+    similarityThreshold = 0.25,
+    topN = 4
   ) {
     const result = {
       contextTexts: [],
@@ -95,7 +96,7 @@ const Weaviate = {
       .withClassName(camelCase(namespace))
       .withFields(`${fields} _additional { id certainty }`)
       .withNearVector({ vector: queryVector })
-      .withLimit(4)
+      .withLimit(topN)
       .do();
 
     const responses = queryResponse?.data?.Get?.[camelCase(namespace)];
@@ -347,6 +348,7 @@ const Weaviate = {
     input = "",
     LLMConnector = null,
     similarityThreshold = 0.25,
+    topN = 4,
   }) {
     if (!namespace || !input || !LLMConnector)
       throw new Error("Invalid request to performSimilaritySearch.");
@@ -365,7 +367,8 @@ const Weaviate = {
       client,
       namespace,
       queryVector,
-      similarityThreshold
+      similarityThreshold,
+      topN
     );
 
     const sources = sourceDocuments.map((metadata, i) => {
diff --git a/server/utils/vectorDbProviders/zilliz/index.js b/server/utils/vectorDbProviders/zilliz/index.js
index 31afab35..159ea809 100644
--- a/server/utils/vectorDbProviders/zilliz/index.js
+++ b/server/utils/vectorDbProviders/zilliz/index.js
@@ -266,6 +266,7 @@ const Zilliz = {
     input = "",
     LLMConnector = null,
     similarityThreshold = 0.25,
+    topN = 4,
   }) {
     if (!namespace || !input || !LLMConnector)
       throw new Error("Invalid request to performSimilaritySearch.");
@@ -284,7 +285,8 @@ const Zilliz = {
       client,
       namespace,
       queryVector,
-      similarityThreshold
+      similarityThreshold,
+      topN
     );
 
     const sources = sourceDocuments.map((metadata, i) => {
@@ -300,7 +302,8 @@ const Zilliz = {
     client,
     namespace,
     queryVector,
-    similarityThreshold = 0.25
+    similarityThreshold = 0.25,
+    topN = 4
   ) {
     const result = {
       contextTexts: [],
@@ -310,6 +313,7 @@ const Zilliz = {
     const response = await client.search({
       collection_name: namespace,
       vectors: queryVector,
+      limit: topN,
     });
     response.results.forEach((match) => {
       if (match.score < similarityThreshold) return;