feat: add tests (#843)

This commit is contained in:
Jeremy 2025-03-04 17:25:46 -08:00 committed by GitHub
parent 95a2af2535
commit 9dd9a9898c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
35 changed files with 1063 additions and 191 deletions

View file

@ -9,7 +9,7 @@ import {
updateChatVisiblityById,
} from '@/lib/db/queries';
import { VisibilityType } from '@/components/visibility-selector';
import { myProvider } from '@/lib/ai/models';
import { myProvider } from '@/lib/ai/providers';
export async function saveChatModelAsCookie(model: string) {
const cookieStore = await cookies();

View file

@ -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) {