@mastra/memory 0.0.2-alpha.61 → 0.0.2-alpha.63

package/CHANGELOG.md

@@ -1,5 +1,18 @@
 # @mastra/memory
 
+## 0.0.2-alpha.63
+
+### Patch Changes
+
+- Updated dependencies [9fb3039]
+  - @mastra/core@0.1.27-alpha.81
+
+## 0.0.2-alpha.62
+
+### Patch Changes
+
+- 7f5b1b2: @mastra/memory tsup bundling
+
 ## 0.0.2-alpha.61
 
 ### Patch Changes

package/dist/index.d.ts

@@ -1,3 +1,2 @@
-export * from './postgres';
-export * from './kv/upstash';
-//# sourceMappingURL=index.d.ts.map
+
+export {  }

package/dist/index.js

@@ -1,8 +1 @@
 
-'use strict'
-
-if (process.env.NODE_ENV === 'production') {
-  module.exports = require('./memory.cjs.production.min.js')
-} else {
-  module.exports = require('./memory.cjs.development.js')
-}

package/dist/kv/upstash.d.ts

@@ -1,14 +1,15 @@
-import { MastraMemory, MessageType as BaseMastraMessageType, ThreadType, MessageResponse, AiMessageType } from '@mastra/core';
+import { MastraMemory, ThreadType, MessageResponse, AiMessageType, MessageType as MessageType$1 } from '@mastra/core';
 import { Redis } from '@upstash/redis';
 import { ToolResultPart } from 'ai';
-interface MessageType extends BaseMastraMessageType {
+
+interface MessageType extends MessageType$1 {
     tokens?: number;
 }
 interface SerializedThreadType extends Omit<ThreadType, 'createdAt' | 'updatedAt'> {
     createdAt: string;
     updatedAt: string;
 }
-export declare class UpstashKVMemory extends MastraMemory {
+declare class UpstashKVMemory extends MastraMemory {
     private prefix;
     kv: Redis;
     constructor(config: {
@@ -68,5 +69,5 @@ export declare class UpstashKVMemory extends MastraMemory {
     parseThread(thread: SerializedThreadType): ThreadType;
     parseMessages(messages: MessageType[]): MessageType[];
 }
-export {};
-//# sourceMappingURL=upstash.d.ts.map
+
+export { UpstashKVMemory };

package/dist/kv/upstash.js

@@ -0,0 +1,204 @@
+import { MastraMemory } from '@mastra/core';
+import { Redis } from '@upstash/redis';
+import crypto from 'crypto';
+
+// src/kv/upstash.ts
+var UpstashKVMemory = class extends MastraMemory {
+  constructor(config) {
+    super();
+    this.prefix = config.prefix || "mastra";
+    this.MAX_CONTEXT_TOKENS = config.maxTokens;
+    this.kv = new Redis({
+      url: config.url,
+      token: config.token
+    });
+  }
+  getThreadKey(threadId) {
+    return `${this.prefix}:thread:${threadId}`;
+  }
+  getMessagesKey(threadId) {
+    return `${this.prefix}:messages:${threadId}`;
+  }
+  getToolCacheKey(hashedArgs) {
+    return `${this.prefix}:tool:${hashedArgs}`;
+  }
+  async getThreadById({ threadId }) {
+    const thread = await this.kv.get(this.getThreadKey(threadId));
+    return thread ? this.parseThread(thread) : null;
+  }
+  async getThreadsByResourceId({ resourceid }) {
+    const pattern = `${this.prefix}:thread:*`;
+    const keys = await this.kv.keys(pattern);
+    const threads = await Promise.all(keys.map((key) => this.kv.get(key)));
+    return threads.filter((thread) => thread?.resourceid === resourceid).map((thread) => this.parseThread(thread));
+  }
+  async saveThread({ thread }) {
+    const key = this.getThreadKey(thread.id);
+    const serializedThread = {
+      ...thread,
+      createdAt: thread.createdAt.toISOString(),
+      updatedAt: thread.updatedAt.toISOString()
+    };
+    await this.kv.set(key, serializedThread);
+    return thread;
+  }
+  async updateThread(id, title, metadata) {
+    const key = this.getThreadKey(id);
+    const thread = await this.kv.get(key);
+    if (!thread) {
+      throw new Error(`Thread ${id} not found`);
+    }
+    const updatedThread = {
+      ...thread,
+      title,
+      metadata,
+      updatedAt: (/* @__PURE__ */ new Date()).toISOString()
+    };
+    await this.kv.set(key, updatedThread);
+    return this.parseThread(updatedThread);
+  }
+  async deleteThread(id) {
+    await this.kv.del(this.getThreadKey(id));
+    await this.kv.del(this.getMessagesKey(id));
+  }
+  /**
+   * Tool Cache
+   */
+  async validateToolCallArgs({ hashedArgs }) {
+    const cacheKey = this.getToolCacheKey(hashedArgs);
+    const cached = await this.kv.get(cacheKey);
+    return !!cached && new Date(cached.expireAt) > /* @__PURE__ */ new Date();
+  }
+  async getToolResult({
+    threadId,
+    toolArgs,
+    toolName
+  }) {
+    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(cacheKey);
+    if (cached && new Date(cached.expireAt) > /* @__PURE__ */ new Date()) {
+      return cached.result || null;
+    }
+    return null;
+  }
+  async getContextWindow({
+    threadId,
+    startDate,
+    endDate,
+    // @ts-ignore
+    format = "raw"
+  }) {
+    const messagesKey = this.getMessagesKey(threadId);
+    const messages = await this.kv.lrange(messagesKey, 0, -1);
+    let filteredMessages = messages;
+    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 = [];
+      for (const message of filteredMessages.reverse()) {
+        const content = message.role === "assistant" ? message.content[0]?.text || "" : message.content;
+        const tokens = Math.ceil(content.length / 4);
+        if (totalTokens + tokens > this.MAX_CONTEXT_TOKENS) {
+          break;
+        }
+        totalTokens += tokens;
+        messagesWithinTokenLimit.unshift({
+          ...message,
+          tokens
+        });
+      }
+      return this.parseMessages(messagesWithinTokenLimit);
+    }
+    return this.parseMessages(filteredMessages);
+  }
+  /**
+   * Messages
+   */
+  async getMessages({
+    threadId
+  }) {
+    const messagesKey = this.getMessagesKey(threadId);
+    const messages = await this.kv.lrange(messagesKey, 0, -1);
+    const parsedMessages = this.parseMessages(messages);
+    const uiMessages = this.convertToUIMessages(parsedMessages);
+    return { messages: parsedMessages, uiMessages };
+  }
+  async saveMessages({ messages }) {
+    const processedMessages = [];
+    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[0]?.text || "" : message.content;
+        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 * 1e3);
+        for (const hashedArg of hashedToolCallArgs) {
+          const cacheKey = this.getToolCacheKey(hashedArg);
+          await this.kv.set(cacheKey, { expireAt: expireAt.toISOString() });
+        }
+      }
+      await this.kv.rpush(messagesKey, processedMessage);
+      processedMessages.push(processedMessage);
+    }
+    return processedMessages;
+  }
+  async deleteMessage(id) {
+    const pattern = `${this.prefix}:messages:*`;
+    const keys = await this.kv.keys(pattern);
+    for (const key of keys) {
+      const messages = await this.kv.lrange(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() {
+    const pattern = `${this.prefix}:*`;
+    const keys = await this.kv.keys(pattern);
+    if (keys.length > 0) {
+      await this.kv.del(...keys);
+    }
+  }
+  parseThread(thread) {
+    return {
+      ...thread,
+      createdAt: new Date(thread.createdAt),
+      updatedAt: new Date(thread.updatedAt)
+    };
+  }
+  parseMessages(messages) {
+    return messages.map((message) => ({
+      ...message,
+      createdAt: new Date(message.createdAt)
+    }));
+  }
+};
+
+export { UpstashKVMemory };

package/dist/postgres/index.d.ts

@@ -1,6 +1,7 @@
-import { MastraMemory, MessageType, ThreadType, AiMessageType, MessageResponse } from '@mastra/core';
+import { MastraMemory, ThreadType, MessageResponse, MessageType, AiMessageType } from '@mastra/core';
 import { ToolResultPart } from 'ai';
-export declare class PgMemory extends MastraMemory {
+
+declare class PgMemory extends MastraMemory {
     private pool;
     hasTables: boolean;
     constructor(config: {
@@ -57,4 +58,5 @@ export declare class PgMemory extends MastraMemory {
     drop(): Promise<void>;
     ensureTablesExist(): Promise<void>;
 }
-//# sourceMappingURL=index.d.ts.map
+
+export { PgMemory };

package/dist/postgres/index.js

@@ -0,0 +1,414 @@
+import { MastraMemory } from '@mastra/core';
+import crypto from 'crypto';
+import pg from 'pg';
+
+// src/postgres/index.ts
+var { Pool } = pg;
+var PgMemory = class extends MastraMemory {
+  constructor(config) {
+    super();
+    this.hasTables = false;
+    this.pool = new Pool({ connectionString: config.connectionString });
+    this.MAX_CONTEXT_TOKENS = config.maxTokens;
+  }
+  /**
+   * Threads
+   */
+  async getThreadById({ threadId }) {
+    await this.ensureTablesExist();
+    const client = await this.pool.connect();
+    try {
+      const result = await client.query(
+        `
+                SELECT id, title, created_at AS createdAt, updated_at AS updatedAt, resourceid, metadata
+                FROM mastra_threads
+                WHERE id = $1
+            `,
+        [threadId]
+      );
+      return result.rows[0] || null;
+    } finally {
+      client.release();
+    }
+  }
+  async getThreadsByResourceId({ resourceid }) {
+    await this.ensureTablesExist();
+    const client = await this.pool.connect();
+    try {
+      const result = await client.query(
+        `
+                SELECT id, title, resourceid, created_at AS createdAt, updated_at AS updatedAt, metadata
+                FROM mastra_threads
+                WHERE resourceid = $1
+            `,
+        [resourceid]
+      );
+      return result.rows;
+    } finally {
+      client.release();
+    }
+  }
+  async saveThread({ thread }) {
+    await this.ensureTablesExist();
+    const client = await this.pool.connect();
+    try {
+      const { id, title, createdAt, updatedAt, resourceid, metadata } = thread;
+      const result = await client.query(
+        `
+                INSERT INTO mastra_threads (id, title, created_at, updated_at, resourceid, metadata)
+                VALUES ($1, $2, $3, $4, $5, $6)
+                ON CONFLICT (id) DO UPDATE SET title = $2, updated_at = $4, resourceid = $5, metadata = $6
+                RETURNING id, title, created_at AS createdAt, updated_at AS updatedAt, resourceid, metadata
+            `,
+        [id, title, createdAt, updatedAt, resourceid, JSON.stringify(metadata)]
+      );
+      return result?.rows?.[0];
+    } finally {
+      client.release();
+    }
+  }
+  async updateThread(id, title, metadata) {
+    const client = await this.pool.connect();
+    try {
+      const result = await client.query(
+        `
+                UPDATE mastra_threads
+                SET title = $1, metadata = $2, updated_at = NOW()
+                WHERE id = $3
+                RETURNING *
+                `,
+        [title, JSON.stringify(metadata), id]
+      );
+      return result?.rows?.[0];
+    } finally {
+      client.release();
+    }
+  }
+  async deleteThread(id) {
+    const client = await this.pool.connect();
+    try {
+      await client.query(
+        `
+                DELETE FROM mastra_messages
+                WHERE thread_id = $1
+                `,
+        [id]
+      );
+      await client.query(
+        `
+                DELETE FROM mastra_threads
+                WHERE id = $1
+                `,
+        [id]
+      );
+    } finally {
+      client.release();
+    }
+  }
+  /**
+   * Tool Cache
+   */
+  async validateToolCallArgs({ hashedArgs }) {
+    await this.ensureTablesExist();
+    const client = await this.pool.connect();
+    try {
+      const toolArgsResult = await client.query(
+        ` SELECT tool_call_ids as toolCallIds, 
+                tool_call_args as toolCallArgs,
+                created_at AS createdAt
+         FROM mastra_messages
+         WHERE tool_call_args::jsonb @> $1
+         AND tool_call_args_expire_at > $2
+         ORDER BY created_at ASC
+         LIMIT 1`,
+        [JSON.stringify([hashedArgs]), (/* @__PURE__ */ new Date()).toISOString()]
+      );
+      return toolArgsResult.rows.length > 0;
+    } catch (error) {
+      console.log("error checking if valid arg exists====", error);
+      return false;
+    } finally {
+      client.release();
+    }
+  }
+  async getToolResult({
+    threadId,
+    toolArgs,
+    toolName
+  }) {
+    await this.ensureTablesExist();
+    console.log("checking for cached tool result====", JSON.stringify(toolArgs, null, 2));
+    const client = await this.pool.connect();
+    try {
+      const hashedToolArgs = crypto.createHash("sha256").update(JSON.stringify({ args: toolArgs, threadId, toolName })).digest("hex");
+      const toolArgsResult = await client.query(
+        `SELECT tool_call_ids, 
+                tool_call_args,
+                created_at
+         FROM mastra_messages
+         WHERE tool_call_args::jsonb @> $1
+         AND tool_call_args_expire_at > $2
+         ORDER BY created_at ASC
+         LIMIT 1`,
+        [JSON.stringify([hashedToolArgs]), (/* @__PURE__ */ new Date()).toISOString()]
+      );
+      if (toolArgsResult.rows.length > 0) {
+        const toolCallArgs = JSON.parse(toolArgsResult.rows[0]?.tool_call_args);
+        const toolCallIds = JSON.parse(toolArgsResult.rows[0]?.tool_call_ids);
+        const createdAt = toolArgsResult.rows[0]?.created_at;
+        const toolCallArgsIndex = toolCallArgs.findIndex((arg) => arg === hashedToolArgs);
+        const correspondingToolCallId = toolCallIds[toolCallArgsIndex];
+        const toolResult = await client.query(
+          `SELECT content 
+         FROM mastra_messages 
+         WHERE thread_id = $1
+         AND tool_call_ids ILIKE $2
+         AND type = 'tool-result'
+         AND created_at = $3
+         LIMIT 1`,
+          [threadId, `%${correspondingToolCallId}%`, new Date(createdAt).toISOString()]
+        );
+        if (toolResult.rows.length === 0) {
+          console.log("no tool result found");
+          return null;
+        }
+        const toolResultContent = JSON.parse(toolResult.rows[0]?.content);
+        const requiredToolResult = toolResultContent.find((part) => part.toolCallId === correspondingToolCallId);
+        if (requiredToolResult) {
+          return requiredToolResult.result;
+        }
+      }
+      return null;
+    } catch (error) {
+      console.log("error getting cached tool result====", error);
+      return null;
+    } finally {
+      client.release();
+    }
+  }
+  async getContextWindow({
+    threadId,
+    startDate,
+    endDate,
+    format = "raw"
+  }) {
+    await this.ensureTablesExist();
+    const client = await this.pool.connect();
+    try {
+      if (this.MAX_CONTEXT_TOKENS) {
+        const result2 = await client.query(
+          `WITH RankedMessages AS (
+             SELECT *,
+                    SUM(tokens) OVER (ORDER BY created_at DESC) as running_total
+             FROM mastra_messages
+             WHERE thread_id = $1
+             ${startDate ? `AND created_at >= '${startDate.toISOString()}'` : ""}
+             ${endDate ? `AND created_at <= '${endDate.toISOString()}'` : ""}
+             ORDER BY created_at DESC
+           )
+           SELECT id, 
+                  content, 
+                  role, 
+                  type,
+                  created_at AS createdAt, 
+                  thread_id AS threadId
+           FROM RankedMessages
+           WHERE running_total <= $2
+           ORDER BY created_at ASC`,
+          [threadId, this.MAX_CONTEXT_TOKENS]
+        );
+        console.log("Format", format);
+        return this.parseMessages(result2.rows);
+      }
+      const result = await client.query(
+        `SELECT id, 
+                content, 
+                role, 
+                type,
+                created_at AS createdAt, 
+                thread_id AS threadId
+           FROM mastra_messages
+           WHERE thread_id = $1
+           ${startDate ? `AND created_at >= '${startDate.toISOString()}'` : ""}
+           ${endDate ? `AND created_at <= '${endDate.toISOString()}'` : ""}
+           ORDER BY created_at ASC`,
+        [threadId]
+      );
+      console.log("Format", format);
+      return this.parseMessages(result.rows);
+    } catch (error) {
+      console.log("error getting context window====", error);
+      return [];
+    } finally {
+      client.release();
+    }
+  }
+  /**
+   * Messages
+   */
+  async getMessages({
+    threadId
+  }) {
+    await this.ensureTablesExist();
+    const client = await this.pool.connect();
+    try {
+      const result = await client.query(
+        `
+                SELECT 
+                    id, 
+                    content, 
+                    role, 
+                    type,
+                    created_at AS createdAt, 
+                    thread_id AS threadId
+                FROM mastra_messages
+                WHERE thread_id = $1
+                ORDER BY created_at ASC
+            `,
+        [threadId]
+      );
+      const messages = this.parseMessages(result.rows);
+      const uiMessages = this.convertToUIMessages(messages);
+      return { messages, uiMessages };
+    } finally {
+      client.release();
+    }
+  }
+  async saveMessages({ messages }) {
+    await this.ensureTablesExist();
+    const client = await this.pool.connect();
+    try {
+      await client.query("BEGIN");
+      for (const message of messages) {
+        const { id, content, role, createdAt, threadId, toolCallIds, toolCallArgs, type, toolNames } = message;
+        let tokens = null;
+        if (type === "text") {
+          const contentMssg = role === "assistant" ? content[0]?.text || "" : content;
+          tokens = this.estimateTokens(contentMssg);
+        }
+        const hashedToolCallArgs = toolCallArgs ? toolCallArgs.map(
+          (args, index) => crypto.createHash("sha256").update(JSON.stringify({ args, threadId, toolName: toolNames?.[index] })).digest("hex")
+        ) : null;
+        let validArgExists = false;
+        if (hashedToolCallArgs?.length) {
+          validArgExists = true;
+          for (let i = 0; i < hashedToolCallArgs.length; i++) {
+            const isValid = await this.validateToolCallArgs({
+              hashedArgs: hashedToolCallArgs[i]
+            });
+            if (!isValid) {
+              validArgExists = false;
+              break;
+            }
+          }
+        }
+        const toolCallArgsExpireAt = !toolCallArgs ? null : validArgExists ? createdAt : new Date(createdAt.getTime() + 5 * 60 * 1e3);
+        await client.query(
+          `
+                    INSERT INTO mastra_messages (id, content, role, created_at, thread_id, tool_call_ids, tool_call_args, type, tokens, tool_call_args_expire_at)
+                    VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)
+                `,
+          [
+            id,
+            JSON.stringify(content),
+            role,
+            createdAt.toISOString(),
+            threadId,
+            JSON.stringify(toolCallIds),
+            JSON.stringify(hashedToolCallArgs),
+            type,
+            tokens,
+            toolCallArgsExpireAt?.toISOString()
+          ]
+        );
+      }
+      await client.query("COMMIT");
+      return messages;
+    } catch (error) {
+      await client.query("ROLLBACK");
+      throw error;
+    } finally {
+      client.release();
+    }
+  }
+  async deleteMessage(id) {
+    const client = await this.pool.connect();
+    try {
+      await client.query(
+        `
+                DELETE FROM mastra_messages
+                WHERE id = $1
+                `,
+        [id]
+      );
+    } finally {
+      client.release();
+    }
+  }
+  /**
+   * Table Management
+   */
+  async drop() {
+    const client = await this.pool.connect();
+    await client.query("DELETE FROM mastra_messages");
+    await client.query("DELETE FROM mastra_threads");
+    client.release();
+    await this.pool.end();
+  }
+  async ensureTablesExist() {
+    if (this.hasTables) {
+      return;
+    }
+    const client = await this.pool.connect();
+    try {
+      const threadsResult = await client.query(`
+                SELECT EXISTS (
+                    SELECT 1
+                    FROM information_schema.tables
+                    WHERE table_name = 'mastra_threads'
+                );
+            `);
+      if (!threadsResult?.rows?.[0]?.exists) {
+        await client.query(`
+                    CREATE TABLE IF NOT EXISTS mastra_threads (
+                        id UUID PRIMARY KEY,
+                        resourceid TEXT,
+                        title TEXT,
+                        created_at TIMESTAMP WITH TIME ZONE NOT NULL,
+                        updated_at TIMESTAMP WITH TIME ZONE NOT NULL,
+                        metadata JSONB
+                    );
+                `);
+      }
+      const messagesResult = await client.query(`
+                SELECT EXISTS (
+                    SELECT 1
+                    FROM information_schema.tables
+                    WHERE table_name = 'mastra_messages'
+                );
+            `);
+      if (!messagesResult?.rows?.[0]?.exists) {
+        await client.query(`
+                    CREATE TABLE IF NOT EXISTS mastra_messages (
+                        id UUID PRIMARY KEY,
+                        content TEXT NOT NULL,
+                        role VARCHAR(20) NOT NULL,
+                        created_at TIMESTAMP WITH TIME ZONE NOT NULL,
+                        tool_call_ids TEXT DEFAULT NULL,
+                        tool_call_args TEXT DEFAULT NULL,
+                        tool_call_args_expire_at TIMESTAMP WITH TIME ZONE DEFAULT NULL,
+                        type VARCHAR(20) NOT NULL,
+                        tokens INTEGER DEFAULT NULL,
+                        thread_id UUID NOT NULL,
+                        FOREIGN KEY (thread_id) REFERENCES mastra_threads(id)
+                    );
+                `);
+      }
+    } finally {
+      client.release();
+      this.hasTables = true;
+    }
+  }
+};
+
+export { PgMemory };