Update to AI SDK v5 and latest dependencies

JustAnOkapi · JustAnOkapi · commit e83080995b92 · 2026-01-20T15:28:31.000-06:00
- Upgrade ai package from 4.x to 5.x
- Upgrade Next.js to 15.5.9
- Upgrade @assistant-ui/react to 0.11.57 and @assistant-ui/react-ai-sdk to 1.1.21
- Upgrade @mem0/vercel-ai-provider to 2.0.5
- Migrate API route to use createUIMessageStream and convertToModelMessages
- Update client to use DefaultChatTransport for useChatRuntime
- Update memory-ui to read DataMessagePart from message content
diff --git a/app/api/chat/route.ts b/app/api/chat/route.ts
@@ -1,10 +1,16 @@
 /* eslint-disable @typescript-eslint/no-explicit-any */
 
-import { createDataStreamResponse, jsonSchema, streamText } from "ai";
+import {
+  streamText,
+  UIMessage,
+  convertToModelMessages,
+  createUIMessageStream,
+  createUIMessageStreamResponse,
+} from "ai";
 import { addMemories, getMemories } from "@mem0/vercel-ai-provider";
 import { openai } from "@ai-sdk/openai";
+import { frontendTools } from "@assistant-ui/react-ai-sdk";
 
-export const runtime = "edge";
 export const maxDuration = 30;
 
 const retrieveMemories = (memories: any) => {
@@ -20,46 +26,75 @@ const retrieveMemories = (memories: any) => {
   return `System Message: ${systemPrompt} ${memoriesText}`;
 };
 
+// Extract text content from messages for mem0
+const extractTextFromMessages = (messages: UIMessage[]): string => {
+  return messages
+    .map((m) => {
+      if (Array.isArray(m.parts)) {
+        return m.parts
+          .filter((p): p is { type: "text"; text: string } => p.type === "text")
+          .map((p) => p.text)
+          .join(" ");
+      }
+      return "";
+    })
+    .join("\n");
+};
+
 export async function POST(req: Request) {
-  const { messages, system, tools, userId } = await req.json();
+  const {
+    messages,
+    system,
+    tools,
+    userId,
+  } = (await req.json()) as {
+    messages: UIMessage[];
+    system?: string;
+    tools?: Parameters<typeof frontendTools>[0];
+    userId: string;
+  };
+
+  // Convert UIMessages to model messages for streamText
+  const modelMessages = convertToModelMessages(messages);
 
-  const memories = await getMemories(messages, { user_id: userId });
+  // Extract text for mem0 (it expects string or LanguageModelV2Prompt)
+  const textPrompt = extractTextFromMessages(messages);
+
+  const memories = await getMemories(textPrompt, { user_id: userId });
   const mem0Instructions = retrieveMemories(memories);
 
-  const result = streamText({
-    model: openai("gpt-4o"),
-    messages,
-    // forward system prompt and tools from the frontend
-    system: [system, mem0Instructions].filter(Boolean).join("\n"),
-    tools: Object.fromEntries(
-      Object.entries<{ parameters: unknown }>(tools).map(([name, tool]) => [
-        name,
-        {
-          parameters: jsonSchema(tool.parameters!),
-        },
-      ])
-    ),
-  });
+  // addMemories expects LanguageModelV2Prompt, use modelMessages with type cast
+  const addMemoriesTask = addMemories(modelMessages as any, { user_id: userId });
 
-  const addMemoriesTask = addMemories(messages, { user_id: userId });
-  return createDataStreamResponse({
-    execute: async (writer) => {
+  const stream = createUIMessageStream({
+    execute: async ({ writer }) => {
+      // Write memory annotation as data part with custom type
       if (memories.length > 0) {
-        writer.writeMessageAnnotation({
-          type: "mem0-get",
-          memories,
+        writer.write({
+          type: "data-mem0-get",
+          data: memories,
         });
       }
 
-      result.mergeIntoDataStream(writer);
+      const result = streamText({
+        model: openai("gpt-4o"),
+        messages: modelMessages,
+        system: [system, mem0Instructions].filter(Boolean).join("\n"),
+        tools: tools ? (frontendTools(tools) as any) : undefined,
+      });
 
+      writer.merge(result.toUIMessageStream());
+
+      // Add memories after stream completes
       const newMemories = await addMemoriesTask;
       if (newMemories.length > 0) {
-        writer.writeMessageAnnotation({
-          type: "mem0-update",
-          memories: newMemories,
+        writer.write({
+          type: "data-mem0-update",
+          data: newMemories,
         });
       }
     },
   });
+
+  return createUIMessageStreamResponse({ stream });
 }
diff --git a/app/assistant.tsx b/app/assistant.tsx
@@ -4,12 +4,14 @@ import { AssistantRuntimeProvider } from "@assistant-ui/react";
 import { useChatRuntime } from "@assistant-ui/react-ai-sdk";
 import { Thread } from "@/components/assistant-ui/thread";
 import { ThreadList } from "@/components/assistant-ui/thread-list";
-import { useEffect, useState } from "react";
+import { useEffect, useState, useMemo } from "react";
 import { v4 as uuidv4 } from "uuid";
 import Cookies from "js-cookie";
+import { DefaultChatTransport } from "ai";
 
 const useUserId = () => {
-  const [userId, setUserId] = useState<string>("");
+  // null = loading, string = loaded
+  const [userId, setUserId] = useState<string | null>(null);
 
   useEffect(() => {
     let id = Cookies.get("userId");
@@ -23,12 +25,18 @@ const useUserId = () => {
   return userId;
 };
 
-export const Assistant = () => {
-  const userId = useUserId();
-  const runtime = useChatRuntime({
-    api: "/api/chat",
-    body: { userId },
-  });
+// Separate component that only renders when userId is ready
+const AssistantInner = ({ userId }: { userId: string }) => {
+  const transport = useMemo(
+    () =>
+      new DefaultChatTransport({
+        api: "/api/chat",
+        body: { userId },
+      }),
+    [userId]
+  );
+
+  const runtime = useChatRuntime({ transport });
 
   return (
     <AssistantRuntimeProvider runtime={runtime}>
@@ -39,3 +47,18 @@ export const Assistant = () => {
     </AssistantRuntimeProvider>
   );
 };
+
+export const Assistant = () => {
+  const userId = useUserId();
+
+  // Don't render until userId is loaded to prevent race conditions
+  if (!userId) {
+    return (
+      <div className="h-dvh flex items-center justify-center">
+        <div className="text-muted-foreground">Loading...</div>
+      </div>
+    );
+  }
+
+  return <AssistantInner userId={userId} />;
+};
diff --git a/components/assistant-ui/memory-ui.tsx b/components/assistant-ui/memory-ui.tsx
@@ -3,70 +3,98 @@ import { FC, useMemo } from "react";
 import MemoryIndicator, { Memory } from "./memory-indicator";
 
 type RetrievedMemory = {
-  isNew: boolean;
   id: string;
   memory: string;
   user_id: string;
   categories: readonly string[];
-  immutable: boolean;
   created_at: string;
   updated_at: string;
   score: number;
 };
 
 type NewMemory = {
-  id: string;
-  data: {
+  id?: string;
+  memory?: string;
+  data?: {
     memory: string;
   };
-  event: "ADD" | "DELETE";
+  event?: "ADD" | "DELETE";
+  // For async response
+  message?: string;
+  status?: string;
+};
+
+// DataMessagePart from assistant-ui: type: "data" with name field
+type DataMessagePart<T = unknown> = {
+  readonly type: "data";
+  readonly name: string;
+  readonly data: T;
 };
 
-type NewMemoryAnnotation = {
-  readonly type: "mem0-update";
-  readonly memories: readonly NewMemory[];
+type Mem0GetDataPart = DataMessagePart<readonly RetrievedMemory[]> & {
+  readonly name: "mem0-get";
 };
 
-type GetMemoryAnnotation = {
-  readonly type: "mem0-get";
-  readonly memories: readonly RetrievedMemory[];
+type Mem0UpdateDataPart = DataMessagePart<readonly NewMemory[]> & {
+  readonly name: "mem0-update";
 };
 
-type MemoryAnnotation = NewMemoryAnnotation | GetMemoryAnnotation;
+const isMem0GetPart = (p: unknown): p is Mem0GetDataPart =>
+  typeof p === "object" &&
+  p != null &&
+  "type" in p &&
+  p.type === "data" &&
+  "name" in p &&
+  p.name === "mem0-get" &&
+  "data" in p;
 
-const isMemoryAnnotation = (a: unknown): a is MemoryAnnotation =>
-  typeof a === "object" &&
-  a != null &&
-  "type" in a &&
-  (a.type === "mem0-update" || a.type === "mem0-get");
+const isMem0UpdatePart = (p: unknown): p is Mem0UpdateDataPart =>
+  typeof p === "object" &&
+  p != null &&
+  "type" in p &&
+  p.type === "data" &&
+  "name" in p &&
+  p.name === "mem0-update" &&
+  "data" in p;
 
 const useMemories = (): Memory[] => {
-  const annotations = useMessage((m) => m.metadata.unstable_annotations);
-  console.log("annotations", annotations);
-  return useMemo(
-    () =>
-      annotations?.filter(isMemoryAnnotation).flatMap((a) => {
-        if (a.type === "mem0-update") {
-          return a.memories.map(
-            (m): Memory => ({
-              event: m.event,
-              id: m.id,
-              memory: m.data.memory,
-              score: 1,
-            })
-          );
-        } else if (a.type === "mem0-get") {
-          return a.memories.map((m) => ({
+  const content = useMessage((m) => m.content);
+
+  return useMemo(() => {
+    if (!content) return [];
+
+    const memories: Memory[] = [];
+
+    for (const part of content as unknown[]) {
+      if (isMem0GetPart(part)) {
+        for (const m of part.data) {
+          memories.push({
             event: "GET",
             id: m.id,
             memory: m.memory,
             score: m.score,
-          }));
+          });
         }
-        throw new Error("Unexpected annotation: " + JSON.stringify(a));
-      }) ?? [],
-    [annotations]
-  );
+      } else if (isMem0UpdatePart(part)) {
+        for (const m of part.data) {
+          // Skip async pending responses
+          if (m.status === "PENDING") continue;
+
+          const memoryText = m.memory ?? m.data?.memory ?? "";
+          if (!memoryText) continue;
+
+          memories.push({
+            event: m.event ?? "ADD",
+            id: m.id ?? "",
+            memory: memoryText,
+            score: 1,
+          });
+        }
+      }
+    }
+
+    return memories;
+  }, [content]);
 };
 
 export const MemoryUI: FC = () => {
diff --git a/package.json b/package.json
@@ -9,40 +9,41 @@
     "lint": "next lint"
   },
   "dependencies": {
-    "@ai-sdk/openai": "^1.1.15",
-    "@assistant-ui/react": "^0.8.2",
-    "@assistant-ui/react-ai-sdk": "^0.8.0",
-    "@assistant-ui/react-markdown": "^0.8.0",
-    "@mem0/vercel-ai-provider": "^0.0.14",
-    "@radix-ui/react-avatar": "^1.1.3",
-    "@radix-ui/react-popover": "^1.1.6",
-    "@radix-ui/react-slot": "^1.1.2",
-    "@radix-ui/react-tooltip": "^1.1.8",
+    "@ai-sdk/openai": "^2.0.88",
+    "@assistant-ui/react": "^0.11.57",
+    "@assistant-ui/react-ai-sdk": "^1.1.21",
+    "@assistant-ui/react-markdown": "^0.11.9",
+    "@mem0/vercel-ai-provider": "^2.0.5",
+    "@radix-ui/react-avatar": "^1.1.11",
+    "@radix-ui/react-popover": "^1.1.15",
+    "@radix-ui/react-slot": "^1.2.4",
+    "@radix-ui/react-tooltip": "^1.2.8",
     "@types/js-cookie": "^3.0.6",
     "@types/uuid": "^10.0.0",
-    "ai": "^4.1.46",
+    "ai": "^5.0.116",
     "class-variance-authority": "^0.7.1",
     "clsx": "^2.1.1",
     "js-cookie": "^3.0.5",
-    "lucide-react": "^0.477.0",
-    "next": "15.2.0",
-    "react": "^19.0.0",
-    "react-dom": "^19.0.0",
+    "lucide-react": "^0.562.0",
+    "next": "^15.2.8",
+    "react": "^19.2.3",
+    "react-dom": "^19.2.3",
     "remark-gfm": "^4.0.1",
-    "tailwind-merge": "^3.0.2",
+    "tailwind-merge": "^3.4.0",
     "tailwindcss-animate": "^1.0.7",
-    "uuid": "^11.1.0"
+    "uuid": "^11.1.0",
+    "zod": "^3.25.0"
   },
   "devDependencies": {
-    "@eslint/eslintrc": "^3",
+    "@eslint/eslintrc": "^3.3.3",
     "@types/node": "^22",
-    "@types/react": "^19",
-    "@types/react-dom": "^19",
-    "eslint": "^9",
-    "eslint-config-next": "15.2.0",
-    "postcss": "^8",
-    "tailwindcss": "^3.4.1",
-    "typescript": "^5"
+    "@types/react": "^19.2.9",
+    "@types/react-dom": "^19.2.3",
+    "eslint": "^9.39.2",
+    "eslint-config-next": "^15.2.8",
+    "postcss": "^8.5.6",
+    "tailwindcss": "^3.4.17",
+    "typescript": "^5.9.3"
   },
   "packageManager": "pnpm@10.5.2"
 }
diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml
