smartai2dian0/src/modules/conversation-chat/use-websocket-chat.js

import { ref } from 'vue';
import { kuky_Authorization } from '../../base/storage';

/**
 * 获取websocket连接的地址
 * @param {string} type - 对话类型，可选值：report/excel/chat
 * @param {string} agent_id - 对话使用的agent ID
 * @param {string} chat_id - 对话记录的ID
 * @returns
 */
function getWSUrl(type, agent_id, chat_id) {
  const host = location.host;
  const token = kuky_Authorization.get();

  if (type === 'report') {
    return `ws://${host}/api/${type}/ws/${agent_id}/${chat_id}?token=${token}`;
  } else if (type === 'chat') {
    return `ws://${host}/api/${type}/ws/${agent_id}/${chat_id}?token=${token}`;
  } else if (type === 'excel') {
    return `ws://${host}/api/document/ws/excel`;
  } else {
    return `ws://${host}/api/chat/ws/${agent_id}/${chat_id}?token=${token}`;
  }
}

/**
 * 确保一个值是函数
 * 如果值本身是函数，则返回值本身
 * 如果值本身不是函数，则返回一个空函数
 * @param {function} fn
 * @returns
 */
function protectedFunction(fn) {
  if (typeof fn === 'function') {
    return fn;
  }

  return Function.prototype;
}

/**
 * 使用websocket进行对话
 * @see https://developer.mozilla.org/zh-CN/docs/Web/API/WebSockets_API
 * @param {string} agent_id - agent id
 * @param {string} chat_id - 对话ID
 * @param {string} type - 对话类型，可选值：report/excel/chat
 * @param {object} trigger - 触发的方法
 * @param {function} trigger.onMessage - 当收到消息时触发的方法
 * @param {function} trigger.onClose - 当websocket关闭时触发的方法
 * @param {function} trigger.onError - 当websocket出错时触发的方法
 * @returns
 */
export default function useWebSocketChat(trigger) {
  let delay_message_list = [];

  const onMessage = protectedFunction(trigger.onMessage);
  const onClose = protectedFunction(trigger.onClose);
  const onError = protectedFunction(trigger.onError);

  /**
   * websocket当前状态
   * @see https://developer.mozilla.org/zh-CN/docs/Web/API/WebSocket/readyState
   * 0 CONNECTING 套接字已经创建，但是连接尚未打开
   * 1 OPEN 连接已经打开，准备进行通信
   * 2 CLOSING 连接正在关闭中
   * 3 CLOSED 连接已关闭或者无法打开
   * -1 此处自定义值，未启动
   */
  const websocket_status = ref(-1);

  /**
   * 创建一个websocket连接
   * @returns
   */
  function createWSConnection(agent_id, chat_id, type) {
    const wsUrl = getWSUrl(type, agent_id, chat_id);

    const wskt = new WebSocket(wsUrl);

    // 当 WebSocket 打开时...
    wskt.onopen = function (event) {
      delay_message_list.forEach((item) => {
        wskt.send(item);
      });

      // 所有消息发送完成后，清空消息队列
      delay_message_list = [];
    };

    // 当 WebSocket 关闭时...
    wskt.onclose = function (event) {
      onClose(event);
    };

    // 当有错误发生时...
    wskt.onerror = function (error) {
      onError(error);
    };

    /**
     * 当收到websocket消息时
     * @param {WebSocketEvent} event
     */
    wskt.onmessage = (event) => {
      onMessage(event);
    };

    return wskt;
  }

  let websocket = null;

  return {
    /**
     * 初始化websocket连接
     * @param {string} agent_id
     * @param {string} chat_id
     * @param {string} type
     */
    initWSConnection(agent_id, chat_id, type) {
      if (websocket === null) {
        // websocket未创建：立即创建
        websocket = createWSConnection(agent_id, chat_id, type);
      }

      return websocket;
    },
    /**
     * 发送会话消息
     * @param {string} msg
     */
    sendWSMessage(msg) {
      if (websocket === null) {
        // websocket未创建：立即创建
        websocket = createWSConnection();
      }

      // 判断socket是否开启
      // 如果开启，则直接发送，如果没有开启，则将数据推入消息队列
      if (websocket.readyState === 0) {
        delay_message_list.push(msg);
      }

      if (websocket.readyState === 1) {
        websocket.send(msg);
      }

      if (websocket.readyState === 2) {
        // 正在关闭
        console.log('web socket正在关闭');
      }

      if (websocket.readyState === 3) {
        delay_message_list.push(msg);

        // 已关闭: 重新创建socket
        websocket = createWSConnection();
      }
    },
    /**
     * 获取websocket实例本身
     * @returns
     */
    getWSInstance() {
      return websocket;
    },
    /**
     * 结束对话，断开websocket连接
     */
    closeWSConnect() {
      if (websocket === null) {
        return;
      }

      websocket.close();
    }
  };
}