totem-llm 0.1.0

package/server/utils/agents/aibitat/utils/dedupe.js

@@ -0,0 +1,152 @@
+// Some models may attempt to call an expensive or annoying function many times and in that case we will want
+// to implement some stateful tracking during that agent session. GPT4 and other more powerful models are smart
+// enough to realize this, but models like 3.5 lack this. Open source models suffer greatly from this issue.
+// eg: "save something to file..."
+//  agent -> saves
+//  agent -> saves
+//  agent -> saves
+//  agent -> saves
+// ... do random # of times.
+// We want to block all the reruns of a plugin, so we can add this to prevent that behavior from
+// spamming the user (or other costly function) that have the exact same signatures.
+
+// Track Run/isDuplicate prevents _exact_ data re-runs based on the SHA of their inputs
+// StartCooldown/isOnCooldown does prevention of _near-duplicate_ runs based on only the function name that is running.
+// isMarkedUnique/markUnique/removeUniqueConstraint prevents one-time functions from re-running. EG: charting.
+const crypto = require("crypto");
+const DEFAULT_COOLDOWN_MS = 30 * 1000;
+
+class Deduplicator {
+  #hashes = {};
+  #cooldowns = {};
+  #uniques = {};
+  constructor() {}
+
+  log(message, ...args) {
+    console.log(`\x1b[36m[Deduplicator]\x1b[0m ${message}`, ...args);
+  }
+
+  trackRun(
+    key,
+    params = {},
+    options = {
+      cooldown: false,
+      cooldownInMs: DEFAULT_COOLDOWN_MS,
+      markUnique: false,
+    }
+  ) {
+    const hash = crypto
+      .createHash("sha256")
+      .update(JSON.stringify({ key, params }))
+      .digest("hex");
+    this.#hashes[hash] = Number(new Date());
+    if (options.cooldown)
+      this.startCooldown(key, { cooldownInMs: options.cooldownInMs });
+    if (options.markUnique) this.markUnique(key);
+  }
+
+  /**
+   * Checks if a key and params are:
+   * - exactly the same as a previous run.
+   * - on cooldown.
+   * - marked as unique.
+   * @param {string} key - The key to check.
+   * @param {Object} params - The parameters to check.
+   * @returns {{isDuplicate: boolean, reason: string}} - The result of the check.
+   */
+  isDuplicate(key, params = {}) {
+    const newSig = crypto
+      .createHash("sha256")
+      .update(JSON.stringify({ key, params }))
+      .digest("hex");
+    if (this.#hashes.hasOwnProperty(newSig))
+      return {
+        isDuplicate: true,
+        reason: `an exact duplicate of previous run of ${key}`,
+      };
+    if (this.isOnCooldown(key))
+      return {
+        isDuplicate: true,
+        reason: `the function is on cooldown for ${key}.`,
+      };
+    if (this.isMarkedUnique(key))
+      return {
+        isDuplicate: true,
+        reason: `the function is marked as unique for ${key}. Can only be called once per agent session.`,
+      };
+    return { isDuplicate: false, reason: "" };
+  }
+
+  /**
+   * Resets the object property for this instance of the Deduplicator class
+   * @param {('runs'|'cooldowns'|'uniques')} type - The type of prop to reset
+   */
+  reset(type = "runs") {
+    switch (type) {
+      case "runs":
+        this.#hashes = {};
+        break;
+      case "cooldowns":
+        this.#cooldowns = {};
+        break;
+      case "uniques":
+        this.#uniques = {};
+        break;
+    }
+    return;
+  }
+
+  /**
+   * Starts a cooldown for a key.
+   * @param {string} key - The key to start the cooldown for (string key of the function name).
+   * @param {Object} parameters - The parameters for the cooldown.
+   * @param {number} parameters.cooldownInMs - The cooldown in milliseconds.
+   */
+  startCooldown(
+    key,
+    parameters = {
+      cooldownInMs: DEFAULT_COOLDOWN_MS,
+    }
+  ) {
+    const cooldownDelay = parameters.cooldownInMs || DEFAULT_COOLDOWN_MS;
+    this.log(`Starting cooldown for ${key} for ${cooldownDelay}ms`);
+    this.#cooldowns[key] = Number(new Date()) + Number(cooldownDelay);
+  }
+
+  /**
+   * Checks if a key is on cooldown.
+   * @param {string} key - The key to check.
+   * @returns {boolean} - True if the key is on cooldown, false otherwise.
+   */
+  isOnCooldown(key) {
+    if (!this.#cooldowns.hasOwnProperty(key)) return false;
+    return Number(new Date()) <= this.#cooldowns[key];
+  }
+
+  /**
+   * Checks if a key is marked as unique and currently tracked by the deduplicator.
+   * @param {string} key - The key to check.
+   * @returns {boolean} - True if the key is marked as unique, false otherwise.
+   */
+  isMarkedUnique(key) {
+    return this.#uniques.hasOwnProperty(key);
+  }
+
+  /**
+   * Removes the unique constraint for a key.
+   * @param {string} key - The key to remove the unique constraint for.
+   */
+  removeUniqueConstraint(key) {
+    delete this.#uniques[key];
+  }
+
+  /**
+   * Marks a key as unique and currently tracked by the deduplicator.
+   * @param {string} key - The key to mark as unique.
+   */
+  markUnique(key) {
+    this.#uniques[key] = Number(new Date());
+  }
+}
+
+module.exports.Deduplicator = Deduplicator;

package/server/utils/agents/aibitat/utils/summarize.js

@@ -0,0 +1,63 @@
+const { loadSummarizationChain } = require("langchain/chains");
+const { PromptTemplate } = require("@langchain/core/prompts");
+const { RecursiveCharacterTextSplitter } = require("@langchain/textsplitters");
+const Provider = require("../providers/ai-provider");
+/**
+ * @typedef {Object} LCSummarizationConfig
+ * @property {string} provider The LLM to use for summarization (inherited)
+ * @property {string} model The LLM Model to use for summarization (inherited)
+ * @property {AbortController['signal']} controllerSignal Abort controller to stop recursive summarization
+ * @property {string} content The text content of the text to summarize
+ */
+
+/**
+ * Summarize content using LLM LC-Chain call
+ * @param {LCSummarizationConfig} The LLM to use for summarization (inherited)
+ * @returns {Promise<string>} The summarized content.
+ */
+async function summarizeContent({
+  provider = "openai",
+  model = null,
+  controllerSignal,
+  content,
+}) {
+  const llm = Provider.LangChainChatModel(provider, {
+    temperature: 0,
+    model: model,
+  });
+
+  const textSplitter = new RecursiveCharacterTextSplitter({
+    separators: ["\n\n", "\n"],
+    chunkSize: 10000,
+    chunkOverlap: 500,
+  });
+  const docs = await textSplitter.createDocuments([content]);
+
+  const mapPrompt = `
+      Write a detailed summary of the following text for a research purpose:
+      "{text}"
+      SUMMARY:
+      `;
+
+  const mapPromptTemplate = new PromptTemplate({
+    template: mapPrompt,
+    inputVariables: ["text"],
+  });
+
+  // This convenience function creates a document chain prompted to summarize a set of documents.
+  const chain = loadSummarizationChain(llm, {
+    type: "map_reduce",
+    combinePrompt: mapPromptTemplate,
+    combineMapPrompt: mapPromptTemplate,
+    verbose: process.env.NODE_ENV === "development",
+  });
+
+  const res = await chain.call({
+    ...(controllerSignal ? { signal: controllerSignal } : {}),
+    input_documents: docs,
+  });
+
+  return res.text;
+}
+
+module.exports = { summarizeContent };

package/server/utils/agents/aibitat/utils/toolReranker.js

@@ -0,0 +1,225 @@
+const { TokenManager } = require("../../../helpers/tiktoken");
+const {
+  NativeEmbeddingReranker,
+} = require("../../../EmbeddingRerankers/native");
+
+const CHUNK_SIZE = 25;
+const MAX_TEXT_LENGTH = 1000;
+
+class ToolReranker {
+  /**
+   * The default number of top tools to keep after reranking
+   * @type {number}
+   */
+  static defaultTopN = 15;
+
+  static instance = null;
+
+  constructor() {
+    if (ToolReranker.instance) return ToolReranker.instance;
+    ToolReranker.instance = this;
+    this.tokenManager = new TokenManager();
+    this.reranker = null;
+  }
+
+  log(text, ...args) {
+    console.log(`\x1b[33m[IntelligentSkillSelector]\x1b[0m ${text}`, ...args);
+  }
+
+  /**
+   * Check if tool reranking is enabled via environment variable
+   * @returns {boolean}
+   */
+  static isEnabled() {
+    if (!("AGENT_SKILL_RERANKER_ENABLED" in process.env)) return false;
+    if (process.env.AGENT_SKILL_RERANKER_ENABLED === "false") return false;
+    return true;
+  }
+
+  /**
+   * Get the configured topN value from environment or use default
+   * @returns {number}
+   */
+  static getTopN() {
+    const envTopN = parseInt(process.env.AGENT_SKILL_RERANKER_TOP_N, 10);
+    return !isNaN(envTopN) && envTopN > 0 ? envTopN : ToolReranker.defaultTopN;
+  }
+
+  /**
+   * Truncate text to max length, trying to break at word boundary
+   */
+  #truncateText(text, maxLength = MAX_TEXT_LENGTH) {
+    if (!text || text.length <= maxLength) return text;
+    const truncated = text.slice(0, maxLength);
+    const lastSpace = truncated.lastIndexOf(" ");
+    return lastSpace > maxLength * 0.8
+      ? truncated.slice(0, lastSpace)
+      : truncated;
+  }
+
+  /**
+   * Get or initialize the reranker instance
+   */
+  async #getReranker() {
+    if (!this.reranker) {
+      this.reranker = new NativeEmbeddingReranker();
+      await this.reranker.initClient();
+    }
+    return this.reranker;
+  }
+
+  /**
+   * Process documents in chunks and merge results to get global top K
+   */
+  async #chunkedRerank(query, documents, topK) {
+    const reranker = await this.#getReranker();
+    const totalDocs = documents.length;
+
+    if (totalDocs <= CHUNK_SIZE) {
+      return await reranker.rerank(query, documents, { topK });
+    }
+
+    this.log(`Processing ${totalDocs} documents in chunks of ${CHUNK_SIZE}...`);
+    const allScored = [];
+
+    for (let i = 0; i < totalDocs; i += CHUNK_SIZE) {
+      const chunk = documents.slice(i, i + CHUNK_SIZE);
+      const chunkNum = Math.floor(i / CHUNK_SIZE) + 1;
+      const totalChunks = Math.ceil(totalDocs / CHUNK_SIZE);
+
+      this.log(
+        `Processing chunk ${chunkNum}/${totalChunks} (${chunk.length} docs)...`
+      );
+
+      const chunkResults = await reranker.rerank(query, chunk, {
+        topK: chunk.length,
+      });
+
+      chunkResults.forEach((result) => {
+        allScored.push({
+          ...result,
+          rerank_corpus_id: result.rerank_corpus_id + i,
+        });
+      });
+    }
+
+    allScored.sort((a, b) => b.rerank_score - a.rerank_score);
+    return allScored.slice(0, topK);
+  }
+
+  /**
+   * Convert a tool/function definition to a text representation for reranking.
+   * @param {Object} tool - The tool definition object
+   * @returns {{text: string, toolName: string, tool: Object, tokens: number}}
+   */
+  #toolToDocument(tool) {
+    const parts = [];
+    if (!tool || !tool.name)
+      return { text: null, toolName: null, tool: null, tokens: 0 };
+
+    parts.push(tool.name);
+    if (tool.description) parts.push(tool.description);
+
+    if (tool.parameters?.properties) {
+      const paramNames = Object.keys(tool.parameters.properties);
+      if (paramNames.length > 0) {
+        const paramDescriptions = paramNames.map((name) => {
+          const prop = tool.parameters.properties[name];
+          return prop.description ? `${name}: ${prop.description}` : name;
+        });
+        parts.push(paramDescriptions.join(", "));
+      }
+    }
+
+    if (
+      tool.examples &&
+      Array.isArray(tool.examples) &&
+      tool.examples.length > 0
+    ) {
+      const examplePrompts = tool.examples
+        .map((ex) => ex.prompt)
+        .filter(Boolean);
+      if (examplePrompts.length > 0) {
+        parts.push(examplePrompts.join("; "));
+      }
+    }
+
+    const textContent = parts.join("\n");
+    return {
+      text: textContent,
+      toolName: tool.name,
+      tool,
+      tokens: this.tokenManager.countFromString(textContent),
+    };
+  }
+
+  /**
+   * Rerank tools based on the user prompt and return the top N most relevant tools.
+   * Uses chunked processing to handle large numbers of tools efficiently.
+   * @param {string} userPrompt - The user's query/prompt
+   * @param {Object[]} tools - Array of tool/function definitions from aibitat.functions
+   * @param {Object} options - Options for reranking
+   * @param {number} options.topN - Number of top tools to return
+   * @returns {Promise<Object[]>} - Array of reranked tools (top N)
+   */
+  async rerank(userPrompt, tools = [], options = {}) {
+    if (!ToolReranker.isEnabled()) return tools;
+    if (!tools || tools.length === 0) return tools;
+    const { topN = ToolReranker.getTopN() } = options;
+
+    if (tools.length <= topN) {
+      this.log(`${tools.length} tools <= ${topN}, skipping reranking`);
+      return tools;
+    }
+
+    try {
+      this.log(`Starting tool reranking for ${tools.length} tools...`);
+      const documents = tools.map((tool) => this.#toolToDocument(tool));
+      const originalTokenCount = documents.reduce(
+        (acc, doc) => acc + doc.tokens,
+        0
+      );
+
+      const startTime = Date.now();
+      // Truncate and format documents for reranking
+      const rerankDocs = documents.map((doc) => ({
+        text: this.#truncateText(doc.text),
+      }));
+
+      const reranked = await this.#chunkedRerank(userPrompt, rerankDocs, topN);
+      const elapsedMs = Date.now() - startTime;
+
+      const rerankedIndices = reranked.map((doc) => ({
+        index: doc.rerank_corpus_id,
+        score: doc.rerank_score,
+      }));
+
+      const rerankedTools = rerankedIndices.map(
+        ({ index }) => documents[index].tool
+      );
+      const newTokenCount = rerankedIndices.reduce(
+        (acc, { index }) => acc + documents[index].tokens,
+        0
+      );
+      const percentSaved = Math.round(
+        ((originalTokenCount - newTokenCount) / originalTokenCount) * 100
+      );
+      this.log(`
+Identified top ${rerankedTools.length} of ${tools.length} tools in ${elapsedMs}ms
+${originalTokenCount.toLocaleString()} -> ${newTokenCount.toLocaleString()} tokens \x1b[0;93m(${percentSaved}% reduction)\x1b[0m`);
+
+      let logText = "Selected tools:\n";
+      rerankedIndices.forEach(({ index }, i) => {
+        logText += `  ${i + 1}. ${documents[index].toolName}\n`;
+      });
+      this.log(logText);
+      return rerankedTools;
+    } catch (error) {
+      this.log(`Error during tool reranking: ${error.message}`);
+      this.log("Falling back to original tool set");
+      return tools;
+    }
+  }
+}
+
+module.exports = { ToolReranker };

package/server/utils/agents/defaults.js

@@ -0,0 +1,196 @@
+const AgentPlugins = require("./aibitat/plugins");
+const { SystemSettings } = require("../../models/systemSettings");
+const { safeJsonParse } = require("../http");
+const Provider = require("./aibitat/providers/ai-provider");
+const ImportedPlugin = require("./imported");
+const { AgentFlows } = require("../agentFlows");
+const MCPCompatibilityLayer = require("../MCP");
+
+// This is a list of skills that are built-in and default enabled.
+const DEFAULT_SKILLS = [
+  AgentPlugins.memory.name,
+  AgentPlugins.docSummarizer.name,
+  AgentPlugins.webScraping.name,
+];
+
+/**
+ * Configuration for agent skills that require availability checks and disabled sub-skill lists.
+ * Each entry maps a skill name to its availability checker and disabled skills list key.
+ */
+const SKILL_FILTER_CONFIG = {
+  "filesystem-agent": {
+    getAvailability: () =>
+      require("./aibitat/plugins/filesystem/lib").isToolAvailable(),
+    disabledSettingKey: "disabled_filesystem_skills",
+  },
+  "create-files-agent": {
+    getAvailability: () =>
+      require("./aibitat/plugins/create-files/lib").isToolAvailable(),
+    disabledSettingKey: "disabled_create_files_skills",
+  },
+  "gmail-agent": {
+    getAvailability: async () =>
+      require("./aibitat/plugins/gmail/lib").GmailBridge.isToolAvailable(),
+    disabledSettingKey: "disabled_gmail_skills",
+  },
+  "outlook-agent": {
+    getAvailability: async () =>
+      require("./aibitat/plugins/outlook/lib").OutlookBridge.isToolAvailable(),
+    disabledSettingKey: "disabled_outlook_skills",
+  },
+};
+
+const USER_AGENT = {
+  name: "USER",
+  getDefinition: () => {
+    return {
+      interrupt: "ALWAYS",
+      role: "I am the human monitor and oversee this chat. Any questions on action or decision making should be directed to me.",
+    };
+  },
+};
+
+const WORKSPACE_AGENT = {
+  name: "@agent",
+  /**
+   * Get the definition for the workspace agent with its role (prompt) and functions in Aibitat format
+   * @param {string} provider
+   * @param {import("@prisma/client").workspaces | null} workspace
+   * @param {import("@prisma/client").users | null} user
+   * @param {string} [prompt] - Current user message for memory reranking
+   * @returns {Promise<{ role: string, functions: object[] }>}
+   */
+  getDefinition: async (
+    provider = null,
+    workspace = null,
+    user = null,
+    prompt = ""
+  ) => {
+    let [role, clarifyingQuestionsSkills] = await Promise.all([
+      Provider.systemPrompt({
+        provider,
+        workspace,
+        user,
+        prompt,
+      }),
+      clarifyingQuestionsSkillIfEnabled(),
+    ]);
+
+    // If clarifying questions tools are enabled, add a note to the role that the user must use the request-user-input tool to ask questions.
+    if (!!clarifyingQuestionsSkills?.length)
+      role +=
+        "\n\nWhen you need information from the user (URLs, file paths, preferences, choices, etc.), you MUST use the request-user-input tool. Do not ask questions in your text response - the user cannot reply to text. Only the tool can collect user input.";
+
+    return {
+      role,
+      functions: [
+        ...(await agentSkillsFromSystemSettings()),
+        ...clarifyingQuestionsSkills,
+        ...ImportedPlugin.activeImportedPlugins(),
+        ...AgentFlows.activeFlowPlugins(),
+        ...(await new MCPCompatibilityLayer().activeMCPServers()),
+      ],
+    };
+  },
+};
+
+/**
+ * Conditionally include the request-user-input sub-tools in the workspace agent's
+ * function list when the admin has enabled clarifying questions.
+ * Returns an empty array when disabled so the tools aren't visible to the LLM.
+ * Names use the parent#child convention so #attachPlugins loads each sub-tool.
+ * @returns {Promise<string[]>}
+ */
+async function clarifyingQuestionsSkillIfEnabled() {
+  const enabled =
+    (await SystemSettings.getValueOrFallback(
+      { label: "agent_clarifying_questions_enabled" },
+      "false"
+    )) === "true";
+  if (!enabled) return [];
+
+  const parentName = AgentPlugins.requestUserInput.name;
+  const subPlugins = AgentPlugins.requestUserInput.plugin;
+  if (!Array.isArray(subPlugins)) return [];
+  return subPlugins.map((sub) => `${parentName}#${sub.name}`);
+}
+
+/**
+ * Fetches and preloads the names/identifiers for plugins that will be dynamically
+ * loaded later
+ * @returns {Promise<string[]>}
+ */
+async function agentSkillsFromSystemSettings() {
+  const systemFunctions = [];
+
+  // Load non-imported built-in skills that are configurable, but are default enabled.
+  const _disabledDefaultSkills = safeJsonParse(
+    await SystemSettings.getValueOrFallback(
+      { label: "disabled_agent_skills" },
+      "[]"
+    ),
+    []
+  );
+  DEFAULT_SKILLS.forEach((skill) => {
+    if (!_disabledDefaultSkills.includes(skill))
+      systemFunctions.push(AgentPlugins[skill].name);
+  });
+
+  // Load non-imported built-in skills that are configurable.
+  const _setting = safeJsonParse(
+    await SystemSettings.getValueOrFallback(
+      { label: "default_agent_skills" },
+      "[]"
+    ),
+    []
+  );
+
+  // Pre-load disabled sub-skills and availability for configured skills
+  const skillFilterState = {};
+  for (const skillName of Object.keys(SKILL_FILTER_CONFIG)) {
+    if (!_setting.includes(skillName)) continue;
+    const config = SKILL_FILTER_CONFIG[skillName];
+    skillFilterState[skillName] = {
+      available: await config.getAvailability(),
+      disabledSubSkills: safeJsonParse(
+        await SystemSettings.getValueOrFallback(
+          { label: config.disabledSettingKey },
+          "[]"
+        ),
+        []
+      ),
+    };
+  }
+
+  for (const skillName of _setting) {
+    if (!AgentPlugins.hasOwnProperty(skillName)) continue;
+
+    // This is a plugin module with many sub-children plugins who
+    // need to be named via `${parent}#${child}` naming convention
+    if (Array.isArray(AgentPlugins[skillName].plugin)) {
+      for (const subPlugin of AgentPlugins[skillName].plugin) {
+        // Check if this skill has filter configuration
+        const filterState = skillFilterState[skillName];
+        if (filterState) {
+          if (!filterState.available) continue;
+          if (filterState.disabledSubSkills.includes(subPlugin.name)) continue;
+        }
+
+        systemFunctions.push(
+          `${AgentPlugins[skillName].name}#${subPlugin.name}`
+        );
+      }
+      continue;
+    }
+
+    // This is normal single-stage plugin
+    systemFunctions.push(AgentPlugins[skillName].name);
+  }
+  return systemFunctions;
+}
+
+module.exports = {
+  USER_AGENT,
+  WORKSPACE_AGENT,
+  agentSkillsFromSystemSettings,
+};