From ac18fc11817353cb3907566e7f188ab363e46777 Mon Sep 17 00:00:00 2001 From: Louis Date: Sun, 11 Feb 2024 20:22:12 +0700 Subject: [PATCH] fix: decouple thread summary update (#1994) --- web/containers/Providers/EventHandler.tsx | 144 +++++++++++++++------- 1 file changed, 99 insertions(+), 45 deletions(-) diff --git a/web/containers/Providers/EventHandler.tsx b/web/containers/Providers/EventHandler.tsx index 7f8bd261c..28a47c284 100644 --- a/web/containers/Providers/EventHandler.tsx +++ b/web/containers/Providers/EventHandler.tsx @@ -14,6 +14,7 @@ import { MessageEvent, MessageRequestType, ModelEvent, + Thread, } from '@janhq/core' import { useAtomValue, useSetAtom } from 'jotai' import { ulid } from 'ulid' @@ -80,39 +81,7 @@ export default function EventHandler({ children }: { children: ReactNode }) { const onNewMessageResponse = useCallback( (message: ThreadMessage) => { - const thread = threadsRef.current?.find((e) => e.id == message.thread_id) - // If this is the first ever prompt in the thread - if (thread && thread.title.trim() == 'New Thread') { - // This is the first time message comes in on a new thread - // Summarize the first message, and make that the title of the Thread - // 1. Get the summary of the first prompt using whatever engine user is currently using - const firstPrompt = messagesRef?.current[0].content[0].text.value.trim() - const summarizeFirstPrompt = - 'Summarize "' + firstPrompt + '" in 5 words as a title' - - // Prompt: Given this query from user {query}, return to me the summary in 5 words as the title - const msgId = ulid() - const messages: ChatCompletionMessage[] = [ - { - role: ChatCompletionRole.User, - content: summarizeFirstPrompt, - } as ChatCompletionMessage, - ] - - const firstPromptRequest: MessageRequest = { - id: msgId, - threadId: message.thread_id, - type: MessageRequestType.Summary, - messages, - model: activeModelRef?.current, - } - - // 2. Update the title with the result of the inference - // the title will be updated as part of the `EventName.OnFirstPromptUpdate` - events.emit(MessageEvent.OnMessageSent, firstPromptRequest) - } - - if (message.type !== MessageRequestType.Summary) { + if (message.type === MessageRequestType.Thread) { addNewMessage(message) } }, @@ -158,7 +127,34 @@ export default function EventHandler({ children }: { children: ReactNode }) { [setStateModel, setQueuedMessage, setLoadModelError] ) - const onMessageResponseUpdate = useCallback( + const updateThreadTitle = useCallback( + (message: ThreadMessage) => { + // Update only when it's finished + if (message.status === MessageStatus.Pending) { + return + } + + const thread = threadsRef.current?.find((e) => e.id == message.thread_id) + const messageContent = message.content[0]?.text?.value + if (thread && messageContent) { + // Update the Thread title with the response of the inference on the 1st prompt + updateThread({ + ...thread, + title: messageContent, + metadata: thread.metadata, + }) + + extensionManager + .get(ExtensionTypeEnum.Conversational) + ?.saveThread({ + ...thread, + }) + } + }, + [updateThread] + ) + + const updateThreadMessage = useCallback( (message: ThreadMessage) => { updateMessage( message.id, @@ -186,11 +182,6 @@ export default function EventHandler({ children }: { children: ReactNode }) { ...(messageContent && { lastMessage: messageContent }), } - // Update the Thread title with the response of the inference on the 1st prompt - if (message.type === MessageRequestType.Summary) { - thread.title = messageContent - } - updateThread({ ...thread, metadata, @@ -204,16 +195,79 @@ export default function EventHandler({ children }: { children: ReactNode }) { }) // If this is not the summary of the Thread, don't need to add it to the Thread - if (message.type !== MessageRequestType.Summary) { - extensionManager - .get(ExtensionTypeEnum.Conversational) - ?.addNewMessage(message) - } + extensionManager + .get(ExtensionTypeEnum.Conversational) + ?.addNewMessage(message) + + // Attempt to generate the title of the Thread when needed + generateThreadTitle(message, thread) } }, - [updateMessage, updateThreadWaiting, setIsGeneratingResponse, updateThread] + [setIsGeneratingResponse, updateMessage, updateThread, updateThreadWaiting] ) + const onMessageResponseUpdate = useCallback( + (message: ThreadMessage) => { + switch (message.type) { + case MessageRequestType.Summary: + updateThreadTitle(message) + break + default: + updateThreadMessage(message) + } + }, + [updateThreadMessage, updateThreadTitle] + ) + + const generateThreadTitle = (message: ThreadMessage, thread: Thread) => { + // If this is the first ever prompt in the thread + if ( + thread && + thread.title?.trim() === 'New Thread' && + activeModelRef.current + ) { + // This is the first time message comes in on a new thread + // Summarize the first message, and make that the title of the Thread + // 1. Get the summary of the first prompt using whatever engine user is currently using + const threadMessages = messagesRef?.current + const summarizeFirstPrompt = + 'Summarize the conversation above in 5 words as a title' + + // Prompt: Given this query from user {query}, return to me the summary in 5 words as the title + const msgId = ulid() + const messages: ChatCompletionMessage[] = [ + ...threadMessages.map((msg) => { + return { + role: msg.role, + content: msg.content[0]?.text.value, + } + }), + { + role: ChatCompletionRole.User, + content: summarizeFirstPrompt, + } as ChatCompletionMessage, + ] + + const firstPromptRequest: MessageRequest = { + id: msgId, + threadId: message.thread_id, + type: MessageRequestType.Summary, + messages, + model: { + ...activeModelRef.current, + parameters: { + stream: false, + }, + }, + } + + // 2. Update the title with the result of the inference + setTimeout(() => { + events.emit(MessageEvent.OnMessageSent, firstPromptRequest) + }, 1000) + } + } + useEffect(() => { console.log('Registering events') if (window.core?.events) {