From 8472321a812381c911235b58e33728006b458686 Mon Sep 17 00:00:00 2001
From: Ettore Di Giacinto <mudler@users.noreply.github.com>
Date: Sat, 31 May 2025 08:50:46 +0200
Subject: [PATCH] feat(ui): display thinking tags appropriately (#5540)

* fix(streaming): stream complete runes

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* feat(ui): display thinking tags separately

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

---------

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
Signed-off-by: Ettore Di Giacinto <mudler@users.noreply.github.com>
---
 core/http/static/chat.js  | 57 +++++++++++++++++++++++++++++++++++++--
 core/http/views/chat.html | 15 ++++++++++-
 2 files changed, 69 insertions(+), 3 deletions(-)

diff --git a/core/http/static/chat.js b/core/http/static/chat.js
index 8653689c..db5c884e 100644
--- a/core/http/static/chat.js
+++ b/core/http/static/chat.js
@@ -280,6 +280,9 @@ async function promptGPT(systemPrompt, input) {
 
   let buffer = "";
   let contentBuffer = [];
+  let thinkingContent = "";
+  let isThinking = false;
+  let lastThinkingMessageIndex = -1;
 
   try {
     while (true) {
@@ -303,7 +306,44 @@ async function promptGPT(systemPrompt, input) {
             const token = jsonData.choices[0].delta.content;
 
             if (token) {
-              contentBuffer.push(token);
+              // Check for thinking tags
+              if (token.includes("<thinking>") || token.includes("<think>")) {
+                isThinking = true;
+                thinkingContent = "";
+                lastThinkingMessageIndex = -1;
+                return;
+              }
+              if (token.includes("</thinking>") || token.includes("</think>")) {
+                isThinking = false;
+                if (thinkingContent.trim()) {
+                  // Only add the final thinking message if we don't already have one
+                  if (lastThinkingMessageIndex === -1) {
+                    Alpine.store("chat").add("thinking", thinkingContent);
+                  }
+                }
+                return;
+              }
+
+              // Handle content based on thinking state
+              if (isThinking) {
+                thinkingContent += token;
+                // Update the last thinking message or create a new one
+                if (lastThinkingMessageIndex === -1) {
+                  // Create new thinking message
+                  Alpine.store("chat").add("thinking", thinkingContent);
+                  lastThinkingMessageIndex = Alpine.store("chat").history.length - 1;
+                } else {
+                  // Update existing thinking message
+                  const chatStore = Alpine.store("chat");
+                  const lastMessage = chatStore.history[lastThinkingMessageIndex];
+                  if (lastMessage && lastMessage.role === "thinking") {
+                    lastMessage.content = thinkingContent;
+                    lastMessage.html = DOMPurify.sanitize(marked.parse(thinkingContent));
+                  }
+                }
+              } else {
+                contentBuffer.push(token);
+              }
             }
           } catch (error) {
             console.error("Failed to parse line:", line, error);
@@ -322,6 +362,9 @@ async function promptGPT(systemPrompt, input) {
     if (contentBuffer.length > 0) {
       addToChat(contentBuffer.join(""));
     }
+    if (thinkingContent.trim() && lastThinkingMessageIndex === -1) {
+      Alpine.store("chat").add("thinking", thinkingContent);
+    }
 
     // Highlight all code blocks once at the end
     hljs.highlightAll();
@@ -375,7 +418,17 @@ document.addEventListener("alpine:init", () => {
     },
     add(role, content, image, audio) {
       const N = this.history.length - 1;
-      if (this.history.length && this.history[N].role === role) {
+      // For thinking messages, always create a new message
+      if (role === "thinking") {
+        let c = "";
+        const lines = content.split("\n");
+        lines.forEach((line) => {
+          c += DOMPurify.sanitize(marked.parse(line));
+        });
+        this.history.push({ role, content, html: c, image, audio });
+      }
+      // For other messages, merge if same role
+      else if (this.history.length && this.history[N].role === role) {
         this.history[N].content += content;
         this.history[N].html = DOMPurify.sanitize(
           marked.parse(this.history[N].content)
diff --git a/core/http/views/chat.html b/core/http/views/chat.html
index 4198c2aa..3683e1bc 100644
--- a/core/http/views/chat.html
+++ b/core/http/views/chat.html
@@ -257,7 +257,20 @@ SOFTWARE.
                     </div>
                   </div>
                 </template>
-                <template x-if="message.role != 'user'">
+                <template x-if="message.role === 'thinking'">
+                  <div class="flex items-center space-x-2 w-full">
+                    <div class="flex flex-col flex-1">
+                      <div class="p-2 flex-1 rounded bg-blue-900/50 text-blue-100 border border-blue-700/50">
+                        <div class="flex items-center space-x-2">
+                          <i class="fa-solid fa-brain text-blue-400"></i>
+                          <span class="text-xs font-semibold text-blue-300">Thinking</span>
+                        </div>
+                        <div class="mt-1" x-html="message.html"></div>
+                      </div>
+                    </div>
+                  </div>
+                </template>
+                <template x-if="message.role != 'user' && message.role != 'thinking'">
                   <div class="flex items-center space-x-2">
                     {{ if $galleryConfig }}
                     {{ if $galleryConfig.Icon }}<img src="{{$galleryConfig.Icon}}" class="rounded-lg mt-2 max-w-8 max-h-8">{{end}}