From 538078747df4b6f566f58c884c157e88677402bd Mon Sep 17 00:00:00 2001
From: Timothy Carambat <rambat1010@gmail.com>
Date: Thu, 12 Dec 2024 10:12:32 -0800
Subject: [PATCH] Add vector search API endpoint (#2815)
* Add vector search API endpoint
* Add missing import
* Modify the data that is returned
* Change similarityThreshold to scoreThreshold
As this is what is actually returned by the search
* Removing logging (oops!)
* chore: regen swagger docs for new endpoint
fix: update function to sanity check values to prevent crashes during search
---------
Co-authored-by: Scott Bowler <scott@dcsdigital.co.uk>
---
server/endpoints/api/workspace/index.js | 133 +++++++++++++++++++++++-
server/swagger/openapi.json | 74 +++++++++++++
2 files changed, 206 insertions(+), 1 deletion(-)
diff --git a/server/endpoints/api/workspace/index.js b/server/endpoints/api/workspace/index.js
index ada7f2fb7..2b0ad577d 100644
--- 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,137 @@ 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)
+ return response.status(400).json({
+ message: `Workspace ${slug} is not a valid workspace.`,
+ });
+
+ if (!query?.length)
+ return response.status(400).json({
+ message: "Query parameter cannot be empty.",
+ });
+
+ const VectorDb = getVectorDbClass();
+ const hasVectorizedSpace = await VectorDb.hasNamespace(workspace.slug);
+ const embeddingsCount = await VectorDb.namespaceCount(workspace.slug);
+
+ if (!hasVectorizedSpace || embeddingsCount === 0)
+ return response
+ .status(200)
+ .json({
+ results: [],
+ message: "No embeddings found for this workspace.",
+ });
+
+ const parseSimilarityThreshold = () => {
+ let input = parseFloat(scoreThreshold);
+ if (isNaN(input) || input < 0 || input > 1)
+ return workspace?.similarityThreshold ?? 0.25;
+ return input;
+ };
+
+ const parseTopN = () => {
+ let input = Number(topN);
+ if (isNaN(input) || input < 1) return workspace?.topN ?? 4;
+ return input;
+ };
+
+ const results = await VectorDb.performSimilaritySearch({
+ namespace: workspace.slug,
+ input: String(query),
+ LLMConnector: getLLMProvider(),
+ similarityThreshold: parseSimilarityThreshold(),
+ topN: parseTopN(),
+ });
+
+ 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 };
diff --git a/server/swagger/openapi.json b/server/swagger/openapi.json
index a849da028..230398ada 100644
--- a/server/swagger/openapi.json
+++ b/server/swagger/openapi.json
@@ -2056,6 +2056,80 @@
}
}
},
+ "/v1/workspace/{slug}/vector-search": {
+ "post": {
+ "tags": [
+ "Workspaces"
+ ],
+ "description": "Perform a vector similarity search in a workspace",
+ "parameters": [
+ {
+ "name": "slug",
+ "in": "path",
+ "required": true,
+ "schema": {
+ "type": "string"
+ },
+ "description": "Unique slug of workspace to search in"
+ }
+ ],
+ "responses": {
+ "200": {
+ "description": "OK",
+ "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
+ }
+ ]
+ }
+ }
+ }
+ }
+ },
+ "400": {
+ "description": "Bad Request"
+ },
+ "403": {
+ "description": "Forbidden"
+ },
+ "500": {
+ "description": "Internal Server Error"
+ }
+ },
+ "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
+ }
+ }
+ }
+ }
+ }
+ },
"/v1/system/env-dump": {
"get": {
"tags": [