Mintplex-Labs · scottybo · Dec 12, 2024 · Dec 12, 2024 · Dec 12, 2024 · Dec 12, 2024
diff --git a/server/endpoints/api/workspace/index.js b/server/endpoints/api/workspace/index.js
@@ -4,7 +4,7 @@ const { Telemetry } = require("../../../models/telemetry");
 const { DocumentVectors } = require("../../../models/vectors");
 const { Workspace } = require("../../../models/workspace");
 const { WorkspaceChats } = require("../../../models/workspaceChats");
-const { getVectorDbClass } = require("../../../utils/helpers");
+const { getVectorDbClass, getLLMProvider } = require("../../../utils/helpers");
 const { multiUserMode, reqBody } = require("../../../utils/http");
 const { validApiKey } = require("../../../utils/middleware/validApiKey");
 const { VALID_CHAT_MODE } = require("../../../utils/chats/stream");
@@ -841,6 +841,126 @@ function apiWorkspaceEndpoints(app) {
       }
     }
   );
+
+  app.post(
+    "/v1/workspace/:slug/vector-search",
+    [validApiKey],
+    async (request, response) => {
+      /*
+    #swagger.tags = ['Workspaces']
+    #swagger.description = 'Perform a vector similarity search in a workspace'
+    #swagger.parameters['slug'] = {
+        in: 'path',
+        description: 'Unique slug of workspace to search in',
+        required: true,
+        type: 'string'
+    }
+    #swagger.requestBody = {
+      description: 'Query to perform vector search with and optional parameters',
+      required: true,
+      content: {
+        "application/json": {
+          example: {
+            query: "What is the meaning of life?",
+            topN: 4,
+            scoreThreshold: 0.75
+          }
+        }
+      }
+    }
+    #swagger.responses[200] = {
+      content: {
+        "application/json": {
+          schema: {
+            type: 'object',
+            example: {
+              results: [
+                {
+                  id: "5a6bee0a-306c-47fc-942b-8ab9bf3899c4",
+                  text: "Document chunk content...",
+                  metadata: {
+                    url: "file://document.txt",
+                    title: "document.txt",
+                    author: "no author specified",
+                    description: "no description found",
+                    docSource: "post:123456",
+                    chunkSource: "document.txt",
+                    published: "12/1/2024, 11:39:39 AM",
+                    wordCount: 8,
+                    tokenCount: 9
+                  },
+                  distance: 0.541887640953064,
+                  score: 0.45811235904693604
+                }
+              ]
+            }
+          }
+        }
+      }
+    }
+    */
+      try {
+        const { slug } = request.params;
+        const { query, topN, scoreThreshold } = reqBody(request);
+        const workspace = await Workspace.get({ slug: String(slug) });
+
+        if (!workspace) {
+          response.status(400).json({
+            message: `Workspace ${slug} is not a valid workspace.`,
+          });
+          return;
+        }
+
+        if (!query?.length) {
+          response.status(400).json({
+            message: "Query parameter cannot be empty.",
+          });
+          return;
+        }
+
+        const VectorDb = getVectorDbClass();
+        const hasVectorizedSpace = await VectorDb.hasNamespace(workspace.slug);
+        const embeddingsCount = await VectorDb.namespaceCount(workspace.slug);
+
+        if (!hasVectorizedSpace || embeddingsCount === 0) {
+          response.status(200).json({ results: [] });
+          return;
+        }
+
+        const LLMConnector = getLLMProvider();
+        const results = await VectorDb.performSimilaritySearch({
+          namespace: workspace.slug,
+          input: query,
+          LLMConnector,
+          similarityThreshold: scoreThreshold ?? workspace?.similarityThreshold,
+          topN: topN ?? workspace?.topN ?? 4,
+        });
+
+        response.status(200).json({
+          results: results.sources.map((source) => ({
+            id: source.id,
+            text: source.text,
+            metadata: {
+              url: source.url,
+              title: source.title,
+              author: source.docAuthor,
+              description: source.description,
+              docSource: source.docSource,
+              chunkSource: source.chunkSource,
+              published: source.published,
+              wordCount: source.wordCount,
+              tokenCount: source.token_count_estimate
+            },
+            distance: source._distance,
+            score: source.score
+          })),
+        });
+      } catch (e) {
+        console.error(e.message, e);
+        response.sendStatus(500).end();
+      }
+    }
+  );
 }
 
 module.exports = { apiWorkspaceEndpoints };