feat: receive message for openai #48

Signed-off-by: seven <[email protected]>

src/electron/chatBotApi.ts

@@ -53,24 +53,44 @@ const chatBotApi = {
     assistantId,
     threadId,
     question,
+    mainWindow,
   }: {
     openai: OpenAI;
     assistantId: string;
     threadId: string;
     question: string;
+    mainWindow: Electron.BrowserWindow;
   }) => {
     await openai.beta.threads.messages.create(threadId, { role: 'user', content: question });
 
     openai.beta.threads.runs
       .stream(threadId, { assistant_id: assistantId })
-      .on('textCreated', text => console.log('textCreated, text:', text))
-      .on('textDelta', (textDelta, snapshot) =>
-        console.log('textDelta, textDelta:', JSON.stringify({ textDelta, snapshot })),
+      .on('messageCreated', message =>
+        mainWindow.webContents.send('chat-bot-api-message-delta', {
+          msgEvent: 'messageCreated',
+          message,
+        }),
+      )
+      .on('messageDelta', (delta, snapshot) => {
+        console.log('messageDelta, delta:', delta, 'snapshot:', snapshot);
+        mainWindow.webContents.send('chat-bot-api-message-delta', {
+          msgEvent: 'messageDelta',
+          delta,
+        });
+      })
+      .on('messageDone', message =>
+        mainWindow.webContents.send('chat-bot-api-message-delta', {
+          msgEvent: 'messageDone',
+          message,
+        }),
       );
   },
 };
 
-const registerChatBotApiListener = (ipcMain: Electron.IpcMain) => {
+const registerChatBotApiListener = (
+  ipcMain: Electron.IpcMain,
+  mainWindow: Electron.BrowserWindow,
+) => {
   let openai: OpenAI;
 
   ipcMain.handle(
@@ -79,7 +99,6 @@ const registerChatBotApiListener = (ipcMain: Electron.IpcMain) => {
       _,
       { method, question, apiKey, prompt, model, assistantId, threadId }: ChatBotApiInput,
     ) => {
-      console.log(`chatBotApi method: ${method}`);
       if (method === ChatBotApiMethods.INITIALIZE) {
         openai = createOpenaiClient({ apiKey });
 
@@ -95,7 +114,7 @@ const registerChatBotApiListener = (ipcMain: Electron.IpcMain) => {
         if (!openai) {
           openai = createOpenaiClient({ apiKey });
         }
-        await chatBotApi.ask({ openai, assistantId, threadId, question: question });
+        await chatBotApi.ask({ openai, assistantId, threadId, question: question, mainWindow });
       }
     },
   );

src/electron/main.ts

@@ -1,4 +1,4 @@
-import { app, BrowserWindow, ipcMain, shell } from 'electron';
+import { app, BrowserWindow, ipcMain, IpcMain, shell } from 'electron';
 import path from 'path';
 import { createMenu } from './menu';
 import { debug } from '../common';
@@ -30,19 +30,22 @@ const loadDevTools = async () => {
   }
 };
 
-const createWindow = async () => {
-  // Create the browser window.
-  const mainWindow = new BrowserWindow({
-    width: 1200,
-    minWidth: 900,
-    height: 750,
-    minHeight: 600,
-    webPreferences: {
-      preload: path.join(__dirname, 'preload.js'),
-      webSecurity: false,
-    },
-    icon: path.resolve(__dirname, '../../dockit.png'),
-  });
+const registerListeners = (ipcMain: IpcMain, mainWindow: BrowserWindow) => {
+  registerStoreApiListener(ipcMain);
+  registerSourceFileApiListener(ipcMain);
+  registerFetchApiListener(ipcMain);
+  registerChatBotApiListener(ipcMain, mainWindow);
+
+  ipcMain.handle('versions', () => ({
+    node: process.versions.chrome,
+    chrome: process.versions.chrome,
+    electron: process.versions.electron,
+    version: app.getVersion(),
+    name: app.getName(),
+  }));
+};
+
+const renderMainWindow = async (mainWindow: BrowserWindow) => {
   createMenu(mainWindow);
   // and load the index.html of the app.
   if (MAIN_WINDOW_VITE_DEV_SERVER_URL) {
@@ -65,21 +68,30 @@ const createWindow = async () => {
       path.join(__dirname, `../renderer/${MAIN_WINDOW_VITE_NAME}/index.html`),
     );
   }
+};
 
-  ipcMain.handle('versions', () => ({
-    node: process.versions.chrome,
-    chrome: process.versions.chrome,
-    electron: process.versions.electron,
-    version: app.getVersion(),
-    name: app.getName(),
-  }));
+const createWindow = () => {
+  // Create the browser window.
+  return new BrowserWindow({
+    width: 1200,
+    minWidth: 900,
+    height: 750,
+    minHeight: 600,
+    webPreferences: {
+      preload: path.join(__dirname, 'preload.js'),
+      webSecurity: false,
+    },
+    icon: path.resolve(__dirname, '../../dockit.png'),
+  });
 };
 
 // This method will be called when Electron has finished
 // initialization and is ready to create browser windows.
 // Some APIs can only be used after this event occurs.
 app.on('ready', async () => {
-  await createWindow();
+  const mainWindow = createWindow();
+  await renderMainWindow(mainWindow);
+  registerListeners(ipcMain, mainWindow);
   await loadDevTools();
 });
 
@@ -96,7 +108,8 @@ app.on('window-all-closed', () => {
 // dock icon is clicked and there are no other windows open.
 app.on('activate', async () => {
   if (BrowserWindow.getAllWindows().length === 0) {
-    await createWindow();
+    const mainWindow = createWindow();
+    await renderMainWindow(mainWindow);
   }
 });
 
@@ -107,10 +120,5 @@ ipcMain.on('open-link', (_event, link: string) => {
   shell.openExternal(link);
 });
 
-registerStoreApiListener(ipcMain);
-registerSourceFileApiListener(ipcMain);
-registerFetchApiListener(ipcMain);
-registerChatBotApiListener(ipcMain);
-
 // In this file you can include the rest of your app's specific main process
 // code. You can also put them in separate files and import them here.

src/electron/preload.ts

@@ -46,4 +46,6 @@ contextBridge.exposeInMainWorld('chatBotApi', {
     threadId: string;
   }) =>
     ipcRenderer.invoke('chatBotApi', { method: 'ASK', question, apiKey, assistantId, threadId }),
+  onMessageReceived: (callback: (value: unknown) => void) =>
+    ipcRenderer.on('chat-bot-api-message-delta', (_event, value) => callback(value)),
 });

src/layout/components/tool-bar-right.vue

@@ -88,7 +88,7 @@ import { ChatMessageRole, useChatStore } from '../../store';
 import { storeToRefs } from 'pinia';
 
 const chatStore = useChatStore();
-const { sendMessage } = chatStore;
+const { sendMessage, fetchChats } = chatStore;
 const { chats } = storeToRefs(chatStore);
 const selectedItemId = ref(-1);
 const chatBot = ref({ active: false });
@@ -145,6 +145,7 @@ onMounted(() => {
 onUnmounted(() => {
   window.removeEventListener('resize', updateHeight);
 });
+fetchChats();
 </script>
 
 <style scoped>

src/store/chatStore.ts

@@ -7,13 +7,24 @@ enum MessageStatus {
   SENDING = 'SENDING',
   SENT = 'SENT',
   FAILED = 'FAILED',
+  RECEIVED = 'RECEIVED',
 }
 export enum ChatMessageRole {
   USER = 'USER',
   BOT = 'BOT',
 }
 const { chatBotApi, storeAPI } = window;
 
+let receiveRegistration = false;
+
+const getOpenAiConfig = async () => {
+  const { openAi } = await storeAPI.getSecret('aigcConfig', { openAi: undefined });
+  if (!openAi) {
+    throw new Error(lang.global.t('setting.ai.missing'));
+  }
+  return openAi;
+};
+
 export const useChatStore = defineStore('chat', {
   state: (): {
     chats: Array<{
@@ -28,30 +39,45 @@ export const useChatStore = defineStore('chat', {
       assistantId: string;
       threadId: string;
     }>;
-    openaiConfig: { apiKey: string; model: string; prompt: string };
   } => {
     return {
       chats: [],
-      openaiConfig: { apiKey: '', model: '', prompt: '' },
     };
   },
   actions: {
+    async fetchChats() {
+      const chats = await storeAPI.get('chats', undefined);
+      this.chats = chats ?? [];
+    },
     async sendMessage(content: string) {
-      const chat = this.chats[0];
+      if (!receiveRegistration) {
+        console.log('register onMessageReceived');
+        chatBotApi.onMessageReceived(({ delta, msgEvent }) => {
+          console.log('onMessageReceived', delta, msgEvent);
+          if (msgEvent === 'messageCreated') {
+            this.chats[0].messages.push({
+              id: ulid(),
+              status: MessageStatus.RECEIVED,
+              role: ChatMessageRole.BOT,
+              content: '',
+            });
+          } else if (msgEvent === 'messageDelta') {
+            const messageChunk = delta.content.map(({ text }) => text.value).join('');
+            this.chats[0].messages[this.chats[0].messages.length - 1].content += messageChunk;
+          } else if (msgEvent === 'messageDone') {
+            storeAPI.set('chats', pureObject(this.chats));
+          }
+        });
+        receiveRegistration = true;
+      }
 
-      if (!chat) {
-        const { openAi } = await storeAPI.getSecret('aigcConfig', { openAi: undefined });
-        console.log('openAi', openAi);
-        if (!openAi) {
-          throw new Error(lang.global.t('setting.ai.missing'));
-        }
-        this.openaiConfig = openAi;
+      if (!this.chats[0]) {
         const chats = await storeAPI.get('chats', undefined);
         if (chats) {
           this.chats = chats;
         } else {
           try {
-            const { assistantId, threadId } = await chatBotApi.initialize(openAi);
+            const { assistantId, threadId } = await chatBotApi.initialize(await getOpenAiConfig());
             this.chats.push({ id: ulid(), type: 'openai', messages: [], assistantId, threadId });
             await storeAPI.set('chats', pureObject(this.chats));
           } catch (err) {
@@ -69,11 +95,12 @@ export const useChatStore = defineStore('chat', {
       });
       await storeAPI.set('chats', pureObject(this.chats));
       try {
+        const openaiConfig = await getOpenAiConfig();
         await chatBotApi.ask({
           question: content,
           assistantId,
           threadId,
-          apiKey: this.openaiConfig.apiKey,
+          apiKey: openaiConfig.apiKey,
         });
       } catch (err) {
         messages[messages.length - 1].status = MessageStatus.FAILED;

src/vite-env.d.ts

@@ -46,6 +46,22 @@ export interface IChatBotAPI {
     threadId: string;
     apiKey: string;
   }) => Promise<void>;
+  onMessageReceived: (
+    fn: (value: {
+      msgEvent: string;
+      delta?: {
+        content: Array<{
+          index: number;
+          type: string;
+          text: {
+            value: string;
+            annotations: Array<string>;
+          };
+        }>;
+      };
+      message: string;
+    }) => void,
+  ) => void;
 }
 
 export interface ISourceFileAPI {