chat.ts 15 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555
  1. import { create } from "zustand";
  2. import { persist } from "zustand/middleware";
  3. import { trimTopic } from "../utils";
  4. import Locale from "../locales";
  5. import { showToast } from "../components/ui-lib";
  6. import { ModelType } from "./config";
  7. import { createEmptyMask, Mask } from "./mask";
  8. import { StoreKey } from "../constant";
  9. import { api, RequestMessage } from "../client/api";
  10. import { ChatControllerPool } from "../client/controller";
  11. import { prettyObject } from "../utils/format";
  12. import { estimateTokenLength } from "../utils/token";
  13. export type ChatMessage = RequestMessage & {
  14. date: string;
  15. streaming?: boolean;
  16. isError?: boolean;
  17. id?: number;
  18. model?: ModelType;
  19. };
  20. export function createMessage(override: Partial<ChatMessage>): ChatMessage {
  21. return {
  22. id: Date.now(),
  23. date: new Date().toLocaleString(),
  24. role: "user",
  25. content: "",
  26. ...override,
  27. };
  28. }
  29. export interface ChatStat {
  30. tokenCount: number;
  31. wordCount: number;
  32. charCount: number;
  33. }
  34. export interface ChatSession {
  35. id: number;
  36. topic: string;
  37. memoryPrompt: string;
  38. messages: ChatMessage[];
  39. stat: ChatStat;
  40. lastUpdate: number;
  41. lastSummarizeIndex: number;
  42. clearContextIndex?: number;
  43. mask: Mask;
  44. }
  45. export const DEFAULT_TOPIC = Locale.Store.DefaultTopic;
  46. export const BOT_HELLO: ChatMessage = createMessage({
  47. role: "assistant",
  48. content: Locale.Store.BotHello,
  49. });
  50. function createEmptySession(): ChatSession {
  51. return {
  52. id: Date.now() + Math.random(),
  53. topic: DEFAULT_TOPIC,
  54. memoryPrompt: "",
  55. messages: [],
  56. stat: {
  57. tokenCount: 0,
  58. wordCount: 0,
  59. charCount: 0,
  60. },
  61. lastUpdate: Date.now(),
  62. lastSummarizeIndex: 0,
  63. mask: createEmptyMask(),
  64. };
  65. }
  66. interface ChatStore {
  67. sessions: ChatSession[];
  68. currentSessionIndex: number;
  69. globalId: number;
  70. clearSessions: () => void;
  71. moveSession: (from: number, to: number) => void;
  72. selectSession: (index: number) => void;
  73. newSession: (mask?: Mask) => void;
  74. deleteSession: (index: number) => void;
  75. currentSession: () => ChatSession;
  76. onNewMessage: (message: ChatMessage) => void;
  77. onUserInput: (content: string) => Promise<void>;
  78. summarizeSession: () => void;
  79. updateStat: (message: ChatMessage) => void;
  80. updateCurrentSession: (updater: (session: ChatSession) => void) => void;
  81. updateMessage: (
  82. sessionIndex: number,
  83. messageIndex: number,
  84. updater: (message?: ChatMessage) => void,
  85. ) => void;
  86. resetSession: () => void;
  87. getMessagesWithMemory: () => ChatMessage[];
  88. getMemoryPrompt: () => ChatMessage;
  89. clearAllData: () => void;
  90. }
  91. function countMessages(msgs: ChatMessage[]) {
  92. return msgs.reduce((pre, cur) => pre + estimateTokenLength(cur.content), 0);
  93. }
  94. export const useChatStore = create<ChatStore>()(
  95. persist(
  96. (set, get) => ({
  97. sessions: [createEmptySession()],
  98. currentSessionIndex: 0,
  99. globalId: 0,
  100. clearSessions() {
  101. set(() => ({
  102. sessions: [createEmptySession()],
  103. currentSessionIndex: 0,
  104. }));
  105. },
  106. selectSession(index: number) {
  107. set({
  108. currentSessionIndex: index,
  109. });
  110. },
  111. moveSession(from: number, to: number) {
  112. set((state) => {
  113. const { sessions, currentSessionIndex: oldIndex } = state;
  114. // move the session
  115. const newSessions = [...sessions];
  116. const session = newSessions[from];
  117. newSessions.splice(from, 1);
  118. newSessions.splice(to, 0, session);
  119. // modify current session id
  120. let newIndex = oldIndex === from ? to : oldIndex;
  121. if (oldIndex > from && oldIndex <= to) {
  122. newIndex -= 1;
  123. } else if (oldIndex < from && oldIndex >= to) {
  124. newIndex += 1;
  125. }
  126. return {
  127. currentSessionIndex: newIndex,
  128. sessions: newSessions,
  129. };
  130. });
  131. },
  132. newSession(mask) {
  133. const session = createEmptySession();
  134. set(() => ({ globalId: get().globalId + 1 }));
  135. session.id = get().globalId;
  136. if (mask) {
  137. session.mask = { ...mask };
  138. session.topic = mask.name;
  139. }
  140. set((state) => ({
  141. currentSessionIndex: 0,
  142. sessions: [session].concat(state.sessions),
  143. }));
  144. },
  145. deleteSession(index) {
  146. const deletingLastSession = get().sessions.length === 1;
  147. const deletedSession = get().sessions.at(index);
  148. if (!deletedSession) return;
  149. const sessions = get().sessions.slice();
  150. sessions.splice(index, 1);
  151. const currentIndex = get().currentSessionIndex;
  152. let nextIndex = Math.min(
  153. currentIndex - Number(index < currentIndex),
  154. sessions.length - 1,
  155. );
  156. if (deletingLastSession) {
  157. nextIndex = 0;
  158. sessions.push(createEmptySession());
  159. }
  160. // for undo delete action
  161. const restoreState = {
  162. currentSessionIndex: get().currentSessionIndex,
  163. sessions: get().sessions.slice(),
  164. };
  165. set(() => ({
  166. currentSessionIndex: nextIndex,
  167. sessions,
  168. }));
  169. showToast(
  170. Locale.Home.DeleteToast,
  171. {
  172. text: Locale.Home.Revert,
  173. onClick() {
  174. set(() => restoreState);
  175. },
  176. },
  177. 5000,
  178. );
  179. },
  180. currentSession() {
  181. let index = get().currentSessionIndex;
  182. const sessions = get().sessions;
  183. if (index < 0 || index >= sessions.length) {
  184. index = Math.min(sessions.length - 1, Math.max(0, index));
  185. set(() => ({ currentSessionIndex: index }));
  186. }
  187. const session = sessions[index];
  188. return session;
  189. },
  190. onNewMessage(message) {
  191. get().updateCurrentSession((session) => {
  192. session.lastUpdate = Date.now();
  193. });
  194. get().updateStat(message);
  195. get().summarizeSession();
  196. },
  197. async onUserInput(content) {
  198. const session = get().currentSession();
  199. const modelConfig = session.mask.modelConfig;
  200. const userMessage: ChatMessage = createMessage({
  201. role: "user",
  202. content,
  203. });
  204. const botMessage: ChatMessage = createMessage({
  205. role: "assistant",
  206. streaming: true,
  207. id: userMessage.id! + 1,
  208. model: modelConfig.model,
  209. });
  210. const systemInfo = createMessage({
  211. role: "system",
  212. content: `IMPORTANT: You are a virtual assistant powered by the ${
  213. modelConfig.model
  214. } model, now time is ${new Date().toLocaleString()}}`,
  215. id: botMessage.id! + 1,
  216. });
  217. // get recent messages
  218. const systemMessages = [];
  219. // if user define a mask with context prompts, wont send system info
  220. if (session.mask.context.length === 0) {
  221. systemMessages.push(systemInfo);
  222. }
  223. const recentMessages = get().getMessagesWithMemory();
  224. const sendMessages = systemMessages.concat(
  225. recentMessages.concat(userMessage),
  226. );
  227. const sessionIndex = get().currentSessionIndex;
  228. const messageIndex = get().currentSession().messages.length + 1;
  229. // save user's and bot's message
  230. get().updateCurrentSession((session) => {
  231. session.messages.push(userMessage);
  232. session.messages.push(botMessage);
  233. });
  234. // make request
  235. console.log("[User Input] ", sendMessages);
  236. api.llm.chat({
  237. messages: sendMessages,
  238. config: { ...modelConfig, stream: true },
  239. onUpdate(message) {
  240. botMessage.streaming = true;
  241. if (message) {
  242. botMessage.content = message;
  243. }
  244. set(() => ({}));
  245. },
  246. onFinish(message) {
  247. botMessage.streaming = false;
  248. if (message) {
  249. botMessage.content = message;
  250. get().onNewMessage(botMessage);
  251. }
  252. ChatControllerPool.remove(
  253. sessionIndex,
  254. botMessage.id ?? messageIndex,
  255. );
  256. set(() => ({}));
  257. },
  258. onError(error) {
  259. const isAborted = error.message.includes("aborted");
  260. botMessage.content =
  261. "\n\n" +
  262. prettyObject({
  263. error: true,
  264. message: error.message,
  265. });
  266. botMessage.streaming = false;
  267. userMessage.isError = !isAborted;
  268. botMessage.isError = !isAborted;
  269. set(() => ({}));
  270. ChatControllerPool.remove(
  271. sessionIndex,
  272. botMessage.id ?? messageIndex,
  273. );
  274. console.error("[Chat] failed ", error);
  275. },
  276. onController(controller) {
  277. // collect controller for stop/retry
  278. ChatControllerPool.addController(
  279. sessionIndex,
  280. botMessage.id ?? messageIndex,
  281. controller,
  282. );
  283. },
  284. });
  285. },
  286. getMemoryPrompt() {
  287. const session = get().currentSession();
  288. return {
  289. role: "system",
  290. content:
  291. session.memoryPrompt.length > 0
  292. ? Locale.Store.Prompt.History(session.memoryPrompt)
  293. : "",
  294. date: "",
  295. } as ChatMessage;
  296. },
  297. getMessagesWithMemory() {
  298. const session = get().currentSession();
  299. const modelConfig = session.mask.modelConfig;
  300. // wont send cleared context messages
  301. const clearedContextMessages = session.messages.slice(
  302. session.clearContextIndex ?? 0,
  303. );
  304. const messages = clearedContextMessages.filter((msg) => !msg.isError);
  305. const n = messages.length;
  306. const context = session.mask.context.slice();
  307. // long term memory
  308. if (
  309. modelConfig.sendMemory &&
  310. session.memoryPrompt &&
  311. session.memoryPrompt.length > 0
  312. ) {
  313. const memoryPrompt = get().getMemoryPrompt();
  314. context.push(memoryPrompt);
  315. }
  316. // get short term and unmemoried long term memory
  317. const shortTermMemoryMessageIndex = Math.max(
  318. 0,
  319. n - modelConfig.historyMessageCount,
  320. );
  321. const longTermMemoryMessageIndex = session.lastSummarizeIndex;
  322. const mostRecentIndex = Math.max(
  323. shortTermMemoryMessageIndex,
  324. longTermMemoryMessageIndex,
  325. );
  326. const threshold = modelConfig.compressMessageLengthThreshold * 2;
  327. // get recent messages as many as possible
  328. const reversedRecentMessages = [];
  329. for (
  330. let i = n - 1, count = 0;
  331. i >= mostRecentIndex && count < threshold;
  332. i -= 1
  333. ) {
  334. const msg = messages[i];
  335. if (!msg || msg.isError) continue;
  336. count += msg.content.length;
  337. reversedRecentMessages.push(msg);
  338. }
  339. // concat
  340. const recentMessages = context.concat(reversedRecentMessages.reverse());
  341. return recentMessages;
  342. },
  343. updateMessage(
  344. sessionIndex: number,
  345. messageIndex: number,
  346. updater: (message?: ChatMessage) => void,
  347. ) {
  348. const sessions = get().sessions;
  349. const session = sessions.at(sessionIndex);
  350. const messages = session?.messages;
  351. updater(messages?.at(messageIndex));
  352. set(() => ({ sessions }));
  353. },
  354. resetSession() {
  355. get().updateCurrentSession((session) => {
  356. session.messages = [];
  357. session.memoryPrompt = "";
  358. });
  359. },
  360. summarizeSession() {
  361. const session = get().currentSession();
  362. // remove error messages if any
  363. const messages = session.messages;
  364. // should summarize topic after chating more than 50 words
  365. const SUMMARIZE_MIN_LEN = 50;
  366. if (
  367. session.topic === DEFAULT_TOPIC &&
  368. countMessages(messages) >= SUMMARIZE_MIN_LEN
  369. ) {
  370. const topicMessages = messages.concat(
  371. createMessage({
  372. role: "user",
  373. content: Locale.Store.Prompt.Topic,
  374. }),
  375. );
  376. api.llm.chat({
  377. messages: topicMessages,
  378. config: {
  379. model: "gpt-3.5-turbo",
  380. },
  381. onFinish(message) {
  382. get().updateCurrentSession(
  383. (session) =>
  384. (session.topic =
  385. message.length > 0 ? trimTopic(message) : DEFAULT_TOPIC),
  386. );
  387. },
  388. });
  389. }
  390. const modelConfig = session.mask.modelConfig;
  391. const summarizeIndex = Math.max(
  392. session.lastSummarizeIndex,
  393. session.clearContextIndex ?? 0,
  394. );
  395. let toBeSummarizedMsgs = messages
  396. .filter((msg) => !msg.isError)
  397. .slice(summarizeIndex);
  398. const historyMsgLength = countMessages(toBeSummarizedMsgs);
  399. if (historyMsgLength > modelConfig?.max_tokens ?? 4000) {
  400. const n = toBeSummarizedMsgs.length;
  401. toBeSummarizedMsgs = toBeSummarizedMsgs.slice(
  402. Math.max(0, n - modelConfig.historyMessageCount),
  403. );
  404. }
  405. // add memory prompt
  406. toBeSummarizedMsgs.unshift(get().getMemoryPrompt());
  407. const lastSummarizeIndex = session.messages.length;
  408. console.log(
  409. "[Chat History] ",
  410. toBeSummarizedMsgs,
  411. historyMsgLength,
  412. modelConfig.compressMessageLengthThreshold,
  413. );
  414. if (
  415. historyMsgLength > modelConfig.compressMessageLengthThreshold &&
  416. modelConfig.sendMemory
  417. ) {
  418. api.llm.chat({
  419. messages: toBeSummarizedMsgs.concat({
  420. role: "system",
  421. content: Locale.Store.Prompt.Summarize,
  422. date: "",
  423. }),
  424. config: { ...modelConfig, stream: true },
  425. onUpdate(message) {
  426. session.memoryPrompt = message;
  427. },
  428. onFinish(message) {
  429. console.log("[Memory] ", message);
  430. session.lastSummarizeIndex = lastSummarizeIndex;
  431. },
  432. onError(err) {
  433. console.error("[Summarize] ", err);
  434. },
  435. });
  436. }
  437. },
  438. updateStat(message) {
  439. get().updateCurrentSession((session) => {
  440. session.stat.charCount += message.content.length;
  441. // TODO: should update chat count and word count
  442. });
  443. },
  444. updateCurrentSession(updater) {
  445. const sessions = get().sessions;
  446. const index = get().currentSessionIndex;
  447. updater(sessions[index]);
  448. set(() => ({ sessions }));
  449. },
  450. clearAllData() {
  451. localStorage.clear();
  452. location.reload();
  453. },
  454. }),
  455. {
  456. name: StoreKey.Chat,
  457. version: 2,
  458. migrate(persistedState, version) {
  459. const state = persistedState as any;
  460. const newState = JSON.parse(JSON.stringify(state)) as ChatStore;
  461. if (version < 2) {
  462. newState.globalId = 0;
  463. newState.sessions = [];
  464. const oldSessions = state.sessions;
  465. for (const oldSession of oldSessions) {
  466. const newSession = createEmptySession();
  467. newSession.topic = oldSession.topic;
  468. newSession.messages = [...oldSession.messages];
  469. newSession.mask.modelConfig.sendMemory = true;
  470. newSession.mask.modelConfig.historyMessageCount = 4;
  471. newSession.mask.modelConfig.compressMessageLengthThreshold = 1000;
  472. newState.sessions.push(newSession);
  473. }
  474. }
  475. return newState;
  476. },
  477. },
  478. ),
  479. );