jan/web/helpers/EventHandler.tsx

import { addNewMessageAtom, updateMessageAtom } from './atoms/ChatMessage.atom'
import { toChatMessage } from '@models/ChatMessage'
import { events, EventName, NewMessageResponse, DataService } from '@janhq/core'
import { useSetAtom } from 'jotai'
import { ReactNode, useEffect } from 'react'
import useGetBots from '@hooks/useGetBots'
import useGetUserConversations from '@hooks/useGetUserConversations'
import {
  updateConversationAtom,
  updateConversationWaitingForResponseAtom,
} from './atoms/Conversation.atom'
import { executeSerial } from '../../electron/core/plugin-manager/execution/extension-manager'
import { debounce } from 'lodash'
import { setDownloadStateAtom, setDownloadStateSuccessAtom } from "./atoms/DownloadState.atom";
import { downloadedModelAtom } from "./atoms/DownloadedModel.atom";
import { ModelManagementService } from "@janhq/core";
import { getDownloadedModels } from "../hooks/useGetDownloadedModels";

let currentConversation: Conversation | undefined = undefined

const debouncedUpdateConversation = debounce(
  async (updatedConv: Conversation) => {
    await executeSerial(DataService.UpdateConversation, updatedConv)
  },
  1000
)


export default function EventHandler({ children }: { children: ReactNode }) {
  const addNewMessage = useSetAtom(addNewMessageAtom)
  const updateMessage = useSetAtom(updateMessageAtom)
  const updateConversation = useSetAtom(updateConversationAtom)
  const { getBotById } = useGetBots()
  const { getConversationById } = useGetUserConversations()

  const updateConvWaiting = useSetAtom(updateConversationWaitingForResponseAtom)
  const setDownloadState = useSetAtom(setDownloadStateAtom);
  const setDownloadStateSuccess = useSetAtom(setDownloadStateSuccessAtom);
  const setDownloadedModels = useSetAtom(downloadedModelAtom);

  async function handleNewMessageResponse(message: NewMessageResponse) {
    if (message.conversationId) {
      const convo = await getConversationById(message.conversationId)
      const botId = convo?.botId
      console.debug('botId', botId)
      if (botId) {
        const bot = await getBotById(botId)
        const newResponse = toChatMessage(message, bot)
        addNewMessage(newResponse)
      } else {
        const newResponse = toChatMessage(message)
        addNewMessage(newResponse)
      }
    }
  }
  async function handleMessageResponseUpdate(
    messageResponse: NewMessageResponse
  ) {
    if (
      messageResponse.conversationId &&
      messageResponse._id &&
      messageResponse.message
    ) {
      updateMessage(
        messageResponse._id,
        messageResponse.conversationId,
        messageResponse.message
      )
    }

    if (messageResponse.conversationId) {
      if (
        !currentConversation ||
        currentConversation._id !== messageResponse.conversationId
      ) {
        currentConversation = await getConversationById(
          messageResponse.conversationId
        )
      }

      const updatedConv: Conversation = {
        ...currentConversation,
        lastMessage: messageResponse.message,
      }

      updateConversation(updatedConv)
      debouncedUpdateConversation(updatedConv)
    }
  }

  async function handleMessageResponseFinished(
    messageResponse: NewMessageResponse
  ) {
    if (!messageResponse.conversationId) return
    console.debug('handleMessageResponseFinished', messageResponse)
    updateConvWaiting(messageResponse.conversationId, false)
  }

  function handleDownloadUpdate(state: any) {
    if (!state) return;
    setDownloadState(state);
  }

  function handleDownloadSuccess(state: any) {
    if (state && state.fileName && state.success === true) {
      setDownloadStateSuccess(state.fileName);
      executeSerial(ModelManagementService.UpdateFinishedDownloadAt, state.fileName).then(() => {
        getDownloadedModels().then((models) => {
          setDownloadedModels(models);
        });
      });
    }
  }

  useEffect(() => {
    if (window.corePlugin.events) {
      events.on(EventName.OnNewMessageResponse, handleNewMessageResponse)
      events.on(EventName.OnMessageResponseUpdate, handleMessageResponseUpdate)
      events.on(
        "OnMessageResponseFinished",
        // EventName.OnMessageResponseFinished,
        handleMessageResponseFinished
      )
      events.on(EventName.OnDownloadUpdate, handleDownloadUpdate);
      events.on(EventName.OnDownloadSuccess, handleDownloadSuccess);
    }
  }, [])

  useEffect(() => {
    return () => {
      events.off(EventName.OnNewMessageResponse, handleNewMessageResponse)
      events.off(EventName.OnMessageResponseUpdate, handleMessageResponseUpdate)
      events.off(
        "OnMessageResponseFinished",
        // EventName.OnMessageResponseFinished,
        handleMessageResponseFinished
      )
      events.off(EventName.OnDownloadUpdate, handleDownloadUpdate);
      events.off(EventName.OnDownloadSuccess, handleDownloadSuccess);
    }
  }, [])
  return <>{children}</>
}