From 42026d100b91fd9314bd6d1c4d60727eb9a49a35 Mon Sep 17 00:00:00 2001
From: Makito <i@maki.to>
Date: Mon, 21 Jul 2025 15:43:56 +0900
Subject: [PATCH] feat(stage-ui): better tts input chunking (#306)

* feat(stage-ui): better tts input chunking
* chore: remove unused audio analyser
* fix: words count
* fix: use ctx.data instead

---------

Co-authored-by: Neko <neko@ayaka.moe>
---
 cspell.config.yaml                            |   1 +
 packages/stage-ui/package.json                |   1 +
 .../Scenarios/Providers/SpeechPlayground.vue  |   8 +
 .../Providers/SpeechStreamingPlayground.vue   | 163 ++++++++++++++++++
 packages/stage-ui/src/composables/queues.ts   |  16 +-
 packages/stage-ui/src/utils/tts.ts            | 142 +++++++++++++++
 pnpm-lock.yaml                                |   9 +
 7 files changed, 327 insertions(+), 13 deletions(-)
 create mode 100644 packages/stage-ui/src/components/Scenarios/Providers/SpeechStreamingPlayground.vue
 create mode 100644 packages/stage-ui/src/utils/tts.ts

diff --git a/cspell.config.yaml b/cspell.config.yaml
index 89cf277c5..d837a4025 100644
--- a/cspell.config.yaml
+++ b/cspell.config.yaml
@@ -33,6 +33,7 @@ words:
   - cientos
   - cjkfonts
   - clippy
+  - clustr
   - collectblock
   - colorjs
   - Comfortaa
diff --git a/packages/stage-ui/package.json b/packages/stage-ui/package.json
index f59e5ab76..b77cdd9d0 100644
--- a/packages/stage-ui/package.json
+++ b/packages/stage-ui/package.json
@@ -139,6 +139,7 @@
     "@types/three": "^0.178.1",
     "@unocss/reset": "^66.3.3",
     "@vitejs/plugin-vue": "^6.0.0",
+    "clustr": "^0.0.3",
     "histoire": "1.0.0-alpha.2",
     "unocss": "^66.3.3",
     "unplugin-yaml": "^3.0.1",
diff --git a/packages/stage-ui/src/components/Scenarios/Providers/SpeechPlayground.vue b/packages/stage-ui/src/components/Scenarios/Providers/SpeechPlayground.vue
index eaa8a40dc..3f982bcff 100644
--- a/packages/stage-ui/src/components/Scenarios/Providers/SpeechPlayground.vue
+++ b/packages/stage-ui/src/components/Scenarios/Providers/SpeechPlayground.vue
@@ -5,6 +5,8 @@ import { FieldCheckbox, FieldSelect } from '@proj-airi/ui'
 import { computed, onUnmounted, ref, watch } from 'vue'
 import { useI18n } from 'vue-i18n'
 
+import SpeechStreamingPlayground from './SpeechStreamingPlayground.vue'
+
 import { TestDummyMarker } from '../../Gadgets'
 
 const props = defineProps<{
@@ -204,6 +206,12 @@ defineExpose({
         {{ errorMessage }}
       </div>
       <audio v-if="audioUrl" ref="audioPlayer" :src="audioUrl" controls class="mt-2 w-full" />
+
+      <SpeechStreamingPlayground
+        :text="testText"
+        :voice="selectedVoice"
+        :generate-speech="generateSpeech"
+      />
     </div>
     <!-- Slot for additional provider-specific UI in the playground -->
     <slot />
diff --git a/packages/stage-ui/src/components/Scenarios/Providers/SpeechStreamingPlayground.vue b/packages/stage-ui/src/components/Scenarios/Providers/SpeechStreamingPlayground.vue
new file mode 100644
index 000000000..7be0b8adf
--- /dev/null
+++ b/packages/stage-ui/src/components/Scenarios/Providers/SpeechStreamingPlayground.vue
@@ -0,0 +1,163 @@
+<script setup lang="ts">
+import type { TTSInputChunk } from '../../../utils/tts'
+
+import { animate } from 'animejs'
+import { ref } from 'vue'
+
+import { useQueue } from '../../../composables/queue'
+import { useMessageContentQueue } from '../../../composables/queues'
+import { useAudioContext } from '../../../stores/audio'
+import { chunkTTSInput } from '../../../utils/tts'
+
+const props = defineProps<{
+  text: string
+  // Provider-specific handlers (provided from parent)
+  generateSpeech: (input: string, voice: string, useSSML: boolean) => Promise<ArrayBuffer>
+  voice: string
+}>()
+
+const { audioContext } = useAudioContext()
+const nowSpeaking = ref(false)
+const ttsInputChunks = ref<TTSInputChunk[]>([])
+const speechGenerationIndex = ref(-1)
+
+const audioQueue = useQueue<{ audioBuffer: AudioBuffer, text: string }>({
+  handlers: [
+    (ctx) => {
+      return new Promise((resolve) => {
+        // Create an AudioBufferSourceNode
+        const source = audioContext.createBufferSource()
+        source.buffer = ctx.data.audioBuffer
+
+        // Connect the source to the AudioContext's destination (the speakers)
+        source.connect(audioContext.destination)
+
+        // Start playing the audio
+        nowSpeaking.value = true
+        source.start(0)
+        source.onended = () => {
+          nowSpeaking.value = false
+          resolve()
+        }
+      })
+    },
+  ],
+})
+
+async function handleSpeechGeneration(ctx: { data: string }) {
+  speechGenerationIndex.value++
+
+  try {
+    const input = ctx.data
+    const res = await props.generateSpeech(input, props.voice, false)
+
+    // Decode the ArrayBuffer into an AudioBuffer
+    const audioBuffer = await audioContext.decodeAudioData(res)
+    await audioQueue.add({ audioBuffer, text: ctx.data })
+  }
+  catch (error) {
+    console.error('Speech generation failed:', error)
+  }
+}
+
+const ttsQueue = useQueue<string>({
+  handlers: [
+    handleSpeechGeneration,
+  ],
+})
+
+const messageContentQueue = useMessageContentQueue(ttsQueue)
+
+async function testStreaming() {
+  await messageContentQueue.add(props.text)
+}
+
+async function testChunking() {
+  const chunks = []
+  for await (const chunk of chunkTTSInput(props.text, { boost: 1, minimumWords: 4, maximumWords: 12 })) {
+    chunks.push(chunk)
+  }
+  ttsInputChunks.value = chunks
+}
+</script>
+
+<template>
+  <div class="flex items-center gap-1 text-sm font-medium">
+    Streaming Playground
+  </div>
+  <div flex="~ row" gap-4>
+    <button
+      border="neutral-800 dark:neutral-200 solid 2" transition="border duration-250 ease-in-out"
+      rounded-lg px-4 text="neutral-100 dark:neutral-900" py-2 text-sm
+      bg="neutral-700 dark:neutral-300" @click="testChunking"
+    >
+      <div flex="~ row" items-center gap-2>
+        <div i-solar:round-double-alt-arrow-right-bold-duotone />
+        <span>Test chunking</span>
+      </div>
+    </button>
+
+    <button
+      v-if="ttsInputChunks.length > 0"
+      border="neutral-800 dark:neutral-200 solid 2" transition="border duration-250 ease-in-out"
+      rounded-lg px-4 text="neutral-100 dark:neutral-900" py-2 text-sm
+      bg="neutral-700 dark:neutral-300" @click="testStreaming"
+    >
+      <div flex="~ row" items-center gap-2>
+        <div i-solar:round-double-alt-arrow-right-bold-duotone />
+        <span>Test streaming</span>
+      </div>
+    </button>
+  </div>
+
+  <div flex="~ col gap-2 items-start" py-4>
+    <div
+      v-for="(chunk, i) in ttsInputChunks"
+      :key="i"
+      flex="~ row gap-2 items-center"
+    >
+      <div
+        flex="~ row gap-2 items-center"
+        rounded-xl px-2 py-1.5
+        :class="{
+          'bg-neutral-100 dark:bg-neutral-800': speechGenerationIndex < i,
+          'bg-neutral-200 dark:bg-neutral-700': speechGenerationIndex >= i,
+        }"
+      >
+        <span ml-1>{{ chunk.text }}</span>
+        <span
+          rounded-full px-2 py-.5 text-nowrap text-xs
+          b="~ dashed"
+          :class="{
+            'b-green text-green': chunk.reason === 'boost',
+            'b-orange text-orange': chunk.reason === 'limit',
+            'b-red text-red': chunk.reason === 'hard',
+            'b-purple text-purple': chunk.reason === 'flush',
+          }"
+        >
+          {{ chunk.words }} words,
+          {{ chunk.reason }}
+        </span>
+      </div>
+      <Transition
+        :css="false"
+        @enter="(el) => animate(el, {
+          opacity: [0, 1],
+          translateX: [10, 0],
+          duration: 200,
+          ease: 'inOut',
+        })"
+      >
+        <div
+          v-if="speechGenerationIndex >= i"
+          tag="div"
+          flex="~ row items-center gap-1"
+          text-sm
+        >
+          <div i-solar-check-circle-line-duotone />
+          <div>Queued</div>
+        </div>
+      </Transition>
+    </div>
+  </div>
+</template>
diff --git a/packages/stage-ui/src/composables/queues.ts b/packages/stage-ui/src/composables/queues.ts
index 01e0870f6..84167fd6b 100644
--- a/packages/stage-ui/src/composables/queues.ts
+++ b/packages/stage-ui/src/composables/queues.ts
@@ -6,6 +6,7 @@ import { ref } from 'vue'
 
 import { llmInferenceEndToken } from '../constants'
 import { EMOTION_VALUES } from '../constants/emotions'
+import { chunkTTSInput } from '../utils/tts'
 import { useQueue } from './queue'
 
 export function useEmotionsMessageQueue(emotionsQueue: UseQueueReturn<Emotion>) {
@@ -116,19 +117,8 @@ export function useMessageContentQueue(ttsQueue: UseQueueReturn<string>) {
           return
         }
 
-        const endMarker = /[.?!]/
-        processed.value += ctx.data
-
-        while (processed.value) {
-          const endMarkerExecArray = endMarker.exec(processed.value)
-          if (!endMarkerExecArray || typeof endMarkerExecArray.index === 'undefined')
-            break
-
-          const before = processed.value.slice(0, endMarkerExecArray.index + 1)
-          const after = processed.value.slice(endMarkerExecArray.index + 1)
-
-          await ttsQueue.add(before)
-          processed.value = after
+        for await (const chunk of chunkTTSInput(ctx.data)) {
+          await ttsQueue.add(chunk.text)
         }
       },
     ],
diff --git a/packages/stage-ui/src/utils/tts.ts b/packages/stage-ui/src/utils/tts.ts
new file mode 100644
index 000000000..03e1f20d0
--- /dev/null
+++ b/packages/stage-ui/src/utils/tts.ts
@@ -0,0 +1,142 @@
+import type { ReaderLike } from 'clustr'
+
+import { readGraphemeClusters } from 'clustr'
+
+const keptPunctuations = new Set('?？!！')
+const hardPunctuations = new Set('.。?？!！…⋯～~「」\n\t\r')
+const softPunctuations = new Set(',，、–—:：;；《》')
+
+export interface TTSInputChunk {
+  text: string
+  words: number
+  reason: 'boost' | 'limit' | 'hard' | 'flush'
+}
+
+export interface TTSInputChunkOptions {
+  boost?: number
+  minimumWords?: number
+  maximumWords?: number
+}
+
+/**
+ * Processes the input string or UTF-8 byte stream reader into chunks suitable for TTS synthesis.
+ *
+ * @param input A string or a ReaderLike object that reads from an underlying UTF-8 byte stream.
+ * @param options
+ * @param options.boost Specifies the number of chunks to yield using greedier rules. This may help
+ *                      reduce the initial delay when processing long input text.
+ * @param options.minimumWords Minimum number of words in a chunk.
+ * @param options.maximumWords Maximum number of words in a chunk.
+ */
+export async function* chunkTTSInput(input: string | ReaderLike, options?: TTSInputChunkOptions): AsyncGenerator<TTSInputChunk, void, unknown> {
+  const {
+    boost = 2,
+    minimumWords = 4,
+    maximumWords = 12,
+  } = options ?? {}
+
+  const iterator = readGraphemeClusters(
+    typeof input === 'string'
+      ? new ReadableStream({
+          start(controller) {
+            controller.enqueue(new TextEncoder().encode(input))
+            controller.close()
+          },
+        }).getReader()
+      : input,
+  )
+
+  const segmenter = new Intl.Segmenter(undefined, { granularity: 'word' }) // I love Intl.Segmenter
+
+  let yieldCount = 0
+  let buffer = ''
+  let chunk = ''
+  let chunkWordsCount = 0
+
+  let previousValue: string | undefined
+  let current = await iterator.next()
+
+  while (!current.done) {
+    const value = current.value
+
+    if (value.length > 1) {
+      previousValue = value
+      current = await iterator.next()
+      continue
+    }
+
+    const hard = hardPunctuations.has(value)
+    const soft = softPunctuations.has(value)
+    const kept = keptPunctuations.has(value)
+
+    if (hard || soft) {
+      switch (value) {
+        case '.':
+        case ',': {
+          if (previousValue !== undefined && /\d/.test(previousValue)) {
+            const next = await iterator.next()
+            if (!next.done && next.value && /\d/.test(next.value)) {
+              // This dot could be a decimal point, so we skip it
+              previousValue = next.value
+              current = next
+              continue
+            }
+          }
+        }
+      }
+
+      if (buffer.length === 0) {
+        previousValue = value
+        current = await iterator.next()
+        continue
+      }
+
+      const words = [...segmenter.segment(buffer)].filter(w => w.isWordLike)
+
+      if (chunkWordsCount > minimumWords && chunkWordsCount + words.length > maximumWords) {
+        const text = kept ? chunk.trim() + value : chunk.trim()
+        yield {
+          text,
+          words: chunkWordsCount,
+          reason: 'limit',
+        }
+        yieldCount++
+        chunk = ''
+        chunkWordsCount = 0
+      }
+
+      chunk += buffer + value
+      chunkWordsCount += words.length
+      buffer = ''
+
+      if (hard || chunkWordsCount > maximumWords || yieldCount < boost) {
+        const text = chunk.trim()
+        yield {
+          text,
+          words: chunkWordsCount,
+          reason: hard ? 'hard' : chunkWordsCount > maximumWords ? 'limit' : 'boost',
+        }
+        yieldCount++
+        chunk = ''
+        chunkWordsCount = 0
+      }
+
+      previousValue = value
+      current = await iterator.next()
+      continue
+    }
+
+    buffer += value
+    previousValue = value
+    current = await iterator.next()
+  }
+
+  if (chunk.length > 0 || buffer.length > 0) {
+    const text = (chunk + buffer).trim()
+    yield {
+      text,
+      words: chunkWordsCount + [...segmenter.segment(buffer)].filter(w => w.isWordLike).length,
+      reason: 'flush',
+    }
+  }
+}
diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml
index acf6838fd..65ae02cab 100644
--- a/pnpm-lock.yaml
+++ b/pnpm-lock.yaml
@@ -1419,6 +1419,9 @@ importers:
       '@vitejs/plugin-vue':
         specifier: ^6.0.0
         version: 6.0.0(vite@6.3.5(@types/node@24.0.13)(jiti@2.4.2)(less@4.3.0)(lightningcss@1.30.1)(terser@5.43.1)(tsx@4.20.3)(yaml@2.8.0))(vue@3.5.17(typescript@5.8.3))
+      clustr:
+        specifier: ^0.0.3
+        version: 0.0.3
       histoire:
         specifier: 1.0.0-alpha.2
         version: 1.0.0-alpha.2(@types/node@24.0.13)(bufferutil@4.0.9)(less@4.3.0)(lightningcss@1.30.1)(terser@5.43.1)(utf-8-validate@5.0.10)(vite@6.3.5(@types/node@24.0.13)(jiti@2.4.2)(less@4.3.0)(lightningcss@1.30.1)(terser@5.43.1)(tsx@4.20.3)(yaml@2.8.0))
@@ -7047,6 +7050,10 @@ packages:
     resolution: {integrity: sha512-eYm0QWBtUrBWZWG0d386OGAw16Z995PiOVo2B7bjWSbHedGl5e0ZWaq65kOGgUSNesEIDkB9ISbTg/JK9dhCZA==}
     engines: {node: '>=6'}
 
+  clustr@0.0.3:
+    resolution: {integrity: sha512-nGJZKUSi2K83migJ7sJj/JYNVbT0L81b7fFNcQ3XIgkXqkeG1tpO1NtISDjr3Unnjqsztzvp/p+8JgBhpAc87g==}
+    engines: {node: '>=22.10.0'}
+
   color-convert@2.0.1:
     resolution: {integrity: sha512-RRECPsj7iu/xb5oKYcsFHSppFNnsj/52OVTRKb4zP5onXwVF3zVmmToNcOfGC+CRDpfK/U584fMg38ZHCaElKQ==}
     engines: {node: '>=7.0.0'}
@@ -19426,6 +19433,8 @@ snapshots:
   clsx@2.1.1:
     optional: true
 
+  clustr@0.0.3: {}
+
   color-convert@2.0.1:
     dependencies:
       color-name: 1.1.4