@mastra/memory 0.0.2-alpha.8 → 0.1.0-alpha.65

package/package.json

@@ -1,22 +1,29 @@
 {
   "name": "@mastra/memory",
-  "version": "0.0.2-alpha.8",
+  "version": "0.1.0-alpha.65",
   "description": "",
   "type": "module",
-  "main": "dist/index.js",
-  "module": "dist/memory.esm.js",
-  "types": "dist/index.d.ts",
+  "main": "./dist/index.js",
+  "types": "./dist/index.d.ts",
   "exports": {
     ".": {
       "import": {
-        "types": "./dist/index.d.ts",
-        "default": "./dist/memory.esm.js"
-      },
-      "require": {
         "types": "./dist/index.d.ts",
         "default": "./dist/index.js"
       }
     },
+    "./kv-upstash": {
+      "import": {
+        "types": "./dist/kv/upstash.d.ts",
+        "default": "./dist/kv/upstash.js"
+      }
+    },
+    "./postgres": {
+      "import": {
+        "types": "./dist/postgres/index.d.ts",
+        "default": "./dist/postgres/index.js"
+      }
+    },
     "./package.json": "./package.json"
   },
   "keywords": [],
@@ -29,7 +36,7 @@
     "pg-pool": "^3.7.0",
     "postgres": "^3.4.5",
     "redis": "^4.7.0",
-    "@mastra/core": "0.1.27-alpha.29"
+    "@mastra/core": "0.2.0-alpha.83"
   },
   "devDependencies": {
     "@babel/preset-env": "^7.26.0",
@@ -37,13 +44,13 @@
     "@tsconfig/recommended": "^1.0.7",
     "@types/node": "^22.9.0",
     "@types/pg": "^8.11.10",
-    "dts-cli": "^2.0.5",
-    "jest": "^29.7.0",
-    "ts-jest": "^29.2.5"
+    "tsup": "^8.0.1",
+    "vitest": "^3.0.4"
   },
   "scripts": {
-    "build": "dts build",
-    "build:dev": "dts watch",
-    "test": "node --experimental-vm-modules node_modules/jest/bin/jest.js"
+    "build": "tsup src/index.ts src/postgres/index.ts src/kv/upstash.ts --format esm --dts --clean --treeshake",
+    "dev": "tsup src/index.ts src/postgres/index.ts src/kv/upstash.ts --format esm --dts --clean --watch",
+    "test": "vitest",
+    "test:watch": "vitest watch"
   }
 }

package/src/index.test.ts

@@ -0,0 +1,7 @@
+import { describe, it, expect } from 'vitest';
+
+describe('temp', () => {
+  it(`temp`, () => {
+    expect(true).toBe(true);
+  });
+});

package/src/index.ts

@@ -1,2 +1,170 @@
-export * from './postgres';
-export * from './kv/upstash';
+import { CoreMessage } from '@mastra/core';
+import { MastraMemory, MessageType, MemoryConfig, SharedMemoryConfig, StorageThreadType } from '@mastra/core/memory';
+import { StorageGetMessagesArg } from '@mastra/core/storage';
+import { Message as AiMessage } from 'ai';
+
+/**
+ * Concrete implementation of MastraMemory that adds support for thread configuration
+ * and message injection.
+ */
+export class Memory extends MastraMemory {
+  constructor(config: SharedMemoryConfig) {
+    super({ name: 'Memory', ...config });
+  }
+
+  async getMessages({
+    threadId,
+    selectBy,
+    threadConfig,
+  }: StorageGetMessagesArg): Promise<{ messages: CoreMessage[]; uiMessages: AiMessage[] }> {
+    let vectorResults:
+      | null
+      | {
+          id: string;
+          score: number;
+          metadata?: Record<string, any>;
+          vector?: number[];
+        }[] = null;
+
+    this.logger.info(`Memory getMessages() with:`, {
+      threadId,
+      selectBy,
+      threadConfig,
+    });
+
+    const config = this.getMergedThreadConfig(threadConfig || {});
+
+    const vectorConfig =
+      typeof config?.injectVectorHistorySearch === `boolean`
+        ? {
+            includeResults: 2,
+            includePrevious: 2,
+            includeNext: 2,
+          }
+        : {
+            includeResults: config?.injectVectorHistorySearch?.includeResults || 2,
+            includePrevious: config?.injectVectorHistorySearch?.includePrevious || 2,
+            includeNext: config?.injectVectorHistorySearch?.includeNext || 2,
+          };
+
+    if (selectBy?.vectorSearchString && this.vector) {
+      const { embeddings } = await this.vector.embed(selectBy.vectorSearchString, this.parseEmbeddingOptions());
+
+      await this.vector.createIndex('memory_messages', 1536);
+      vectorResults = await this.vector.query('memory_messages', embeddings[0]!, vectorConfig.includeResults, {
+        thread_id: threadId,
+      });
+    }
+
+    // Get raw messages from storage
+    const rawMessages = await this.storage.__getMessages({
+      threadId,
+      selectBy: {
+        ...selectBy,
+        ...(vectorResults?.length
+          ? {
+              include: vectorResults.map(r => ({
+                id: r.metadata?.message_id,
+                withNextMessages: vectorConfig.includeNext,
+                withPreviousMessages: vectorConfig.includePrevious,
+              })),
+            }
+          : {}),
+      },
+      threadConfig: config,
+    });
+
+    // Parse and convert messages
+    const messages = this.parseMessages(rawMessages);
+    const uiMessages = this.convertToUIMessages(rawMessages);
+
+    return { messages, uiMessages };
+  }
+
+  async rememberMessages({
+    threadId,
+    vectorMessageSearch,
+    config,
+  }: {
+    threadId: string;
+    vectorMessageSearch?: string;
+    config?: MemoryConfig;
+  }) {
+    const threadConfig = this.getMergedThreadConfig(config || {});
+
+    if (!threadConfig.injectRecentMessages && !threadConfig.injectVectorHistorySearch) {
+      return {
+        messages: [],
+        uiMessages: [],
+      } satisfies Awaited<ReturnType<typeof this.getMessages>>;
+    }
+
+    const messages = await this.getMessages({
+      threadId,
+      selectBy: {
+        last: threadConfig.injectRecentMessages,
+        vectorSearchString:
+          threadConfig.injectVectorHistorySearch && vectorMessageSearch ? vectorMessageSearch : undefined,
+      },
+      threadConfig: config,
+    });
+
+    this.logger.info(`Remembered message history includes ${messages.messages.length} messages.`);
+    return messages;
+  }
+
+  async getThreadById({ threadId }: { threadId: string }): Promise<StorageThreadType | null> {
+    return this.storage.__getThreadById({ threadId });
+  }
+
+  async getThreadsByResourceId({ resourceId }: { resourceId: string }): Promise<StorageThreadType[]> {
+    return this.storage.__getThreadsByResourceId({ resourceId });
+  }
+
+  async saveThread({ thread }: { thread: StorageThreadType }): Promise<StorageThreadType> {
+    return this.storage.__saveThread({ thread });
+  }
+
+  async updateThread({
+    id,
+    title,
+    metadata,
+  }: {
+    id: string;
+    title: string;
+    metadata: Record<string, unknown>;
+  }): Promise<StorageThreadType> {
+    return this.storage.__updateThread({
+      id,
+      title,
+      metadata,
+    });
+  }
+
+  async saveMessages({ messages }: { messages: MessageType[] }): Promise<MessageType[]> {
+    if (this.vector) {
+      await this.vector.createIndex('memory_messages', 1536);
+      for (const message of messages) {
+        if (typeof message.content !== `string`) continue;
+        const { embeddings } = await this.vector.embed(message.content, this.parseEmbeddingOptions());
+        await this.vector.upsert('memory_messages', embeddings, [
+          {
+            text: message.content,
+            message_id: message.id,
+            thread_id: message.threadId,
+          },
+        ]);
+      }
+    }
+    return this.storage.__saveMessages({ messages });
+  }
+
+  async deleteThread(threadId: string): Promise<void> {
+    await this.storage.__deleteThread({ threadId });
+
+    // TODO: Also clean up vector storage if it exists
+    // if (this.vector) {
+    //   await this.vector.deleteThread(threadId); ?? filter by thread attributes and delete all returned messages?
+    // }
+  }
+}

package/vitest.config.ts

@@ -0,0 +1,8 @@
+import { defineConfig } from 'vitest/config';
+
+export default defineConfig({
+  test: {
+    environment: 'node',
+    include: ['src/**/*.test.ts'],
+  },
+});

package/dist/kv/upstash.d.ts

@@ -1,71 +0,0 @@
-import { MastraMemory, MessageType as BaseMastraMessageType, ThreadType, MessageResponse } from '@mastra/core';
-import { Redis } from '@upstash/redis';
-import { ToolResultPart, Message as AiMessage } from 'ai';
-interface MessageType extends BaseMastraMessageType {
-    tokens?: number;
-}
-interface SerializedThreadType extends Omit<ThreadType, 'createdAt' | 'updatedAt'> {
-    createdAt: string;
-    updatedAt: string;
-}
-export declare class UpstashKVMemory extends MastraMemory {
-    private prefix;
-    kv: Redis;
-    constructor(config: {
-        url: string;
-        token: string;
-        prefix?: string;
-        maxTokens?: number;
-    });
-    private getThreadKey;
-    private getMessagesKey;
-    private getToolCacheKey;
-    getThreadById({ threadId }: {
-        threadId: string;
-    }): Promise<ThreadType | null>;
-    getThreadsByResourceId({ resourceid }: {
-        resourceid: string;
-    }): Promise<ThreadType[]>;
-    saveThread({ thread }: {
-        thread: ThreadType;
-    }): Promise<ThreadType>;
-    updateThread(id: string, title: string, metadata: Record<string, unknown>): Promise<ThreadType>;
-    deleteThread(id: string): Promise<void>;
-    /**
-     * Tool Cache
-     */
-    validateToolCallArgs({ hashedArgs }: {
-        hashedArgs: string;
-    }): Promise<boolean>;
-    getToolResult({ threadId, toolArgs, toolName, }: {
-        threadId: string;
-        toolArgs: Record<string, unknown>;
-        toolName: string;
-    }): Promise<ToolResultPart['result'] | null>;
-    getContextWindow<T extends 'raw' | 'core_message'>({ threadId, startDate, endDate, format, }: {
-        format?: T;
-        threadId: string;
-        startDate?: Date;
-        endDate?: Date;
-    }): Promise<MessageResponse<T>>;
-    /**
-     * Messages
-     */
-    getMessages({ threadId }: {
-        threadId: string;
-    }): Promise<{
-        messages: MessageType[];
-        uiMessages: AiMessage[];
-    }>;
-    saveMessages({ messages }: {
-        messages: MessageType[];
-    }): Promise<MessageType[]>;
-    deleteMessage(id: string): Promise<void>;
-    /**
-     * Cleanup
-     */
-    drop(): Promise<void>;
-    parseThread(thread: SerializedThreadType): ThreadType;
-    parseMessages(messages: MessageType[]): MessageType[];
-}
-export {};