feat: support Vercel AI SDK 5 Beta

apps/web/app/api/chat/route.ts

@@ -1,17 +1,21 @@
-import { Message, LlamaIndexAdapter, StreamData } from 'ai'
+import { fakeStreamText } from '@/app/utils'
+import { toUIMessageStream } from '@ai-sdk/llamaindex'
+import { UIMessage as Message } from '@ai-sdk/react'
 import {
-  ChatMessage,
+  MessageContentDetail,
   OpenAI,
   OpenAIEmbedding,
   Settings,
   SimpleChatEngine,
 } from 'llamaindex'
 import { NextResponse, type NextRequest } from 'next/server'
-import { fakeStreamText } from '@/app/utils'
 
 export const runtime = 'nodejs'
 export const dynamic = 'force-dynamic'
 
+const DATA_PREFIX = 'data: '
+const TOKEN_DELAY = 100
+
 Settings.llm = new OpenAI({ model: 'gpt-4o-mini' })
 Settings.embedModel = new OpenAIEmbedding({
   model: 'text-embedding-3-large',
@@ -24,32 +28,77 @@ export async function POST(request: NextRequest) {
     const messages = body.messages
     const lastMessage = messages[messages.length - 1]
 
-    const vercelStreamData = new StreamData()
-
     if (!process.env.OPENAI_API_KEY) {
       // Return fake stream if API key is not set
       return new Response(fakeStreamText(), {
         headers: {
-          'Content-Type': 'text/plain',
+          'Content-Type': 'text/event-stream',
           Connection: 'keep-alive',
         },
       })
     }
 
     const chatEngine = new SimpleChatEngine()
 
+    const messageContent = (lastMessage.parts[0] as { text: string }).text ?? ''
+
     const response = await chatEngine.chat({
-      message: lastMessage.content,
-      chatHistory: messages as ChatMessage[],
+      message: messageContent,
+      chatHistory: messages.map(message => ({
+        role: message.role,
+        content: message.parts as MessageContentDetail[],
+      })),
       stream: true,
     })
 
-    return LlamaIndexAdapter.toDataStreamResponse(response, {
-      data: vercelStreamData,
-      callbacks: {
-        onCompletion: async () => {
-          await vercelStreamData.close()
-        },
+    // TODO: `toUIMessageStream` not working as expected now
+    // const sseStream = toUIMessageStream(response)
+
+    const sseStream = new ReadableStream({
+      async start(controller) {
+        const encoder = new TextEncoder()
+
+        // New SSE format with text chunks
+        function writeStream(chunk: {
+          id: string
+          type: string
+          delta?: string
+        }) {
+          controller.enqueue(
+            encoder.encode(`${DATA_PREFIX}${JSON.stringify(chunk)}\n\n`)
+          )
+        }
+
+        // Generate a unique message id
+        const messageId = crypto.randomUUID()
+
+        // Start the text chunk
+        const startChunk = { id: messageId, type: 'text-start' }
+        writeStream(startChunk)
+
+        // Consume the response and write the chunks to the controller
+        for await (const chunk of response) {
+          writeStream({
+            id: messageId,
+            type: 'text-delta',
+            delta: chunk.delta,
+          })
+        }
+
+        // End the text chunk
+        const endChunk = { id: messageId, type: 'text-end' }
+        writeStream(endChunk)
+
+        controller.close()
+      },
+    })
+
+    return new Response(sseStream, {
+      status: 200,
+      statusText: 'OK',
+      headers: {
+        'content-type': 'text/event-stream',
+        connection: 'keep-alive',
       },
     })
   } catch (error) {

apps/web/app/demo/canvas/code-preview/page.tsx

@@ -9,19 +9,35 @@ import {
   useChatCanvas,
 } from '@llamaindex/chat-ui'
 import { DynamicComponent } from '@llamaindex/dynamic-ui'
-import { Message, useChat } from 'ai/react'
+import { UIMessage as Message, useChat } from '@ai-sdk/react'
 
 const initialMessages: Message[] = [
   {
     id: 'code-gen1',
     role: 'user',
-    content: 'Generate a simple calculator',
+    parts: [{ type: 'text', text: 'Generate a simple calculator' }],
   },
   {
     id: 'code-gen2',
     role: 'assistant',
-    content:
-      '\n```annotation\n{"type":"artifact","data":{"type":"code","created_at":1752124365106,"data":{"language":"typescript","file_name":"calculator.tsx","code":"import React, { useState } from \\"react\\"\\nimport { Button } from \\"@/components/ui/button\\"\\nimport { Card } from \\"@/components/ui/card\\"\\nimport { cn } from \\"@/lib/utils\\"\\n\\nconst buttons = [\\n  [\\"7\\", \\"8\\", \\"9\\", \\"/\\"],\\n  [\\"4\\", \\"5\\", \\"6\\", \\"*\\"],\\n  [\\"1\\", \\"2\\", \\"3\\", \\"-\\"],\\n  [\\"0\\", \\"C\\", \\"=\\", \\"+\\"],\\n]\\n\\nexport default function Calculator() {\\n  const [input, setInput] = useState<string>(\\"\\")\\n  const [result, setResult] = useState<string | null>(null)\\n\\n  const handleButtonClick = (value: string) => {\\n    if (value === \\"C\\") {\\n      setInput(\\"\\")\\n      setResult(null)\\n      return\\n    }\\n    if (value === \\"=\\") {\\n      try {\\n        // eslint-disable-next-line no-eval\\n        const evalResult = eval(input)\\n        setResult(evalResult.toString())\\n      } catch {\\n        setResult(\\"Error\\")\\n      }\\n      return\\n    }\\n    if (result !== null) {\\n      setInput(value.match(/[0-9.]/) ? value : result + value)\\n      setResult(null)\\n    } else {\\n      setInput((prev) => prev + value)\\n    }\\n  }\\n\\n  return (\\n    <div className=\\"flex items-center justify-center min-h-screen bg-muted\\">\\n      <Card className=\\"w-[320px] p-6 shadow-lg\\">\\n        <div className={cn(\\"mb-4 h-16 bg-background rounded flex items-end justify-end px-4 text-2xl font-mono border\\", result && \\"text-muted-foreground\\")}>\\n          {result !== null ? result : input || \\"0\\"}\\n        </div>\\n        <div className=\\"grid grid-cols-4 gap-3\\">\\n          {buttons.flat().map((btn, idx) => (\\n            <Button\\n              key={idx}\\n              variant={btn === \\"C\\" ? \\"destructive\\" : btn === \\"=\\" ? \\"default\\" : \\"outline\\"}\\n              className={cn(\\n                \\"h-12 text-xl\\",\\n                btn === \\"=\\" && \\"col-span-1 bg-primary text-primary-foreground\\",\\n                btn === \\"C\\" && \\"col-span-1\\"\\n              )}\\n              onClick={() => handleButtonClick(btn)}\\n            >\\n              {btn}\\n            </Button>\\n          ))}\\n        </div>\\n      </Card>\\n    </div>\\n  )\\n}"}}}\n```\nHere\'s how the simple calculator works:\n\n- The calculator displays the current input or the result at the top.\n- You can click the number buttons (0-9) and the operators (+, -, *, /) to build your calculation.\n- Pressing the = button evaluates the expression and shows the result.\n- Pressing the C button clears the input and resets the calculator.\n- If you get an error (like dividing by zero or entering an invalid expression), "Error" will be displayed.\n\nYou can further customize the calculator\'s appearance or add more features as needed! If you have any questions about how the code works or want to add more functionality, let me know!',
+    parts: [
+      {
+        type: 'text',
+        text: "Here's the simple calculator:",
+      },
+      {
+        type: 'data-artifact',
+        data: {
+          type: 'code',
+          created_at: 1752124365106,
+          data: {
+            language: 'typescript',
+            file_name: 'calculator.tsx',
+            code: 'import React, { useState } from "react"\nimport { Button } from "@/components/ui/button"\nimport { Card } from "@/components/ui/card"\nimport { cn } from "@/lib/utils"\n\nconst buttons = [\n  ["7", "8", "9", "/"],\n  ["4", "5", "6", "*"],\n  ["1", "2", "3", "-"],\n  ["0", "C", "=", "+"],\n]\n\nexport default function Calculator() {\n  const [input, setInput] = useState<string>("")\n  const [result, setResult] = useState<string | null>(null)\n\n  const handleButtonClick = (value: string) => {\n    if (value === "C") {\n      setInput("")\n      setResult(null)\n      return\n    }\n    if (value === "=") {\n      try {\n        // eslint-disable-next-line no-eval\n        const evalResult = eval(input)\n        setResult(evalResult.toString())\n      } catch {\n        setResult("Error")\n      }\n      return\n    }\n    if (result !== null) {\n      setInput(value.match(/[0-9.]/) ? value : result + value)\n      setResult(null)\n    } else {\n      setInput((prev) => prev + value)\n    }\n  }\n\n  return (\n    <div className="flex items-center justify-center min-h-screen bg-muted">\n      <Card className="w-[320px] p-6 shadow-lg">\n        <div className={cn("mb-4 h-16 bg-background rounded flex items-end justify-end px-4 text-2xl font-mono border", result && "text-muted-foreground")}>\n          {result !== null ? result : input || "0"}\n        </div>\n        <div className="grid grid-cols-4 gap-3">\n          {buttons.flat().map((btn, idx) => (\n            <Button\n              key={idx}\n              variant={btn === "C" ? "destructive" : btn === "=" ? "default" : "outline"}\n              className={cn(\n                "h-12 text-xl",\n                btn === "=" && "col-span-1 bg-primary text-primary-foreground",\n                btn === "C" && "col-span-1"\n              )}\n              onClick={() => handleButtonClick(btn)}\n            >\n              {btn}\n            </Button>\n          ))}\n        </div>\n      </Card>\n    </div>\n  )\n}',
+          },
+        },
+      },
+    ],
   },
 ]
 
@@ -30,7 +46,7 @@ export default function Page(): JSX.Element {
 }
 
 function CustomChat() {
-  const handler = useChat({ initialMessages })
+  const handler = useChat({ messages: initialMessages })
 
   return (
     <ChatSection

apps/web/app/demo/canvas/custom/page.tsx

@@ -11,7 +11,7 @@ import {
   useChatCanvas,
   useChatUI,
 } from '@llamaindex/chat-ui'
-import { Message, useChat } from 'ai/react'
+import { UIMessage as Message, useChat } from '@ai-sdk/react'
 import { Image } from 'lucide-react'
 
 const code = `
@@ -25,11 +25,11 @@ import {
   useChatCanvas,
   useChatUI,
 } from '@llamaindex/chat-ui'
-import { useChat } from 'ai/react'
+import { useChat } from '@ai-sdk/react'
 import { Image } from 'lucide-react'
 
 export function CustomChat() {
-  const handler = useChat({ initialMessages: [] })
+  const handler = useChat()
 
   return (
     <ChatSection
@@ -103,11 +103,8 @@ function CustomChatMessages() {
           >
             <ChatMessage.Avatar />
             <ChatMessage.Content>
-              <ChatMessage.Content.Markdown
-                annotationRenderers={{
-                  artifact: CustomArtifactCard,
-                }}
-              />
+              <ChatMessage.Part.Markdown />
+              <ChatMessage.Part.Artifact />
             </ChatMessage.Content>
             <ChatMessage.Actions />
           </ChatMessage>
@@ -116,32 +113,24 @@ function CustomChatMessages() {
     </ChatMessages>
   )
 }
-
-// custom artifact card for image artifacts
-function CustomArtifactCard({ data }: { data: Artifact }) {
-  return (
-    <ChatCanvas.Artifact
-      data={data}
-      getTitle={artifact => (artifact as ImageArtifact).data.caption}
-      iconMap={{ image: Image }}
-    />
-  )
-}
 `
 
 const initialMessages: Message[] = [
   {
     id: '1',
     role: 'user',
-    content: 'Generate an image of a cat',
+    parts: [{ type: 'text', text: 'Generate an image of a cat' }],
   },
   {
     id: '2',
     role: 'assistant',
-    content:
-      'Here is a cat image named Millie.' +
-      `\n\`\`\`annotation\n${JSON.stringify({
-        type: 'artifact',
+    parts: [
+      {
+        type: 'text',
+        text: 'Here is a cat image named Millie.',
+      },
+      {
+        type: 'data-artifact',
         data: {
           type: 'image',
           data: {
@@ -150,21 +139,24 @@ const initialMessages: Message[] = [
           },
           created_at: 1745480281756,
         },
-      })}
-      \n\`\`\`\n`,
+      },
+    ],
   },
   {
     id: '3',
     role: 'user',
-    content: 'Please generate a black cat image',
+    parts: [{ type: 'text', text: 'Please generate a black cat image' }],
   },
   {
     id: '4',
     role: 'assistant',
-    content:
-      'Here is a black cat image named Poppy.' +
-      `\n\`\`\`annotation\n${JSON.stringify({
-        type: 'artifact',
+    parts: [
+      {
+        type: 'text',
+        text: 'Here is a black cat image named Poppy.',
+      },
+      {
+        type: 'data-artifact',
         data: {
           type: 'image',
           data: {
@@ -173,8 +165,8 @@ const initialMessages: Message[] = [
           },
           created_at: 1745480281999,
         },
-      })}
-      \n\`\`\`\n`,
+      },
+    ],
   },
 ]
 
@@ -184,7 +176,7 @@ export default function Page(): JSX.Element {
 
 function CustomChat() {
   const { copyToClipboard, isCopied } = useCopyToClipboard({ timeout: 2000 })
-  const handler = useChat({ initialMessages })
+  const handler = useChat({ messages: initialMessages })
 
   return (
     <ChatSection
@@ -279,11 +271,8 @@ function CustomChatMessages() {
           >
             <ChatMessage.Avatar />
             <ChatMessage.Content>
-              <ChatMessage.Content.Markdown
-                annotationRenderers={{
-                  artifact: CustomArtifactCard,
-                }}
-              />
+              <ChatMessage.Part.Markdown />
+              <ChatMessage.Part.Artifact />
             </ChatMessage.Content>
             <ChatMessage.Actions />
           </ChatMessage>
@@ -292,14 +281,3 @@ function CustomChatMessages() {
     </ChatMessages>
   )
 }
-
-// custom artifact card for image artifacts
-function CustomArtifactCard({ data }: { data: Artifact }) {
-  return (
-    <ChatCanvas.Artifact
-      data={data}
-      getTitle={artifact => (artifact as ImageArtifact).data.caption}
-      iconMap={{ image: Image }}
-    />
-  )
-}