smartai2dian0/src/modules/conversation-dialog/use-sse-chat.js

import { ref } from 'vue';

import fetchEventSource from '../../3rd-libs/microsoft-feat-event-source';
import { kuky_Authorization } from '../../base/storage';

import { getConfigField } from '../../config/index';

/**
 * 区分智能体类型
 * 参数：mode
 * 1.聊天助手 [agent-dialog，advanced-chat]
 * 2.Agent [agent-chat]
 * 3.工作流 [workflow]
 * */

/**
 * 获取sse连接的地址
 * @param {string} menu_id - 对话类型，可选值：1/2/3/4/5/6/7/8/9
 * @param {string} agent_id - 对话使用的agent ID
 * @param {string} chat_id - 对话记录的ID
 * @param {string} mode - 对话类型
 * [agent-dialog], ragflow 智能体
 * [agent-chat，chat，'advanced-chat']，dify智能体
 * [workflow] dify智能体
 * [agent-basic]
 * @returns
 */
function getSseUrl(agent_id, mode) {
  const chat_id = agent_id;
  let base_url = '';

  // 获取环境配置并动态替换 URL
  if (import.meta.env.MODE === 'production') {
    base_url = getConfigField('API_BASE_URL');
  } else {
    base_url = '';
  }
  if (['agent-dialog'].includes(mode)) {
    return base_url + `/api/v1/chat/${chat_id}/completions`;
  }

  if (['advanced-chat', 'agent-chat', 'chat'].includes(mode)) {
    return base_url + `/api/v1/agent/${chat_id}/completions`;
  }
  if (['workflow'].includes(mode)) {
    return base_url + `/api/v1/workflow/${chat_id}/completions`;
  }
  if (['agent-basic'].includes(mode)) {
    return base_url + `/api/v1/complex/${chat_id}/completions`;
  }

  console.log(base_url);
}

/**
 *获取token
 */
function getToken() {
  return kuky_Authorization.get();
}

export default function useFetchEventSource() {
  const messages_list = ref([]); // 存储接收到的消息
  const isLoading = ref(false); // 是否正在连接
  const error = ref(null); // 错误信息
  const isConnected = ref(false); // 是否已连接
  const abortController = new AbortController(); // 用于取消请求

  /**
   * queryParams: sse接口需要传递的参数
   * */
  const startStream = (queryParams, onMessage = defaultOnMessage) => {
    isLoading.value = true;
    isConnected.value = false;
    error.value = null;
    const { agent_id, session_id, mode, query, files, inputs } = queryParams;

    const sse_url = getSseUrl(agent_id, mode);
    if (!sse_url) {
      return;
    }
    fetchEventSource(sse_url, {
      method: 'POST',
      headers: {
        Authorization: `Bearer ${getToken()}`,
        'Content-Type': 'application/json',
        Accept: 'text/event-stream'
      },
      body: JSON.stringify({ sessionId: session_id, query, files, inputs }),
      signal: abortController.signal, // 支持取消请求messages
      // openWhenHidden: true, //页面退至后台后保持连接
      onmessage(ev) {
        onMessage(ev);
      },
      onclose() {
        console.log('Connection closed by server');
        isConnected.value = false;
        isLoading.value = false;
      },
      onerror(err) {
        console.error('Error received:', err);
        error.value = err;
        isLoading.value = false;
        isConnected.value = false;
      }
    });
  };
  /**
   * 终止当前的流连接
   */
  const stopStream = () => {
    abortController.abort();
    isLoading.value = false;
    isConnected.value = false;
    console.log('Stream aborted');
  };
  /**
   * 默认的消息处理函数
   */
  const defaultOnMessage = (ev) => {
    try {
      const message = JSON.parse(ev.data); // 假设返回的是JSON数据
      console.log('Received message:', message);
      // messages_list.value.push(message); // 将消息存入 messages_list
    } catch (e) {
      console.error('Failed to parse message:', e);
    }
  };

  return {
    startStream,
    stopStream,
    messages_list,
    isLoading,
    error,
    isConnected
  };
}