YG_FT/frontend/src/composables/useStreamChat.ts

import { ref } from 'vue'
import { streamChat, streamChatReal } from '@/api/modules/compare'

export interface StreamMessage {
  /** 用户问题 */
  question: string
  /** 完整回答（含 think 标签原始内容） */
  fullContent: string
  /** 去除 think 标签后的展示内容 */
  displayContent: string
  /** 思考过程内容 */
  thinkContent: string
  /** 是否正在思考（think 标签未闭合） */
  isThinking: boolean
  /** 是否流式中 */
  isStreaming: boolean
  /** 是否已完成 */
  done: boolean
  /** 错误信息 */
  error?: string
}

export interface SendOptions {
  /** 是否使用 mock 模式（默认 true，向后兼容） */
  useMock?: boolean
}

/**
 * 流式对话 composable
 * 移植自原 model-chat.html：
 * - fetch + body.getReader() + TextDecoder
 * - <think>...</think> 标签解析（思考过程可折叠）
 * - 50ms 节流更新
 */
export function useStreamChat() {
  const message = ref<StreamMessage>({
    question: '',
    fullContent: '',
    displayContent: '',
    thinkContent: '',
    isThinking: false,
    isStreaming: false,
    done: false,
  })
  const loading = ref(false)

  /** 从 SSE 帧中提取错误信息（后端/计算节点错误以 data: {"error": "..."} 形式下发） */
  function extractSseError(buffer: string): string | null {
    const trimmed = buffer.trim()
    if (!trimmed.startsWith('data: ')) return null
    const lines = trimmed.split(/\r?\n/)
    for (let i = lines.length - 1; i >= 0; i--) {
      const line = lines[i].trim()
      if (!line.startsWith('data: ')) continue
      try {
        const obj = JSON.parse(line.slice(6))
        if (obj && typeof obj.error === 'string' && obj.error) return obj.error
      } catch {
        /* 非 JSON 的 data 行忽略 */
      }
    }
    return trimmed
  }

  /** 从内容中解析 think 标签 */
  function parseContent(content: string) {
    const thinkRegex = /<think>([\s\S]*?)(<\/think>)?/g
    let think = ''
    let display = content
    let isThinking = false

    let match
    // 检查是否有未闭合的 think 标签
    const openTags = (content.match(/<think>/g) || []).length
    const closeTags = (content.match(/<\/think>/g) || []).length
    isThinking = openTags > closeTags

    // 提取所有 think 内容
    while ((match = thinkRegex.exec(content)) !== null) {
      think += match[1]
    }
    // 去除 think 标签得到展示内容
    display = content.replace(/<think>[\s\S]*?(<\/think>|$)/g, '').trim()

    return { think: think.trim(), display, isThinking }
  }

  /**
   * 发起流式对话
   * @param payload 后端请求体 { port, model_name, model_path, system_prompt, user_question, ... }
   * @param options 可选配置 { useMock?: boolean }
   */
  async function send(payload: any, options?: SendOptions) {
    const useMock = options?.useMock ?? true
    loading.value = true
    message.value = {
      question: payload.user_question || '',
      fullContent: '',
      displayContent: '',
      thinkContent: '',
      isThinking: false,
      isStreaming: true,
      done: false,
    }

    let lastUpdate = 0
    const UPDATE_INTERVAL = 50 // 50ms 节流

    try {
      const response = useMock
        ? await streamChat(payload)
        : await streamChatReal(payload)

      if (!response.ok) {
        throw new Error(`HTTP ${response.status}`)
      }

      const reader = response.body?.getReader()
      if (!reader) throw new Error('无法读取响应流')

      const decoder = new TextDecoder('utf-8')
      let buffer = ''

      while (true) {
        const { done, value } = await reader.read()
        if (done) break

        buffer += decoder.decode(value, { stream: true })
        message.value.fullContent = buffer

        const now = Date.now()
        if (now - lastUpdate >= UPDATE_INTERVAL) {
          lastUpdate = now
          const parsed = parseContent(buffer)
          message.value.thinkContent = parsed.think
          message.value.displayContent = parsed.display
          message.value.isThinking = parsed.isThinking
        }
      }

      // 最终更新
      // 若整段响应是 SSE 错误帧，提取 error 字段以干净文案展示
      const sseError = extractSseError(buffer)
      if (sseError) {
        message.value.isThinking = false
        message.value.isStreaming = false
        message.value.done = true
        message.value.error = sseError
        message.value.displayContent = sseError
        return
      }
      const parsed = parseContent(buffer)
      message.value.thinkContent = parsed.think
      message.value.displayContent = parsed.display
      message.value.isThinking = false
      message.value.isStreaming = false
      message.value.done = true
    } catch (e: any) {
      message.value.isStreaming = false
      message.value.done = true
      message.value.error = e.message || '流式请求失败'
      message.value.displayContent = message.value.fullContent || message.value.error || '请求失败'
    } finally {
      loading.value = false
    }
  }

  function reset() {
    message.value = {
      question: '',
      fullContent: '',
      displayContent: '',
      thinkContent: '',
      isThinking: false,
      isStreaming: false,
      done: false,
    }
  }

  return { message, loading, send, reset }
}