feat(vercel-ai-sdk): migrate to Vercel AI SDK v6 (#4741)
Co-authored-by: kartik-mem0 <kartik.labhshetwar@mem0.ai>
This commit is contained in:
Abhishek Chauhan
co-authored by
kartik-mem0
kartik-mem0
parent
b819d95d18
commit
b4a50e3dc8
@@ -1,14 +1,14 @@
|
||||
/* eslint-disable camelcase */
|
||||
import {
|
||||
LanguageModelV2CallOptions,
|
||||
LanguageModelV2Message,
|
||||
LanguageModelV2Source
|
||||
LanguageModelV3,
|
||||
LanguageModelV3CallOptions,
|
||||
LanguageModelV3GenerateResult,
|
||||
LanguageModelV3Message,
|
||||
LanguageModelV3Source,
|
||||
LanguageModelV3StreamResult,
|
||||
} from '@ai-sdk/provider';
|
||||
|
||||
import { LanguageModelV2 } from '@ai-sdk/provider';
|
||||
// streaming uses provider-native doStream; no middleware needed
|
||||
|
||||
import { Mem0ChatConfig, Mem0ChatModelId, Mem0ChatSettings, Mem0ConfigSettings, Mem0StreamResponse } from "./mem0-types";
|
||||
import { Mem0ChatConfig, Mem0ChatModelId, Mem0ChatSettings, Mem0ConfigSettings } from "./mem0-types";
|
||||
import { Mem0ClassSelector } from "./mem0-provider-selector";
|
||||
import { Mem0ProviderSettings } from "./mem0-provider";
|
||||
import { addMemories, getMemories } from "./mem0-utils";
|
||||
@@ -17,12 +17,8 @@ const generateRandomId = () => {
|
||||
return Math.random().toString(36).substring(2, 15) + Math.random().toString(36).substring(2, 15);
|
||||
}
|
||||
|
||||
export class Mem0GenericLanguageModel implements LanguageModelV2 {
|
||||
readonly specificationVersion = "v2";
|
||||
readonly defaultObjectGenerationMode = "json";
|
||||
// We don't support images for now
|
||||
readonly supportsImageUrls = false;
|
||||
// Allow All Media Types for now
|
||||
export class Mem0GenericLanguageModel implements LanguageModelV3 {
|
||||
readonly specificationVersion = "v3";
|
||||
readonly supportedUrls: Record<string, RegExp[]> = {
|
||||
'*': [/.*/]
|
||||
};
|
||||
@@ -38,148 +34,108 @@ export class Mem0GenericLanguageModel implements LanguageModelV2 {
|
||||
|
||||
provider: string;
|
||||
|
||||
private async processMemories(messagesPrompts: LanguageModelV2Message[], mem0Config: Mem0ConfigSettings) {
|
||||
private async processMemories(messagesPrompts: LanguageModelV3Message[], mem0Config: Mem0ConfigSettings) {
|
||||
try {
|
||||
// Add New Memories
|
||||
addMemories(messagesPrompts, mem0Config).then((res) => {
|
||||
return res;
|
||||
}).catch((e) => {
|
||||
console.error("Error while adding memories");
|
||||
return { memories: [], messagesPrompts: [] };
|
||||
// Add new memories — await to ensure writes complete before returning
|
||||
try {
|
||||
await addMemories(messagesPrompts, mem0Config);
|
||||
} catch (e) {
|
||||
console.error("Error while adding memories");
|
||||
}
|
||||
|
||||
// Get memories (always returns an array since graph support is removed)
|
||||
const memories: any[] = await getMemories(messagesPrompts, mem0Config) ?? [];
|
||||
|
||||
const mySystemPrompt = "These are the memories I have stored. Give more weightage to the question by users and try to answer that first. You have to modify your answer based on the memories I have provided. If the memories are irrelevant you can ignore them. Also don't reply to this section of the prompt, or the memories, they are only for your reference. The System prompt starts after text System Message: \n\n";
|
||||
|
||||
let memoriesText = "";
|
||||
try {
|
||||
memoriesText = memories
|
||||
?.map((memory: any) => `Memory: ${memory?.memory}\n\n`)
|
||||
.join("\n\n");
|
||||
} catch (e) {
|
||||
console.error("Error while parsing memories");
|
||||
}
|
||||
|
||||
const memoriesPrompt = `System Message: ${mySystemPrompt} ${memoriesText} `;
|
||||
|
||||
// Clone the prompt array to avoid mutating the caller's reference on retries
|
||||
const updatedPrompts = [...messagesPrompts];
|
||||
|
||||
if (memories?.length > 0) {
|
||||
const systemPrompt: LanguageModelV3Message = {
|
||||
role: "system",
|
||||
content: memoriesPrompt
|
||||
};
|
||||
updatedPrompts.unshift(systemPrompt);
|
||||
}
|
||||
|
||||
return { memories, messagesPrompts: updatedPrompts };
|
||||
} catch (e) {
|
||||
console.error("Error while processing memories");
|
||||
return { memories: [], messagesPrompts: [...messagesPrompts] };
|
||||
}
|
||||
}
|
||||
|
||||
async doGenerate(options: LanguageModelV3CallOptions): Promise<LanguageModelV3GenerateResult> {
|
||||
const provider = this.config.provider;
|
||||
const mem0_api_key = this.config.mem0ApiKey;
|
||||
|
||||
const settings: Mem0ProviderSettings = {
|
||||
provider: provider,
|
||||
mem0ApiKey: mem0_api_key,
|
||||
apiKey: this.config.apiKey,
|
||||
modelType: this.config.modelType,
|
||||
}
|
||||
|
||||
const mem0Config: Mem0ConfigSettings = {
|
||||
mem0ApiKey: mem0_api_key,
|
||||
...this.config.mem0Config,
|
||||
...this.settings,
|
||||
}
|
||||
|
||||
const selector = new Mem0ClassSelector(this.modelId, settings, this.provider_config);
|
||||
|
||||
const { memories, messagesPrompts: updatedPrompts } = await this.processMemories(options.prompt, mem0Config);
|
||||
|
||||
const model = selector.createProvider();
|
||||
|
||||
const ans = await model.doGenerate({
|
||||
...options,
|
||||
prompt: updatedPrompts,
|
||||
});
|
||||
|
||||
// Get Memories
|
||||
let memories = await getMemories(messagesPrompts, mem0Config);
|
||||
|
||||
const mySystemPrompt = "These are the memories I have stored. Give more weightage to the question by users and try to answer that first. You have to modify your answer based on the memories I have provided. If the memories are irrelevant you can ignore them. Also don't reply to this section of the prompt, or the memories, they are only for your reference. The System prompt starts after text System Message: \n\n";
|
||||
|
||||
const isGraphEnabled = mem0Config?.enable_graph;
|
||||
|
||||
let memoriesText = "";
|
||||
let memoriesText2 = "";
|
||||
try {
|
||||
// @ts-ignore
|
||||
if (isGraphEnabled) {
|
||||
memoriesText = memories?.results?.map((memory: any) => {
|
||||
return `Memory: ${memory?.memory}\n\n`;
|
||||
}).join("\n\n");
|
||||
|
||||
memoriesText2 = memories?.relations?.map((memory: any) => {
|
||||
return `Relation: ${memory?.source} -> ${memory?.relationship} -> ${memory?.target} \n\n`;
|
||||
}).join("\n\n");
|
||||
} else {
|
||||
memoriesText = memories?.map((memory: any) => {
|
||||
return `Memory: ${memory?.memory}\n\n`;
|
||||
}).join("\n\n");
|
||||
}
|
||||
} catch(e) {
|
||||
console.error("Error while parsing memories");
|
||||
if (!memories || memories?.length === 0) {
|
||||
return ans;
|
||||
}
|
||||
|
||||
let graphPrompt = "";
|
||||
if (isGraphEnabled) {
|
||||
graphPrompt = `HERE ARE THE GRAPHS RELATIONS FOR THE PREFERENCES OF THE USER:\n\n ${memoriesText2}`;
|
||||
}
|
||||
|
||||
const memoriesPrompt = `System Message: ${mySystemPrompt} ${memoriesText} ${graphPrompt} `;
|
||||
|
||||
// System Prompt - The memories go as a system prompt
|
||||
const systemPrompt: LanguageModelV2Message = {
|
||||
role: "system",
|
||||
content: memoriesPrompt
|
||||
const mem0Source: LanguageModelV3Source = {
|
||||
type: "source",
|
||||
sourceType: "url",
|
||||
id: "mem0-" + generateRandomId(),
|
||||
url: "https://app.mem0.ai?utm_source=oss&utm_medium=vercel-ai-sdk-src",
|
||||
title: "Mem0 Memories",
|
||||
providerMetadata: {
|
||||
mem0: {
|
||||
memories: memories,
|
||||
memoriesText: memories
|
||||
?.map((memory: any) => memory?.memory)
|
||||
.join("\n\n"),
|
||||
},
|
||||
},
|
||||
};
|
||||
|
||||
// Add the system prompt to the beginning of the messages if there are memories
|
||||
if (memories?.length > 0) {
|
||||
messagesPrompts.unshift(systemPrompt);
|
||||
}
|
||||
|
||||
if (isGraphEnabled) {
|
||||
memories = memories?.results;
|
||||
}
|
||||
|
||||
return { memories, messagesPrompts };
|
||||
} catch(e) {
|
||||
console.error("Error while processing memories");
|
||||
return { memories: [], messagesPrompts };
|
||||
}
|
||||
return {
|
||||
...ans,
|
||||
content: [...(ans.content ?? []), mem0Source],
|
||||
};
|
||||
}
|
||||
|
||||
async doGenerate(options: LanguageModelV2CallOptions): Promise<Awaited<ReturnType<LanguageModelV2['doGenerate']>>> {
|
||||
try {
|
||||
const provider = this.config.provider;
|
||||
const mem0_api_key = this.config.mem0ApiKey;
|
||||
|
||||
const settings: Mem0ProviderSettings = {
|
||||
provider: provider,
|
||||
mem0ApiKey: mem0_api_key,
|
||||
apiKey: this.config.apiKey,
|
||||
}
|
||||
|
||||
const mem0Config: Mem0ConfigSettings = {
|
||||
mem0ApiKey: mem0_api_key,
|
||||
...this.config.mem0Config,
|
||||
...this.settings,
|
||||
}
|
||||
|
||||
const selector = new Mem0ClassSelector(this.modelId, settings, this.provider_config);
|
||||
|
||||
let messagesPrompts = options.prompt;
|
||||
|
||||
// Process memories and update prompts
|
||||
const { memories, messagesPrompts: updatedPrompts } = await this.processMemories(messagesPrompts, mem0Config);
|
||||
|
||||
const model = selector.createProvider();
|
||||
|
||||
const ans = await model.doGenerate({
|
||||
...options,
|
||||
prompt: updatedPrompts,
|
||||
});
|
||||
|
||||
// If there are no memories, return the original response
|
||||
if (!memories || memories?.length === 0) {
|
||||
return ans;
|
||||
}
|
||||
|
||||
try {
|
||||
// Create sources array with existing sources
|
||||
const sources: LanguageModelV2Source[] = [
|
||||
{
|
||||
type: "source",
|
||||
title: "Mem0 Memories",
|
||||
sourceType: "url",
|
||||
id: "mem0-" + generateRandomId(),
|
||||
url: "https://app.mem0.ai?utm_source=oss&utm_medium=vercel-ai-sdk-src",
|
||||
providerMetadata: {
|
||||
mem0: {
|
||||
memories: memories,
|
||||
memoriesText: memories
|
||||
?.map((memory: any) => memory?.memory)
|
||||
.join("\n\n"),
|
||||
},
|
||||
},
|
||||
},
|
||||
];
|
||||
} catch (e) {
|
||||
console.error("Error while creating sources");
|
||||
}
|
||||
|
||||
return {
|
||||
...ans,
|
||||
// sources
|
||||
};
|
||||
} catch (error) {
|
||||
// Handle errors properly
|
||||
console.error("Error in doGenerate:", error);
|
||||
throw new Error("Failed to generate response.");
|
||||
}
|
||||
}
|
||||
|
||||
async doStream(options: LanguageModelV2CallOptions): Promise<Awaited<ReturnType<LanguageModelV2['doStream']>>> {
|
||||
async doStream(options: LanguageModelV3CallOptions): Promise<LanguageModelV3StreamResult> {
|
||||
try {
|
||||
const provider = this.config.provider;
|
||||
const mem0_api_key = this.config.mem0ApiKey;
|
||||
|
||||
|
||||
const settings: Mem0ProviderSettings = {
|
||||
provider: provider,
|
||||
mem0ApiKey: mem0_api_key,
|
||||
@@ -194,31 +150,18 @@ export class Mem0GenericLanguageModel implements LanguageModelV2 {
|
||||
}
|
||||
|
||||
const selector = new Mem0ClassSelector(this.modelId, settings, this.provider_config);
|
||||
|
||||
let messagesPrompts = options.prompt;
|
||||
|
||||
// Process memories and update prompts
|
||||
const { memories, messagesPrompts: updatedPrompts } = await this.processMemories(messagesPrompts, mem0Config);
|
||||
|
||||
const { messagesPrompts: updatedPrompts } = await this.processMemories(options.prompt, mem0Config);
|
||||
|
||||
const baseModel = selector.createProvider();
|
||||
|
||||
// Use the provider's native streaming directly to avoid buffering
|
||||
const streamResponse = await baseModel.doStream({
|
||||
...options,
|
||||
prompt: updatedPrompts,
|
||||
});
|
||||
|
||||
// If there are no memories, return the original stream
|
||||
if (!memories || memories?.length === 0) {
|
||||
return streamResponse;
|
||||
}
|
||||
|
||||
// Return stream untouched for true streaming behavior
|
||||
return {
|
||||
stream: streamResponse.stream,
|
||||
request: streamResponse.request,
|
||||
response: streamResponse.response,
|
||||
};
|
||||
// Return the full stream response, preserving all V3 fields (warnings, etc.)
|
||||
return streamResponse;
|
||||
} catch (error) {
|
||||
console.error("Error in doStream:", error);
|
||||
throw new Error("Streaming failed or method not implemented.");
|
||||
|
||||
Reference in New Issue
Block a user