feat: add reasoning model (#750)

Co-authored-by: Matt Apperson <me@mattapperson.com>

lib/ai/custom-middleware.ts

@@ -1,3 +0,0 @@
-import type { Experimental_LanguageModelV1Middleware } from 'ai';
-
-export const customMiddleware: Experimental_LanguageModelV1Middleware = {};

lib/ai/index.ts

@@ -1,13 +0,0 @@
-import { openai } from '@ai-sdk/openai';
-import { experimental_wrapLanguageModel as wrapLanguageModel } from 'ai';
-
-import { customMiddleware } from './custom-middleware';
-
-export const customModel = (apiIdentifier: string) => {
-  return wrapLanguageModel({
-    model: openai(apiIdentifier),
-    middleware: customMiddleware,
-  });
-};
-
-export const imageGenerationModel = openai.image('dall-e-3');

lib/ai/models.ts

@@ -1,25 +1,49 @@
-// Define your models here.
+import { openai } from '@ai-sdk/openai';
+import { fireworks } from '@ai-sdk/fireworks';
+import {
+  customProvider,
+  extractReasoningMiddleware,
+  wrapLanguageModel,
+} from 'ai';
 
-export interface Model {
+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'),
+    'block-model': openai('gpt-4o-mini'),
+  },
+  imageModels: {
+    'small-model': openai.image('dall-e-3'),
+  },
+});
+
+interface ChatModel {
   id: string;
-  label: string;
-  apiIdentifier: string;
+  name: string;
   description: string;
 }
 
-export const models: Array<Model> = [
+export const chatModels: Array<ChatModel> = [
   {
-    id: 'gpt-4o-mini',
-    label: 'GPT 4o mini',
-    apiIdentifier: 'gpt-4o-mini',
+    id: 'chat-model-small',
+    name: 'Small model',
     description: 'Small model for fast, lightweight tasks',
   },
   {
-    id: 'gpt-4o',
-    label: 'GPT 4o',
-    apiIdentifier: 'gpt-4o',
-    description: 'For complex, multi-step tasks',
+    id: 'chat-model-large',
+    name: 'Large model',
+    description: 'Large model for complex, multi-step tasks',
   },
-] as const;
-
-export const DEFAULT_MODEL_NAME: string = 'gpt-4o-mini';
+  {
+    id: 'chat-model-reasoning',
+    name: 'Reasoning model',
+    description: 'Uses advanced reasoning',
+  },
+];

lib/ai/tools/create-document.ts

@@ -8,23 +8,17 @@ import {
   tool,
 } from 'ai';
 import { z } from 'zod';
-import { customModel, imageGenerationModel } from '..';
 import { codePrompt, sheetPrompt } from '../prompts';
 import { saveDocument } from '@/lib/db/queries';
 import { Session } from 'next-auth';
-import { Model } from '../models';
+import { myProvider } from '../models';
 
 interface CreateDocumentProps {
-  model: Model;
   session: Session;
   dataStream: DataStreamWriter;
 }
 
-export const createDocument = ({
-  model,
-  session,
-  dataStream,
-}: CreateDocumentProps) =>
+export const createDocument = ({ session, dataStream }: CreateDocumentProps) =>
   tool({
     description:
       'Create a document for a writing or content creation activities. This tool will call other functions that will generate the contents of the document based on the title and kind.',
@@ -58,7 +52,7 @@ export const createDocument = ({
 
       if (kind === 'text') {
         const { fullStream } = streamText({
-          model: customModel(model.apiIdentifier),
+          model: myProvider.languageModel('block-model'),
           system:
             'Write about the given topic. Markdown is supported. Use headings wherever appropriate.',
           experimental_transform: smoothStream({ chunking: 'word' }),
@@ -82,7 +76,7 @@ export const createDocument = ({
         dataStream.writeData({ type: 'finish', content: '' });
       } else if (kind === 'code') {
         const { fullStream } = streamObject({
-          model: customModel(model.apiIdentifier),
+          model: myProvider.languageModel('block-model'),
           system: codePrompt,
           prompt: title,
           schema: z.object({
@@ -111,7 +105,7 @@ export const createDocument = ({
         dataStream.writeData({ type: 'finish', content: '' });
       } else if (kind === 'image') {
         const { image } = await experimental_generateImage({
-          model: imageGenerationModel,
+          model: myProvider.imageModel('small-model'),
           prompt: title,
           n: 1,
         });
@@ -126,7 +120,7 @@ export const createDocument = ({
         dataStream.writeData({ type: 'finish', content: '' });
       } else if (kind === 'sheet') {
         const { fullStream } = streamObject({
-          model: customModel(model.apiIdentifier),
+          model: myProvider.languageModel('block-model'),
           system: sheetPrompt,
           prompt: title,
           schema: z.object({

lib/ai/tools/request-suggestions.ts

@@ -1,20 +1,17 @@
 import { z } from 'zod';
-import { Model } from '../models';
 import { Session } from 'next-auth';
 import { DataStreamWriter, streamObject, tool } from 'ai';
 import { getDocumentById, saveSuggestions } from '@/lib/db/queries';
 import { Suggestion } from '@/lib/db/schema';
-import { customModel } from '..';
 import { generateUUID } from '@/lib/utils';
+import { myProvider } from '../models';
 
 interface RequestSuggestionsProps {
-  model: Model;
   session: Session;
   dataStream: DataStreamWriter;
 }
 
 export const requestSuggestions = ({
-  model,
   session,
   dataStream,
 }: RequestSuggestionsProps) =>
@@ -39,7 +36,7 @@ export const requestSuggestions = ({
       > = [];
 
       const { elementStream } = streamObject({
-        model: customModel(model.apiIdentifier),
+        model: myProvider.languageModel('block-model'),
         system:
           'You are a help writing assistant. Given a piece of writing, please offer suggestions to improve the piece of writing and describe the change. It is very important for the edits to contain full sentences instead of just words. Max 5 suggestions.',
         prompt: document.content,

lib/ai/tools/update-document.ts

@@ -6,24 +6,18 @@ import {
   streamText,
   tool,
 } from 'ai';
-import { Model } from '../models';
 import { Session } from 'next-auth';
 import { z } from 'zod';
 import { getDocumentById, saveDocument } from '@/lib/db/queries';
-import { customModel, imageGenerationModel } from '..';
 import { updateDocumentPrompt } from '../prompts';
+import { myProvider } from '../models';
 
 interface UpdateDocumentProps {
-  model: Model;
   session: Session;
   dataStream: DataStreamWriter;
 }
 
-export const updateDocument = ({
-  model,
-  session,
-  dataStream,
-}: UpdateDocumentProps) =>
+export const updateDocument = ({ session, dataStream }: UpdateDocumentProps) =>
   tool({
     description: 'Update a document with the given description.',
     parameters: z.object({
@@ -51,7 +45,7 @@ export const updateDocument = ({
 
       if (document.kind === 'text') {
         const { fullStream } = streamText({
-          model: customModel(model.apiIdentifier),
+          model: myProvider.languageModel('block-model'),
           system: updateDocumentPrompt(currentContent, 'text'),
           experimental_transform: smoothStream({ chunking: 'word' }),
           prompt: description,
@@ -82,7 +76,7 @@ export const updateDocument = ({
         dataStream.writeData({ type: 'finish', content: '' });
       } else if (document.kind === 'code') {
         const { fullStream } = streamObject({
-          model: customModel(model.apiIdentifier),
+          model: myProvider.languageModel('block-model'),
           system: updateDocumentPrompt(currentContent, 'code'),
           prompt: description,
           schema: z.object({
@@ -111,7 +105,7 @@ export const updateDocument = ({
         dataStream.writeData({ type: 'finish', content: '' });
       } else if (document.kind === 'image') {
         const { image } = await experimental_generateImage({
-          model: imageGenerationModel,
+          model: myProvider.imageModel('image-model'),
           prompt: description,
           n: 1,
         });
@@ -126,7 +120,7 @@ export const updateDocument = ({
         dataStream.writeData({ type: 'finish', content: '' });
       } else if (document.kind === 'sheet') {
         const { fullStream } = streamObject({
-          model: customModel(model.apiIdentifier),
+          model: myProvider.languageModel('block-model'),
           system: updateDocumentPrompt(currentContent, 'sheet'),
           prompt: description,
           schema: z.object({

lib/utils.ts

@@ -1,9 +1,10 @@
 import type {
   CoreAssistantMessage,
-  CoreMessage,
   CoreToolMessage,
   Message,
+  TextStreamPart,
   ToolInvocation,
+  ToolSet,
 } from 'ai';
 import { type ClassValue, clsx } from 'clsx';
 import { twMerge } from 'tailwind-merge';
@@ -96,6 +97,7 @@ export function convertToUIMessages(
     }
 
     let textContent = '';
+    let reasoning: string | undefined = undefined;
     const toolInvocations: Array<ToolInvocation> = [];
 
     if (typeof message.content === 'string') {
@@ -111,6 +113,8 @@ export function convertToUIMessages(
             toolName: content.toolName,
             args: content.args,
           });
+        } else if (content.type === 'reasoning') {
+          reasoning = content.reasoning;
         }
       }
     }
@@ -119,6 +123,7 @@ export function convertToUIMessages(
       id: message.id,
       role: message.role as Message['role'],
       content: textContent,
+      reasoning,
       toolInvocations,
     });
 
@@ -129,9 +134,13 @@ export function convertToUIMessages(
 type ResponseMessageWithoutId = CoreToolMessage | CoreAssistantMessage;
 type ResponseMessage = ResponseMessageWithoutId & { id: string };
 
-export function sanitizeResponseMessages(
-  messages: Array<ResponseMessage>,
-): Array<ResponseMessage> {
+export function sanitizeResponseMessages({
+  messages,
+  reasoning,
+}: {
+  messages: Array<ResponseMessage>;
+  reasoning: string | undefined;
+}) {
   const toolResultIds: Array<string> = [];
 
   for (const message of messages) {
@@ -157,6 +166,11 @@ export function sanitizeResponseMessages(
           : true,
     );
 
+    if (reasoning) {
+      // @ts-expect-error: reasoning message parts in sdk is wip
+      sanitizedContent.push({ type: 'reasoning', reasoning });
+    }
+
     return {
       ...message,
       content: sanitizedContent,