shalu-system-designer/apps/designer/src/components/ai/AIChat.tsx

import React, { useEffect, useMemo, useRef, useState } from "react";
import {
  PlusOutlined,
  CloseOutlined,
  FieldTimeOutlined,
  UserOutlined,
  SendOutlined,
  LoadingOutlined,
  EditOutlined,
  DeleteOutlined,
} from "@ant-design/icons";
import {
  Button,
  Tooltip,
  Input,
  Avatar,
  Form,
  Dropdown,
  MenuProps,
} from "antd";
import { Message } from "ai/react";
import MarkdownViewer from "./MarkdownViewer";
import { uuid } from "@repo/utils";
import { useLocalStorageState } from "ahooks";
import { useModel } from "umi";

// 用户消息
function UserMessage({ message }: { message: Message }) {
  return (
    <>
      <div className="rounded-8px bg-#eff6ff p-8px leading-1.5em">
        {message.content ?? ""}
      </div>

      <Avatar
        className="flex-shrink-0"
        size={32}
        icon={<UserOutlined />}
      ></Avatar>
    </>
  );
}

// 助手消息
function AssistantMessage({ message }: { message: Message }) {
  return (
    <>
      <Avatar
        size={32}
        className="flex-shrink-0"
        icon={
          <svg className="icon h-32px w-32px" aria-hidden="true">
            <use xlinkHref="#icon-AI1"></use>
          </svg>
        }
      ></Avatar>

      <div
        className="rounded-8px bg-#fff p-8px leading-1.5em overflow-x-auto"
        style={{ overflowX: "auto" }}
      >
        <MarkdownViewer content={message.content ?? ""} />
      </div>
    </>
  );
}

// 消息信息
function MessageInfo({ message }: { message: Message }) {
  return (
    <div
      key={message.id}
      className={`flex items-start space-x-2 mb-4 ${message.role === "user" ? "justify-end" : "justify-start"}`}
    >
      {message.role === "assistant" && (
        <AssistantMessage message={message}></AssistantMessage>
      )}
      {message.role === "user" && <UserMessage message={message}></UserMessage>}
    </div>
  );
}

// 消息列表
function MessageList({ messages }: { messages: Message[] }) {
  return messages.map((message) => (
    <MessageInfo key={message.id} message={message}></MessageInfo>
  ));
}

interface ChatHistoryItem {
  id: string;
  messages: Message[];
  createdAt: number;
  updatedAt: number;
  title: string;
}

interface DateGroups {
  today: ChatHistoryItem[];
  yesterday: ChatHistoryItem[];
  last7Days: ChatHistoryItem[];
  last30Days: ChatHistoryItem[];
  older: ChatHistoryItem[];
}

// 对历史记录进行分组
function groupDataByDate(data: ChatHistoryItem[]): DateGroups {
  const normalizeDate = (date: Date | string): Date => {
    const d = new Date(date);
    d.setHours(0, 0, 0, 0);
    return d;
  };

  const now = normalizeDate(new Date()); // 当前日期归一化

  const groups: DateGroups = {
    today: [],
    yesterday: [],
    last7Days: [],
    last30Days: [],
    older: [],
  };

  data.forEach((item: ChatHistoryItem) => {
    const itemDate = normalizeDate(new Date(item.updatedAt));
    const diffTime = now.getTime() - itemDate.getTime();
    const diffDays = Math.floor(diffTime / (1000 * 3600 * 24));

    if (diffDays === 0) {
      groups.today.push(item);
    } else if (diffDays === 1) {
      groups.yesterday.push(item);
    } else if (diffDays <= 6) {
      groups.last7Days.push(item);
    } else if (diffDays <= 29) {
      groups.last30Days.push(item);
    } else {
      groups.older.push(item);
    }
  });

  return groups;
}

export default function AIChat(props: { onClose?: () => void }) {
  const [focused, setFocused] = useState(false);
  const [chatStarted, setChatStarted] = useState(false);
  const [scrollHeight, setScrollHeight] = useState(0);
  const scrollAreaRef = useRef<HTMLDivElement>(null);
  const observer = useRef<ResizeObserver | null>(null);
  const [messages, setMessages] = useState<Message[]>([]);
  const [inputVal, setInputVal] = useState("");
  // 对话历史
  const [history, setHistory] = useLocalStorageState<ChatHistoryItem[]>(
    "chat-history",
    { defaultValue: [] }
  );
  const { loading, agent, cancel } = useModel("aiModel");

  const [chatId, setChatId] = React.useState(uuid());

  useEffect(() => {
    // 判断messages是否存在消息
    if (!messages.length) {
      return;
    }
    // 再判断历史记录是否存在该聊天,存在则更新记录
    if (history?.find((item) => item.id === chatId)) {
      setHistory(
        history?.map((item) => {
          if (item.id === chatId) {
            return {
              ...item,
              messages,
              updatedAt: Date.now(),
            };
          }
          return item;
        })
      );
    } else {
      setHistory((history) => {
        const title =
          messages.find((message) => message.role === "user")?.content ||
          "新对话";
        const newData = {
          id: chatId,
          messages,
          createdAt: Date.now(),
          updatedAt: Date.now(),
          title,
        };
        return history ? [newData, ...history] : [newData];
      });
    }
  }, [messages, chatId]);

  // 发起请求
  const onRequest = () => {
    agent.request(
      {
        // 应用名称
        app_name: "app1",
        // 会话内容
        chat_query: inputVal,
        // 会话名称 第一次
        chat_name: "新会话",
        // 会话id 后续会话带入
        // conversation_id: ;
      },
      {
        onSuccess: (msg) => {
          console.log("success", msg);
          setMessages((messages) => {
            const arr = [...messages];
            return arr;
          });
        },
        onError: (error) => {
          console.log("err:", error);
        },
        onUpdate: (msg) => {
          console.log("update", msg);
          setMessages((messages) => {
            const arr = [...messages];
            arr[messages.length - 1].content += msg.answer;
            arr[messages.length - 1].id = msg.message_id;
            setChatId(msg.conversation_id);
            return arr;
          });
        },
      }
    );
    setInputVal("");
  };

  // 处理提交
  const onSubmit = () => {
    if (inputVal.trim()) {
      if (!chatStarted) setChatStarted(true);
    }
    setMessages((arr) => {
      const index = arr.length;
      return [
        ...arr,
        { id: index + "", content: inputVal, role: "user" },
        {
          id: index + 1 + "",
          content: "",
          role: "assistant",
        },
      ];
    });
    onRequest();
  };

  // 开启新的对话
  const onNewChat = () => {
    setChatId(uuid());
    setChatStarted(false);
  };

  const stop = () => {
    cancel();
  };

  React.useEffect(() => {
    return () => {
      // 取消所有进行中的请求
      const controller = new AbortController();
      controller.abort();
    };
  }, []);

  React.useEffect(() => {
    const scrollElement = scrollAreaRef.current;
    if (scrollElement) {
      observer.current = new ResizeObserver((entries) => {
        for (let entry of entries) {
          if (entry.target === scrollElement) {
            setScrollHeight(entry.target.scrollHeight);
            entry.target.scrollTop = entry.target.scrollHeight;
          }
        }
      });

      observer.current.observe(scrollElement);

      return () => {
        observer.current?.disconnect();
      };
    }
  }, [messages]);

  // 添加一个 useEffect 来监听 chatId 的变化
  useEffect(() => {
    // 当 chatId 变化时，从历史记录中找到对应的消息
    const currentChat = history?.find((item) => item.id === chatId);
    if (currentChat) {
      observer.current?.disconnect();
      // 清空可滚动高度
      setScrollHeight(0);
      setTimeout(() => {
        setMessages(currentChat.messages);
        setChatStarted(true);
      }, 100);
    }
  }, [chatId]);

  const [hoverId, setHoverId] = useState<string | null>(null);

  const items = useMemo(() => {
    const hasCurrentChat = history?.find((item) => item.id === chatId);
    const groups = groupDataByDate(
      hasCurrentChat
        ? history || []
        : [
            {
              id: chatId,
              messages,
              createdAt: Date.now(),
              updatedAt: Date.now(),
              title:
                messages.find((message) => message.role === "user")?.content ||
                "新对话",
            },
            ...(history || []),
          ]
    );

    // 获取items
    const getItems = (list: ChatHistoryItem[]) => {
      return (list || []).map((item) => {
        return {
          key: item.id,
          label: (
            <div className="w-180px relative">
              <div className="w-full flex">
                <span
                  className="truncate"
                  style={{
                    overflow: "hidden",
                    whiteSpace: "nowrap",
                    textOverflow: "ellipsis",
                  }}
                >
                  <Tooltip title={item.title}>{item.title}</Tooltip>
                </span>
                {item.id === chatId ? (
                  <span className="text-12px color-#999 flex-shrink-0">
                    (当前)
                  </span>
                ) : null}
              </div>
              {/* TODO: 添加编辑删除按钮 */}
              {hoverId === item.id && (
                <div className="h-full w-50px text-right absolute right-0 top-0 bg-#fff">
                  <EditOutlined
                    onClick={(e) => {
                      e.stopPropagation();
                    }}
                  />
                  <DeleteOutlined />
                </div>
              )}
            </div>
          ),
          onClick: () => {
            if (item.id === chatId) return;

            setChatId(item.id);
          },
          onMouseOver: () => {
            setHoverId(item.id);
          },
        };
      });
    };

    const today = groups.today.length
      ? [
          {
            key: "today",
            label: "今天",
            disabled: true,
          },
          {
            key: "today-divider",
            type: "divider",
          },
          ...getItems(groups.today),
        ]
      : [];

    const yesterday = groups.yesterday.length
      ? [
          {
            key: "yesterday",
            label: "昨天",
            disabled: true,
          },
          {
            key: "yesterday-divider",
            type: "divider",
          },
          ...getItems(groups.yesterday),
        ]
      : [];

    const last7Days = groups.last7Days.length
      ? [
          {
            key: "last7Days",
            label: "最近7天",
            disabled: true,
          },
          {
            key: "last7Days-divider",
            type: "divider",
          },
          ...getItems(groups.last7Days),
        ]
      : [];

    const last30Days = groups.last30Days.length
      ? [
          {
            key: "last30Days",
            label: "最近30天",
            disabled: true,
          },
          {
            key: "last30Days-divider",
            type: "divider",
          },
          ...getItems(groups.last30Days),
        ]
      : [];

    const older = groups.older.length
      ? [
          {
            key: "older",
            label: "更早",
            disabled: true,
          },
          {
            key: "older-divider",
            type: "divider",
          },
          ...getItems(groups.older),
        ]
      : [];

    return [
      ...today,
      ...yesterday,
      ...last7Days,
      ...last30Days,
      ...older,
    ] as MenuProps["items"];
  }, [messages, chatId]);

  const handleInputChange = (str: string) => {
    setInputVal(str);
  };

  return (
    <div className="flex-1 h-full flex flex-col">
      <div className="chat-head w-full h-40px px-10px color-#333 flex items-center justify-between">
        <i className="iconfont icon-duihua"></i>
        <span>
          <Tooltip title="新建会话">
            <Button
              type="text"
              size="small"
              icon={<PlusOutlined />}
              onClick={onNewChat}
            ></Button>
          </Tooltip>
          <Dropdown
            menu={{ items }}
            trigger={["click"]}
            placement="bottomLeft"
            arrow
          >
            <Tooltip title="历史记录">
              <Button
                type="text"
                size="small"
                icon={<FieldTimeOutlined />}
              ></Button>
            </Tooltip>
          </Dropdown>
          <Button
            type="text"
            size="small"
            icon={<CloseOutlined />}
            onClick={() => props.onClose?.()}
          ></Button>
        </span>
      </div>

      <div
        className="chat-content flex-1 bg-#f5f5f5 px-10px overflow-y-auto mt-12px"
        ref={scrollAreaRef}
      >
        <div style={{ minHeight: `${scrollHeight}px` }}>
          {!chatStarted && (
            <>
              <div className="text-center pt-200px">
                <svg className="icon h-40px! w-40px!" aria-hidden="true">
                  <use xlinkHref="#icon-AI1"></use>
                </svg>
              </div>
              <h2 className="text-center">询问AI助手</h2>
              <p className="text-center">我是AI助手，有什么可以帮您的吗？</p>
            </>
          )}

          {chatStarted && (
            <div className="overflow-y-auto h-full">
              <MessageList messages={messages}></MessageList>
              <div className="flex justify-center items-center h-40px">
                {loading && (
                  <Button
                    type="primary"
                    icon={<LoadingOutlined />}
                    loading={loading}
                  >
                    思考中...
                  </Button>
                )}
              </div>
              {/* {error && (
                <div>
                  <div className="text-center">请求失败:{error.message}</div>
                  <div className="flex justify-center items-center h-40px">
                    <Button type="primary" onClick={() => reload()}>
                      重试
                    </Button>
                  </div>
                </div>
              )} */}
            </div>
          )}
        </div>
      </div>

      <div
        style={{
          borderColor: focused ? "#1890ff" : "#ddd",
        }}
        className="chat-foot bg-#f3f4f6 rounded-10px border border-solid border-1px m-10px"
      >
        <Form onFinish={onSubmit}>
          <Input.TextArea
            rows={3}
            autoSize={{ maxRows: 3, minRows: 3 }}
            placeholder="输入询问内容..."
            variant="borderless"
            onFocus={() => setFocused(true)}
            onBlur={() => setFocused(false)}
            value={inputVal}
            onChange={(e) => handleInputChange(e.target.value)}
            disabled={loading}
            onPressEnter={onSubmit}
          />
          <div className="float-right p-10px">
            {loading ? (
              <Tooltip title="停止生成">
                <Button
                  type="primary"
                  shape="circle"
                  icon={<i className="iconfont icon-stopcircle" />}
                  onClick={stop}
                ></Button>
              </Tooltip>
            ) : (
              <Button
                type="primary"
                icon={<SendOutlined />}
                disabled={!inputVal.trim()}
                htmlType="submit"
              >
                发送
              </Button>
            )}
          </div>
        </Form>
      </div>
    </div>
  );
}