JDSS/mem0
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Vercel AI SDK migration to V5 (#3223)

Browse Source 
Co-authored-by: Saket Aryan <94069182+whysosaket@users.noreply.github.com>
This commit is contained in:
Parshva Daftari
2025-08-05 20:27:57 +05:30
committed by GitHub
co-authored by Saket Aryan
parent 42e60d6724
commit 0e18f54d36
23 changed files with 321 additions and 423 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vercel-ai-sdk/src/mem0-generic-language-model.ts
+63 -64
View File
@@ -1,9 +1,13 @@
/* eslint-disable camelcase */
import {
LanguageModelV1,
LanguageModelV1CallOptions,
LanguageModelV1Message,
} from "@ai-sdk/provider";
LanguageModelV2CallOptions,
LanguageModelV2Message,
LanguageModelV2Source,
LanguageModelV2StreamPart
} from '@ai-sdk/provider';
import { LanguageModelV2 } from '@ai-sdk/provider';
import { simulateStreamingMiddleware, wrapLanguageModel } from 'ai';
import { Mem0ChatConfig, Mem0ChatModelId, Mem0ChatSettings, Mem0ConfigSettings, Mem0StreamResponse } from "./mem0-types";
import { Mem0ClassSelector } from "./mem0-provider-selector";
@@ -14,10 +18,15 @@ const generateRandomId = () => {
return Math.random().toString(36).substring(2, 15) + Math.random().toString(36).substring(2, 15);
}
export class Mem0GenericLanguageModel implements LanguageModelV1 {
readonly specificationVersion = "v1";
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
readonly supportedUrls: Record<string, RegExp[]> = {
'*': [/.*/]
};
constructor(
public readonly modelId: Mem0ChatModelId,
@@ -30,7 +39,7 @@ export class Mem0GenericLanguageModel implements LanguageModelV1 {
provider: string;
private async processMemories(messagesPrompts: LanguageModelV1Message[], mem0Config: Mem0ConfigSettings) {
private async processMemories(messagesPrompts: LanguageModelV2Message[], mem0Config: Mem0ConfigSettings) {
try {
// Add New Memories
addMemories(messagesPrompts, mem0Config).then((res) => {
@@ -76,7 +85,7 @@ export class Mem0GenericLanguageModel implements LanguageModelV1 {
const memoriesPrompt = `System Message: ${mySystemPrompt} ${memoriesText} ${graphPrompt} `;
// System Prompt - The memories go as a system prompt
const systemPrompt: LanguageModelV1Message = {
const systemPrompt: LanguageModelV2Message = {
role: "system",
content: memoriesPrompt
};
@@ -97,7 +106,7 @@ export class Mem0GenericLanguageModel implements LanguageModelV1 {
}
}
async doGenerate(options: LanguageModelV1CallOptions): Promise<Awaited<ReturnType<LanguageModelV1['doGenerate']>>> {
async doGenerate(options: LanguageModelV2CallOptions): Promise<Awaited<ReturnType<LanguageModelV2['doGenerate']>>> {
try {
const provider = this.config.provider;
const mem0_api_key = this.config.mem0ApiKey;
@@ -133,44 +142,32 @@ export class Mem0GenericLanguageModel implements LanguageModelV1 {
return ans;
}
// Create sources array with existing sources
const sources = [...(ans.sources || [])];
// Add a combined source with all memories
if (Array.isArray(memories) && memories?.length > 0) {
sources.push({
title: "Mem0 Memories",
sourceType: "url",
id: "mem0-" + generateRandomId(),
url: "https://app.mem0.ai",
providerMetadata: {
mem0: {
memories: memories,
memoriesText: memories?.map((memory: any) => memory?.memory).join("\n\n")
}
}
});
// Add individual memory sources for more detailed information
memories?.forEach((memory: any) => {
sources.push({
title: memory.title || "Memory",
try {
// Create sources array with existing sources
const sources: LanguageModelV2Source[] = [
{
type: "source",
title: "Mem0 Memories",
sourceType: "url",
id: "mem0-memory-" + generateRandomId(),
id: "mem0-" + generateRandomId(),
url: "https://app.mem0.ai",
providerMetadata: {
mem0: {
memory: memory,
memoryText: memory?.memory
}
}
});
});
memories: memories,
memoriesText: memories
?.map((memory: any) => memory?.memory)
.join("\n\n"),
},
},
},
];
} catch (e) {
console.error("Error while creating sources");
}
return {
...ans,
sources
// sources
};
} catch (error) {
// Handle errors properly
@@ -179,7 +176,7 @@ export class Mem0GenericLanguageModel implements LanguageModelV1 {
}
}
async doStream(options: LanguageModelV1CallOptions): Promise<Awaited<ReturnType<LanguageModelV1['doStream']>>> {
async doStream(options: LanguageModelV2CallOptions): Promise<Awaited<ReturnType<LanguageModelV2['doStream']>>> {
try {
const provider = this.config.provider;
const mem0_api_key = this.config.mem0ApiKey;
@@ -204,7 +201,13 @@ export class Mem0GenericLanguageModel implements LanguageModelV1 {
// Process memories and update prompts
const { memories, messagesPrompts: updatedPrompts } = await this.processMemories(messagesPrompts, mem0Config);
const model = selector.createProvider();
const baseModel = selector.createProvider();
// Wrap the model with streaming middleware using the new Vercel AI SDK 5.0 approach
const model = wrapLanguageModel({
model: baseModel,
middleware: simulateStreamingMiddleware(),
});
const streamResponse = await model.doStream({
...options,
@@ -228,16 +231,15 @@ export class Mem0GenericLanguageModel implements LanguageModelV1 {
// Create a single source that contains all memories
controller.enqueue({
type: 'source',
source: {
title: "Mem0 Memories",
sourceType: "url",
id: "mem0-" + generateRandomId(),
url: "https://app.mem0.ai",
providerMetadata: {
mem0: {
memories: memories,
memoriesText: memories?.map((memory: any) => memory?.memory).join("\n\n")
}
title: "Mem0 Memories",
sourceType: "url",
id: "mem0-" + generateRandomId(),
url: "https://app.mem0.ai",
providerOptions: {
mem0: {
memories: memories,
memoriesText: memories?.map((memory: any) => memory?.memory).join("\n\n")
}
}
});
@@ -246,16 +248,15 @@ export class Mem0GenericLanguageModel implements LanguageModelV1 {
memories?.forEach((memory: any) => {
controller.enqueue({
type: 'source',
source: {
title: memory?.title || "Memory",
sourceType: "url",
id: "mem0-memory-" + generateRandomId(),
url: "https://app.mem0.ai",
providerMetadata: {
mem0: {
memory: memory,
memoryText: memory?.memory
}
title: memory?.title || "Memory",
sourceType: "url",
id: "mem0-memory-" + generateRandomId(),
url: "https://app.mem0.ai",
providerOptions: {
mem0: {
memory: memory,
memoryText: memory?.memory
}
}
});
@@ -277,10 +278,8 @@ export class Mem0GenericLanguageModel implements LanguageModelV1 {
// Return a new stream response with our enhanced stream
return {
stream: enhancedStream,
rawCall: streamResponse.rawCall,
rawResponse: streamResponse.rawResponse,
request: streamResponse.request,
warnings: streamResponse.warnings
response: streamResponse.response,
};
} catch (error) {
console.error("Error in doStream:", error);
vercel-ai-sdk/src/mem0-provider-selector.ts
+2 -2
View File
@@ -1,6 +1,6 @@
import { Mem0ProviderSettings } from "./mem0-provider";
import Mem0AITextGenerator, { ProviderSettings } from "./provider-response-provider";
import { LanguageModelV1 } from "ai";
import { LanguageModelV2 } from '@ai-sdk/provider';
class Mem0ClassSelector {
modelId: string;
@@ -24,7 +24,7 @@ class Mem0ClassSelector {
}
}
createProvider(): LanguageModelV1 {
createProvider(): LanguageModelV2 {
return new Mem0AITextGenerator(this.modelId, this.config , this.provider_config || {});
}
}
vercel-ai-sdk/src/mem0-provider.ts
+11 -15
View File
@@ -1,27 +1,23 @@
import { LanguageModelV1, ProviderV1 } from "@ai-sdk/provider";
import { loadApiKey, withoutTrailingSlash } from "@ai-sdk/provider-utils";
import { ProviderV2 } from '@ai-sdk/provider';
import { LanguageModelV2 } from '@ai-sdk/provider';
import { withoutTrailingSlash } from "@ai-sdk/provider-utils";
import { Mem0ChatModelId, Mem0ChatSettings, Mem0Config } from "./mem0-types";
import { OpenAIProviderSettings } from "@ai-sdk/openai";
import { Mem0GenericLanguageModel } from "./mem0-generic-language-model";
import { OpenAIChatSettings } from "@ai-sdk/openai/internal";
import { AnthropicMessagesSettings } from "@ai-sdk/anthropic/internal";
import { AnthropicProviderSettings } from "@ai-sdk/anthropic";
import { LLMProviderSettings } from "./mem0-types";
export interface Mem0Provider extends ProviderV1 {
(modelId: Mem0ChatModelId, settings?: Mem0ChatSettings): LanguageModelV1;
export interface Mem0Provider extends ProviderV2 {
(modelId: Mem0ChatModelId, settings?: Mem0ChatSettings): LanguageModelV2;
chat(modelId: Mem0ChatModelId, settings?: Mem0ChatSettings): LanguageModelV1;
completion(modelId: Mem0ChatModelId, settings?: Mem0ChatSettings): LanguageModelV1;
chat(modelId: Mem0ChatModelId, settings?: Mem0ChatSettings): LanguageModelV2;
completion(modelId: Mem0ChatModelId, settings?: Mem0ChatSettings): LanguageModelV2;
languageModel(
modelId: Mem0ChatModelId,
settings?: Mem0ChatSettings
): LanguageModelV1;
): LanguageModelV2;
}
export interface Mem0ProviderSettings
extends OpenAIChatSettings,
AnthropicMessagesSettings {
export interface Mem0ProviderSettings {
baseURL?: string;
/**
* Custom fetch implementation. You can use it as a middleware to intercept
@@ -46,7 +42,7 @@ export interface Mem0ProviderSettings
/**
* The configuration for the provider.
*/
config?: OpenAIProviderSettings | AnthropicProviderSettings;
config?: LLMProviderSettings ;
}
export function createMem0(
vercel-ai-sdk/src/mem0-types.ts
+10 -9
View File
@@ -1,10 +1,9 @@
import { Mem0ProviderSettings } from "./mem0-provider";
import { OpenAIChatSettings } from "@ai-sdk/openai/internal";
import { AnthropicMessagesSettings } from "@ai-sdk/anthropic/internal";
import {
LanguageModelV1
} from "@ai-sdk/provider";
import { OpenAIProviderSettings } from "@ai-sdk/openai";
import { AnthropicProviderSettings } from "@ai-sdk/anthropic";
import { LanguageModelV2 } from '@ai-sdk/provider';
import { CohereProviderSettings } from "@ai-sdk/cohere";
import { GroqProviderSettings } from "@ai-sdk/groq";
export type Mem0ChatModelId =
| (string & NonNullable<unknown>);
@@ -33,9 +32,11 @@ export interface Mem0ConfigSettings {
export interface Mem0ChatConfig extends Mem0ConfigSettings, Mem0ProviderSettings {}
export interface Mem0Config extends Mem0ConfigSettings {}
export interface Mem0ChatSettings extends OpenAIChatSettings, AnthropicMessagesSettings, Mem0ConfigSettings {}
export interface LLMProviderSettings extends OpenAIProviderSettings, AnthropicProviderSettings, CohereProviderSettings, GroqProviderSettings {}
export interface Mem0StreamResponse extends Awaited<ReturnType<LanguageModelV1['doStream']>> {
export interface Mem0Config extends Mem0ConfigSettings {}
export interface Mem0ChatSettings extends Mem0ConfigSettings {}
export interface Mem0StreamResponse extends Awaited<ReturnType<LanguageModelV2['doStream']>> {
memories: any;
}
vercel-ai-sdk/src/mem0-utils.ts
+7 -7
View File
@@ -1,4 +1,4 @@
import { LanguageModelV1Prompt } from 'ai';
import { LanguageModelV2Prompt } from '@ai-sdk/provider';
import { Mem0ConfigSettings } from './mem0-types';
import { loadApiKey } from '@ai-sdk/provider-utils';
interface Message {
@@ -6,7 +6,7 @@ interface Message {
content: string | Array<{type: string, text: string}>;
}
const flattenPrompt = (prompt: LanguageModelV1Prompt) => {
const flattenPrompt = (prompt: LanguageModelV2Prompt) => {
try {
return prompt.map((part) => {
if (part.role === "user") {
@@ -23,7 +23,7 @@ const flattenPrompt = (prompt: LanguageModelV1Prompt) => {
}
}
const convertToMem0Format = (messages: LanguageModelV1Prompt) => {
const convertToMem0Format = (messages: LanguageModelV2Prompt) => {
try {
return messages.flatMap((message: any) => {
try {
@@ -127,7 +127,7 @@ const searchInternalMemories = async (query: string, config?: Mem0ConfigSettings
}
}
const addMemories = async (messages: LanguageModelV1Prompt, config?: Mem0ConfigSettings) => {
const addMemories = async (messages: LanguageModelV2Prompt, config?: Mem0ConfigSettings) => {
try {
let finalMessages: Array<Message> = [];
if (typeof messages === "string") {
@@ -172,7 +172,7 @@ const updateMemories = async (messages: Array<Message>, config?: Mem0ConfigSetti
}
}
const retrieveMemories = async (prompt: LanguageModelV1Prompt | string, config?: Mem0ConfigSettings) => {
const retrieveMemories = async (prompt: LanguageModelV2Prompt | string, config?: Mem0ConfigSettings) => {
try {
const message = typeof prompt === 'string' ? prompt : flattenPrompt(prompt);
const systemPrompt = "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";
@@ -208,7 +208,7 @@ const retrieveMemories = async (prompt: LanguageModelV1Prompt | string, config?:
}
}
const getMemories = async (prompt: LanguageModelV1Prompt | string, config?: Mem0ConfigSettings) => {
const getMemories = async (prompt: LanguageModelV2Prompt | string, config?: Mem0ConfigSettings) => {
try {
const message = typeof prompt === 'string' ? prompt : flattenPrompt(prompt);
const memories = await searchInternalMemories(message, config);
@@ -223,7 +223,7 @@ const getMemories = async (prompt: LanguageModelV1Prompt | string, config?: Mem0
}
}
const searchMemories = async (prompt: LanguageModelV1Prompt | string, config?: Mem0ConfigSettings) => {
const searchMemories = async (prompt: LanguageModelV2Prompt | string, config?: Mem0ConfigSettings) => {
try {
const message = typeof prompt === 'string' ? prompt : flattenPrompt(prompt);
const memories = await searchInternalMemories(message, config);
vercel-ai-sdk/src/provider-response-provider.ts
+10 -8
View File
@@ -1,4 +1,4 @@
import { LanguageModelV1, LanguageModelV1CallOptions } from "ai";
import { LanguageModelV2, LanguageModelV2CallOptions } from "@ai-sdk/provider";
import { Mem0ProviderSettings } from "./mem0-provider";
import { createOpenAI, OpenAIProviderSettings } from "@ai-sdk/openai";
import { CohereProviderSettings, createCohere } from "@ai-sdk/cohere";
@@ -7,13 +7,15 @@ import { createGoogleGenerativeAI, GoogleGenerativeAIProviderSettings } from "@a
import { createGroq, GroqProviderSettings } from "@ai-sdk/groq";
// Define a private provider field
class Mem0AITextGenerator implements LanguageModelV1 {
readonly specificationVersion = "v1";
class Mem0AITextGenerator implements LanguageModelV2 {
readonly specificationVersion = "v2";
readonly defaultObjectGenerationMode = "json";
readonly supportsImageUrls = false;
readonly modelId: string;
readonly provider = "mem0";
readonly supportedUrls: Record<string, RegExp[]> = {
'*': [/.*/]
};
private languageModel: any; // Use any type to avoid version conflicts
constructor(modelId: string, config: Mem0ProviderSettings, provider_config: ProviderSettings) {
@@ -68,14 +70,14 @@ class Mem0AITextGenerator implements LanguageModelV1 {
}
}
async doGenerate(options: LanguageModelV1CallOptions): Promise<Awaited<ReturnType<LanguageModelV1['doGenerate']>>> {
async doGenerate(options: LanguageModelV2CallOptions): Promise<Awaited<ReturnType<LanguageModelV2['doGenerate']>>> {
const result = await this.languageModel.doGenerate(options);
return result as Awaited<ReturnType<LanguageModelV1['doGenerate']>>;
return result as Awaited<ReturnType<LanguageModelV2['doGenerate']>>;
}
async doStream(options: LanguageModelV1CallOptions): Promise<Awaited<ReturnType<LanguageModelV1['doStream']>>> {
async doStream(options: LanguageModelV2CallOptions): Promise<Awaited<ReturnType<LanguageModelV2['doStream']>>> {
const result = await this.languageModel.doStream(options);
return result as Awaited<ReturnType<LanguageModelV1['doStream']>>;
return result as Awaited<ReturnType<LanguageModelV2['doStream']>>;
}
}