mirror of
https://github.com/supermemoryai/supermemory.git
synced 2026-08-28 05:25:33 +00:00
feat: ai sdk language model withSupermemory (#446)
This commit is contained in:
MaheshtheDev
parent
20c41706e1
commit
35ac9e086b
12 changed files with 7264 additions and 9 deletions
10
.github/workflows/claude-code-review.yml
vendored
10
.github/workflows/claude-code-review.yml
vendored
|
|
@ -17,14 +17,14 @@ jobs:
|
|||
# github.event.pull_request.user.login == 'external-contributor' ||
|
||||
# github.event.pull_request.user.login == 'new-developer' ||
|
||||
# github.event.pull_request.author_association == 'FIRST_TIME_CONTRIBUTOR'
|
||||
|
||||
|
||||
runs-on: ubuntu-latest
|
||||
permissions:
|
||||
contents: read
|
||||
pull-requests: read
|
||||
issues: read
|
||||
id-token: write
|
||||
|
||||
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@v4
|
||||
|
|
@ -36,6 +36,7 @@ jobs:
|
|||
uses: anthropics/claude-code-action@v1
|
||||
with:
|
||||
anthropic_api_key: ${{ secrets.ANTHROPIC_API_KEY }}
|
||||
use_sticky_comment: true
|
||||
prompt: |
|
||||
Please review this pull request and provide feedback on:
|
||||
- Code quality and best practices
|
||||
|
|
@ -43,12 +44,11 @@ jobs:
|
|||
- Performance considerations
|
||||
- Security concerns
|
||||
- Test coverage
|
||||
|
||||
|
||||
Use the repository's CLAUDE.md for guidance on style and conventions. Be constructive and helpful in your feedback.
|
||||
|
||||
Use `gh pr comment` with your Bash tool to leave your review as a comment on the PR.
|
||||
|
||||
|
||||
# See https://github.com/anthropics/claude-code-action/blob/main/docs/usage.md
|
||||
# or https://docs.anthropic.com/en/docs/claude-code/sdk#command-line for available options
|
||||
claude_args: '--allowed-tools "Bash(gh issue view:*),Bash(gh search:*),Bash(gh issue list:*),Bash(gh pr comment:*),Bash(gh pr diff:*),Bash(gh pr view:*),Bash(gh pr list:*)"'
|
||||
|
||||
|
|
|
|||
|
|
@ -22,13 +22,13 @@
|
|||
"@ai-sdk/anthropic": "^1.2.12",
|
||||
"@ai-sdk/cerebras": "^0.2.16",
|
||||
"@ai-sdk/google": "^1.2.22",
|
||||
"@ai-sdk/openai": "^1.3.23",
|
||||
"@ai-sdk/openai": "^2.0.42",
|
||||
"@anthropic-ai/sdk": "^0.55.1",
|
||||
"@google/genai": "^1.10.0",
|
||||
"@google/generative-ai": "^0.24.1",
|
||||
"@hono/zod-validator": "^0.7.1",
|
||||
"@scalar/hono-api-reference": "^0.9.11",
|
||||
"ai": "^4.3.19",
|
||||
"ai": "^5.0.59",
|
||||
"alchemy": "^0.55.2",
|
||||
"atmn": "^0.0.16",
|
||||
"better-auth": "^1.3.3",
|
||||
|
|
@ -41,7 +41,6 @@
|
|||
"drizzle-zod": "~0.7.1",
|
||||
"file-type": "^21.0.0",
|
||||
"hono-openapi": "^0.4.8",
|
||||
|
||||
"nanoid": "^5.1.5",
|
||||
"neverthrow": "^8.2.0",
|
||||
"pg": "^8.16.3",
|
||||
|
|
|
|||
|
|
@ -60,6 +60,124 @@ const addTool = addMemoryTool(process.env.SUPERMEMORY_API_KEY!, {
|
|||
})
|
||||
```
|
||||
|
||||
#### AI SDK Middleware with Supermemory
|
||||
|
||||
> [!CAUTION]
|
||||
> `withSupermemory` is in beta
|
||||
|
||||
- `withSupermemory` will take advantage supermemory profile v4 endpoint personalized based on container tag
|
||||
- Make sure you have `SUPERMEMORY_API_KEY` in env
|
||||
|
||||
```typescript
|
||||
import { generateText } from "ai"
|
||||
import { withSupermemory } from "@supermemory/tools/ai-sdk"
|
||||
import { openai } from "@ai-sdk/openai"
|
||||
|
||||
const modelWithMemory = withSupermemory(openai("gpt-5"), "user_id_life")
|
||||
|
||||
const result = await generateText({
|
||||
model: modelWithMemory,
|
||||
messages: [{ role: "user", content: "where do i live?" }],
|
||||
})
|
||||
|
||||
console.log(result.text)
|
||||
```
|
||||
|
||||
#### Verbose Mode
|
||||
|
||||
Enable verbose logging to see detailed information about memory search and transformation:
|
||||
|
||||
```typescript
|
||||
import { generateText } from "ai"
|
||||
import { withSupermemory } from "@supermemory/tools/ai-sdk"
|
||||
import { openai } from "@ai-sdk/openai"
|
||||
|
||||
const modelWithMemory = withSupermemory(openai("gpt-5"), "user_id_life", {
|
||||
verbose: true
|
||||
})
|
||||
|
||||
const result = await generateText({
|
||||
model: modelWithMemory,
|
||||
messages: [{ role: "user", content: "where do i live?" }],
|
||||
})
|
||||
|
||||
console.log(result.text)
|
||||
```
|
||||
|
||||
When verbose mode is enabled, you'll see console output like:
|
||||
```
|
||||
[supermemory] Searching memories for container: user_id_life
|
||||
[supermemory] User message: where do i live?
|
||||
[supermemory] System prompt exists: false
|
||||
[supermemory] Found 3 memories
|
||||
[supermemory] Memory content: You live in San Francisco, California. Your address is 123 Main Street...
|
||||
[supermemory] Creating new system prompt with memories
|
||||
```
|
||||
|
||||
#### Memory Search Modes
|
||||
|
||||
The middleware supports different modes for memory retrieval:
|
||||
|
||||
**Profile Mode (Default)** - Retrieves user profile memories without query filtering:
|
||||
```typescript
|
||||
import { generateText } from "ai"
|
||||
import { withSupermemory } from "@supermemory/tools/ai-sdk"
|
||||
import { openai } from "@ai-sdk/openai"
|
||||
|
||||
// Uses profile mode by default - gets all user profile memories
|
||||
const modelWithMemory = withSupermemory(openai("gpt-4"), "user-123")
|
||||
|
||||
// Explicitly specify profile mode
|
||||
const modelWithProfile = withSupermemory(openai("gpt-4"), "user-123", {
|
||||
mode: "profile"
|
||||
})
|
||||
|
||||
const result = await generateText({
|
||||
model: modelWithMemory,
|
||||
messages: [{ role: "user", content: "What do you know about me?" }],
|
||||
})
|
||||
```
|
||||
|
||||
**Query Mode** - Searches memories based on the user's message:
|
||||
```typescript
|
||||
import { generateText } from "ai"
|
||||
import { withSupermemory } from "@supermemory/tools/ai-sdk"
|
||||
import { openai } from "@ai-sdk/openai"
|
||||
|
||||
const modelWithQuery = withSupermemory(openai("gpt-4"), "user-123", {
|
||||
mode: "query"
|
||||
})
|
||||
|
||||
const result = await generateText({
|
||||
model: modelWithQuery,
|
||||
messages: [{ role: "user", content: "What's my favorite programming language?" }],
|
||||
})
|
||||
```
|
||||
|
||||
**Full Mode** - Combines both profile and query results:
|
||||
```typescript
|
||||
import { generateText } from "ai"
|
||||
import { withSupermemory } from "@supermemory/tools/ai-sdk"
|
||||
import { openai } from "@ai-sdk/openai"
|
||||
|
||||
const modelWithFull = withSupermemory(openai("gpt-4"), "user-123", {
|
||||
mode: "full"
|
||||
})
|
||||
|
||||
const result = await generateText({
|
||||
model: modelWithFull,
|
||||
messages: [{ role: "user", content: "Tell me about my preferences" }],
|
||||
})
|
||||
```
|
||||
|
||||
**Combined Options** - Use verbose logging with specific modes:
|
||||
```typescript
|
||||
const modelWithOptions = withSupermemory(openai("gpt-4"), "user-123", {
|
||||
mode: "profile",
|
||||
verbose: true
|
||||
})
|
||||
```
|
||||
|
||||
### OpenAI Function Calling Usage
|
||||
|
||||
```typescript
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
{
|
||||
"name": "@supermemory/tools",
|
||||
"type": "module",
|
||||
"version": "1.1.12",
|
||||
"version": "1.2.0",
|
||||
"description": "Memory tools for AI SDK and OpenAI function calling with supermemory",
|
||||
"scripts": {
|
||||
"build": "tsdown",
|
||||
|
|
@ -11,6 +11,7 @@
|
|||
"test:watch": "vitest --watch --testTimeout 100000"
|
||||
},
|
||||
"dependencies": {
|
||||
"@ai-sdk/anthropic": "^2.0.25",
|
||||
"@ai-sdk/openai": "^2.0.23",
|
||||
"@ai-sdk/provider": "^2.0.0",
|
||||
"ai": "^5.0.29",
|
||||
|
|
|
|||
|
|
@ -119,3 +119,5 @@ export function supermemoryTools(
|
|||
addMemory: addMemoryTool(apiKey, config),
|
||||
}
|
||||
}
|
||||
|
||||
export { withSupermemory } from "./vercel"
|
||||
|
|
|
|||
59
packages/tools/src/vercel/index.ts
Normal file
59
packages/tools/src/vercel/index.ts
Normal file
|
|
@ -0,0 +1,59 @@
|
|||
import type { LanguageModelV2 } from "@ai-sdk/provider"
|
||||
import { wrapLanguageModel } from "ai"
|
||||
import { createSupermemoryMiddleware } from "./middleware"
|
||||
|
||||
/**
|
||||
* Wraps a language model with supermemory middleware to automatically inject relevant memories
|
||||
* into the system prompt based on the user's message content.
|
||||
*
|
||||
* This middleware searches the supermemory API for relevant memories using the container tag
|
||||
* and user message, then either appends memories to an existing system prompt or creates
|
||||
* a new system prompt with the memories.
|
||||
*
|
||||
* @param model - The language model to wrap with supermemory capabilities
|
||||
* @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.verbose - Optional flag to enable detailed logging of memory search and injection process (default: false)
|
||||
* @param options.mode - Optional mode for memory search: "profile" (default), "query", or "full"
|
||||
*
|
||||
* @returns A wrapped language model that automatically includes relevant memories in prompts
|
||||
*
|
||||
* @example
|
||||
* ```typescript
|
||||
* import { withSupermemory } from "@supermemory/tools/ai-sdk"
|
||||
* import { openai } from "@ai-sdk/openai"
|
||||
*
|
||||
* const modelWithMemory = withSupermemory(openai("gpt-4"), "user-123")
|
||||
*
|
||||
* const result = await generateText({
|
||||
* model: modelWithMemory,
|
||||
* messages: [{ role: "user", content: "What's my favorite programming language?" }]
|
||||
* })
|
||||
* ```
|
||||
*
|
||||
* @throws {Error} When SUPERMEMORY_API_KEY environment variable is not set
|
||||
* @throws {Error} When supermemory API request fails
|
||||
*/
|
||||
const wrapVercelLanguageModel = (
|
||||
model: LanguageModelV2,
|
||||
containerTag: string,
|
||||
options?: { verbose?: boolean; mode?: "profile" | "query" | "full" },
|
||||
): LanguageModelV2 => {
|
||||
const SUPERMEMORY_API_KEY = process.env.SUPERMEMORY_API_KEY
|
||||
|
||||
if (!SUPERMEMORY_API_KEY) {
|
||||
throw new Error("SUPERMEMORY_API_KEY is not set")
|
||||
}
|
||||
|
||||
const verbose = options?.verbose ?? false
|
||||
const mode = options?.mode ?? "profile"
|
||||
|
||||
const wrappedModel = wrapLanguageModel({
|
||||
model,
|
||||
middleware: createSupermemoryMiddleware(containerTag, verbose, mode),
|
||||
})
|
||||
|
||||
return wrappedModel
|
||||
}
|
||||
|
||||
export { wrapVercelLanguageModel as withSupermemory }
|
||||
44
packages/tools/src/vercel/logger.ts
Normal file
44
packages/tools/src/vercel/logger.ts
Normal file
|
|
@ -0,0 +1,44 @@
|
|||
export interface Logger {
|
||||
debug: (message: string, data?: unknown) => void
|
||||
info: (message: string, data?: unknown) => void
|
||||
warn: (message: string, data?: unknown) => void
|
||||
error: (message: string, data?: unknown) => void
|
||||
}
|
||||
|
||||
export const createLogger = (verbose: boolean): Logger => {
|
||||
if (!verbose) {
|
||||
return {
|
||||
debug: () => {},
|
||||
info: () => {},
|
||||
warn: () => {},
|
||||
error: () => {},
|
||||
}
|
||||
}
|
||||
|
||||
return {
|
||||
debug: (message: string, data?: unknown) => {
|
||||
console.log(
|
||||
`[supermemory] ${message}`,
|
||||
data ? JSON.stringify(data, null, 2) : "",
|
||||
)
|
||||
},
|
||||
info: (message: string, data?: unknown) => {
|
||||
console.log(
|
||||
`[supermemory] ${message}`,
|
||||
data ? JSON.stringify(data, null, 2) : "",
|
||||
)
|
||||
},
|
||||
warn: (message: string, data?: unknown) => {
|
||||
console.warn(
|
||||
`[supermemory] ${message}`,
|
||||
data ? JSON.stringify(data, null, 2) : "",
|
||||
)
|
||||
},
|
||||
error: (message: string, data?: unknown) => {
|
||||
console.error(
|
||||
`[supermemory] ${message}`,
|
||||
data ? JSON.stringify(data, null, 2) : "",
|
||||
)
|
||||
},
|
||||
}
|
||||
}
|
||||
171
packages/tools/src/vercel/middleware.ts
Normal file
171
packages/tools/src/vercel/middleware.ts
Normal file
|
|
@ -0,0 +1,171 @@
|
|||
import type {
|
||||
LanguageModelV2CallOptions,
|
||||
LanguageModelV2Middleware,
|
||||
LanguageModelV2Message,
|
||||
} from "@ai-sdk/provider"
|
||||
import { createLogger, type Logger } from "./logger"
|
||||
import { convertProfileToMarkdown, type ProfileStructure } from "./util"
|
||||
|
||||
const getLastUserMessage = (params: LanguageModelV2CallOptions) => {
|
||||
const lastUserMessage = params.prompt
|
||||
.reverse()
|
||||
.find((prompt: LanguageModelV2Message) => prompt.role === "user")
|
||||
const memories = lastUserMessage?.content
|
||||
.filter((content) => content.type === "text")
|
||||
.map((content) => content.text)
|
||||
.join(" ")
|
||||
return memories
|
||||
}
|
||||
|
||||
const supermemoryprofilesearch = async (
|
||||
containerTag: string,
|
||||
queryText: string,
|
||||
): Promise<ProfileStructure> => {
|
||||
const SUPERMEMORY_API_KEY = process.env.SUPERMEMORY_API_KEY
|
||||
|
||||
if (!SUPERMEMORY_API_KEY) {
|
||||
throw new Error("SUPERMEMORY_API_KEY is not set")
|
||||
}
|
||||
|
||||
const payload = queryText
|
||||
? JSON.stringify({
|
||||
q: queryText,
|
||||
containerTag: containerTag,
|
||||
})
|
||||
: JSON.stringify({
|
||||
containerTag: containerTag,
|
||||
})
|
||||
|
||||
try {
|
||||
const response = await fetch("https://api.supermemory.ai/v4/profile", {
|
||||
method: "POST",
|
||||
headers: {
|
||||
"Content-Type": "application/json",
|
||||
Authorization: `Bearer ${SUPERMEMORY_API_KEY}`,
|
||||
},
|
||||
body: payload,
|
||||
})
|
||||
|
||||
if (!response.ok) {
|
||||
const errorText = await response.text().catch(() => "Unknown error")
|
||||
throw new Error(
|
||||
`Supermemory profile search failed: ${response.status} ${response.statusText}. ${errorText}`,
|
||||
)
|
||||
}
|
||||
|
||||
return await response.json()
|
||||
} catch (error) {
|
||||
if (error instanceof Error) {
|
||||
throw error
|
||||
}
|
||||
throw new Error(`Supermemory API request failed: ${error}`)
|
||||
}
|
||||
}
|
||||
|
||||
const addSystemPrompt = async (
|
||||
params: LanguageModelV2CallOptions,
|
||||
containerTag: string,
|
||||
logger: Logger,
|
||||
mode: "profile" | "query" | "full",
|
||||
) => {
|
||||
const systemPromptExists = params.prompt.some(
|
||||
(prompt) => prompt.role === "system",
|
||||
)
|
||||
|
||||
const queryText =
|
||||
mode !== "profile"
|
||||
? params.prompt
|
||||
.reverse()
|
||||
.find((prompt) => prompt.role === "user")
|
||||
?.content?.filter((content) => content.type === "text")
|
||||
?.map((content) => (content.type === "text" ? content.text : ""))
|
||||
?.join(" ") || ""
|
||||
: ""
|
||||
|
||||
const memoriesResponse = await supermemoryprofilesearch(
|
||||
containerTag,
|
||||
queryText,
|
||||
)
|
||||
|
||||
const memoryCountStatic = memoriesResponse.profile.static?.length || 0
|
||||
const memoryCountDynamic = memoriesResponse.profile.dynamic?.length || 0
|
||||
|
||||
logger.info("Memory search completed", {
|
||||
containerTag,
|
||||
memoryCountStatic,
|
||||
memoryCountDynamic,
|
||||
queryText:
|
||||
queryText.substring(0, 100) + (queryText.length > 100 ? "..." : ""),
|
||||
mode,
|
||||
})
|
||||
|
||||
const profileData =
|
||||
mode !== "query" ? convertProfileToMarkdown(memoriesResponse) : ""
|
||||
const searchResultsMemories =
|
||||
mode !== "profile"
|
||||
? `Search results for user's recent message: \n${memoriesResponse.searchResults.results
|
||||
.map((result) => `- ${result.memory}`)
|
||||
.join("\n")}`
|
||||
: ""
|
||||
|
||||
const memories = `${profileData}\n${searchResultsMemories}`.trim()
|
||||
if (memories) {
|
||||
logger.debug("Memory content preview", {
|
||||
content: memories.substring(0, 200),
|
||||
fullLength: memories.length,
|
||||
})
|
||||
}
|
||||
|
||||
if (systemPromptExists) {
|
||||
logger.debug("Appending memories to existing system prompt")
|
||||
return {
|
||||
...params,
|
||||
prompt: params.prompt.map((prompt) =>
|
||||
prompt.role === "system"
|
||||
? { ...prompt, content: `${prompt.content} \n ${memories}` }
|
||||
: prompt,
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
logger.debug(
|
||||
"System prompt does not exist, creating system prompt with memories",
|
||||
)
|
||||
return {
|
||||
...params,
|
||||
prompt: [{ role: "system" as const, content: memories }, ...params.prompt],
|
||||
}
|
||||
}
|
||||
|
||||
export const createSupermemoryMiddleware = (
|
||||
containerTag: string,
|
||||
verbose = false,
|
||||
mode: "profile" | "query" | "full" = "profile",
|
||||
): LanguageModelV2Middleware => {
|
||||
const logger = createLogger(verbose)
|
||||
|
||||
return {
|
||||
transformParams: async ({ params }) => {
|
||||
if (mode !== "profile") {
|
||||
const lastUserMessage = getLastUserMessage(params)
|
||||
if (!lastUserMessage) {
|
||||
logger.debug("No user message found, skipping memory search")
|
||||
return params
|
||||
}
|
||||
}
|
||||
|
||||
logger.info("Starting memory search", {
|
||||
containerTag,
|
||||
mode,
|
||||
})
|
||||
|
||||
const transformedParams = await addSystemPrompt(
|
||||
params,
|
||||
containerTag,
|
||||
logger,
|
||||
mode,
|
||||
)
|
||||
return transformedParams
|
||||
},
|
||||
}
|
||||
}
|
||||
35
packages/tools/src/vercel/util.ts
Normal file
35
packages/tools/src/vercel/util.ts
Normal file
|
|
@ -0,0 +1,35 @@
|
|||
export interface ProfileStructure {
|
||||
profile: {
|
||||
static?: string[]
|
||||
dynamic?: string[]
|
||||
},
|
||||
searchResults: {
|
||||
results: [
|
||||
{
|
||||
memory: string,
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Convert ProfileStructure to markdown
|
||||
* based on profile.static and profile.dynamic properties
|
||||
* @param data ProfileStructure
|
||||
* @returns Markdown string
|
||||
*/
|
||||
export function convertProfileToMarkdown(data: ProfileStructure): string {
|
||||
const sections: string[] = []
|
||||
|
||||
if (data.profile.static && data.profile.static.length > 0) {
|
||||
sections.push("## Static Profile")
|
||||
sections.push(data.profile.static.map((item) => `- ${item}`).join("\n"))
|
||||
}
|
||||
|
||||
if (data.profile.dynamic && data.profile.dynamic.length > 0) {
|
||||
sections.push("## Dynamic Profile")
|
||||
sections.push(data.profile.dynamic.map((item) => `- ${item}`).join("\n"))
|
||||
}
|
||||
|
||||
return sections.join("\n\n")
|
||||
}
|
||||
16
packages/tools/test/ai-sdk-test.ts
Normal file
16
packages/tools/test/ai-sdk-test.ts
Normal file
|
|
@ -0,0 +1,16 @@
|
|||
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", {
|
||||
verbose: true,
|
||||
mode: "query", // options are profile, query, full
|
||||
})
|
||||
|
||||
const result = await generateText({
|
||||
model: modelWithMemory,
|
||||
system: "You are an AI Girlfriend",
|
||||
messages: [{ role: "user", content: "Where do i live?" }],
|
||||
})
|
||||
|
||||
console.log(result.text)
|
||||
525
packages/tools/test/vercel.test.ts
Normal file
525
packages/tools/test/vercel.test.ts
Normal file
|
|
@ -0,0 +1,525 @@
|
|||
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],
|
||||
})
|
||||
})
|
||||
})
|
||||
})
|
||||
Loading…
Add table
Reference in a new issue