feat: add tests (#843)
This commit is contained in:
parent
95a2af2535
commit
9dd9a9898c
35 changed files with 1063 additions and 191 deletions
|
|
@ -4,9 +4,7 @@ import {
|
|||
smoothStream,
|
||||
streamText,
|
||||
} from 'ai';
|
||||
|
||||
import { auth } from '@/app/(auth)/auth';
|
||||
import { myProvider } from '@/lib/ai/models';
|
||||
import { systemPrompt } from '@/lib/ai/prompts';
|
||||
import {
|
||||
deleteChatById,
|
||||
|
|
@ -19,113 +17,124 @@ import {
|
|||
getMostRecentUserMessage,
|
||||
sanitizeResponseMessages,
|
||||
} from '@/lib/utils';
|
||||
|
||||
import { generateTitleFromUserMessage } from '../../actions';
|
||||
import { createDocument } from '@/lib/ai/tools/create-document';
|
||||
import { updateDocument } from '@/lib/ai/tools/update-document';
|
||||
import { requestSuggestions } from '@/lib/ai/tools/request-suggestions';
|
||||
import { getWeather } from '@/lib/ai/tools/get-weather';
|
||||
import { isProductionEnvironment } from '@/lib/constants';
|
||||
import { NextResponse } from 'next/server';
|
||||
import { myProvider } from '@/lib/ai/providers';
|
||||
|
||||
export const maxDuration = 60;
|
||||
|
||||
export async function POST(request: Request) {
|
||||
const {
|
||||
id,
|
||||
messages,
|
||||
selectedChatModel,
|
||||
}: { id: string; messages: Array<Message>; selectedChatModel: string } =
|
||||
await request.json();
|
||||
try {
|
||||
const {
|
||||
id,
|
||||
messages,
|
||||
selectedChatModel,
|
||||
}: {
|
||||
id: string;
|
||||
messages: Array<Message>;
|
||||
selectedChatModel: string;
|
||||
} = await request.json();
|
||||
|
||||
const session = await auth();
|
||||
const session = await auth();
|
||||
|
||||
if (!session || !session.user || !session.user.id) {
|
||||
return new Response('Unauthorized', { status: 401 });
|
||||
}
|
||||
if (!session || !session.user || !session.user.id) {
|
||||
return new Response('Unauthorized', { status: 401 });
|
||||
}
|
||||
|
||||
const userMessage = getMostRecentUserMessage(messages);
|
||||
const userMessage = getMostRecentUserMessage(messages);
|
||||
|
||||
if (!userMessage) {
|
||||
return new Response('No user message found', { status: 400 });
|
||||
}
|
||||
if (!userMessage) {
|
||||
return new Response('No user message found', { status: 400 });
|
||||
}
|
||||
|
||||
const chat = await getChatById({ id });
|
||||
const chat = await getChatById({ id });
|
||||
|
||||
if (!chat) {
|
||||
const title = await generateTitleFromUserMessage({ message: userMessage });
|
||||
await saveChat({ id, userId: session.user.id, title });
|
||||
}
|
||||
if (!chat) {
|
||||
const title = await generateTitleFromUserMessage({
|
||||
message: userMessage,
|
||||
});
|
||||
await saveChat({ id, userId: session.user.id, title });
|
||||
}
|
||||
|
||||
await saveMessages({
|
||||
messages: [{ ...userMessage, createdAt: new Date(), chatId: id }],
|
||||
});
|
||||
await saveMessages({
|
||||
messages: [{ ...userMessage, createdAt: new Date(), chatId: id }],
|
||||
});
|
||||
|
||||
return createDataStreamResponse({
|
||||
execute: (dataStream) => {
|
||||
const result = streamText({
|
||||
model: myProvider.languageModel(selectedChatModel),
|
||||
system: systemPrompt({ selectedChatModel }),
|
||||
messages,
|
||||
maxSteps: 5,
|
||||
experimental_activeTools:
|
||||
selectedChatModel === 'chat-model-reasoning'
|
||||
? []
|
||||
: [
|
||||
'getWeather',
|
||||
'createDocument',
|
||||
'updateDocument',
|
||||
'requestSuggestions',
|
||||
],
|
||||
experimental_transform: smoothStream({ chunking: 'word' }),
|
||||
experimental_generateMessageId: generateUUID,
|
||||
tools: {
|
||||
getWeather,
|
||||
createDocument: createDocument({ session, dataStream }),
|
||||
updateDocument: updateDocument({ session, dataStream }),
|
||||
requestSuggestions: requestSuggestions({
|
||||
session,
|
||||
dataStream,
|
||||
}),
|
||||
},
|
||||
onFinish: async ({ response, reasoning }) => {
|
||||
if (session.user?.id) {
|
||||
try {
|
||||
const sanitizedResponseMessages = sanitizeResponseMessages({
|
||||
messages: response.messages,
|
||||
reasoning,
|
||||
});
|
||||
return createDataStreamResponse({
|
||||
execute: (dataStream) => {
|
||||
const result = streamText({
|
||||
model: myProvider.languageModel(selectedChatModel),
|
||||
system: systemPrompt({ selectedChatModel }),
|
||||
messages,
|
||||
maxSteps: 5,
|
||||
experimental_activeTools:
|
||||
selectedChatModel === 'chat-model-reasoning'
|
||||
? []
|
||||
: [
|
||||
'getWeather',
|
||||
'createDocument',
|
||||
'updateDocument',
|
||||
'requestSuggestions',
|
||||
],
|
||||
experimental_transform: smoothStream({ chunking: 'word' }),
|
||||
experimental_generateMessageId: generateUUID,
|
||||
tools: {
|
||||
getWeather,
|
||||
createDocument: createDocument({ session, dataStream }),
|
||||
updateDocument: updateDocument({ session, dataStream }),
|
||||
requestSuggestions: requestSuggestions({
|
||||
session,
|
||||
dataStream,
|
||||
}),
|
||||
},
|
||||
onFinish: async ({ response, reasoning }) => {
|
||||
if (session.user?.id) {
|
||||
try {
|
||||
const sanitizedResponseMessages = sanitizeResponseMessages({
|
||||
messages: response.messages,
|
||||
reasoning,
|
||||
});
|
||||
|
||||
await saveMessages({
|
||||
messages: sanitizedResponseMessages.map((message) => {
|
||||
return {
|
||||
id: message.id,
|
||||
chatId: id,
|
||||
role: message.role,
|
||||
content: message.content,
|
||||
createdAt: new Date(),
|
||||
};
|
||||
}),
|
||||
});
|
||||
} catch (error) {
|
||||
console.error('Failed to save chat');
|
||||
await saveMessages({
|
||||
messages: sanitizedResponseMessages.map((message) => {
|
||||
return {
|
||||
id: message.id,
|
||||
chatId: id,
|
||||
role: message.role,
|
||||
content: message.content,
|
||||
createdAt: new Date(),
|
||||
};
|
||||
}),
|
||||
});
|
||||
} catch (error) {
|
||||
console.error('Failed to save chat');
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
experimental_telemetry: {
|
||||
isEnabled: true,
|
||||
functionId: 'stream-text',
|
||||
},
|
||||
});
|
||||
},
|
||||
experimental_telemetry: {
|
||||
isEnabled: isProductionEnvironment,
|
||||
functionId: 'stream-text',
|
||||
},
|
||||
});
|
||||
|
||||
result.consumeStream();
|
||||
result.consumeStream();
|
||||
|
||||
result.mergeIntoDataStream(dataStream, {
|
||||
sendReasoning: true,
|
||||
});
|
||||
},
|
||||
onError: () => {
|
||||
return 'Oops, an error occured!';
|
||||
},
|
||||
});
|
||||
result.mergeIntoDataStream(dataStream, {
|
||||
sendReasoning: true,
|
||||
});
|
||||
},
|
||||
onError: () => {
|
||||
return 'Oops, an error occured!';
|
||||
},
|
||||
});
|
||||
} catch (error) {
|
||||
return NextResponse.json({ error }, { status: 400 });
|
||||
}
|
||||
}
|
||||
|
||||
export async function DELETE(request: Request) {
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue