feat: add reasoning model (#750)

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

.env.example

@@ -1,6 +1,9 @@
-# Get your OpenAI API Key here: https://platform.openai.com/account/api-keys
+# Get your OpenAI API Key here for chat models: https://platform.openai.com/account/api-keys
 OPENAI_API_KEY=****
 
+# Get your Fireworks AI API Key here for reasoning models: https://fireworks.ai/account/api-keys
+FIREWORKS_API_KEY=****
+
 # Generate a random secret: https://generate-secret.vercel.app/32 or `openssl rand -base64 32`
 AUTH_SECRET=****
 

app/(chat)/actions.ts

@@ -1,19 +1,19 @@
 'use server';
 
-import { type CoreUserMessage, generateText, Message } from 'ai';
+import { generateText, Message } from 'ai';
 import { cookies } from 'next/headers';
 
-import { customModel } from '@/lib/ai';
 import {
   deleteMessagesByChatIdAfterTimestamp,
   getMessageById,
   updateChatVisiblityById,
 } from '@/lib/db/queries';
 import { VisibilityType } from '@/components/visibility-selector';
+import { myProvider } from '@/lib/ai/models';
 
-export async function saveModelId(model: string) {
+export async function saveChatModelAsCookie(model: string) {
   const cookieStore = await cookies();
-  cookieStore.set('model-id', model);
+  cookieStore.set('chat-model', model);
 }
 
 export async function generateTitleFromUserMessage({
@@ -22,7 +22,7 @@ export async function generateTitleFromUserMessage({
   message: Message;
 }) {
   const { text: title } = await generateText({
-    model: customModel('gpt-4o-mini'),
+    model: myProvider.languageModel('title-model'),
     system: `\n
     - you will generate a short title based on the first message a user begins a conversation with
     - ensure it is not more than 80 characters long

app/(chat)/api/chat/route.ts

@@ -3,11 +3,11 @@ import {
   createDataStreamResponse,
   smoothStream,
   streamText,
+  wrapLanguageModel,
 } from 'ai';
 
 import { auth } from '@/app/(auth)/auth';
-import { customModel } from '@/lib/ai';
+import { myProvider } from '@/lib/ai/models';
-import { models } from '@/lib/ai/models';
 import { systemPrompt } from '@/lib/ai/prompts';
 import {
   deleteChatById,
@@ -48,8 +48,8 @@ export async function POST(request: Request) {
   const {
     id,
     messages,
-    modelId,
+    selectedChatModel,
-  }: { id: string; messages: Array<Message>; modelId: string } =
+  }: { id: string; messages: Array<Message>; selectedChatModel: string } =
     await request.json();
 
   const session = await auth();
@@ -58,12 +58,6 @@ export async function POST(request: Request) {
     return new Response('Unauthorized', { status: 401 });
   }
 
-  const model = models.find((model) => model.id === modelId);
-
-  if (!model) {
-    return new Response('Model not found', { status: 404 });
-  }
-
   const userMessage = getMostRecentUserMessage(messages);
 
   if (!userMessage) {
@@ -84,7 +78,7 @@ export async function POST(request: Request) {
   return createDataStreamResponse({
     execute: (dataStream) => {
       const result = streamText({
-        model: customModel(model.apiIdentifier),
+        model: myProvider.languageModel(selectedChatModel),
         system: systemPrompt,
         messages,
         maxSteps: 5,
@@ -93,32 +87,31 @@ export async function POST(request: Request) {
         experimental_generateMessageId: generateUUID,
         tools: {
           getWeather,
-          createDocument: createDocument({ session, dataStream, model }),
+          createDocument: createDocument({ session, dataStream }),
-          updateDocument: updateDocument({ session, dataStream, model }),
+          updateDocument: updateDocument({ session, dataStream }),
           requestSuggestions: requestSuggestions({
             session,
             dataStream,
-            model,
           }),
         },
-        onFinish: async ({ response }) => {
+        onFinish: async ({ response, reasoning }) => {
           if (session.user?.id) {
             try {
-              const responseMessagesWithoutIncompleteToolCalls =
+              const sanitizedResponseMessages = sanitizeResponseMessages({
-                sanitizeResponseMessages(response.messages);
+                messages: response.messages,
+                reasoning,
+              });
 
               await saveMessages({
-                messages: responseMessagesWithoutIncompleteToolCalls.map(
+                messages: sanitizedResponseMessages.map((message) => {
-                  (message) => {
+                  return {
-                    return {
+                    id: message.id,
-                      id: message.id,
+                    chatId: id,
-                      chatId: id,
+                    role: message.role,
-                      role: message.role,
+                    content: message.content,
-                      content: message.content,
+                    createdAt: new Date(),
-                      createdAt: new Date(),
+                  };
-                    };
+                }),
-                  },
-                ),
               });
             } catch (error) {
               console.error('Failed to save chat');
@@ -131,7 +124,12 @@ export async function POST(request: Request) {
         },
       });
 
-      result.mergeIntoDataStream(dataStream);
+      result.mergeIntoDataStream(dataStream, {
+        sendReasoning: true,
+      });
+    },
+    onError: (error) => {
+      return 'Oops, an error occured!';
     },
   });
 }

app/(chat)/chat/[id]/page.tsx

@@ -3,10 +3,10 @@ import { notFound } from 'next/navigation';
 
 import { auth } from '@/app/(auth)/auth';
 import { Chat } from '@/components/chat';
-import { DEFAULT_MODEL_NAME, models } from '@/lib/ai/models';
 import { getChatById, getMessagesByChatId } from '@/lib/db/queries';
 import { convertToUIMessages } from '@/lib/utils';
 import { DataStreamHandler } from '@/components/data-stream-handler';
+import { DEFAULT_CHAT_MODEL } from '@/lib/ai/models';
 
 export default async function Page(props: { params: Promise<{ id: string }> }) {
   const params = await props.params;
@@ -34,17 +34,29 @@ export default async function Page(props: { params: Promise<{ id: string }> }) {
   });
 
   const cookieStore = await cookies();
-  const modelIdFromCookie = cookieStore.get('model-id')?.value;
+  const chatModelFromCookie = cookieStore.get('chat-model');
-  const selectedModelId =
+
-    models.find((model) => model.id === modelIdFromCookie)?.id ||
+  if (!chatModelFromCookie) {
-    DEFAULT_MODEL_NAME;
+    return (
+      <>
+        <Chat
+          id={chat.id}
+          initialMessages={convertToUIMessages(messagesFromDb)}
+          selectedChatModel={DEFAULT_CHAT_MODEL}
+          selectedVisibilityType={chat.visibility}
+          isReadonly={session?.user?.id !== chat.userId}
+        />
+        <DataStreamHandler id={id} />
+      </>
+    );
+  }
 
   return (
     <>
       <Chat
         id={chat.id}
         initialMessages={convertToUIMessages(messagesFromDb)}
-        selectedModelId={selectedModelId}
+        selectedChatModel={chatModelFromCookie.value}
         selectedVisibilityType={chat.visibility}
         isReadonly={session?.user?.id !== chat.userId}
       />

app/(chat)/page.tsx

@@ -1,7 +1,7 @@
 import { cookies } from 'next/headers';
 
 import { Chat } from '@/components/chat';
-import { DEFAULT_MODEL_NAME, models } from '@/lib/ai/models';
+import { DEFAULT_CHAT_MODEL } from '@/lib/ai/models';
 import { generateUUID } from '@/lib/utils';
 import { DataStreamHandler } from '@/components/data-stream-handler';
 
@@ -9,11 +9,23 @@ export default async function Page() {
   const id = generateUUID();
 
   const cookieStore = await cookies();
-  const modelIdFromCookie = cookieStore.get('model-id')?.value;
+  const modelIdFromCookie = cookieStore.get('chat-model');
 
-  const selectedModelId =
+  if (!modelIdFromCookie) {
-    models.find((model) => model.id === modelIdFromCookie)?.id ||
+    return (
-    DEFAULT_MODEL_NAME;
+      <>
+        <Chat
+          key={id}
+          id={id}
+          initialMessages={[]}
+          selectedChatModel={DEFAULT_CHAT_MODEL}
+          selectedVisibilityType="private"
+          isReadonly={false}
+        />
+        <DataStreamHandler id={id} />
+      </>
+    );
+  }
 
   return (
     <>
@@ -21,7 +33,7 @@ export default async function Page() {
         key={id}
         id={id}
         initialMessages={[]}
-        selectedModelId={selectedModelId}
+        selectedChatModel={modelIdFromCookie.value}
         selectedVisibilityType="private"
         isReadonly={false}
       />

components/chat.tsx

@@ -14,17 +14,18 @@ import { MultimodalInput } from './multimodal-input';
 import { Messages } from './messages';
 import { VisibilityType } from './visibility-selector';
 import { useBlockSelector } from '@/hooks/use-block';
+import { toast } from 'sonner';
 
 export function Chat({
   id,
   initialMessages,
-  selectedModelId,
+  selectedChatModel,
   selectedVisibilityType,
   isReadonly,
 }: {
   id: string;
   initialMessages: Array<Message>;
-  selectedModelId: string;
+  selectedChatModel: string;
   selectedVisibilityType: VisibilityType;
   isReadonly: boolean;
 }) {
@@ -42,7 +43,7 @@ export function Chat({
     reload,
   } = useChat({
     id,
-    body: { id, modelId: selectedModelId },
+    body: { id, selectedChatModel: selectedChatModel },
     initialMessages,
     experimental_throttle: 100,
     sendExtraMessageFields: true,
@@ -50,6 +51,10 @@ export function Chat({
     onFinish: () => {
       mutate('/api/history');
     },
+    onError: (error) => {
+      console.log(error);
+      toast.error('An error occured, please try again!');
+    },
   });
 
   const { data: votes } = useSWR<Array<Vote>>(
@@ -65,7 +70,7 @@ export function Chat({
       <div className="flex flex-col min-w-0 h-dvh bg-background">
         <ChatHeader
           chatId={id}
-          selectedModelId={selectedModelId}
+          selectedModelId={selectedChatModel}
           selectedVisibilityType={selectedVisibilityType}
           isReadonly={isReadonly}
         />

components/markdown.tsx

@@ -1,5 +1,5 @@
 import Link from 'next/link';
-import React, { memo, useMemo, useState } from 'react';
+import React, { memo } from 'react';
 import ReactMarkdown, { type Components } from 'react-markdown';
 import remarkGfm from 'remark-gfm';
 import { CodeBlock } from './code-block';

components/message-reasoning.tsx

@@ -0,0 +1,75 @@
+'use client';
+
+import { useState } from 'react';
+import { ChevronDownIcon, LoaderIcon } from './icons';
+import { motion, AnimatePresence } from 'framer-motion';
+import { Markdown } from './markdown';
+
+interface MessageReasoningProps {
+  isLoading: boolean;
+  reasoning: string;
+}
+
+export function MessageReasoning({
+  isLoading,
+  reasoning,
+}: MessageReasoningProps) {
+  const [isExpanded, setIsExpanded] = useState(true);
+
+  const variants = {
+    collapsed: {
+      height: 0,
+      opacity: 0,
+      marginTop: 0,
+      marginBottom: 0,
+    },
+    expanded: {
+      height: 'auto',
+      opacity: 1,
+      marginTop: '1rem',
+      marginBottom: '0.5rem',
+    },
+  };
+
+  return (
+    <div className="flex flex-col">
+      {isLoading ? (
+        <div className="flex flex-row gap-2 items-center">
+          <div className="font-medium">Reasoning</div>
+          <div className="animate-spin">
+            <LoaderIcon />
+          </div>
+        </div>
+      ) : (
+        <div className="flex flex-row gap-2 items-center">
+          <div className="font-medium">Reasoned for a few seconds</div>
+          <div
+            className="cursor-pointer"
+            onClick={() => {
+              setIsExpanded(!isExpanded);
+            }}
+          >
+            <ChevronDownIcon />
+          </div>
+        </div>
+      )}
+
+      <AnimatePresence initial={false}>
+        {isExpanded && (
+          <motion.div
+            key="content"
+            initial="collapsed"
+            animate="expanded"
+            exit="collapsed"
+            variants={variants}
+            transition={{ duration: 0.2, ease: 'easeInOut' }}
+            style={{ overflow: 'hidden' }}
+            className="pl-4 text-zinc-600 dark:text-zinc-400 border-l flex flex-col gap-4"
+          >
+            <Markdown>{reasoning}</Markdown>
+          </motion.div>
+        )}
+      </AnimatePresence>
+    </div>
+  );
+}

components/message.tsx

@@ -8,7 +8,12 @@ import { memo, useMemo, useState } from 'react';
 import type { Vote } from '@/lib/db/schema';
 
 import { DocumentToolCall, DocumentToolResult } from './document';
-import { PencilEditIcon, SparklesIcon } from './icons';
+import {
+  ChevronDownIcon,
+  LoaderIcon,
+  PencilEditIcon,
+  SparklesIcon,
+} from './icons';
 import { Markdown } from './markdown';
 import { MessageActions } from './message-actions';
 import { PreviewAttachment } from './preview-attachment';
@@ -19,6 +24,7 @@ import { Button } from './ui/button';
 import { Tooltip, TooltipContent, TooltipTrigger } from './ui/tooltip';
 import { MessageEditor } from './message-editor';
 import { DocumentPreview } from './document-preview';
+import { MessageReasoning } from './message-reasoning';
 
 const PurePreviewMessage = ({
   chatId,
@@ -68,7 +74,7 @@ const PurePreviewMessage = ({
             </div>
           )}
 
-          <div className="flex flex-col gap-2 w-full">
+          <div className="flex flex-col gap-4 w-full">
             {message.experimental_attachments && (
               <div className="flex flex-row justify-end gap-2">
                 {message.experimental_attachments.map((attachment) => (
@@ -80,7 +86,14 @@ const PurePreviewMessage = ({
               </div>
             )}
 
-            {message.content && mode === 'view' && (
+            {message.reasoning && (
+              <MessageReasoning
+                isLoading={isLoading}
+                reasoning={message.reasoning}
+              />
+            )}
+
+            {(message.content || message.reasoning) && mode === 'view' && (
               <div className="flex flex-row gap-2 items-start">
                 {message.role === 'user' && !isReadonly && (
                   <Tooltip>
@@ -209,6 +222,8 @@ export const PreviewMessage = memo(
   PurePreviewMessage,
   (prevProps, nextProps) => {
     if (prevProps.isLoading !== nextProps.isLoading) return false;
+    if (prevProps.message.reasoning !== nextProps.message.reasoning)
+      return false;
     if (prevProps.message.content !== nextProps.message.content) return false;
     if (
       !equal(

components/model-selector.tsx

@@ -2,7 +2,7 @@
 
 import { startTransition, useMemo, useOptimistic, useState } from 'react';
 
-import { saveModelId } from '@/app/(chat)/actions';
+import { saveChatModelAsCookie } from '@/app/(chat)/actions';
 import { Button } from '@/components/ui/button';
 import {
   DropdownMenu,
@@ -10,7 +10,7 @@ import {
   DropdownMenuItem,
   DropdownMenuTrigger,
 } from '@/components/ui/dropdown-menu';
-import { models } from '@/lib/ai/models';
+import { chatModels } from '@/lib/ai/models';
 import { cn } from '@/lib/utils';
 
 import { CheckCircleFillIcon, ChevronDownIcon } from './icons';
@@ -25,8 +25,8 @@ export function ModelSelector({
   const [optimisticModelId, setOptimisticModelId] =
     useOptimistic(selectedModelId);
 
-  const selectedModel = useMemo(
+  const selectedChatModel = useMemo(
-    () => models.find((model) => model.id === optimisticModelId),
+    () => chatModels.find((chatModel) => chatModel.id === optimisticModelId),
     [optimisticModelId],
   );
 
@@ -40,38 +40,41 @@ export function ModelSelector({
         )}
       >
         <Button variant="outline" className="md:px-2 md:h-[34px]">
-          {selectedModel?.label}
+          {selectedChatModel?.name}
           <ChevronDownIcon />
         </Button>
       </DropdownMenuTrigger>
       <DropdownMenuContent align="start" className="min-w-[300px]">
-        {models.map((model) => (
+        {chatModels.map((chatModel) => {
-          <DropdownMenuItem
+          const { id } = chatModel;
-            key={model.id}
-            onSelect={() => {
-              setOpen(false);
 
-              startTransition(() => {
+          return (
-                setOptimisticModelId(model.id);
+            <DropdownMenuItem
-                saveModelId(model.id);
+              key={id}
-              });
+              onSelect={() => {
-            }}
+                setOpen(false);
-            className="gap-4 group/item flex flex-row justify-between items-center"
+
-            data-active={model.id === optimisticModelId}
+                startTransition(() => {
-          >
+                  setOptimisticModelId(id);
-            <div className="flex flex-col gap-1 items-start">
+                  saveChatModelAsCookie(id);
-              {model.label}
+                });
-              {model.description && (
+              }}
+              className="gap-4 group/item flex flex-row justify-between items-center"
+              data-active={id === optimisticModelId}
+            >
+              <div className="flex flex-col gap-1 items-start">
+                <div>{chatModel.name}</div>
                 <div className="text-xs text-muted-foreground">
-                  {model.description}
+                  {chatModel.description}
                 </div>
-              )}
+              </div>
-            </div>
+
-            <div className="text-foreground dark:text-foreground opacity-0 group-data-[active=true]/item:opacity-100">
+              <div className="text-foreground dark:text-foreground opacity-0 group-data-[active=true]/item:opacity-100">
-              <CheckCircleFillIcon />
+                <CheckCircleFillIcon />
-            </div>
+              </div>
-          </DropdownMenuItem>
+            </DropdownMenuItem>
-        ))}
+          );
+        })}
       </DropdownMenuContent>
     </DropdownMenu>
   );

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;
+  name: string;
-  apiIdentifier: string;
   description: string;
 }
 
-export const models: Array<Model> = [
+export const chatModels: Array<ChatModel> = [
   {
-    id: 'gpt-4o-mini',
+    id: 'chat-model-small',
-    label: 'GPT 4o mini',
+    name: 'Small model',
-    apiIdentifier: 'gpt-4o-mini',
     description: 'Small model for fast, lightweight tasks',
   },
   {
-    id: 'gpt-4o',
+    id: 'chat-model-large',
-    label: 'GPT 4o',
+    name: 'Large model',
-    apiIdentifier: 'gpt-4o',
+    description: 'Large model for complex, multi-step tasks',
-    description: 'For complex, multi-step tasks',
   },
-] as const;
+  {
-
+    id: 'chat-model-reasoning',
-export const DEFAULT_MODEL_NAME: string = 'gpt-4o-mini';
+    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 = ({
+export const createDocument = ({ session, dataStream }: CreateDocumentProps) =>
-  model,
-  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 = ({
+export const updateDocument = ({ session, dataStream }: UpdateDocumentProps) =>
-  model,
-  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(
+export function sanitizeResponseMessages({
-  messages: Array<ResponseMessage>,
+  messages,
-): Array<ResponseMessage> {
+  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,

package.json

@@ -18,7 +18,8 @@
     "db:up": "drizzle-kit up"
   },
   "dependencies": {
-    "@ai-sdk/openai": "1.0.19",
+    "@ai-sdk/fireworks": "^0.1.8",
+    "@ai-sdk/openai": "1.1.9",
     "@codemirror/lang-javascript": "^6.2.2",
     "@codemirror/lang-python": "^6.1.6",
     "@codemirror/state": "^6.5.0",
@@ -37,7 +38,7 @@
     "@vercel/analytics": "^1.3.1",
     "@vercel/blob": "^0.24.1",
     "@vercel/postgres": "^0.10.0",
-    "ai": "4.1.1",
+    "ai": "4.1.17",
     "bcrypt-ts": "^5.0.2",
     "class-variance-authority": "^0.7.0",
     "classnames": "^2.5.1",

pnpm-lock.yaml

@@ -8,9 +8,12 @@ importers:
 
   .:
     dependencies:
+      '@ai-sdk/fireworks':
+        specifier: ^0.1.8
+        version: 0.1.8(zod@3.23.8)
       '@ai-sdk/openai':
-        specifier: 1.0.19
+        specifier: 1.1.9
-        version: 1.0.19(zod@3.23.8)
+        version: 1.1.9(zod@3.23.8)
       '@codemirror/lang-javascript':
         specifier: ^6.2.2
         version: 6.2.2
@@ -66,8 +69,8 @@ importers:
         specifier: ^0.10.0
         version: 0.10.0
       ai:
-        specifier: 4.1.1
+        specifier: 4.1.17
-        version: 4.1.1(react@19.0.0-rc-45804af1-20241021)(zod@3.23.8)
+        version: 4.1.17(react@19.0.0-rc-45804af1-20241021)(zod@3.23.8)
       bcrypt-ts:
         specifier: ^5.0.2
         version: 5.0.2
@@ -249,14 +252,26 @@ importers:
 
 packages:
 
-  '@ai-sdk/openai@1.0.19':
+  '@ai-sdk/fireworks@0.1.8':
-    resolution: {integrity: sha512-7qmLgppWpGUhSgrH0a6CtgD9hZeRh2hARppl1B7fNhVbekYftSMucsdCiVlKbQzSKPxox0vkNMmwjKa/7xf8bQ==}
+    resolution: {integrity: sha512-3/xrROoWl3K6YA9VMej4p1M2eMS54xz7Z0x7QC7+q9ULLNhiZOpkbHVxDjsBayOlwiXH1MyZlKyDd4jJsgDzgg==}
     engines: {node: '>=18'}
     peerDependencies:
       zod: ^3.0.0
 
-  '@ai-sdk/provider-utils@2.0.7':
+  '@ai-sdk/openai-compatible@0.1.8':
-    resolution: {integrity: sha512-4sfPlKEALHPXLmMFcPlYksst3sWBJXmCDZpIBJisRrmwGG6Nn3mq0N1Zu/nZaGcrWZoOY+HT2Wbxla1oTElYHQ==}
+    resolution: {integrity: sha512-o2WeZmkOgaaEHAIZfAdlAASotNemhWBzupUp7ql/vBKIrPDctbQS4K7XOvG7EZ1dshn0TxB+Ur7Q8HoWSeWPmA==}
+    engines: {node: '>=18'}
+    peerDependencies:
+      zod: ^3.0.0
+
+  '@ai-sdk/openai@1.1.9':
+    resolution: {integrity: sha512-t/CpC4TLipdbgBJTMX/otzzqzCMBSPQwUOkYPGbT/jyuC86F+YO9o+LS0Ty2pGUE1kyT+B3WmJ318B16ZCg4hw==}
+    engines: {node: '>=18'}
+    peerDependencies:
+      zod: ^3.0.0
+
+  '@ai-sdk/provider-utils@2.1.6':
+    resolution: {integrity: sha512-Pfyaj0QZS22qyVn5Iz7IXcJ8nKIKlu2MeSAdKJzTwkAks7zdLaKVB+396Rqcp1bfQnxl7vaduQVMQiXUrgK8Gw==}
     engines: {node: '>=18'}
     peerDependencies:
       zod: ^3.0.0
@@ -264,25 +279,12 @@ packages:
       zod:
         optional: true
 
-  '@ai-sdk/provider-utils@2.1.1':
+  '@ai-sdk/provider@1.0.7':
-    resolution: {integrity: sha512-+FRXSAdzPJFJN6TpyvyGWLo7WJuoBKI1g66UL+sli1HrxlldXSwxRPeb8tMMmNcyi3VKQogg2VsoJjlt4ort5w==}
+    resolution: {integrity: sha512-q1PJEZ0qD9rVR+8JFEd01/QM++csMT5UVwYXSN2u54BrVw/D8TZLTeg2FEfKK00DgAx0UtWd8XOhhwITP9BT5g==}
-    engines: {node: '>=18'}
-    peerDependencies:
-      zod: ^3.0.0
-    peerDependenciesMeta:
-      zod:
-        optional: true
-
-  '@ai-sdk/provider@1.0.4':
-    resolution: {integrity: sha512-lJi5zwDosvvZER3e/pB8lj1MN3o3S7zJliQq56BRr4e9V3fcRyFtwP0JRxaRS5vHYX3OJ154VezVoQNrk0eaKw==}
     engines: {node: '>=18'}
 
-  '@ai-sdk/provider@1.0.5':
+  '@ai-sdk/react@1.1.8':
-    resolution: {integrity: sha512-KATFp9CNXtMEzs8KBwLYK2+rGkkeED6p1+4koQveszyscIavObXIRW7vjr0MoZ9HFIHOUlrcak+3s/Xt3UXmAg==}
+    resolution: {integrity: sha512-buHm7hP21xEOksnRQtJX9fKbi7cAUwanEBa5niddTDibCDKd+kIXP2vaJGy8+heB3rff+XSW3BWlA8pscK+n1g==}
-    engines: {node: '>=18'}
-
-  '@ai-sdk/react@1.1.1':
-    resolution: {integrity: sha512-7LX/YF8sis8UM7p8ftUcu0xySG86/TBddcB42w/+mWOXL6hjYzcuGD8G121TobHsnxVxbsBlF/ykps/GYVvLNg==}
     engines: {node: '>=18'}
     peerDependencies:
       react: ^18 || ^19 || ^19.0.0-rc
@@ -293,8 +295,8 @@ packages:
       zod:
         optional: true
 
-  '@ai-sdk/ui-utils@1.1.1':
+  '@ai-sdk/ui-utils@1.1.8':
-    resolution: {integrity: sha512-lkTxGoebnEgs8HtKeWut0AglXN7zpWQwYmun4yuhpiup7DxPWTmt3vGiYvqQTBOFAmyoea3uzIKjHwRHuayr2w==}
+    resolution: {integrity: sha512-nbok53K1EalO2sZjBLFB33cqs+8SxiL6pe7ekZ7+5f2MJTwdvpShl6d9U4O8fO3DnZ9pYLzaVC0XNMxnJt030Q==}
     engines: {node: '>=18'}
     peerDependencies:
       zod: ^3.0.0
@@ -1638,8 +1640,8 @@ packages:
     engines: {node: '>=0.4.0'}
     hasBin: true
 
-  ai@4.1.1:
+  ai@4.1.17:
-    resolution: {integrity: sha512-oZTzQfrvrXuXnAJhoCsGcLUxSMWWYKkqrk2LfCzcukvB8us7ZUnFBgs9drhXuFP3JkWVbeIZHXjDexZIZcbi8g==}
+    resolution: {integrity: sha512-5SW15tXDuxE/wlEOjRKxLxTOUIGD4C9bIee+FCFvXHTTAZhHiQjViC2s7RtMUW+hbFtGya302jUHY1Pe8A/YuQ==}
     engines: {node: '>=18'}
     peerDependencies:
       react: ^18 || ^19 || ^19.0.0-rc
@@ -3772,52 +3774,52 @@ packages:
 
 snapshots:
 
-  '@ai-sdk/openai@1.0.19(zod@3.23.8)':
+  '@ai-sdk/fireworks@0.1.8(zod@3.23.8)':
     dependencies:
-      '@ai-sdk/provider': 1.0.4
+      '@ai-sdk/openai-compatible': 0.1.8(zod@3.23.8)
-      '@ai-sdk/provider-utils': 2.0.7(zod@3.23.8)
+      '@ai-sdk/provider': 1.0.7
+      '@ai-sdk/provider-utils': 2.1.6(zod@3.23.8)
       zod: 3.23.8
 
-  '@ai-sdk/provider-utils@2.0.7(zod@3.23.8)':
+  '@ai-sdk/openai-compatible@0.1.8(zod@3.23.8)':
     dependencies:
-      '@ai-sdk/provider': 1.0.4
+      '@ai-sdk/provider': 1.0.7
+      '@ai-sdk/provider-utils': 2.1.6(zod@3.23.8)
+      zod: 3.23.8
+
+  '@ai-sdk/openai@1.1.9(zod@3.23.8)':
+    dependencies:
+      '@ai-sdk/provider': 1.0.7
+      '@ai-sdk/provider-utils': 2.1.6(zod@3.23.8)
+      zod: 3.23.8
+
+  '@ai-sdk/provider-utils@2.1.6(zod@3.23.8)':
+    dependencies:
+      '@ai-sdk/provider': 1.0.7
       eventsource-parser: 3.0.0
       nanoid: 3.3.8
       secure-json-parse: 2.7.0
     optionalDependencies:
       zod: 3.23.8
 
-  '@ai-sdk/provider-utils@2.1.1(zod@3.23.8)':
+  '@ai-sdk/provider@1.0.7':
-    dependencies:
-      '@ai-sdk/provider': 1.0.5
-      eventsource-parser: 3.0.0
-      nanoid: 3.3.8
-      secure-json-parse: 2.7.0
-    optionalDependencies:
-      zod: 3.23.8
-
-  '@ai-sdk/provider@1.0.4':
     dependencies:
       json-schema: 0.4.0
 
-  '@ai-sdk/provider@1.0.5':
+  '@ai-sdk/react@1.1.8(react@19.0.0-rc-45804af1-20241021)(zod@3.23.8)':
     dependencies:
-      json-schema: 0.4.0
+      '@ai-sdk/provider-utils': 2.1.6(zod@3.23.8)
-
+      '@ai-sdk/ui-utils': 1.1.8(zod@3.23.8)
-  '@ai-sdk/react@1.1.1(react@19.0.0-rc-45804af1-20241021)(zod@3.23.8)':
-    dependencies:
-      '@ai-sdk/provider-utils': 2.1.1(zod@3.23.8)
-      '@ai-sdk/ui-utils': 1.1.1(zod@3.23.8)
       swr: 2.2.5(react@19.0.0-rc-45804af1-20241021)
       throttleit: 2.1.0
     optionalDependencies:
       react: 19.0.0-rc-45804af1-20241021
       zod: 3.23.8
 
-  '@ai-sdk/ui-utils@1.1.1(zod@3.23.8)':
+  '@ai-sdk/ui-utils@1.1.8(zod@3.23.8)':
     dependencies:
-      '@ai-sdk/provider': 1.0.5
+      '@ai-sdk/provider': 1.0.7
-      '@ai-sdk/provider-utils': 2.1.1(zod@3.23.8)
+      '@ai-sdk/provider-utils': 2.1.6(zod@3.23.8)
       zod-to-json-schema: 3.24.1(zod@3.23.8)
     optionalDependencies:
       zod: 3.23.8
@@ -4921,12 +4923,12 @@ snapshots:
 
   acorn@8.14.0: {}
 
-  ai@4.1.1(react@19.0.0-rc-45804af1-20241021)(zod@3.23.8):
+  ai@4.1.17(react@19.0.0-rc-45804af1-20241021)(zod@3.23.8):
     dependencies:
-      '@ai-sdk/provider': 1.0.5
+      '@ai-sdk/provider': 1.0.7
-      '@ai-sdk/provider-utils': 2.1.1(zod@3.23.8)
+      '@ai-sdk/provider-utils': 2.1.6(zod@3.23.8)
-      '@ai-sdk/react': 1.1.1(react@19.0.0-rc-45804af1-20241021)(zod@3.23.8)
+      '@ai-sdk/react': 1.1.8(react@19.0.0-rc-45804af1-20241021)(zod@3.23.8)
-      '@ai-sdk/ui-utils': 1.1.1(zod@3.23.8)
+      '@ai-sdk/ui-utils': 1.1.8(zod@3.23.8)
       '@opentelemetry/api': 1.9.0
       jsondiffpatch: 0.6.0
     optionalDependencies: