How to extend message with custom attributes for AI SDK UI?

[Godrules500]

With RSC I was able to extend messages to add any attribute that I wanted this
messages[index].feedback= true

Now that I'm switching to ai sdk ui, how do I handle the same thing? I see data and annotations, but I remove the typescript benefits by doing messages[index].data = {feedback:true}. I could mispell it, or put a string instead of a boolean and etc.

So what is the best way to extend it so that I can add custom attributes to the message property for it to go back and forth with the api calls?

I get I could do something like this and set sendExtraMessageFields to true, but I was curious if there was an easier or better way.

import { Message as AIMessage } from 'ai'
import { UseChatHelpers } from 'ai/react'
import { AssistantFeedback } from './AssistantFeedback'

export type CustomMessage = AIMessage & {
  feedback?: AssistantFeedback
}

export type CustomUseChat = Omit<UseChatHelpers, 'messages'> & {
  messages: CustomMessage[]
}

And then call it like

let { messages, setMessages, handleSubmit, isLoading } = useChat({ id: chatId }) as CustomUseChat
messages[index].feedback = true.

Thanks!

[lgrammel]

Would https://sdk.vercel.ai/docs/ai-sdk-ui/streaming-data be an option?

[lgrammel]

When you initialize the chat, you can set messages that have annotations? would that work?

[Godrules500]

I'm not sure if it would. My chat object is

chat:{
	chatId
	messages
	userId
	isDeleted
	title
	etc
}

So I'm not sure if adding the annotations to the message would fit here unless I misunderstand how they work.

[lgrammel]

ah so you want to set the data on the chat not the message. Have you seen setData? You can use it to set the data initially when you restore. It will contain the stream data from the request otherwise. https://sdk.vercel.ai/docs/ai-sdk-ui/streaming-data#updating-and-clearing-data

[Godrules500]

@lgrammel Hey, I've tried using message annotations, but I'm not sure what's wrong, but the client is not picking it up.

I'm trying to handle error handling. Before, when an error happened, I appended a new message with certain error fields set. I'm trying to do this with the UI sdk and update some properties for the AI message. I'd rather just attach it to the AI response properties (is that possible). But I'm trying to use it with message annotations but it is coming back with nothing.

const data = new StreamData()
 
     const result = await streamText({
   temperature: 1,
   tools,
   model: setupVertex({ useSearchGrounding: useGrounding }),
   system: system_prompt,
   messages: convertToCoreMessages(secureMessages),
   maxSteps: 10,
   frequencyPenalty: 0.15,
   onFinish: async finishResult => {
     let done = true
     let text = finishResult.text
     let title = await titlePromise
     // textStream.done()
     let { error, showReportAnIssue } = getFinishedErrorText(finishResult.finishReason, text)
     if (error) {
       let message = messages[messages.length - 1]
       loggerWithFields({ user: user, message: message.content, promptViolation: true }).info('PROMPT_VIOLATION')
       const errorMessage = await handleMessageError(messages, new Error(error), showReportAnIssue, title)
       data.appendMessageAnnotation({
         content: errorMessage.content,
         name: errorMessage.name
       })
       data.close()
     } else {
 	// Here for testing
       // throw new Error('This is a test error');
       data.appendMessageAnnotation({
         id: generateId(),
         content: `I'm an annotation`,
         name: `error`
       })
       data.close()
     }

     if (done) {
       // await streamFinished(user, text, userFiles, title)
     }
   }
 })

And then on the client I have
annotations - {messages?.map((m: Message) => <div key={m.id}>{m.annotations && <>{JSON.stringify(m.annotations)}</>}</div>)}

Which is coming back with []

[lgrammel]

This is strange. Can you check out the reference example and see if that works for you?

• https://github.com/vercel/ai/blob/main/examples/next-openai/app/use-chat-streamdata/page.tsx
• https://github.com/vercel/ai/blob/main/examples/next-openai/app/api/use-chat-streamdata/route.ts