@mastra/memory 0.0.2-alpha.2 → 0.0.2-alpha.21

package/src/kv/upstash.test.ts

@@ -0,0 +1,253 @@
+import { MessageType, ThreadType } from '@mastra/core';
+import { randomUUID } from 'crypto';
+import dotenv from 'dotenv';
+
+import { UpstashKVMemory } from './upstash';
+
+dotenv.config();
+
+// Ensure environment variables are set
+if (!process.env.KV_REST_API_URL || !process.env.KV_REST_API_TOKEN) {
+  throw new Error('Required Vercel KV environment variables are not set');
+}
+
+describe('KVMemory Integration Tests', () => {
+  let memory: UpstashKVMemory;
+  const testPrefix = `test_${Date.now()}`;
+
+  beforeAll(() => {
+    memory = new UpstashKVMemory({
+      url: process.env.KV_REST_API_URL!,
+      token: process.env.KV_REST_API_TOKEN!,
+      prefix: testPrefix,
+      maxTokens: 1000,
+    });
+  });
+
+  afterAll(async () => {
+    await memory.drop();
+  });
+
+  describe('Thread Operations', () => {
+    let testThread: ThreadType;
+
+    beforeEach(() => {
+      testThread = {
+        id: randomUUID(),
+        title: 'Integration Test Thread',
+        createdAt: new Date(),
+        updatedAt: new Date(),
+        resourceid: 'test-resource',
+        metadata: { test: true },
+      };
+    });
+
+    it('should create and retrieve a thread', async () => {
+      const saved = await memory.saveThread({ thread: testThread });
+      expect(saved).toEqual(testThread);
+
+      const retrieved = await memory.getThreadById({ threadId: testThread.id });
+      expect(retrieved).toEqual(testThread);
+    });
+
+    it('should find threads by resource ID', async () => {
+      const thread1 = { ...testThread, id: randomUUID() };
+      const thread2 = {
+        ...testThread,
+        id: randomUUID(),
+        resourceid: 'different-resource',
+      };
+
+      await memory.saveThread({ thread: thread1 });
+      await memory.saveThread({ thread: thread2 });
+
+      const threads = await memory.getThreadsByResourceId({
+        resourceid: 'test-resource',
+      });
+
+      expect(threads.length).toBeGreaterThanOrEqual(1);
+      expect(threads.some(t => t.id === thread1.id)).toBe(true);
+      expect(threads.some(t => t.id === thread2.id)).toBe(false);
+    });
+
+    it('should update thread title and metadata', async () => {
+      await memory.saveThread({ thread: testThread });
+
+      const updatedTitle = 'Updated Title';
+      const updatedMetadata = { updated: true };
+
+      const updated = await memory.updateThread(testThread.id, updatedTitle, updatedMetadata);
+
+      expect(updated.title).toBe(updatedTitle);
+      expect(updated.metadata).toEqual(updatedMetadata);
+      expect(updated.updatedAt).not.toEqual(testThread.updatedAt);
+    });
+
+    it('should delete a thread', async () => {
+      await memory.saveThread({ thread: testThread });
+      await memory.deleteThread(testThread.id);
+
+      const retrieved = await memory.getThreadById({
+        threadId: testThread.id,
+      });
+      expect(retrieved).toBeNull();
+    });
+  });
+
+  describe('Message Operations', () => {
+    let testThread: ThreadType;
+    let testMessage: MessageType;
+
+    beforeEach(async () => {
+      testThread = {
+        id: randomUUID(),
+        title: 'Test Thread for Messages',
+        createdAt: new Date(),
+        updatedAt: new Date(),
+        resourceid: 'test-resource',
+        metadata: {},
+      };
+
+      testMessage = {
+        id: randomUUID(),
+        content: 'Test message content',
+        role: 'user' as const,
+        type: 'text' as const,
+        createdAt: new Date(),
+        threadId: testThread.id,
+      };
+
+      await memory.saveThread({ thread: testThread });
+    });
+
+    it('should save and retrieve messages', async () => {
+      await memory.saveMessages({ messages: [testMessage] });
+
+      const { messages, uiMessages } = await memory.getMessages({
+        threadId: testThread.id,
+      });
+
+      expect(messages.length).toBe(1);
+      expect(messages?.[0]?.content).toBe(testMessage.content);
+      expect(uiMessages.length).toBe(1);
+    });
+
+    it('should handle message deletion', async () => {
+      const message1 = { ...testMessage, id: randomUUID() };
+      const message2 = { ...testMessage, id: randomUUID() };
+
+      await memory.saveMessages({ messages: [message1, message2] });
+      await memory.deleteMessage(message1.id);
+
+      const { messages } = await memory.getMessages({
+        threadId: testThread.id,
+      });
+
+      expect(messages.length).toBe(1);
+      expect(messages?.[0]?.id).toBe(message2.id);
+    });
+
+    it('should respect token limits in context window', async () => {
+      // Create a very long message that should exceed our small token limit
+      const longMessage: MessageType = {
+        ...testMessage,
+        content: 'a'.repeat(100), // This should be about 25 tokens
+        id: randomUUID(),
+      };
+
+      const shortMessage: MessageType = {
+        ...testMessage,
+        content: 'Short message', // This should be about 3 tokens
+        id: randomUUID(),
+      };
+
+      await memory.saveMessages({ messages: [longMessage, shortMessage] });
+
+      // Create a new memory instance with a very small token limit
+      const lowTokenMemory = new UpstashKVMemory({
+        url: process.env.KV_REST_API_URL!,
+        token: process.env.KV_REST_API_TOKEN!,
+        prefix: testPrefix,
+        maxTokens: 5, // Only allow 5 tokens
+      });
+
+      const context = await lowTokenMemory.getContextWindow({
+        threadId: testThread.id,
+        format: 'raw',
+      });
+
+      // Should only get the short message
+      expect(context.length).toBe(1);
+      expect(context?.[0]?.id).toBe(shortMessage.id);
+    });
+
+    it('should filter messages by date range', async () => {
+      const oldMessage: MessageType = {
+        ...testMessage,
+        id: randomUUID(),
+        createdAt: new Date('2023-01-01'),
+      };
+
+      const newMessage: MessageType = {
+        ...testMessage,
+        id: randomUUID(),
+        createdAt: new Date('2024-01-01'),
+      };
+
+      await memory.saveMessages({ messages: [oldMessage, newMessage] });
+
+      const context = await memory.getContextWindow({
+        threadId: testThread.id,
+        format: 'raw',
+        startDate: new Date('2023-12-31'),
+        endDate: new Date('2024-12-31'),
+      });
+
+      expect(context.length).toBe(1);
+      expect(context?.[0]?.id).toBe(newMessage.id);
+    });
+  });
+
+  describe('Tool Cache Operations', () => {
+    let testThread: ThreadType;
+
+    beforeEach(async () => {
+      testThread = {
+        id: randomUUID(),
+        title: 'Test Thread for Tool Cache',
+        createdAt: new Date(),
+        updatedAt: new Date(),
+        resourceid: 'test-resource',
+        metadata: {},
+      };
+
+      await memory.saveThread({ thread: testThread });
+    });
+
+    it('should cache and validate tool call arguments', async () => {
+      const toolArgs = { test: true };
+      const toolName = 'testTool';
+
+      const message: MessageType = {
+        id: randomUUID(),
+        content: 'Tool test',
+        role: 'assistant' as const,
+        type: 'text' as const,
+        createdAt: new Date(),
+        threadId: testThread.id,
+        toolCallArgs: [toolArgs],
+        toolNames: [toolName],
+      };
+
+      await memory.saveMessages({ messages: [message] });
+
+      const result = await memory.getToolResult({
+        threadId: testThread.id,
+        toolArgs,
+        toolName,
+      });
+
+      expect(result).toBeNull();
+    });
+  });
+});

package/src/kv/upstash.ts

@@ -0,0 +1,298 @@
+import { MastraMemory, MessageType as BaseMastraMessageType, ThreadType, MessageResponse } from '@mastra/core';
+import { Redis } from '@upstash/redis';
+import { ToolResultPart, Message as AiMessage, TextPart } from 'ai';
+import crypto from 'crypto';
+
+interface ToolCacheData {
+  expireAt: string;
+  result?: ToolResultPart['result'];
+}
+
+interface MessageType extends BaseMastraMessageType {
+  tokens?: number;
+}
+
+// Internal type for serialized thread data
+interface SerializedThreadType extends Omit<ThreadType, 'createdAt' | 'updatedAt'> {
+  createdAt: string;
+  updatedAt: string;
+}
+
+export class UpstashKVMemory extends MastraMemory {
+  private prefix: string;
+
+  kv: Redis;
+
+  constructor(config: { url: string; token: string; prefix?: string; maxTokens?: number }) {
+    super();
+    this.prefix = config.prefix || 'mastra';
+    this.MAX_CONTEXT_TOKENS = config.maxTokens;
+
+    this.kv = new Redis({
+      url: config.url,
+      token: config.token,
+    });
+  }
+
+  private getThreadKey(threadId: string): string {
+    return `${this.prefix}:thread:${threadId}`;
+  }
+
+  private getMessagesKey(threadId: string): string {
+    return `${this.prefix}:messages:${threadId}`;
+  }
+
+  private getToolCacheKey(hashedArgs: string): string {
+    return `${this.prefix}:tool:${hashedArgs}`;
+  }
+
+  async getThreadById({ threadId }: { threadId: string }): Promise<ThreadType | null> {
+    const thread = await this.kv.get<SerializedThreadType>(this.getThreadKey(threadId));
+    return thread ? this.parseThread(thread) : null;
+  }
+
+  async getThreadsByResourceId({ resourceid }: { resourceid: string }): Promise<ThreadType[]> {
+    const pattern = `${this.prefix}:thread:*`;
+    const keys = await this.kv.keys(pattern);
+
+    const threads = await Promise.all(keys.map(key => this.kv.get<SerializedThreadType>(key)));
+
+    return threads
+      .filter(thread => thread?.resourceid === resourceid)
+      .map(thread => this.parseThread(thread as SerializedThreadType));
+  }
+
+  async saveThread({ thread }: { thread: ThreadType }): Promise<ThreadType> {
+    const key = this.getThreadKey(thread.id);
+    const serializedThread: SerializedThreadType = {
+      ...thread,
+      createdAt: thread.createdAt.toISOString(),
+      updatedAt: thread.updatedAt.toISOString(),
+    };
+    await this.kv.set(key, serializedThread);
+    return thread;
+  }
+
+  async updateThread(id: string, title: string, metadata: Record<string, unknown>): Promise<ThreadType> {
+    const key = this.getThreadKey(id);
+    const thread = await this.kv.get<SerializedThreadType>(key);
+
+    if (!thread) {
+      throw new Error(`Thread ${id} not found`);
+    }
+
+    const updatedThread: SerializedThreadType = {
+      ...thread,
+      title,
+      metadata,
+      updatedAt: new Date().toISOString(),
+    };
+
+    await this.kv.set(key, updatedThread);
+    return this.parseThread(updatedThread);
+  }
+
+  async deleteThread(id: string): Promise<void> {
+    await this.kv.del(this.getThreadKey(id));
+    await this.kv.del(this.getMessagesKey(id));
+  }
+
+  /**
+   * Tool Cache
+   */
+
+  async validateToolCallArgs({ hashedArgs }: { hashedArgs: string }): Promise<boolean> {
+    const cacheKey = this.getToolCacheKey(hashedArgs);
+    const cached = await this.kv.get<ToolCacheData>(cacheKey);
+    return !!cached && new Date(cached.expireAt) > new Date();
+  }
+
+  async getToolResult({
+    threadId,
+    toolArgs,
+    toolName,
+  }: {
+    threadId: string;
+    toolArgs: Record<string, unknown>;
+    toolName: string;
+  }): Promise<ToolResultPart['result'] | null> {
+    const hashedToolArgs = crypto
+      .createHash('sha256')
+      .update(JSON.stringify({ args: toolArgs, threadId, toolName }))
+      .digest('hex');
+
+    const cacheKey = this.getToolCacheKey(hashedToolArgs);
+    const cached = await this.kv.get<ToolCacheData>(cacheKey);
+
+    if (cached && new Date(cached.expireAt) > new Date()) {
+      return cached.result || null;
+    }
+
+    return null;
+  }
+
+  async getContextWindow<T extends 'raw' | 'core_message'>({
+    threadId,
+    startDate,
+    endDate,
+    // @ts-ignore
+    format = 'raw' as T,
+  }: {
+    format?: T;
+    threadId: string;
+    startDate?: Date;
+    endDate?: Date;
+  }) {
+    const messagesKey = this.getMessagesKey(threadId);
+    const messages = await this.kv.lrange<MessageType>(messagesKey, 0, -1);
+
+    let filteredMessages = messages.filter(msg => msg.type === 'text' || msg.type === 'tool-result');
+
+    if (startDate) {
+      filteredMessages = filteredMessages.filter(msg => new Date(msg.createdAt) >= startDate);
+    }
+
+    if (endDate) {
+      filteredMessages = filteredMessages.filter(msg => new Date(msg.createdAt) <= endDate);
+    }
+
+    if (this.MAX_CONTEXT_TOKENS) {
+      let totalTokens = 0;
+      const messagesWithinTokenLimit: MessageType[] = [];
+
+      // Process messages from newest to oldest
+      for (const message of filteredMessages.reverse()) {
+        const content =
+          message.role === 'assistant'
+            ? (message.content as Array<TextPart>)[0]?.text || ''
+            : (message.content as string);
+
+        // Use a more aggressive token estimation
+        // Roughly estimate 1 token per 4 characters
+        const tokens = Math.ceil(content.length / 4);
+
+        // Check if adding this message would exceed the token limit
+        if (totalTokens + tokens > this.MAX_CONTEXT_TOKENS) {
+          break;
+        }
+
+        totalTokens += tokens;
+        messagesWithinTokenLimit.unshift({
+          ...message,
+          tokens,
+        });
+      }
+
+      // Return messages in chronological order
+      return this.parseMessages(messagesWithinTokenLimit) as MessageResponse<T>;
+    }
+
+    return this.parseMessages(filteredMessages) as MessageResponse<T>;
+  }
+
+  /**
+   * Messages
+   */
+
+  async getMessages({ threadId }: { threadId: string }): Promise<{ messages: MessageType[]; uiMessages: AiMessage[] }> {
+    const messagesKey = this.getMessagesKey(threadId);
+    const messages = await this.kv.lrange<MessageType>(messagesKey, 0, -1);
+    const parsedMessages = this.parseMessages(messages);
+    const uiMessages = this.convertToUIMessages(parsedMessages);
+
+    return { messages: parsedMessages, uiMessages };
+  }
+
+  async saveMessages({ messages }: { messages: MessageType[] }): Promise<MessageType[]> {
+    const processedMessages: MessageType[] = [];
+
+    for (const message of messages) {
+      const { threadId, toolCallArgs, toolNames, createdAt } = message;
+      const messagesKey = this.getMessagesKey(threadId);
+
+      const processedMessage = { ...message };
+
+      if (message.type === 'text') {
+        const content =
+          message.role === 'assistant'
+            ? (message.content as Array<TextPart>)[0]?.text || ''
+            : (message.content as string);
+        processedMessage.tokens = this.estimateTokens(content);
+      }
+
+      if (toolCallArgs?.length) {
+        const hashedToolCallArgs = toolCallArgs.map((args, index) =>
+          crypto
+            .createHash('sha256')
+            .update(JSON.stringify({ args, threadId, toolName: toolNames?.[index] }))
+            .digest('hex'),
+        );
+
+        let validArgExists = true;
+        for (const hashedArg of hashedToolCallArgs) {
+          const isValid = await this.validateToolCallArgs({ hashedArgs: hashedArg });
+          if (!isValid) {
+            validArgExists = false;
+            break;
+          }
+        }
+
+        const expireAt = validArgExists ? createdAt : new Date(createdAt.getTime() + 5 * 60 * 1000); // 5 minutes
+
+        for (const hashedArg of hashedToolCallArgs) {
+          const cacheKey = this.getToolCacheKey(hashedArg);
+          await this.kv.set(cacheKey, { expireAt: expireAt.toISOString() } as ToolCacheData);
+        }
+      }
+
+      await this.kv.rpush(messagesKey, processedMessage);
+      processedMessages.push(processedMessage);
+    }
+
+    return processedMessages;
+  }
+
+  async deleteMessage(id: string): Promise<void> {
+    const pattern = `${this.prefix}:messages:*`;
+    const keys = await this.kv.keys(pattern);
+
+    for (const key of keys) {
+      const messages = await this.kv.lrange<MessageType>(key, 0, -1);
+      const filteredMessages = messages.filter(msg => msg.id !== id);
+
+      if (messages.length !== filteredMessages.length) {
+        await this.kv.del(key);
+        if (filteredMessages.length > 0) {
+          await this.kv.rpush(key, ...filteredMessages);
+        }
+      }
+    }
+  }
+
+  /**
+   * Cleanup
+   */
+
+  async drop(): Promise<void> {
+    const pattern = `${this.prefix}:*`;
+    const keys = await this.kv.keys(pattern);
+    if (keys.length > 0) {
+      await this.kv.del(...keys);
+    }
+  }
+
+  parseThread(thread: SerializedThreadType): ThreadType {
+    return {
+      ...thread,
+      createdAt: new Date(thread.createdAt),
+      updatedAt: new Date(thread.updatedAt),
+    };
+  }
+
+  parseMessages(messages: MessageType[]): MessageType[] {
+    return messages.map(message => ({
+      ...message,
+      createdAt: new Date(message.createdAt),
+    }));
+  }
+}