Skip to content

Commit

Permalink
fix: decouple thread summary update (janhq#1994)
Browse files Browse the repository at this point in the history
  • Loading branch information
louis-menlo authored Feb 11, 2024
1 parent bcf6103 commit ac18fc1
Showing 1 changed file with 99 additions and 45 deletions.
144 changes: 99 additions & 45 deletions web/containers/Providers/EventHandler.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ import {
MessageEvent,
MessageRequestType,
ModelEvent,
Thread,
} from '@janhq/core'
import { useAtomValue, useSetAtom } from 'jotai'
import { ulid } from 'ulid'
Expand Down Expand Up @@ -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)
}
},
Expand Down Expand Up @@ -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<ConversationalExtension>(ExtensionTypeEnum.Conversational)
?.saveThread({
...thread,
})
}
},
[updateThread]
)

const updateThreadMessage = useCallback(
(message: ThreadMessage) => {
updateMessage(
message.id,
Expand Down Expand Up @@ -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,
Expand All @@ -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<ConversationalExtension>(ExtensionTypeEnum.Conversational)
?.addNewMessage(message)
}
extensionManager
.get<ConversationalExtension>(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) {
Expand Down

0 comments on commit ac18fc1

Please sign in to comment.