chore: update to ai sdk v5 beta (#1074)

lib/ai/models.test.ts

@@ -1,13 +1,14 @@
 import { simulateReadableStream } from 'ai';
-import { MockLanguageModelV1 } from 'ai/test';
+import { MockLanguageModelV2 } from 'ai/test';
 import { getResponseChunksByPrompt } from '@/tests/prompts/utils';
 
-export const chatModel = new MockLanguageModelV1({
+export const chatModel = new MockLanguageModelV2({
   doGenerate: async () => ({
     rawCall: { rawPrompt: null, rawSettings: {} },
     finishReason: 'stop',
-    usage: { promptTokens: 10, completionTokens: 20 },
-    text: `Hello, world!`,
+    usage: { inputTokens: 10, outputTokens: 20, totalTokens: 30 },
+    content: [{ type: 'text', text: 'Hello, world!' }],
+    warnings: [],
   }),
   doStream: async ({ prompt }) => ({
     stream: simulateReadableStream({
@@ -19,12 +20,13 @@ export const chatModel = new MockLanguageModelV1({
   }),
 });
 
-export const reasoningModel = new MockLanguageModelV1({
+export const reasoningModel = new MockLanguageModelV2({
   doGenerate: async () => ({
     rawCall: { rawPrompt: null, rawSettings: {} },
     finishReason: 'stop',
-    usage: { promptTokens: 10, completionTokens: 20 },
-    text: `Hello, world!`,
+    usage: { inputTokens: 10, outputTokens: 20, totalTokens: 30 },
+    content: [{ type: 'text', text: 'Hello, world!' }],
+    warnings: [],
   }),
   doStream: async ({ prompt }) => ({
     stream: simulateReadableStream({
@@ -36,24 +38,26 @@ export const reasoningModel = new MockLanguageModelV1({
   }),
 });
 
-export const titleModel = new MockLanguageModelV1({
+export const titleModel = new MockLanguageModelV2({
   doGenerate: async () => ({
     rawCall: { rawPrompt: null, rawSettings: {} },
     finishReason: 'stop',
-    usage: { promptTokens: 10, completionTokens: 20 },
-    text: `This is a test title`,
+    usage: { inputTokens: 10, outputTokens: 20, totalTokens: 30 },
+    content: [{ type: 'text', text: 'This is a test title' }],
+    warnings: [],
   }),
   doStream: async () => ({
     stream: simulateReadableStream({
       chunkDelayInMs: 500,
       initialDelayInMs: 1000,
       chunks: [
-        { type: 'text-delta', textDelta: 'This is a test title' },
+        { id: '1', type: 'text-start' },
+        { id: '1', type: 'text-delta', delta: 'This is a test title' },
+        { id: '1', type: 'text-end' },
         {
           type: 'finish',
           finishReason: 'stop',
-          logprobs: undefined,
-          usage: { completionTokens: 10, promptTokens: 3 },
+          usage: { inputTokens: 3, outputTokens: 10, totalTokens: 13 },
         },
       ],
     }),
@@ -61,12 +65,13 @@ export const titleModel = new MockLanguageModelV1({
   }),
 });
 
-export const artifactModel = new MockLanguageModelV1({
+export const artifactModel = new MockLanguageModelV2({
   doGenerate: async () => ({
     rawCall: { rawPrompt: null, rawSettings: {} },
     finishReason: 'stop',
-    usage: { promptTokens: 10, completionTokens: 20 },
-    text: `Hello, world!`,
+    usage: { inputTokens: 10, outputTokens: 20, totalTokens: 30 },
+    content: [{ type: 'text', text: 'Hello, world!' }],
+    warnings: [],
   }),
   doStream: async ({ prompt }) => ({
     stream: simulateReadableStream({

lib/ai/providers.ts

@@ -4,13 +4,13 @@ import {
   wrapLanguageModel,
 } from 'ai';
 import { xai } from '@ai-sdk/xai';
-import { isTestEnvironment } from '../constants';
 import {
   artifactModel,
   chatModel,
   reasoningModel,
   titleModel,
 } from './models.test';
+import { isTestEnvironment } from '../constants';
 
 export const myProvider = isTestEnvironment
   ? customProvider({
@@ -32,6 +32,6 @@ export const myProvider = isTestEnvironment
         'artifact-model': xai('grok-2-1212'),
       },
       imageModels: {
-        'small-model': xai.image('grok-2-image'),
+        'small-model': xai.imageModel('grok-2-image'),
       },
     });

lib/ai/tools/create-document.ts

@@ -1,46 +1,51 @@
 import { generateUUID } from '@/lib/utils';
-import { DataStreamWriter, tool } from 'ai';
+import { tool, type UIMessageStreamWriter } from 'ai';
 import { z } from 'zod';
-import { Session } from 'next-auth';
+import type { Session } from 'next-auth';
 import {
   artifactKinds,
   documentHandlersByArtifactKind,
 } from '@/lib/artifacts/server';
+import type { ChatMessage } from '@/lib/types';
 
 interface CreateDocumentProps {
   session: Session;
-  dataStream: DataStreamWriter;
+  dataStream: UIMessageStreamWriter<ChatMessage>;
 }
 
 export const createDocument = ({ session, dataStream }: CreateDocumentProps) =>
   tool({
     description:
       'Create a document for a writing or content creation activities. This tool will call other functions that will generate the contents of the document based on the title and kind.',
-    parameters: z.object({
+    inputSchema: z.object({
       title: z.string(),
       kind: z.enum(artifactKinds),
     }),
     execute: async ({ title, kind }) => {
       const id = generateUUID();
 
-      dataStream.writeData({
-        type: 'kind',
-        content: kind,
+      dataStream.write({
+        type: 'data-kind',
+        data: kind,
+        transient: true,
       });
 
-      dataStream.writeData({
-        type: 'id',
-        content: id,
+      dataStream.write({
+        type: 'data-id',
+        data: id,
+        transient: true,
       });
 
-      dataStream.writeData({
-        type: 'title',
-        content: title,
+      dataStream.write({
+        type: 'data-title',
+        data: title,
+        transient: true,
       });
 
-      dataStream.writeData({
-        type: 'clear',
-        content: '',
+      dataStream.write({
+        type: 'data-clear',
+        data: null,
+        transient: true,
       });
 
       const documentHandler = documentHandlersByArtifactKind.find(
@@ -59,7 +64,7 @@ export const createDocument = ({ session, dataStream }: CreateDocumentProps) =>
         session,
       });
 
-      dataStream.writeData({ type: 'finish', content: '' });
+      dataStream.write({ type: 'data-finish', data: null, transient: true });
 
       return {
         id,

lib/ai/tools/get-weather.ts

@@ -3,7 +3,7 @@ import { z } from 'zod';
 
 export const getWeather = tool({
   description: 'Get the current weather at a location',
-  parameters: z.object({
+  inputSchema: z.object({
     latitude: z.number(),
     longitude: z.number(),
   }),

lib/ai/tools/request-suggestions.ts

@@ -1,14 +1,15 @@
 import { z } from 'zod';
-import { Session } from 'next-auth';
-import { DataStreamWriter, streamObject, tool } from 'ai';
+import type { Session } from 'next-auth';
+import { streamObject, tool, type UIMessageStreamWriter } from 'ai';
 import { getDocumentById, saveSuggestions } from '@/lib/db/queries';
-import { Suggestion } from '@/lib/db/schema';
+import type { Suggestion } from '@/lib/db/schema';
 import { generateUUID } from '@/lib/utils';
 import { myProvider } from '../providers';
+import type { ChatMessage } from '@/lib/types';
 
 interface RequestSuggestionsProps {
   session: Session;
-  dataStream: DataStreamWriter;
+  dataStream: UIMessageStreamWriter<ChatMessage>;
 }
 
 export const requestSuggestions = ({
@@ -17,7 +18,7 @@ export const requestSuggestions = ({
 }: RequestSuggestionsProps) =>
   tool({
     description: 'Request suggestions for a document',
-    parameters: z.object({
+    inputSchema: z.object({
       documentId: z
         .string()
         .describe('The ID of the document to request edits'),
@@ -49,7 +50,8 @@ export const requestSuggestions = ({
       });
 
       for await (const element of elementStream) {
-        const suggestion = {
+        // @ts-ignore todo: fix type
+        const suggestion: Suggestion = {
           originalText: element.originalSentence,
           suggestedText: element.suggestedSentence,
           description: element.description,
@@ -58,9 +60,10 @@ export const requestSuggestions = ({
           isResolved: false,
         };
 
-        dataStream.writeData({
-          type: 'suggestion',
-          content: suggestion,
+        dataStream.write({
+          type: 'data-suggestion',
+          data: suggestion,
+          transient: true,
         });
 
         suggestions.push(suggestion);

lib/ai/tools/update-document.ts

@@ -1,18 +1,19 @@
-import { DataStreamWriter, tool } from 'ai';
-import { Session } from 'next-auth';
+import { tool, type UIMessageStreamWriter } from 'ai';
+import type { Session } from 'next-auth';
 import { z } from 'zod';
-import { getDocumentById, saveDocument } from '@/lib/db/queries';
+import { getDocumentById } from '@/lib/db/queries';
 import { documentHandlersByArtifactKind } from '@/lib/artifacts/server';
+import type { ChatMessage } from '@/lib/types';
 
 interface UpdateDocumentProps {
   session: Session;
-  dataStream: DataStreamWriter;
+  dataStream: UIMessageStreamWriter<ChatMessage>;
 }
 
 export const updateDocument = ({ session, dataStream }: UpdateDocumentProps) =>
   tool({
     description: 'Update a document with the given description.',
-    parameters: z.object({
+    inputSchema: z.object({
       id: z.string().describe('The ID of the document to update'),
       description: z
         .string()
@@ -27,9 +28,10 @@ export const updateDocument = ({ session, dataStream }: UpdateDocumentProps) =>
         };
       }
 
-      dataStream.writeData({
-        type: 'clear',
-        content: document.title,
+      dataStream.write({
+        type: 'data-clear',
+        data: null,
+        transient: true,
       });
 
       const documentHandler = documentHandlersByArtifactKind.find(
@@ -48,7 +50,7 @@ export const updateDocument = ({ session, dataStream }: UpdateDocumentProps) =>
         session,
       });
 
-      dataStream.writeData({ type: 'finish', content: '' });
+      dataStream.write({ type: 'data-finish', data: null, transient: true });
 
       return {
         id,