Lg/g6 i yb da

examples/next/app/api/chat/route.ts

@@ -1,31 +1,140 @@
-import { MyUIMessage } from '@/util/chat-schema';
+import { callWeatherApi } from '@/util/call-weather-api';
+import { Message } from '@/util/chat-schema';
 import { readChat, saveChat } from '@util/chat-store';
-import { convertToModelMessages, streamText } from 'ai';
+import {
+  convertToModelMessages,
+  createUIMessageStream,
+  createUIMessageStreamResponse,
+  streamText,
+  tool,
+} from 'ai';
+import z from 'zod/v4';
 
 export async function POST(req: Request) {
-  const { message, id }: { message: MyUIMessage; id: string } =
-    await req.json();
+  const { message, id }: { message: Message; id: string } = await req.json();
 
   const chat = await readChat(id);
   const messages = [...chat.messages, message];
 
-  const result = streamText({
-    model: 'openai/gpt-4o-mini',
-    messages: convertToModelMessages(messages),
-  });
+  const stream = createUIMessageStream({
+    execute: async ({ writer }) => {
+      // go through all messages, create a copy, replace data parts with
+      // tool invocations
+      const mappedMessages = messages.map(message => {
+        return {
+          ...message,
+          parts: message.parts.flatMap(part => {
+            if (part.type === 'data-weather') {
+              if (
+                part.data.status === 'generating' ||
+                part.data.status === 'calling api'
+              ) {
+                return []; // ignore generating parts
+              }
 
-  result.consumeStream(); // TODO always consume the stream even when the client disconnects
+              const weather = part.data.weather;
 
-  return result.toUIMessageStreamResponse({
-    originalMessages: messages,
-    messageMetadata: ({ part }) => {
-      if (part.type === 'start') {
-        return { createdAt: Date.now() };
-      }
+              return [
+                {
+                  type: 'tool-getWeatherInformation' as const,
+                  toolCallId: part.id!,
+                  state: 'output-available' as const,
+                  input: { city: weather.city },
+                  output: part.data.weather,
+                },
+                // {
+                //   type: 'text' as const,
+                //   text: `The weather in ${weather.city} is currently ${weather.weather}, with a temperature of ${weather.temperatureInCelsius}°C.`,
+                // },
+              ];
+            }
+            return [part];
+          }),
+        };
+      });
+
+      // TODO introduce a data part mapping in convertToModelMessages
+      const prompt = convertToModelMessages(mappedMessages);
+
+      console.log('prompt', JSON.stringify(prompt, null, 2));
+
+      const result = streamText({
+        model: 'vertex/gemini-2.0-flash-001',
+        prompt,
+        tools: {
+          getWeather: tool({
+            description: 'show the weather in a given city to the user',
+            inputSchema: z.object({ city: z.string() }),
+
+            onInputStart(options) {
+              writer.write({
+                type: 'data-weather',
+                id: options.toolCallId,
+                data: { status: 'generating' },
+              });
+            },
+
+            onInputAvailable(options) {
+              writer.write({
+                type: 'data-weather',
+                id: options.toolCallId,
+                data: { status: 'calling api' },
+              });
+            },
+
+            async execute({ city }, { toolCallId }) {
+              const weather = await callWeatherApi({ city });
+
+              writer.write({
+                type: 'data-weather',
+                id: toolCallId,
+                data: { status: 'available', weather },
+              });
+
+              return weather;
+            },
+          }),
+        },
+      });
+
+      result.consumeStream(); // TODO always consume the stream even when the client disconnects
+
+      writer.merge(
+        result
+          .toUIMessageStream({
+            messageMetadata: ({ part }) => {
+              if (part.type === 'start') {
+                return { createdAt: Date.now() };
+              }
+            },
+          })
+          // TODO this will go away in the future:
+          .pipeThrough(
+            new TransformStream({
+              transform(chunk, controller) {
+                // filter out tool related information
+                if (
+                  chunk.type === 'tool-output-available' ||
+                  chunk.type === 'tool-input-available' ||
+                  chunk.type === 'tool-input-delta' ||
+                  chunk.type === 'tool-input-start'
+                ) {
+                  return;
+                }
+
+                controller.enqueue(chunk as any);
+              },
+            }),
+          ),
+      );
     },
+
+    // save the chat when the stream is finished
+    originalMessages: messages,
     onFinish: ({ messages }) => {
-      // TODO fix type safety
-      saveChat({ id, messages: messages as MyUIMessage[] });
+      saveChat({ id, messages });
     },
   });
+
+  return createUIMessageStreamResponse({ stream });
 }

examples/next/app/chat/[chatId]/chat.tsx

@@ -1,18 +1,18 @@
 'use client';
 
 import { invalidateRouterCache } from '@/app/actions';
-import { MyUIMessage } from '@/util/chat-schema';
+import { Message } from '@/util/chat-schema';
 import { useChat } from '@ai-sdk/react';
 import { DefaultChatTransport } from 'ai';
 import { useEffect, useRef } from 'react';
 import ChatInput from './chat-input';
-import Message from './message';
+import MessageView from './message-view';
 
 export default function ChatComponent({
   chatData,
   isNewChat = false,
 }: {
-  chatData: { id: string; messages: MyUIMessage[] };
+  chatData: { id: string; messages: Message[] };
   isNewChat?: boolean;
 }) {
   const inputRef = useRef<HTMLInputElement>(null);
@@ -45,10 +45,12 @@ export default function ChatComponent({
     inputRef.current?.focus();
   }, []);
 
+  console.log('messages', messages);
+
   return (
     <div className="flex flex-col w-full max-w-md py-24 mx-auto stretch">
       {messages.map(message => (
-        <Message key={message.id} message={message} />
+        <MessageView key={message.id} message={message} />
       ))}
       <ChatInput
         status={status}

examples/next/app/chat/[chatId]/message-view.tsx

@@ -0,0 +1,69 @@
+import { Message } from '@/util/chat-schema';
+
+export default function MessageView({ message }: { message: Message }) {
+  const date = message.metadata?.createdAt
+    ? new Date(message.metadata.createdAt).toLocaleString()
+    : '';
+  const isUser = message.role === 'user';
+
+  return (
+    <div
+      className={`whitespace-pre-wrap my-2 p-3 rounded-lg shadow
+        ${isUser ? 'bg-blue-100 text-right ml-10' : 'bg-gray-100 text-left mr-10'}`}
+    >
+      <div className="mb-1 text-xs text-gray-500">{date}</div>
+      <div className="font-semibold">{isUser ? 'User:' : 'AI:'}</div>
+      <div>
+        {message.parts.map((part, index) => {
+          if (part.type === 'text') {
+            return part.text;
+          }
+
+          if (part.type === 'data-weather') {
+            if (part.data.status === 'generating') {
+              return (
+                <div
+                  key={index}
+                  className="p-2 mt-2 border border-gray-200 rounded bg-gray-50"
+                >
+                  <div className="text-gray-600">
+                    Generating weather data...
+                  </div>
+                </div>
+              );
+            }
+
+            if (part.data.status === 'calling api') {
+              return (
+                <div
+                  key={index}
+                  className="p-2 mt-2 border border-gray-200 rounded bg-gray-50"
+                >
+                  <div className="text-gray-600">Calling weather API...</div>
+                </div>
+              );
+            }
+
+            const { temperatureInCelsius, weather, city } = part.data.weather;
+            return (
+              <div
+                key={index}
+                className="p-2 mt-2 border border-blue-200 rounded bg-blue-50"
+              >
+                <div className="font-medium text-blue-800">
+                  Weather in {city}
+                </div>
+                <div className="text-blue-600">{weather}</div>
+                <div className="font-semibold text-blue-700">
+                  {temperatureInCelsius}°C
+                </div>
+              </div>
+            );
+          }
+
+          return '';
+        })}
+      </div>
+    </div>
+  );
+}

examples/next/util/call-weather-api.ts

@@ -0,0 +1,20 @@
+const weatherOptions = ['sunny', 'cloudy', 'rainy', 'windy'];
+
+export async function callWeatherApi({ city }: { city: string }) {
+  // Add artificial delay
+  await new Promise(resolve => setTimeout(resolve, 1000));
+
+  let weather =
+    weatherOptions[Math.floor(Math.random() * weatherOptions.length)];
+  const temperatureInCelsius = Math.floor(Math.random() * 50 - 15);
+
+  if (weather === 'rainy' && temperatureInCelsius < 0) {
+    weather = 'snowy';
+  }
+
+  return {
+    city,
+    weather,
+    temperatureInCelsius,
+  };
+}

examples/next/util/chat-schema.ts

@@ -1,16 +1,38 @@
-import { UIDataTypes, UIMessage } from 'ai';
-import { z } from 'zod';
+import { UIMessage } from 'ai';
 
-export const myMessageMetadataSchema = z.object({
-  createdAt: z.number(),
-});
+export type WeatherDataPart =
+  | { status: 'generating' }
+  | { status: 'calling api' }
+  | {
+      status: 'available';
+      weather: {
+        city: string;
+        weather: string;
+        temperatureInCelsius: number;
+      };
+    };
 
-export type MyMessageMetadata = z.infer<typeof myMessageMetadataSchema>;
-
-export type MyUIMessage = UIMessage<MyMessageMetadata, UIDataTypes>;
+export type Message = UIMessage<
+  {
+    createdAt: number;
+  },
+  {
+    weather: WeatherDataPart;
+  },
+  {
+    getWeatherInformation: {
+      input: { city: string };
+      output: {
+        city: string;
+        weather: string;
+        temperatureInCelsius: number;
+      };
+    };
+  }
+>;
 
 export type ChatData = {
   id: string;
-  messages: MyUIMessage[];
+  messages: Message[];
   createdAt: number;
 };

examples/next/util/chat-store.ts

@@ -2,7 +2,7 @@ import { generateId } from 'ai';
 import { existsSync, mkdirSync } from 'fs';
 import { readdir, readFile, writeFile } from 'fs/promises';
 import path from 'path';
-import { ChatData, MyUIMessage } from './chat-schema';
+import { ChatData, Message } from './chat-schema';
 
 // example implementation for demo purposes
 // in a real app, you would save the chat to a database
@@ -19,7 +19,7 @@ export async function saveChat({
   messages,
 }: {
   id: string;
-  messages: MyUIMessage[];
+  messages: Message[];
 }): Promise<void> {
   const chat = await readChat(id);
   chat.messages = messages;
@@ -31,7 +31,7 @@ export async function appendMessageToChat({
   message,
 }: {
   id: string;
-  message: MyUIMessage;
+  message: Message;
 }): Promise<void> {
   const chat = await readChat(id);
   chat.messages.push(message);