@mastra/memory 0.0.0-a2a-20250421213654

package/dist/processors/index.js

@@ -0,0 +1,154 @@
+import { MemoryProcessor } from '@mastra/core/memory';
+import { Tiktoken } from 'js-tiktoken/lite';
+import o200k_base from 'js-tiktoken/ranks/o200k_base';
+import { MemoryProcessor as MemoryProcessor$1 } from '@mastra/core';
+
+// src/processors/token-limiter.ts
+var TokenLimiter = class extends MemoryProcessor {
+  encoder;
+  maxTokens;
+  // Token overheads per OpenAI's documentation
+  // See: https://cookbook.openai.com/examples/how_to_count_tokens_with_tiktoken#6-counting-tokens-for-chat-completions-api-calls
+  // Every message follows <|start|>{role/name}\n{content}<|end|>
+  TOKENS_PER_MESSAGE = 3;
+  // tokens added for each message (start & end tokens)
+  TOKENS_PER_TOOL = 2;
+  // empirical adjustment for tool calls
+  TOKENS_PER_CONVERSATION = 25;
+  // fixed overhead for the conversation
+  /**
+   * Create a token limiter for messages.
+   * @param options Either a number (token limit) or a configuration object
+   */
+  constructor(options) {
+    super({
+      name: "TokenLimiter"
+    });
+    if (typeof options === "number") {
+      this.maxTokens = options;
+      this.encoder = new Tiktoken(o200k_base);
+    } else {
+      this.maxTokens = options.limit;
+      this.encoder = new Tiktoken(options.encoding || o200k_base);
+    }
+  }
+  process(messages, { systemMessage, memorySystemMessage, newMessages } = {}) {
+    let totalTokens = 0;
+    totalTokens += this.TOKENS_PER_CONVERSATION;
+    if (systemMessage) {
+      totalTokens += this.countTokens(systemMessage);
+      totalTokens += this.TOKENS_PER_MESSAGE;
+    }
+    if (memorySystemMessage) {
+      totalTokens += this.countTokens(memorySystemMessage);
+      totalTokens += this.TOKENS_PER_MESSAGE;
+    }
+    const allMessages = [...messages, ...newMessages || []];
+    const result = [];
+    for (let i = allMessages.length - 1; i >= 0; i--) {
+      const message = allMessages[i];
+      if (!message) continue;
+      const messageTokens = this.countTokens(message);
+      if (totalTokens + messageTokens <= this.maxTokens) {
+        result.unshift(message);
+        totalTokens += messageTokens;
+      } else {
+        this.logger.info(
+          `filtering ${allMessages.length - result.length}/${allMessages.length} messages, token limit of ${this.maxTokens} exceeded`
+        );
+        break;
+      }
+    }
+    return result;
+  }
+  countTokens(message) {
+    if (typeof message === `string`) {
+      return this.encoder.encode(message).length;
+    }
+    let tokenString = message.role;
+    if (typeof message.content === "string") {
+      tokenString += message.content;
+    } else if (Array.isArray(message.content)) {
+      for (const part of message.content) {
+        tokenString += part.type;
+        if (part.type === "text") {
+          tokenString += part.text;
+        } else if (part.type === "tool-call") {
+          tokenString += part.toolName;
+          if (part.args) {
+            tokenString += typeof part.args === "string" ? part.args : JSON.stringify(part.args);
+          }
+        } else if (part.type === "tool-result") {
+          if (part.result !== void 0) {
+            tokenString += typeof part.result === "string" ? part.result : JSON.stringify(part.result);
+          }
+        } else {
+          tokenString += JSON.stringify(part);
+        }
+      }
+    }
+    const messageOverhead = this.TOKENS_PER_MESSAGE;
+    let toolOverhead = 0;
+    if (Array.isArray(message.content)) {
+      for (const part of message.content) {
+        if (part.type === "tool-call" || part.type === "tool-result") {
+          toolOverhead += this.TOKENS_PER_TOOL;
+        }
+      }
+    }
+    const totalMessageOverhead = messageOverhead + toolOverhead;
+    return this.encoder.encode(tokenString).length + totalMessageOverhead;
+  }
+};
+var ToolCallFilter = class extends MemoryProcessor$1 {
+  exclude;
+  /**
+   * Create a filter for tool calls and results.
+   * @param options Configuration options
+   * @param options.exclude List of specific tool names to exclude. If not provided, all tool calls are excluded.
+   */
+  constructor(options = {}) {
+    super({ name: "ToolCallFilter" });
+    if (!options || !options.exclude) {
+      this.exclude = "all";
+    } else {
+      this.exclude = Array.isArray(options.exclude) ? options.exclude : [];
+    }
+  }
+  process(messages) {
+    if (this.exclude === "all") {
+      return messages.filter((message) => {
+        if (Array.isArray(message.content)) {
+          return !message.content.some((part) => part.type === "tool-call" || part.type === "tool-result");
+        }
+        return true;
+      });
+    }
+    if (this.exclude.length > 0) {
+      const excludedToolCallIds = /* @__PURE__ */ new Set();
+      return messages.filter((message) => {
+        if (!Array.isArray(message.content)) return true;
+        if (message.role === "assistant") {
+          let shouldExclude = false;
+          for (const part of message.content) {
+            if (part.type === "tool-call" && this.exclude.includes(part.toolName)) {
+              excludedToolCallIds.add(part.toolCallId);
+              shouldExclude = true;
+            }
+          }
+          return !shouldExclude;
+        }
+        if (message.role === "tool") {
+          const shouldExclude = message.content.some(
+            (part) => part.type === "tool-result" && excludedToolCallIds.has(part.toolCallId)
+          );
+          return !shouldExclude;
+        }
+        return true;
+      });
+    }
+    return messages;
+  }
+};
+
+export { TokenLimiter, ToolCallFilter };

package/eslint.config.js

@@ -0,0 +1,12 @@
+import { createConfig } from '@internal/lint/eslint';
+
+const config = await createConfig();
+
+/** @type {import("eslint").Linter.Config[]} */
+export default [
+  ...config,
+  {
+    files: ['integration-tests/**/*'],
+    ...(await import('typescript-eslint')).configs.disableTypeChecked,
+  },
+];

package/package.json

@@ -0,0 +1,67 @@
+{
+  "name": "@mastra/memory",
+  "version": "0.0.0-a2a-20250421213654",
+  "description": "",
+  "type": "module",
+  "main": "./dist/index.js",
+  "types": "./dist/index.d.ts",
+  "exports": {
+    ".": {
+      "import": {
+        "types": "./dist/index.d.ts",
+        "default": "./dist/index.js"
+      },
+      "require": {
+        "types": "./dist/index.d.cts",
+        "default": "./dist/index.cjs"
+      }
+    },
+    "./processors": {
+      "import": {
+        "types": "./dist/processors/index.d.ts",
+        "default": "./dist/processors/index.js"
+      },
+      "require": {
+        "types": "./dist/processors/index.d.cts",
+        "default": "./dist/processors/index.cjs"
+      }
+    },
+    "./package.json": "./package.json"
+  },
+  "keywords": [],
+  "author": "",
+  "license": "Elastic-2.0",
+  "dependencies": {
+    "@upstash/redis": "^1.34.5",
+    "ai": "^4.2.2",
+    "js-tiktoken": "^1.0.19",
+    "pg": "^8.13.3",
+    "pg-pool": "^3.7.1",
+    "postgres": "^3.4.5",
+    "redis": "^4.7.0",
+    "xxhash-wasm": "^1.1.0",
+    "zod": "^3.24.2",
+    "@mastra/core": "0.0.0-a2a-20250421213654"
+  },
+  "devDependencies": {
+    "@ai-sdk/openai": "^1.3.3",
+    "@microsoft/api-extractor": "^7.52.1",
+    "@types/node": "^20.17.27",
+    "@types/pg": "^8.11.11",
+    "eslint": "^9.23.0",
+    "tsup": "^8.4.0",
+    "typescript": "^5.8.2",
+    "typescript-eslint": "^8.26.1",
+    "vitest": "^3.0.9",
+    "@internal/lint": "0.0.2"
+  },
+  "scripts": {
+    "check": "tsc --noEmit",
+    "build": "pnpm run check && tsup src/index.ts src/processors/index.ts --format esm,cjs --experimental-dts --clean --treeshake=smallest --splitting",
+    "build:watch": "pnpm build --watch",
+    "test:integration": "cd integration-tests && pnpm run test",
+    "test:unit": "pnpm vitest run ./src/*",
+    "test": "pnpm test:integration && pnpm test:unit",
+    "lint": "eslint ."
+  }
+}