shalu-marketplace/src/hooks/useChat.ts

import { useXAgent, XStream } from "@ant-design/x";
import { useEffect, useRef, useState } from "react";
import { useSessionStorageState } from "ahooks";
import { GetSessionList, GetSessionMessageList } from "@/api/ai";

import type { ConversationsProps } from "@ant-design/x";
import type { ReactNode } from "react";

// 消息格式
type MessageItem = {
  id: string;
  content: string | ReactNode;
  role: "user" | "assistant" | "system";
  status: "loading" | "done" | "error" | "stop";
  loading?: boolean;
  footer?: ReactNode;
};

// 后端返回格式
type ResponseMessageItem = {
  answer: string;
  conversation_id: string;
  created_at: number;
  event: "message" | "message_end" | "message_error" | "ping";
  message_id: string;
  task_id: string;
};

type ChatParams = {
  // 应用名称
  app_name: string;
  // 会话内容
  chat_query: string;
  // 会话名称 第一次
  chat_name?: string;
  // 会话id 后续会话带入
  conversation_id?: string;
};

type ChatProps = {
  // 应用名称
  app_name: string;
  // 会话id 后续会话带入
  conversation_id?: string;
  // 成功获取会话内容
  onSuccess?: (data: ResponseMessageItem) => void;
  // 更新流式消息内容
  onUpdate: (data: ResponseMessageItem) => void;
  // 异常
  onError?: (error: Error) => void;
};

const defaultConversation = {
  // 会话id
  key: "1",
  label: "新的对话",
};

export function useChat({ app_name, onSuccess, onUpdate, onError }: ChatProps) {
  /**
   * 发送消息加载状态
   */
  const [loading, setLoading] = useState(false);

  /**
   * 加载会话记录列表
   */
  const [loadingSession, setLoadingSession] = useState(false);

  /**
   * 加载消息列表
   */
  const [loadingMessages, setLoadingMessages] = useState(false);

  // 用于停止对话
  const abortController = useRef<AbortController | null>(null);

  /**
   * 消息列表
   */
  const [messages, setMessages] = useState<Array<MessageItem>>([]);

  // 会话列表
  const [conversationList, setConversationList] = useState<
    ConversationsProps["items"]
  >([{ ...defaultConversation }]);

  // 活动对话
  const [activeConversation, setActiveConversation] = useState("1");

  // 当前智能体对象
  const [currentAgent, setCurrentAgent] = useSessionStorageState("agent-map");

  useEffect(() => {
    setLoadingSession(true);
    GetSessionList({
      app_name,
      page_index: 1,
    })
      .then((res) => {
        setConversationList([
          { ...defaultConversation },
          ...(res?.result?.model || []).map((item: any) => ({
            ...item,
            key: item.sessionId,
            label: item.name,
          })),
        ]);
      })
      .finally(() => {
        setLoadingSession(false);
      });
  }, [app_name]);

  /**
   * 切换会话
   * @param key 会话id
   * @returns 
   */
  const changeConversation = async (key: string) => {
    setActiveConversation(key);
    if (key === "1") {
      setMessages([]);
      return;
    }
    setLoadingMessages(true);
    // 获取会话内容
    try {
      const res = await GetSessionMessageList({
        app_name,
        session_id: key,
        page_index: 1,
      });

      const list: MessageItem[] = [];
      (res?.result?.model || []).forEach((item: any) => {
        list.push(
          {
            id: item.id + "_query",
            content: item.query,
            role: "user",
            status: "done",
          },
          {
            id: item.id + "_query",
            content: item.answer,
            role: "assistant",
            status: "done",
          }
        );
      });
      setMessages(list);
    } finally {
      setLoadingMessages(false);
    }
  };

  /**
   * 封装智能体
   */
  const [agent] = useXAgent<ResponseMessageItem>({
    request: async (message, { onError, onSuccess, onUpdate }) => {
      abortController.current = new AbortController();
      const signal = abortController.current.signal;
      try {
        setLoading(true);
        const response = await fetch(
          "https://design.shalu.com/api/ai/chat-message",
          {
            method: "POST",
            body: JSON.stringify(message),
            headers: {
              Authorization: localStorage.getItem("token_a") || "",
              "Content-Type": "application/json",
            },
            signal,
          }
        );

        // 判断当前是否流式返回
        if(response.headers.get('content-type')?.includes('text/event-stream')) {
          if (response.body) {
            for await (const chunk of XStream({
              readableStream: response.body,
            })) {
              const data = JSON.parse(chunk.data);
              if (data?.event === "message") {
                onUpdate(data);
              }
              if (data?.event === "message_end") {
                onSuccess(data);
              }
              if (data?.event === "message_error") {
                onError(data);
              }
              if (data?.event === "ping") {
                console.log(">>>> stream start <<<<");
              }
            }
          }
        } else {
          // 接口异常处理
          response.json().then(res => {
            if(res.code === 0 ) {
              onError?.(Error(res?.error || '请求失败'));
              cancel();
            }
          });
        }
      } catch (error) {
        // 判断是不是 abort 错误
        if (signal.aborted) {
          return;
        }
        onError(error as Error);
      } finally {
        setLoading(false);
      }
    },
  });

  /**
   * 发起请求
   * @param chat_query 对话内容
   */
  const onRequest = (chat_query: string) => {
    setConversationList((list) => {
      return list?.map((item) => {
        return {
          ...item,
          label: item.key === "1" ? chat_query : item.label,
        };
      });
    });
    agent.request(
      {
        app_name,
        chat_query,
        chat_name: activeConversation === "1" ? chat_query : undefined,
        conversation_id:
          activeConversation === "1" ? undefined : activeConversation,
      },
      {
        onSuccess: (data) => {
          onSuccess?.(data);
        },
        onUpdate: (data) => {
          onUpdate(data);
          // 更新会话相关信息
          if (activeConversation === "1") {
            setConversationList((list) => {
              return list?.map((item) => {
                return {
                  ...item,
                  // 更新当前会话id
                  key: item.key === "1" ? data.conversation_id : item.key,
                };
              });
            });
            setActiveConversation(data.conversation_id);
          }
        },
        onError: (error) => {
          console.log("error", error);
          onError?.(error);
        },
      }
    );
  };

  /**
   * 停止对话
   */
  const cancel = () => {
    abortController.current?.abort();
  };

  /**
   * 新增会话
   */
  const addConversation = () => {
    setMessages([]);
    setActiveConversation("1");
    // 还没产生对话时 直接清除当前对话
    if (!conversationList?.find((item) => item.key === "1")) {
      setConversationList([
        {
          ...defaultConversation,
        },
        ...(conversationList || []),
      ]);
    }
  };

  return {
    agent,
    loading,
    loadingMessages,
    loadingSession,
    cancel,
    messages,
    setMessages,
    conversationList,
    setConversationList,
    activeConversation,
    setActiveConversation,
    onRequest,
    addConversation,
    changeConversation,
  };
}