totem-llm 0.1.0

package/server/utils/chats/stream.js

@@ -0,0 +1,376 @@
+const { v4: uuidv4 } = require("uuid");
+const { DocumentManager } = require("../DocumentManager");
+const { WorkspaceChats } = require("../../models/workspaceChats");
+const { WorkspaceParsedFiles } = require("../../models/workspaceParsedFiles");
+const { getVectorDbClass, resolveProviderConnector } = require("../helpers");
+const { writeResponseChunk } = require("../helpers/chat/responses");
+const { grepAgents } = require("./agents");
+const {
+  grepCommand,
+  VALID_COMMANDS,
+  chatPrompt,
+  recentChatHistory,
+  sourceIdentifier,
+} = require("./index");
+
+const VALID_CHAT_MODE = ["automatic", "chat", "query"];
+
+async function streamChatWithWorkspace(
+  response,
+  workspace,
+  message,
+  chatMode = "automatic",
+  user = null,
+  thread = null,
+  attachments = []
+) {
+  const uuid = uuidv4();
+  const updatedMessage = await grepCommand(message, user);
+
+  if (Object.keys(VALID_COMMANDS).includes(updatedMessage)) {
+    const data = await VALID_COMMANDS[updatedMessage](
+      workspace,
+      message,
+      uuid,
+      user,
+      thread
+    );
+    writeResponseChunk(response, data);
+    return;
+  }
+
+  // If is agent enabled chat we will exit this flow early.
+  const isAgentChat = await grepAgents({
+    uuid,
+    response,
+    message: updatedMessage,
+    user,
+    workspace,
+    thread,
+    attachments,
+  });
+  if (isAgentChat) return;
+
+  const {
+    connector: LLMConnector,
+    routingMetadata,
+    prefetchedContext,
+    error: routerError,
+  } = await resolveLLMConnector({
+    workspace,
+    message: updatedMessage,
+    user,
+    thread,
+    attachments,
+  });
+
+  if (routerError) {
+    return writeResponseChunk(response, {
+      id: uuid,
+      type: "abort",
+      textResponse: null,
+      sources: [],
+      close: true,
+      error: routerError,
+    });
+  }
+
+  if (routingMetadata?.routedTo?.shouldNotify) {
+    writeResponseChunk(response, {
+      uuid: `${uuid}:route`,
+      type: "modelRouteNotification",
+      routedTo: routingMetadata.routedTo,
+    });
+  }
+
+  const VectorDb = getVectorDbClass();
+
+  const messageLimit = workspace?.openAiHistory || 20;
+  const hasVectorizedSpace = await VectorDb.hasNamespace(workspace.slug);
+  const embeddingsCount = await VectorDb.namespaceCount(workspace.slug);
+
+  // User is trying to query-mode chat a workspace that has no data in it - so
+  // we should exit early as no information can be found under these conditions.
+  if ((!hasVectorizedSpace || embeddingsCount === 0) && chatMode === "query") {
+    const textResponse =
+      workspace?.queryRefusalResponse ??
+      "There is no relevant information in this workspace to answer your query.";
+    writeResponseChunk(response, {
+      id: uuid,
+      type: "textResponse",
+      textResponse,
+      sources: [],
+      attachments,
+      close: true,
+      error: null,
+    });
+    await WorkspaceChats.new({
+      workspaceId: workspace.id,
+      prompt: message,
+      response: {
+        text: textResponse,
+        sources: [],
+        type: chatMode,
+        attachments,
+      },
+      threadId: thread?.id || null,
+      include: false,
+      user,
+    });
+    return;
+  }
+
+  // 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 completeText;
+  let metrics = {};
+  let contextTexts = [];
+  let sources = [];
+  let pinnedDocIdentifiers = [];
+
+  // If the router pre-fetched context we can reuse it; otherwise fetch fresh.
+  const {
+    rawHistory,
+    chatHistory,
+    pinnedDocs: prefetchedPinnedDocs,
+    parsedFiles: prefetchedParsedFiles,
+  } = prefetchedContext ??
+  (await recentChatHistory({ user, workspace, thread, messageLimit }));
+
+  // Pinned docs — reuse pre-fetched if available, otherwise fetch with token cap.
+  const pinnedDocs =
+    prefetchedPinnedDocs ??
+    (await new DocumentManager({
+      workspace,
+      maxTokens: LLMConnector.promptWindowLimit(),
+    }).pinnedDocs());
+  pinnedDocs.forEach((doc) => {
+    const { pageContent, ...metadata } = doc;
+    pinnedDocIdentifiers.push(sourceIdentifier(doc));
+    contextTexts.push(doc.pageContent);
+    sources.push({
+      text:
+        pageContent.slice(0, 1_000) + "...continued on in source document...",
+      ...metadata,
+    });
+  });
+
+  // Parsed files — reuse pre-fetched if available, otherwise fetch fresh.
+  const parsedFiles =
+    prefetchedParsedFiles ??
+    (await WorkspaceParsedFiles.getContextFiles(
+      workspace,
+      thread || null,
+      user || null
+    ));
+  parsedFiles.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: updatedMessage,
+          LLMConnector,
+          similarityThreshold: workspace?.similarityThreshold,
+          topN: workspace?.topN,
+          filterIdentifiers: pinnedDocIdentifiers,
+          rerank: workspace?.vectorSearchMode === "rerank",
+        })
+      : {
+          contextTexts: [],
+          sources: [],
+          message: null,
+        };
+
+  // Failed similarity search if it was run at all and failed.
+  if (!!vectorSearchResults.message) {
+    writeResponseChunk(response, {
+      id: uuid,
+      type: "abort",
+      textResponse: null,
+      sources: [],
+      close: true,
+      error: vectorSearchResults.message,
+    });
+    return;
+  }
+
+  const { fillSourceWindow } = require("../helpers/chat");
+  const filledSources = fillSourceWindow({
+    nDocs: workspace?.topN || 4,
+    searchResults: vectorSearchResults.sources,
+    history: rawHistory,
+    filterIdentifiers: pinnedDocIdentifiers,
+  });
+
+  // Why does contextTexts get all the info, but sources only get current search?
+  // This is to give the ability of the LLM to "comprehend" a contextual response without
+  // populating the Citations under a response with documents the user "thinks" are irrelevant
+  // due to how we manage backfilling of the context to keep chats with the LLM more correct in responses.
+  // If a past citation was used to answer the question - that is visible in the history so it logically makes sense
+  // and does not appear to the user that a new response used information that is otherwise irrelevant for a given prompt.
+  // TLDR; reduces GitHub issues for "LLM citing document that has no answer in it" while keep answers highly accurate.
+  contextTexts = [...contextTexts, ...filledSources.contextTexts];
+  sources = [...sources, ...vectorSearchResults.sources];
+
+  // If in query mode and no context chunks are found from search, backfill, or pins -  do not
+  // let the LLM try to hallucinate a response or use general knowledge and exit early
+  if (chatMode === "query" && contextTexts.length === 0) {
+    const textResponse =
+      workspace?.queryRefusalResponse ??
+      "There is no relevant information in this workspace to answer your query.";
+    writeResponseChunk(response, {
+      id: uuid,
+      type: "textResponse",
+      textResponse,
+      sources: [],
+      close: true,
+      error: null,
+    });
+
+    await WorkspaceChats.new({
+      workspaceId: workspace.id,
+      prompt: message,
+      response: {
+        text: textResponse,
+        sources: [],
+        type: chatMode,
+        attachments,
+      },
+      threadId: thread?.id || null,
+      include: false,
+      user,
+    });
+    return;
+  }
+
+  // Compress & Assemble message to ensure prompt passes token limit with room for response
+  // and build system messages based on inputs and history.
+  // Reuse the system prompt from routing pre-fetch when available.
+  const systemPrompt =
+    prefetchedContext?.systemPrompt ??
+    (await chatPrompt(workspace, user, {
+      prompt: updatedMessage,
+      rawHistory,
+    }));
+  const messages = await LLMConnector.compressMessages(
+    {
+      systemPrompt,
+      userPrompt: updatedMessage,
+      contextTexts,
+      chatHistory,
+      attachments,
+    },
+    rawHistory
+  );
+
+  // If streaming is not explicitly enabled for connector
+  // we do regular waiting of a response and send a single chunk.
+  if (LLMConnector.streamingEnabled() !== true) {
+    console.log(
+      `\x1b[31m[STREAMING DISABLED]\x1b[0m Streaming is not available for ${LLMConnector.constructor.name}. Will use regular chat method.`
+    );
+    const { textResponse, metrics: performanceMetrics } =
+      await LLMConnector.getChatCompletion(messages, {
+        temperature: workspace?.openAiTemp ?? LLMConnector.defaultTemp,
+        user: user,
+      });
+
+    completeText = textResponse;
+    metrics = performanceMetrics;
+    writeResponseChunk(response, {
+      uuid,
+      sources,
+      type: "textResponseChunk",
+      textResponse: completeText,
+      close: true,
+      error: false,
+      metrics,
+    });
+  } else {
+    const stream = await LLMConnector.streamGetChatCompletion(messages, {
+      temperature: workspace?.openAiTemp ?? LLMConnector.defaultTemp,
+      user: user,
+    });
+    completeText = await LLMConnector.handleStream(response, stream, {
+      uuid,
+      sources,
+    });
+    metrics = stream.metrics;
+  }
+
+  if (completeText?.length > 0) {
+    const { chat } = await WorkspaceChats.new({
+      workspaceId: workspace.id,
+      prompt: message,
+      response: {
+        text: completeText,
+        sources,
+        type: chatMode,
+        attachments,
+        metrics,
+      },
+      threadId: thread?.id || null,
+      user,
+    });
+
+    writeResponseChunk(response, {
+      uuid,
+      type: "finalizeResponseStream",
+      close: true,
+      error: false,
+      chatId: chat.id,
+      metrics,
+    });
+    return;
+  }
+
+  writeResponseChunk(response, {
+    uuid,
+    type: "finalizeResponseStream",
+    close: true,
+    error: false,
+    metrics,
+  });
+  return;
+}
+
+async function resolveLLMConnector({
+  workspace,
+  message,
+  user,
+  thread,
+  attachments,
+}) {
+  try {
+    const result = await resolveProviderConnector({
+      workspace,
+      prompt: message,
+      user,
+      thread,
+      attachments,
+    });
+    return { ...result, error: null };
+  } catch (routerError) {
+    return {
+      connector: null,
+      routingMetadata: null,
+      prefetchedContext: null,
+      error: `Model router error: ${routerError.message}`,
+    };
+  }
+}
+
+module.exports = {
+  VALID_CHAT_MODE,
+  streamChatWithWorkspace,
+};

package/server/utils/collectorApi/index.js

@@ -0,0 +1,319 @@
+const { EncryptionManager } = require("../EncryptionManager");
+const { Agent } = require("undici");
+
+/**
+ * @typedef {Object} CollectorOptions
+ * @property {string} whisperProvider - The provider to use for whisper, defaults to "local"
+ * @property {string} WhisperModelPref - The model to use for whisper if set.
+ * @property {string} openAiKey - The API key to use for OpenAI interfacing, mostly passed to OAI Whisper provider.
+ * @property {Object} ocr - The OCR options
+ * @property {{allowAnyIp: "true"|null|undefined}} runtimeSettings - The runtime settings that are passed to the collector. Persisted across requests.
+ */
+
+// When running locally will occupy the 0.0.0.0 hostname space but when deployed inside
+// of docker this endpoint is not exposed so it is only on the Docker instances internal network
+// so no additional security is needed on the endpoint directly. Auth is done however by the express
+// middleware prior to leaving the node-side of the application so that is good enough >:)
+class CollectorApi {
+  /** @type {number} - The default collector port */
+  static DEFAULT_COLLECTOR_PORT = 8888;
+
+  /** @type {number} - The maximum timeout for extension requests in milliseconds */
+  extensionRequestTimeout = 15 * 60_000; // 15 minutes
+  /** @type {Agent} - The agent for extension requests */
+  extensionRequestAgent = new Agent({
+    headersTimeout: this.extensionRequestTimeout,
+    bodyTimeout: this.extensionRequestTimeout,
+  });
+
+  /**
+   * Gets the collector port from the environment variables.
+   * If the port is not set, it will fall back to the default port.
+   * If the port is invalid, it will log a warning and return the default port.
+   * @returns {number}
+   */
+  static getCollectorPort() {
+    if (!("COLLECTOR_PORT" in process.env)) return this.DEFAULT_COLLECTOR_PORT;
+    const port = Number(
+      process.env.COLLECTOR_PORT || this.DEFAULT_COLLECTOR_PORT
+    );
+    if (Number.isInteger(port) && port > 0 && port <= 65535) return port;
+
+    console.warn(
+      `Invalid COLLECTOR_PORT "${process.env.COLLECTOR_PORT}". Falling back to ${this.DEFAULT_COLLECTOR_PORT}.`
+    );
+    return this.DEFAULT_COLLECTOR_PORT;
+  }
+
+  constructor() {
+    const { CommunicationKey } = require("../comKey");
+    this.comkey = new CommunicationKey();
+    this.endpoint = `http://0.0.0.0:${CollectorApi.getCollectorPort()}`;
+  }
+
+  log(text, ...args) {
+    console.log(`\x1b[36m[CollectorApi]\x1b[0m ${text}`, ...args);
+  }
+
+  /**
+   * Attach options to the request passed to the collector API
+   * @returns {CollectorOptions}
+   */
+  #attachOptions() {
+    return {
+      whisperProvider: process.env.WHISPER_PROVIDER || "local",
+      WhisperModelPref: process.env.WHISPER_MODEL_PREF,
+      openAiKey: process.env.OPEN_AI_KEY || null,
+      ocr: {
+        langList: process.env.TARGET_OCR_LANG || "eng",
+      },
+      runtimeSettings: {
+        allowAnyIp: process.env.COLLECTOR_ALLOW_ANY_IP ?? "false",
+        browserLaunchArgs: process.env.ANYTHINGLLM_CHROMIUM_ARGS ?? [],
+      },
+    };
+  }
+
+  async online() {
+    return await fetch(this.endpoint)
+      .then((res) => res.ok)
+      .catch(() => false);
+  }
+
+  async acceptedFileTypes() {
+    return await fetch(`${this.endpoint}/accepts`)
+      .then((res) => {
+        if (!res.ok) throw new Error("failed to GET /accepts");
+        return res.json();
+      })
+      .then((res) => res)
+      .catch((e) => {
+        this.log(e.message);
+        return null;
+      });
+  }
+
+  /**
+   * Process a document
+   * - Will append the options and optional metadata to the request body
+   * @param {string} filename - The filename of the document to process
+   * @param {Object} metadata - Optional metadata key:value pairs
+   * @returns {Promise<Object>} - The response from the collector API
+   */
+  async processDocument(filename = "", metadata = {}) {
+    if (!filename) return false;
+
+    const data = JSON.stringify({
+      filename,
+      metadata,
+      options: this.#attachOptions(),
+    });
+
+    return await fetch(`${this.endpoint}/process`, {
+      method: "POST",
+      headers: {
+        "Content-Type": "application/json",
+        "X-Integrity": this.comkey.sign(data),
+        "X-Payload-Signer": this.comkey.encrypt(
+          new EncryptionManager().xPayload
+        ),
+      },
+      body: data,
+      dispatcher: new Agent({ headersTimeout: 600000 }),
+    })
+      .then((res) => {
+        if (!res.ok) throw new Error("Response could not be completed");
+        return res.json();
+      })
+      .then((res) => res)
+      .catch((e) => {
+        this.log(e.message);
+        return { success: false, reason: e.message, documents: [] };
+      });
+  }
+
+  /**
+   * Process a link
+   * - Will append the options to the request body
+   * @param {string} link - The link to process
+   * @param {{[key: string]: string}} scraperHeaders - Custom headers to apply to the web-scraping request URL
+   * @param {[key: string]: string} metadata - Optional metadata to attach to the document
+   * @returns {Promise<Object>} - The response from the collector API
+   */
+  async processLink(link = "", scraperHeaders = {}, metadata = {}) {
+    if (!link) return false;
+
+    const data = JSON.stringify({
+      link,
+      scraperHeaders,
+      options: this.#attachOptions(),
+      metadata: metadata,
+    });
+
+    return await fetch(`${this.endpoint}/process-link`, {
+      method: "POST",
+      headers: {
+        "Content-Type": "application/json",
+        "X-Integrity": this.comkey.sign(data),
+        "X-Payload-Signer": this.comkey.encrypt(
+          new EncryptionManager().xPayload
+        ),
+      },
+      body: data,
+    })
+      .then((res) => {
+        if (!res.ok) throw new Error("Response could not be completed");
+        return res.json();
+      })
+      .then((res) => res)
+      .catch((e) => {
+        this.log(e.message);
+        return { success: false, reason: e.message, documents: [] };
+      });
+  }
+
+  /**
+   * Process raw text as a document for the collector
+   * - Will append the options to the request body
+   * @param {string} textContent - The text to process
+   * @param {[key: string]: string} metadata - The metadata to process
+   * @returns {Promise<Object>} - The response from the collector API
+   */
+  async processRawText(textContent = "", metadata = {}) {
+    const data = JSON.stringify({
+      textContent,
+      metadata,
+      options: this.#attachOptions(),
+    });
+    return await fetch(`${this.endpoint}/process-raw-text`, {
+      method: "POST",
+      headers: {
+        "Content-Type": "application/json",
+        "X-Integrity": this.comkey.sign(data),
+        "X-Payload-Signer": this.comkey.encrypt(
+          new EncryptionManager().xPayload
+        ),
+      },
+      body: data,
+    })
+      .then((res) => {
+        if (!res.ok) throw new Error("Response could not be completed");
+        return res.json();
+      })
+      .then((res) => res)
+      .catch((e) => {
+        this.log(e.message);
+        return { success: false, reason: e.message, documents: [] };
+      });
+  }
+
+  // We will not ever expose the document processor to the frontend API so instead we relay
+  // all requests through the server. You can use this function to directly expose a specific endpoint
+  // on the document processor.
+  async forwardExtensionRequest({ endpoint, method, body }) {
+    const data = typeof body === "string" ? body : JSON.stringify(body);
+    return await fetch(`${this.endpoint}${endpoint}`, {
+      method,
+      body: data,
+      headers: {
+        "Content-Type": "application/json",
+        "X-Integrity": this.comkey.sign(data),
+        "X-Payload-Signer": this.comkey.encrypt(
+          new EncryptionManager().xPayload
+        ),
+      },
+      // Extensions do a lot of work, and may take a while to complete so we need to increase the timeout
+      // substantially so that they do not show a failure to the user early.
+      dispatcher: this.extensionRequestAgent,
+    })
+      .then((res) => {
+        if (!res.ok) throw new Error("Response could not be completed");
+        return res.json();
+      })
+      .then((res) => res)
+      .catch((e) => {
+        this.log(e.message);
+        return { success: false, data: {}, reason: e.message };
+      });
+  }
+
+  /**
+   * Get the content of a link only in a specific format
+   * - Will append the options to the request body
+   * @param {string} link - The link to get the content of
+   * @param {"text"|"html"} captureAs - The format to capture the content as
+   * @returns {Promise<Object>} - The response from the collector API
+   */
+  async getLinkContent(link = "", captureAs = "text") {
+    if (!link) return false;
+
+    const data = JSON.stringify({
+      link,
+      captureAs,
+      options: this.#attachOptions(),
+    });
+    return await fetch(`${this.endpoint}/util/get-link`, {
+      method: "POST",
+      headers: {
+        "Content-Type": "application/json",
+        "X-Integrity": this.comkey.sign(data),
+        "X-Payload-Signer": this.comkey.encrypt(
+          new EncryptionManager().xPayload
+        ),
+      },
+      body: data,
+    })
+      .then((res) => {
+        if (!res.ok) throw new Error("Response could not be completed");
+        return res.json();
+      })
+      .then((res) => res)
+      .catch((e) => {
+        this.log(e.message);
+        return { success: false, content: null };
+      });
+  }
+
+  /**
+   * Parse a document without processing it
+   * - Will append the options to the request body
+   * @param {string} filename - The filename of the document to parse
+   * @param {Object} parseOptions - Additional options for parsing
+   * @param {string} parseOptions.absolutePath - If provided, use this absolute path instead of looking in the hotdir
+   * @returns {Promise<Object>} - The response from the collector API
+   */
+  async parseDocument(filename = "", parseOptions = {}) {
+    if (!filename) return false;
+
+    const data = JSON.stringify({
+      filename,
+      options: {
+        ...this.#attachOptions(),
+        absolutePath: parseOptions.absolutePath || null,
+      },
+    });
+
+    return await fetch(`${this.endpoint}/parse`, {
+      method: "POST",
+      headers: {
+        "Content-Type": "application/json",
+        "X-Integrity": this.comkey.sign(data),
+        "X-Payload-Signer": this.comkey.encrypt(
+          new EncryptionManager().xPayload
+        ),
+      },
+      body: data,
+    })
+      .then((res) => {
+        if (!res.ok) throw new Error("Response could not be completed");
+        return res.json();
+      })
+      .then((res) => res)
+      .catch((e) => {
+        this.log(e.message);
+        return { success: false, reason: e.message, documents: [] };
+      });
+  }
+}
+
+module.exports.CollectorApi = CollectorApi;