vue3笔记（41-2）websocket

需求：实现 AI Agent 聊天机器人，通过 websocket 与后端通信。
本篇记录 websocket 的使用方法。

socket.ts 文件：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
export interface TCWSRequestData {
  content: string // 输入内容, 必填，非空
  output_len?: number // 输出长度，可选，默认128
  repetition_penalty?: number // 重复惩罚，可选，默认1.0
}
export interface TCWSResponseData {
  code: string
  type: string
  message: string
  data: string // 推理文本结果
  ftl: number // 首字延迟(ms)
  latency: number // 总延迟(ms)
  end?: boolean
}
export const initSocket = (url, onMessage, options: any = {}): any => {
  const handleMessage = (messagePayload: any) => {
    return messagePayload
  }
  const sendRequestRaw = (json: string) => {
    _socket.send(json)
  }
  const sendRequest = (
    content: string,
    output_len: number = 128,
    repetition_penalty: number = 1.0
  ) => {
    const msg = {} as TCWSRequestData
    msg.content = content
    msg.output_len = output_len
    msg.repetition_penalty = repetition_penalty
    _socket.send(JSON.stringify(msg))
  }

  let _socket = null
  let _doRetry = false
  let _retryCount = -1
  const closeSocket = () => {
    console.log('socket 手动关闭')
    _doRetry = false
    if (_socket) {
      _socket.close()
      _socket = null
    }
  }
  const doRetry = () => {
    setTimeout(() => {
      console.log('socket 重试')
      ElMessage.warning('Websocket 重连中...')
      _socket = getSocket()
    }, 3000)
  }
  const getSocket = () => {
    const socket = new WebSocket(url)
    socket.binaryType = 'arraybuffer'
    socket.onopen = () => {
      console.log('socket 连接成功')
      options.onopen && options.onopen()
    }
    socket.onmessage = (data: any) => {
      onMessage(handleMessage(data))
    }
    socket.onclose = () => {
      console.log('socket 关闭，retry：', _doRetry, _retryCount)
      if (_doRetry) {
        if (_retryCount !== 0) {
          doRetry()
          if (_retryCount > 0) _retryCount--
        }
      }
    }
    socket.onerror = () => {
      console.log('socket 出错，retry：', _doRetry, _retryCount)
    }
    return socket
  }
  const getState = () => {
    return (_socket && _socket.readyState) || WebSocket.CLOSED
  }
  _retryCount = options.retryCount || _retryCount
  _doRetry = options.retry || _doRetry
  _socket = getSocket()
  return { sendRequestRaw, sendRequest, closeSocket, getState }
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
import { initSocket, TCWSResponseData } from './socket'
let socket = null

let wsTimer = null
const onMessage = (res) => {
  clearTimeout(wsTimer)

  chatObj.responseData = JSON.parse(res.data)
  if (chatObj.responseData?.end) { // 结束会话
    wsTimer = setTimeout(() => {
      closeWs()
    }, 5000)
    return
  }
  resType.value = 'text'
  resData.value = res.data
  textData.value =
    chatObj.responseData.code == '1' ? chatObj.responseData.data : chatObj.responseData.message

  wsTimer = setTimeout(() => {
    closeWs()
  }, 5000)
}
const openWs = (callback) => {
  if (!chatObj.socket || chatObj.socket.getState() == WebSocket.CLOSED)
    chatObj.socket = initSocket(baseUrl.value + urlPath.value, onMessage, { onopen: callback })
  else callback()
}

const closeWs = () => {
  chatObj.socket && chatObj.socket.closeSocket()
}

onBeforeUnmount(() => {
  closeWs()
})
1
2
3
4
5
6
7
8
9
10
const submitData = () => {
  // json格式
  openWs(() => { 
    socket.sendRequestRaw(JSON.parse(rowData.value))
  })
  // 非json格式
  openWs(() => {
    socket.sendRequest(rowData.value)
  }) 
}

javascript