@lobehub/chat 1.0.0

package/src/store/session/slices/chat/action.ts

@@ -0,0 +1,210 @@
+import { StateCreator } from 'zustand/vanilla';
+
+import { fetchChatModel } from '@/services/chatModel';
+import { SessionStore, agentSelectors, chatSelectors, sessionSelectors } from '@/store/session';
+import { ChatMessage } from '@/types/chatMessage';
+import { FetchSSEOptions, fetchSSE } from '@/utils/fetch';
+import { nanoid } from '@/utils/uuid';
+
+import { MessageDispatch, messagesReducer } from './messageReducer';
+
+const LOADING_FLAT = '...';
+
+/**
+ * 聊天操作
+ */
+export interface ChatAction {
+  /**
+   * 清除消息
+   */
+  clearMessage: () => void;
+  /**
+   * 创建或发送消息
+   * @param text - 消息文本
+   */
+  createOrSendMsg: (text: string) => Promise<void>;
+  /**
+   * 删除消息
+   * @param id - 消息 ID
+   */
+  deleteMessage: (id: string) => void;
+  /**
+   * 分发消息
+   * @param payload - 消息分发参数
+   */
+  dispatchMessage: (payload: MessageDispatch) => void;
+  /**
+   * 生成消息
+   * @param messages - 聊天消息数组
+   * @param options - 获取 SSE 选项
+   */
+  generateMessage: (messages: ChatMessage[], options: FetchSSEOptions) => Promise<void>;
+  /**
+   * 处理消息编辑
+   * @param messageId - 消息 ID，可选
+   */
+  handleMessageEditing: (messageId: string | undefined) => void;
+  /**
+   * 实际获取 AI 响应
+   *
+   * @param messages - 聊天消息数组
+   * @param parentId - 父消息 ID，可选
+   */
+  realFetchAIResponse: (messages: ChatMessage[], parentId?: string) => Promise<void>;
+  /**
+   * 重新发送消息
+   * @param id - 消息 ID
+   */
+  resendMessage: (id: string) => Promise<void>;
+  /**
+   * 发送消息
+   * @param text - 消息文本
+   */
+  sendMessage: (text: string) => Promise<void>;
+}
+
+export const createChatSlice: StateCreator<
+  SessionStore,
+  [['zustand/devtools', never]],
+  [],
+  ChatAction
+> = (set, get) => ({
+  clearMessage: () => {
+    get().dispatchMessage({ type: 'resetMessages' });
+  },
+
+  createOrSendMsg: async (message) => {
+    if (!message) return;
+
+    const { sendMessage, createSession } = get();
+    const session = sessionSelectors.currentSession(get());
+
+    if (!session) {
+      await createSession();
+    }
+
+    sendMessage(message);
+  },
+
+  deleteMessage: (id) => {
+    get().dispatchMessage({ id, type: 'deleteMessage' });
+  },
+
+  dispatchMessage: (payload) => {
+    const { activeId } = get();
+    const session = sessionSelectors.currentSession(get());
+    if (!activeId || !session) return;
+
+    const chats = messagesReducer(session.chats, payload);
+
+    get().dispatchSession({ chats, id: activeId, type: 'updateSessionChat' });
+  },
+
+  generateMessage: async (messages, options) => {
+    set({ chatLoading: true });
+    const config = agentSelectors.currentAgentConfigSafe(get());
+
+    const fetcher = () => fetchChatModel({ messages, model: config.model, ...config.params });
+
+    await fetchSSE(fetcher, options);
+
+    set({ chatLoading: false });
+  },
+
+  handleMessageEditing: (messageId) => {
+    set({ editingMessageId: messageId });
+  },
+
+  realFetchAIResponse: async (messages: ChatMessage[], parentId?: string) => {
+    const { dispatchMessage, generateMessage } = get();
+
+    // 添加 systemRole
+    const { systemRole } = agentSelectors.currentAgentConfigSafe(get());
+    if (systemRole) {
+      messages.unshift({ content: systemRole, role: 'system' } as ChatMessage);
+    }
+
+    // 再添加一个空的信息用于放置 ai 响应，注意顺序不能反
+    // 因为如果顺序反了，messages 中将包含新增的 ai message
+    const assistantId = nanoid();
+    const userId = parentId ?? nanoid();
+
+    dispatchMessage({
+      id: assistantId,
+      message: LOADING_FLAT,
+      parentId: userId,
+      role: 'assistant',
+      type: 'addMessage',
+    });
+
+    let output = '';
+    // 生成 ai message
+    await generateMessage(messages, {
+      onErrorHandle: (error) => {
+        dispatchMessage({ id: assistantId, key: 'error', type: 'updateMessage', value: error });
+      },
+      onMessageHandle: (text) => {
+        output += text;
+
+        dispatchMessage({
+          id: assistantId,
+          key: 'content',
+          type: 'updateMessage',
+          value: output,
+        });
+
+        // 滚动到最后一条消息
+        const item = document.querySelector('#for-loading');
+        if (!item) return;
+
+        item.scrollIntoView({ behavior: 'smooth' });
+      },
+    });
+  },
+
+  resendMessage: async (messageId) => {
+    const session = sessionSelectors.currentSession(get());
+
+    if (!session) return;
+
+    // 1. 构造所有相关的历史记录
+    const chats = chatSelectors.currentChats(get());
+
+    const currentIndex = chats.findIndex((c) => c.id === messageId);
+
+    const histories = chats
+      .slice(0, currentIndex + 1)
+      // 如果点击重新发送的 message 其 role 是 assistant，那么需要移除
+      // 如果点击重新发送的 message 其 role 是 user，则不需要移除
+      .filter((c) => !(c.role === 'assistant' && c.id === messageId));
+
+    if (histories.length <= 0) return;
+
+    const { realFetchAIResponse } = get();
+
+    const latestMsg = histories.filter((s) => s.role === 'user').at(-1);
+
+    if (!latestMsg) return;
+
+    await realFetchAIResponse(histories, latestMsg.id);
+  },
+
+  sendMessage: async (message) => {
+    const { dispatchMessage, realFetchAIResponse, autocompleteSessionAgentMeta } = get();
+    const session = sessionSelectors.currentSession(get());
+    if (!session || !message) return;
+
+    const userId = nanoid();
+    dispatchMessage({ id: userId, message, role: 'user', type: 'addMessage' });
+
+    // 先拿到当前的 messages
+    const messages = chatSelectors.currentChats(get());
+
+    await realFetchAIResponse(messages);
+
+    const chats = chatSelectors.currentChats(get());
+    if (chats.length >= 4) {
+      autocompleteSessionAgentMeta(session.id);
+    }
+  },
+});

package/src/store/session/slices/chat/index.ts

@@ -0,0 +1,3 @@
+export * from './action';
+export * from './initialState';
+export * from './selectors';

package/src/store/session/slices/chat/initialState.ts

@@ -0,0 +1,12 @@
+export interface ChatState {
+  chatLoading: boolean;
+  editingMessageId?: string;
+}
+
+export const initialChatState: ChatState = {
+  chatLoading: false,
+
+  // activeId: null,
+  // searchKeywords: '',
+  //
+};

package/src/store/session/slices/chat/messageReducer.test.ts

@@ -0,0 +1,70 @@
+test('placeholder', () => {});
+// describe('messagesReducer', () => {
+//   let initialState: ChatMessage[];
+//
+//   beforeEach(() => {
+//     initialState = [
+//       { role: 'user', content: 'Hello!' },
+//       { role: 'assistant', content: 'Hi there!' },
+//     ];
+//   });
+//
+//   it('should add a message', () => {
+//     const newMessage: ChatMessage = { role: 'user', content: 'How are you?' };
+//     const action: MessageDispatch = { type: 'addMessage', message: newMessage };
+//     const newState = messagesReducer(initialState, action);
+//     expect(newState).toEqual([...initialState, newMessage]);
+//   });
+//
+//   it('should delete a message', () => {
+//     const action: MessageDispatch = { type: 'deleteMessage', index: 1 };
+//     const newState = messagesReducer(initialState, action);
+//     expect(newState).toEqual([{ role: 'user', content: 'Hello!' }]);
+//   });
+//
+//   it('should update a message', () => {
+//     const action: MessageDispatch = { type: 'updateMessage', index: 1, message: 'I am fine!' };
+//     const newState = messagesReducer(initialState, action);
+//     expect(newState).toEqual([
+//       { role: 'user', content: 'Hello!' },
+//       { role: 'assistant', content: 'I am fine!' },
+//     ]);
+//   });
+//
+//   it('should add a user message', () => {
+//     const action: MessageDispatch = { type: 'addUserMessage', message: 'Goodbye!' };
+//     const newState = messagesReducer(initialState, action);
+//     expect(newState).toEqual([
+//       { role: 'user', content: 'Hello!' },
+//       { role: 'assistant', content: 'Hi there!' },
+//       { role: 'user', content: 'Goodbye!' },
+//     ]);
+//   });
+//
+//   it('should set error message correctly', () => {
+//     const action: MessageDispatch = {
+//       type: 'setErrorMessage',
+//       error: { message: 'Not Found', status: 404, type: 'chatbot' },
+//       index: 0,
+//     };
+//     const newState = messagesReducer(initialState, action);
+//     expect(newState).toEqual([
+//       {
+//         role: 'user',
+//         content: 'Hello!',
+//         error: { message: 'Not Found', status: 404, type: 'chatbot' },
+//       },
+//       { role: 'assistant', content: 'Hi there!' },
+//     ]);
+//   });
+//
+//   it('should update the latest bot message', () => {
+//     const responseStream = ['I', ' am', ' a', ' bot.'];
+//     const action: MessageDispatch = { type: 'updateLatestBotMessage', responseStream };
+//     const newState = messagesReducer(initialState, action);
+//     expect(newState).toEqual([
+//       { role: 'user', content: 'Hello!' },
+//       { role: 'assistant', content: 'I am a bot.' },
+//     ]);
+//   });
+// });

package/src/store/session/slices/chat/messageReducer.ts

@@ -0,0 +1,84 @@
+import { produce } from 'immer';
+
+import { ChatMessage, ChatMessageMap } from '@/types/chatMessage';
+import { LLMRoleType } from '@/types/llm';
+import { MetaData } from '@/types/meta';
+import { nanoid } from '@/utils/uuid';
+
+interface AddMessage {
+  id?: string;
+  message: string;
+  meta?: MetaData;
+  parentId?: string;
+  quotaId?: string;
+  role: LLMRoleType;
+  type: 'addMessage';
+}
+
+interface DeleteMessage {
+  id: string;
+  type: 'deleteMessage';
+}
+
+interface ResetMessages {
+  type: 'resetMessages';
+}
+
+interface UpdateMessage {
+  id: string;
+  key: keyof ChatMessage;
+  type: 'updateMessage';
+  value: ChatMessage[keyof ChatMessage];
+}
+
+export type MessageDispatch = AddMessage | DeleteMessage | ResetMessages | UpdateMessage;
+
+export const messagesReducer = (
+  state: ChatMessageMap,
+  payload: MessageDispatch,
+): ChatMessageMap => {
+  switch (payload.type) {
+    case 'addMessage': {
+      return produce(state, (draftState) => {
+        const mid = payload.id || nanoid();
+
+        draftState[mid] = {
+          content: payload.message,
+          createAt: Date.now(),
+          id: mid,
+          meta: payload.meta || {},
+          parentId: payload.parentId,
+          quotaId: payload.quotaId,
+          role: payload.role,
+          updateAt: Date.now(),
+        };
+      });
+    }
+
+    case 'deleteMessage': {
+      return produce(state, (draftState) => {
+        delete draftState[payload.id];
+      });
+    }
+
+    case 'updateMessage': {
+      return produce(state, (draftState) => {
+        const { id, key, value } = payload;
+        const message = draftState[id];
+        if (!message) return;
+
+        // @ts-ignore
+        message[key] = value;
+        message.updateAt = Date.now();
+      });
+    }
+
+    case 'resetMessages': {
+      return {};
+    }
+
+    default: {
+      throw new Error('暂未实现的 type，请检查 reducer');
+    }
+  }
+};

package/src/store/session/slices/chat/selectors.ts

@@ -0,0 +1,83 @@
+import { encode } from 'gpt-tokenizer';
+
+import { agentSelectors } from '@/store/session';
+import { ChatMessage } from '@/types/chatMessage';
+
+import type { SessionStore } from '../../store';
+import { sessionSelectors } from '../session';
+
+// 展示在聊天框中的消息
+const currentChats = (s: SessionStore): ChatMessage[] => {
+  const session = sessionSelectors.currentSession(s);
+  if (!session) return [];
+
+  const basic = Object.values<ChatMessage>(session.chats)
+    // 首先按照时间顺序排序，越早的在越前面
+    .sort((pre, next) => pre.createAt - next.createAt)
+    // 过滤掉已归档的消息，归档消息不应该出现在聊天框中
+    .filter((m) => !m.archive)
+    // 映射头像关系
+    .map((m) => {
+      return {
+        ...m,
+        meta:
+          m.role === 'assistant'
+            ? {
+                avatar: agentSelectors.currentAgentAvatar(s),
+                title: session.meta.title,
+              }
+            : m.meta,
+      };
+    });
+
+  const finalList: ChatMessage[] = [];
+
+  const addItem = (item: ChatMessage) => {
+    const isExist = finalList.findIndex((i) => item.id === i.id) > -1;
+    if (!isExist) {
+      finalList.push(item);
+    }
+  };
+
+  for (const item of basic) {
+    // 先判存在与否，不存在就加入
+    addItem(item);
+
+    for (const another of basic) {
+      if (another.parentId === item.id) {
+        addItem(another);
+      }
+    }
+  }
+
+  return finalList;
+};
+
+const systemRoleSel = (s: SessionStore): string => {
+  const config = agentSelectors.currentAgentConfigSafe(s);
+
+  return config.systemRole;
+};
+
+const totalTokens = (s: SessionStore): number[] => {
+  const chats = currentChats(s);
+  return encode(chats.map((m) => m.content).join(''));
+};
+
+const systemRoleTokens = (s: SessionStore): number[] => {
+  const systemRole = systemRoleSel(s);
+
+  return encode(systemRole || '');
+};
+
+const totalTokenCount = (s: SessionStore) => totalTokens(s).length;
+
+const systemRoleTokenCount = (s: SessionStore) => systemRoleTokens(s).length;
+
+export const chatSelectors = {
+  currentChats,
+  systemRoleTokenCount,
+  systemRoleTokens,
+  totalTokenCount,
+  totalTokens,
+};

package/src/store/session/slices/session/action.ts

@@ -0,0 +1,118 @@
+import Router from 'next/router';
+import { StateCreator } from 'zustand/vanilla';
+
+import { SessionStore } from '@/store/session';
+import { LanguageModel } from '@/types/llm';
+import { LobeAgentSession, LobeSessionType } from '@/types/session';
+import { uuid } from '@/utils/uuid';
+
+import { SessionDispatch, sessionsReducer } from './reducers/session';
+
+export interface SessionAction {
+  activeSession: (sessionId: string) => void;
+  /**
+   * @title 添加会话
+   * @param session - 会话信息
+   * @returns void
+   */
+  createSession: () => Promise<string>;
+
+  /**
+   * 分发聊天记录
+   * @param payload - 聊天记录
+   */
+  dispatchSession: (payload: SessionDispatch) => void;
+  /**
+   * @title 删除会话
+   * @param index - 会话索引
+   * @returns void
+   */
+  removeSession: (sessionId: string) => void;
+
+  /**
+   * @title 切换会话
+   * @param sessionId - 会话索引
+   * @returns void
+   */
+  switchSession: (sessionId?: string | 'new') => Promise<void>;
+
+  /**
+   * 生成压缩后的消息
+   * @returns 压缩后的消息
+   */
+  // genShareUrl: () => string;
+}
+
+export const createSessionSlice: StateCreator<
+  SessionStore,
+  [['zustand/devtools', never]],
+  [],
+  SessionAction
+> = (set, get) => ({
+  activeSession: (sessionId) => {
+    set({ activeId: sessionId });
+  },
+
+  createSession: async () => {
+    const { dispatchSession, switchSession } = get();
+
+    const timestamp = Date.now();
+
+    const newSession: LobeAgentSession = {
+      chats: {},
+      config: {
+        model: LanguageModel.GPT3_5,
+        params: {
+          temperature: 0.6,
+        },
+        systemRole: '',
+      },
+      createAt: timestamp,
+      id: uuid(),
+      meta: {},
+      type: LobeSessionType.Agent,
+      updateAt: timestamp,
+    };
+
+    dispatchSession({ session: newSession, type: 'addSession' });
+
+    await switchSession(newSession.id);
+
+    return newSession.id;
+  },
+
+  dispatchSession: (payload) => {
+    const { type, ...res } = payload;
+    set({ sessions: sessionsReducer(get().sessions, payload) }, false, {
+      payload: res,
+      type: `dispatchChat/${type}`,
+    });
+  },
+
+  removeSession: (sessionId) => {
+    get().dispatchSession({ id: sessionId, type: 'removeSession' });
+
+    if (sessionId === get().activeId) {
+      Router.push('/');
+    }
+  },
+
+  switchSession: async (sessionId) => {
+    if (get().activeId === sessionId) return;
+
+    if (sessionId) {
+      get().activeSession(sessionId);
+    }
+
+    // 新会话
+    await Router.push(`/chat/${sessionId}`);
+  },
+
+  // genShareUrl: () => {
+  //   const session = sessionSelectors.currentSession(get());
+  //   if (!session) return '';
+  //
+  //   const agent = session.config;
+  //   return genShareMessagesUrl(session.chats, agent.systemRole);
+  // },
+});

package/src/store/session/slices/session/index.ts

@@ -0,0 +1,3 @@
+export * from './action';
+export * from './initialState';
+export * from './selectors';

package/src/store/session/slices/session/initialState.ts

@@ -0,0 +1,31 @@
+import { LobeAgentSession, LobeSessionType } from '@/types/session';
+
+import { initialLobeAgentConfig } from '../agentConfig';
+
+export interface SessionState {
+  /**
+   * @title 当前活动的会话
+   * @description 当前正在编辑或查看的会话
+   * @default null
+   */
+  activeId: string | null;
+  searchKeywords: string;
+  sessions: Record<string, LobeAgentSession>;
+}
+
+export const initLobeSession: LobeAgentSession = {
+  chats: {},
+  config: initialLobeAgentConfig,
+  createAt: Date.now(),
+  id: '',
+  meta: {},
+  type: LobeSessionType.Agent,
+  updateAt: Date.now(),
+};
+
+export const initialSessionState: SessionState = {
+  activeId: null,
+
+  searchKeywords: '',
+  sessions: {},
+};