chore: update to ai sdk v5 beta (#1074)
This commit is contained in:
Jeremy
GitHub
parent
7d8e71383f
commit
4c281fe09d
54 changed files with 1372 additions and 1060 deletions
|
|
@ -1,8 +1,9 @@
|
|||
import {
|
||||
appendClientMessage,
|
||||
appendResponseMessages,
|
||||
createDataStream,
|
||||
convertToModelMessages,
|
||||
createUIMessageStream,
|
||||
JsonToSseTransformStream,
|
||||
smoothStream,
|
||||
stepCountIs,
|
||||
streamText,
|
||||
} from 'ai';
|
||||
import { auth, type UserType } from '@/app/(auth)/auth';
|
||||
|
|
@ -13,11 +14,10 @@ import {
|
|||
getChatById,
|
||||
getMessageCountByUserId,
|
||||
getMessagesByChatId,
|
||||
getStreamIdsByChatId,
|
||||
saveChat,
|
||||
saveMessages,
|
||||
} from '@/lib/db/queries';
|
||||
import { generateUUID, getTrailingMessageId } from '@/lib/utils';
|
||||
import { convertToUIMessages, generateUUID } from '@/lib/utils';
|
||||
import { generateTitleFromUserMessage } from '../../actions';
|
||||
import { createDocument } from '@/lib/ai/tools/create-document';
|
||||
import { updateDocument } from '@/lib/ai/tools/update-document';
|
||||
|
|
@ -33,15 +33,16 @@ import {
|
|||
type ResumableStreamContext,
|
||||
} from 'resumable-stream';
|
||||
import { after } from 'next/server';
|
||||
import type { Chat } from '@/lib/db/schema';
|
||||
import { differenceInSeconds } from 'date-fns';
|
||||
import { ChatSDKError } from '@/lib/errors';
|
||||
import type { ChatMessage } from '@/lib/types';
|
||||
import type { ChatModel } from '@/lib/ai/models';
|
||||
import type { VisibilityType } from '@/components/visibility-selector';
|
||||
|
||||
export const maxDuration = 60;
|
||||
|
||||
let globalStreamContext: ResumableStreamContext | null = null;
|
||||
|
||||
function getStreamContext() {
|
||||
export function getStreamContext() {
|
||||
if (!globalStreamContext) {
|
||||
try {
|
||||
globalStreamContext = createResumableStreamContext({
|
||||
|
|
@ -72,8 +73,17 @@ export async function POST(request: Request) {
|
|||
}
|
||||
|
||||
try {
|
||||
const { id, message, selectedChatModel, selectedVisibilityType } =
|
||||
requestBody;
|
||||
const {
|
||||
id,
|
||||
message,
|
||||
selectedChatModel,
|
||||
selectedVisibilityType,
|
||||
}: {
|
||||
id: string;
|
||||
message: ChatMessage;
|
||||
selectedChatModel: ChatModel['id'];
|
||||
selectedVisibilityType: VisibilityType;
|
||||
} = requestBody;
|
||||
|
||||
const session = await auth();
|
||||
|
||||
|
|
@ -111,13 +121,8 @@ export async function POST(request: Request) {
|
|||
}
|
||||
}
|
||||
|
||||
const previousMessages = await getMessagesByChatId({ id });
|
||||
|
||||
const messages = appendClientMessage({
|
||||
// @ts-expect-error: todo add type conversion from DBMessage[] to UIMessage[]
|
||||
messages: previousMessages,
|
||||
message,
|
||||
});
|
||||
const messagesFromDb = await getMessagesByChatId({ id });
|
||||
const uiMessages = [message, ...convertToUIMessages(messagesFromDb)];
|
||||
|
||||
const { longitude, latitude, city, country } = geolocation(request);
|
||||
|
||||
|
|
@ -135,7 +140,7 @@ export async function POST(request: Request) {
|
|||
id: message.id,
|
||||
role: 'user',
|
||||
parts: message.parts,
|
||||
attachments: message.experimental_attachments ?? [],
|
||||
attachments: [],
|
||||
createdAt: new Date(),
|
||||
},
|
||||
],
|
||||
|
|
@ -144,13 +149,13 @@ export async function POST(request: Request) {
|
|||
const streamId = generateUUID();
|
||||
await createStreamId({ streamId, chatId: id });
|
||||
|
||||
const stream = createDataStream({
|
||||
execute: (dataStream) => {
|
||||
const stream = createUIMessageStream({
|
||||
execute: ({ writer: dataStream }) => {
|
||||
const result = streamText({
|
||||
model: myProvider.languageModel(selectedChatModel),
|
||||
system: systemPrompt({ selectedChatModel, requestHints }),
|
||||
messages,
|
||||
maxSteps: 5,
|
||||
messages: convertToModelMessages(uiMessages),
|
||||
stopWhen: stepCountIs(5),
|
||||
experimental_activeTools:
|
||||
selectedChatModel === 'chat-model-reasoning'
|
||||
? []
|
||||
|
|
@ -161,7 +166,6 @@ export async function POST(request: Request) {
|
|||
'requestSuggestions',
|
||||
],
|
||||
experimental_transform: smoothStream({ chunking: 'word' }),
|
||||
experimental_generateMessageId: generateUUID,
|
||||
tools: {
|
||||
getWeather,
|
||||
createDocument: createDocument({ session, dataStream }),
|
||||
|
|
@ -171,42 +175,6 @@ export async function POST(request: Request) {
|
|||
dataStream,
|
||||
}),
|
||||
},
|
||||
onFinish: async ({ response }) => {
|
||||
if (session.user?.id) {
|
||||
try {
|
||||
const assistantId = getTrailingMessageId({
|
||||
messages: response.messages.filter(
|
||||
(message) => message.role === 'assistant',
|
||||
),
|
||||
});
|
||||
|
||||
if (!assistantId) {
|
||||
throw new Error('No assistant message found!');
|
||||
}
|
||||
|
||||
const [, assistantMessage] = appendResponseMessages({
|
||||
messages: [message],
|
||||
responseMessages: response.messages,
|
||||
});
|
||||
|
||||
await saveMessages({
|
||||
messages: [
|
||||
{
|
||||
id: assistantId,
|
||||
chatId: id,
|
||||
role: assistantMessage.role,
|
||||
parts: assistantMessage.parts,
|
||||
attachments:
|
||||
assistantMessage.experimental_attachments ?? [],
|
||||
createdAt: new Date(),
|
||||
},
|
||||
],
|
||||
});
|
||||
} catch (_) {
|
||||
console.error('Failed to save chat');
|
||||
}
|
||||
}
|
||||
},
|
||||
experimental_telemetry: {
|
||||
isEnabled: isProductionEnvironment,
|
||||
functionId: 'stream-text',
|
||||
|
|
@ -215,11 +183,27 @@ export async function POST(request: Request) {
|
|||
|
||||
result.consumeStream();
|
||||
|
||||
result.mergeIntoDataStream(dataStream, {
|
||||
sendReasoning: true,
|
||||
dataStream.merge(
|
||||
result.toUIMessageStream({
|
||||
sendReasoning: true,
|
||||
}),
|
||||
);
|
||||
},
|
||||
generateId: generateUUID,
|
||||
onFinish: async ({ messages }) => {
|
||||
await saveMessages({
|
||||
messages: messages.map((message) => ({
|
||||
id: message.id,
|
||||
role: message.role,
|
||||
parts: message.parts,
|
||||
createdAt: new Date(),
|
||||
attachments: [],
|
||||
chatId: id,
|
||||
})),
|
||||
});
|
||||
},
|
||||
onError: () => {
|
||||
onError: (error) => {
|
||||
console.log(error);
|
||||
return 'Oops, an error occurred!';
|
||||
},
|
||||
});
|
||||
|
|
@ -228,7 +212,9 @@ export async function POST(request: Request) {
|
|||
|
||||
if (streamContext) {
|
||||
return new Response(
|
||||
await streamContext.resumableStream(streamId, () => stream),
|
||||
await streamContext.resumableStream(streamId, () =>
|
||||
stream.pipeThrough(new JsonToSseTransformStream()),
|
||||
),
|
||||
);
|
||||
} else {
|
||||
return new Response(stream);
|
||||
|
|
@ -240,101 +226,6 @@ export async function POST(request: Request) {
|
|||
}
|
||||
}
|
||||
|
||||
export async function GET(request: Request) {
|
||||
const streamContext = getStreamContext();
|
||||
const resumeRequestedAt = new Date();
|
||||
|
||||
if (!streamContext) {
|
||||
return new Response(null, { status: 204 });
|
||||
}
|
||||
|
||||
const { searchParams } = new URL(request.url);
|
||||
const chatId = searchParams.get('chatId');
|
||||
|
||||
if (!chatId) {
|
||||
return new ChatSDKError('bad_request:api').toResponse();
|
||||
}
|
||||
|
||||
const session = await auth();
|
||||
|
||||
if (!session?.user) {
|
||||
return new ChatSDKError('unauthorized:chat').toResponse();
|
||||
}
|
||||
|
||||
let chat: Chat;
|
||||
|
||||
try {
|
||||
chat = await getChatById({ id: chatId });
|
||||
} catch {
|
||||
return new ChatSDKError('not_found:chat').toResponse();
|
||||
}
|
||||
|
||||
if (!chat) {
|
||||
return new ChatSDKError('not_found:chat').toResponse();
|
||||
}
|
||||
|
||||
if (chat.visibility === 'private' && chat.userId !== session.user.id) {
|
||||
return new ChatSDKError('forbidden:chat').toResponse();
|
||||
}
|
||||
|
||||
const streamIds = await getStreamIdsByChatId({ chatId });
|
||||
|
||||
if (!streamIds.length) {
|
||||
return new ChatSDKError('not_found:stream').toResponse();
|
||||
}
|
||||
|
||||
const recentStreamId = streamIds.at(-1);
|
||||
|
||||
if (!recentStreamId) {
|
||||
return new ChatSDKError('not_found:stream').toResponse();
|
||||
}
|
||||
|
||||
const emptyDataStream = createDataStream({
|
||||
execute: () => {},
|
||||
});
|
||||
|
||||
const stream = await streamContext.resumableStream(
|
||||
recentStreamId,
|
||||
() => emptyDataStream,
|
||||
);
|
||||
|
||||
/*
|
||||
* For when the generation is streaming during SSR
|
||||
* but the resumable stream has concluded at this point.
|
||||
*/
|
||||
if (!stream) {
|
||||
const messages = await getMessagesByChatId({ id: chatId });
|
||||
const mostRecentMessage = messages.at(-1);
|
||||
|
||||
if (!mostRecentMessage) {
|
||||
return new Response(emptyDataStream, { status: 200 });
|
||||
}
|
||||
|
||||
if (mostRecentMessage.role !== 'assistant') {
|
||||
return new Response(emptyDataStream, { status: 200 });
|
||||
}
|
||||
|
||||
const messageCreatedAt = new Date(mostRecentMessage.createdAt);
|
||||
|
||||
if (differenceInSeconds(resumeRequestedAt, messageCreatedAt) > 15) {
|
||||
return new Response(emptyDataStream, { status: 200 });
|
||||
}
|
||||
|
||||
const restoredStream = createDataStream({
|
||||
execute: (buffer) => {
|
||||
buffer.writeData({
|
||||
type: 'append-message',
|
||||
message: JSON.stringify(mostRecentMessage),
|
||||
});
|
||||
},
|
||||
});
|
||||
|
||||
return new Response(restoredStream, { status: 200 });
|
||||
}
|
||||
|
||||
return new Response(stream, { status: 200 });
|
||||
}
|
||||
|
||||
export async function DELETE(request: Request) {
|
||||
const { searchParams } = new URL(request.url);
|
||||
const id = searchParams.get('id');
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue