Merge pull request #458 from janhq/allow-switching-model

fix: allow switching models when switch between conversations

web/app/_components/HistoryItem/index.tsx

@@ -1,6 +1,5 @@
 import React from 'react'
 import { useAtomValue, useSetAtom } from 'jotai'
-import { ModelManagementService } from '@janhq/core'
 import {
   getActiveConvoIdAtom,
   setActiveConvoIdAtom,
@@ -12,11 +11,13 @@ import {
 } from '@helpers/atoms/MainView.atom'
 import { displayDate } from '@utils/datetime'
 import { twMerge } from 'tailwind-merge'
-import { executeSerial } from '@services/pluginService'
+import { activeAssistantModelAtom } from '@helpers/atoms/Model.atom'
+import { switchingModelConfirmationModalPropsAtom } from '@helpers/atoms/Modal.atom'
+import useStartStopModel from '@hooks/useStartStopModel'
+import useGetModelById from '@hooks/useGetModelById'
 
 type Props = {
   conversation: Conversation
-  avatarUrl?: string
   name: string
   summary?: string
   updatedAt?: string
@@ -24,22 +25,41 @@ type Props = {
 
 const HistoryItem: React.FC<Props> = ({
   conversation,
-  avatarUrl,
   name,
   summary,
   updatedAt,
 }) => {
-  const setMainViewState = useSetAtom(setMainViewStateAtom)
   const activeConvoId = useAtomValue(getActiveConvoIdAtom)
+  const isSelected = activeConvoId === conversation._id
+  const activeModel = useAtomValue(activeAssistantModelAtom)
+  const { startModel } = useStartStopModel()
+  const { getModelById } = useGetModelById()
+
+  const setMainViewState = useSetAtom(setMainViewStateAtom)
   const setActiveConvoId = useSetAtom(setActiveConvoIdAtom)
   const updateConvWaiting = useSetAtom(updateConversationWaitingForResponseAtom)
-  const isSelected = activeConvoId === conversation._id
+  const setConfirmationModalProps = useSetAtom(
+    switchingModelConfirmationModalPropsAtom
+  )
 
   const onClick = async () => {
-    const model = await executeSerial(
-      ModelManagementService.GetModelById,
-      conversation.modelId
-    )
+    if (conversation.modelId == null) {
+      console.debug('modelId is undefined')
+      return
+    }
+
+    const model = await getModelById(conversation.modelId)
+    if (model != null) {
+      if (activeModel == null) {
+        // if there's no active model, we simply load conversation's model
+        startModel(model._id)
+      } else if (activeModel._id !== model._id) {
+        // display confirmation modal
+        setConfirmationModalProps({
+          replacingModel: model,
+        })
+      }
+    }
 
     if (conversation._id) updateConvWaiting(conversation._id, true)
 

web/app/_components/HistoryList/index.tsx

@@ -1,5 +1,5 @@
 import HistoryItem from '../HistoryItem'
-import { useEffect, useState } from 'react'
+import { useEffect } from 'react'
 import ExpandableHeader from '../ExpandableHeader'
 import { useAtomValue } from 'jotai'
 import { searchAtom } from '@helpers/JotaiWrapper'
@@ -33,7 +33,6 @@ const HistoryList: React.FC = () => {
                 key={convo._id}
                 conversation={convo}
                 summary={convo.summary}
-                avatarUrl={convo.image}
                 name={convo.name || 'Jan'}
                 updatedAt={convo.updatedAt ?? ''}
               />

web/app/_components/InputToolbar/index.tsx

@@ -3,66 +3,27 @@
 
 import BasicPromptInput from '../BasicPromptInput'
 import BasicPromptAccessories from '../BasicPromptAccessories'
-import { useAtomValue, useSetAtom } from 'jotai'
+import { useAtomValue } from 'jotai'
 import SecondaryButton from '../SecondaryButton'
-import { useEffect, useState } from 'react'
 import { PlusIcon } from '@heroicons/react/24/outline'
 import useCreateConversation from '@hooks/useCreateConversation'
 import { activeAssistantModelAtom } from '@helpers/atoms/Model.atom'
 import {
-  currentConversationAtom,
   currentConvoStateAtom,
+  getActiveConvoIdAtom,
 } from '@helpers/atoms/Conversation.atom'
-import useGetBots from '@hooks/useGetBots'
-import { activeBotAtom } from '@helpers/atoms/Bot.atom'
-import { useGetDownloadedModels } from '@hooks/useGetDownloadedModels'
+import useGetInputState from '@hooks/useGetInputState'
+import { Button } from '../../../uikit/button'
+import useStartStopModel from '@hooks/useStartStopModel'
 
 const InputToolbar: React.FC = () => {
   const activeModel = useAtomValue(activeAssistantModelAtom)
-  const { requestCreateConvo } = useCreateConversation()
   const currentConvoState = useAtomValue(currentConvoStateAtom)
-  const currentConvo = useAtomValue(currentConversationAtom)
+  const { inputState, currentConvo } = useGetInputState()
+  const { requestCreateConvo } = useCreateConversation()
+  const { startModel } = useStartStopModel()
 
-  const setActiveBot = useSetAtom(activeBotAtom)
-  const { getBotById } = useGetBots()
-  const [inputState, setInputState] = useState<
-    'available' | 'disabled' | 'loading'
-  >()
-  const [error, setError] = useState<string | undefined>()
-  const { downloadedModels } = useGetDownloadedModels()
-
-  useEffect(() => {
-    const getReplyState = async () => {
-      setInputState('loading')
-      if (currentConvo && currentConvo.botId && currentConvo.botId.length > 0) {
-        // if botId is set, check if bot is available
-        const bot = await getBotById(currentConvo.botId)
-        console.debug('Found bot', JSON.stringify(bot, null, 2))
-        if (bot) {
-          setActiveBot(bot)
-        }
-        setInputState(bot ? 'available' : 'disabled')
-        setError(
-          bot
-            ? undefined
-            : `Bot ${currentConvo.botId} has been deleted by its creator. Your chat history is saved but you won't be able to send new messages.`
-        )
-      } else {
-        const model = downloadedModels.find(
-          (model) => model._id === activeModel?._id
-        )
-
-        setInputState(model ? 'available' : 'disabled')
-        setError(
-          model
-            ? undefined
-            : `Model ${activeModel?._id} cannot be found. Your chat history is saved but you won't be able to send new messages.`
-        )
-      }
-    }
-    getReplyState()
-    // eslint-disable-next-line react-hooks/exhaustive-deps
-  }, [currentConvo])
+  const activeConvoId = useAtomValue(getActiveConvoIdAtom)
 
   const onNewConversationClick = () => {
     if (activeModel) {
@@ -70,16 +31,42 @@ const InputToolbar: React.FC = () => {
     }
   }
 
-  if (inputState === 'loading') return <div>Loading..</div>
+  const onStartModelClick = () => {
+    const modelId = currentConvo?.modelId
+    if (!modelId) return
+    startModel(modelId)
+  }
 
-  if (inputState === 'disabled')
+  if (!activeConvoId) {
+    return null
+  }
+
+  if (inputState === 'model-mismatch' || inputState === 'loading') {
+    const message = inputState === 'loading' ? 'Loading..' : 'Model mismatch!'
+    return (
+      <div className="sticky bottom-0 flex items-center justify-center bg-background/90">
+        <div className="mb-2">
+          <p className="mx-auto my-5 line-clamp-2 text-ellipsis text-center italic text-gray-600">
+            {message}
+          </p>
+          <Button onClick={onStartModelClick}>
+            Load {currentConvo?.modelId}
+          </Button>
+        </div>
+      </div>
+    )
+  }
+
+  if (inputState === 'model-not-found') {
     return (
       <div className="sticky bottom-0 flex items-center justify-center bg-background/90">
         <p className="mx-auto my-5 line-clamp-2 text-ellipsis text-center italic text-gray-600">
-          {error}
+          Model {currentConvo?.modelId} not found! Please re-download the model
+          first.
         </p>
       </div>
     )
+  }
 
   return (
     <div className="sticky bottom-0 w-full bg-background/90 px-5 py-0">

web/app/_components/SwitchingModelConfirmationModal/index.tsx

@@ -0,0 +1,129 @@
+import React, { Fragment } from 'react'
+import { Dialog, Transition } from '@headlessui/react'
+import { ExclamationTriangleIcon, XMarkIcon } from '@heroicons/react/24/outline'
+import { switchingModelConfirmationModalPropsAtom } from '@helpers/atoms/Modal.atom'
+import { useAtom, useAtomValue } from 'jotai'
+import { activeAssistantModelAtom } from '@helpers/atoms/Model.atom'
+import useStartStopModel from '@hooks/useStartStopModel'
+
+export type SwitchingModelConfirmationModalProps = {
+  replacingModel: AssistantModel
+}
+
+const SwitchingModelConfirmationModal: React.FC = () => {
+  const [props, setProps] = useAtom(switchingModelConfirmationModalPropsAtom)
+  const activeModel = useAtomValue(activeAssistantModelAtom)
+  const { startModel } = useStartStopModel()
+
+  const onConfirmSwitchModelClick = () => {
+    const modelId = props?.replacingModel._id
+    if (modelId) {
+      startModel(modelId)
+    }
+    setProps(undefined)
+  }
+
+  return (
+    <Transition.Root show={props != null} as={Fragment}>
+      <Dialog
+        as="div"
+        className="relative z-10"
+        onClose={() => setProps(undefined)}
+      >
+        <Transition.Child
+          as={Fragment}
+          enter="ease-out duration-300"
+          enterFrom="opacity-0"
+          enterTo="opacity-100"
+          leave="ease-in duration-200"
+          leaveFrom="opacity-100"
+          leaveTo="opacity-0"
+        >
+          <div className="fixed inset-0 bg-gray-500 bg-opacity-75 transition-opacity" />
+        </Transition.Child>
+
+        <div className="fixed inset-0 z-10 w-screen overflow-y-auto">
+          <div className="flex min-h-full items-end justify-center p-4 text-center sm:items-center sm:p-0">
+            <Transition.Child
+              as={Fragment}
+              enter="ease-out duration-300"
+              enterFrom="opacity-0 translate-y-4 sm:translate-y-0 sm:scale-95"
+              enterTo="opacity-100 translate-y-0 sm:scale-100"
+              leave="ease-in duration-200"
+              leaveFrom="opacity-100 translate-y-0 sm:scale-100"
+              leaveTo="opacity-0 translate-y-4 sm:translate-y-0 sm:scale-95"
+            >
+              <Dialog.Panel className="relative transform overflow-hidden rounded-lg bg-white px-4 pb-4 pt-5 text-left shadow-xl transition-all sm:my-8 sm:w-full sm:max-w-lg sm:p-6">
+                <div className="absolute right-0 top-0 hidden pr-4 pt-4 sm:block">
+                  <button
+                    type="button"
+                    className="rounded-md bg-white text-gray-400 hover:text-gray-500 focus:outline-none focus:ring-2 focus:ring-indigo-500 focus:ring-offset-2"
+                    onClick={() => setProps(undefined)}
+                  >
+                    <span className="sr-only">Close</span>
+                    <XMarkIcon className="h-6 w-6" aria-hidden="true" />
+                  </button>
+                </div>
+                <div className="sm:flex sm:items-start">
+                  <div className="mx-auto flex h-12 w-12 flex-shrink-0 items-center justify-center rounded-full bg-red-100 sm:mx-0 sm:h-10 sm:w-10">
+                    <ExclamationTriangleIcon
+                      className="h-6 w-6 text-red-600"
+                      aria-hidden="true"
+                    />
+                  </div>
+                  <div className="mt-3 text-center sm:ml-4 sm:mt-0 sm:text-left">
+                    <Dialog.Title
+                      as="h3"
+                      className="text-base font-semibold leading-6 text-gray-900"
+                    >
+                      Switching model
+                    </Dialog.Title>
+                    <div className="mt-2 flex flex-col">
+                      <p className="text-sm text-gray-500">
+                        Selected conversation is using model{' '}
+                        <span className="font-semibold text-black">
+                          {props?.replacingModel._id}
+                        </span>
+                        , but the active model is using{' '}
+                        <span className="font-semibold text-black">
+                          {activeModel?._id}
+                        </span>
+                        .
+                      </p>
+                      <br />
+                      <p className="text-sm text-gray-500">
+                        Switch to
+                        <span className="font-semibold text-black">
+                          {' '}
+                          {props?.replacingModel._id}?
+                        </span>
+                      </p>
+                    </div>
+                  </div>
+                </div>
+                <div className="mt-5 sm:mt-4 sm:flex sm:flex-row-reverse">
+                  <button
+                    type="button"
+                    className="inline-flex w-full justify-center rounded-md bg-red-600 px-3 py-2 text-sm font-semibold text-white shadow-sm hover:bg-red-500 sm:ml-3 sm:w-auto"
+                    onClick={onConfirmSwitchModelClick}
+                  >
+                    Switch
+                  </button>
+                  <button
+                    type="button"
+                    className="mt-3 inline-flex w-full justify-center rounded-md bg-white px-3 py-2 text-sm font-semibold text-gray-900 shadow-sm ring-1 ring-inset ring-gray-300 hover:bg-gray-50 sm:mt-0 sm:w-auto"
+                    onClick={() => setProps(undefined)}
+                  >
+                    Cancel
+                  </button>
+                </div>
+              </Dialog.Panel>
+            </Transition.Child>
+          </div>
+        </div>
+      </Dialog>
+    </Transition.Root>
+  )
+}
+
+export default SwitchingModelConfirmationModal

web/containers/Providers/index.tsx

@@ -69,17 +69,17 @@ const Providers = (props: PropsWithChildren) => {
   return (
     <JotaiWrapper>
       {setupCore && (
-        <EventListenerWrapper>
         <ThemeWrapper>
           {activated ? (
+            <EventListenerWrapper>
               <ModalWrapper>{children}</ModalWrapper>
+            </EventListenerWrapper>
           ) : (
-              <div className="bg-background flex h-screen w-screen items-center justify-center">
+            <div className="flex h-screen w-screen items-center justify-center bg-background">
               <CompactLogo width={56} height={56} />
             </div>
           )}
         </ThemeWrapper>
-        </EventListenerWrapper>
       )}
     </JotaiWrapper>
   )

web/helpers/EventHandler.tsx

@@ -3,15 +3,30 @@ import { toChatMessage } from '@models/ChatMessage'
 import { events, EventName, NewMessageResponse } from '@janhq/core'
 import { useSetAtom } from 'jotai'
 import { ReactNode, useEffect } from 'react'
+import useGetBots from '@hooks/useGetBots'
+import useGetUserConversations from '@hooks/useGetUserConversations'
 
 export default function EventHandler({ children }: { children: ReactNode }) {
   const addNewMessage = useSetAtom(addNewMessageAtom)
   const updateMessage = useSetAtom(updateMessageAtom)
+  const { getBotById } = useGetBots()
+  const { getConversationById } = useGetUserConversations()
 
-  function handleNewMessageResponse(message: NewMessageResponse) {
+  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
   ) {
@@ -40,5 +55,5 @@ export default function EventHandler({ children }: { children: ReactNode }) {
       events.off(EventName.OnMessageResponseUpdate, handleMessageResponseUpdate)
     }
   }, [])
-  return <> {children}</>
+  return <>{children}</>
 }

web/helpers/ModalWrapper.tsx

@@ -5,6 +5,7 @@ import ConfirmDeleteConversationModal from '@/_components/ConfirmDeleteConversat
 import ConfirmDeleteModelModal from '@/_components/ConfirmDeleteModelModal'
 import ConfirmSignOutModal from '@/_components/ConfirmSignOutModal'
 import MobileMenuPane from '@/_components/MobileMenuPane'
+import SwitchingModelConfirmationModal from '@/_components/SwitchingModelConfirmationModal'
 import { ReactNode } from 'react'
 
 type Props = {
@@ -18,6 +19,7 @@ export const ModalWrapper: React.FC<Props> = ({ children }) => (
     <ConfirmSignOutModal />
     <ConfirmDeleteModelModal />
     <BotListModal />
+    <SwitchingModelConfirmationModal />
     {children}
   </>
 )

web/helpers/atoms/Modal.atom.ts

@@ -1,3 +1,4 @@
+import { SwitchingModelConfirmationModalProps } from '@/_components/SwitchingModelConfirmationModal'
 import { atom } from 'jotai'
 
 export const showConfirmDeleteConversationModalAtom = atom(false)
@@ -7,3 +8,7 @@ export const showingAdvancedPromptAtom = atom<boolean>(false)
 export const showingProductDetailAtom = atom<boolean>(false)
 export const showingMobilePaneAtom = atom<boolean>(false)
 export const showingBotListModalAtom = atom<boolean>(false)
+
+export const switchingModelConfirmationModalPropsAtom = atom<
+  SwitchingModelConfirmationModalProps | undefined
+>(undefined)

web/hooks/useChatMessages.ts

@@ -3,11 +3,15 @@ import { executeSerial } from '@services/pluginService'
 import { useAtomValue, useSetAtom } from 'jotai'
 import { useEffect } from 'react'
 import { DataService } from '@janhq/core'
-import { getActiveConvoIdAtom } from '@helpers/atoms/Conversation.atom'
+import {
+  getActiveConvoIdAtom,
+  userConversationsAtom,
+} from '@helpers/atoms/Conversation.atom'
 import {
   getCurrentChatMessagesAtom,
   setCurrentChatMessagesAtom,
 } from '@helpers/atoms/ChatMessage.atom'
+import useGetBots from './useGetBots'
 
 /**
  * Custom hooks to get chat messages for current(active) conversation
@@ -16,6 +20,8 @@ const useChatMessages = () => {
   const setMessages = useSetAtom(setCurrentChatMessagesAtom)
   const messages = useAtomValue(getCurrentChatMessagesAtom)
   const activeConvoId = useAtomValue(getActiveConvoIdAtom)
+  const userConversations = useAtomValue(userConversationsAtom)
+  const { getBotById } = useGetBots()
 
   const getMessages = async (convoId: string) => {
     const data: any = await executeSerial(
@@ -26,6 +32,12 @@ const useChatMessages = () => {
       return []
     }
 
+    const convo = userConversations.find((c) => c._id === convoId)
+    if (convo && convo.botId) {
+      const bot = await getBotById(convo.botId)
+      return parseMessages(data, bot)
+    }
+
     return parseMessages(data)
   }
 
@@ -47,10 +59,10 @@ const useChatMessages = () => {
   return { messages }
 }
 
-function parseMessages(messages: RawMessage[]): ChatMessage[] {
+function parseMessages(messages: RawMessage[], bot?: Bot): ChatMessage[] {
   const newMessages: ChatMessage[] = []
   for (const m of messages) {
-    const chatMessage = toChatMessage(m)
+    const chatMessage = toChatMessage(m, bot)
     newMessages.push(chatMessage)
   }
   return newMessages

web/hooks/useCreateConversation.ts

@@ -6,19 +6,18 @@ import {
   setActiveConvoIdAtom,
   addNewConversationStateAtom,
 } from '@helpers/atoms/Conversation.atom'
+import useGetModelById from './useGetModelById'
 
 const useCreateConversation = () => {
   const [userConversations, setUserConversations] = useAtom(
     userConversationsAtom
   )
+  const { getModelById } = useGetModelById()
   const setActiveConvoId = useSetAtom(setActiveConvoIdAtom)
   const addNewConvoState = useSetAtom(addNewConversationStateAtom)
 
   const createConvoByBot = async (bot: Bot) => {
-    const model = await executeSerial(
-      ModelManagementService.GetModelById,
-      bot.modelId
-    )
+    const model = await getModelById(bot.modelId)
 
     if (!model) {
       alert(

web/hooks/useGetInputState.ts

@@ -0,0 +1,52 @@
+import { currentConversationAtom } from '@helpers/atoms/Conversation.atom'
+import { activeAssistantModelAtom } from '@helpers/atoms/Model.atom'
+import { useAtomValue } from 'jotai'
+import { useEffect, useState } from 'react'
+import { useGetDownloadedModels } from './useGetDownloadedModels'
+
+export default function useGetInputState() {
+  const [inputState, setInputState] = useState<InputType>('loading')
+  const currentConvo = useAtomValue(currentConversationAtom)
+  const activeModel = useAtomValue(activeAssistantModelAtom)
+  const { downloadedModels } = useGetDownloadedModels()
+
+  const handleInputState = (
+    convo: Conversation | undefined,
+    currentModel: AssistantModel | undefined,
+    models: AssistantModel[]
+  ) => {
+    if (convo == null) return
+    if (currentModel == null) {
+      setInputState('loading')
+      return
+    }
+
+    // check if convo model id is in downloaded models
+    const isModelAvailable = downloadedModels.some(
+      (model) => model._id === convo.modelId
+    )
+
+    if (!isModelAvailable) {
+      // can't find model in downloaded models
+      setInputState('model-not-found')
+      return
+    }
+
+    if (convo.modelId !== currentModel._id) {
+      // in case convo model and active model is different,
+      // ask user to init the required model
+      setInputState('model-mismatch')
+      return
+    }
+
+    setInputState('available')
+  }
+
+  useEffect(() => {
+    handleInputState(currentConvo, activeModel, downloadedModels)
+  }, [currentConvo, activeModel, downloadedModels])
+
+  return { inputState, currentConvo }
+}
+
+type InputType = 'available' | 'loading' | 'model-mismatch' | 'model-not-found'

web/hooks/useGetModelById.ts

@@ -0,0 +1,23 @@
+import { ModelManagementService } from '@janhq/core'
+import { executeSerial } from '../../electron/core/plugin-manager/execution/extension-manager'
+
+export default function useGetModelById() {
+  const getModelById = async (
+    modelId: string
+  ): Promise<AssistantModel | undefined> => {
+    return queryModelById(modelId)
+  }
+
+  return { getModelById }
+}
+
+const queryModelById = async (
+  modelId: string
+): Promise<AssistantModel | undefined> => {
+  const model = await executeSerial(
+    ModelManagementService.GetModelById,
+    modelId
+  )
+
+  return model
+}

web/hooks/useGetUserConversations.ts

@@ -29,7 +29,14 @@ const useGetUserConversations = () => {
     }
   }
 
+  const getConversationById = async (
+    id: string
+  ): Promise<Conversation | undefined> => {
+    return await executeSerial(DataService.GetConversationById, id)
+  }
+
   return {
+    getConversationById,
     getUserConversations,
   }
 }

web/hooks/useStartStopModel.ts

@@ -1,11 +1,12 @@
 import { executeSerial } from '@services/pluginService'
-import { ModelManagementService, InferenceService } from '@janhq/core'
+import { InferenceService } from '@janhq/core'
 import { useAtom, useSetAtom } from 'jotai'
 import { activeAssistantModelAtom, stateModel } from '@helpers/atoms/Model.atom'
-import { useState } from 'react'
+import useGetModelById from './useGetModelById'
 
 export default function useStartStopModel() {
   const [activeModel, setActiveModel] = useAtom(activeAssistantModelAtom)
+  const { getModelById } = useGetModelById()
   const setStateModel = useSetAtom(stateModel)
 
   const startModel = async (modelId: string) => {
@@ -16,25 +17,27 @@ export default function useStartStopModel() {
 
     setStateModel({ state: 'start', loading: true, model: modelId })
 
-    const model = await executeSerial(
-      ModelManagementService.GetModelById,
-      modelId
-    )
+    const model = await getModelById(modelId)
+
     if (!model) {
       alert(`Model ${modelId} not found! Please re-download the model first.`)
       setStateModel((prev) => ({ ...prev, loading: false }))
+      return
     }
+
     const currentTime = Date.now()
     console.debug('Init model: ', model._id)
 
-    const res = await executeSerial(InferenceService.InitModel, model._id)
+    const res = await initModel(model._id)
     if (res?.error) {
       const errorMessage = `Failed to init model: ${res.error}`
       console.error(errorMessage)
       alert(errorMessage)
     } else {
       console.debug(
-        `Init model successfully!, take ${Date.now() - currentTime}ms`
+        `Init model ${modelId} successfully!, take ${
+          Date.now() - currentTime
+        }ms`
       )
       setActiveModel(model)
     }
@@ -52,3 +55,7 @@ export default function useStartStopModel() {
 
   return { startModel, stopModel }
 }
+
+const initModel = async (modelId: string): Promise<any> => {
+  return executeSerial(InferenceService.InitModel, modelId)
+}

web/models/ChatMessage.ts

@@ -41,7 +41,8 @@ export interface RawMessage {
 }
 
 export const toChatMessage = (
-  m: RawMessage | NewMessageResponse
+  m: RawMessage | NewMessageResponse,
+  bot?: Bot
 ): ChatMessage => {
   const createdAt = new Date(m.createdAt ?? '').getTime()
   const imageUrls: string[] = []
@@ -56,18 +57,18 @@ export const toChatMessage = (
 
   const content = m.message ?? ''
 
+  let senderName = m.user === 'user' ? 'You' : 'Assistant'
+  if (senderName === 'Assistant' && bot) {
+    senderName = bot.name
+  }
+
   return {
     id: (m._id ?? 0).toString(),
     conversationId: (m.conversationId ?? 0).toString(),
     messageType: messageType,
     messageSenderType: messageSenderType,
     senderUid: m.user?.toString() || '0',
-    senderName:
-      m.user === 'user'
-        ? 'You'
-        : m.user && m.user !== 'ai' && m.user !== 'assistant'
-        ? m.user
-        : 'Assistant',
+    senderName: senderName,
     senderAvatarUrl: m.avatar
       ? m.avatar
       : m.user === 'user'