totem-llm 0.1.0

package/server/models/workspaceParsedFiles.js

@@ -0,0 +1,246 @@
+const prisma = require("../utils/prisma");
+const { EventLogs } = require("./eventLogs");
+const { Document } = require("./documents");
+const { documentsPath, directUploadsPath } = require("../utils/files");
+const { safeJsonParse } = require("../utils/http");
+const fs = require("fs");
+const path = require("path");
+
+const WorkspaceParsedFiles = {
+  create: async function ({
+    filename,
+    workspaceId,
+    userId = null,
+    threadId = null,
+    metadata = null,
+    tokenCountEstimate = 0,
+  }) {
+    try {
+      const file = await prisma.workspace_parsed_files.create({
+        data: {
+          filename,
+          workspaceId: parseInt(workspaceId),
+          userId: userId ? parseInt(userId) : null,
+          threadId: threadId ? parseInt(threadId) : null,
+          metadata,
+          tokenCountEstimate,
+        },
+      });
+
+      await EventLogs.logEvent(
+        "workspace_file_uploaded",
+        {
+          filename,
+          workspaceId,
+        },
+        userId
+      );
+
+      return { file, error: null };
+    } catch (error) {
+      console.error("FAILED TO CREATE PARSED FILE RECORD.", error.message);
+      return { file: null, error: error.message };
+    }
+  },
+
+  /**
+   * Gets a parsed file by its ID or a clause.
+   * @param {object} clause - The clause to filter the parsed files.
+   * @returns {Promise<import("@prisma/client").workspace_parsed_files | null>} The parsed file.
+   */
+  get: async function (clause = {}) {
+    try {
+      const file = await prisma.workspace_parsed_files.findFirst({
+        where: clause,
+      });
+      return file;
+    } catch (error) {
+      console.error(error.message);
+      return null;
+    }
+  },
+
+  where: async function (
+    clause = {},
+    limit = null,
+    orderBy = null,
+    select = null
+  ) {
+    try {
+      const files = await prisma.workspace_parsed_files.findMany({
+        where: clause,
+        ...(limit !== null ? { take: limit } : {}),
+        ...(orderBy !== null ? { orderBy } : {}),
+        ...(select !== null ? { select } : {}),
+      });
+      return files;
+    } catch (error) {
+      console.error(error.message);
+      return [];
+    }
+  },
+
+  delete: async function (clause = {}) {
+    try {
+      const result = await prisma.workspace_parsed_files.deleteMany({
+        where: clause,
+      });
+      return result.count > 0;
+    } catch (error) {
+      console.error(error.message);
+      return false;
+    }
+  },
+
+  totalTokenCount: async function (clause = {}) {
+    const { _sum } = await prisma.workspace_parsed_files.aggregate({
+      where: clause,
+      _sum: { tokenCountEstimate: true },
+    });
+    return _sum.tokenCountEstimate || 0;
+  },
+
+  /**
+   * Moves a parsed file to the documents and embeds it.
+   * @param {import("@prisma/client").users | null} user - The user performing the operation.
+   * @param {number} fileId - The ID of the parsed file.
+   * @param {import("@prisma/client").workspaces} workspace - The workspace the file belongs to.
+   * @returns {Promise<{ success: boolean, error: string | null, document: import("@prisma/client").workspace_documents | null }>} The result of the operation.
+   */
+  moveToDocumentsAndEmbed: async function (user = null, fileId, workspace) {
+    try {
+      const parsedFile = await this.get({
+        id: parseInt(fileId),
+        ...(user ? { userId: user.id } : {}),
+        workspaceId: workspace.id,
+      });
+      if (!parsedFile) throw new Error("File not found");
+
+      // Get file location from metadata
+      const metadata = safeJsonParse(parsedFile.metadata, {});
+      const location = metadata.location;
+      if (!location) throw new Error("No file location in metadata");
+
+      // Get file from metadata location
+      const sourceFile = path.join(directUploadsPath, path.basename(location));
+      if (!fs.existsSync(sourceFile)) throw new Error("Source file not found");
+
+      // Move to custom-documents
+      const customDocsPath = path.join(documentsPath, "custom-documents");
+      if (!fs.existsSync(customDocsPath))
+        fs.mkdirSync(customDocsPath, { recursive: true });
+
+      // Copy the file to custom-documents
+      const targetPath = path.join(customDocsPath, path.basename(location));
+      fs.copyFileSync(sourceFile, targetPath);
+      fs.unlinkSync(sourceFile);
+
+      const {
+        failedToEmbed = [],
+        errors = [],
+        embedded = [],
+      } = await Document.addDocuments(
+        workspace,
+        [`custom-documents/${path.basename(location)}`],
+        parsedFile.userId
+      );
+
+      if (failedToEmbed.length > 0)
+        throw new Error(errors[0] || "Failed to embed document");
+
+      const document = await Document.get({
+        workspaceId: workspace.id,
+        docpath: embedded[0],
+      });
+      return { success: true, error: null, document };
+    } catch (error) {
+      console.error("Failed to move and embed file:", error);
+      return { success: false, error: error.message, document: null };
+    } finally {
+      // Always delete the file after processing
+      await this.delete({ id: parseInt(fileId) });
+    }
+  },
+
+  getContextMetadataAndLimits: async function (
+    workspace,
+    thread = null,
+    user = null
+  ) {
+    try {
+      if (!workspace) throw new Error("Workspace is required");
+      const files = await this.where({
+        workspaceId: workspace.id,
+        threadId: thread?.id || null,
+        ...(user ? { userId: user.id } : {}),
+      });
+
+      const results = [];
+      let totalTokens = 0;
+
+      for (const file of files) {
+        const metadata = safeJsonParse(file.metadata, {});
+        totalTokens += file.tokenCountEstimate || 0;
+        results.push({
+          id: file.id,
+          title: metadata.title || metadata.location,
+          location: metadata.location,
+          token_count_estimate: file.tokenCountEstimate,
+        });
+      }
+
+      return {
+        files: results,
+        contextWindow: workspace.contextWindow,
+        currentContextTokenCount: totalTokens,
+      };
+    } catch (error) {
+      console.error("Failed to get context metadata:", error);
+      return {
+        files: [],
+        contextWindow: Infinity,
+        currentContextTokenCount: 0,
+      };
+    }
+  },
+
+  getContextFiles: async function (workspace, thread = null, user = null) {
+    try {
+      const files = await this.where({
+        workspaceId: workspace.id,
+        threadId: thread?.id || null,
+        ...(user ? { userId: user.id } : {}),
+      });
+
+      const results = [];
+      for (const file of files) {
+        const metadata = safeJsonParse(file.metadata, {});
+        const location = metadata.location;
+        if (!location) continue;
+
+        const sourceFile = path.join(
+          directUploadsPath,
+          path.basename(location)
+        );
+        if (!fs.existsSync(sourceFile)) continue;
+
+        const content = fs.readFileSync(sourceFile, "utf-8");
+        const data = safeJsonParse(content, null);
+        if (!data?.pageContent) continue;
+
+        results.push({
+          pageContent: data.pageContent,
+          token_count_estimate: file.tokenCountEstimate,
+          ...metadata,
+        });
+      }
+
+      return results;
+    } catch (error) {
+      console.error("Failed to get context files:", error);
+      return [];
+    }
+  },
+};
+
+module.exports = { WorkspaceParsedFiles };

package/server/models/workspaceThread.js

@@ -0,0 +1,151 @@
+const prisma = require("../utils/prisma");
+const slugifyModule = require("slugify");
+const { v4: uuidv4 } = require("uuid");
+const truncate = require("truncate");
+
+const WorkspaceThread = {
+  defaultName: "Thread",
+  writable: ["name"],
+
+  /**
+   * The default Slugify module requires some additional mapping to prevent downstream issues
+   * if the user is able to define a slug externally. We have to block non-escapable URL chars
+   * so that is the slug is rendered it doesn't break the URL or UI when visited.
+   * @param  {...any} args - slugify args for npm package.
+   * @returns {string}
+   */
+  slugify: function (...args) {
+    slugifyModule.extend({
+      "+": " plus ",
+      "!": " bang ",
+      "@": " at ",
+      "*": " splat ",
+      ".": " dot ",
+      ":": "",
+      "~": "",
+      "(": "",
+      ")": "",
+      "'": "",
+      '"': "",
+      "|": "",
+    });
+    return slugifyModule(...args);
+  },
+
+  new: async function (workspace, userId = null, data = {}) {
+    try {
+      const thread = await prisma.workspace_threads.create({
+        data: {
+          name: data.name ? String(data.name) : this.defaultName,
+          slug: data.slug
+            ? this.slugify(data.slug, { lowercase: true })
+            : uuidv4(),
+          user_id: userId ? Number(userId) : null,
+          workspace_id: workspace.id,
+        },
+      });
+
+      return { thread, message: null };
+    } catch (error) {
+      console.error(error.message);
+      return { thread: null, message: error.message };
+    }
+  },
+
+  update: async function (prevThread = null, data = {}) {
+    if (!prevThread) throw new Error("No thread id provided for update");
+
+    const validData = {};
+    Object.entries(data).forEach(([key, value]) => {
+      if (!this.writable.includes(key)) return;
+      validData[key] = value;
+    });
+
+    if (Object.keys(validData).length === 0)
+      return { thread: prevThread, message: "No valid fields to update!" };
+
+    try {
+      const thread = await prisma.workspace_threads.update({
+        where: { id: prevThread.id },
+        data: validData,
+      });
+      return { thread, message: null };
+    } catch (error) {
+      console.error(error.message);
+      return { thread: null, message: error.message };
+    }
+  },
+
+  get: async function (clause = {}) {
+    try {
+      const thread = await prisma.workspace_threads.findFirst({
+        where: clause,
+      });
+
+      return thread || null;
+    } catch (error) {
+      console.error(error.message);
+      return null;
+    }
+  },
+
+  delete: async function (clause = {}) {
+    try {
+      await prisma.workspace_threads.deleteMany({
+        where: clause,
+      });
+      return true;
+    } catch (error) {
+      console.error(error.message);
+      return false;
+    }
+  },
+
+  where: async function (
+    clause = {},
+    limit = null,
+    orderBy = null,
+    include = null
+  ) {
+    try {
+      const results = await prisma.workspace_threads.findMany({
+        where: clause,
+        ...(limit !== null ? { take: limit } : {}),
+        ...(orderBy !== null ? { orderBy } : {}),
+        ...(include !== null ? { include } : {}),
+      });
+      return results;
+    } catch (error) {
+      console.error(error.message);
+      return [];
+    }
+  },
+
+  // Will fire on first message (included or not) for a thread and rename the thread based on the prompt.
+  autoRenameThread: async function ({
+    workspace = null,
+    thread = null,
+    user = null,
+    prompt = null,
+    onRename = null,
+  }) {
+    if (!workspace || !thread || !prompt) return false;
+    if (thread.name !== this.defaultName) return false; // don't rename if already named.
+
+    const { WorkspaceChats } = require("./workspaceChats");
+    const chatCount = await WorkspaceChats.count({
+      workspaceId: workspace.id,
+      user_id: user?.id || null,
+      thread_id: thread.id,
+    });
+    if (chatCount !== 1) return { renamed: false, thread };
+    const { thread: updatedThread } = await this.update(thread, {
+      name: truncate(prompt, 22),
+    });
+
+    onRename?.(updatedThread);
+    return true;
+  },
+};
+
+module.exports = { WorkspaceThread };

package/server/models/workspaceUsers.js

@@ -0,0 +1,103 @@
+const prisma = require("../utils/prisma");
+
+const WorkspaceUser = {
+  createMany: async function (userId, workspaceIds = []) {
+    if (workspaceIds.length === 0) return;
+    try {
+      await prisma.$transaction(
+        workspaceIds.map((workspaceId) =>
+          prisma.workspace_users.create({
+            data: { user_id: userId, workspace_id: workspaceId },
+          })
+        )
+      );
+    } catch (error) {
+      console.error(error.message);
+    }
+    return;
+  },
+
+  /**
+   * Create many workspace users.
+   * @param {Array<number>} userIds - An array of user IDs to create workspace users for.
+   * @param {number} workspaceId - The ID of the workspace to create workspace users for.
+   * @returns {Promise<void>} A promise that resolves when the workspace users are created.
+   */
+  createManyUsers: async function (userIds = [], workspaceId) {
+    if (userIds.length === 0) return;
+    try {
+      await prisma.$transaction(
+        userIds.map((userId) =>
+          prisma.workspace_users.create({
+            data: {
+              user_id: Number(userId),
+              workspace_id: Number(workspaceId),
+            },
+          })
+        )
+      );
+    } catch (error) {
+      console.error(error.message);
+    }
+    return;
+  },
+
+  create: async function (userId = 0, workspaceId = 0) {
+    try {
+      await prisma.workspace_users.create({
+        data: { user_id: Number(userId), workspace_id: Number(workspaceId) },
+      });
+      return true;
+    } catch (error) {
+      console.error(
+        "FAILED TO CREATE WORKSPACE_USER RELATIONSHIP.",
+        error.message
+      );
+      return false;
+    }
+  },
+
+  get: async function (clause = {}) {
+    try {
+      const result = await prisma.workspace_users.findFirst({ where: clause });
+      return result || null;
+    } catch (error) {
+      console.error(error.message);
+      return null;
+    }
+  },
+
+  where: async function (clause = {}, limit = null) {
+    try {
+      const results = await prisma.workspace_users.findMany({
+        where: clause,
+        ...(limit !== null ? { take: limit } : {}),
+      });
+      return results;
+    } catch (error) {
+      console.error(error.message);
+      return [];
+    }
+  },
+
+  count: async function (clause = {}) {
+    try {
+      const count = await prisma.workspace_users.count({ where: clause });
+      return count;
+    } catch (error) {
+      console.error(error.message);
+      return 0;
+    }
+  },
+
+  delete: async function (clause = {}) {
+    try {
+      await prisma.workspace_users.deleteMany({ where: clause });
+    } catch (error) {
+      console.error(error.message);
+    }
+    return;
+  },
+};
+
+module.exports.WorkspaceUser = WorkspaceUser;

package/server/models/workspacesSuggestedMessages.js

@@ -0,0 +1,83 @@
+const prisma = require("../utils/prisma");
+
+const WorkspaceSuggestedMessages = {
+  get: async function (clause = {}) {
+    try {
+      const message = await prisma.workspace_suggested_messages.findFirst({
+        where: clause,
+      });
+      return message || null;
+    } catch (error) {
+      console.error(error.message);
+      return null;
+    }
+  },
+
+  where: async function (clause = {}, limit) {
+    try {
+      const messages = await prisma.workspace_suggested_messages.findMany({
+        where: clause,
+        take: limit || undefined,
+      });
+      return messages;
+    } catch (error) {
+      console.error(error.message);
+      return [];
+    }
+  },
+
+  saveAll: async function (messages, workspaceSlug) {
+    try {
+      const workspace = await prisma.workspaces.findUnique({
+        where: { slug: workspaceSlug },
+      });
+
+      if (!workspace) throw new Error("Workspace not found");
+
+      // Delete all existing messages for the workspace
+      await prisma.workspace_suggested_messages.deleteMany({
+        where: { workspaceId: workspace.id },
+      });
+
+      // Create new messages
+      // We create each message individually because prisma
+      // with sqlite does not support createMany()
+      for (const message of messages) {
+        await prisma.workspace_suggested_messages.create({
+          data: {
+            workspaceId: workspace.id,
+            heading: message.heading,
+            message: message.message,
+          },
+        });
+      }
+    } catch (error) {
+      console.error("Failed to save all messages", error.message);
+    }
+  },
+
+  getMessages: async function (workspaceSlug) {
+    try {
+      const workspace = await prisma.workspaces.findUnique({
+        where: { slug: workspaceSlug },
+      });
+
+      if (!workspace) throw new Error("Workspace not found");
+
+      const messages = await prisma.workspace_suggested_messages.findMany({
+        where: { workspaceId: workspace.id },
+        orderBy: { createdAt: "asc" },
+      });
+
+      return messages.map((msg) => ({
+        heading: msg.heading,
+        message: msg.message,
+      }));
+    } catch (error) {
+      console.error("Failed to get all messages", error.message);
+      return [];
+    }
+  },
+};
+
+module.exports.WorkspaceSuggestedMessages = WorkspaceSuggestedMessages;