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?

[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

[ppetrick]

Running into the same problem. Issue happens for appendMessageAnnotation when I await any function that returns a promise inside onFinish (in my case, a nested generateObject call). Would appear OP is doing the same with the titlePromise call in the code they posted. Thoughts on how to address?

[lgrammel]

@ppetrick have you updated the ai library to the latest version? (also all @ai-sdk/* packages)

[ppetrick]

Updated to latest (3.4.33) from 3.4.14. Seems to have fixed the problem. Thanks!