feat: add tests (#843)

lib/ai/models.test.ts

@@ -0,0 +1,226 @@
+import { CoreMessage, FinishReason, simulateReadableStream } from 'ai';
+import { MockLanguageModelV1 } from 'ai/test';
+
+interface TextDeltaChunk {
+  type: 'text-delta';
+  textDelta: string;
+}
+
+interface FinishChunk {
+  type: 'finish';
+  finishReason: FinishReason;
+  logprobs: undefined;
+  usage: { completionTokens: number; promptTokens: number };
+}
+
+type Chunk = TextDeltaChunk | FinishChunk;
+
+const getResponseChunksByPrompt = (prompt: CoreMessage[]): Array<Chunk> => {
+  const userMessage = prompt.at(-1);
+
+  if (!userMessage) {
+    throw new Error('No user message found');
+  }
+
+  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 [];
+};
+
+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 () => ({
+    stream: simulateReadableStream({
+      chunks: [
+        { type: 'text-delta', textDelta: 'test' },
+        {
+          type: 'finish',
+          finishReason: 'stop',
+          logprobs: undefined,
+          usage: { completionTokens: 10, promptTokens: 3 },
+        },
+      ],
+    }),
+    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 () => ({
+    stream: simulateReadableStream({
+      chunks: [
+        { type: 'text-delta', textDelta: 'test' },
+        {
+          type: 'finish',
+          finishReason: 'stop',
+          logprobs: undefined,
+          usage: { completionTokens: 10, promptTokens: 3 },
+        },
+      ],
+    }),
+    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;
+}

lib/ai/models.ts

@@ -1,30 +1,5 @@
-import { openai } from '@ai-sdk/openai';
-import { fireworks } from '@ai-sdk/fireworks';
-import {
-  customProvider,
-  extractReasoningMiddleware,
-  wrapLanguageModel,
-} from 'ai';
-
 export const DEFAULT_CHAT_MODEL: string = 'chat-model-small';
 
-export const myProvider = customProvider({
-  languageModels: {
-    'chat-model-small': openai('gpt-4o-mini'),
-    'chat-model-large': openai('gpt-4o'),
-    'chat-model-reasoning': wrapLanguageModel({
-      model: fireworks('accounts/fireworks/models/deepseek-r1'),
-      middleware: extractReasoningMiddleware({ tagName: 'think' }),
-    }),
-    'title-model': openai('gpt-4-turbo'),
-    'artifact-model': openai('gpt-4o-mini'),
-  },
-  imageModels: {
-    'small-model': openai.image('dall-e-2'),
-    'large-model': openai.image('dall-e-3'),
-  },
-});
-
 interface ChatModel {
   id: string;
   name: string;

lib/ai/providers.ts

@@ -0,0 +1,41 @@
+import {
+  customProvider,
+  extractReasoningMiddleware,
+  wrapLanguageModel,
+} from 'ai';
+import { openai } from '@ai-sdk/openai';
+import { fireworks } from '@ai-sdk/fireworks';
+import { isTestEnvironment } from '../constants';
+import {
+  artifactModel,
+  chatModel,
+  reasoningModel,
+  titleModel,
+} from './models.test';
+
+export const myProvider = isTestEnvironment
+  ? customProvider({
+      languageModels: {
+        'chat-model-small': chatModel,
+        'chat-model-large': chatModel,
+        'chat-model-reasoning': reasoningModel,
+        'title-model': titleModel,
+        'artifact-model': artifactModel,
+      },
+    })
+  : customProvider({
+      languageModels: {
+        'chat-model-small': openai('gpt-4o-mini'),
+        'chat-model-large': openai('gpt-4o'),
+        'chat-model-reasoning': wrapLanguageModel({
+          model: fireworks('accounts/fireworks/models/deepseek-r1'),
+          middleware: extractReasoningMiddleware({ tagName: 'think' }),
+        }),
+        'title-model': openai('gpt-4-turbo'),
+        'artifact-model': openai('gpt-4o-mini'),
+      },
+      imageModels: {
+        'small-model': openai.image('dall-e-2'),
+        'large-model': openai.image('dall-e-3'),
+      },
+    });

lib/ai/tools/request-suggestions.ts

@@ -4,7 +4,7 @@ import { DataStreamWriter, streamObject, tool } from 'ai';
 import { getDocumentById, saveSuggestions } from '@/lib/db/queries';
 import { Suggestion } from '@/lib/db/schema';
 import { generateUUID } from '@/lib/utils';
-import { myProvider } from '../models';
+import { myProvider } from '../providers';
 
 interface RequestSuggestionsProps {
   session: Session;