refactor: messages completion

src/renderer/src/pages/history/components/SearchMessage.tsx

@@ -24,7 +24,7 @@ const SearchMessage: FC<Props> = ({ message, ...props }) => {
   return (
     <MessagesContainer {...props}>
       <ContainerWrapper style={{ paddingTop: 20, paddingBottom: 20, position: 'relative' }}>
-        <MessageItem message={message} showMenu={false} />
+        <MessageItem message={message} />
         <Button
           type="text"
           size="middle"

src/renderer/src/pages/history/components/TopicMessages.tsx

@@ -38,7 +38,7 @@ const TopicMessages: FC<Props> = ({ topic, ...props }) => {
       <ContainerWrapper style={{ paddingTop: 30, paddingBottom: 30 }}>
         {topic?.messages.map((message) => (
           <div key={message.id} style={{ position: 'relative' }}>
-            <MessageItem message={message} showMenu={false} />
+            <MessageItem message={message} />
             <Button
               type="text"
               size="middle"

src/renderer/src/pages/home/Inputbar/Inputbar.tsx

@@ -87,6 +87,7 @@ const Inputbar: FC<Props> = ({ assistant, setActiveTopic }) => {
       assistantId: assistant.id,
       topicId: assistant.topics[0].id || uuid(),
       createdAt: dayjs().format('YYYY-MM-DD HH:mm:ss'),
+      type: 'text',
       status: 'success'
     }
 

src/renderer/src/pages/home/Messages/Message.tsx

@@ -1,11 +1,15 @@
 import { FONT_FAMILY } from '@renderer/config/constant'
+import db from '@renderer/databases'
 import { useAssistant } from '@renderer/hooks/useAssistant'
 import { useModel } from '@renderer/hooks/useModel'
 import { useSettings } from '@renderer/hooks/useSettings'
+import { fetchChatCompletion } from '@renderer/services/api'
 import { EVENT_NAMES, EventEmitter } from '@renderer/services/event'
-import { Message } from '@renderer/types'
+import { estimateMessageUsage } from '@renderer/services/tokens'
+import { Message, Topic } from '@renderer/types'
+import { runAsyncFunction } from '@renderer/utils'
 import { Divider } from 'antd'
-import { FC, memo, useEffect, useMemo, useRef } from 'react'
+import { Dispatch, FC, memo, SetStateAction, useEffect, useMemo, useRef, useState } from 'react'
 import { useTranslation } from 'react-i18next'
 import styled from 'styled-components'
 
@@ -16,31 +20,32 @@ import MessgeTokens from './MessageTokens'
 
 interface Props {
   message: Message
+  topic?: Topic
   index?: number
   total?: number
-  lastMessage?: boolean
-  showMenu?: boolean
   hidePresetMessages?: boolean
-  onEditMessage?: (message: Message) => void
+  onGetMessages?: () => Message[]
+  onSetMessages?: Dispatch<SetStateAction<Message[]>>
   onDeleteMessage?: (message: Message) => void
 }
 
 const MessageItem: FC<Props> = ({
-  message,
+  message: _message,
+  topic,
   index,
-  lastMessage,
-  showMenu = true,
   hidePresetMessages,
-  onEditMessage,
+  onDeleteMessage,
-  onDeleteMessage
+  onSetMessages,
+  onGetMessages
 }) => {
+  const [message, setMessage] = useState(_message)
   const { t } = useTranslation()
   const { assistant, setModel } = useAssistant(message.assistantId)
   const model = useModel(message.modelId)
   const { showMessageDivider, messageFont, fontSize } = useSettings()
-  const messageRef = useRef<HTMLDivElement>(null)
+  const messageContainerRef = useRef<HTMLDivElement>(null)
 
-  const isLastMessage = lastMessage || index === 0
+  const isLastMessage = index === 0
   const isAssistantMessage = message.role === 'assistant'
 
   const fontFamily = useMemo(() => {
@@ -49,16 +54,23 @@ const MessageItem: FC<Props> = ({
 
   const messageBorder = showMessageDivider ? undefined : 'none'
 
+  const onEditMessage = (msg: Message) => {
+    setMessage(msg)
+    const messages = onGetMessages?.().map((m) => (m.id === message.id ? message : m))
+    messages && onSetMessages?.(messages)
+    topic && db.topics.update(topic.id, { messages })
+  }
+
   useEffect(() => {
     const unsubscribes = [
       EventEmitter.on(EVENT_NAMES.LOCATE_MESSAGE + ':' + message.id, (highlight: boolean = true) => {
-        if (messageRef.current) {
+        if (messageContainerRef.current) {
-          messageRef.current.scrollIntoView({ behavior: 'smooth' })
+          messageContainerRef.current.scrollIntoView({ behavior: 'smooth' })
           if (highlight) {
             setTimeout(() => {
-              messageRef.current?.classList.add('message-highlight')
+              messageContainerRef.current?.classList.add('message-highlight')
               setTimeout(() => {
-                messageRef.current?.classList.remove('message-highlight')
+                messageContainerRef.current?.classList.remove('message-highlight')
               }, 2500)
             }, 500)
           }
@@ -68,6 +80,40 @@ const MessageItem: FC<Props> = ({
     return () => unsubscribes.forEach((unsub) => unsub())
   }, [message])
 
+  useEffect(() => {
+    if (message.role === 'user' && !message.usage) {
+      runAsyncFunction(async () => {
+        const usage = await estimateMessageUsage(message)
+        setMessage({ ...message, usage })
+        const topic = await db.topics.get({ id: message.topicId })
+        const messages = topic?.messages.map((m) => (m.id === message.id ? { ...m, usage } : m))
+        db.topics.update(message.topicId, { messages })
+      })
+    }
+  }, [message])
+
+  useEffect(() => {
+    if (topic && onGetMessages && onSetMessages) {
+      if (message.status === 'sending') {
+        const messages = onGetMessages()
+        fetchChatCompletion({
+          message,
+          messages: [...messages, message],
+          assistant,
+          topic,
+          onResponse: (msg) => {
+            setMessage(msg)
+            if (msg.status === 'success') {
+              const _messages = messages.map((m) => (m.id === msg.id ? msg : m))
+              onSetMessages(_messages)
+              db.topics.update(topic.id, { messages: _messages })
+            }
+          }
+        })
+      }
+    }
+  }, [])
+
   if (hidePresetMessages && message.isPreset) {
     return null
   }
@@ -81,11 +127,11 @@ const MessageItem: FC<Props> = ({
   }
 
   return (
-    <MessageContainer key={message.id} className="message" ref={messageRef}>
+    <MessageContainer key={message.id} className="message" ref={messageContainerRef}>
       <MessageHeader message={message} assistant={assistant} model={model} />
       <MessageContentContainer style={{ fontFamily, fontSize }}>
         <MessageContent message={message} model={model} />
-        {!lastMessage && showMenu && (
+        {!message.status.includes('ing') && (
           <MessageFooter style={{ border: messageBorder, flexDirection: isLastMessage ? 'row-reverse' : undefined }}>
             <MessgeTokens message={message} isLastMessage={isLastMessage} />
             <MessageMenubar

src/renderer/src/pages/home/Messages/Messages.tsx

@@ -3,11 +3,17 @@ import db from '@renderer/databases'
 import { useAssistant } from '@renderer/hooks/useAssistant'
 import { useSettings } from '@renderer/hooks/useSettings'
 import { getTopic, TopicManager } from '@renderer/hooks/useTopic'
-import { fetchChatCompletion, fetchMessagesSummary } from '@renderer/services/api'
+import { fetchMessagesSummary } from '@renderer/services/api'
 import { getDefaultTopic } from '@renderer/services/assistant'
 import { EVENT_NAMES, EventEmitter } from '@renderer/services/event'
-import { deleteMessageFiles, filterMessages, getContextCount } from '@renderer/services/messages'
+import {
-import { estimateHistoryTokens, estimateMessageUsage } from '@renderer/services/tokens'
+  deleteMessageFiles,
+  filterMessages,
+  getAssistantMessage,
+  getContextCount,
+  getUserMessage
+} from '@renderer/services/messages'
+import { estimateHistoryTokens } from '@renderer/services/tokens'
 import { Assistant, Message, Model, Topic } from '@renderer/types'
 import { captureScrollableDiv, runAsyncFunction, uuid } from '@renderer/utils'
 import { t } from 'i18next'
@@ -27,11 +33,13 @@ interface Props {
 
 const Messages: FC<Props> = ({ assistant, topic, setActiveTopic }) => {
   const [messages, setMessages] = useState<Message[]>([])
-  const [lastMessage, setLastMessage] = useState<Message | null>(null)
   const containerRef = useRef<HTMLDivElement>(null)
   const { updateTopic, addTopic } = useAssistant(assistant.id)
   const { showTopics, topicPosition, showAssistants } = useSettings()
 
+  const messagesRef = useRef(messages)
+  messagesRef.current = messages
+
   const maxWidth = useMemo(() => {
     const showRightTopics = showTopics && topicPosition === 'right'
     const minusAssistantsWidth = showAssistants ? '- var(--assistants-width)' : ''
@@ -39,22 +47,23 @@ const Messages: FC<Props> = ({ assistant, topic, setActiveTopic }) => {
     return `calc(100vw - var(--sidebar-width) ${minusAssistantsWidth} ${minusRightTopicsWidth} - 5px)`
   }, [showAssistants, showTopics, topicPosition])
 
+  const scrollToBottom = useCallback(() => {
+    setTimeout(() => containerRef.current?.scrollTo({ top: containerRef.current.scrollHeight, behavior: 'auto' }), 10)
+  }, [])
+
   const onSendMessage = useCallback(
     async (message: Message) => {
-      if (message.role === 'user') {
+      const assistantMessage = getAssistantMessage({ assistant, topic })
-        estimateMessageUsage(message).then((usage) => {
+
-          setMessages((prev) => {
+      setMessages((prev) => {
-            const _messages = prev.map((m) => (m.id === message.id ? { ...m, usage } : m))
+        const messages = prev.concat([message, assistantMessage])
-            db.topics.update(topic.id, { messages: _messages })
+        db.topics.put({ id: topic.id, messages })
-            return _messages
+        return messages
-          })
+      })
-        })
+
-      }
+      scrollToBottom()
-      const _messages = [...messages, message]
-      setMessages(_messages)
-      db.topics.put({ id: topic.id, messages: _messages })
     },
-    [messages, topic.id]
+    [assistant, scrollToBottom, topic]
   )
 
   const autoRenameTopic = useCallback(async () => {
@@ -79,56 +88,19 @@ const Messages: FC<Props> = ({ assistant, topic, setActiveTopic }) => {
     [messages, topic.id]
   )
 
-  const onEditMessage = useCallback(
+  const onGetMessages = useCallback(() => {
-    (message: Message) => {
+    return messagesRef.current
-      const _messages = messages.map((m) => (m.id === message.id ? message : m))
-      setMessages(_messages)
-      db.topics.update(topic.id, { messages: _messages })
-    },
-    [messages, topic.id]
-  )
-
-  const scrollToBottom = useCallback(() => {
-    setTimeout(() => containerRef.current?.scrollTo({ top: containerRef.current.scrollHeight, behavior: 'auto' }), 10)
   }, [])
 
   useEffect(() => {
     const unsubscribes = [
-      EventEmitter.on(EVENT_NAMES.SEND_MESSAGE, async (msg: Message) => {
+      EventEmitter.on(EVENT_NAMES.SEND_MESSAGE, onSendMessage),
-        await onSendMessage(msg)
+      EventEmitter.on(EVENT_NAMES.RECEIVE_MESSAGE, async () => {
-
-        // Scroll to bottom
-        scrollToBottom()
-
-        // Fetch completion
-        fetchChatCompletion({
-          assistant,
-          messages: [...messages, msg],
-          topic,
-          onResponse: setLastMessage
-        })
-      }),
-      EventEmitter.on(EVENT_NAMES.RECEIVE_MESSAGE, async (msg: Message) => {
-        setLastMessage(null)
-        onSendMessage(msg)
         setTimeout(() => EventEmitter.emit(EVENT_NAMES.AI_AUTO_RENAME), 100)
       }),
       EventEmitter.on(EVENT_NAMES.REGENERATE_MESSAGE, async (model: Model) => {
         const lastUserMessage = last(filterMessages(messages).filter((m) => m.role === 'user'))
-        if (lastUserMessage) {
+        lastUserMessage && onSendMessage({ ...lastUserMessage, id: uuid(), type: '@', modelId: model.id })
-          onSendMessage({
-            ...lastUserMessage,
-            id: uuid(),
-            type: '@',
-            modelId: model.id
-          })
-          fetchChatCompletion({
-            assistant,
-            topic,
-            messages: [...messages, lastUserMessage],
-            onResponse: setLastMessage
-          })
-        }
       }),
       EventEmitter.on(EVENT_NAMES.AI_AUTO_RENAME, autoRenameTopic),
       EventEmitter.on(EVENT_NAMES.CLEAR_MESSAGES, () => {
@@ -156,16 +128,11 @@ const Messages: FC<Props> = ({ assistant, topic, setActiveTopic }) => {
           return
         }
 
-        onSendMessage({
+        setMessages((prev) => {
-          id: uuid(),
+          const messages = prev.concat([getUserMessage({ assistant, topic, type: 'clear' })])
-          assistantId: assistant.id,
+          db.topics.put({ id: topic.id, messages })
-          role: 'user',
+          return messages
-          content: '',
+        })
-          topicId: topic.id,
-          createdAt: new Date().toISOString(),
-          status: 'success',
-          type: 'clear'
-        } as Message)
 
         scrollToBottom()
       }),
@@ -219,6 +186,8 @@ const Messages: FC<Props> = ({ assistant, topic, setActiveTopic }) => {
     })
   }, [assistant, messages])
 
+  const memoizedMessages = useMemo(() => reverse([...messages]), [messages])
+
   return (
     <Container
       id="messages"
@@ -226,16 +195,17 @@ const Messages: FC<Props> = ({ assistant, topic, setActiveTopic }) => {
       key={assistant.id}
       ref={containerRef}
       right={topicPosition === 'left'}>
-      <Suggestions assistant={assistant} messages={messages} lastMessage={lastMessage} />
+      <Suggestions assistant={assistant} messages={messages} />
-      {lastMessage && <MessageItem key={lastMessage.id} message={lastMessage} lastMessage />}
+      {memoizedMessages.map((message, index) => (
-      {reverse([...messages]).map((message, index) => (
         <MessageItem
           key={message.id}
           message={message}
+          topic={topic}
           index={index}
           hidePresetMessages={assistant.settings?.hideMessages}
-          onEditMessage={onEditMessage}
+          onSetMessages={setMessages}
           onDeleteMessage={onDeleteMessage}
+          onGetMessages={onGetMessages}
         />
       ))}
       <Prompt assistant={assistant} key={assistant.prompt} />

src/renderer/src/pages/home/components/Suggestions.tsx

@@ -3,6 +3,7 @@ import { EVENT_NAMES, EventEmitter } from '@renderer/services/event'
 import { Assistant, Message, Suggestion } from '@renderer/types'
 import { uuid } from '@renderer/utils'
 import dayjs from 'dayjs'
+import { last } from 'lodash'
 import { FC, useEffect, useState } from 'react'
 import BeatLoader from 'react-spinners/BeatLoader'
 import styled from 'styled-components'
@@ -10,12 +11,11 @@ import styled from 'styled-components'
 interface Props {
   assistant: Assistant
   messages: Message[]
-  lastMessage: Message | null
 }
 
 const suggestionsMap = new Map<string, Suggestion[]>()
 
-const Suggestions: FC<Props> = ({ assistant, messages, lastMessage }) => {
+const Suggestions: FC<Props> = ({ assistant, messages }) => {
   const [suggestions, setSuggestions] = useState<Suggestion[]>(
     suggestionsMap.get(messages[messages.length - 1]?.id) || []
   )
@@ -29,6 +29,7 @@ const Suggestions: FC<Props> = ({ assistant, messages, lastMessage }) => {
       assistantId: assistant.id,
       topicId: assistant.topics[0].id || uuid(),
       createdAt: dayjs().format('YYYY-MM-DD HH:mm:ss'),
+      type: 'text',
       status: 'success'
     }
 
@@ -54,7 +55,7 @@ const Suggestions: FC<Props> = ({ assistant, messages, lastMessage }) => {
     setSuggestions(suggestionsMap.get(messages[messages.length - 1]?.id) || [])
   }, [messages])
 
-  if (lastMessage) {
+  if (last(messages)?.status !== 'success') {
     return null
   }
 

src/renderer/src/pages/translate/TranslatePage.tsx

@@ -116,6 +116,7 @@ const TranslatePage: FC = () => {
       topicId: uuid(),
       modelId: translateModel.id,
       createdAt: new Date().toISOString(),
+      type: 'text',
       status: 'sending'
     }
 

src/renderer/src/providers/AiProvider.ts

@@ -10,6 +10,10 @@ export default class AiProvider {
     this.sdk = ProviderFactory.create(provider)
   }
 
+  public async fakeCompletions(params: CompletionsParams): Promise<void> {
+    return this.sdk.fakeCompletions(params)
+  }
+
   public async completions({ messages, assistant, onChunk, onFilterMessages }: CompletionsParams): Promise<void> {
     return this.sdk.completions({ messages, assistant, onChunk, onFilterMessages })
   }

src/renderer/src/providers/BaseProvider.ts

@@ -1,5 +1,6 @@
 import { getOllamaKeepAliveTime } from '@renderer/hooks/useOllama'
 import { Assistant, Message, Provider, Suggestion } from '@renderer/types'
+import { delay } from '@renderer/utils'
 import OpenAI from 'openai'
 
 export default abstract class BaseProvider {
@@ -20,6 +21,13 @@ export default abstract class BaseProvider {
     return this.provider.id === 'ollama' ? getOllamaKeepAliveTime() : undefined
   }
 
+  public async fakeCompletions({ onChunk }: CompletionsParams) {
+    for (let i = 0; i < 100; i++) {
+      await delay(0.01)
+      onChunk({ text: i + '\n', usage: { completion_tokens: 0, prompt_tokens: 0, total_tokens: 0 } })
+    }
+  }
+
   abstract completions({ messages, assistant, onChunk, onFilterMessages }: CompletionsParams): Promise<void>
   abstract translate(message: Message, assistant: Assistant): Promise<string>
   abstract summaries(messages: Message[], assistant: Assistant): Promise<string>

src/renderer/src/services/api.ts

@@ -2,8 +2,6 @@ import i18n from '@renderer/i18n'
 import store from '@renderer/store'
 import { setGenerating } from '@renderer/store/runtime'
 import { Assistant, Message, Provider, Suggestion, Topic } from '@renderer/types'
-import { uuid } from '@renderer/utils'
-import dayjs from 'dayjs'
 import { isEmpty } from 'lodash'
 
 import AiProvider from '../providers/AiProvider'
@@ -19,11 +17,12 @@ import { filterMessages } from './messages'
 import { estimateMessagesUsage } from './tokens'
 
 export async function fetchChatCompletion({
+  message,
   messages,
-  topic,
   assistant,
   onResponse
 }: {
+  message: Message
   messages: Message[]
   topic: Topic
   assistant: Assistant
@@ -32,23 +31,10 @@ export async function fetchChatCompletion({
   window.keyv.set(EVENT_NAMES.CHAT_COMPLETION_PAUSED, false)
 
   const provider = getAssistantProvider(assistant)
-  const defaultModel = getDefaultModel()
-  const model = assistant.model || defaultModel
   const AI = new AiProvider(provider)
 
   store.dispatch(setGenerating(true))
 
-  const message: Message = {
-    id: uuid(),
-    role: 'assistant',
-    content: '',
-    assistantId: assistant.id,
-    topicId: topic.id,
-    modelId: model.id,
-    createdAt: dayjs().format('YYYY-MM-DD HH:mm:ss'),
-    status: 'sending'
-  }
-
   onResponse({ ...message })
 
   // Handle paused state
@@ -102,6 +88,7 @@ export async function fetchChatCompletion({
 
   // Emit chat completion event
   EventEmitter.emit(EVENT_NAMES.RECEIVE_MESSAGE, message)
+  onResponse(message)
 
   // Reset generating state
   store.dispatch(setGenerating(false))

src/renderer/src/services/messages.ts

@@ -1,10 +1,11 @@
 import { DEFAULT_CONEXTCOUNT } from '@renderer/config/constant'
 import { getTopicById } from '@renderer/hooks/useTopic'
-import { Assistant, Message } from '@renderer/types'
+import { Assistant, Message, Topic } from '@renderer/types'
+import { uuid } from '@renderer/utils'
 import { isEmpty, takeRight } from 'lodash'
 import { NavigateFunction } from 'react-router'
 
-import { getAssistantById } from './assistant'
+import { getAssistantById, getDefaultModel } from './assistant'
 import { EVENT_NAMES, EventEmitter } from './event'
 import FileManager from './file'
 
@@ -48,3 +49,45 @@ export async function locateToMessage(navigate: NavigateFunction, message: Messa
   setTimeout(() => EventEmitter.emit(EVENT_NAMES.SHOW_TOPIC_SIDEBAR), 0)
   setTimeout(() => EventEmitter.emit(EVENT_NAMES.LOCATE_MESSAGE + ':' + message.id), 300)
 }
+
+export function getUserMessage({
+  assistant,
+  topic,
+  type
+}: {
+  assistant: Assistant
+  topic: Topic
+  type: Message['type']
+}): Message {
+  const defaultModel = getDefaultModel()
+  const model = assistant.model || defaultModel
+
+  return {
+    id: uuid(),
+    role: 'user',
+    content: '',
+    assistantId: assistant.id,
+    topicId: topic.id,
+    modelId: model.id,
+    createdAt: new Date().toISOString(),
+    type,
+    status: 'success'
+  }
+}
+
+export function getAssistantMessage({ assistant, topic }: { assistant: Assistant; topic: Topic }): Message {
+  const defaultModel = getDefaultModel()
+  const model = assistant.model || defaultModel
+
+  return {
+    id: uuid(),
+    role: 'assistant',
+    content: '',
+    assistantId: assistant.id,
+    topicId: topic.id,
+    modelId: model.id,
+    createdAt: new Date().toISOString(),
+    type: 'text',
+    status: 'sending'
+  }
+}

src/renderer/src/types/index.ts

@@ -43,7 +43,7 @@ export type Message = {
   files?: FileType[]
   images?: string[]
   usage?: OpenAI.Completions.CompletionUsage
-  type?: 'text' | '@' | 'clear'
+  type: 'text' | '@' | 'clear'
   isPreset?: boolean
 }