whai/cherry-studio
1
1
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

refactor: messages completion

Browse Source
This commit is contained in:
kangfenmao 2024-10-26 17:07:19 +08:00
parent 7c99621558
commit ee966010e1
12 changed files with 176 additions and 115 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
src/renderer/src/pages/history/components/SearchMessage.tsx
View File

@ -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"

2
src/renderer/src/pages/history/components/TopicMessages.tsx
View File

@ -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"

1
src/renderer/src/pages/home/Inputbar/Inputbar.tsx
View File

@ -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'
}

82
src/renderer/src/pages/home/Messages/Message.tsx
View File

@ -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) {
messageRef.current.scrollIntoView({ behavior: 'smooth' })
if (messageContainerRef.current) {
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

116
src/renderer/src/pages/home/Messages/Messages.tsx
View File

@ -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 { estimateHistoryTokens, estimateMessageUsage } from '@renderer/services/tokens'
import {
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') {
estimateMessageUsage(message).then((usage) => {
setMessages((prev) => {
const _messages = prev.map((m) => (m.id === message.id ? { ...m, usage } : m))
db.topics.update(topic.id, { messages: _messages })
return _messages
})
})
}
const _messages = [...messages, message]
setMessages(_messages)
db.topics.put({ id: topic.id, messages: _messages })
const assistantMessage = getAssistantMessage({ assistant, topic })
setMessages((prev) => {
const messages = prev.concat([message, assistantMessage])
db.topics.put({ id: topic.id, messages })
return messages
})
scrollToBottom()
},
[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(
(message: Message) => {
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)
const onGetMessages = useCallback(() => {
return messagesRef.current
}, [])
useEffect(() => {
const unsubscribes = [
EventEmitter.on(EVENT_NAMES.SEND_MESSAGE, async (msg: Message) => {
await onSendMessage(msg)
// 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)
EventEmitter.on(EVENT_NAMES.SEND_MESSAGE, onSendMessage),
EventEmitter.on(EVENT_NAMES.RECEIVE_MESSAGE, async () => {
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) {
onSendMessage({
...lastUserMessage,
id: uuid(),
type: '@',
modelId: model.id
})
fetchChatCompletion({
assistant,
topic,
messages: [...messages, lastUserMessage],
onResponse: setLastMessage
})
}
lastUserMessage && onSendMessage({ ...lastUserMessage, id: uuid(), type: '@', modelId: model.id })
}),
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({
id: uuid(),
assistantId: assistant.id,
role: 'user',
content: '',
topicId: topic.id,
createdAt: new Date().toISOString(),
status: 'success',
type: 'clear'
} as Message)
setMessages((prev) => {
const messages = prev.concat([getUserMessage({ assistant, topic, type: 'clear' })])
db.topics.put({ id: topic.id, messages })
return messages
})
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} />
{lastMessage && <MessageItem key={lastMessage.id} message={lastMessage} lastMessage />}
{reverse([...messages]).map((message, index) => (
<Suggestions assistant={assistant} messages={messages} />
{memoizedMessages.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} />

7
src/renderer/src/pages/home/components/Suggestions.tsx
View File

@ -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
}

1
src/renderer/src/pages/translate/TranslatePage.tsx
View File

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

4
src/renderer/src/providers/AiProvider.ts
View File

@ -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 })
}

8
src/renderer/src/providers/BaseProvider.ts
View File

@ -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>

19
src/renderer/src/services/api.ts
View File

@ -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))

47
src/renderer/src/services/messages.ts
View File

@ -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'
}
}

2
src/renderer/src/types/index.ts
View File

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