@ant-design/x/lib/use-x-chat/index.js

Craft AI-driven interfaces effortlessly

"use strict";

var _interopRequireDefault = require("@babel/runtime/helpers/interopRequireDefault").default;
Object.defineProperty(exports, "__esModule", {
  value: true
});
exports.default = useXChat;
var _rcUtil = require("rc-util");
var _react = _interopRequireDefault(require("react"));
var _useSyncState = _interopRequireDefault(require("./useSyncState"));
function toArray(item) {
  return Array.isArray(item) ? item : [item];
}
function useXChat(config) {
  const {
    defaultMessages,
    agent,
    requestFallback,
    requestPlaceholder,
    parser,
    transformMessage,
    transformStream,
    resolveAbortController
  } = config;

  // ========================= Agent Messages =========================
  const idRef = _react.default.useRef(0);
  const [messages, setMessages, getMessages] = (0, _useSyncState.default)(() => (defaultMessages || []).map((info, index) => ({
    id: `default_${index}`,
    status: 'local',
    ...info
  })));
  const createMessage = (message, status) => {
    const msg = {
      id: `msg_${idRef.current}`,
      message,
      status
    };
    idRef.current += 1;
    return msg;
  };

  // ========================= BubbleMessages =========================
  const parsedMessages = _react.default.useMemo(() => {
    const list = [];
    messages.forEach(agentMsg => {
      const rawParsedMsg = parser ? parser(agentMsg.message) : agentMsg.message;
      const bubbleMsgs = toArray(rawParsedMsg);
      bubbleMsgs.forEach((bubbleMsg, bubbleMsgIndex) => {
        let key = agentMsg.id;
        if (bubbleMsgs.length > 1) {
          key = `${key}_${bubbleMsgIndex}`;
        }
        list.push({
          id: key,
          message: bubbleMsg,
          status: agentMsg.status
        });
      });
    });
    return list;
  }, [messages]);

  // ============================ Request =============================
  const getFilteredMessages = msgs => msgs.filter(info => info.status !== 'loading' && info.status !== 'error').map(info => info.message);

  // For agent to use. Will filter out loading and error message
  const getRequestMessages = () => getFilteredMessages(getMessages());
  const getTransformMessage = params => {
    const {
      chunk,
      chunks,
      originMessage
    } = params;
    if (typeof transformMessage === 'function') {
      return transformMessage(params);
    }

    // Compatible bug fixes, onSuccess(output) => onSuccess(output[]),before v1.1.1

    if (chunk) {
      return chunk;
    }
    if (Array.isArray(chunks)) {
      const chunk = chunks?.length > 0 ? chunks?.[chunks?.length - 1] : undefined;
      return originMessage ? originMessage : chunk;
    }
    return chunks;
  };
  const onRequest = (0, _rcUtil.useEvent)(requestParams => {
    if (!agent) throw new Error('The agent parameter is required when using the onRequest method in an agent generated by useXAgent.');
    let loadingMsgId = null;
    let message;
    let otherRequestParams = {};
    if (requestParams && typeof requestParams === 'object' && 'message' in requestParams) {
      const {
        message: requestParamsMessage,
        ...other
      } = requestParams;
      message = requestParamsMessage;
      otherRequestParams = other;
    } else {
      message = requestParams;
    }
    // Add placeholder message
    setMessages(ori => {
      let nextMessages = [...ori, createMessage(message, 'local')];
      if (requestPlaceholder) {
        let placeholderMsg;
        if (typeof requestPlaceholder === 'function') {
          // typescript has bug that not get real return type when use `typeof function` check
          placeholderMsg = requestPlaceholder(message, {
            messages: getFilteredMessages(nextMessages)
          });
        } else {
          placeholderMsg = requestPlaceholder;
        }
        const loadingMsg = createMessage(placeholderMsg, 'loading');
        loadingMsgId = loadingMsg.id;
        nextMessages = [...nextMessages, loadingMsg];
      }
      return nextMessages;
    });

    // Request
    let updatingMsgId = null;
    const updateMessage = (status, chunk, chunks) => {
      let msg = getMessages().find(info => info.id === updatingMsgId);
      if (!msg) {
        // Create if not exist
        const transformData = getTransformMessage({
          chunk,
          status,
          chunks
        });
        msg = createMessage(transformData, status);
        setMessages(ori => {
          const oriWithoutPending = ori.filter(info => info.id !== loadingMsgId);
          return [...oriWithoutPending, msg];
        });
        updatingMsgId = msg.id;
      } else {
        // Update directly
        setMessages(ori => {
          return ori.map(info => {
            if (info.id === updatingMsgId) {
              const transformData = getTransformMessage({
                originMessage: info.message,
                chunk,
                chunks,
                status
              });
              return {
                ...info,
                message: transformData,
                status
              };
            }
            return info;
          });
        });
      }
      return msg;
    };
    agent.request({
      message,
      messages: getRequestMessages(),
      ...otherRequestParams
    }, {
      onUpdate: chunk => {
        updateMessage('loading', chunk, []);
      },
      onSuccess: chunks => {
        updateMessage('success', undefined, chunks);
      },
      onError: async error => {
        if (requestFallback) {
          let fallbackMsg;
          // Update as error
          if (typeof requestFallback === 'function') {
            // typescript has bug that not get real return type when use `typeof function` check
            fallbackMsg = await requestFallback(message, {
              error,
              messages: getRequestMessages()
            });
          } else {
            fallbackMsg = requestFallback;
          }
          setMessages(ori => [...ori.filter(info => info.id !== loadingMsgId && info.id !== updatingMsgId), createMessage(fallbackMsg, 'error')]);
        } else {
          // Remove directly
          setMessages(ori => {
            return ori.filter(info => info.id !== loadingMsgId && info.id !== updatingMsgId);
          });
        }
      },
      onStream: controller => {
        resolveAbortController?.(controller);
      }
    }, transformStream);
  });
  return {
    onRequest,
    messages,
    parsedMessages,
    setMessages
  };
}