diff --git a/packages/tools/src/vercel/index.ts b/packages/tools/src/vercel/index.ts index ffd44603..582b46af 100644 --- a/packages/tools/src/vercel/index.ts +++ b/packages/tools/src/vercel/index.ts @@ -15,6 +15,8 @@ import type { PromptTemplate, MemoryPromptData } from "./memory-prompt" const DEFAULT_MEMORY_RETRIEVAL_TIMEOUT_MS = 5000 interface WrapVercelLanguageModelOptions { + /** The container tag/identifier for memory search (e.g., user ID, project ID) */ + containerTag: string /** Optional conversation ID to group messages for contextual memory generation */ conversationId?: string /** Enable detailed logging of memory search and injection */ @@ -73,8 +75,8 @@ interface WrapVercelLanguageModelOptions { * detection of `model.specificationVersion`. * * @param model - The language model to wrap with supermemory capabilities (V2 or V3) - * @param containerTag - The container tag/identifier for memory search (e.g., user ID, project ID) - * @param options - Optional configuration options for the middleware + * @param options - Configuration options for Supermemory integration + * @param options.containerTag - Required. The container tag/identifier for memory search (e.g., user ID, project ID) * @param options.conversationId - Optional conversation ID to group messages into a single document for contextual memory generation * @param options.verbose - Optional flag to enable detailed logging of memory search and injection process (default: false) * @param options.mode - Optional mode for memory search: "profile", "query", or "full" (default: "profile") @@ -90,7 +92,8 @@ interface WrapVercelLanguageModelOptions { * import { withSupermemory } from "@supermemory/tools/ai-sdk" * import { openai } from "@ai-sdk/openai" * - * const modelWithMemory = withSupermemory(openai("gpt-4"), "user-123", { + * const modelWithMemory = withSupermemory(openai("gpt-4"), { + * containerTag: "user-123", * conversationId: "conversation-456", * mode: "full", * addMemory: "always" @@ -107,10 +110,9 @@ interface WrapVercelLanguageModelOptions { */ const wrapVercelLanguageModel = ( model: T, - containerTag: string, - options?: WrapVercelLanguageModelOptions, + options: WrapVercelLanguageModelOptions, ): T => { - const providedApiKey = options?.apiKey ?? process.env.SUPERMEMORY_API_KEY + const providedApiKey = options.apiKey ?? process.env.SUPERMEMORY_API_KEY if (!providedApiKey) { throw new Error( @@ -119,18 +121,18 @@ const wrapVercelLanguageModel = ( } const ctx = createSupermemoryContext({ - containerTag, + containerTag: options.containerTag, apiKey: providedApiKey, - conversationId: options?.conversationId, - verbose: options?.verbose ?? false, - mode: options?.mode ?? "profile", - addMemory: options?.addMemory ?? "never", - baseUrl: options?.baseUrl, - promptTemplate: options?.promptTemplate, + conversationId: options.conversationId, + verbose: options.verbose ?? false, + mode: options.mode ?? "profile", + addMemory: options.addMemory ?? "never", + baseUrl: options.baseUrl, + promptTemplate: options.promptTemplate, memoryRetrievalTimeoutMs: DEFAULT_MEMORY_RETRIEVAL_TIMEOUT_MS, }) - const skipMemoryOnError = options?.skipMemoryOnError ?? true + const skipMemoryOnError = options.skipMemoryOnError ?? true // Proxy keeps prototype/getter fields (e.g. provider, modelId) that `{ ...model }` drops. return new Proxy(model, { diff --git a/packages/tools/test/ai-sdk-test.ts b/packages/tools/test/ai-sdk-test.ts index 3d9f13ea..19bb272d 100644 --- a/packages/tools/test/ai-sdk-test.ts +++ b/packages/tools/test/ai-sdk-test.ts @@ -2,7 +2,8 @@ import { generateText } from "ai" import { withSupermemory } from "../src/ai-sdk" import { openai } from "@ai-sdk/openai" -const modelWithMemory = withSupermemory(openai("gpt-5"), "user_id_life", { +const modelWithMemory = withSupermemory(openai("gpt-5"), { + containerTag: "user_id_life", verbose: true, mode: "query", // options are profile, query, full (default is profile) addMemory: "always", // options are always, never (default is never) diff --git a/packages/tools/test/chatapp/app/api/chat/route.ts b/packages/tools/test/chatapp/app/api/chat/route.ts index c8f7a8be..ded72745 100644 --- a/packages/tools/test/chatapp/app/api/chat/route.ts +++ b/packages/tools/test/chatapp/app/api/chat/route.ts @@ -1,7 +1,8 @@ import { gateway, streamText, type ModelMessage } from "ai" import { withSupermemory } from "@supermemory/tools/ai-sdk" -const model = withSupermemory(gateway("google/gemini-2.5-flash"), "user-1", { +const model = withSupermemory(gateway("google/gemini-2.5-flash"), { + containerTag: "user-1", apiKey: process.env.SUPERMEMORY_API_KEY ?? "", mode: "full", addMemory: "always", diff --git a/packages/tools/test/chatapp/app/api/stream/route.ts b/packages/tools/test/chatapp/app/api/stream/route.ts index 5889a549..a86c3941 100644 --- a/packages/tools/test/chatapp/app/api/stream/route.ts +++ b/packages/tools/test/chatapp/app/api/stream/route.ts @@ -8,6 +8,7 @@ const SUPERMEMORY_USER_ID = "user-1" const gatewayModel = gateway("google/gemini-2.5-flash") const supermemoryOptions = { + containerTag: SUPERMEMORY_USER_ID, apiKey: process.env.SUPERMEMORY_API_KEY ?? "", mode: "full" as const, addMemory: "always" as const, @@ -31,11 +32,7 @@ export async function POST(req: Request) { }) : gatewayModel - const model = withSupermemory( - innerModel, - SUPERMEMORY_USER_ID, - supermemoryOptions, - ) + const model = withSupermemory(innerModel, supermemoryOptions) const result = streamText({ model, diff --git a/packages/tools/test/with-supermemory/integration.test.ts b/packages/tools/test/with-supermemory/integration.test.ts index 74263c19..13ea59c6 100644 --- a/packages/tools/test/with-supermemory/integration.test.ts +++ b/packages/tools/test/with-supermemory/integration.test.ts @@ -96,14 +96,11 @@ describe.skipIf(!shouldRunIntegration)( const { model, getCapturedGenerateParams } = createIntegrationMockModel() - const wrapped = withSupermemory( - model, - INTEGRATION_CONFIG.containerTag, - { - apiKey: INTEGRATION_CONFIG.apiKey, - mode: "profile", - }, - ) + const wrapped = withSupermemory(model, { + containerTag: INTEGRATION_CONFIG.containerTag, + apiKey: INTEGRATION_CONFIG.apiKey, + mode: "profile", + }) await wrapped.doGenerate({ prompt: [ @@ -127,16 +124,13 @@ describe.skipIf(!shouldRunIntegration)( const conversationId = `test-generate-${Date.now()}` - const wrapped = withSupermemory( - model, - INTEGRATION_CONFIG.containerTag, - { - apiKey: INTEGRATION_CONFIG.apiKey, - mode: "profile", - addMemory: "always", - conversationId, - }, - ) + const wrapped = withSupermemory(model, { + containerTag: INTEGRATION_CONFIG.containerTag, + apiKey: INTEGRATION_CONFIG.apiKey, + mode: "profile", + addMemory: "always", + conversationId, + }) await wrapped.doGenerate({ prompt: [ @@ -172,15 +166,12 @@ describe.skipIf(!shouldRunIntegration)( const conversationId = `test-conversation-${Date.now()}` - const wrapped = withSupermemory( - model, - INTEGRATION_CONFIG.containerTag, - { - apiKey: INTEGRATION_CONFIG.apiKey, - mode: "profile", - conversationId, - }, - ) + const wrapped = withSupermemory(model, { + containerTag: INTEGRATION_CONFIG.containerTag, + apiKey: INTEGRATION_CONFIG.apiKey, + mode: "profile", + conversationId, + }) await wrapped.doGenerate({ prompt: [ @@ -203,14 +194,11 @@ describe.skipIf(!shouldRunIntegration)( it("should fetch memories and stream response", async () => { const { model, getCapturedStreamParams } = createIntegrationMockModel() - const wrapped = withSupermemory( - model, - INTEGRATION_CONFIG.containerTag, - { - apiKey: INTEGRATION_CONFIG.apiKey, - mode: "profile", - }, - ) + const wrapped = withSupermemory(model, { + containerTag: INTEGRATION_CONFIG.containerTag, + apiKey: INTEGRATION_CONFIG.apiKey, + mode: "profile", + }) const { stream } = await wrapped.doStream({ prompt: [ @@ -242,16 +230,13 @@ describe.skipIf(!shouldRunIntegration)( const conversationId = `test-stream-${Date.now()}` - const wrapped = withSupermemory( - model, - INTEGRATION_CONFIG.containerTag, - { - apiKey: INTEGRATION_CONFIG.apiKey, - mode: "profile", - addMemory: "always", - conversationId, - }, - ) + const wrapped = withSupermemory(model, { + containerTag: INTEGRATION_CONFIG.containerTag, + apiKey: INTEGRATION_CONFIG.apiKey, + mode: "profile", + addMemory: "always", + conversationId, + }) const { stream } = await wrapped.doStream({ prompt: [ @@ -286,14 +271,11 @@ describe.skipIf(!shouldRunIntegration)( it("should handle text-delta chunks correctly", async () => { const { model } = createIntegrationMockModel() - const wrapped = withSupermemory( - model, - INTEGRATION_CONFIG.containerTag, - { - apiKey: INTEGRATION_CONFIG.apiKey, - mode: "profile", - }, - ) + const wrapped = withSupermemory(model, { + containerTag: INTEGRATION_CONFIG.containerTag, + apiKey: INTEGRATION_CONFIG.apiKey, + mode: "profile", + }) const { stream } = await wrapped.doStream({ prompt: [ @@ -327,14 +309,11 @@ describe.skipIf(!shouldRunIntegration)( const { model } = createIntegrationMockModel() const fetchSpy = vi.spyOn(globalThis, "fetch") - const wrapped = withSupermemory( - model, - INTEGRATION_CONFIG.containerTag, - { - apiKey: INTEGRATION_CONFIG.apiKey, - mode: "profile", - }, - ) + const wrapped = withSupermemory(model, { + containerTag: INTEGRATION_CONFIG.containerTag, + apiKey: INTEGRATION_CONFIG.apiKey, + mode: "profile", + }) await wrapped.doGenerate({ prompt: [ @@ -368,14 +347,11 @@ describe.skipIf(!shouldRunIntegration)( const { model } = createIntegrationMockModel() const fetchSpy = vi.spyOn(globalThis, "fetch") - const wrapped = withSupermemory( - model, - INTEGRATION_CONFIG.containerTag, - { - apiKey: INTEGRATION_CONFIG.apiKey, - mode: "query", - }, - ) + const wrapped = withSupermemory(model, { + containerTag: INTEGRATION_CONFIG.containerTag, + apiKey: INTEGRATION_CONFIG.apiKey, + mode: "query", + }) await wrapped.doGenerate({ prompt: [ @@ -409,14 +385,11 @@ describe.skipIf(!shouldRunIntegration)( const { model } = createIntegrationMockModel() const fetchSpy = vi.spyOn(globalThis, "fetch") - const wrapped = withSupermemory( - model, - INTEGRATION_CONFIG.containerTag, - { - apiKey: INTEGRATION_CONFIG.apiKey, - mode: "full", - }, - ) + const wrapped = withSupermemory(model, { + containerTag: INTEGRATION_CONFIG.containerTag, + apiKey: INTEGRATION_CONFIG.apiKey, + mode: "full", + }) await wrapped.doGenerate({ prompt: [ @@ -456,15 +429,12 @@ describe.skipIf(!shouldRunIntegration)( generalSearchMemories: string }) => `${data.userMemories}` - const wrapped = withSupermemory( - model, - INTEGRATION_CONFIG.containerTag, - { - apiKey: INTEGRATION_CONFIG.apiKey, - mode: "profile", - promptTemplate: customTemplate, - }, - ) + const wrapped = withSupermemory(model, { + containerTag: INTEGRATION_CONFIG.containerTag, + apiKey: INTEGRATION_CONFIG.apiKey, + mode: "profile", + promptTemplate: customTemplate, + }) await wrapped.doGenerate({ prompt: [ @@ -485,15 +455,12 @@ describe.skipIf(!shouldRunIntegration)( const { model, getCapturedGenerateParams } = createIntegrationMockModel() - const wrapped = withSupermemory( - model, - INTEGRATION_CONFIG.containerTag, - { - apiKey: INTEGRATION_CONFIG.apiKey, - mode: "profile", - verbose: true, - }, - ) + const wrapped = withSupermemory(model, { + containerTag: INTEGRATION_CONFIG.containerTag, + apiKey: INTEGRATION_CONFIG.apiKey, + mode: "profile", + verbose: true, + }) await wrapped.doGenerate({ prompt: [ @@ -514,15 +481,12 @@ describe.skipIf(!shouldRunIntegration)( const fetchSpy = vi.spyOn(globalThis, "fetch") // Use the configured base URL (or default) - const wrapped = withSupermemory( - model, - INTEGRATION_CONFIG.containerTag, - { - apiKey: INTEGRATION_CONFIG.apiKey, - mode: "profile", - baseUrl: INTEGRATION_CONFIG.baseUrl, - }, - ) + const wrapped = withSupermemory(model, { + containerTag: INTEGRATION_CONFIG.containerTag, + apiKey: INTEGRATION_CONFIG.apiKey, + mode: "profile", + baseUrl: INTEGRATION_CONFIG.baseUrl, + }) await wrapped.doGenerate({ prompt: [ @@ -556,14 +520,11 @@ describe.skipIf(!shouldRunIntegration)( new Error("Model error"), ) - const wrapped = withSupermemory( - model, - INTEGRATION_CONFIG.containerTag, - { - apiKey: INTEGRATION_CONFIG.apiKey, - mode: "profile", - }, - ) + const wrapped = withSupermemory(model, { + containerTag: INTEGRATION_CONFIG.containerTag, + apiKey: INTEGRATION_CONFIG.apiKey, + mode: "profile", + }) await expect( wrapped.doGenerate({ @@ -581,14 +542,11 @@ describe.skipIf(!shouldRunIntegration)( const { model, getCapturedGenerateParams } = createIntegrationMockModel() - const wrapped = withSupermemory( - model, - INTEGRATION_CONFIG.containerTag, - { - apiKey: "invalid-api-key-12345", - mode: "profile", - }, - ) + const wrapped = withSupermemory(model, { + containerTag: INTEGRATION_CONFIG.containerTag, + apiKey: "invalid-api-key-12345", + mode: "profile", + }) await wrapped.doGenerate({ prompt: [ @@ -606,15 +564,12 @@ describe.skipIf(!shouldRunIntegration)( it("should reject on invalid API key when skipMemoryOnError is false", async () => { const { model } = createIntegrationMockModel() - const wrapped = withSupermemory( - model, - INTEGRATION_CONFIG.containerTag, - { - apiKey: "invalid-api-key-12345", - mode: "profile", - skipMemoryOnError: false, - }, - ) + const wrapped = withSupermemory(model, { + containerTag: INTEGRATION_CONFIG.containerTag, + apiKey: "invalid-api-key-12345", + mode: "profile", + skipMemoryOnError: false, + }) await expect( wrapped.doGenerate({ diff --git a/packages/tools/test/with-supermemory/unit.test.ts b/packages/tools/test/with-supermemory/unit.test.ts index dcefc54e..8461fb97 100644 --- a/packages/tools/test/with-supermemory/unit.test.ts +++ b/packages/tools/test/with-supermemory/unit.test.ts @@ -73,7 +73,7 @@ describe("Unit: withSupermemory", () => { const mockModel = createMockLanguageModel() expect(() => { - withSupermemory(mockModel, TEST_CONFIG.containerTag) + withSupermemory(mockModel, { containerTag: TEST_CONFIG.containerTag }) }).toThrow("SUPERMEMORY_API_KEY is not set") }) @@ -81,7 +81,9 @@ describe("Unit: withSupermemory", () => { process.env.SUPERMEMORY_API_KEY = "test-key" const mockModel = createMockLanguageModel() - const wrappedModel = withSupermemory(mockModel, TEST_CONFIG.containerTag) + const wrappedModel = withSupermemory(mockModel, { + containerTag: TEST_CONFIG.containerTag, + }) expect(wrappedModel).toBeDefined() expect(wrappedModel.specificationVersion).toBe("v2") @@ -99,7 +101,9 @@ describe("Unit: withSupermemory", () => { doStream: vi.fn(), } const inner = Object.create(proto) as LanguageModelV2 - const wrappedModel = withSupermemory(inner, TEST_CONFIG.containerTag) + const wrappedModel = withSupermemory(inner, { + containerTag: TEST_CONFIG.containerTag, + }) expect(wrappedModel.specificationVersion).toBe("v2") expect(wrappedModel.provider).toBe("gateway") @@ -414,7 +418,8 @@ describe("Unit: withSupermemory", () => { warnings: [], }) - const wrapped = withSupermemory(inner, TEST_CONFIG.containerTag, { + const wrapped = withSupermemory(inner, { + containerTag: TEST_CONFIG.containerTag, apiKey: "k", }) @@ -436,7 +441,8 @@ describe("Unit: withSupermemory", () => { }) const inner = createMockLanguageModel() - const wrapped = withSupermemory(inner, TEST_CONFIG.containerTag, { + const wrapped = withSupermemory(inner, { + containerTag: TEST_CONFIG.containerTag, apiKey: "k", skipMemoryOnError: false, }) @@ -475,7 +481,8 @@ describe("Unit: withSupermemory", () => { warnings: [], }) - const wrapped = withSupermemory(inner, TEST_CONFIG.containerTag, { + const wrapped = withSupermemory(inner, { + containerTag: TEST_CONFIG.containerTag, apiKey: "k", })