totem-llm 0.1.0

package/collector/utils/extensions/YoutubeTranscript/index.js

@@ -0,0 +1,219 @@
+const fs = require("fs");
+const path = require("path");
+const { default: slugify } = require("slugify");
+const { v4 } = require("uuid");
+const {
+  writeToServerDocuments,
+  sanitizeFileName,
+  documentsFolder,
+  isWithin,
+} = require("../../files");
+const { tokenizeString } = require("../../tokenizer");
+const { YoutubeLoader } = require("./YoutubeLoader");
+const { validYoutubeVideoUrl } = require("../../url");
+
+/**
+ * Fetch the transcript content for a YouTube video
+ * @param {string} url - The URL of the YouTube video
+ * @returns {Promise<{success: boolean, reason: string|null, content: string|null, metadata: TranscriptMetadata}>} - The transcript content for the YouTube video
+ */
+async function fetchVideoTranscriptContent({ url }) {
+  if (!validYoutubeVideoUrl(url)) {
+    return {
+      success: false,
+      reason: "Invalid URL. Should be youtu.be or youtube.com/watch.",
+      content: null,
+      metadata: {},
+    };
+  }
+
+  console.log(`-- Working YouTube ${url} --`);
+  const loader = YoutubeLoader.createFromUrl(url, { addVideoInfo: true });
+  const { docs, error } = await loader
+    .load()
+    .then((docs) => ({ docs, error: null }))
+    .catch((e) => ({
+      docs: [],
+      error: e.message?.split("Error:")?.[1] || e.message,
+    }));
+
+  if (!docs.length || !!error) {
+    return {
+      success: false,
+      reason: error ?? "No transcript found for that YouTube video.",
+      content: null,
+      metadata: {},
+    };
+  }
+
+  const metadata = docs[0].metadata;
+  const content = docs[0].pageContent;
+  if (!content.length) {
+    return {
+      success: false,
+      reason: "No transcript could be parsed for that YouTube video.",
+      content: null,
+      metadata: {},
+    };
+  }
+
+  return {
+    success: true,
+    reason: null,
+    content,
+    metadata,
+  };
+}
+
+/**
+ * @typedef {Object} TranscriptMetadata
+ * @property {string} title - The title of the video
+ * @property {string} author - The author of the video
+ * @property {string} description - The description of the video
+ * @property {string} view_count - The view count of the video
+ * @property {string} source - The source of the video (videoId)
+ */
+
+/**
+ * @typedef {Object} TranscriptAsDocument
+ * @property {boolean} success - Whether the transcript was successful
+ * @property {string|null} reason - The reason for the transcript
+ * @property {TranscriptMetadata} metadata - The metadata from the transcript
+ */
+
+/**
+ * @typedef {Object} TranscriptAsContent
+ * @property {boolean} success - Whether the transcript was successful
+ * @property {string|null} reason - The reason for the transcript
+ * @property {string|null} content - The content of the transcript
+ * @property {Object[]} documents - The documents from the transcript
+ * @property {boolean} saveAsDocument - Whether to save the transcript as a document
+ */
+
+/**
+ * Load the transcript content for a YouTube video as well as save it to the server documents
+ * @param {Object} params - The parameters for the YouTube transcript
+ * @param {string} params.url - The URL of the YouTube video
+ * @param {Object} options - The options for the YouTube transcript
+ * @param {boolean} options.parseOnly - Whether to parse the transcript content only or save it to the server documents
+ * @returns {Promise<TranscriptAsDocument | TranscriptAsContent>} - The transcript content for the YouTube video
+ */
+async function loadYouTubeTranscript({ url }, options = { parseOnly: false }) {
+  const transcriptResults = await fetchVideoTranscriptContent({ url });
+  if (!transcriptResults.success) {
+    return {
+      success: false,
+      reason:
+        transcriptResults.reason ||
+        "An unknown error occurred during transcription retrieval",
+      documents: [],
+      content: null,
+      saveAsDocument: options.parseOnly,
+      data: {},
+    };
+  }
+
+  const { content, metadata } = transcriptResults;
+
+  if (options.parseOnly) {
+    return {
+      success: true,
+      reason: null,
+      content: buildTranscriptContentWithMetadata(content, metadata),
+      documents: [],
+      saveAsDocument: options.parseOnly,
+      data: {},
+    };
+  }
+
+  const outFolder = sanitizeFileName(
+    slugify(`${metadata.author} YouTube transcripts`).toLowerCase()
+  );
+  const outFolderPath = path.resolve(documentsFolder, outFolder);
+  const uuid = v4();
+  const fileName = sanitizeFileName(`${slugify(metadata.title)}-${uuid}`);
+
+  if (!isWithin(documentsFolder, path.resolve(outFolderPath, fileName))) {
+    console.error(
+      `[YouTube Loader]: Invalid file path ${path.resolve(
+        outFolderPath,
+        fileName
+      )} is not within the documents folder ${documentsFolder}`
+    );
+    return {
+      success: false,
+      reason: `[YouTube Loader]: Invalid file path ${path.resolve(
+        outFolderPath,
+        fileName
+      )} is not within the documents folder ${documentsFolder}`,
+      documents: [],
+      data: {},
+    };
+  }
+
+  if (!fs.existsSync(outFolderPath))
+    fs.mkdirSync(outFolderPath, { recursive: true });
+  const data = {
+    id: uuid,
+    url: url + ".youtube",
+    title: metadata.title || url,
+    docAuthor: metadata.author,
+    description: metadata.description,
+    docSource: url,
+    chunkSource: `youtube://${url}`,
+    published: new Date().toLocaleString(),
+    wordCount: content.split(" ").length,
+    pageContent: content,
+    token_count_estimate: tokenizeString(content),
+  };
+
+  console.log(`[YouTube Loader]: Saving ${metadata.title} to ${outFolder}`);
+  const document = writeToServerDocuments({
+    data,
+    filename: fileName,
+    destinationOverride: outFolderPath,
+  });
+
+  return {
+    success: true,
+    reason: null,
+    documents: [document],
+    data: {
+      title: metadata.title,
+      author: metadata.author,
+      destination: outFolder,
+    },
+  };
+}
+
+/**
+ * Generate the transcript content and metadata into a single string
+ *
+ * Why? For ephemeral documents where we just want the content, we want to include the metadata as keys in the content
+ * so that the LLM has context about the video, this gives it a better understanding of the video
+ * and allows it to use the metadata in the conversation if relevant.
+ * Examples:
+ * - How many views does <LINK> have?
+ * - Checkout <LINK> and tell me the key points and if it is performing well
+ * - Summarize this video <LINK>? -> description could have links and references
+ * @param {string} content - The content of the transcript
+ * @param {TranscriptMetadata} metadata - The metadata from the transcript
+ * @returns {string} - The concatenated transcript content and metadata
+ */
+function buildTranscriptContentWithMetadata(content = "", metadata = {}) {
+  const VALID_METADATA_KEYS = ["title", "author", "description", "view_count"];
+  if (!content || !metadata || Object.keys(metadata).length === 0)
+    return content;
+
+  let contentWithMetadata = "";
+  VALID_METADATA_KEYS.forEach((key) => {
+    if (!metadata[key]) return;
+    contentWithMetadata += `<${key}>${metadata[key]}</${key}>`;
+  });
+  return `${contentWithMetadata}\nTranscript:\n${content}`;
+}
+
+module.exports = {
+  loadYouTubeTranscript,
+  fetchVideoTranscriptContent,
+};

package/collector/utils/files/index.js

@@ -0,0 +1,240 @@
+const fs = require("fs");
+const path = require("path");
+const { MimeDetector } = require("./mime");
+
+/**
+ * The folder where documents are stored to be stored when
+ * processed by the collector.
+ */
+const documentsFolder =
+  process.env.NODE_ENV === "development"
+    ? path.resolve(__dirname, `../../../server/storage/documents`)
+    : path.resolve(process.env.STORAGE_DIR, `documents`);
+
+/**
+ * The folder where direct uploads are stored to be stored when
+ * processed by the collector. These are files that were DnD'd into UI
+ * and are not to be embedded or selectable from the file picker.
+ */
+const directUploadsFolder =
+  process.env.NODE_ENV === "development"
+    ? path.resolve(__dirname, `../../../server/storage/direct-uploads`)
+    : path.resolve(process.env.STORAGE_DIR, `direct-uploads`);
+
+/**
+ * Checks if a file is text by checking the mime type and then falling back to buffer inspection.
+ * This way we can capture all the cases where the mime type is not known but still parseable as text
+ * without having to constantly add new mime type overrides.
+ * @param {string} filepath - The path to the file.
+ * @returns {boolean} - Returns true if the file is text, false otherwise.
+ */
+function isTextType(filepath) {
+  if (!fs.existsSync(filepath)) return false;
+  const result = isKnownTextMime(filepath);
+  if (result.valid) return true; // Known text type - return true.
+  if (result.reason !== "generic") return false; // If any other reason than generic - return false.
+  return parseableAsText(filepath); // Fallback to parsing as text via buffer inspection.
+}
+
+/**
+ * Checks if a file is known to be text by checking the mime type.
+ * @param {string} filepath - The path to the file.
+ * @returns {boolean} - Returns true if the file is known to be text, false otherwise.
+ */
+function isKnownTextMime(filepath) {
+  try {
+    const mimeLib = new MimeDetector();
+    const mime = mimeLib.getType(filepath);
+    if (mimeLib.badMimes.includes(mime))
+      return { valid: false, reason: "bad_mime" };
+
+    const type = mime.split("/")[0];
+    if (mimeLib.nonTextTypes.includes(type))
+      return { valid: false, reason: "non_text_mime" };
+    return { valid: true, reason: "valid_mime" };
+  } catch {
+    return { valid: false, reason: "generic" };
+  }
+}
+
+/**
+ * Checks if a file is parseable as text by forcing it to be read as text in utf8 encoding.
+ * If the file looks too much like a binary file, it will return false.
+ * @param {string} filepath - The path to the file.
+ * @returns {boolean} - Returns true if the file is parseable as text, false otherwise.
+ */
+function parseableAsText(filepath) {
+  try {
+    const fd = fs.openSync(filepath, "r");
+    const buffer = Buffer.alloc(1024); // Read first 1KB of the file synchronously
+    const bytesRead = fs.readSync(fd, buffer, 0, 1024, 0);
+    fs.closeSync(fd);
+
+    const content = buffer.subarray(0, bytesRead).toString("utf8");
+    const nullCount = (content.match(/\0/g) || []).length;
+    //eslint-disable-next-line
+    const controlCount = (content.match(/[\x00-\x08\x0B\x0C\x0E-\x1F]/g) || [])
+      .length;
+
+    const threshold = bytesRead * 0.1;
+    return nullCount + controlCount < threshold;
+  } catch {
+    return false;
+  }
+}
+
+function trashFile(filepath) {
+  if (!fs.existsSync(filepath)) return;
+
+  try {
+    const isDir = fs.lstatSync(filepath).isDirectory();
+    if (isDir) return;
+  } catch {
+    return;
+  }
+
+  fs.rmSync(filepath);
+  return;
+}
+
+function createdDate(filepath) {
+  try {
+    const { birthtimeMs, birthtime } = fs.statSync(filepath);
+    if (birthtimeMs === 0) throw new Error("Invalid stat for file!");
+    return birthtime.toLocaleString();
+  } catch {
+    return "unknown";
+  }
+}
+
+/**
+ * Writes a document to the server documents folder.
+ * @param {Object} params - The parameters for the function.
+ * @param {Object} params.data - The data to write to the file. Must look like a document object.
+ * @param {string} params.filename - The name of the file to write to.
+ * @param {string|null} params.destinationOverride - A forced destination to write to - will be honored if provided.
+ * @param {Object} params.options - The options for the function.
+ * @param {boolean} params.options.parseOnly - If true, the file will be written to the direct uploads folder instead of the documents folder. Will be ignored if destinationOverride is provided.
+ * @returns {Object} - The data with the location added.
+ */
+function writeToServerDocuments({
+  data = {},
+  filename,
+  destinationOverride = null,
+  options = {},
+}) {
+  if (!filename) throw new Error("Filename is required!");
+
+  let destination = null;
+  if (destinationOverride) destination = path.resolve(destinationOverride);
+  else if (options.parseOnly) destination = path.resolve(directUploadsFolder);
+  else destination = path.resolve(documentsFolder, "custom-documents");
+
+  if (!fs.existsSync(destination))
+    fs.mkdirSync(destination, { recursive: true });
+  const safeFilename = sanitizeFileName(filename);
+  const destinationFilePath = normalizePath(
+    path.resolve(destination, safeFilename) + ".json"
+  );
+
+  fs.writeFileSync(destinationFilePath, JSON.stringify(data, null, 4), {
+    encoding: "utf-8",
+  });
+
+  return {
+    ...data,
+    // relative location string that can be passed into the /update-embeddings api
+    // that will work since we know the location exists and since we only allow
+    // 1-level deep folders this will always work. This still works for integrations like GitHub and YouTube.
+    location: destinationFilePath.split("/").slice(-2).join("/"),
+    isDirectUpload: options.parseOnly || false,
+  };
+}
+
+// When required we can wipe the entire collector hotdir and tmp storage in case
+// there were some large file failures that we unable to be removed a reboot will
+// force remove them.
+async function wipeCollectorStorage() {
+  const cleanHotDir = new Promise((resolve) => {
+    const directory = path.resolve(__dirname, "../../hotdir");
+    fs.readdir(directory, (err, files) => {
+      if (err) resolve();
+
+      for (const file of files) {
+        if (file === "__HOTDIR__.md") continue;
+        try {
+          fs.rmSync(path.join(directory, file));
+        } catch {}
+      }
+      resolve();
+    });
+  });
+
+  const cleanTmpDir = new Promise((resolve) => {
+    const directory = path.resolve(__dirname, "../../storage/tmp");
+    fs.readdir(directory, (err, files) => {
+      if (err) resolve();
+
+      for (const file of files) {
+        if (file === ".placeholder") continue;
+        try {
+          fs.rmSync(path.join(directory, file));
+        } catch {}
+      }
+      resolve();
+    });
+  });
+
+  await Promise.all([cleanHotDir, cleanTmpDir]);
+  console.log(`Collector hot directory and tmp storage wiped!`);
+  return;
+}
+
+/**
+ * Checks if a given path is within another path.
+ * @param {string} outer - The outer path (should be resolved).
+ * @param {string} inner - The inner path (should be resolved).
+ * @returns {boolean} - Returns true if the inner path is within the outer path, false otherwise.
+ */
+function isWithin(outer, inner) {
+  if (outer === inner) return false;
+  const rel = path.relative(outer, inner);
+  return !rel.startsWith("../") && rel !== "..";
+}
+
+function normalizePath(filepath = "") {
+  const result = path
+    .normalize(filepath.trim())
+    .replace(/^(\.\.(\/|\\|$))+/, "")
+    .trim();
+  if (["..", ".", "/"].includes(result)) throw new Error("Invalid path.");
+  return result;
+}
+
+/**
+ * Strips characters that are illegal in Windows filenames, including Unicode
+ * quotation marks (U+201C, U+201D, etc.) that can get corrupted into ASCII
+ * double-quotes during charset conversion in the upload pipeline.
+ * @param {string} fileName - The filename to sanitize.
+ * @returns {string} - The sanitized filename.
+ */
+function sanitizeFileName(fileName) {
+  if (!fileName) return fileName;
+  return fileName.replace(
+    /[<>:"/\\|?*\u201C\u201D\u201E\u201F\u2018\u2019\u201A\u201B]/g,
+    ""
+  );
+}
+
+module.exports = {
+  trashFile,
+  isTextType,
+  createdDate,
+  writeToServerDocuments,
+  wipeCollectorStorage,
+  normalizePath,
+  isWithin,
+  sanitizeFileName,
+  documentsFolder,
+  directUploadsFolder,
+};

package/collector/utils/files/mime.js

@@ -0,0 +1,64 @@
+const MimeLib = require("mime");
+class MimeDetector {
+  nonTextTypes = ["multipart", "model", "audio", "video", "font"];
+  badMimes = [
+    "application/octet-stream",
+    "application/zip",
+    "application/pkcs8",
+    "application/vnd.microsoft.portable-executable",
+    "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", // XLSX are binaries and need to be handled explicitly.
+    "application/x-msdownload",
+  ];
+
+  constructor() {
+    this.lib = MimeLib;
+    this.setOverrides();
+  }
+
+  setOverrides() {
+    // the .ts extension maps to video/mp2t because of https://en.wikipedia.org/wiki/MPEG_transport_stream
+    // which has had this extension far before TS was invented. So need to force re-map this MIME map.
+    this.lib.define(
+      {
+        "text/plain": [
+          "ts",
+          "tsx",
+          "py",
+          "opts",
+          "lock",
+          "jsonl",
+          "qml",
+          "sh",
+          "c",
+          "cs",
+          "h",
+          "js",
+          "lua",
+          "pas",
+          "r",
+          "go",
+          "ino",
+          "hpp",
+          "linq",
+          "cs",
+        ],
+      },
+      true
+    );
+  }
+
+  /**
+   * Returns the MIME type of the file. If the file has no extension found, it will be processed as a text file.
+   * @param {string} filepath
+   * @returns {string}
+   */
+  getType(filepath) {
+    const parsedMime = this.lib.getType(filepath);
+    if (!!parsedMime) return parsedMime;
+    return null;
+  }
+}
+
+module.exports = {
+  MimeDetector,
+};

package/collector/utils/http/index.js

@@ -0,0 +1,54 @@
+process.env.NODE_ENV === "development"
+  ? require("dotenv").config({ path: `.env.${process.env.NODE_ENV}` })
+  : require("dotenv").config();
+const DEFAULT_COLLECTOR_PORT = 8888;
+
+function reqBody(request) {
+  return typeof request.body === "string"
+    ? JSON.parse(request.body)
+    : request.body;
+}
+
+function queryParams(request) {
+  return request.query;
+}
+
+/**
+ * Validates if the provided baseUrl is a valid URL at all.
+ * - Does not validate if the URL is reachable or accessible.
+ * - Does not do any further validation of the URL like `validURL` in `utils/url/index.js`
+ * @param {string} baseUrl
+ * @returns {boolean}
+ */
+function validBaseUrl(baseUrl) {
+  try {
+    new URL(baseUrl);
+    return true;
+  } catch {
+    return false;
+  }
+}
+
+/**
+ * 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}
+ */
+function getCollectorPort() {
+  if (!("COLLECTOR_PORT" in process.env)) return DEFAULT_COLLECTOR_PORT;
+  const port = Number(process.env.COLLECTOR_PORT);
+  if (Number.isInteger(port) && port > 0 && port <= 65535) return port;
+
+  console.warn(
+    `Invalid COLLECTOR_PORT "${process.env.COLLECTOR_PORT}". Falling back to ${DEFAULT_COLLECTOR_PORT}.`
+  );
+  return DEFAULT_COLLECTOR_PORT;
+}
+
+module.exports = {
+  reqBody,
+  queryParams,
+  validBaseUrl,
+  getCollectorPort,
+};

package/collector/utils/logger/index.js

@@ -0,0 +1,68 @@
+const winston = require("winston");
+
+class Logger {
+  logger = console;
+  static _instance;
+  constructor() {
+    if (Logger._instance) return Logger._instance;
+    this.logger =
+      process.env.NODE_ENV === "production" ? this.getWinstonLogger() : console;
+    Logger._instance = this;
+  }
+
+  getWinstonLogger() {
+    const logger = winston.createLogger({
+      level: "info",
+      defaultMeta: { service: "collector" },
+      transports: [
+        new winston.transports.Console({
+          format: winston.format.combine(
+            winston.format.colorize(),
+            winston.format.printf(
+              ({ level, message, service, origin = "" }) => {
+                return `\x1b[36m[${service}]\x1b[0m${
+                  origin ? `\x1b[33m[${origin}]\x1b[0m` : ""
+                } ${level}: ${message}`;
+              }
+            )
+          ),
+        }),
+      ],
+    });
+
+    function formatArgs(args) {
+      return args
+        .map((arg) => {
+          if (arg instanceof Error) {
+            return arg.stack; // If argument is an Error object, return its stack trace
+          } else if (typeof arg === "object") {
+            return JSON.stringify(arg); // Convert objects to JSON string
+          } else {
+            return arg; // Otherwise, return as-is
+          }
+        })
+        .join(" ");
+    }
+
+    console.log = function (...args) {
+      logger.info(formatArgs(args));
+    };
+    console.error = function (...args) {
+      logger.error(formatArgs(args));
+    };
+    console.info = function (...args) {
+      logger.warn(formatArgs(args));
+    };
+    return logger;
+  }
+}
+
+/**
+ * Sets and overrides Console methods for logging when called.
+ * This is a singleton method and will not create multiple loggers.
+ * @returns {winston.Logger | console} - instantiated logger interface.
+ */
+function setLogger() {
+  return new Logger().logger;
+}
+module.exports = setLogger;