totem-llm 0.1.0

package/server/utils/BackgroundWorkers/index.js

@@ -0,0 +1,432 @@
+const path = require("path");
+const Graceful = require("@ladjs/graceful");
+const Bree = require("@mintplex-labs/bree");
+const later = require("@breejs/later");
+const PQueue = require("p-queue").default;
+const setLogger = require("../logger");
+
+// Use UTC time for cron interpretation. This ensures consistent behavior
+// regardless of server timezone (e.g., when running in containers).
+later.date.UTC();
+
+class BackgroundService {
+  name = "BackgroundWorkerService";
+  static _instance = null;
+  documentSyncEnabled = false;
+  memoryExtractionEnabled = false;
+  #root = path.resolve(__dirname, "../../jobs");
+  #scheduledJobTimers = new Map();
+  #scheduledJobQueue = new PQueue({
+    concurrency: Number(process.env.SCHEDULED_JOB_MAX_CONCURRENT) || 1,
+  });
+  // Tracks in-flight worker processes per scheduled jobId so we can kill any
+  // active runs when the job is deleted. Without this, a running worker
+  // outlives the cascade-delete of its scheduled_job_runs row and throws when
+  // it tries to write the result back (prisma.update on a missing row).
+  #scheduledJobWorkers = new Map();
+
+  #alwaysRunJobs = [
+    {
+      name: "cleanup-orphan-documents",
+      timeout: "1m",
+      interval: "12hr",
+    },
+    {
+      name: "cleanup-generated-files",
+      timeout: "5m",
+      interval: "8hr",
+    },
+  ];
+
+  #memoryJobs = [
+    {
+      name: "extract-memories",
+      interval: process.env.MEMORY_EXTRACTION_INTERVAL || "3hr",
+    },
+  ];
+
+  #documentSyncJobs = [
+    // Job for auto-sync of documents
+    // https://github.com/breejs/bree
+    {
+      name: "sync-watched-documents",
+      interval: "1hr",
+    },
+  ];
+
+  constructor() {
+    if (BackgroundService._instance) {
+      this.#log("SINGLETON LOCK: Using existing BackgroundService.");
+      return BackgroundService._instance;
+    }
+
+    this.logger = setLogger();
+    BackgroundService._instance = this;
+  }
+
+  #log(text, ...args) {
+    console.log(`\x1b[36m[${this.name}]\x1b[0m ${text}`, ...args);
+  }
+
+  /**
+   * Returns the root path where job files are located.
+   * Handles the difference between development and production (bundled) environments.
+   * @returns {string}
+   */
+  get jobsRoot() {
+    return this.#root;
+  }
+
+  /**
+   * Wraps the logger so that IPC messages carrying `silent: true` are
+   * suppressed. Bree unconditionally calls `logger.info(message)` for
+   * every IPC message from forked processes, so this is the only
+   * interception point.
+   */
+  #makeBreeLogger() {
+    const base = this.logger;
+    const isSilent = (args) => args.length === 1 && args[0]?.silent === true;
+
+    const wrapped = Object.create(base);
+    wrapped.info = (...args) => {
+      if (!isSilent(args)) base.info(...args);
+    };
+    wrapped.log = (...args) => {
+      if (!isSilent(args)) base.log(...args);
+    };
+    return wrapped;
+  }
+
+  async boot() {
+    const { DocumentSyncQueue } = require("../../models/documentSyncQueue");
+    const { SystemSettings } = require("../../models/systemSettings");
+    const { ScheduledJobRun } = require("../../models/scheduledJobRun");
+
+    this.documentSyncEnabled = await DocumentSyncQueue.enabled();
+    this.memoryExtractionEnabled = await SystemSettings.autoMemoriesEnabled();
+
+    // Mark any orphaned scheduled job runs as failed (server crashed mid-execution)
+    const orphanedCount = await ScheduledJobRun.failOrphanedRuns();
+    if (orphanedCount > 0) {
+      this.#log(
+        `Marked ${orphanedCount} orphaned scheduled job run(s) as failed`
+      );
+    }
+
+    const jobsToRun = this.jobs();
+
+    this.#log("Starting...");
+    this.bree = new Bree({
+      logger: this.#makeBreeLogger(),
+      root: this.#root,
+      jobs: jobsToRun,
+      errorHandler: this.onError,
+      workerMessageHandler: this.onWorkerMessageHandler,
+      runJobsAs: "process",
+    });
+    this.graceful = new Graceful({ brees: [this.bree], logger: this.logger });
+    this.graceful.listen();
+
+    this.bree.start();
+    this.#log(
+      `Service started with ${jobsToRun.length} jobs`,
+      jobsToRun.map((j) => j.name)
+    );
+
+    await this.#bootScheduledJobs();
+  }
+
+  /**
+   * Cleanup scheduled jobs (in-process cron timers + p-queue)
+   */
+  #cleanupScheduledJobs() {
+    for (const [id, timer] of this.#scheduledJobTimers) {
+      timer.clear();
+      this.#scheduledJobTimers.delete(id);
+    }
+    this.#scheduledJobQueue.clear();
+  }
+
+  async stop() {
+    this.#log("Stopping...");
+    this.#cleanupScheduledJobs();
+    if (!!this.graceful && !!this.bree) this.graceful.stopBree(this.bree, 0);
+    this.bree = null;
+    this.graceful = null;
+    this.#log("Service stopped");
+  }
+
+  /** @returns {import("@mintplex-labs/bree").Job[]} */
+  jobs() {
+    const activeJobs = [...this.#alwaysRunJobs];
+    if (this.memoryExtractionEnabled) activeJobs.push(...this.#memoryJobs);
+    if (this.documentSyncEnabled) activeJobs.push(...this.#documentSyncJobs);
+    return activeJobs;
+  }
+
+  /**
+   * Sync the memory extraction job based on current settings.
+   * Called when memory_enabled or memory_auto_extraction changes.
+   * @param {boolean} enabled - The desired state (should the extraction job run?)
+   */
+  async syncMemoryJob(enabled) {
+    if (!this.bree) return;
+
+    const jobName = this.#memoryJobs[0].name;
+    const isCurrentlyRunning = this.bree.config.jobs.some(
+      (j) => j.name === jobName
+    );
+
+    if (enabled && !isCurrentlyRunning) {
+      this.memoryExtractionEnabled = true;
+      await this.bree.add(this.#memoryJobs[0]);
+      await this.bree.start(jobName);
+      this.#log(`Added and started ${jobName} job`);
+    } else if (!enabled && isCurrentlyRunning) {
+      this.memoryExtractionEnabled = false;
+      await this.bree.stop(jobName);
+      await this.bree.remove(jobName);
+      this.#log(`Stopped and removed ${jobName} job`);
+    }
+  }
+
+  onError(error, _workerMetadata) {
+    this.logger.error(`${error.message}`, {
+      service: "bg-worker",
+      origin: error.name,
+    });
+  }
+
+  onWorkerMessageHandler(message, _workerMetadata) {
+    if (message?.silent || message?.message?.silent) return;
+    this.logger.info(`${message.message}`, {
+      service: "bg-worker",
+      origin: message.name,
+    });
+  }
+
+  /**
+   * Spawn a one-off Bree worker process for the given script.
+   * @param {string} scriptPath - Absolute path to the worker JS file
+   * @returns {Promise<{ worker: ChildProcess, jobId: string }>}
+   */
+  async spawnWorker(scriptPath) {
+    if (!this.bree)
+      throw new Error("BackgroundService has not been booted yet");
+
+    const jobId = `${path.basename(scriptPath, ".js")}-${Date.now()}`;
+
+    await this.bree.add({
+      name: jobId,
+      path: scriptPath,
+    });
+
+    await this.bree.run(jobId);
+    const worker = this.bree.workers.get(jobId);
+
+    if (!worker) throw new Error("Failed to get worker reference from Bree");
+
+    return { worker, jobId };
+  }
+
+  /**
+   * Remove a one-off Bree job registration so stale entries don't accumulate.
+   * @param {string} jobId
+   */
+  async removeJob(jobId) {
+    if (!jobId) return;
+    try {
+      if (this.bree) await this.bree.remove(jobId);
+    } catch {
+      /* Job may already be removed */
+    }
+  }
+
+  // ---------------------------------------------------------------
+  // Scheduled Jobs — in-process cron timers + p-queue
+  //
+  // Bree tightly couples scheduling with worker spawning — when a
+  // Bree cron fires, it directly calls run() which immediately
+  // spawns a child process with no way to intercept it. We manage
+  // our own cron timers (via later.setInterval) to decouple
+  // scheduling from execution so we can route jobs through p-queue
+  // for global concurrency control before spawning workers.
+  //
+  // Per-job dedup lives in the database, not in process memory: any
+  // non-terminal row (`queued` or `running`) in scheduled_job_runs means
+  // the job has a run in flight. ScheduledJobRun.start() does the check +
+  // insert atomically and creates the row in `queued` status. The worker
+  // transitions it to `running` once it actually begins executing, so
+  // `startedAt` reflects execution start rather than queue-claim time.
+  // Cron-fired and manually-triggered enqueues use the same rule —
+  // at most one in-flight run per job, regardless of source.
+  // ---------------------------------------------------------------
+
+  /**
+   * Register cron timers for all enabled scheduled jobs on startup.
+   */
+  async #bootScheduledJobs() {
+    const { ScheduledJob } = require("../../models/scheduledJob");
+    const enabledJobs = await ScheduledJob.allEnabled();
+
+    for (const job of enabledJobs) {
+      await ScheduledJob.recomputeNextRunAt(job.id);
+      this.addScheduledJob(job);
+    }
+
+    if (enabledJobs.length > 0) {
+      this.#log(
+        `Registered ${enabledJobs.length} scheduled job(s) (max concurrent: ${this.#scheduledJobQueue.concurrency})`,
+        enabledJobs.map((j) => `${j.name} (${j.schedule})`)
+      );
+    }
+  }
+
+  /**
+   * Register an in-process cron timer for a scheduled job.
+   * When the cron fires, the jobId is enqueued for execution.
+   * @param {object} job - scheduled_jobs DB record
+   */
+  addScheduledJob(job) {
+    this.removeScheduledJob(job.id);
+    const sched = later.parse.cron(job.schedule);
+    const timer = later.setInterval(() => {
+      this.enqueueScheduledJob(job.id);
+    }, sched);
+    this.#scheduledJobTimers.set(job.id, timer);
+  }
+
+  /**
+   * Remove an in-process cron timer for a scheduled job and kill any in-flight
+   * worker processes for it. Killing in-flight workers prevents them from
+   * writing results back to a scheduled_job_runs row that the FK cascade (from
+   * a subsequent ScheduledJob.delete) is about to remove.
+   * @param {number} jobId - scheduled_jobs.id
+   */
+  removeScheduledJob(jobId) {
+    const timer = this.#scheduledJobTimers.get(jobId);
+    if (timer) timer.clear();
+    this.#scheduledJobTimers.delete(jobId);
+
+    const workers = this.#scheduledJobWorkers.get(jobId);
+    if (workers) {
+      for (const worker of workers) {
+        try {
+          worker.kill("SIGTERM");
+        } catch {
+          /* worker may have already exited */
+        }
+      }
+      this.#scheduledJobWorkers.delete(jobId);
+    }
+  }
+
+  /**
+   * Re-sync a scheduled job's cron timer after an update.
+   * Removes the old timer and re-adds if still enabled.
+   * @param {number} jobId - scheduled_jobs.id
+   */
+  async syncScheduledJob(jobId) {
+    const { ScheduledJob } = require("../../models/scheduledJob");
+    this.removeScheduledJob(jobId);
+    const job = await ScheduledJob.get({ id: Number(jobId) });
+    if (job && job.enabled) {
+      this.addScheduledJob(job);
+    }
+  }
+
+  /**
+   * Kill a specific run's worker process. This terminates the worker but does
+   * not update the database — the caller should use ScheduledJobRun.kill()
+   * before or after calling this to mark the run as failed.
+   *
+   * @param {number} jobId - scheduled_jobs.id (parent job)
+   * @param {number} runId - scheduled_job_runs.id (not directly used, but for
+   *   future multi-run support; currently we kill all workers for the jobId)
+   * @returns {boolean} true if a worker was found and killed, false otherwise
+   */
+  killRun(jobId, _runId) {
+    const workers = this.#scheduledJobWorkers.get(Number(jobId));
+    if (!workers || workers.size === 0) return false;
+
+    let killed = false;
+    for (const worker of workers) {
+      try {
+        worker.kill("SIGTERM");
+        killed = true;
+      } catch {
+        /* worker may have already exited */
+      }
+    }
+    return killed;
+  }
+
+  /**
+   * Enqueue a scheduled job for execution. Called by both the cron timer
+   * (in addScheduledJob) and the manual trigger endpoint. ScheduledJobRun.start()
+   * atomically rejects the call if the job already has a run in flight.
+   *
+   * @param {number} jobId - scheduled_jobs.id
+   * @returns {Promise<object|null>} the created run row, or null if skipped
+   *   because a run is already in flight for this job.
+   */
+  async enqueueScheduledJob(jobId) {
+    const { ScheduledJobRun } = require("../../models/scheduledJobRun");
+
+    const run = await ScheduledJobRun.start(jobId);
+    // if start returns null, skip enqueuing, schueduled job already has a run in flight
+    if (!run) return null;
+
+    this.#scheduledJobQueue.add(() =>
+      this.#runScheduledJobWorker(jobId, run.id).catch(async (err) => {
+        this.#log(`Scheduled job ${jobId} failed: ${err.message}`);
+        await ScheduledJobRun.failIfNotTerminal(run.id, err.message);
+      })
+    );
+    return run;
+  }
+
+  /**
+   * Spawn the run-scheduled-job worker for a given run and resolve when it
+   * exits so p-queue can advance. The worker reads its payload from an IPC
+   * `process.on("message", ...)` handler.
+   *
+   * @param {number} jobId
+   * @param {number} runId
+   * @returns {Promise<void>}
+   */
+  async #runScheduledJobWorker(jobId, runId) {
+    const scriptPath = path.resolve(this.jobsRoot, "run-scheduled-job.js");
+    const { worker, jobId: workerId } = await this.spawnWorker(scriptPath);
+
+    if (!this.#scheduledJobWorkers.has(jobId)) {
+      this.#scheduledJobWorkers.set(jobId, new Set());
+    }
+    this.#scheduledJobWorkers.get(jobId).add(worker);
+
+    try {
+      worker.send({ jobId, runId });
+      await new Promise((resolve, reject) => {
+        worker.on("exit", (code, signal) => {
+          // SIGTERM is sent by removeScheduledJob when the job is deleted
+          // mid-run; treat that as a normal exit rather than a worker failure.
+          if (code === 0 || code == null || signal === "SIGTERM") {
+            resolve();
+          } else {
+            reject(new Error(`Worker exited with code ${code}`));
+          }
+        });
+        worker.on("error", reject);
+      });
+    } finally {
+      const workers = this.#scheduledJobWorkers.get(jobId);
+      if (workers) {
+        workers.delete(worker);
+        if (workers.size === 0) this.#scheduledJobWorkers.delete(jobId);
+      }
+      await this.removeJob(workerId).catch(() => {});
+    }
+  }
+}
+
+module.exports.BackgroundService = BackgroundService;

package/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;

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

@@ -0,0 +1,115 @@
+const { toChunks, reportEmbeddingProgress } = require("../../helpers");
+
+class AzureOpenAiEmbedder {
+  constructor() {
+    const { AzureOpenAI } = require("openai");
+    if (!process.env.AZURE_OPENAI_ENDPOINT)
+      throw new Error("No Azure API endpoint was set.");
+    if (!process.env.AZURE_OPENAI_KEY)
+      throw new Error("No Azure API key was set.");
+
+    this.className = "AzureOpenAiEmbedder";
+    this.apiVersion = "2024-12-01-preview";
+    const openai = new AzureOpenAI({
+      apiKey: process.env.AZURE_OPENAI_KEY,
+      endpoint: process.env.AZURE_OPENAI_ENDPOINT,
+      apiVersion: this.apiVersion,
+    });
+
+    // We cannot assume the model fallback since the model is based on the deployment name
+    // and not the model name - so this will throw on embedding if the model is not defined.
+    this.model = process.env.EMBEDDING_MODEL_PREF;
+    this.openai = openai;
+
+    // Limit of how many strings we can process in a single pass to stay with resource or network limits
+    // https://learn.microsoft.com/en-us/azure/ai-services/openai/faq#i-am-trying-to-use-embeddings-and-received-the-error--invalidrequesterror--too-many-inputs--the-max-number-of-inputs-is-1---how-do-i-fix-this-:~:text=consisting%20of%20up%20to%2016%20inputs%20per%20API%20request
+    this.maxConcurrentChunks = 16;
+
+    // https://learn.microsoft.com/en-us/answers/questions/1188074/text-embedding-ada-002-token-context-length
+    this.embeddingMaxChunkLength = 2048;
+  }
+
+  log(text, ...args) {
+    console.log(`\x1b[36m[${this.className}]\x1b[0m ${text}`, ...args);
+  }
+
+  async embedTextInput(textInput) {
+    const result = await this.embedChunks(
+      Array.isArray(textInput) ? textInput : [textInput]
+    );
+    return result?.[0] || [];
+  }
+
+  async embedChunks(textChunks = []) {
+    if (!this.model) throw new Error("No Embedding Model preference defined.");
+
+    this.log(`Embedding ${textChunks.length} chunks...`);
+    // Because there is a limit on how many chunks can be sent at once to Azure OpenAI
+    // we concurrently execute each max batch of text chunks possible.
+    // Refer to constructor maxConcurrentChunks for more info.
+    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,
+            })
+            .then((res) => {
+              chunksProcessed += chunk.length;
+              reportEmbeddingProgress(chunksProcessed, textChunks.length);
+              resolve({ data: res.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 Azure 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(`Azure OpenAI Failed to embed: ${error}`);
+    return data.length > 0 &&
+      data.every((embd) => embd.hasOwnProperty("embedding"))
+      ? data.map((embd) => embd.embedding)
+      : null;
+  }
+}
+
+module.exports = {
+  AzureOpenAiEmbedder,
+};

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

@@ -0,0 +1,91 @@
+const { toChunks, reportEmbeddingProgress } = require("../../helpers");
+
+class CohereEmbedder {
+  constructor() {
+    if (!process.env.COHERE_API_KEY)
+      throw new Error("No Cohere API key was set.");
+
+    const { CohereClient } = require("cohere-ai");
+    const cohere = new CohereClient({
+      token: process.env.COHERE_API_KEY,
+    });
+
+    this.cohere = cohere;
+    this.model = process.env.EMBEDDING_MODEL_PREF || "embed-english-v3.0";
+    this.inputType = "search_document";
+
+    // Limit of how many strings we can process in a single pass to stay with resource or network limits
+    this.maxConcurrentChunks = 96; // Cohere's limit per request is 96
+    this.embeddingMaxChunkLength = 1945; // https://docs.cohere.com/docs/embed-2 - assume a token is roughly 4 letters with some padding
+  }
+
+  async embedTextInput(textInput) {
+    this.inputType = "search_query";
+    const result = await this.embedChunks([textInput]);
+    return result?.[0] || [];
+  }
+
+  async embedChunks(textChunks = []) {
+    const embeddingRequests = [];
+    this.inputType = "search_document";
+    let chunksProcessed = 0;
+
+    for (const chunk of toChunks(textChunks, this.maxConcurrentChunks)) {
+      embeddingRequests.push(
+        new Promise((resolve) => {
+          this.cohere
+            .embed({
+              texts: chunk,
+              model: this.model,
+              inputType: this.inputType,
+            })
+            .then((res) => {
+              chunksProcessed += chunk.length;
+              reportEmbeddingProgress(chunksProcessed, textChunks.length);
+              resolve({ data: res.embeddings, 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) => {
+      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(`Cohere Failed to embed: ${error}`);
+
+    return data.length > 0 ? data : null;
+  }
+}
+
+module.exports = {
+  CohereEmbedder,
+};