Thinking is now working as expected

Author: sh1ftred

components/chat/ChatMessages.tsx

@@ -8,6 +8,7 @@ import { RefObject, useState } from 'react';
 interface ChatMessagesProps {
   messages: Message[];
   streamingContent: string;
+  thinkingContent: string;
   editingMessageIndex: number | null;
   editingContent: string;
   setEditingContent: (content: string) => void;
@@ -22,6 +23,7 @@ interface ChatMessagesProps {
 export default function ChatMessages({
   messages,
   streamingContent,
+  thinkingContent,
   editingMessageIndex,
   editingContent,
   setEditingContent,
@@ -148,7 +150,7 @@ export default function ChatMessages({
                 </div>
               ) : (
                 <div className="flex flex-col items-start mb-6 group">
-                  {message.thinking && (
+                  {(message.thinking) && (
                     <ThinkingSection thinking={message.thinking} />
                   )}
                   <div className="max-w-[95%] text-gray-100 py-2 px-0.5">
@@ -201,6 +203,10 @@ export default function ChatMessages({
           ))
         )}
 
+        {thinkingContent && (
+          <ThinkingSection thinkingContent={thinkingContent} isStreaming={streamingContent==''}/>
+        )}
+
         {streamingContent && (
           <div className="flex flex-col items-start mb-6">
             <div className="max-w-[95%] text-gray-100 py-2 px-0.5">

components/chat/MainChatArea.tsx

@@ -19,6 +19,7 @@ const MainChatArea: React.FC = () => {
     messages,
     setMessages,
     streamingContent,
+    thinkingContent,
     editingMessageIndex,
     editingContent,
     setEditingContent,
@@ -101,6 +102,7 @@ const MainChatArea: React.FC = () => {
       <ChatMessages
         messages={messages}
         streamingContent={streamingContent}
+        thinkingContent={thinkingContent}
         editingMessageIndex={editingMessageIndex}
         editingContent={editingContent}
         setEditingContent={setEditingContent}

components/ui/ThinkingSection.tsx

@@ -1,20 +1,41 @@
 'use client';
 
-import { useState } from 'react';
+import { useState, useEffect } from 'react';
 import { motion, AnimatePresence } from 'motion/react';
 import { ChevronDown, Brain, Copy, Check } from 'lucide-react';
 import MarkdownRenderer from '@/components/MarkdownRenderer';
 
 interface ThinkingSectionProps {
   thinking?: string;
+  thinkingContent?: string;
   isStreaming?: boolean;
 }
 
-export default function ThinkingSection({ thinking, isStreaming = false }: ThinkingSectionProps) {
+export default function ThinkingSection({ thinking, thinkingContent, isStreaming = false }: ThinkingSectionProps) {
   const [isExpanded, setIsExpanded] = useState(false);
   const [copied, setCopied] = useState(false);
+  const [thinkingText, setThinkingText] = useState('');
 
-  if (!thinking && !isStreaming) return null;
+  useEffect(() => {
+    if (thinkingContent !== '') {
+      setThinkingText(thinkingContent!);
+      setIsExpanded(true);
+    } else {
+      setThinkingText(thinking || '');
+      setIsExpanded(false);
+    }
+  }, [thinking, thinkingContent]);
+  
+  useEffect(() => {
+    if (isStreaming){
+      setIsExpanded(true)
+    }
+    else{
+      setIsExpanded(false)
+    }
+  }, [isStreaming]);
+
+  if (!thinking && !isStreaming && !thinkingContent) return null;
 
   const handleCopy = async () => {
     if (!thinking) return;
@@ -64,8 +85,8 @@ export default function ThinkingSection({ thinking, isStreaming = false }: Think
                 </button>
               )}
               <div className="text-xs text-gray-300 font-mono leading-relaxed">
-                {thinking ? (
-                  <MarkdownRenderer content={thinking} />
+                {thinkingText ? (
+                  <MarkdownRenderer content={thinkingText} />
                 ) : (
                   <div className="flex items-center gap-2">
                     <div className="w-1 h-1 bg-gray-400 rounded-full animate-pulse" />

hooks/useChatActions.ts

@@ -17,6 +17,7 @@ export interface UseChatActionsReturn {
   inputMessage: string;
   isLoading: boolean;
   streamingContent: string;
+  thinkingContent: string;
   balance: number;
   isBalanceLoading: boolean;
   uploadedImages: string[];
@@ -72,6 +73,7 @@ export const useChatActions = (): UseChatActionsReturn => {
   const [inputMessage, setInputMessage] = useState('');
   const [isLoading, setIsLoading] = useState(false);
   const [streamingContent, setStreamingContent] = useState('');
+  const [thinkingContent, setThinkingContent] = useState('');
   const [balance, setBalance] = useState(0);
   const [isBalanceLoading, setIsBalanceLoading] = useState(true);
   const [uploadedImages, setUploadedImages] = useState<string[]>([]);
@@ -180,7 +182,7 @@ export const useChatActions = (): UseChatActionsReturn => {
     if (messagesEndRef.current) {
       messagesEndRef.current.scrollIntoView({ behavior: 'smooth' });
     }
-  }, [streamingContent]);
+  }, [streamingContent, thinkingContent]);
 
   const setTransactionHistory = useCallback((value: React.SetStateAction<TransactionHistory[]>) => {
     setTransactionHistoryState(prev => {
@@ -285,6 +287,7 @@ export const useChatActions = (): UseChatActionsReturn => {
   ) => {
     setIsLoading(true);
     setStreamingContent('');
+    setThinkingContent('');
 
     // Create a ref to track current messages during the API call
     let currentMessages = messageHistory;
@@ -305,6 +308,7 @@ export const useChatActions = (): UseChatActionsReturn => {
         receiveToken,
         activeMintUrl: cashuStore.activeMintUrl,
         onStreamingUpdate: setStreamingContent,
+        onThinkingUpdate: setThinkingContent,
         onMessagesUpdate: updateMessages,
         onMessageAppend: (message) => {
           // Append to current messages state
@@ -326,13 +330,15 @@ export const useChatActions = (): UseChatActionsReturn => {
     } finally {
       setIsLoading(false);
       setStreamingContent('');
+      setThinkingContent('');
     }
   }, [usingNip60, balance, sendToken, receiveToken, cashuStore.activeMintUrl, transactionHistory, setPendingCashuAmountState]);
 
   return {
     inputMessage,
     isLoading,
     streamingContent,
+    thinkingContent,
     balance,
     isBalanceLoading,
     uploadedImages,

utils/apiUtils.ts

@@ -16,6 +16,7 @@ export interface FetchAIResponseParams {
   receiveToken: (token: string) => Promise<any[]>;
   activeMintUrl?: string | null;
   onStreamingUpdate: (content: string) => void;
+  onThinkingUpdate: (content: string) => void;
   onMessagesUpdate: (messages: Message[]) => void;
   onMessageAppend: (message: Message) => void;
   onBalanceUpdate: (balance: number) => void;
@@ -42,6 +43,7 @@ export const fetchAIResponse = async (params: FetchAIResponseParams): Promise<vo
     receiveToken,
     activeMintUrl,
     onStreamingUpdate,
+    onThinkingUpdate,
     onMessagesUpdate,
     onMessageAppend,
     onBalanceUpdate,
@@ -126,7 +128,7 @@ export const fetchAIResponse = async (params: FetchAIResponseParams): Promise<vo
       throw new Error('Response body is not available');
     }
 
-    const streamingResult = await processStreamingResponse(response, onStreamingUpdate, selectedModel?.id);
+    const streamingResult = await processStreamingResponse(response, onStreamingUpdate, onThinkingUpdate, selectedModel?.id);
 
     if (streamingResult.content) {
       const assistantMessage = createTextMessage('assistant', streamingResult.content);
@@ -146,6 +148,7 @@ export const fetchAIResponse = async (params: FetchAIResponseParams): Promise<vo
     }
 
     onStreamingUpdate('');
+    onThinkingUpdate('');
 
     // Handle refund and balance update
     await handlePostResponseRefund({
@@ -290,13 +293,15 @@ interface StreamingResult {
 async function processStreamingResponse(
   response: Response,
   onStreamingUpdate: (content: string) => void,
+  onThinkingUpdate: (content: string) => void,
   modelId?: string
 ): Promise<StreamingResult> {
   const reader = response.body!.getReader();
   const decoder = new TextDecoder('utf-8');
   let accumulatedContent = '';
   let accumulatedThinking = '';
   let isInThinking = false;
+  let isInContent = false;
   let usage: StreamingResult['usage'];
   let model: string | undefined;
   let finish_reason: string | undefined;
@@ -324,18 +329,52 @@ async function processStreamingResponse(
           try {
             const parsedData = JSON.parse(jsonData);
 
-            // Handle content delta
+            // Handle reasoning delta. OpenRouter does this. 
             if (parsedData.choices &&
+              parsedData.choices[0] &&
+              parsedData.choices[0].delta &&
+              parsedData.choices[0].delta.reasoning) {
+              
+                let newContent;
+                if (!isInThinking) {
+                  newContent = "<thinking> " + parsedData.choices[0].delta.reasoning;
+                  isInThinking = true;
+                }
+                else {
+                  newContent = parsedData.choices[0].delta.reasoning;
+                }
+                const thinkingResult = extractThinkingFromStream(newContent, accumulatedThinking);
+                accumulatedThinking = thinkingResult.thinking;
+                onThinkingUpdate(accumulatedThinking)
+              }
+
+            // Handle content delta
+            else if (parsedData.choices &&
               parsedData.choices[0] &&
               parsedData.choices[0].delta &&
               parsedData.choices[0].delta.content) {
 
+              if (isInThinking && !isInContent) {
+                const newContent = "</thinking>";
+                const thinkingResult = extractThinkingFromStream(newContent, accumulatedThinking);
+                accumulatedThinking = thinkingResult.thinking;
+                onThinkingUpdate(accumulatedThinking);
+                
+                if (thinkingResult.content) {
+                  accumulatedContent += thinkingResult.content;
+                  onStreamingUpdate(accumulatedContent);
+                }
+                isInThinking = false;
+                isInContent = true;
+              }
+
               const newContent = parsedData.choices[0].delta.content;
 
               if (modelId && isThinkingCapableModel(modelId)) {
                 const thinkingResult = extractThinkingFromStream(newContent, accumulatedThinking);
                 accumulatedThinking = thinkingResult.thinking;
                 isInThinking = thinkingResult.isInThinking;
+                onThinkingUpdate(accumulatedThinking);
 
                 if (thinkingResult.content) {
                   accumulatedContent += thinkingResult.content;
@@ -379,7 +418,7 @@ async function processStreamingResponse(
 
   return {
     content: accumulatedContent,
-    thinking: (modelId && isThinkingCapableModel(modelId) && accumulatedThinking) ? accumulatedThinking : undefined,
+    thinking: (modelId && accumulatedThinking) ? accumulatedThinking : undefined,
     usage,
     model,
     finish_reason