feat: add artifact tests (#859)

This commit is contained in:
Jeremy 2025-03-11 14:39:36 -07:00 committed by GitHub
parent 9628c54755
commit 8e561dced4
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
13 changed files with 669 additions and 255 deletions

View file

@ -1,180 +1,6 @@
import { CoreMessage, FinishReason, simulateReadableStream } from 'ai';
import { simulateReadableStream } from 'ai';
import { MockLanguageModelV1 } from 'ai/test';
interface ReasoningChunk {
type: 'reasoning';
textDelta: string;
}
interface TextDeltaChunk {
type: 'text-delta';
textDelta: string;
}
interface FinishChunk {
type: 'finish';
finishReason: FinishReason;
logprobs: undefined;
usage: { completionTokens: number; promptTokens: number };
}
type Chunk = TextDeltaChunk | ReasoningChunk | FinishChunk;
const getResponseChunksByPrompt = (
prompt: CoreMessage[],
isReasoningEnabled: boolean = false,
): Array<Chunk> => {
const userMessage = prompt.at(-1);
if (!userMessage) {
throw new Error('No user message found');
}
if (isReasoningEnabled) {
if (
compareMessages(userMessage, {
role: 'user',
content: [{ type: 'text', text: 'why is the sky blue?' }],
})
) {
return [
{ type: 'reasoning', textDelta: 'the ' },
{ type: 'reasoning', textDelta: 'sky ' },
{ type: 'reasoning', textDelta: 'is ' },
{ type: 'reasoning', textDelta: 'blue ' },
{ type: 'reasoning', textDelta: 'because ' },
{ type: 'reasoning', textDelta: 'of ' },
{ type: 'reasoning', textDelta: 'rayleigh ' },
{ type: 'reasoning', textDelta: 'scattering! ' },
{ type: 'text-delta', textDelta: "it's " },
{ type: 'text-delta', textDelta: 'just ' },
{ type: 'text-delta', textDelta: 'blue ' },
{ type: 'text-delta', textDelta: 'duh!' },
{
type: 'finish',
finishReason: 'stop',
logprobs: undefined,
usage: { completionTokens: 10, promptTokens: 3 },
},
];
} else if (
compareMessages(userMessage, {
role: 'user',
content: [{ type: 'text', text: 'why is grass green?' }],
})
) {
return [
{ type: 'reasoning', textDelta: 'grass ' },
{ type: 'reasoning', textDelta: 'is ' },
{ type: 'reasoning', textDelta: 'green ' },
{ type: 'reasoning', textDelta: 'because ' },
{ type: 'reasoning', textDelta: 'of ' },
{ type: 'reasoning', textDelta: 'chlorophyll ' },
{ type: 'reasoning', textDelta: 'absorption! ' },
{ type: 'text-delta', textDelta: "it's " },
{ type: 'text-delta', textDelta: 'just ' },
{ type: 'text-delta', textDelta: 'green ' },
{ type: 'text-delta', textDelta: 'duh!' },
{
type: 'finish',
finishReason: 'stop',
logprobs: undefined,
usage: { completionTokens: 10, promptTokens: 3 },
},
];
}
}
if (
compareMessages(userMessage, {
role: 'user',
content: [{ type: 'text', text: 'why is grass green?' }],
})
) {
return [
{ type: 'text-delta', textDelta: "it's " },
{ type: 'text-delta', textDelta: 'just ' },
{ type: 'text-delta', textDelta: 'green ' },
{ type: 'text-delta', textDelta: 'duh!' },
{
type: 'finish',
finishReason: 'stop',
logprobs: undefined,
usage: { completionTokens: 10, promptTokens: 3 },
},
];
} else if (
compareMessages(userMessage, {
role: 'user',
content: [{ type: 'text', text: 'why is the sky blue?' }],
})
) {
return [
{ type: 'text-delta', textDelta: "it's " },
{ type: 'text-delta', textDelta: 'just ' },
{ type: 'text-delta', textDelta: 'blue ' },
{ type: 'text-delta', textDelta: 'duh!' },
{
type: 'finish',
finishReason: 'stop',
logprobs: undefined,
usage: { completionTokens: 10, promptTokens: 3 },
},
];
} else if (
compareMessages(userMessage, {
role: 'user',
content: [
{ type: 'text', text: 'What are the advantages of using Next.js?' },
],
})
) {
return [
{ type: 'text-delta', textDelta: 'with ' },
{ type: 'text-delta', textDelta: 'next.js ' },
{ type: 'text-delta', textDelta: 'you ' },
{ type: 'text-delta', textDelta: 'can ' },
{ type: 'text-delta', textDelta: 'ship ' },
{ type: 'text-delta', textDelta: 'fast! ' },
{
type: 'finish',
finishReason: 'stop',
logprobs: undefined,
usage: { completionTokens: 10, promptTokens: 3 },
},
];
} else if (
compareMessages(userMessage, {
role: 'user',
content: [
{
type: 'text',
text: 'who painted this?',
},
{
type: 'image',
image: '...',
},
],
})
) {
return [
{ type: 'text-delta', textDelta: 'this ' },
{ type: 'text-delta', textDelta: 'painting ' },
{ type: 'text-delta', textDelta: 'is ' },
{ type: 'text-delta', textDelta: 'by ' },
{ type: 'text-delta', textDelta: 'monet!' },
{
type: 'finish',
finishReason: 'stop',
logprobs: undefined,
usage: { completionTokens: 10, promptTokens: 3 },
},
];
}
return [];
};
import { getResponseChunksByPrompt } from '@/tests/prompts/utils';
export const chatModel = new MockLanguageModelV1({
doGenerate: async () => ({
@ -236,46 +62,10 @@ export const artifactModel = new MockLanguageModelV1({
usage: { promptTokens: 10, completionTokens: 20 },
text: `Hello, world!`,
}),
doStream: async () => ({
doStream: async ({ prompt }) => ({
stream: simulateReadableStream({
chunks: [
{ type: 'text-delta', textDelta: 'test' },
{
type: 'finish',
finishReason: 'stop',
logprobs: undefined,
usage: { completionTokens: 10, promptTokens: 3 },
},
],
chunks: getResponseChunksByPrompt(prompt),
}),
rawCall: { rawPrompt: null, rawSettings: {} },
}),
});
function compareMessages(msg1: CoreMessage, msg2: CoreMessage): boolean {
if (msg1.role !== msg2.role) return false;
if (!Array.isArray(msg1.content) || !Array.isArray(msg2.content)) {
return false;
}
if (msg1.content.length !== msg2.content.length) return false;
for (let i = 0; i < msg1.content.length; i++) {
const item1 = msg1.content[i];
const item2 = msg2.content[i];
if (item1.type !== item2.type) return false;
if (item1.type === 'image' && item2.type === 'image') {
// if (item1.image.toString() !== item2.image.toString()) return false;
// if (item1.mimeType !== item2.mimeType) return false;
} else if (item1.type === 'text' && item2.type === 'text') {
if (item1.text !== item2.text) return false;
} else {
return false;
}
}
return true;
}