@mastra/memory 0.14.3-alpha.1 → 0.14.3-alpha.2

package/src/processors/index.test.ts

@@ -1,246 +0,0 @@
-import { openai } from '@ai-sdk/openai';
-import { Agent } from '@mastra/core/agent';
-import type { MastraMessageV1 } from '@mastra/core/memory';
-import { createTool } from '@mastra/core/tools';
-import type { CoreMessage } from 'ai';
-import cl100k_base from 'js-tiktoken/ranks/cl100k_base';
-import { describe, it, expect, vi } from 'vitest';
-import { z } from 'zod';
-import { generateConversationHistory } from '../../integration-tests/src/test-utils';
-import { TokenLimiter, ToolCallFilter } from './index';
-
-vi.setConfig({ testTimeout: 20_000, hookTimeout: 20_000 });
-
-describe('TokenLimiter', () => {
-  it('should limit messages to the specified token count', () => {
-    // Create messages with predictable token counts (approximately 25 tokens each)
-    const { fakeCore } = generateConversationHistory({
-      threadId: '1',
-      messageCount: 5,
-      toolNames: [],
-      toolFrequency: 0,
-    });
-
-    const limiter = new TokenLimiter(200);
-    const result = limiter.process(fakeCore) as MastraMessageV1[];
-
-    // Should prioritize newest messages (higher ids)
-    expect(result.length).toBe(2);
-    expect(result[0].id).toBe('message-8');
-    expect(result[1].id).toBe('message-9');
-  });
-
-  it('should handle empty messages array', () => {
-    const limiter = new TokenLimiter(1000);
-    const result = limiter.process([]);
-    expect(result).toEqual([]);
-  });
-
-  it('should use different encodings based on configuration', () => {
-    const { fakeCore } = generateConversationHistory({
-      threadId: '6',
-      messageCount: 1,
-      toolNames: [],
-      toolFrequency: 0,
-    });
-
-    // Create limiters with different encoding settings
-    const defaultLimiter = new TokenLimiter(1000);
-    const customLimiter = new TokenLimiter({
-      limit: 1000,
-      encoding: cl100k_base,
-    });
-
-    // All should process fakeCore successfully but potentially with different token counts
-    const defaultResult = defaultLimiter.process(fakeCore);
-    const customResult = customLimiter.process(fakeCore);
-
-    // Each should return the same fakeCore but with potentially different token counts
-    expect(defaultResult.length).toBe(fakeCore.length);
-    expect(customResult.length).toBe(fakeCore.length);
-  });
-
-  function estimateTokens(messages: MastraMessageV1[]) {
-    // Create a TokenLimiter just for counting tokens
-    const testLimiter = new TokenLimiter(Infinity);
-
-    let estimatedTokens = testLimiter.TOKENS_PER_CONVERSATION;
-
-    // Count tokens for each message including all overheads
-    for (const message of messages) {
-      // Base token count from the countTokens method
-      estimatedTokens += testLimiter.countTokens(message as CoreMessage); // TODO: this is really actually a MastraMessageV1 but in previous implementations we were casting V1 to CoreMessage which is almost the same but not exactly
-    }
-
-    return Number(estimatedTokens.toFixed(2));
-  }
-
-  function percentDifference(a: number, b: number) {
-    const difference = Number(((Math.abs(a - b) / b) * 100).toFixed(2));
-    console.log(`${a} and ${b} are ${difference}% different`);
-    return difference;
-  }
-
-  async function expectTokenEstimate(config: Parameters<typeof generateConversationHistory>[0], agent: Agent) {
-    const { messages, fakeCore, counts } = generateConversationHistory(config);
-
-    const estimate = estimateTokens(messages);
-    const used = (await agent.generate(fakeCore)).usage.promptTokens;
-
-    console.log(`Estimated ${estimate} tokens, used ${used} tokens.\n`, counts);
-
-    // Check if within 2% margin
-    expect(percentDifference(estimate, used)).toBeLessThanOrEqual(2);
-  }
-
-  const calculatorTool = createTool({
-    id: 'calculator',
-    description: 'Perform a simple calculation',
-    inputSchema: z.object({
-      expression: z.string().describe('The mathematical expression to calculate'),
-    }),
-    execute: async ({ context: { expression } }) => {
-      // Don't actually eval the expression. The model is dumb and sometimes passes "banana" as the expression because that's one of the sample tokens we're using in input messages lmao
-      return `The result of ${expression} is 10`;
-    },
-  });
-
-  const agent = new Agent({
-    name: 'token estimate agent',
-    model: openai('gpt-4o-mini'),
-    instructions: ``,
-    tools: { calculatorTool },
-  });
-
-  describe.concurrent(`98% accuracy`, () => {
-    it(`20 messages, no tools`, async () => {
-      await expectTokenEstimate(
-        {
-          messageCount: 10,
-          toolFrequency: 0,
-          threadId: '2',
-        },
-        agent,
-      );
-    });
-
-    it(`60 messages, no tools`, async () => {
-      await expectTokenEstimate(
-        {
-          messageCount: 30,
-          toolFrequency: 0,
-          threadId: '3',
-        },
-        agent,
-      );
-    });
-
-    it(`20 messages, 0 tools`, async () => {
-      await expectTokenEstimate(
-        {
-          messageCount: 10,
-          toolFrequency: 0,
-          threadId: '3',
-        },
-        agent,
-      );
-    });
-
-    it(`20 messages, 2 tool messages`, async () => {
-      await expectTokenEstimate(
-        {
-          messageCount: 10,
-          toolFrequency: 5,
-          threadId: '3',
-        },
-        agent,
-      );
-    });
-
-    it(`40 messages, 6 tool messages`, async () => {
-      await expectTokenEstimate(
-        {
-          messageCount: 20,
-          toolFrequency: 5,
-          threadId: '4',
-        },
-        agent,
-      );
-    });
-
-    it(`100 messages, 24 tool messages`, async () => {
-      await expectTokenEstimate(
-        {
-          messageCount: 50,
-          toolFrequency: 4,
-          threadId: '5',
-        },
-        agent,
-      );
-    });
-
-    it(
-      `101 messages, 49 tool calls`,
-      async () => {
-        await expectTokenEstimate(
-          {
-            messageCount: 50,
-            toolFrequency: 1,
-            threadId: '5',
-          },
-          agent,
-        );
-      },
-      {
-        // for some reason AI SDK randomly returns 2x token count here
-        retry: 3,
-      },
-    );
-  });
-});
-
-describe.concurrent('ToolCallFilter', () => {
-  it('should exclude all tool calls when created with no arguments', () => {
-    const { fakeCore } = generateConversationHistory({
-      threadId: '3',
-      toolNames: ['weather', 'calculator', 'search'],
-      messageCount: 1,
-    });
-    const filter = new ToolCallFilter();
-    const result = filter.process(fakeCore) as MastraMessageV1[];
-
-    // Should only keep the text message and assistant res
-    expect(result.length).toBe(2);
-    expect(result[0].id).toBe('message-0');
-  });
-
-  it('should exclude specific tool calls by name', () => {
-    const { fakeCore } = generateConversationHistory({
-      threadId: '4',
-      toolNames: ['weather', 'calculator'],
-      messageCount: 2,
-    });
-    const filter = new ToolCallFilter({ exclude: ['weather'] });
-    const result = filter.process(fakeCore) as MastraMessageV1[];
-
-    // Should keep text message, assistant reply, calculator tool call, and calculator result
-    expect(result.length).toBe(4);
-    expect(result[0].id).toBe('message-0');
-    expect(result[1].id).toBe('message-1');
-    expect(result[2].id).toBe('message-2');
-    expect(result[3].id).toBe('message-3');
-  });
-
-  it('should keep all messages when exclude list is empty', () => {
-    const { fakeCore } = generateConversationHistory({
-      threadId: '5',
-      toolNames: ['weather', 'calculator'],
-    });
-
-    const filter = new ToolCallFilter({ exclude: [] });
-    const result = filter.process(fakeCore);
-
-    // Should keep all messages
-    expect(result.length).toBe(fakeCore.length);
-  });
-});

package/src/processors/index.ts

@@ -1,2 +0,0 @@
-export * from './token-limiter';
-export * from './tool-call-filter';

package/src/processors/token-limiter.ts

@@ -1,159 +0,0 @@
-import type { CoreMessage } from '@mastra/core/llm';
-import { MemoryProcessor } from '@mastra/core/memory';
-import type { MemoryProcessorOpts } from '@mastra/core/memory';
-
-import { Tiktoken } from 'js-tiktoken/lite';
-import type { TiktokenBPE } from 'js-tiktoken/lite';
-import o200k_base from 'js-tiktoken/ranks/o200k_base';
-
-/**
- * Configuration options for TokenLimiter
- */
-interface TokenLimiterOptions {
-  /** Maximum number of tokens to allow */
-  limit: number;
-  /** Optional encoding to use (defaults to o200k_base which is used by gpt-4o) */
-  encoding?: TiktokenBPE;
-}
-
-/**
- * Limits the total number of tokens in the messages.
- * Uses js-tiktoken with o200k_base encoding by default for accurate token counting with modern models.
- */
-export class TokenLimiter extends MemoryProcessor {
-  private encoder: Tiktoken;
-  private maxTokens: number;
-
-  // 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|>
-  public TOKENS_PER_MESSAGE = 3.8; // tokens added for each message (start & end tokens)
-  public TOKENS_PER_CONVERSATION = 24; // fixed overhead for the conversation
-
-  /**
-   * Create a token limiter for messages.
-   * @param options Either a number (token limit) or a configuration object
-   */
-  constructor(options: number | TokenLimiterOptions) {
-    super({
-      name: 'TokenLimiter',
-    });
-
-    if (typeof options === 'number') {
-      // Simple number format - just the token limit with default encoding
-      this.maxTokens = options;
-      this.encoder = new Tiktoken(o200k_base);
-    } else {
-      // Object format with limit and optional encoding
-      this.maxTokens = options.limit;
-      this.encoder = new Tiktoken(options.encoding || o200k_base);
-    }
-  }
-
-  process(
-    messages: CoreMessage[],
-    { systemMessage, memorySystemMessage, newMessages }: MemoryProcessorOpts = {},
-  ): CoreMessage[] {
-    // Messages are already chronologically ordered - take most recent ones up to the token limit
-    let totalTokens = 0;
-
-    // Start with the conversation overhead
-    totalTokens += this.TOKENS_PER_CONVERSATION;
-
-    if (systemMessage) {
-      totalTokens += this.countTokens(systemMessage);
-      totalTokens += this.TOKENS_PER_MESSAGE; // Add message overhead for system message
-    }
-
-    if (memorySystemMessage) {
-      totalTokens += this.countTokens(memorySystemMessage);
-      totalTokens += this.TOKENS_PER_MESSAGE; // Add message overhead for memory system message
-    }
-
-    const allMessages = [...messages, ...(newMessages || [])];
-
-    const result: CoreMessage[] = [];
-
-    // Process messages in reverse (newest first) so that we stop estimating tokens on old messages. Once we get to our limit of tokens there's no reason to keep processing older messages
-    for (let i = allMessages.length - 1; i >= 0; i--) {
-      const message = allMessages[i];
-
-      // Skip undefined messages (shouldn't happen, but TypeScript is concerned)
-      if (!message) continue;
-
-      const messageTokens = this.countTokens(message);
-
-      if (totalTokens + messageTokens <= this.maxTokens) {
-        // Insert at the beginning to maintain chronological order, but only if it's not a new message
-        if (i < messages.length) {
-          // less than messages.length because we're iterating in reverse. If the index is greater than messages.length it's a new message
-          result.unshift(message);
-        }
-        totalTokens += messageTokens;
-      } else {
-        this.logger.info(
-          `filtering ${allMessages.length - result.length}/${allMessages.length} messages, token limit of ${this.maxTokens} exceeded`,
-        );
-        // If we can't fit the message, we stop
-        break;
-      }
-    }
-
-    return result;
-  }
-
-  public countTokens(message: string | CoreMessage): number {
-    if (typeof message === `string`) {
-      return this.encoder.encode(message).length;
-    }
-
-    let tokenString = message.role;
-    let overhead = 0;
-
-    if (typeof message.content === 'string' && message.content) {
-      tokenString += message.content;
-    } else if (Array.isArray(message.content)) {
-      // Calculate tokens for each content part
-      for (const part of message.content) {
-        if (part.type === 'text') {
-          tokenString += part.text;
-        } else if (part.type === 'tool-call' || part.type === `tool-result`) {
-          if (`args` in part && part.args && part.type === `tool-call`) {
-            tokenString += part.toolName as any;
-            if (typeof part.args === 'string') {
-              tokenString += part.args;
-            } else {
-              tokenString += JSON.stringify(part.args);
-              // minus some tokens for JSON
-              overhead -= 12;
-            }
-          }
-          // Token cost for result if present
-          if (`result` in part && part.result !== undefined && part.type === `tool-result`) {
-            if (typeof part.result === 'string') {
-              tokenString += part.result;
-            } else {
-              tokenString += JSON.stringify(part.result);
-              // minus some tokens for JSON
-              overhead -= 12;
-            }
-          }
-        } else {
-          tokenString += JSON.stringify(part);
-        }
-      }
-    }
-
-    if (
-      typeof message.content === `string` ||
-      // if the message included non-tool parts, add our message overhead
-      message.content.some(p => p.type !== `tool-call` && p.type !== `tool-result`)
-    ) {
-      // Ensure we account for message formatting tokens
-      // See: https://cookbook.openai.com/examples/how_to_count_tokens_with_tiktoken#6-counting-tokens-for-chat-completions-api-calls
-      overhead += this.TOKENS_PER_MESSAGE;
-    }
-
-    return this.encoder.encode(tokenString).length + overhead;
-  }
-}

package/src/processors/tool-call-filter.ts

@@ -1,77 +0,0 @@
-import type { CoreMessage } from '@mastra/core/llm';
-import { MemoryProcessor } from '@mastra/core/memory';
-
-/**
- * Filters out tool calls and results from messages.
- * By default (with no arguments), excludes all tool calls and their results.
- * Can be configured to exclude only specific tools by name.
- */
-export class ToolCallFilter extends MemoryProcessor {
-  private exclude: string[] | 'all';
-
-  /**
-   * 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: { exclude?: string[] } = {}) {
-    super({ name: 'ToolCallFilter' });
-    // If no options or exclude is provided, exclude all tools
-    if (!options || !options.exclude) {
-      this.exclude = 'all'; // Exclude all tools
-    } else {
-      // Exclude specific tools
-      this.exclude = Array.isArray(options.exclude) ? options.exclude : [];
-    }
-  }
-
-  process(messages: CoreMessage[]): CoreMessage[] {
-    // Case 1: Exclude all tool calls and tool results
-    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;
-      });
-    }
-
-    // Case 2: Exclude specific tools by name
-    if (this.exclude.length > 0) {
-      // Single pass approach - track excluded tool call IDs while filtering
-      const excludedToolCallIds = new Set<string>();
-
-      return messages.filter(message => {
-        if (!Array.isArray(message.content)) return true;
-
-        // For assistant messages, check for excluded tool calls and track their IDs
-        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;
-        }
-
-        // For tool messages, filter out results for excluded tool calls
-        if (message.role === 'tool') {
-          const shouldExclude = message.content.some(
-            part => part.type === 'tool-result' && excludedToolCallIds.has(part.toolCallId),
-          );
-
-          return !shouldExclude;
-        }
-
-        return true;
-      });
-    }
-
-    // Case 3: Empty exclude array, return original messages
-    return messages;
-  }
-}

package/src/tools/working-memory.ts

@@ -1,154 +0,0 @@
-import type { MemoryConfig } from '@mastra/core/memory';
-import { createTool } from '@mastra/core/tools';
-import { convertSchemaToZod } from '@mastra/schema-compat';
-import type { Schema } from 'ai';
-import { z, ZodObject } from 'zod';
-import type { ZodType } from 'zod';
-
-export const updateWorkingMemoryTool = (memoryConfig?: MemoryConfig) => {
-  const schema = memoryConfig?.workingMemory?.schema;
-
-  let inputSchema: ZodType = z.object({
-    memory: z
-      .string()
-      .describe(`The Markdown formatted working memory content to store. This MUST be a string. Never pass an object.`),
-  });
-
-  if (schema) {
-    inputSchema = z.object({
-      memory:
-        schema instanceof ZodObject
-          ? schema
-          : (convertSchemaToZod({ jsonSchema: schema } as Schema).describe(
-              `The JSON formatted working memory content to store.`,
-            ) as ZodObject<any>),
-    });
-  }
-
-  return createTool({
-    id: 'update-working-memory',
-    description: `Update the working memory with new information. Any data not included will be overwritten.${schema ? ' Always pass data as string to the memory field. Never pass an object.' : ''}`,
-    inputSchema,
-    execute: async params => {
-      const { context, threadId, memory, resourceId } = params;
-      if (!threadId || !memory || !resourceId) {
-        throw new Error('Thread ID, Memory instance, and resourceId are required for working memory updates');
-      }
-
-      let thread = await memory.getThreadById({ threadId });
-
-      if (!thread) {
-        thread = await memory.createThread({
-          threadId,
-          resourceId,
-          memoryConfig,
-        });
-      }
-
-      if (thread.resourceId && thread.resourceId !== resourceId) {
-        throw new Error(`Thread with id ${threadId} resourceId does not match the current resourceId ${resourceId}`);
-      }
-
-      const workingMemory = typeof context.memory === 'string' ? context.memory : JSON.stringify(context.memory);
-
-      // Use the new updateWorkingMemory method which handles both thread and resource scope
-      await memory.updateWorkingMemory({
-        threadId,
-        resourceId,
-        workingMemory,
-        memoryConfig,
-      });
-
-      return { success: true };
-    },
-  });
-};
-
-export const __experimental_updateWorkingMemoryToolVNext = (config: MemoryConfig) => {
-  return createTool({
-    id: 'update-working-memory',
-    description: 'Update the working memory with new information.',
-    inputSchema: z.object({
-      newMemory: z
-        .string()
-        .optional()
-        .describe(
-          `The ${config.workingMemory?.schema ? 'JSON' : 'Markdown'} formatted working memory content to store`,
-        ),
-      searchString: z
-        .string()
-        .optional()
-        .describe(
-          "The working memory string to find. Will be replaced with the newMemory string. If this is omitted or doesn't exist, the newMemory string will be appended to the end of your working memory. Replacing single lines at a time is encouraged for greater accuracy. If updateReason is not 'append-new-memory', this search string must be provided or the tool call will be rejected.",
-        ),
-      updateReason: z
-        .enum(['append-new-memory', 'clarify-existing-memory', 'replace-irrelevant-memory'])
-        .optional()
-        .describe(
-          "The reason you're updating working memory. Passing any value other than 'append-new-memory' requires a searchString to be provided. Defaults to append-new-memory",
-        ),
-    }),
-    execute: async params => {
-      const { context, threadId, memory, resourceId } = params;
-      if (!threadId || !memory || !resourceId) {
-        throw new Error('Thread ID, Memory instance, and resourceId are required for working memory updates');
-      }
-
-      let thread = await memory.getThreadById({ threadId });
-
-      if (!thread) {
-        thread = await memory.createThread({
-          threadId,
-          resourceId,
-          memoryConfig: config,
-        });
-      }
-
-      if (thread.resourceId && thread.resourceId !== resourceId) {
-        throw new Error(`Thread with id ${threadId} resourceId does not match the current resourceId ${resourceId}`);
-      }
-
-      const workingMemory = context.newMemory || '';
-      if (!context.updateReason) context.updateReason = `append-new-memory`;
-
-      if (
-        context.searchString &&
-        config.workingMemory?.scope === `resource` &&
-        context.updateReason === `replace-irrelevant-memory`
-      ) {
-        // don't allow replacements due to something not being relevant to the current conversation
-        // if there's no searchString, then we will append.
-        context.searchString = undefined;
-      }
-
-      if (context.updateReason === `append-new-memory` && context.searchString) {
-        // do not find/replace when append-new-memory is selected
-        // some models get confused and pass a search string even when they don't want to replace it.
-        // TODO: maybe they're trying to add new info after the search string?
-        context.searchString = undefined;
-      }
-
-      if (context.updateReason !== `append-new-memory` && !context.searchString) {
-        return {
-          success: false,
-          reason: `updateReason was ${context.updateReason} but no searchString was provided. Unable to replace undefined with "${context.newMemory}"`,
-        };
-      }
-
-      // Use the new updateWorkingMemory method which handles both thread and resource scope
-      const result = await memory.__experimental_updateWorkingMemoryVNext({
-        threadId,
-        resourceId,
-        workingMemory: workingMemory,
-        searchString: context.searchString,
-        memoryConfig: config,
-      });
-
-      if (result) {
-        return result;
-      }
-
-      return { success: true };
-    },
-  });
-};

package/tsconfig.build.json

@@ -1,9 +0,0 @@
-{
-  "extends": ["./tsconfig.json", "../../tsconfig.build.json"],
-  "compilerOptions": {
-    "outDir": "./dist",
-    "rootDir": "./src"
-  },
-  "include": ["src/**/*"],
-  "exclude": ["node_modules", "**/*.test.ts", "src/**/*.mock.ts"]
-}

package/tsconfig.json

@@ -1,5 +0,0 @@
-{
-  "extends": "../../tsconfig.node.json",
-  "include": ["src/**/*", "tsup.config.ts"],
-  "exclude": ["node_modules", "**/*.test.ts"]
-}

package/tsup.config.ts

@@ -1,17 +0,0 @@
-import { generateTypes } from '@internal/types-builder';
-import { defineConfig } from 'tsup';
-
-export default defineConfig({
-  entry: ['src/index.ts', 'src/processors/index.ts'],
-  format: ['esm', 'cjs'],
-  clean: true,
-  dts: false,
-  splitting: true,
-  treeshake: {
-    preset: 'smallest',
-  },
-  sourcemap: true,
-  onSuccess: async () => {
-    await generateTypes(process.cwd());
-  },
-});

package/vitest.config.ts

@@ -1,11 +0,0 @@
-import { defineConfig } from 'vitest/config';
-
-export default defineConfig({
-  test: {
-    environment: 'node',
-    include: ['src/**/*.test.ts'],
-    // smaller output to save token space when LLMs run tests
-    reporters: 'dot',
-    bail: 1,
-  },
-});