useChat to support streamText with structured output

[mrasoahaingo]

Feature Description

Thanks to the recent version of AI SDK we can now set the experimental_output to streamText function.

According to the docs we can then get the streamed object from experimental_partialOutputStream.

If we implement the useChat response we can do something like this:

export async function POST(request: Request) {
	const { messages } = await request.json();

	return createDataStreamResponse({
		status: 200,
		statusText: "OK",
		headers: {
			"Custom-Header": "value",
		},
		async execute(dataStream) {
			const result = streamText({
				experimental_output: Output.object({
					schema: z.object({
						answer: z.string().describe("The answer of the assistant"),
						suggestedUserAnswers: z
							.array(z.string())
							.describe("The suggested user answers"),
					}),
				}),
				model,
				system,
				messages: convertToCoreMessages(messages),
				tools: {
					... // Some tools
				},
			});

			result.mergeIntoDataStream(dataStream);
		},
		onError: (error: unknown) => `Custom error: ${(error as Error).message}`,
	});
}

But in the front, you end by displaying stream stringified json which cannot be parsed because the json is incomplete.

So I found a work around where I can stream the answer text and send the suggestedUserAnswers through annotation...

export async function POST(request: Request) {
	const { messages } = await request.json();

	return createDataStreamResponse({
		status: 200,
		statusText: "OK",
		headers: {
			"Custom-Header": "value",
		},
		async execute(dataStream) {
			const result = streamText({
				experimental_output: Output.object({
					schema: z.object({
						answer: z.string().describe("The answer of the assistant"),
						suggestedUserAnswers: z
							.array(z.string())
							.describe("The suggested user answers"),
					}),
				}),
				model,
				system,
				messages: convertToCoreMessages(messages),
				tools: {
					... // Some tools
				},
			});

			const { experimental_partialOutputStream: partialOutputStream } = result;

			let previousText = '';
			let previousSuggestedUserAnswers: string[] = [];

			const stream = new ReadableStream({
				async start(controller) {
					for await (const chunk of partialOutputStream) {
						const { answer, suggestedUserAnswers } = chunk;
						const textDiff = getTextDiff(answer ?? '', previousText);

						controller.enqueue(formatDataStreamPart('text', textDiff));
						
						previousText = answer ?? '';
						previousSuggestedUserAnswers = suggestedUserAnswers?.filter((answer): answer is string => answer !== undefined) ?? [];
					}

					dataStream.writeMessageAnnotation({
						suggestedUserAnswers: previousSuggestedUserAnswers,
					});

					controller.close();
				},
			});

			dataStream.merge(stream);
		},
		onError: (error: unknown) => `Custom error: ${(error as Error).message}`,
	});
}

It works... but it's very hacky...

Use Cases

No response

Additional context

No response

[lgrammel]

Would useObject work w/ your use case?

[mrasoahaingo]

Hello, actually useObject doesn't have messages like useChat.
I also tried to set data on message but it's only possible when you use assistant (with part code = 6)

I don't know if it's possible to improve useChat because currently the message content only receive the deltas that's why I had to create a function getTextDiff

[mrasoahaingo]

Maybe if it can identify if it's a plain text then it should display the append the chunk (like actual useChat), or if it's an object it should replace the object (like useObject). so the message.content can be a text OR an object? or maybe push into message.data like useAssistant

[mrasoahaingo]

Actually my workaround completely break the tool calls and other type of message

[mrasoahaingo]

OMG, I just found that chunked json can be safety parse on front side with import { parsePartialJson } from '@ai-sdk/ui-utils';

Now I can remove all the work around on backend side and simply result.mergeIntoDataStream(dataStream); as usual...

Thanks !