mirror of
https://github.com/supermemoryai/supermemory.git
synced 2026-05-18 14:48:12 +00:00
549 lines
14 KiB
TypeScript
549 lines
14 KiB
TypeScript
import { describe, it, expect, beforeEach, vi, afterEach } from "vitest"
|
|
import { withSupermemory } from "../src/vercel"
|
|
import { createSupermemoryMiddleware } from "../src/vercel/middleware"
|
|
import type {
|
|
LanguageModelV2,
|
|
LanguageModelV2CallOptions,
|
|
} from "@ai-sdk/provider"
|
|
import Supermemory from "supermemory"
|
|
import "dotenv/config"
|
|
|
|
// Test configuration
|
|
const TEST_CONFIG = {
|
|
apiKey: process.env.SUPERMEMORY_API_KEY || "test-api-key",
|
|
baseURL: process.env.SUPERMEMORY_BASE_URL,
|
|
containerTag: "test-vercel-wrapper",
|
|
}
|
|
|
|
// Mock language model for testing
|
|
const createMockLanguageModel = (): LanguageModelV2 => ({
|
|
specificationVersion: "v2",
|
|
provider: "test-provider",
|
|
modelId: "test-model",
|
|
supportedUrls: {},
|
|
doGenerate: vi.fn(),
|
|
doStream: vi.fn(),
|
|
})
|
|
|
|
// Mock supermemory search response
|
|
const createMockSearchResponse = (contents: string[]) => ({
|
|
results: contents.map((content) => ({
|
|
chunks: [{ content }],
|
|
})),
|
|
})
|
|
|
|
// Helper to call transformParams with proper signature
|
|
const callTransformParams = async (
|
|
middleware: ReturnType<typeof createSupermemoryMiddleware>,
|
|
params: LanguageModelV2CallOptions,
|
|
) => {
|
|
const mockModel = createMockLanguageModel()
|
|
return middleware.transformParams?.({
|
|
type: "generate",
|
|
params,
|
|
model: mockModel,
|
|
})
|
|
}
|
|
|
|
describe("withSupermemory / wrapVercelLanguageModel", () => {
|
|
let originalEnv: string | undefined
|
|
|
|
beforeEach(() => {
|
|
originalEnv = process.env.SUPERMEMORY_API_KEY
|
|
vi.clearAllMocks()
|
|
})
|
|
|
|
afterEach(() => {
|
|
if (originalEnv) {
|
|
process.env.SUPERMEMORY_API_KEY = originalEnv
|
|
} else {
|
|
delete process.env.SUPERMEMORY_API_KEY
|
|
}
|
|
})
|
|
|
|
describe("Environment validation", () => {
|
|
it("should throw error if SUPERMEMORY_API_KEY is not set", () => {
|
|
delete process.env.SUPERMEMORY_API_KEY
|
|
|
|
const mockModel = createMockLanguageModel()
|
|
|
|
expect(() => {
|
|
withSupermemory(mockModel, TEST_CONFIG.containerTag)
|
|
}).toThrow("SUPERMEMORY_API_KEY is not set")
|
|
})
|
|
|
|
it("should successfully create wrapped model with valid API key", () => {
|
|
process.env.SUPERMEMORY_API_KEY = "test-key"
|
|
|
|
const mockModel = createMockLanguageModel()
|
|
const wrappedModel = withSupermemory(mockModel, TEST_CONFIG.containerTag)
|
|
|
|
expect(wrappedModel).toBeDefined()
|
|
expect(wrappedModel.specificationVersion).toBe("v2")
|
|
})
|
|
})
|
|
|
|
describe("createSupermemoryMiddleware", () => {
|
|
// biome-ignore lint/suspicious/noExplicitAny: Mock object for testing
|
|
let mockSupermemory: any
|
|
|
|
beforeEach(() => {
|
|
mockSupermemory = {
|
|
search: {
|
|
execute: vi.fn(),
|
|
},
|
|
}
|
|
})
|
|
|
|
it("should return params unchanged when there is no user message", async () => {
|
|
const middleware = createSupermemoryMiddleware(
|
|
mockSupermemory,
|
|
TEST_CONFIG.containerTag,
|
|
)
|
|
|
|
const params: LanguageModelV2CallOptions = {
|
|
prompt: [
|
|
{
|
|
role: "system",
|
|
content: "You are a helpful assistant",
|
|
},
|
|
],
|
|
}
|
|
|
|
const result = await callTransformParams(middleware, params)
|
|
|
|
expect(result).toEqual(params)
|
|
expect(mockSupermemory.search.execute).not.toHaveBeenCalled()
|
|
})
|
|
|
|
it("should extract last user message with text content", async () => {
|
|
mockSupermemory.search.execute.mockResolvedValue(
|
|
createMockSearchResponse([]),
|
|
)
|
|
|
|
const middleware = createSupermemoryMiddleware(
|
|
mockSupermemory,
|
|
TEST_CONFIG.containerTag,
|
|
)
|
|
|
|
const params: LanguageModelV2CallOptions = {
|
|
prompt: [
|
|
{
|
|
role: "user",
|
|
content: [{ type: "text", text: "Hello, how are you?" }],
|
|
},
|
|
],
|
|
}
|
|
|
|
await callTransformParams(middleware, params)
|
|
|
|
expect(mockSupermemory.search.execute).toHaveBeenCalledWith({
|
|
q: "Hello, how are you?",
|
|
containerTags: [TEST_CONFIG.containerTag],
|
|
})
|
|
})
|
|
|
|
it("should handle multiple user messages and extract the last one", async () => {
|
|
mockSupermemory.search.execute.mockResolvedValue(
|
|
createMockSearchResponse([]),
|
|
)
|
|
|
|
const middleware = createSupermemoryMiddleware(
|
|
mockSupermemory,
|
|
TEST_CONFIG.containerTag,
|
|
)
|
|
|
|
const params: LanguageModelV2CallOptions = {
|
|
prompt: [
|
|
{
|
|
role: "user",
|
|
content: [{ type: "text", text: "First message" }],
|
|
},
|
|
{
|
|
role: "assistant",
|
|
content: [{ type: "text", text: "Response" }],
|
|
},
|
|
{
|
|
role: "user",
|
|
content: [{ type: "text", text: "Last message" }],
|
|
},
|
|
],
|
|
}
|
|
|
|
await callTransformParams(middleware, params)
|
|
|
|
expect(mockSupermemory.search.execute).toHaveBeenCalledWith({
|
|
q: "Last message",
|
|
containerTags: [TEST_CONFIG.containerTag],
|
|
})
|
|
})
|
|
|
|
it("should concatenate multiple text parts in user message", async () => {
|
|
mockSupermemory.search.execute.mockResolvedValue(
|
|
createMockSearchResponse([]),
|
|
)
|
|
|
|
const middleware = createSupermemoryMiddleware(
|
|
mockSupermemory,
|
|
TEST_CONFIG.containerTag,
|
|
)
|
|
|
|
const params: LanguageModelV2CallOptions = {
|
|
prompt: [
|
|
{
|
|
role: "user",
|
|
content: [
|
|
{ type: "text", text: "Part 1" },
|
|
{ type: "text", text: "Part 2" },
|
|
{ type: "text", text: "Part 3" },
|
|
],
|
|
},
|
|
],
|
|
}
|
|
|
|
await callTransformParams(middleware, params)
|
|
|
|
expect(mockSupermemory.search.execute).toHaveBeenCalledWith({
|
|
q: "Part 1 Part 2 Part 3",
|
|
containerTags: [TEST_CONFIG.containerTag],
|
|
})
|
|
})
|
|
|
|
it("should create new system prompt when none exists", async () => {
|
|
mockSupermemory.search.execute.mockResolvedValue(
|
|
createMockSearchResponse([
|
|
"Memory 1: User likes TypeScript",
|
|
"Memory 2: User prefers clean code",
|
|
]),
|
|
)
|
|
|
|
const middleware = createSupermemoryMiddleware(
|
|
mockSupermemory,
|
|
TEST_CONFIG.containerTag,
|
|
)
|
|
|
|
const params: LanguageModelV2CallOptions = {
|
|
prompt: [
|
|
{
|
|
role: "user",
|
|
content: [{ type: "text", text: "Tell me about TypeScript" }],
|
|
},
|
|
],
|
|
}
|
|
|
|
const result = await callTransformParams(middleware, params)
|
|
|
|
expect(result?.prompt).toHaveLength(2)
|
|
expect(result?.prompt[0]?.role).toBe("system")
|
|
expect(result?.prompt[0]?.content).toContain(
|
|
"Memory 1: User likes TypeScript Memory 2: User prefers clean code",
|
|
)
|
|
expect(result?.prompt[1]?.role).toBe("user")
|
|
})
|
|
|
|
it("should append memories to existing system prompt", async () => {
|
|
mockSupermemory.search.execute.mockResolvedValue(
|
|
createMockSearchResponse(["Memory: User is an expert developer"]),
|
|
)
|
|
|
|
const middleware = createSupermemoryMiddleware(
|
|
mockSupermemory,
|
|
TEST_CONFIG.containerTag,
|
|
)
|
|
|
|
const params: LanguageModelV2CallOptions = {
|
|
prompt: [
|
|
{
|
|
role: "system",
|
|
content: "You are a helpful coding assistant",
|
|
},
|
|
{
|
|
role: "user",
|
|
content: [{ type: "text", text: "Help me code" }],
|
|
},
|
|
],
|
|
}
|
|
|
|
const result = await callTransformParams(middleware, params)
|
|
|
|
expect(result?.prompt).toHaveLength(2)
|
|
expect(result?.prompt[0]?.role).toBe("system")
|
|
expect(result?.prompt[0]?.content).toContain(
|
|
"You are a helpful coding assistant",
|
|
)
|
|
expect(result?.prompt[0]?.content).toContain(
|
|
"Memory: User is an expert developer",
|
|
)
|
|
})
|
|
|
|
it("should handle empty memory results", async () => {
|
|
mockSupermemory.search.execute.mockResolvedValue(
|
|
createMockSearchResponse([]),
|
|
)
|
|
|
|
const middleware = createSupermemoryMiddleware(
|
|
mockSupermemory,
|
|
TEST_CONFIG.containerTag,
|
|
)
|
|
|
|
const params: LanguageModelV2CallOptions = {
|
|
prompt: [
|
|
{
|
|
role: "user",
|
|
content: [{ type: "text", text: "Hello" }],
|
|
},
|
|
],
|
|
}
|
|
|
|
const result = await callTransformParams(middleware, params)
|
|
|
|
// Should still create system prompt even if memories are empty
|
|
expect(result?.prompt).toHaveLength(2)
|
|
expect(result?.prompt[0]?.role).toBe("system")
|
|
})
|
|
|
|
it("should filter out non-text content from user message", async () => {
|
|
mockSupermemory.search.execute.mockResolvedValue(
|
|
createMockSearchResponse([]),
|
|
)
|
|
|
|
const middleware = createSupermemoryMiddleware(
|
|
mockSupermemory,
|
|
TEST_CONFIG.containerTag,
|
|
)
|
|
|
|
const params: LanguageModelV2CallOptions = {
|
|
prompt: [
|
|
{
|
|
role: "user",
|
|
content: [
|
|
{ type: "text", text: "Text part" },
|
|
// File part is non-text content
|
|
{ type: "file", data: "base64...", mimeType: "image/png" },
|
|
{ type: "text", text: "Another text part" },
|
|
],
|
|
},
|
|
],
|
|
}
|
|
|
|
await callTransformParams(middleware, params)
|
|
|
|
// Should only extract text content
|
|
expect(mockSupermemory.search.execute).toHaveBeenCalledWith({
|
|
q: "Text part Another text part",
|
|
containerTags: [TEST_CONFIG.containerTag],
|
|
})
|
|
})
|
|
|
|
it("should handle multiple memory chunks correctly", async () => {
|
|
mockSupermemory.search.execute.mockResolvedValue({
|
|
results: [
|
|
{
|
|
chunks: [
|
|
{ content: "Chunk 1" },
|
|
{ content: "Chunk 2" },
|
|
{ content: "Chunk 3" },
|
|
],
|
|
},
|
|
{
|
|
chunks: [{ content: "Chunk 4" }, { content: "Chunk 5" }],
|
|
},
|
|
],
|
|
})
|
|
|
|
const middleware = createSupermemoryMiddleware(
|
|
mockSupermemory,
|
|
TEST_CONFIG.containerTag,
|
|
)
|
|
|
|
const params: LanguageModelV2CallOptions = {
|
|
prompt: [
|
|
{
|
|
role: "user",
|
|
content: [{ type: "text", text: "Query" }],
|
|
},
|
|
],
|
|
}
|
|
|
|
const result = await callTransformParams(middleware, params)
|
|
|
|
const systemContent = result?.prompt[0]?.content as string
|
|
// Chunks from same result should be joined with space
|
|
expect(systemContent).toContain("Chunk 1 Chunk 2 Chunk 3")
|
|
// Results should be joined with newline
|
|
expect(systemContent).toContain("Chunk 4 Chunk 5")
|
|
})
|
|
})
|
|
|
|
describe("Integration with real Supermemory", () => {
|
|
// Skip these tests if no API key is available
|
|
const shouldRunIntegration = !!process.env.SUPERMEMORY_API_KEY
|
|
|
|
it.skipIf(!shouldRunIntegration)(
|
|
"should work with real Supermemory API",
|
|
async () => {
|
|
const supermemory = new Supermemory({
|
|
apiKey: process.env.SUPERMEMORY_API_KEY ?? "",
|
|
baseURL: TEST_CONFIG.baseURL,
|
|
})
|
|
|
|
const middleware = createSupermemoryMiddleware(
|
|
supermemory,
|
|
TEST_CONFIG.containerTag,
|
|
)
|
|
|
|
const params: LanguageModelV2CallOptions = {
|
|
prompt: [
|
|
{
|
|
role: "user",
|
|
content: [{ type: "text", text: "Tell me about programming" }],
|
|
},
|
|
],
|
|
}
|
|
|
|
const result = await callTransformParams(middleware, params)
|
|
|
|
expect(result?.prompt).toBeDefined()
|
|
expect(result?.prompt.length).toBeGreaterThanOrEqual(1)
|
|
},
|
|
)
|
|
|
|
it.skipIf(!shouldRunIntegration)(
|
|
"should create wrapped model and use it",
|
|
async () => {
|
|
process.env.SUPERMEMORY_API_KEY = TEST_CONFIG.apiKey
|
|
|
|
const mockModel = createMockLanguageModel()
|
|
const wrappedModel = withSupermemory(
|
|
mockModel,
|
|
TEST_CONFIG.containerTag,
|
|
)
|
|
|
|
expect(wrappedModel).toBeDefined()
|
|
expect(wrappedModel.provider).toBe("test-provider")
|
|
expect(wrappedModel.modelId).toBe("test-model")
|
|
},
|
|
)
|
|
})
|
|
|
|
describe("Edge cases", () => {
|
|
// biome-ignore lint/suspicious/noExplicitAny: Mock object for testing
|
|
let mockSupermemory: any
|
|
|
|
beforeEach(() => {
|
|
mockSupermemory = {
|
|
search: {
|
|
execute: vi.fn(),
|
|
},
|
|
}
|
|
})
|
|
|
|
it("should handle Supermemory API errors gracefully", async () => {
|
|
mockSupermemory.search.execute.mockRejectedValue(new Error("API Error"))
|
|
|
|
const middleware = createSupermemoryMiddleware(
|
|
mockSupermemory,
|
|
TEST_CONFIG.containerTag,
|
|
)
|
|
|
|
const params: LanguageModelV2CallOptions = {
|
|
prompt: [
|
|
{
|
|
role: "user",
|
|
content: [{ type: "text", text: "Hello" }],
|
|
},
|
|
],
|
|
}
|
|
|
|
await expect(callTransformParams(middleware, params)).rejects.toThrow(
|
|
"API Error",
|
|
)
|
|
})
|
|
|
|
it("should handle empty prompt array", async () => {
|
|
const middleware = createSupermemoryMiddleware(
|
|
mockSupermemory,
|
|
TEST_CONFIG.containerTag,
|
|
)
|
|
|
|
const params: LanguageModelV2CallOptions = {
|
|
prompt: [],
|
|
}
|
|
|
|
const result = await callTransformParams(middleware, params)
|
|
|
|
expect(result).toEqual(params)
|
|
expect(mockSupermemory.search.execute).not.toHaveBeenCalled()
|
|
})
|
|
|
|
it("should handle user message with empty content array", async () => {
|
|
const middleware = createSupermemoryMiddleware(
|
|
mockSupermemory,
|
|
TEST_CONFIG.containerTag,
|
|
)
|
|
|
|
const params: LanguageModelV2CallOptions = {
|
|
prompt: [
|
|
{
|
|
role: "user",
|
|
content: [],
|
|
},
|
|
],
|
|
}
|
|
|
|
const result = await callTransformParams(middleware, params)
|
|
|
|
expect(result).toEqual(params)
|
|
expect(mockSupermemory.search.execute).not.toHaveBeenCalled()
|
|
})
|
|
|
|
it("should use correct container tag", async () => {
|
|
mockSupermemory.search.execute.mockResolvedValue(
|
|
createMockSearchResponse([]),
|
|
)
|
|
|
|
const customTag = "my-custom-project"
|
|
const middleware = createSupermemoryMiddleware(mockSupermemory, customTag)
|
|
|
|
const params: LanguageModelV2CallOptions = {
|
|
prompt: [
|
|
{
|
|
role: "user",
|
|
content: [{ type: "text", text: "Query" }],
|
|
},
|
|
],
|
|
}
|
|
|
|
await callTransformParams(middleware, params)
|
|
|
|
expect(mockSupermemory.search.execute).toHaveBeenCalledWith({
|
|
q: "Query",
|
|
containerTags: [customTag],
|
|
})
|
|
})
|
|
|
|
it("should not mutate the original params.prompt array", async () => {
|
|
mockSupermemory.search.execute.mockResolvedValue(
|
|
createMockSearchResponse([]),
|
|
)
|
|
|
|
const middleware = createSupermemoryMiddleware(
|
|
mockSupermemory,
|
|
TEST_CONFIG.containerTag,
|
|
)
|
|
|
|
const originalPrompt = [
|
|
{ role: "user" as const, content: [{ type: "text" as const, text: "First" }] },
|
|
{ role: "user" as const, content: [{ type: "text" as const, text: "Last" }] }
|
|
]
|
|
const params: LanguageModelV2CallOptions = { prompt: [...originalPrompt] }
|
|
|
|
await callTransformParams(middleware, params)
|
|
|
|
// Verify order is unchanged
|
|
expect(params.prompt[0]?.content[0]).toBe("First")
|
|
expect(params.prompt[1]?.content[0]).toBe("Last")
|
|
})
|
|
|
|
})
|
|
})
|