cmc-oa/oa-ui/src/api/llm/rag.js

/*
 * @Author: ysh
 * @Date: 2025-07-08 16:43:22
 * @LastEditors: Please set LastEditors
 * @LastEditTime: 2025-07-16 15:49:45
 */
import request from '@/utils/request'
import { getToken } from '@/utils/auth'

// 查询cmc聊天记录列表
export function getAnswer(question, collectionName) {
  return request({
    url: '/llm/rag/answer',
    method: 'get',
    params: { question, collectionName }
  })
}

// 查询上下文引用文件
export function getContextFile(question, collectionName) {
  return request({
    url: '/llm/rag/context',
    method: 'get',
    params: { question, collectionName }
  })
}

function parseSseEvents(buffer) {
  // SSE 事件以空行分隔：\n\n（兼容 \r\n）
  const events = []
  const normalized = buffer.replace(/\r\n/g, '\n')
  const parts = normalized.split('\n\n')
  // 最后一段可能是不完整事件，留给下次拼接
  const rest = parts.pop() ?? ''

  for (const part of parts) {
    if (!part.trim()) continue
    const lines = part.split('\n')
    const dataLines = []
    for (const line of lines) {
      if (line.startsWith('data:')) {
        dataLines.push(line.slice(5).trimStart())
      }
    }
    const data = dataLines.join('\n')
    if (data !== '') events.push(data)
  }

  return { events, rest }
}

function normalizeSseData(data) {
  if (!data) return ''
  const trimmed = data.trim()
  // 移除可能的JSON包装（如果有的话）
  if (trimmed.startsWith('{') && trimmed.endsWith('}')) {
    try {
      const parsed = JSON.parse(trimmed)
      if (parsed.content) return parsed.content
      if (parsed.resultContent) return parsed.resultContent
    } catch (e) {
      // 不是JSON，直接返回
    }
  }
  return trimmed
}

function stripToolCallStream(text, state) {
  if (!text) return ''
  const s = String(text)
  if (!state.inToolCall) {
    const toolStart = s.indexOf('<tool_call>')
    if (toolStart === -1) {
      return s
    } else {
      state.inToolCall = true
      return s.slice(0, toolStart)
    }
  } else {
    const toolEnd = s.indexOf('</tool_call>')
    if (toolEnd === -1) {
      return ''
    } else {
      state.inToolCall = false
      return s.slice(toolEnd + 11)
    }
  }
}

function streamFetchSse(url, onMessage, onError, onComplete) {
  const controller = new AbortController()
  const toolCallState = { inToolCall: false }
  fetch(url, {
    method: 'GET',
    headers: {
      'Authorization': 'Bearer ' + getToken(),
      'Accept': 'text/event-stream',
      'Cache-Control': 'no-cache'
    },
    signal: controller.signal
  })
    .then(async (response) => {
      if (!response.ok) throw new Error(`HTTP error! status: ${response.status}`)
      if (!response.body) throw new Error('ReadableStream not supported')

      const reader = response.body.getReader()
      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 })
        const parsed = parseSseEvents(buffer)
        buffer = parsed.rest

        for (const data of parsed.events) {
          const normalized = normalizeSseData(data)
          if (normalized === '[DONE]') {
            onComplete()
            controller.abort()
            return
          }
          const visible = stripToolCallStream(normalized, toolCallState)
          if (visible !== '') onMessage(visible)
        }
      }

      // 兜底：流结束但没收到 [DONE]
      if (buffer.trim()) {
        const parsed = parseSseEvents(buffer + '\n\n')
        for (const data of parsed.events) {
          const normalized = normalizeSseData(data)
          if (normalized !== '' && normalized !== '[DONE]') {
            const visible = stripToolCallStream(normalized, toolCallState)
            if (visible !== '') onMessage(visible)
          }
        }
      }
      onComplete()
    })
    .catch((error) => {
      if (error.name === 'AbortError') return
      console.error('流式请求错误:', error)
      onError(error)
    })

  return controller
}

// 流式回答API（SSE）
export function getAnswerStream(question, collectionName, onMessage, onError, onComplete) {
  const baseURL = process.env.VUE_APP_BASE_API
  const url = `${baseURL}/llm/rag/answer?question=${encodeURIComponent(question)}&collectionName=${encodeURIComponent(collectionName)}`
  return streamFetchSse(url, onMessage, onError, onComplete)
}