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
2024-06-30 14:40:08 +02:00 · 2024-02-21 13:15:45 -08:00 · 2024-02-21 13:15:45 -08:00 · 791c0ee9dc
commit 791c0ee9dc
parent e63c426223
17 changed files with 591 additions and 152 deletions
--- 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({
    </div>
  );
 }
+
+export default memo(Directory);
--- 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({
      <p className="col-span-2 pl-2 uppercase overflow-x-hidden">
        {getFileExtension(item.url)}
      </p>
-      <div className="col-span-2 flex justify-end items-center">
-        {item?.cached && (
-          <div className="bg-white/10 rounded-3xl">
-            <p className="text-xs px-2 py-0.5">Cached</p>
-          </div>
-        )}
+      <div className="col-span-2 flex justify-center items-center">
        {hasChanges ? (
          <div className="w-4 h-4 ml-2 flex-shrink-0" />
        ) : (
-          <ArrowUUpLeft
-            onClick={onRemoveClick}
-            className="text-base font-bold w-4 h-4 ml-2 flex-shrink-0 cursor-pointer"
-          />
+          <div className="flex gap-x-2 items-center">
+            <PinItemToWorkspace
+              workspace={workspace}
+              docPath={`${folderName}/${item.name}`} // how to find documents during pin/unpin
+              item={item}
+            />
+            <RemoveItemFromWorkspace item={item} onClick={onRemoveClick} />
+          </div>
        )}
      </div>
    </div>
  );
 }
+
+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 <div />;
+
+  const PinIcon = pinned ? PushPinSlash : PushPin;
+  return (
+    <div
+      onMouseEnter={() => setHover(true)}
+      onMouseLeave={() => setHover(false)}
+    >
+      <PinIcon
+        data-tooltip-id={`pin-${item.id}`}
+        data-tooltip-content={
+          pinned ? "Unpin document from workspace" : "Pin document to workspace"
+        }
+        onClick={updatePinStatus}
+        weight={hover ? "fill" : "regular"}
+        className={`outline-none text-base font-bold w-4 h-4 ml-2 flex-shrink-0 cursor-pointer ${
+          pinned ? "hover:text-red-300" : ""
+        }`}
+      />
+      <Tooltip
+        id={`pin-${item.id}`}
+        place="bottom"
+        delayShow={300}
+        className="tooltip !text-xs"
+      />
+    </div>
+  );
+});
+
+const RemoveItemFromWorkspace = ({ item, onClick }) => {
+  return (
+    <div>
+      <ArrowUUpLeft
+        data-tooltip-id={`remove-${item.id}`}
+        data-tooltip-content="Remove document from workspace"
+        onClick={onClick}
+        className="text-base font-bold w-4 h-4 ml-2 flex-shrink-0 cursor-pointer"
+      />
+      <Tooltip
+        id={`remove-${item.id}`}
+        place="bottom"
+        delayShow={300}
+        className="tooltip !text-xs"
+      />
+    </div>
+  );
+};
--- 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({
            <p className="col-span-5">Name</p>
            <p className="col-span-3">Date</p>
            <p className="col-span-2">Kind</p>
-            <p className="col-span-2">Cached</p>
+            <p className="col-span-2" />
          </div>
          <div className="w-full h-full flex items-center justify-center flex-col gap-y-5">
            <PreLoader />
@ -43,78 +47,145 @@ export default function WorkspaceDirectory({
  }

  return (
-    <div className="px-8">
-      <div className="flex items-center justify-start w-[560px]">
-        <h3 className="text-white text-base font-bold ml-5">
-          {workspace.name}
-        </h3>
-      </div>
-      <div
-        className={`relative w-[560px] h-[445px] bg-zinc-900 rounded-2xl mt-5 overflow-y-auto border-4 transition-all duration-300 ${
-          highlightWorkspace ? "border-cyan-300/80" : "border-transparent"
-        }`}
-      >
-        <div className="text-white/80 text-xs grid grid-cols-12 py-2 px-8 border-b border-white/20 bg-zinc-900 sticky top-0 z-10">
-          <p className="col-span-5">Name</p>
-          <p className="col-span-3">Date</p>
-          <p className="col-span-2">Kind</p>
-          <p className="col-span-2">Cached</p>
+    <>
+      <div className="px-8">
+        <div className="flex items-center justify-start w-[560px]">
+          <h3 className="text-white text-base font-bold ml-5">
+            {workspace.name}
+          </h3>
        </div>
-        <div className="w-full h-full flex flex-col z-0">
-          {Object.values(files.items).some(
-            (folder) => folder.items.length > 0
-          ) || movedItems.length > 0 ? (
-            <>
-              {files.items.map((folder) =>
-                folder.items.map((item, index) => (
-                  <WorkspaceFileRow
-                    key={index}
-                    item={item}
-                    folderName={folder.name}
-                    workspace={workspace}
-                    setLoading={setLoading}
-                    setLoadingMessage={setLoadingMessage}
-                    fetchKeys={fetchKeys}
-                    hasChanges={hasChanges}
-                    movedItems={movedItems}
-                  />
-                ))
-              )}
-            </>
-          ) : (
-            <div className="w-full h-full flex items-center justify-center">
-              <p className="text-white text-opacity-40 text-sm font-medium">
-                No Documents
+        <div
+          className={`relative w-[560px] h-[445px] bg-zinc-900 rounded-2xl mt-5 overflow-y-auto border-4 transition-all duration-300 ${
+            highlightWorkspace ? "border-cyan-300/80" : "border-transparent"
+          }`}
+        >
+          <div className="text-white/80 text-xs grid grid-cols-12 py-2 px-8 border-b border-white/20 bg-zinc-900 sticky top-0 z-10">
+            <p className="col-span-5">Name</p>
+            <p className="col-span-3">Date</p>
+            <p className="col-span-2">Kind</p>
+            <p className="col-span-2" />
+          </div>
+          <div className="w-full h-full flex flex-col z-0">
+            {Object.values(files.items).some(
+              (folder) => folder.items.length > 0
+            ) || movedItems.length > 0 ? (
+              <>
+                {files.items.map((folder) =>
+                  folder.items.map((item, index) => (
+                    <WorkspaceFileRow
+                      key={index}
+                      item={item}
+                      folderName={folder.name}
+                      workspace={workspace}
+                      setLoading={setLoading}
+                      setLoadingMessage={setLoadingMessage}
+                      fetchKeys={fetchKeys}
+                      hasChanges={hasChanges}
+                      movedItems={movedItems}
+                    />
+                  ))
+                )}
+              </>
+            ) : (
+              <div className="w-full h-full flex items-center justify-center">
+                <p className="text-white text-opacity-40 text-sm font-medium">
+                  No Documents
+                </p>
+              </div>
+            )}
+          </div>
+        </div>
+        {hasChanges && (
+          <div className="flex items-center justify-between py-6 transition-all duration-300">
+            <div className="text-white/80">
+              <p className="text-sm font-semibold">
+                {embeddingCosts === 0
+                  ? ""
+                  : `Estimated Cost: ${
+                      embeddingCosts < 0.01
+                        ? `< $0.01`
+                        : dollarFormat(embeddingCosts)
+                    }`}
+              </p>
+              <p className="mt-2 text-xs italic" hidden={embeddingCosts === 0}>
+                *One time cost for embeddings
              </p>
            </div>
-          )}
-        </div>
+
+            <button
+              onClick={saveChanges}
+              className="transition-all duration-300 border border-slate-200 px-5 py-2.5 rounded-lg text-white text-sm items-center flex gap-x-2 hover:bg-slate-200 hover:text-slate-800 focus:ring-gray-800"
+            >
+              Save and Embed
+            </button>
+          </div>
+        )}
      </div>
-      {hasChanges && (
-        <div className="flex items-center justify-between py-6 transition-all duration-300">
-          <div className="text-white/80">
-            <p className="text-sm font-semibold">
-              {embeddingCosts === 0
-                ? ""
-                : `Estimated Cost: ${
-                    embeddingCosts < 0.01
-                      ? `< $0.01`
-                      : dollarFormat(embeddingCosts)
-                  }`}
+      <PinAlert />
+    </>
+  );
+}
+
+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 (
+    <ModalWrapper isOpen={showAlert}>
+      <div className="relative w-full max-w-2xl max-h-full">
+        <div className="relative bg-main-gradient rounded-lg shadow">
+          <div className="flex items-start justify-between p-4 border-b rounded-t border-gray-500/50">
+            <div className="flex items-center gap-2">
+              <PushPin className="text-red-600 text-lg w-6 h-6" weight="fill" />
+              <h3 className="text-xl font-semibold text-white">
+                What is document pinning?
+              </h3>
+            </div>
+          </div>
+          <div className="w-full p-6 text-white text-md flex flex-col gap-y-2">
+            <p>
+              When you <b>pin</b> a document in AnythingLLM we will inject the
+              entire content of the document into your prompt window for your
+              LLM to fully comprehend.
            </p>
-            <p className="mt-2 text-xs italic" hidden={embeddingCosts === 0}>
-              *One time cost for embeddings
+            <p>
+              This works best with <b>large-context models</b> or small files
+              that are critical to its knowledge-base.
+            </p>
+            <p>
+              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.
            </p>
          </div>

-          <button
-            onClick={saveChanges}
-            className="transition-all duration-300 border border-slate-200 px-5 py-2.5 rounded-lg text-white text-sm items-center flex gap-x-2 hover:bg-slate-200 hover:text-slate-800 focus:ring-gray-800"
-          >
-            Save and Embed
-          </button>
+          <div className="flex w-full justify-between items-center p-6 space-x-2 border-t rounded-b border-gray-500/50">
+            <button disabled={true} className="invisible" />
+            <button
+              onClick={dismissAlert}
+              className="transition-all duration-300 border border-slate-200 px-4 py-2 rounded-lg text-white text-sm items-center flex gap-x-2 hover:bg-slate-200 hover:text-slate-800 focus:ring-gray-800"
+            >
+              Okay, got it
+            </button>
+          </div>
        </div>
-      )}
-    </div>
+      </div>
+    </ModalWrapper>
  );
-}
+});
+
+export default memo(WorkspaceDirectory);
--- 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
--- 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,
 };

--- 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";
--- 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 };
--- 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 };
--- a/server/prisma/migrations/20240219211018_init/migration.sql
+++ b/server/prisma/migrations/20240219211018_init/migration.sql
@ -0,0 +1,2 @@
+-- AlterTable
+ALTER TABLE "workspace_documents" ADD COLUMN "pinned" BOOLEAN DEFAULT false;
--- 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])
--- 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;
--- a/server/utils/DocumentManager/index.js
+++ 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;
--- 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) {
--- 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) {
--- 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) {
--- 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);
--- 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);
  });