From 791c0ee9dc9c3ca245c53a24bf005a46e4a3d24f Mon Sep 17 00:00:00 2001 From: Timothy Carambat Date: Wed, 21 Feb 2024 13:15:45 -0800 Subject: [PATCH] Enable ability to do full-text query on documents (#758) * Enable ability to do full-text query on documents Show alert modal on first pin for client Add ability to use pins in stream/chat/embed * typo and copy update * simplify spread of context and sources --- .../Documents/Directory/index.jsx | 6 +- .../WorkspaceFileRow/index.jsx | 115 +++++++++- .../Documents/WorkspaceDirectory/index.jsx | 205 ++++++++++++------ .../Modals/MangeWorkspace/Documents/index.jsx | 2 +- frontend/src/models/workspace.js | 19 ++ frontend/src/utils/constants.js | 1 + server/endpoints/workspaces.js | 27 +++ server/models/documents.js | 57 ++++- .../20240219211018_init/migration.sql | 2 + server/prisma/schema.prisma | 1 + server/utils/AiProviders/openAi/index.js | 14 +- server/utils/DocumentManager/index.js | 72 ++++++ server/utils/chats/embed.js | 61 ++++-- server/utils/chats/index.js | 61 ++++-- server/utils/chats/stream.js | 61 ++++-- server/utils/files/index.js | 5 + server/utils/helpers/chat/index.js | 34 ++- 17 files changed, 591 insertions(+), 152 deletions(-) create mode 100644 server/prisma/migrations/20240219211018_init/migration.sql create mode 100644 server/utils/DocumentManager/index.js diff --git a/frontend/src/components/Modals/MangeWorkspace/Documents/Directory/index.jsx b/frontend/src/components/Modals/MangeWorkspace/Documents/Directory/index.jsx index af8ae32a..557fe418 100644 --- a/frontend/src/components/Modals/MangeWorkspace/Documents/Directory/index.jsx +++ b/frontend/src/components/Modals/MangeWorkspace/Documents/Directory/index.jsx @@ -1,10 +1,10 @@ import UploadFile from "../UploadFile"; import PreLoader from "@/components/Preloader"; -import { useEffect, useState } from "react"; +import { memo, useEffect, useState } from "react"; import FolderRow from "./FolderRow"; import pluralize from "pluralize"; -export default function Directory({ +function Directory({ files, loading, setLoading, @@ -146,3 +146,5 @@ export default function Directory({ ); } + +export default memo(Directory); diff --git a/frontend/src/components/Modals/MangeWorkspace/Documents/WorkspaceDirectory/WorkspaceFileRow/index.jsx b/frontend/src/components/Modals/MangeWorkspace/Documents/WorkspaceDirectory/WorkspaceFileRow/index.jsx index 91e165d4..93b62cef 100644 --- a/frontend/src/components/Modals/MangeWorkspace/Documents/WorkspaceDirectory/WorkspaceFileRow/index.jsx +++ b/frontend/src/components/Modals/MangeWorkspace/Documents/WorkspaceDirectory/WorkspaceFileRow/index.jsx @@ -1,12 +1,19 @@ -import { useState } from "react"; +import { memo, useState } from "react"; import { formatDate, getFileExtension, middleTruncate, } from "@/utils/directories"; -import { ArrowUUpLeft, File } from "@phosphor-icons/react"; +import { + ArrowUUpLeft, + File, + PushPin, + PushPinSlash, +} from "@phosphor-icons/react"; import Workspace from "@/models/workspace"; import debounce from "lodash.debounce"; +import { Tooltip } from "react-tooltip"; +import showToast from "@/utils/toast"; export default function WorkspaceFileRow({ item, @@ -80,21 +87,105 @@ export default function WorkspaceFileRow({

{getFileExtension(item.url)}

-
- {item?.cached && ( -
-

Cached

-
- )} +
{hasChanges ? (
) : ( - +
+ + +
)}
); } + +const PinItemToWorkspace = memo(({ workspace, docPath, item }) => { + const [pinned, setPinned] = useState( + item?.pinnedWorkspaces?.includes(workspace.id) || false + ); + const [hover, setHover] = useState(false); + const pinEvent = new CustomEvent("pinned_document"); + + const updatePinStatus = async () => { + try { + if (!pinned) window.dispatchEvent(pinEvent); + const success = await Workspace.setPinForDocument( + workspace.slug, + docPath, + !pinned + ); + + if (!success) { + showToast(`Failed to ${!pinned ? "pin" : "unpin"} document.`, "error", { + clear: true, + }); + return; + } + + showToast( + `Document ${!pinned ? "pinned to" : "unpinned from"} workspace`, + "success", + { clear: true } + ); + setPinned(!pinned); + } catch (error) { + showToast(`Failed to pin document. ${error.message}`, "error", { + clear: true, + }); + return; + } + }; + + if (!item) return
; + + const PinIcon = pinned ? PushPinSlash : PushPin; + return ( +
setHover(true)} + onMouseLeave={() => setHover(false)} + > + + +
+ ); +}); + +const RemoveItemFromWorkspace = ({ item, onClick }) => { + return ( +
+ + +
+ ); +}; diff --git a/frontend/src/components/Modals/MangeWorkspace/Documents/WorkspaceDirectory/index.jsx b/frontend/src/components/Modals/MangeWorkspace/Documents/WorkspaceDirectory/index.jsx index 12080b9b..2232e746 100644 --- a/frontend/src/components/Modals/MangeWorkspace/Documents/WorkspaceDirectory/index.jsx +++ b/frontend/src/components/Modals/MangeWorkspace/Documents/WorkspaceDirectory/index.jsx @@ -1,8 +1,12 @@ import PreLoader from "@/components/Preloader"; import { dollarFormat } from "@/utils/numbers"; import WorkspaceFileRow from "./WorkspaceFileRow"; +import { memo, useEffect, useState } from "react"; +import ModalWrapper from "@/components/ModalWrapper"; +import { PushPin } from "@phosphor-icons/react"; +import { SEEN_DOC_PIN_ALERT } from "@/utils/constants"; -export default function WorkspaceDirectory({ +function WorkspaceDirectory({ workspace, files, highlightWorkspace, @@ -29,7 +33,7 @@ export default function WorkspaceDirectory({

Name

Date

Kind

-

Cached

+

@@ -43,78 +47,145 @@ export default function WorkspaceDirectory({ } return ( -
-
-

- {workspace.name} -

-
-
-
-

Name

-

Date

-

Kind

-

Cached

+ <> +
+
+

+ {workspace.name} +

-
- {Object.values(files.items).some( - (folder) => folder.items.length > 0 - ) || movedItems.length > 0 ? ( - <> - {files.items.map((folder) => - folder.items.map((item, index) => ( - - )) - )} - - ) : ( -
-

- No Documents +

+
+

Name

+

Date

+

Kind

+

+

+
+ {Object.values(files.items).some( + (folder) => folder.items.length > 0 + ) || movedItems.length > 0 ? ( + <> + {files.items.map((folder) => + folder.items.map((item, index) => ( + + )) + )} + + ) : ( +
+

+ No Documents +

+
+ )} +
+
+ {hasChanges && ( +
+
+

+ {embeddingCosts === 0 + ? "" + : `Estimated Cost: ${ + embeddingCosts < 0.01 + ? `< $0.01` + : dollarFormat(embeddingCosts) + }`} +

+
- )} -
+ + +
+ )}
- {hasChanges && ( -
-
-

- {embeddingCosts === 0 - ? "" - : `Estimated Cost: ${ - embeddingCosts < 0.01 - ? `< $0.01` - : dollarFormat(embeddingCosts) - }`} + + + ); +} + +const PinAlert = memo(() => { + const [showAlert, setShowAlert] = useState(false); + function dismissAlert() { + setShowAlert(false); + window.localStorage.setItem(SEEN_DOC_PIN_ALERT, "1"); + window.removeEventListener(handlePinEvent); + } + + function handlePinEvent() { + if (!!window?.localStorage?.getItem(SEEN_DOC_PIN_ALERT)) return; + setShowAlert(true); + } + + useEffect(() => { + if (!window || !!window?.localStorage?.getItem(SEEN_DOC_PIN_ALERT)) return; + window?.addEventListener("pinned_document", handlePinEvent); + }, []); + + return ( + +

+
+
+
+ +

+ What is document pinning? +

+
+
+
+

+ When you pin a document in AnythingLLM we will inject the + entire content of the document into your prompt window for your + LLM to fully comprehend.

-

+ This works best with large-context models or small files + that are critical to its knowledge-base. +

+

+ If you are not getting the answers you desire from AnythingLLM by + default then pinning is a great way to get higher quality answers + in a click.

- +
+ +
- )} -
+
+ ); -} +}); + +export default memo(WorkspaceDirectory); diff --git a/frontend/src/components/Modals/MangeWorkspace/Documents/index.jsx b/frontend/src/components/Modals/MangeWorkspace/Documents/index.jsx index 705e78fa..e8b63c90 100644 --- a/frontend/src/components/Modals/MangeWorkspace/Documents/index.jsx +++ b/frontend/src/components/Modals/MangeWorkspace/Documents/index.jsx @@ -2,8 +2,8 @@ import { ArrowsDownUp } from "@phosphor-icons/react"; import { useEffect, useState } from "react"; import Workspace from "../../../../models/workspace"; import System from "../../../../models/system"; -import Directory from "./Directory"; import showToast from "../../../../utils/toast"; +import Directory from "./Directory"; import WorkspaceDirectory from "./WorkspaceDirectory"; // OpenAI Cost per token diff --git a/frontend/src/models/workspace.js b/frontend/src/models/workspace.js index d77e2ad5..d0038874 100644 --- a/frontend/src/models/workspace.js +++ b/frontend/src/models/workspace.js @@ -218,6 +218,25 @@ const Workspace = { return { success: false, error: e.message }; }); }, + setPinForDocument: async function (slug, docPath, pinStatus) { + return fetch(`${API_BASE}/workspace/${slug}/update-pin`, { + method: "POST", + headers: baseHeaders(), + body: JSON.stringify({ docPath, pinStatus }), + }) + .then((res) => { + if (!res.ok) { + throw new Error( + res.statusText || "Error setting pin status for document." + ); + } + return true; + }) + .catch((e) => { + console.error(e); + return false; + }); + }, threads: WorkspaceThread, }; diff --git a/frontend/src/utils/constants.js b/frontend/src/utils/constants.js index 2fde1ee0..6fd29534 100644 --- a/frontend/src/utils/constants.js +++ b/frontend/src/utils/constants.js @@ -4,6 +4,7 @@ export const AUTH_USER = "anythingllm_user"; export const AUTH_TOKEN = "anythingllm_authToken"; export const AUTH_TIMESTAMP = "anythingllm_authTimestamp"; export const COMPLETE_QUESTIONNAIRE = "anythingllm_completed_questionnaire"; +export const SEEN_DOC_PIN_ALERT = "anythingllm_pinned_document_alert"; export const USER_BACKGROUND_COLOR = "bg-historical-msg-user"; export const AI_BACKGROUND_COLOR = "bg-historical-msg-system"; diff --git a/server/endpoints/workspaces.js b/server/endpoints/workspaces.js index 82040272..54228bba 100644 --- a/server/endpoints/workspaces.js +++ b/server/endpoints/workspaces.js @@ -395,6 +395,33 @@ function workspaceEndpoints(app) { } } ); + + app.post( + "/workspace/:slug/update-pin", + [ + validatedRequest, + flexUserRoleValid([ROLES.admin, ROLES.manager]), + validWorkspaceSlug, + ], + async (request, response) => { + try { + const { docPath, pinStatus = false } = reqBody(request); + const workspace = response.locals.workspace; + + const document = await Document.get({ + workspaceId: workspace.id, + docpath: docPath, + }); + if (!document) return response.sendStatus(404).end(); + + await Document.update(document.id, { pinned: pinStatus }); + return response.status(200).end(); + } catch (error) { + console.error("Error processing the pin status update:", error); + return response.status(500).end(); + } + } + ); } module.exports = { workspaceEndpoints }; diff --git a/server/models/documents.js b/server/models/documents.js index 9f50aa91..aa62ccf6 100644 --- a/server/models/documents.js +++ b/server/models/documents.js @@ -1,4 +1,3 @@ -const { fileData } = require("../utils/files"); const { v4: uuidv4 } = require("uuid"); const { getVectorDbClass } = require("../utils/helpers"); const prisma = require("../utils/prisma"); @@ -6,6 +5,8 @@ const { Telemetry } = require("./telemetry"); const { EventLogs } = require("./eventLogs"); const Document = { + writable: ["pinned"], + forWorkspace: async function (workspaceId = null) { if (!workspaceId) return []; return await prisma.workspace_documents.findMany({ @@ -23,7 +24,7 @@ const Document = { } }, - firstWhere: async function (clause = {}) { + get: async function (clause = {}) { try { const document = await prisma.workspace_documents.findFirst({ where: clause, @@ -35,9 +36,39 @@ const Document = { } }, + getPins: async function (clause = {}) { + try { + const workspaceIds = await prisma.workspace_documents.findMany({ + where: clause, + select: { + workspaceId: true, + }, + }); + return workspaceIds.map((pin) => pin.workspaceId) || []; + } catch (error) { + console.error(error.message); + return []; + } + }, + + where: async function (clause = {}, limit = null, orderBy = null) { + try { + const results = await prisma.workspace_documents.findMany({ + where: clause, + ...(limit !== null ? { take: limit } : {}), + ...(orderBy !== null ? { orderBy } : {}), + }); + return results; + } catch (error) { + console.error(error.message); + return []; + } + }, + addDocuments: async function (workspace, additions = [], userId = null) { const VectorDb = getVectorDbClass(); if (additions.length === 0) return { failed: [], embedded: [] }; + const { fileData } = require("../utils/files"); const embedded = []; const failedToEmbed = []; const errors = new Set(); @@ -101,7 +132,7 @@ const Document = { if (removals.length === 0) return; for (const path of removals) { - const document = await this.firstWhere({ + const document = await this.get({ docpath: path, workspaceId: workspace.id, }); @@ -151,6 +182,26 @@ const Document = { return 0; } }, + update: async function (id = null, data = {}) { + if (!id) throw new Error("No workspace document id provided for update"); + + const validKeys = Object.keys(data).filter((key) => + this.writable.includes(key) + ); + if (validKeys.length === 0) + return { document: { id }, message: "No valid fields to update!" }; + + try { + const document = await prisma.workspace_documents.update({ + where: { id }, + data, + }); + return { document, message: null }; + } catch (error) { + console.error(error.message); + return { document: null, message: error.message }; + } + }, }; module.exports = { Document }; diff --git a/server/prisma/migrations/20240219211018_init/migration.sql b/server/prisma/migrations/20240219211018_init/migration.sql new file mode 100644 index 00000000..98e8b24a --- /dev/null +++ b/server/prisma/migrations/20240219211018_init/migration.sql @@ -0,0 +1,2 @@ +-- AlterTable +ALTER TABLE "workspace_documents" ADD COLUMN "pinned" BOOLEAN DEFAULT false; diff --git a/server/prisma/schema.prisma b/server/prisma/schema.prisma index 77b25c8d..8cd3a1d3 100644 --- a/server/prisma/schema.prisma +++ b/server/prisma/schema.prisma @@ -30,6 +30,7 @@ model workspace_documents { docpath String workspaceId Int metadata String? + pinned Boolean? @default(false) createdAt DateTime @default(now()) lastUpdatedAt DateTime @default(now()) workspace workspaces @relation(fields: [workspaceId], references: [id]) diff --git a/server/utils/AiProviders/openAi/index.js b/server/utils/AiProviders/openAi/index.js index c3c983f8..d4dc14dc 100644 --- a/server/utils/AiProviders/openAi/index.js +++ b/server/utils/AiProviders/openAi/index.js @@ -195,11 +195,15 @@ class OpenAiLLM { `OpenAI chat: ${this.model} is not valid for chat completion!` ); - const { data } = await this.openai.createChatCompletion({ - model: this.model, - messages, - temperature, - }); + const { data } = await this.openai + .createChatCompletion({ + model: this.model, + messages, + temperature, + }) + .catch((e) => { + throw new Error(e.response.data.error.message); + }); if (!data.hasOwnProperty("choices")) return null; return data.choices[0].message.content; diff --git a/server/utils/DocumentManager/index.js b/server/utils/DocumentManager/index.js new file mode 100644 index 00000000..17fd9860 --- /dev/null +++ b/server/utils/DocumentManager/index.js @@ -0,0 +1,72 @@ +const fs = require("fs"); +const path = require("path"); + +const documentsPath = + process.env.NODE_ENV === "development" + ? path.resolve(__dirname, `../../storage/documents`) + : path.resolve(process.env.STORAGE_DIR, `documents`); + +class DocumentManager { + constructor({ workspace = null, maxTokens = null }) { + this.workspace = workspace; + this.maxTokens = maxTokens || Number.POSITIVE_INFINITY; + this.documentStoragePath = documentsPath; + } + + log(text, ...args) { + console.log(`\x1b[36m[DocumentManager]\x1b[0m ${text}`, ...args); + } + + async pinnedDocuments() { + if (!this.workspace) return []; + const { Document } = require("../../models/documents"); + return await Document.where({ + workspaceId: Number(this.workspace.id), + pinned: true, + }); + } + + async pinnedDocs() { + if (!this.workspace) return []; + const docPaths = (await this.pinnedDocuments()).map((doc) => doc.docpath); + if (docPaths.length === 0) return []; + + let tokens = 0; + const pinnedDocs = []; + for await (const docPath of docPaths) { + try { + const filePath = path.resolve(this.documentStoragePath, docPath); + const data = JSON.parse( + fs.readFileSync(filePath, { encoding: "utf-8" }) + ); + + if ( + !data.hasOwnProperty("pageContent") || + !data.hasOwnProperty("token_count_estimate") + ) { + this.log( + `Skipping document - Could not find page content or token_count_estimate in pinned source.` + ); + continue; + } + + if (tokens >= this.maxTokens) { + this.log( + `Skipping document - Token limit of ${this.maxTokens} has already been exceeded by pinned documents.` + ); + continue; + } + + pinnedDocs.push(data); + tokens += data.token_count_estimate || 0; + } catch {} + } + + this.log( + `Found ${pinnedDocs.length} pinned sources - prepending to content with ~${tokens} tokens of content.` + ); + return pinnedDocs; + } +} + +module.exports.DocumentManager = DocumentManager; diff --git a/server/utils/chats/embed.js b/server/utils/chats/embed.js index 7a4c52d1..f748a3a5 100644 --- a/server/utils/chats/embed.js +++ b/server/utils/chats/embed.js @@ -6,6 +6,7 @@ const { convertToPromptHistory, writeResponseChunk, } = require("../helpers/chat/responses"); +const { DocumentManager } = require("../DocumentManager"); async function streamChatWithForEmbed( response, @@ -64,6 +65,8 @@ async function streamChatWithForEmbed( } let completeText; + let contextTexts = []; + let sources = []; const { rawHistory, chatHistory } = await recentEmbedChatHistory( sessionId, embed, @@ -71,26 +74,43 @@ async function streamChatWithForEmbed( chatMode ); - const { - contextTexts = [], - sources = [], - message: error, - } = embeddingsCount !== 0 // if there no embeddings don't bother searching. - ? await VectorDb.performSimilaritySearch({ - namespace: embed.workspace.slug, - input: message, - LLMConnector, - similarityThreshold: embed.workspace?.similarityThreshold, - topN: embed.workspace?.topN, - }) - : { - contextTexts: [], - sources: [], - message: null, - }; + // Look for pinned documents and see if the user decided to use this feature. We will also do a vector search + // as pinning is a supplemental tool but it should be used with caution since it can easily blow up a context window. + await new DocumentManager({ + workspace: embed.workspace, + maxTokens: LLMConnector.limits.system, + }) + .pinnedDocs() + .then((pinnedDocs) => { + pinnedDocs.forEach((doc) => { + const { pageContent, ...metadata } = doc; + contextTexts.push(doc.pageContent); + sources.push({ + text: + pageContent.slice(0, 1_000) + + "...continued on in source document...", + ...metadata, + }); + }); + }); - // Failed similarity search. - if (!!error) { + const vectorSearchResults = + embeddingsCount !== 0 + ? await VectorDb.performSimilaritySearch({ + namespace: embed.workspace.slug, + input: message, + LLMConnector, + similarityThreshold: embed.workspace?.similarityThreshold, + topN: embed.workspace?.topN, + }) + : { + contextTexts: [], + sources: [], + message: null, + }; + + // Failed similarity search if it was run at all and failed. + if (!!vectorSearchResults.message) { writeResponseChunk(response, { id: uuid, type: "abort", @@ -102,6 +122,9 @@ async function streamChatWithForEmbed( return; } + contextTexts = [...contextTexts, ...vectorSearchResults.contextTexts]; + sources = [...sources, ...vectorSearchResults.sources]; + // If in query mode and no sources are found, do not // let the LLM try to hallucinate a response or use general knowledge if (chatMode === "query" && sources.length === 0) { diff --git a/server/utils/chats/index.js b/server/utils/chats/index.js index 6d8cccf9..10df9983 100644 --- a/server/utils/chats/index.js +++ b/server/utils/chats/index.js @@ -3,6 +3,7 @@ const { WorkspaceChats } = require("../../models/workspaceChats"); const { resetMemory } = require("./commands/reset"); const { getVectorDbClass, getLLMProvider } = require("../helpers"); const { convertToPromptHistory } = require("../helpers/chat/responses"); +const { DocumentManager } = require("../DocumentManager"); const VALID_COMMANDS = { "/reset": resetMemory, @@ -73,6 +74,8 @@ async function chatWithWorkspace( // If we are here we know that we are in a workspace that is: // 1. Chatting in "chat" mode and may or may _not_ have embeddings // 2. Chatting in "query" mode and has at least 1 embedding + let contextTexts = []; + let sources = []; const { rawHistory, chatHistory } = await recentChatHistory({ user, workspace, @@ -81,36 +84,56 @@ async function chatWithWorkspace( chatMode, }); - const { - contextTexts = [], - sources = [], - message: error, - } = embeddingsCount !== 0 // if there no embeddings don't bother searching. - ? await VectorDb.performSimilaritySearch({ - namespace: workspace.slug, - input: message, - LLMConnector, - similarityThreshold: workspace?.similarityThreshold, - topN: workspace?.topN, - }) - : { - contextTexts: [], - sources: [], - message: null, - }; + // Look for pinned documents and see if the user decided to use this feature. We will also do a vector search + // as pinning is a supplemental tool but it should be used with caution since it can easily blow up a context window. + await new DocumentManager({ + workspace, + maxTokens: LLMConnector.limits.system, + }) + .pinnedDocs() + .then((pinnedDocs) => { + pinnedDocs.forEach((doc) => { + const { pageContent, ...metadata } = doc; + contextTexts.push(doc.pageContent); + sources.push({ + text: + pageContent.slice(0, 1_000) + + "...continued on in source document...", + ...metadata, + }); + }); + }); + + const vectorSearchResults = + embeddingsCount !== 0 + ? await VectorDb.performSimilaritySearch({ + namespace: workspace.slug, + input: message, + LLMConnector, + similarityThreshold: workspace?.similarityThreshold, + topN: workspace?.topN, + }) + : { + contextTexts: [], + sources: [], + message: null, + }; // Failed similarity search if it was run at all and failed. - if (!!error) { + if (!!vectorSearchResults.message) { return { id: uuid, type: "abort", textResponse: null, sources: [], close: true, - error, + error: vectorSearchResults.message, }; } + contextTexts = [...contextTexts, ...vectorSearchResults.contextTexts]; + sources = [...sources, ...vectorSearchResults.sources]; + // If in query mode and no sources are found, do not // let the LLM try to hallucinate a response or use general knowledge and exit early if (chatMode === "query" && sources.length === 0) { diff --git a/server/utils/chats/stream.js b/server/utils/chats/stream.js index 4f86c49d..f1a335bc 100644 --- a/server/utils/chats/stream.js +++ b/server/utils/chats/stream.js @@ -1,4 +1,5 @@ const { v4: uuidv4 } = require("uuid"); +const { DocumentManager } = require("../DocumentManager"); const { WorkspaceChats } = require("../../models/workspaceChats"); const { getVectorDbClass, getLLMProvider } = require("../helpers"); const { writeResponseChunk } = require("../helpers/chat/responses"); @@ -74,6 +75,8 @@ async function streamChatWithWorkspace( // 1. Chatting in "chat" mode and may or may _not_ have embeddings // 2. Chatting in "query" mode and has at least 1 embedding let completeText; + let contextTexts = []; + let sources = []; const { rawHistory, chatHistory } = await recentChatHistory({ user, workspace, @@ -82,37 +85,57 @@ async function streamChatWithWorkspace( chatMode, }); - const { - contextTexts = [], - sources = [], - message: error, - } = embeddingsCount !== 0 // if there no embeddings don't bother searching. - ? await VectorDb.performSimilaritySearch({ - namespace: workspace.slug, - input: message, - LLMConnector, - similarityThreshold: workspace?.similarityThreshold, - topN: workspace?.topN, - }) - : { - contextTexts: [], - sources: [], - message: null, - }; + // Look for pinned documents and see if the user decided to use this feature. We will also do a vector search + // as pinning is a supplemental tool but it should be used with caution since it can easily blow up a context window. + await new DocumentManager({ + workspace, + maxTokens: LLMConnector.limits.system, + }) + .pinnedDocs() + .then((pinnedDocs) => { + pinnedDocs.forEach((doc) => { + const { pageContent, ...metadata } = doc; + contextTexts.push(doc.pageContent); + sources.push({ + text: + pageContent.slice(0, 1_000) + + "...continued on in source document...", + ...metadata, + }); + }); + }); + + const vectorSearchResults = + embeddingsCount !== 0 + ? await VectorDb.performSimilaritySearch({ + namespace: workspace.slug, + input: message, + LLMConnector, + similarityThreshold: workspace?.similarityThreshold, + topN: workspace?.topN, + }) + : { + contextTexts: [], + sources: [], + message: null, + }; // Failed similarity search if it was run at all and failed. - if (!!error) { + if (!!vectorSearchResults.message) { writeResponseChunk(response, { id: uuid, type: "abort", textResponse: null, sources: [], close: true, - error, + error: vectorSearchResults.message, }); return; } + contextTexts = [...contextTexts, ...vectorSearchResults.contextTexts]; + sources = [...sources, ...vectorSearchResults.sources]; + // If in query mode and no sources are found, do not // let the LLM try to hallucinate a response or use general knowledge and exit early if (chatMode === "query" && sources.length === 0) { diff --git a/server/utils/files/index.js b/server/utils/files/index.js index dff5bef9..1ba01796 100644 --- a/server/utils/files/index.js +++ b/server/utils/files/index.js @@ -1,6 +1,7 @@ const fs = require("fs"); const path = require("path"); const { v5: uuidv5 } = require("uuid"); +const { Document } = require("../../models/documents"); const documentsPath = process.env.NODE_ENV === "development" ? path.resolve(__dirname, `../../storage/documents`) @@ -55,6 +56,10 @@ async function viewLocalFiles() { type: "file", ...metadata, cached: await cachedVectorInformation(cachefilename, true), + pinnedWorkspaces: await Document.getPins({ + docpath: cachefilename, + pinned: true, + }), }); } directory.items.push(subdocs); diff --git a/server/utils/helpers/chat/index.js b/server/utils/helpers/chat/index.js index 7292c422..84afd516 100644 --- a/server/utils/helpers/chat/index.js +++ b/server/utils/helpers/chat/index.js @@ -85,11 +85,35 @@ async function messageArrayCompressor(llm, messages = [], rawHistory = []) { // Split context from system prompt - cannonball since its over the window. // We assume the context + user prompt is enough tokens to fit. const [prompt, context = ""] = system.content.split("Context:"); - system.content = `${cannonball({ - input: prompt, - targetTokenSize: llm.limits.system, - tiktokenInstance: tokenManager, - })}${context ? `\nContext: ${context}` : ""}`; + let compressedPrompt; + let compressedContext; + + // If the user system prompt contribution's to the system prompt is more than + // 25% of the system limit, we will cannonball it - this favors the context + // over the instruction from the user. + if (tokenManager.countFromString(prompt) >= llm.limits.system * 0.25) { + compressedPrompt = cannonball({ + input: prompt, + targetTokenSize: llm.limits.system * 0.25, + tiktokenInstance: tokenManager, + }); + } else { + compressedPrompt = prompt; + } + + if (tokenManager.countFromString(context) >= llm.limits.system * 0.75) { + compressedContext = cannonball({ + input: context, + targetTokenSize: llm.limits.system * 0.75, + tiktokenInstance: tokenManager, + }); + } else { + compressedContext = context; + } + + system.content = `${compressedPrompt}${ + compressedContext ? `\nContext: ${compressedContext}` : "" + }`; resolve(system); });