From 084da9ebab8b55e5575064da427817326a0e0060 Mon Sep 17 00:00:00 2001
From: kangfenmao <kangfenmao@qq.com>
Date: Wed, 22 Jan 2025 15:08:44 +0800
Subject: [PATCH] feat: enhance message model handling and user display

- Updated Message component to fallback to message.model if model retrieval fails, improving robustness.
- Refactored MessageHeader to utilize getModelName for better user name display based on message role, enhancing clarity.
- Introduced getModelName function in ModelService to streamline model name retrieval, improving code modularity and readability.
---
 src/renderer/src/pages/home/Messages/Message.tsx    |  2 +-
 .../src/pages/home/Messages/MessageHeader.tsx       | 13 ++++++++++---
 src/renderer/src/services/ModelService.ts           |  4 ++++
 3 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/src/renderer/src/pages/home/Messages/Message.tsx b/src/renderer/src/pages/home/Messages/Message.tsx
index 5be93e0f..b740cc82 100644
--- a/src/renderer/src/pages/home/Messages/Message.tsx
+++ b/src/renderer/src/pages/home/Messages/Message.tsx
@@ -53,7 +53,7 @@ const MessageItem: FC<Props> = ({
   const [message, setMessage] = useState(_message)
   const { t } = useTranslation()
   const { assistant, setModel } = useAssistant(message.assistantId)
-  const model = useModel(getMessageModelId(message))
+  const model = useModel(getMessageModelId(message)) || message.model
   const { isBubbleStyle } = useMessageStyle()
   const { showMessageDivider, messageFont, fontSize } = useSettings()
   const messageContainerRef = useRef<HTMLDivElement>(null)
diff --git a/src/renderer/src/pages/home/Messages/MessageHeader.tsx b/src/renderer/src/pages/home/Messages/MessageHeader.tsx
index 745cc74f..385b7bba 100644
--- a/src/renderer/src/pages/home/Messages/MessageHeader.tsx
+++ b/src/renderer/src/pages/home/Messages/MessageHeader.tsx
@@ -6,6 +6,7 @@ import { useTheme } from '@renderer/context/ThemeProvider'
 import useAvatar from '@renderer/hooks/useAvatar'
 import { useMessageStyle, useSettings } from '@renderer/hooks/useSettings'
 import { getMessageModelId } from '@renderer/services/MessagesService'
+import { getModelName } from '@renderer/services/ModelService'
 import { Assistant, Message, Model } from '@renderer/types'
 import { firstLetter, removeLeadingEmoji } from '@renderer/utils'
 import { Avatar } from 'antd'
@@ -35,10 +36,16 @@ const MessageHeader: FC<Props> = memo(({ assistant, model, message }) => {
   const avatarSource = useMemo(() => getAvatarSource(isLocalAi, getMessageModelId(message)), [message])
 
   const getUserName = useCallback(() => {
-    if (isLocalAi && message.role !== 'user') return APP_NAME
-    if (message.role === 'assistant') return model?.name || model?.id || message.modelId || ''
+    if (isLocalAi && message.role !== 'user') {
+      return APP_NAME
+    }
+
+    if (message.role === 'assistant') {
+      return getModelName(model) || getMessageModelId(message) || ''
+    }
+
     return userName || t('common.you')
-  }, [message.modelId, message.role, model?.id, model?.name, t, userName])
+  }, [message, model, t, userName])
 
   const isAssistantMessage = message.role === 'assistant'
   const showMinappIcon = sidebarIcons.visible.includes('minapp')
diff --git a/src/renderer/src/services/ModelService.ts b/src/renderer/src/services/ModelService.ts
index 92279ce2..2c348192 100644
--- a/src/renderer/src/services/ModelService.ts
+++ b/src/renderer/src/services/ModelService.ts
@@ -15,3 +15,7 @@ export const hasModel = (m?: Model) => {
 
   return allModels.find((model) => model.id === m?.id)
 }
+
+export function getModelName(model?: Model) {
+  return model?.name || model?.id || ''
+}