This little exercise has more to do about TypeScript than it has to do with the AI SDK altogether. But while this can seem something that is not that important, let me tell you that it’ll be a crucial part on how you make your applications type-safe while you work with your FrontEnd code!

That’s because being able to bring the same types produced by our backend has always being a fundamental part of also discovering issues between BE and FE in case something has changed.

And as you’ll see we will do this with less code as possible, because we will use the magic of InferUITools and the generics we can pass to UIMessage in order to export and import all the types we need!

So the first thing we have to do in order to make all this magic work is to move tools configuration into a separate tools variables. This is quite simple, so I’ll skip it.

Once we have our tools variable, we want to leverage InferUITools to generate a discriminated union for all the different tool definitions that we have inside.

type MyTools = InferUITools<typeof tools>;

Now that we have a proper type for all the tools that we can use, it’s time to pass it as a third parameter in the UIMessage generics so we can export it and use wherever we need to consume the output of our tools.

export type MyUIMessage = UIMessage<never, never, MyTools>;

This was all the work we had to do inside our api/chet.ts, now it is time to bring all this power inside our FrontEnd, and the first part is right where we call useChat.

As you can imagine, the useChat hook has been created to handle standard UIMessage and this kind of general purpose type for messages is not able to handle our specific responses where a message can have a type name identified by the tool that generated it.

Do you remember in our previous lesson? each message inside the parts generated by the onFinish callback had a type formatted like tool-<toolName>.

Since it is likely that we want to leverage useChat with different type of messages, it make sense that this hook is capable to accept a generic that will customize the structure of the data that it can handle.

const { messages, sendMessage } = useChat<MyUIMessage>({});

Here we leverage our newly created MyUIMessage (don’t forget to import it), and if you try to hover over messages, you will see that it will be typed as MyUIMessage[].

Now inside our interface we map over messages and for each message we pass message.parts to the Message component. Let’s see how it leverages this kind of type.

export const Message = ({
  role,
  parts,
}: {
  role: string;
  parts: MyUIMessage['parts'];
}) => {}

As you can see here, Message takes two props:

• role by which we can show the “user” that send the message (meaning we will write User in case role === 'user' and AI otherwise)
• parts all the parts of the message

First we will check which item in parts is type === 'text' so we can return a proper message ‘typed’ by the AI:

const prefix = role === 'user' ? 'User: ' : 'AI: ';

const text = parts
  .map((part) => {
    if (part.type === 'text') {
      return part.text;
    }
    return '';
  })
  .join('');
  
return (
  <div className="flex flex-col gap-2">
    <div className="prose prose-invert my-6">
      <ReactMarkdown>{prefix + text}</ReactMarkdown>
    </div>
    
    { /* We will loop parts here */ }
  </div>
);

Now that we know how the text part will be displayed, let’s see what we can do with the other type held in each item in parts:

{parts.map((part, index) => {
  if (part.type === 'tool-writeFile') { }
  if (part.type === 'tool-readFile') { }
  if (part.type === 'tool-deletePath') { }
  if (part.type === 'tool-listDirectory') { }
  if (part.type === 'tool-createDirectory') { }
  if (part.type === 'tool-exists') { }
  if (part.type === 'tool-searchFiles') { }
  
  return null;
})}

As you can see, per each part we check if type is set to a value that we can handle. If it does not match with any tool call we simply return null and do not display anything.

But if it does match we have the opportunity of returning a specific JSX capable to handle the tool response. Let’s discover how I’ve implemented the JSX for the tool-writeFile response, also because Matt already filled all the other kind of matches.

if (part.type === 'tool-writeFile') {
  return (
    <div
      key={index}
      className="bg-green-900/20 border border-green-700 rounded p-3 text-sm"
    >
      <div className="font-semibold text-green-300 mb-1">
        📝 Wrote file
      </div>
      <div className="text-green-200">
        Path: {part.input?.path ?? 'Unknown'}
        <br />
        Length: {part.input?.content?.length ?? 'Unkown'}
      </div>
    </div>
  );
}

As you can see, this is just standard JSX where we can leverage the path specific proprs (like path or content) in our response.

I would like to expand a little bit on the reason why the shape of the parts looks exactly the same as the ones we got from the onFinish callback even though we pass the streaming back to our chat.

And in all honesty the reason is pretty simple. Matt decided to introduce us to the shape of messages in the previous lesson just because the kind of array that we also get from toUIMessageStreamResponse is almost the same.

The mayor difference between toUIMessageStream that we used in previous lesson and toUIMessageStreamResponse that we have used now is that the latter will also take care of all the additional information we have to add to stream the messages over HTTP and make them available to useChat.