supermemory/packages/tools/src/ai-sdk.ts

import Supermemory from "supermemory"
import { tool } from "ai"
import { z } from "zod"
import {
	CLIENT_OPTIONS,
	DEFAULT_VALUES,
	PARAMETER_DESCRIPTIONS,
	SEARCH_LIMIT_BOUNDS,
	TOOL_DESCRIPTIONS,
	clampSearchLimit,
	deleteDocumentByIdentifier,
	getContainerTags,
} from "./tools-shared"
import { forgetMemoryRequest } from "./shared/forget-memory"
import type { SupermemoryToolsConfig } from "./types"

function createClient(apiKey: string, config?: SupermemoryToolsConfig) {
	return new Supermemory({
		apiKey,
		...CLIENT_OPTIONS,
		...(config?.baseUrl ? { baseURL: config.baseUrl } : {}),
	})
}

// Export individual tool creators
export const searchMemoriesTool = (
	apiKey: string,
	config?: SupermemoryToolsConfig,
) => {
	const client = createClient(apiKey, config)

	const containerTags = getContainerTags(config)
	const strict = config?.strict ?? false

	return tool({
		description: TOOL_DESCRIPTIONS.searchMemories,
		inputSchema: z.object({
			informationToGet: z
				.string()
				.describe(PARAMETER_DESCRIPTIONS.informationToGet),
			includeFullDocs: strict
				? z
						.boolean()
						.default(DEFAULT_VALUES.includeFullDocs)
						.describe(PARAMETER_DESCRIPTIONS.includeFullDocs)
				: z
						.boolean()
						.optional()
						.default(DEFAULT_VALUES.includeFullDocs)
						.describe(PARAMETER_DESCRIPTIONS.includeFullDocs),
			limit: strict
				? z.coerce
						.number()
						.int()
						.min(SEARCH_LIMIT_BOUNDS.min)
						.max(SEARCH_LIMIT_BOUNDS.max)
						.default(DEFAULT_VALUES.limit)
						.describe(PARAMETER_DESCRIPTIONS.searchLimit)
				: z.coerce
						.number()
						.int()
						.min(SEARCH_LIMIT_BOUNDS.min)
						.max(SEARCH_LIMIT_BOUNDS.max)
						.optional()
						.default(DEFAULT_VALUES.limit)
						.describe(PARAMETER_DESCRIPTIONS.searchLimit),
		}),
		execute: async ({ informationToGet, limit = DEFAULT_VALUES.limit }) => {
			try {
				const response = await client.search({
					q: informationToGet,
					containerTag: containerTags[0],
					limit: clampSearchLimit(limit),
					threshold: DEFAULT_VALUES.searchThreshold,
					searchMode: "hybrid",
				})

				return {
					success: true,
					results: response.results,
					count: response.results?.length || 0,
				}
			} catch (error) {
				return {
					success: false,
					error: error instanceof Error ? error.message : "Unknown error",
				}
			}
		},
	})
}

export const addMemoryTool = (
	apiKey: string,
	config?: SupermemoryToolsConfig,
) => {
	const client = createClient(apiKey, config)

	const containerTags = getContainerTags(config)

	return tool({
		description: TOOL_DESCRIPTIONS.addMemory,
		inputSchema: z.object({
			memory: z.string().describe(PARAMETER_DESCRIPTIONS.memory),
		}),
		execute: async ({ memory }) => {
			try {
				const metadata: Record<string, string | number | boolean> = {}

				const response = await client.add({
					content: memory,
					containerTags,
					...(Object.keys(metadata).length > 0 && { metadata }),
				})

				return {
					success: true,
					memory: response,
				}
			} catch (error) {
				return {
					success: false,
					error: error instanceof Error ? error.message : "Unknown error",
				}
			}
		},
	})
}

export const getProfileTool = (
	apiKey: string,
	config?: SupermemoryToolsConfig,
) => {
	const client = createClient(apiKey, config)

	const containerTags = getContainerTags(config)
	const strict = config?.strict ?? false

	return tool({
		description: TOOL_DESCRIPTIONS.getProfile,
		inputSchema: z.object({
			containerTag: strict
				? z.string().describe(PARAMETER_DESCRIPTIONS.containerTag)
				: z.string().optional().describe(PARAMETER_DESCRIPTIONS.containerTag),
			query: z.string().optional().describe(PARAMETER_DESCRIPTIONS.query),
		}),
		execute: async ({ containerTag, query }) => {
			try {
				const tag = containerTag || containerTags[0]

				const response = await client.profile({
					containerTag: tag,
					...(query && { q: query }),
				})

				return {
					success: true,
					profile: response.profile,
					searchResults: response.searchResults,
				}
			} catch (error) {
				return {
					success: false,
					error: error instanceof Error ? error.message : "Unknown error",
				}
			}
		},
	})
}

export const documentListTool = (
	apiKey: string,
	config?: SupermemoryToolsConfig,
) => {
	const client = createClient(apiKey, config)

	const containerTags = getContainerTags(config)
	const strict = config?.strict ?? false

	return tool({
		description: TOOL_DESCRIPTIONS.documentList,
		inputSchema: z.object({
			containerTag: z
				.string()
				.optional()
				.describe(PARAMETER_DESCRIPTIONS.containerTag),
			limit: strict
				? z.coerce
						.number()
						.default(DEFAULT_VALUES.limit)
						.describe(PARAMETER_DESCRIPTIONS.limit)
				: z.coerce
						.number()
						.optional()
						.default(DEFAULT_VALUES.limit)
						.describe(PARAMETER_DESCRIPTIONS.limit),
			page: z.coerce.number().optional().describe(PARAMETER_DESCRIPTIONS.page),
		}),
		execute: async ({ containerTag, limit, page }) => {
			try {
				const scopeTags: [string, ...string[]] = containerTag
					? [containerTag]
					: containerTags

				const response = await client.documents.list({
					containerTags: scopeTags,
					limit: limit || DEFAULT_VALUES.limit,
					...(page !== undefined && { page }),
				})

				return {
					success: true,
					documents: response.memories,
					pagination: response.pagination,
				}
			} catch (error) {
				return {
					success: false,
					error: error instanceof Error ? error.message : "Unknown error",
				}
			}
		},
	})
}

export const documentDeleteTool = (
	apiKey: string,
	config?: SupermemoryToolsConfig,
) => {
	const client = createClient(apiKey, config)
	const containerTags = getContainerTags(config)
	const strict = config?.strict ?? false

	return tool({
		description: TOOL_DESCRIPTIONS.documentDelete,
		inputSchema: z.object({
			documentId: z.string().describe(PARAMETER_DESCRIPTIONS.documentId),
			containerTag: strict
				? z
						.string()
						.nullable()
						.describe(PARAMETER_DESCRIPTIONS.documentContainerTag)
				: z
						.string()
						.optional()
						.describe(PARAMETER_DESCRIPTIONS.documentContainerTag),
		}),
		execute: async ({ documentId, containerTag }) => {
			try {
				const scopeTags: [string, ...string[]] = containerTag
					? [containerTag]
					: containerTags
				await deleteDocumentByIdentifier(client, documentId, scopeTags)

				return {
					success: true,
					message: `Document ${documentId} deleted successfully`,
				}
			} catch (error) {
				return {
					success: false,
					error: error instanceof Error ? error.message : "Unknown error",
				}
			}
		},
	})
}

export const documentAddTool = (
	apiKey: string,
	config?: SupermemoryToolsConfig,
) => {
	const client = createClient(apiKey, config)

	const containerTags = getContainerTags(config)

	return tool({
		description: TOOL_DESCRIPTIONS.documentAdd,
		inputSchema: z.object({
			content: z.string().describe(PARAMETER_DESCRIPTIONS.content),
			title: z.string().optional().describe(PARAMETER_DESCRIPTIONS.title),
			description: z
				.string()
				.optional()
				.describe(PARAMETER_DESCRIPTIONS.description),
		}),
		execute: async ({ content, title, description }) => {
			try {
				const metadata: Record<string, string> = {}
				if (title) metadata.title = title
				if (description) metadata.description = description

				const response = await client.documents.add({
					content,
					containerTags,
					...(Object.keys(metadata).length > 0 && { metadata }),
				})

				return {
					success: true,
					document: response,
				}
			} catch (error) {
				return {
					success: false,
					error: error instanceof Error ? error.message : "Unknown error",
				}
			}
		},
	})
}

export const memoryForgetTool = (
	apiKey: string,
	config?: SupermemoryToolsConfig,
) => {
	const containerTags = getContainerTags(config)

	return tool({
		description: TOOL_DESCRIPTIONS.memoryForget,
		inputSchema: z.object({
			containerTag: z
				.string()
				.optional()
				.describe(PARAMETER_DESCRIPTIONS.containerTag),
			memoryId: z.string().optional().describe(PARAMETER_DESCRIPTIONS.memoryId),
			memoryContent: z
				.string()
				.optional()
				.describe(PARAMETER_DESCRIPTIONS.memoryContent),
			reason: z.string().optional().describe(PARAMETER_DESCRIPTIONS.reason),
		}),
		execute: async ({ containerTag, memoryId, memoryContent, reason }) => {
			try {
				if (!memoryId && !memoryContent) {
					return {
						success: false,
						error: "Either memoryId or memoryContent must be provided",
					}
				}

				const tag = containerTag || containerTags[0]

				await forgetMemoryRequest(
					apiKey,
					{
						containerTag: tag as string,
						...(memoryId && { id: memoryId }),
						...(memoryContent && { content: memoryContent }),
						...(reason && { reason }),
					},
					config?.baseUrl,
				)

				return {
					success: true,
					message: "Memory forgotten successfully",
				}
			} catch (error) {
				return {
					success: false,
					error: error instanceof Error ? error.message : "Unknown error",
				}
			}
		},
	})
}

/**
 * Create Supermemory tools for AI SDK
 */
export function supermemoryTools(
	apiKey: string,
	config?: SupermemoryToolsConfig,
) {
	return {
		searchMemories: searchMemoriesTool(apiKey, config),
		addMemory: addMemoryTool(apiKey, config),
		getProfile: getProfileTool(apiKey, config),
		documentList: documentListTool(apiKey, config),
		documentDelete: documentDeleteTool(apiKey, config),
		documentAdd: documentAddTool(apiKey, config),
		memoryForget: memoryForgetTool(apiKey, config),
	}
}

// `./vercel` is not a published subpath, so this is the only way consumers reach the middleware types.
export {
	withSupermemory,
	type WithSupermemoryOptions,
	type PromptTemplate,
	type MemoryPromptData,
} from "./vercel"
export { getContainerTags } from "./tools-shared"