feat: add reasoning model (#750)

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

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(
-    () => models.find((model) => model.id === optimisticModelId),
+  const selectedChatModel = useMemo(
+    () => 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) => (
-          <DropdownMenuItem
-            key={model.id}
-            onSelect={() => {
-              setOpen(false);
+        {chatModels.map((chatModel) => {
+          const { id } = chatModel;
 
-              startTransition(() => {
-                setOptimisticModelId(model.id);
-                saveModelId(model.id);
-              });
-            }}
-            className="gap-4 group/item flex flex-row justify-between items-center"
-            data-active={model.id === optimisticModelId}
-          >
-            <div className="flex flex-col gap-1 items-start">
-              {model.label}
-              {model.description && (
+          return (
+            <DropdownMenuItem
+              key={id}
+              onSelect={() => {
+                setOpen(false);
+
+                startTransition(() => {
+                  setOptimisticModelId(id);
+                  saveChatModelAsCookie(id);
+                });
+              }}
+              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 className="text-foreground dark:text-foreground opacity-0 group-data-[active=true]/item:opacity-100">
-              <CheckCircleFillIcon />
-            </div>
-          </DropdownMenuItem>
-        ))}
+              </div>
+
+              <div className="text-foreground dark:text-foreground opacity-0 group-data-[active=true]/item:opacity-100">
+                <CheckCircleFillIcon />
+              </div>
+            </DropdownMenuItem>
+          );
+        })}
       </DropdownMenuContent>
     </DropdownMenu>
   );