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": [