Vercel AI SDK migration to V5 (#3223)
Co-authored-by: Saket Aryan <94069182+whysosaket@users.noreply.github.com>
This commit is contained in:
Parshva Daftari
co-authored by
Saket Aryan
Saket Aryan
parent
42e60d6724
commit
0e18f54d36
@@ -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);
|
||||
|
||||
Reference in New Issue
Block a user