From fdb6b6f365149844d7fb5ade4dd9ea25d71c9ea4 Mon Sep 17 00:00:00 2001 From: Scott Bowler Date: Thu, 12 Dec 2024 07:38:23 +0000 Subject: [PATCH 1/6] Add vector search API endpoint --- server/endpoints/api/workspace/index.js | 103 ++++++++++++++++++++++++ 1 file changed, 103 insertions(+) diff --git a/server/endpoints/api/workspace/index.js b/server/endpoints/api/workspace/index.js index ada7f2fb730..fbc7ae7913c 100644 --- a/server/endpoints/api/workspace/index.js +++ b/server/endpoints/api/workspace/index.js @@ -841,6 +841,109 @@ 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, + similarityThreshold: 0.75 + } + } + } + } + #swagger.responses[200] = { + content: { + "application/json": { + schema: { + type: 'object', + example: { + results: [ + { + text: "Document chunk content...", + metadata: { + title: "document.pdf", + source: "documents/file.pdf" + }, + similarity: 0.89 + } + ] + } + } + } + } + } + */ + try { + const { slug } = request.params; + const { query, topN, similarityThreshold } = 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: + similarityThreshold ?? workspace?.similarityThreshold, + topN: topN ?? workspace?.topN ?? 4, + }); + + response.status(200).json({ + results: results.sources.map((source) => ({ + text: source.text, + metadata: { + title: source.title, + source: source.source, + }, + similarity: source.similarity, + })), + }); + } catch (e) { + console.error(e.message, e); + response.sendStatus(500).end(); + } + } + ); } module.exports = { apiWorkspaceEndpoints }; From 2d80d5f01a3bf71d5f0dafe48ae3a93c33d38e92 Mon Sep 17 00:00:00 2001 From: Scott Bowler Date: Thu, 12 Dec 2024 15:03:05 +0000 Subject: [PATCH 2/6] Add missing import --- server/endpoints/api/workspace/index.js | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/server/endpoints/api/workspace/index.js b/server/endpoints/api/workspace/index.js index fbc7ae7913c..5ef039ce86c 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"); From 3c0353591ef9ab905648cf8cef2a99b29271057d Mon Sep 17 00:00:00 2001 From: Scott Bowler Date: Thu, 12 Dec 2024 15:06:57 +0000 Subject: [PATCH 3/6] Modify the data that is returned --- server/endpoints/api/workspace/index.js | 30 ++++++++++++++++++++----- 1 file changed, 25 insertions(+), 5 deletions(-) diff --git a/server/endpoints/api/workspace/index.js b/server/endpoints/api/workspace/index.js index 5ef039ce86c..9415a47ab64 100644 --- a/server/endpoints/api/workspace/index.js +++ b/server/endpoints/api/workspace/index.js @@ -876,12 +876,21 @@ function apiWorkspaceEndpoints(app) { example: { results: [ { + id: "5a6bee0a-306c-47fc-942b-8ab9bf3899c4", text: "Document chunk content...", metadata: { - title: "document.pdf", - source: "documents/file.pdf" + 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 }, - similarity: 0.89 + similarity: 0.541887640953064, + score: 0.45811235904693604 } ] } @@ -928,14 +937,25 @@ function apiWorkspaceEndpoints(app) { topN: topN ?? workspace?.topN ?? 4, }); + console.log(results); + response.status(200).json({ results: results.sources.map((source) => ({ + id: source.id, text: source.text, metadata: { + url: source.url, title: source.title, - source: source.source, + author: source.docAuthor, + description: source.description, + docSource: source.docSource, + chunkSource: source.chunkSource, + published: source.published, + wordCount: source.wordCount, + tokenCount: source.token_count_estimate }, - similarity: source.similarity, + similarity: source._distance, + score: source.score })), }); } catch (e) { From 2c1549d7909bbf82f0dbb19c6af8690911ef864a Mon Sep 17 00:00:00 2001 From: Scott Bowler Date: Thu, 12 Dec 2024 15:22:25 +0000 Subject: [PATCH 4/6] Change similarityThreshold to scoreThreshold As this is what is actually returned by the search --- server/endpoints/api/workspace/index.js | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/server/endpoints/api/workspace/index.js b/server/endpoints/api/workspace/index.js index 9415a47ab64..bcdb9982cd9 100644 --- a/server/endpoints/api/workspace/index.js +++ b/server/endpoints/api/workspace/index.js @@ -863,7 +863,7 @@ function apiWorkspaceEndpoints(app) { example: { query: "What is the meaning of life?", topN: 4, - similarityThreshold: 0.75 + scoreThreshold: 0.75 } } } @@ -889,7 +889,7 @@ function apiWorkspaceEndpoints(app) { wordCount: 8, tokenCount: 9 }, - similarity: 0.541887640953064, + distance: 0.541887640953064, score: 0.45811235904693604 } ] @@ -901,7 +901,7 @@ function apiWorkspaceEndpoints(app) { */ try { const { slug } = request.params; - const { query, topN, similarityThreshold } = reqBody(request); + const { query, topN, scoreThreshold } = reqBody(request); const workspace = await Workspace.get({ slug: String(slug) }); if (!workspace) { @@ -932,8 +932,7 @@ function apiWorkspaceEndpoints(app) { namespace: workspace.slug, input: query, LLMConnector, - similarityThreshold: - similarityThreshold ?? workspace?.similarityThreshold, + similarityThreshold: scoreThreshold ?? workspace?.similarityThreshold, topN: topN ?? workspace?.topN ?? 4, }); @@ -954,7 +953,7 @@ function apiWorkspaceEndpoints(app) { wordCount: source.wordCount, tokenCount: source.token_count_estimate }, - similarity: source._distance, + distance: source._distance, score: source.score })), }); From 8dcd95a78cfa3fff6b355c7865bcda39394e1df0 Mon Sep 17 00:00:00 2001 From: Scott Bowler Date: Thu, 12 Dec 2024 15:25:15 +0000 Subject: [PATCH 5/6] Removing logging (oops!) --- server/endpoints/api/workspace/index.js | 2 -- 1 file changed, 2 deletions(-) diff --git a/server/endpoints/api/workspace/index.js b/server/endpoints/api/workspace/index.js index bcdb9982cd9..fda26970bed 100644 --- a/server/endpoints/api/workspace/index.js +++ b/server/endpoints/api/workspace/index.js @@ -936,8 +936,6 @@ function apiWorkspaceEndpoints(app) { topN: topN ?? workspace?.topN ?? 4, }); - console.log(results); - response.status(200).json({ results: results.sources.map((source) => ({ id: source.id, From 2e8f9d8c0856707b5da12ef56e531564eb134030 Mon Sep 17 00:00:00 2001 From: timothycarambat Date: Thu, 12 Dec 2024 10:09:28 -0800 Subject: [PATCH 6/6] chore: regen swagger docs for new endpoint fix: update function to sanity check values to prevent crashes during search --- server/endpoints/api/workspace/index.js | 49 +++++++++------- server/swagger/openapi.json | 74 +++++++++++++++++++++++++ 2 files changed, 104 insertions(+), 19 deletions(-) diff --git a/server/endpoints/api/workspace/index.js b/server/endpoints/api/workspace/index.js index fda26970bed..2b0ad577dbb 100644 --- a/server/endpoints/api/workspace/index.js +++ b/server/endpoints/api/workspace/index.js @@ -904,36 +904,47 @@ function apiWorkspaceEndpoints(app) { const { query, topN, scoreThreshold } = reqBody(request); const workspace = await Workspace.get({ slug: String(slug) }); - if (!workspace) { - response.status(400).json({ + if (!workspace) + return response.status(400).json({ message: `Workspace ${slug} is not a valid workspace.`, }); - return; - } - if (!query?.length) { - response.status(400).json({ + if (!query?.length) + return 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; - } + 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 LLMConnector = getLLMProvider(); const results = await VectorDb.performSimilaritySearch({ namespace: workspace.slug, - input: query, - LLMConnector, - similarityThreshold: scoreThreshold ?? workspace?.similarityThreshold, - topN: topN ?? workspace?.topN ?? 4, + input: String(query), + LLMConnector: getLLMProvider(), + similarityThreshold: parseSimilarityThreshold(), + topN: parseTopN(), }); response.status(200).json({ @@ -949,10 +960,10 @@ function apiWorkspaceEndpoints(app) { chunkSource: source.chunkSource, published: source.published, wordCount: source.wordCount, - tokenCount: source.token_count_estimate + tokenCount: source.token_count_estimate, }, distance: source._distance, - score: source.score + score: source.score, })), }); } catch (e) { diff --git a/server/swagger/openapi.json b/server/swagger/openapi.json index a849da02890..230398ada54 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": [