feat(stage-ui): process for reasoning models

packages/stage-ui/package.json

@@ -120,6 +120,7 @@
     "posthog-js": "catalog:",
     "postprocessing": "^6.38.2",
     "rehype-katex": "^7.0.1",
+    "rehype-parse": "^9.0.1",
     "rehype-stringify": "^10.0.1",
     "reka-ui": "^2.7.0",
     "remark-math": "^6.0.0",
@@ -129,6 +130,7 @@
     "three": "^0.182.0",
     "unified": "^11.0.5",
     "unist-builder": "^4.0.0",
+    "unist-util-visit": "^5.0.0",
     "unspeech": "^0.1.7",
     "uuid": "^13.0.0",
     "vaul-vue": "^0.4.1",
@@ -164,7 +166,9 @@
     "@proj-airi/vite-plugin-warpdrive": "workspace:*",
     "@types/audioworklet": "catalog:",
     "@types/culori": "^4.0.1",
+    "@types/hast": "catalog:",
     "@types/three": "^0.182.0",
+    "@types/unist": "catalog:",
     "@unocss/reset": "^66.5.11",
     "@vitejs/plugin-vue": "^6.0.3",
     "clustr": "^1.0.2",

packages/stage-ui/src/components/scenarios/chat/ChatAssistantItem.vue

@@ -4,6 +4,7 @@ import type { ChatAssistantMessage, ChatSlices, ChatSlicesText } from '../../../
 import { computed } from 'vue'
 
 import MarkdownRenderer from '../../markdown/MarkdownRenderer.vue'
+import ChatResponseCategorization from './ChatResponseCategorization.vue'
 import ChatToolCallBlock from './ChatToolCallBlock.vue'
 
 const props = withDefaults(defineProps<{
@@ -66,6 +67,12 @@ const boxClasses = computed(() => [
         </template>
       </div>
       <div v-else-if="showLoader" i-eos-icons:three-dots-loading />
+
+      <ChatResponseCategorization
+        v-if="message.categorization"
+        :message="message"
+        :variant="variant"
+      />
     </div>
   </div>
 </template>

packages/stage-ui/src/components/scenarios/chat/ChatResponseCategorization.vue

@@ -0,0 +1,53 @@
+<script setup lang="ts">
+import type { ChatAssistantMessage } from '../../../types/chat'
+
+import { computed } from 'vue'
+
+import MarkdownRenderer from '../../markdown/MarkdownRenderer.vue'
+import Collapsable from '../../misc/Collapsable.vue'
+
+const props = defineProps<{
+  message: ChatAssistantMessage
+  variant?: 'desktop' | 'mobile'
+}>()
+
+const hasReasoning = computed(() => !!props.message.categorization?.reasoning?.trim())
+
+const containerClasses = computed(() => [
+  'mt-2',
+  props.variant === 'mobile' ? 'text-xs' : 'text-sm',
+])
+</script>
+
+<template>
+  <div v-if="hasReasoning" :class="containerClasses" flex="~ col" gap-1>
+    <Collapsable :default="false">
+      <template #trigger="slotProps">
+        <button
+          class="w-full flex items-center justify-between rounded-lg bg-neutral-100/50 px-2 py-1 text-xs text-neutral-600 outline-none transition-all duration-200 dark:bg-neutral-800/50 hover:bg-neutral-200/50 dark:text-neutral-400 dark:hover:bg-neutral-700/50"
+          @click="slotProps.setVisible(!slotProps.visible)"
+        >
+          <div flex="~ items-center" gap-1.5>
+            <div i-solar:lightbulb-bolt-bold-duotone size-3.5 text-amber-500 dark:text-amber-400 />
+            <span font-medium>Reasoning</span>
+          </div>
+          <div
+            i-solar:alt-arrow-down-linear
+            size-3
+            transition="transform duration-200"
+            :class="{ 'rotate-180': slotProps.visible }"
+          />
+        </button>
+      </template>
+      <div
+        class="mt-1 border border-neutral-200 rounded-md bg-neutral-50/80 px-2 py-1.5 dark:border-neutral-700 dark:bg-neutral-900/80"
+      >
+        <MarkdownRenderer
+          :content="message.categorization?.reasoning ?? ''"
+          class="break-words"
+          text="xs neutral-700 dark:neutral-300"
+        />
+      </div>
+    </Collapsable>
+  </div>
+</template>

packages/stage-ui/src/composables/index.ts

@@ -6,4 +6,5 @@ export * from './micvad'
 export * from './queues'
 export * from './use-analytics'
 export * from './use-build-info'
+export * from './use-chat-memory'
 export * from './whisper'

packages/stage-ui/src/composables/llmmarkerParser.test.ts

@@ -142,4 +142,42 @@ describe('useLlmmarkerParser', async () => {
       expect(collectedSpecials).toEqual(expectedSpecials)
     }
   })
+
+  it('should call onEnd with full text', async () => {
+    const fullText = 'Hello, world!'
+    let endText = ''
+
+    const parser = useLlmmarkerParser({
+      onEnd(text) {
+        endText = text
+      },
+    })
+
+    for (const char of fullText) {
+      await parser.consume(char)
+    }
+
+    await parser.end()
+
+    expect(endText).toBe(fullText)
+  })
+
+  it('should call onEnd with full text including specials', async () => {
+    const fullText = 'Hello <|special|> world!'
+    let endText = ''
+
+    const parser = useLlmmarkerParser({
+      onEnd(text) {
+        endText = text
+      },
+    })
+
+    for (const char of fullText) {
+      await parser.consume(char)
+    }
+
+    await parser.end()
+
+    expect(endText).toBe(fullText)
+  })
 })

packages/stage-ui/src/composables/llmmarkerParser.ts

@@ -18,6 +18,11 @@ const TAG_CLOSE = '|>'
 export function useLlmmarkerParser(options: {
   onLiteral?: (literal: string) => void | Promise<void>
   onSpecial?: (special: string) => void | Promise<void>
+  /**
+   * Called when parsing ends with the full accumulated text.
+   * Useful for final processing like categorization or filtering.
+   */
+  onEnd?: (fullText: string) => void | Promise<void>
   /**
    * The minimum length of text required to emit a literal part.
    * Useful for avoiding emitting literal parts too fast.
@@ -27,6 +32,7 @@ export function useLlmmarkerParser(options: {
   const minLiteralEmitLength = Math.max(1, options.minLiteralEmitLength ?? 1)
   let buffer = ''
   let inTag = false
+  let fullText = ''
 
   return {
     /**
@@ -36,6 +42,7 @@ export function useLlmmarkerParser(options: {
      * @param textPart The chunk of text to consume.
      */
     async consume(textPart: string) {
+      fullText += textPart
       buffer += textPart
 
       while (buffer.length > 0) {
@@ -82,6 +89,7 @@ export function useLlmmarkerParser(options: {
         await options.onLiteral?.(buffer)
         buffer = ''
       }
+      await options.onEnd?.(fullText)
     },
   }
 }