zjugis_OA/client/src/hooks/web/useSSE.ts

import { ref, reactive, onBeforeUnmount } from 'vue'
import { SSEApi } from '@/interface/ai'

interface UseSSEQueryBody {
  mutationFn: (body: SSEApi) => void
  onMessage?: (data: any) => void
  onError?: (err: any) => void
}

interface AResult {
  content: string // 思考过程
  result: object // 最终结果（含有多个对象，通过状态为key区分）
}

const useSSE = (body: UseSSEQueryBody) => {
  const { mutationFn, onMessage, onError } = body
  const data = reactive<AResult>({ content: '', result: {} })
  const loading = ref<boolean>(false)
  const ctrlAbout = new AbortController()

  // 思考过程更新
  const contentChange = (text: string): void => {
    data.content += text
  }

  // 结果更新
  const resultChange = (key: string, result: object): void => {
    data.result[key] = result
  }

  // 对话结束回调
  const answerDone = () => {
    loading.value = false
  }

  const mutation = (paramBody: any) => {
    loading.value = true
    mutationFn?.({
      paramBody,
      signal: ctrlAbout.signal,
      onMessage: (result: any) => {
        onMessage?.(result)
      },
      onError: (err) => {
        console.log('error: ', err)
        ctrlAbout.abort()
        onError?.(err)
      },
      contentChange,
      resultChange,
      answerDone
    })
  }

  onBeforeUnmount(() => {
    ctrlAbout.abort()
  })

  return { data, ctrlAbout, loading, mutation }
}

export default useSSE