71 lines
2.1 KiB
TypeScript
71 lines
2.1 KiB
TypeScript
import { simulateReadableStream } from 'ai';
|
|
import { MockLanguageModelV1 } from 'ai/test';
|
|
import { getResponseChunksByPrompt } from '@/tests/prompts/utils';
|
|
|
|
export const chatModel = new MockLanguageModelV1({
|
|
doGenerate: async () => ({
|
|
rawCall: { rawPrompt: null, rawSettings: {} },
|
|
finishReason: 'stop',
|
|
usage: { promptTokens: 10, completionTokens: 20 },
|
|
text: `Hello, world!`,
|
|
}),
|
|
doStream: async ({ prompt }) => ({
|
|
stream: simulateReadableStream({
|
|
chunks: getResponseChunksByPrompt(prompt),
|
|
}),
|
|
rawCall: { rawPrompt: null, rawSettings: {} },
|
|
}),
|
|
});
|
|
|
|
export const reasoningModel = new MockLanguageModelV1({
|
|
doGenerate: async () => ({
|
|
rawCall: { rawPrompt: null, rawSettings: {} },
|
|
finishReason: 'stop',
|
|
usage: { promptTokens: 10, completionTokens: 20 },
|
|
text: `Hello, world!`,
|
|
}),
|
|
doStream: async ({ prompt }) => ({
|
|
stream: simulateReadableStream({
|
|
chunks: getResponseChunksByPrompt(prompt, true),
|
|
}),
|
|
rawCall: { rawPrompt: null, rawSettings: {} },
|
|
}),
|
|
});
|
|
|
|
export const titleModel = new MockLanguageModelV1({
|
|
doGenerate: async () => ({
|
|
rawCall: { rawPrompt: null, rawSettings: {} },
|
|
finishReason: 'stop',
|
|
usage: { promptTokens: 10, completionTokens: 20 },
|
|
text: `This is a test title`,
|
|
}),
|
|
doStream: async () => ({
|
|
stream: simulateReadableStream({
|
|
chunks: [
|
|
{ type: 'text-delta', textDelta: 'This is a test title' },
|
|
{
|
|
type: 'finish',
|
|
finishReason: 'stop',
|
|
logprobs: undefined,
|
|
usage: { completionTokens: 10, promptTokens: 3 },
|
|
},
|
|
],
|
|
}),
|
|
rawCall: { rawPrompt: null, rawSettings: {} },
|
|
}),
|
|
});
|
|
|
|
export const artifactModel = new MockLanguageModelV1({
|
|
doGenerate: async () => ({
|
|
rawCall: { rawPrompt: null, rawSettings: {} },
|
|
finishReason: 'stop',
|
|
usage: { promptTokens: 10, completionTokens: 20 },
|
|
text: `Hello, world!`,
|
|
}),
|
|
doStream: async ({ prompt }) => ({
|
|
stream: simulateReadableStream({
|
|
chunks: getResponseChunksByPrompt(prompt),
|
|
}),
|
|
rawCall: { rawPrompt: null, rawSettings: {} },
|
|
}),
|
|
});
|