totem-llm 0.1.0

package/server/utils/EmbeddingEngines/localAi/index.js

@@ -0,0 +1,122 @@
+const {
+  toChunks,
+  maximumChunkLength,
+  reportEmbeddingProgress,
+} = require("../../helpers");
+
+class LocalAiEmbedder {
+  constructor() {
+    if (!process.env.EMBEDDING_BASE_PATH)
+      throw new Error("No embedding base path was set.");
+    if (!process.env.EMBEDDING_MODEL_PREF)
+      throw new Error("No embedding model was set.");
+
+    this.className = "LocalAiEmbedder";
+    const { OpenAI: OpenAIApi } = require("openai");
+    this.model = process.env.EMBEDDING_MODEL_PREF;
+    this.openai = new OpenAIApi({
+      baseURL: process.env.EMBEDDING_BASE_PATH,
+      apiKey: process.env.LOCAL_AI_API_KEY ?? null,
+    });
+
+    // Limit of how many strings we can process in a single pass to stay with resource or network limits
+    this.maxConcurrentChunks = 50;
+    this.embeddingMaxChunkLength = maximumChunkLength();
+
+    this.log(
+      `Initialized with ${this.model} - Max Size: ${this.embeddingMaxChunkLength}` +
+        (this.outputDimensions
+          ? ` - Output Dimensions: ${this.outputDimensions}`
+          : " Assuming default output dimensions")
+    );
+  }
+
+  log(text, ...args) {
+    console.log(`\x1b[36m[${this.className}]\x1b[0m ${text}`, ...args);
+  }
+
+  get outputDimensions() {
+    if (
+      process.env.EMBEDDING_OUTPUT_DIMENSIONS &&
+      !isNaN(process.env.EMBEDDING_OUTPUT_DIMENSIONS) &&
+      process.env.EMBEDDING_OUTPUT_DIMENSIONS > 0
+    )
+      return parseInt(process.env.EMBEDDING_OUTPUT_DIMENSIONS);
+    return null;
+  }
+
+  async embedTextInput(textInput) {
+    const result = await this.embedChunks(
+      Array.isArray(textInput) ? textInput : [textInput]
+    );
+    return result?.[0] || [];
+  }
+
+  async embedChunks(textChunks = []) {
+    const embeddingRequests = [];
+    let chunksProcessed = 0;
+    for (const chunk of toChunks(textChunks, this.maxConcurrentChunks)) {
+      embeddingRequests.push(
+        new Promise((resolve) => {
+          this.openai.embeddings
+            .create({
+              model: this.model,
+              input: chunk,
+              dimensions: this.outputDimensions,
+            })
+            .then((result) => {
+              chunksProcessed += chunk.length;
+              reportEmbeddingProgress(chunksProcessed, textChunks.length);
+              resolve({ data: result?.data, error: null });
+            })
+            .catch((e) => {
+              chunksProcessed += chunk.length;
+              reportEmbeddingProgress(chunksProcessed, textChunks.length);
+              e.type =
+                e?.response?.data?.error?.code ||
+                e?.response?.status ||
+                "failed_to_embed";
+              e.message = e?.response?.data?.error?.message || e.message;
+              resolve({ data: [], error: e });
+            });
+        })
+      );
+    }
+
+    const { data = [], error = null } = await Promise.all(
+      embeddingRequests
+    ).then((results) => {
+      // If any errors were returned from LocalAI abort the entire sequence because the embeddings
+      // will be incomplete.
+      const errors = results
+        .filter((res) => !!res.error)
+        .map((res) => res.error)
+        .flat();
+      if (errors.length > 0) {
+        let uniqueErrors = new Set();
+        errors.map((error) =>
+          uniqueErrors.add(`[${error.type}]: ${error.message}`)
+        );
+
+        return {
+          data: [],
+          error: Array.from(uniqueErrors).join(", "),
+        };
+      }
+      return {
+        data: results.map((res) => res?.data || []).flat(),
+        error: null,
+      };
+    });
+
+    if (!!error) throw new Error(`LocalAI Failed to embed: ${error}`);
+    return data.length > 0 &&
+      data.every((embd) => embd.hasOwnProperty("embedding"))
+      ? data.map((embd) => embd.embedding)
+      : null;
+  }
+}
+
+module.exports = {
+  LocalAiEmbedder,
+};

package/server/utils/EmbeddingEngines/mistral/index.js

@@ -0,0 +1,40 @@
+class MistralEmbedder {
+  constructor() {
+    if (!process.env.MISTRAL_API_KEY)
+      throw new Error("No Mistral API key was set.");
+
+    const { OpenAI: OpenAIApi } = require("openai");
+    this.openai = new OpenAIApi({
+      baseURL: "https://api.mistral.ai/v1",
+      apiKey: process.env.MISTRAL_API_KEY ?? null,
+    });
+    this.model = process.env.EMBEDDING_MODEL_PREF || "mistral-embed";
+  }
+
+  async embedTextInput(textInput) {
+    const result = await this.embedChunks(
+      Array.isArray(textInput) ? textInput : [textInput]
+    );
+    return result?.[0] || [];
+  }
+
+  async embedChunks(textChunks = []) {
+    try {
+      const response = await this.openai.embeddings.create({
+        model: this.model,
+        input: textChunks,
+      });
+      const embeddings = response?.data?.map((emb) => emb.embedding) || [];
+      if (embeddings.length === 0)
+        throw new Error("Mistral returned empty embeddings for batch");
+      return embeddings;
+    } catch (error) {
+      console.error("Failed to get embeddings from Mistral.", error.message);
+      throw new Error(`Mistral Failed to embed: ${error.message}`);
+    }
+  }
+}
+
+module.exports = {
+  MistralEmbedder,
+};

package/server/utils/EmbeddingEngines/native/constants.js

@@ -0,0 +1,63 @@
+const SUPPORTED_NATIVE_EMBEDDING_MODELS = {
+  "Xenova/all-MiniLM-L6-v2": {
+    maxConcurrentChunks: 25,
+    // Right now, this is NOT the token length, and is instead the number of characters
+    // that can be processed in a single pass. So we override to 1,000 characters.
+    // roughtly the max number of tokens assuming 2 characters per token. (undershooting)
+    // embeddingMaxChunkLength: 512, (from the model card)
+    embeddingMaxChunkLength: 1_000,
+    chunkPrefix: "",
+    queryPrefix: "",
+    apiInfo: {
+      id: "Xenova/all-MiniLM-L6-v2",
+      name: "all-MiniLM-L6-v2",
+      description:
+        "A lightweight and fast model for embedding text. The default model for AnythingLLM.",
+      lang: "English",
+      size: "23MB",
+      modelCard: "https://huggingface.co/Xenova/all-MiniLM-L6-v2",
+    },
+  },
+  "Xenova/nomic-embed-text-v1": {
+    maxConcurrentChunks: 5,
+    // Right now, this is NOT the token length, and is instead the number of characters
+    // that can be processed in a single pass. So we override to 16,000 characters.
+    // roughtly the max number of tokens assuming 2 characters per token. (undershooting)
+    // embeddingMaxChunkLength: 8192, (from the model card)
+    embeddingMaxChunkLength: 16_000,
+    chunkPrefix: "search_document: ",
+    queryPrefix: "search_query: ",
+    apiInfo: {
+      id: "Xenova/nomic-embed-text-v1",
+      name: "nomic-embed-text-v1",
+      description:
+        "A high-performing open embedding model with a large token context window. Requires more processing power and memory.",
+      lang: "English",
+      size: "139MB",
+      modelCard: "https://huggingface.co/Xenova/nomic-embed-text-v1",
+    },
+  },
+  "MintplexLabs/multilingual-e5-small": {
+    maxConcurrentChunks: 5,
+    // Right now, this is NOT the token length, and is instead the number of characters
+    // that can be processed in a single pass. So we override to 1,000 characters.
+    // roughtly the max number of tokens assuming 2 characters per token. (undershooting)
+    // embeddingMaxChunkLength: 512, (from the model card)
+    embeddingMaxChunkLength: 1_000,
+    chunkPrefix: "passage: ",
+    queryPrefix: "query: ",
+    apiInfo: {
+      id: "MintplexLabs/multilingual-e5-small",
+      name: "multilingual-e5-small",
+      description:
+        "A larger multilingual embedding model that supports 100+ languages. Requires more processing power and memory.",
+      lang: "100+ languages",
+      size: "487MB",
+      modelCard: "https://huggingface.co/intfloat/multilingual-e5-small",
+    },
+  },
+};
+
+module.exports = {
+  SUPPORTED_NATIVE_EMBEDDING_MODELS,
+};

package/server/utils/EmbeddingEngines/native/index.js

@@ -0,0 +1,290 @@
+const path = require("path");
+const fs = require("fs");
+const { toChunks, reportEmbeddingProgress } = require("../../helpers");
+const { v4 } = require("uuid");
+const { SUPPORTED_NATIVE_EMBEDDING_MODELS } = require("./constants");
+
+class NativeEmbedder {
+  static defaultModel = "Xenova/all-MiniLM-L6-v2";
+
+  /**
+   * Supported embedding models for native.
+   * @type {Record<string, {
+   *   chunkPrefix: string;
+   *   queryPrefix: string;
+   *   apiInfo: {
+   *     id: string;
+   *     name: string;
+   *     description: string;
+   *     lang: string;
+   *     size: string;
+   *     modelCard: string;
+   *   };
+   * }>}
+   */
+  static supportedModels = SUPPORTED_NATIVE_EMBEDDING_MODELS;
+
+  // This is a folder that Mintplex Labs hosts for those who cannot capture the HF model download
+  // endpoint for various reasons. This endpoint is not guaranteed to be active or maintained
+  // and may go offline at any time at Mintplex Labs's discretion.
+  #fallbackHost = "https://cdn.anythingllm.com/support/models/";
+
+  constructor() {
+    this.className = "NativeEmbedder";
+    this.model = this.getEmbeddingModel();
+    this.modelInfo = this.getEmbedderInfo();
+    this.cacheDir = path.resolve(
+      process.env.STORAGE_DIR
+        ? path.resolve(process.env.STORAGE_DIR, `models`)
+        : path.resolve(__dirname, `../../../storage/models`)
+    );
+    this.modelPath = path.resolve(this.cacheDir, ...this.model.split("/"));
+    this.modelDownloaded = fs.existsSync(this.modelPath);
+
+    // Limit of how many strings we can process in a single pass to stay with resource or network limits
+    this.maxConcurrentChunks = this.modelInfo.maxConcurrentChunks;
+    this.embeddingMaxChunkLength = this.modelInfo.embeddingMaxChunkLength;
+
+    // Make directory when it does not exist in existing installations
+    if (!fs.existsSync(this.cacheDir)) fs.mkdirSync(this.cacheDir);
+    this.log(`Initialized ${this.model}`);
+  }
+
+  log(text, ...args) {
+    console.log(`\x1b[36m[${this.className}]\x1b[0m ${text}`, ...args);
+  }
+
+  /**
+   * Get the selected model from the environment variable.
+   * @returns {string}
+   */
+  static _getEmbeddingModel() {
+    const envModel =
+      process.env.EMBEDDING_MODEL_PREF ?? NativeEmbedder.defaultModel;
+    if (NativeEmbedder.supportedModels?.[envModel]) return envModel;
+    return NativeEmbedder.defaultModel;
+  }
+
+  get embeddingPrefix() {
+    return NativeEmbedder.supportedModels[this.model]?.chunkPrefix || "";
+  }
+
+  get queryPrefix() {
+    return NativeEmbedder.supportedModels[this.model]?.queryPrefix || "";
+  }
+
+  /**
+   * Get the available models in an API response format
+   * we can use to populate the frontend dropdown.
+   * @returns {{id: string, name: string, description: string, lang: string, size: string, modelCard: string}[]}
+   */
+  static availableModels() {
+    return Object.values(NativeEmbedder.supportedModels).map(
+      (model) => model.apiInfo
+    );
+  }
+
+  /**
+   * Get the embedding model to use.
+   * We only support a few models and will default to the default model if the environment variable is not set or not supported.
+   *
+   * Why only a few? Because we need to mirror them on the CDN so non-US users can download them.
+   * eg: "Xenova/all-MiniLM-L6-v2"
+   * eg: "Xenova/nomic-embed-text-v1"
+   * @returns {string}
+   */
+  getEmbeddingModel() {
+    const envModel =
+      process.env.EMBEDDING_MODEL_PREF ?? NativeEmbedder.defaultModel;
+    if (NativeEmbedder.supportedModels?.[envModel]) return envModel;
+    return NativeEmbedder.defaultModel;
+  }
+
+  /**
+   * Get the embedding model info.
+   *
+   * Will always fallback to the default model if the model is not supported.
+   * @returns {Object}
+   */
+  getEmbedderInfo() {
+    const model = this.getEmbeddingModel();
+    return NativeEmbedder.supportedModels[model];
+  }
+
+  #tempfilePath() {
+    const filename = `${v4()}.tmp`;
+    const tmpPath = process.env.STORAGE_DIR
+      ? path.resolve(process.env.STORAGE_DIR, "tmp")
+      : path.resolve(__dirname, `../../../storage/tmp`);
+    if (!fs.existsSync(tmpPath)) fs.mkdirSync(tmpPath, { recursive: true });
+    return path.resolve(tmpPath, filename);
+  }
+
+  async #writeToTempfile(filePath, data) {
+    try {
+      await fs.promises.appendFile(filePath, data, { encoding: "utf8" });
+    } catch (e) {
+      console.error(`Error writing to tempfile: ${e}`);
+    }
+  }
+
+  async #fetchWithHost(hostOverride = null) {
+    try {
+      // Convert ESM to CommonJS via import so we can load this library.
+      const pipeline = (...args) =>
+        import("@xenova/transformers").then(({ pipeline, env }) => {
+          if (!this.modelDownloaded) {
+            // if model is not downloaded, we will log where we are fetching from.
+            if (hostOverride) {
+              env.remoteHost = hostOverride;
+              env.remotePathTemplate = "{model}/"; // Our S3 fallback url does not support revision File structure.
+            }
+            this.log(`Downloading ${this.model} from ${env.remoteHost}`);
+          }
+          return pipeline(...args);
+        });
+      return {
+        pipeline: await pipeline("feature-extraction", this.model, {
+          cache_dir: this.cacheDir,
+          ...(!this.modelDownloaded
+            ? {
+                // Show download progress if we need to download any files
+                progress_callback: (data) => {
+                  if (!data.hasOwnProperty("progress")) return;
+                  console.log(
+                    `\x1b[36m[NativeEmbedder - Downloading model]\x1b[0m ${
+                      data.file
+                    } ${~~data?.progress}%`
+                  );
+                },
+              }
+            : {}),
+        }),
+        retry: false,
+        error: null,
+      };
+    } catch (error) {
+      return {
+        pipeline: null,
+        retry: hostOverride === null ? this.#fallbackHost : false,
+        error,
+      };
+    }
+  }
+
+  // This function will do a single fallback attempt (not recursive on purpose) to try to grab the embedder model on first embed
+  // since at time, some clients cannot properly download the model from HF servers due to a number of reasons (IP, VPN, etc).
+  // Given this model is critical and nobody reads the GitHub issues before submitting the bug, we get the same bug
+  // report 20 times a day: https://github.com/Mintplex-Labs/anything-llm/issues/821
+  // So to attempt to monkey-patch this we have a single fallback URL to help alleviate duplicate bug reports.
+  async embedderClient() {
+    if (!this.modelDownloaded)
+      this.log(
+        "The native embedding model has never been run and will be downloaded right now. Subsequent runs will be faster. (~23MB)"
+      );
+
+    let fetchResponse = await this.#fetchWithHost();
+    if (fetchResponse.pipeline !== null) {
+      this.modelDownloaded = true;
+      return fetchResponse.pipeline;
+    }
+
+    this.log(
+      `Failed to download model from primary URL. Using fallback ${fetchResponse.retry}`
+    );
+    if (!!fetchResponse.retry)
+      fetchResponse = await this.#fetchWithHost(fetchResponse.retry);
+    if (fetchResponse.pipeline !== null) {
+      this.modelDownloaded = true;
+      return fetchResponse.pipeline;
+    }
+
+    throw fetchResponse.error;
+  }
+
+  /**
+   * Apply the query prefix to the text input if it is required by the model.
+   * eg: nomic-embed-text-v1 requires a query prefix for embedding/searching.
+   * @param {string|string[]} textInput - The text to embed.
+   * @returns {string|string[]} The text with the prefix applied.
+   */
+  #applyQueryPrefix(textInput) {
+    if (!this.queryPrefix) return textInput;
+    if (Array.isArray(textInput))
+      textInput = textInput.map((text) => `${this.queryPrefix}${text}`);
+    else textInput = `${this.queryPrefix}${textInput}`;
+    return textInput;
+  }
+
+  /**
+   * Embed a single text input.
+   * @param {string|string[]} textInput - The text to embed.
+   * @returns {Promise<Array<number>>} The embedded text.
+   */
+  async embedTextInput(textInput) {
+    textInput = this.#applyQueryPrefix(textInput);
+    const result = await this.embedChunks(
+      Array.isArray(textInput) ? textInput : [textInput]
+    );
+    return result?.[0] || [];
+  }
+
+  // If you are thinking you want to edit this function - you probably don't.
+  // This process was benchmarked heavily on a t3.small (2GB RAM 1vCPU)
+  // and without careful memory management for the V8 garbage collector
+  // this function will likely result in an OOM on any resource-constrained deployment.
+  // To help manage very large documents we run a concurrent write-log each iteration
+  // to keep the embedding result out of memory. The `maxConcurrentChunk` is set to 25,
+  // as 50 seems to overflow no matter what. Given the above, memory use hovers around ~30%
+  // during a very large document (>100K words) but can spike up to 70% before gc.
+  // This seems repeatable for all document sizes.
+  // While this does take a while, it is zero set up and is 100% free and on-instance.
+  // It still may crash depending on other elements at play - so no promises it works under all conditions.
+  async embedChunks(textChunks = []) {
+    const tmpFilePath = this.#tempfilePath();
+    const chunks = toChunks(textChunks, this.maxConcurrentChunks);
+    const chunkLen = chunks.length;
+    const totalChunks = textChunks.length;
+
+    for (let [idx, chunk] of chunks.entries()) {
+      if (idx === 0) await this.#writeToTempfile(tmpFilePath, "[");
+      let data;
+      let pipeline = await this.embedderClient();
+      let output = await pipeline(chunk, {
+        pooling: "mean",
+        normalize: true,
+      });
+
+      if (output.length === 0) {
+        pipeline = null;
+        output = null;
+        data = null;
+        continue;
+      }
+
+      data = JSON.stringify(output.tolist());
+      await this.#writeToTempfile(tmpFilePath, data);
+      this.log(`Embedded Chunk Group ${idx + 1} of ${chunkLen}`);
+      if (chunkLen - 1 !== idx) await this.#writeToTempfile(tmpFilePath, ",");
+      if (chunkLen - 1 === idx) await this.#writeToTempfile(tmpFilePath, "]");
+
+      reportEmbeddingProgress(
+        Math.min((idx + 1) * this.maxConcurrentChunks, totalChunks),
+        totalChunks
+      );
+      pipeline = null;
+      output = null;
+      data = null;
+    }
+
+    const embeddingResults = JSON.parse(
+      fs.readFileSync(tmpFilePath, { encoding: "utf-8" })
+    );
+    fs.rmSync(tmpFilePath, { force: true });
+    return embeddingResults.length > 0 ? embeddingResults.flat() : null;
+  }
+}
+
+module.exports = {
+  NativeEmbedder,
+};

package/server/utils/EmbeddingEngines/ollama/index.js

@@ -0,0 +1,136 @@
+const {
+  maximumChunkLength,
+  reportEmbeddingProgress,
+} = require("../../helpers");
+const { Ollama } = require("ollama");
+const { OllamaAILLM } = require("../../AiProviders/ollama");
+
+class OllamaEmbedder {
+  constructor() {
+    if (!process.env.EMBEDDING_BASE_PATH)
+      throw new Error("No embedding base path was set.");
+    if (!process.env.EMBEDDING_MODEL_PREF)
+      throw new Error("No embedding model was set.");
+
+    this.className = "OllamaEmbedder";
+    this.basePath = process.env.EMBEDDING_BASE_PATH;
+    this.model = process.env.EMBEDDING_MODEL_PREF;
+    this.maxConcurrentChunks = process.env.OLLAMA_EMBEDDING_BATCH_SIZE
+      ? Number(process.env.OLLAMA_EMBEDDING_BATCH_SIZE)
+      : 1;
+    this.embeddingMaxChunkLength = maximumChunkLength();
+    this.authToken = process.env.OLLAMA_AUTH_TOKEN;
+
+    const headers = this.authToken
+      ? { Authorization: `Bearer ${this.authToken}` }
+      : {};
+    this.client = new Ollama({
+      host: this.basePath,
+      headers,
+      fetch: OllamaAILLM.applyOllamaFetch(),
+    });
+    this.log(
+      `initialized with model ${this.model} at ${this.basePath}. Batch size: ${this.maxConcurrentChunks}, num_ctx: ${this.embeddingMaxChunkLength}`
+    );
+  }
+
+  log(text, ...args) {
+    console.log(`\x1b[36m[${this.className}]\x1b[0m ${text}`, ...args);
+  }
+
+  /**
+   * Checks if the Ollama service is alive by pinging the base path.
+   * @returns {Promise<boolean>} - A promise that resolves to true if the service is alive, false otherwise.
+   */
+  async #isAlive() {
+    return await fetch(this.basePath)
+      .then((res) => res.ok)
+      .catch((e) => {
+        this.log(e.message);
+        return false;
+      });
+  }
+
+  async embedTextInput(textInput) {
+    const result = await this.embedChunks(
+      Array.isArray(textInput) ? textInput : [textInput]
+    );
+    return result?.[0] || [];
+  }
+
+  /**
+   * This function takes an array of text chunks and embeds them using the Ollama API.
+   * Chunks are processed in batches based on the maxConcurrentChunks setting to balance
+   * resource usage on the Ollama endpoint.
+   *
+   * We will use the num_ctx option to set the maximum context window to the max chunk length defined by the user in the settings
+   * so that the maximum context window is used and content is not truncated.
+   *
+   * We also assume the default keep alive option. This could cause issues with models being unloaded and reloaded
+   * on low memory machines, but that is simply a user-end issue we cannot control. If the LLM and embedder are
+   * constantly being loaded and unloaded, the user should use another LLM or Embedder to avoid this issue.
+   * @param {string[]} textChunks - An array of text chunks to embed.
+   * @returns {Promise<Array<number[]>>} - A promise that resolves to an array of embeddings.
+   */
+  async embedChunks(textChunks = []) {
+    if (!(await this.#isAlive()))
+      throw new Error(
+        `Ollama service could not be reached. Is Ollama running?`
+      );
+    this.log(
+      `Embedding ${textChunks.length} chunks of text with ${this.model} in batches of ${this.maxConcurrentChunks}.`
+    );
+
+    let data = [];
+    let error = null;
+
+    // Process chunks in batches based on maxConcurrentChunks
+    const totalBatches = Math.ceil(
+      textChunks.length / this.maxConcurrentChunks
+    );
+    let currentBatch = 0;
+
+    for (let i = 0; i < textChunks.length; i += this.maxConcurrentChunks) {
+      const batch = textChunks.slice(i, i + this.maxConcurrentChunks);
+      currentBatch++;
+
+      try {
+        // Use input param instead of prompt param to support batch processing
+        const res = await this.client.embed({
+          model: this.model,
+          input: batch,
+          options: {
+            // Always set the num_ctx to the max chunk length defined by the user in the settings
+            // so that the maximum context window is used and content is not truncated.
+            num_ctx: this.embeddingMaxChunkLength,
+          },
+        });
+
+        const { embeddings } = res;
+        if (!Array.isArray(embeddings) || embeddings.length === 0)
+          throw new Error("Ollama returned empty embeddings for batch!");
+
+        // Using prompt param in embed() would return a single embedding (number[])
+        // but input param returns an array of embeddings (number[][]) for batch processing.
+        // This is why we spread the embeddings array into the data array.
+        data.push(...embeddings);
+        reportEmbeddingProgress(data.length, textChunks.length);
+        this.log(
+          `Batch ${currentBatch}/${totalBatches}: Embedded ${embeddings.length} chunks. Total: ${data.length}/${textChunks.length}`
+        );
+      } catch (err) {
+        this.log(err.message);
+        error = err.message;
+        data = [];
+        break;
+      }
+    }
+
+    if (!!error) throw new Error(`Ollama Failed to embed: ${error}`);
+    return data.length > 0 ? data : null;
+  }
+}
+
+module.exports = {
+  OllamaEmbedder,
+};