The done in unstable_onSetAIState doesnt work, it always gives False. Making it hard to save it to a database

[kaarthik108]

Description

unstable_onSetAIState streams the state messages, but unfornutaley not able to keep track of when it finishes streaming.

The done prop doesnt work, always returns false.

Code example

  unstable_onSetAIState: async ({ state, done, key }) => {
    "use server";

    const { chatId, messages } = state;
    console.log("messages------\n\n", messages);

    const title = messages[0].content.substring(0, 100);
    const path = `/chat/${chatId}`;
    const chat: Chat = {
      id: chatId,
      title,
      messages: messages,
      path,
      userId: "test",
    };
    await saveChat(chat);
  },
});

This streams messages making it imposiible to save to a database

`

rgument #0:

'messages------\n\n'

 messages:

Array(10) [
{ id: 'gdSHceT', role: 'user', content: '\n\nhi' },
{ id: 'test', role: 'assistant', content: 'Hello' },
{ id: 'test', role: 'assistant', content: 'Hello!' },
{ id: 'test', role: 'assistant', content: 'Hello! How' },
{ id: 'test', role: 'assistant', content: 'Hello! How can' },
{ id: 'test', role: 'assistant', content: 'Hello! How can I' },
{ id: 'test', role: 'assistant', content: 'Hello! How can I assist' },
{
id: 'test',
role: 'assistant',
content: 'Hello! How can I assist you'
},
{
id: 'test',
role: 'assistant',
content: 'Hello! How can I assist you today'
},
{
id: 'test',
role: 'assistant',
content: 'Hello! How can I assist you today?'
}
]

18:52:12.390 | +47 ms | next.js server

 argument #0:

'messages------\n\n'

`

Additional context

No response

[jeremyphilemon]

@kaarthik108 done returns true only when .done() is called on the mutable ai state. Can you confirm you're seeing the same behavior after marking your state as done?

[kaarthik108]

Hey @jeremyphilemon , I do have aiState.done on all of my function call and text-delta methods. but the done is still giving 'False'

[kaarthik108]

I got a workaround to save it from inside the submit function, but would be good to do it on unstable_onsetaistate instead.

async function submitUserMessage(content: string) {
  "use server";
  const aiState = getMutableAIState();

  aiState.update({
    ...aiState.get(),
    messages: [
      ...aiState.get().messages,
      {
        id: nanoid(),
        role: "user",
        content: `${aiState.get().interactions.join("\n\n")}\n\n${content}`,
      },
    ],
  });

  const history = aiState.get().messages.map((message: any) => ({
    role: message.role,
    content: message.content,
  }));

  const textStream = createStreamableValue("");
  const spinnerStream = createStreamableUI(<BotCard>{spinner}</BotCard>);
  // const skeletonStream = createStreamableUI(<AreaSkeleton />);
  const messageStream = createStreamableUI(null);
  const uiStream = createStreamableUI();

  const getDDL = await getContext(content);

  (async () => {
    try {
      const result = await experimental_streamText({
        model: openai.chat("gpt-4-turbo"),
        temperature: 0,
        tools: {
          query_data: {
            description:
              "Query the data from the sqlite database and return the results.",
            parameters: querySchema,
          },
        },
        system: `
You are a Sqlite data analytics assistant. 
      `,
        messages: [...history],
      });

      let textContent = "";
      spinnerStream.done(null);

      for await (const delta of result.fullStream) {
        const { type } = delta;

        if (type === "text-delta") {
          const { textDelta } = delta;

          textContent += textDelta;

          messageStream.update(<BotMessage content={textContent} />);

          // aiState.update({
          //   ...aiState.get(),
          //   messages: [
          //     ...aiState.get().messages,
          //     {
          //       id: "1",
          //       role: "assistant",
          //       content: textContent,
          //     },
          //   ],
          // });

          aiState.done({
            ...aiState.get(),
            interactions: [],
            messages: [
              ...aiState.get().messages,
              {
                id: Date.now(),
                role: "assistant",
                content: textContent,
              },
            ],
          });
        } else if (type === "tool-call") {
          const { toolName, args } = delta;

          if (toolName === "query_data") {
            const {
              format,
              title,
              timeField,
              categories,
              index,
              yaxis,
              size,
              query,
            } = args;

            uiStream.update(
              <BotCard>
                <ChartWrapper props={args} />
              </BotCard>
            );

            aiState.done({
              ...aiState.get(),
              interactions: [],
              messages: [
                ...aiState.get().messages,
                {
                  id: nanoid(),
                  role: "assistant",
                  content: `[Sqlite query results for code: ${query} and chart format: ${format} with categories: ${categories} and data index: ${index} and yaxis: ${yaxis} and size: ${size}]`,
                  display: {
                    name: "query_data",
                    props: args,
                  },
                },
              ],
            });
          }
        }
      }
      const ms = aiState.get().messages;
      const id = aiState.get().chatId;
      // @ts-ignore
      const latestMessages: _Message[] = Array.from(
        new Map(
          ms.filter(isMessage).map((msg: Message) => [msg.role, msg])
        ).values()
      );
      await saveChat({
        id,
        title: latestMessages.map((msg) => msg.content).join("\n\n"),
        messages: latestMessages,
        path: `/chat/${id}`,
        userId: "test",
      });

      uiStream.done();
      textStream.done();
      messageStream.done();
    } catch (e) {
      console.error(e);

      const error = new Error("The AI got into some problem.");
      uiStream.error(error);
      textStream.error(error);
      messageStream.error(error);
      // @ts-ignore
      aiState.done();
    }
  })();

  return {
    id: nanoid(),
    attachments: uiStream.value,
    spinner: spinnerStream.value,
    display: messageStream.value,
  };
}

[Phoenix-Alpha]

For me done works but saveChat seems to be not working in some cases. I guess it's when the message content is large and I move to different chat page before saving finishes. What is the solution for this?

[MaxLeiter]

Hey @Phoenix-Alpha if this is still an issue, please open a new issue with more details. Closing as we haven't had new reports of onSetAIState not working.

[Godrules500]

@MaxLeiter I am still having this issue.

If I set the state with useAIState, it is not updating properly. It's causing an issue when switching between chats where the last message and prompt is left off because useAIState doesn't seem to be working properly. 364