fix: adapt AI assistant streaming answer parsing (#319)

Co-authored-by: JounQin <admin@1stg.me>

Author: yangxiaolang

.changeset/many-guests-shake.md

@@ -0,0 +1,5 @@
+---
+"@alauda/doom": patch
+---
+
+fix: adapt AI assistant streaming answer parsing

.gitignore

@@ -14,3 +14,20 @@ node_modules
 /test.*
 *.log
 *.tsbuildinfo
+
+# Local AI code agent configs
+.codex
+.codex/
+.claude
+.claude/
+.continue
+.continue/
+.cursor
+.cursor/
+.gemini
+.gemini/
+.roo
+.roo/
+.windsurf
+.windsurf/
+.aider*

packages/doom/src/global/Intelligence/AIAssistant/index.tsx

@@ -17,7 +17,11 @@ import { Preamble } from './Preamble/index.tsx'
 import { ResizableUserInput } from './ResizableUserInput/index.tsx'
 import { Thinking } from './Thinking.tsx'
 import type { ChatMessage } from './types.ts'
-import { parseStreamContent } from './utils.ts'
+import {
+  consumeSSEEvents,
+  getAnswerDelta,
+  parseStreamContent,
+} from './utils.ts'
 
 import CloseIcon from '@alauda/doom/assets/close.svg?react'
 import LogoutIcon from '@alauda/doom/assets/logout.svg?react'
@@ -31,8 +35,6 @@ export interface AIAssistantProps {
   onCleanup?: () => void
 }
 
-let textDecoder: TextDecoder | undefined
-
 export const AIAssistant = ({ open, onOpenChange }: AIAssistantProps) => {
   const t = useTranslation()
 
@@ -132,7 +134,7 @@ export const AIAssistant = ({ open, onOpenChange }: AIAssistantProps) => {
 
     const sessionId = sessionIdRef.current
 
-    const res = await xfetch('/smart/api/smart_answer', {
+    const res = await xfetch('/smart/api/smart_answer_with_search', {
       type: null,
       method: ApiMethod.POST,
       body: {
@@ -141,39 +143,80 @@ export const AIAssistant = ({ open, onOpenChange }: AIAssistantProps) => {
       },
     })
 
-    textDecoder ??= new TextDecoder()
+    const textDecoder = new TextDecoder()
 
     let text = ''
+    let sseBuffer = ''
+    let sseEvent: string | undefined
 
     for await (const chunk_ of res.body! as unknown as AsyncIterable<Uint8Array>) {
       if (sessionId !== sessionIdRef.current) {
         break
       }
 
-      const chunk = textDecoder.decode(chunk_)
+      sseBuffer += textDecoder.decode(chunk_, { stream: true })
 
-      text += chunk
-        .replace(/data: (?:\n\n)?/g, '')
-        .replace(/(?<!\n)\n\n(?!\n)/g, '')
-        .replace(/(?<!\n)\n\n\n\n(?!\n)/g, '\n')
-        .replace(/\n{2,}/g, '\n')
+      const result = consumeSSEEvents(sseBuffer, sseEvent)
+      const { events, remainder } = result
+      sseBuffer = remainder
+      sseEvent = result.event
 
-      const parsed = parseStreamContent(text)
+      let updated = false
+
+      for (const event of events) {
+        const delta = getAnswerDelta(event)
+        if (!delta) {
+          continue
+        }
+
+        text += delta
+        updated = true
+      }
 
-      if (
-        !parsed.refDocs.length &&
-        !parsed.thinkingProcess &&
-        !parsed.content
-      ) {
+      if (!updated) {
         continue
       }
 
+      const parsed = parseStreamContent(text)
+
       flushMessages((messages) => [
         ...messages.slice(0, index),
         { ...messages[index], ...parsed },
         ...messages.slice(index + 1),
       ])
     }
+
+    if (sessionId !== sessionIdRef.current) {
+      return
+    }
+
+    sseBuffer += textDecoder.decode()
+
+    const { events } = consumeSSEEvents(`${sseBuffer}\n`, sseEvent)
+
+    let updated = false
+
+    for (const event of events) {
+      const delta = getAnswerDelta(event)
+      if (!delta) {
+        continue
+      }
+
+      text += delta
+      updated = true
+    }
+
+    if (!updated) {
+      return
+    }
+
+    const parsed = parseStreamContent(text)
+
+    flushMessages((messages) => [
+      ...messages.slice(0, index),
+      { ...messages[index], ...parsed },
+      ...messages.slice(index + 1),
+    ])
   }
 
   const [loading, setLoading] = useState(false)

packages/doom/src/global/Intelligence/AIAssistant/utils.ts

@@ -4,21 +4,161 @@ import { isObject } from 'es-toolkit/compat'
 
 import type { RefDoc } from './types.js'
 
+export interface SSEEvent {
+  data: string
+  event?: string
+}
+
+export interface SSEConsumeResult {
+  event?: string
+  events: SSEEvent[]
+  remainder: string
+}
+
+const IGNORED_SSE_EVENT_TYPES =
+  /(?:^|[._-])(?:metadata|trace|updates?)(?:$|[._-])/i
+
+const STREAM_TEXT_KEYS = ['delta', 'text', 'content', 'answer', 'message']
+
 export const unicodeToString = (unicodeStr: string) =>
   unicodeStr.replace(/\\u([0-9a-fA-F]{4})/g, (_match, p1: string) =>
     String.fromCharCode(Number.parseInt(p1, 16)),
   )
 
+const isRecord = (value: unknown): value is Record<string, unknown> =>
+  isObject(value) && !Array.isArray(value)
+
 export const removePrefix = (
   prefix: string | null | undefined,
   text: string,
 ) => (prefix && text.startsWith(prefix) ? text.slice(prefix.length) : text)
 
+const parseSSEFieldValue = (line: string, field: string) => {
+  const value = line.slice(field.length + 1)
+  return value.startsWith(' ') ? value.slice(1) : value
+}
+
+const getStringValue = (value: unknown): string | undefined => {
+  if (typeof value === 'string') {
+    return value
+  }
+
+  if (Array.isArray(value)) {
+    const text = value.map((item) => getStringValue(item) || '').join('')
+    return text || undefined
+  }
+
+  if (!isRecord(value)) {
+    return
+  }
+
+  for (const key of STREAM_TEXT_KEYS) {
+    const text = getStringValue(value[key])
+    if (text) {
+      return text
+    }
+  }
+}
+
+const getChoicesDelta = (choices: unknown): string | undefined => {
+  if (!Array.isArray(choices)) {
+    return
+  }
+
+  const text = choices.map((choice) => getStringValue(choice) || '').join('')
+
+  return text || undefined
+}
+
+export const consumeSSEEvents = (
+  buffer: string,
+  event?: string,
+): SSEConsumeResult => {
+  const events: SSEEvent[] = []
+
+  let remaining = buffer
+  let currentEvent = event
+
+  let lineBoundary = remaining.match(/\r?\n/)
+
+  while (lineBoundary?.index != null) {
+    const line = remaining.slice(0, lineBoundary.index)
+    remaining = remaining.slice(lineBoundary.index + lineBoundary[0].length)
+
+    if (!line) {
+      currentEvent = undefined
+    } else if (line.startsWith(':')) {
+      //
+    } else if (line.startsWith('event:')) {
+      currentEvent = parseSSEFieldValue(line, 'event')
+    } else if (line.startsWith('data:')) {
+      events.push({
+        event: currentEvent,
+        data: parseSSEFieldValue(line, 'data'),
+      })
+    }
+
+    lineBoundary = remaining.match(/\r?\n/)
+  }
+
+  return {
+    event: currentEvent,
+    events,
+    remainder: remaining,
+  }
+}
+
+export const getAnswerDelta = ({ event, data }: SSEEvent) => {
+  if (!data || data === '[DONE]') {
+    return ''
+  }
+
+  if (event && IGNORED_SSE_EVENT_TYPES.test(event)) {
+    return ''
+  }
+
+  try {
+    const payload = JSON.parse(data)
+
+    if (typeof payload === 'string') {
+      return payload
+    }
+
+    if (!isRecord(payload)) {
+      return ''
+    }
+
+    const type = typeof payload.type === 'string' ? payload.type : undefined
+
+    if (type && IGNORED_SSE_EVENT_TYPES.test(type)) {
+      return ''
+    }
+
+    const choiceDelta = getChoicesDelta(payload.choices)
+
+    if (choiceDelta) {
+      return choiceDelta
+    }
+
+    for (const key of STREAM_TEXT_KEYS) {
+      const text = getStringValue(payload[key])
+      if (text) {
+        return text
+      }
+    }
+
+    return ''
+  } catch {
+    return data
+  }
+}
+
 export function parseStreamContent(text: string) {
-  const matchDocs = text.match(/<docs>([\s\S]*?)<\/docs>/)
-  const docsReferences = matchDocs?.length
-    ? unicodeToString(matchDocs[1])
+  const docsMatches = [...text.matchAll(/<docs>([\s\S]*?)<\/docs>/g)]
+  const docsReferences = docsMatches.length
+    ? unicodeToString(docsMatches.at(-1)![1])
     : '[]'
+
   let refDocs: RefDoc[] = []
   try {
     refDocs = (JSON.parse(docsReferences) as RefDoc[]).filter((r) =>
@@ -28,19 +168,23 @@ export function parseStreamContent(text: string) {
     //
   }
 
-  const matchThinks = text.match(/<think>[\s\S]*?<\/think>|^[\s\S]*?<\/think>/g)
-  const firstMatched = matchThinks?.[0]
-  let matched = firstMatched || matchThinks?.[1]
-  if (matched) {
-    matched = removePrefix(firstMatched || '', matched)
+  const thinkMatches = [...text.matchAll(/<think>([\s\S]*?)<\/think>/g)]
+  let thinkingProcess = thinkMatches.at(-1)?.[1]
+
+  if (!thinkingProcess) {
+    const openThinkMatch = text.match(/<think>([\s\S]*)$/)
+    thinkingProcess = openThinkMatch?.[1]
   }
-  const thinkingProcess = matchThinks && matched
+
+  const content = text
+    .replace(/<docs>[\s\S]*?<\/docs>/g, '')
+    .replace(/<think>[\s\S]*?<\/think>/g, '')
+    .replace(/<docs>[\s\S]*$/g, '')
+    .replace(/<think>[\s\S]*$/g, '')
 
   return {
     refDocs,
-    thinkingProcess,
-    content: unicodeToString(
-      removePrefix(thinkingProcess, removePrefix(matchDocs?.[0], text)),
-    ),
+    thinkingProcess: thinkingProcess && unicodeToString(thinkingProcess),
+    content: unicodeToString(content),
   }
 }