shalu-marketplace/src/models/aiModel.ts

import { useXAgent, XStream } from "@ant-design/x";
import { useRef, useState } from "react";

type MessageItem = {
  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;
};

export default function aiModel() {
  const [loading, setLoading] = useState(false);
  const abortController = useRef<AbortController | null>(null);

  // 封装智能体
  const [agent] = useXAgent<MessageItem>({
    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.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("start");
            }
          }
        }
      } catch (error) {
        // 判断是不是 abort 错误
        if (signal.aborted) {
          return;
        }
        onError(error as Error);
      } finally {
        setLoading(false);
      }
    },
  });

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

  return {
    agent,
    loading,
    setLoading,
    cancel,
  };
}