feat: support for <think></think> tags to allow reasoning tokens form…

…atted in UI (stackblitz-labs#1205)
Stijnus · Jan 28, 2025 · a199295 · a199295
1 parent 32bfdd9
commit a199295
Showing 1 changed file with 22 additions and 0 deletions.
diff --git a/app/utils/markdown.ts b/app/utils/markdown.ts
@@ -54,8 +54,28 @@ export const allowedHTMLElements = [
   'tr',
   'ul',
   'var',
+  'think',
 ];
 
+// Add custom rehype plugin
+function remarkThinkRawContent() {
+  return (tree: any) => {
+    visit(tree, (node: any) => {
+      if (node.type === 'html' && node.value && node.value.startsWith('<think>')) {
+        const cleanedContent = node.value.slice(7);
+        node.value = `<div class="__boltThought__">${cleanedContent}`;
+
+        return;
+      }
+
+      if (node.type === 'html' && node.value && node.value.startsWith('</think>')) {
+        const cleanedContent = node.value.slice(8);
+        node.value = `</div>${cleanedContent}`;
+      }
+    });
+  };
+}
+
 const rehypeSanitizeOptions: RehypeSanitizeOptions = {
   ...defaultSchema,
   tagNames: allowedHTMLElements,
@@ -79,6 +99,8 @@ export function remarkPlugins(limitedMarkdown: boolean) {
     plugins.unshift(limitedMarkdownPlugin);
   }
 
+  plugins.unshift(remarkThinkRawContent);
+
   return plugins;
 }