fix: cleaning a thread should just clear out messages (#4316)

* fix: cleaning a thread should just clear out messages

* chore: clean up
This commit is contained in:
Louis 2024-12-23 12:35:30 +07:00 committed by GitHub
parent abb718c57f
commit 24222679fb
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 68 additions and 70 deletions

View File

@ -147,7 +147,9 @@ export default class CortexConversationalExtension extends ConversationalExtensi
*/
async getThreadAssistant(threadId: string): Promise<ThreadAssistantInfo> {
return this.queue.add(() =>
ky.get(`${API_URL}/v1/assistants/${threadId}?limit=-1`).json<ThreadAssistantInfo>()
ky
.get(`${API_URL}/v1/assistants/${threadId}?limit=-1`)
.json<ThreadAssistantInfo>()
) as Promise<ThreadAssistantInfo>
}
/**
@ -188,7 +190,7 @@ export default class CortexConversationalExtension extends ConversationalExtensi
* Do health check on cortex.cpp
* @returns
*/
healthz(): Promise<void> {
async healthz(): Promise<void> {
return ky
.get(`${API_URL}/healthz`, {
retry: { limit: 20, delay: () => 500, methods: ['get'] },

View File

@ -125,6 +125,26 @@ export const waitingToSendMessage = atom<boolean | undefined>(undefined)
*/
export const isGeneratingResponseAtom = atom<boolean | undefined>(undefined)
/**
* Create a new thread and add it to the thread list
*/
export const createNewThreadAtom = atom(null, (get, set, newThread: Thread) => {
// create thread state for this new thread
const currentState = { ...get(threadStatesAtom) }
const threadState: ThreadState = {
hasMore: false,
waitingForResponse: false,
lastMessage: undefined,
}
currentState[newThread.id] = threadState
set(threadStatesAtom, currentState)
// add the new thread on top of the thread list to the state
const threads = get(threadsAtom)
set(threadsAtom, [newThread, ...threads])
})
/**
* Remove a thread state from the atom
*/

View File

@ -33,29 +33,12 @@ import { activeAssistantAtom } from '@/helpers/atoms/Assistant.atom'
import { selectedModelAtom } from '@/helpers/atoms/Model.atom'
import {
threadsAtom,
threadStatesAtom,
updateThreadAtom,
setThreadModelParamsAtom,
isGeneratingResponseAtom,
createNewThreadAtom,
} from '@/helpers/atoms/Thread.atom'
const createNewThreadAtom = atom(null, (get, set, newThread: Thread) => {
// create thread state for this new thread
const currentState = { ...get(threadStatesAtom) }
const threadState: ThreadState = {
hasMore: false,
waitingForResponse: false,
lastMessage: undefined,
}
currentState[newThread.id] = threadState
set(threadStatesAtom, currentState)
// add the new thread on top of the thread list to the state
const threads = get(threadsAtom)
set(threadsAtom, [newThread, ...threads])
})
export const useCreateNewThread = () => {
const createNewThread = useSetAtom(createNewThreadAtom)
const { setActiveThread } = useSetActiveThread()

View File

@ -55,17 +55,21 @@ describe('useDeleteThread', () => {
const mockCleanMessages = jest.fn()
;(useSetAtom as jest.Mock).mockReturnValue(() => mockCleanMessages)
;(useAtomValue as jest.Mock).mockReturnValue(['thread 1'])
const mockCreateNewThread = jest.fn()
;(useCreateNewThread as jest.Mock).mockReturnValue({
requestCreateNewThread: mockCreateNewThread,
})
const mockSaveThread = jest.fn()
const mockDeleteThread = jest.fn().mockResolvedValue({})
const mockDeleteMessage = jest.fn().mockResolvedValue({})
const mockModifyThread = jest.fn().mockResolvedValue({})
extensionManager.get = jest.fn().mockReturnValue({
saveThread: mockSaveThread,
getThreadAssistant: jest.fn().mockResolvedValue({}),
deleteThread: mockDeleteThread,
listMessages: jest.fn().mockResolvedValue([
{
id: 'message1',
text: 'Message 1',
},
]),
deleteMessage: mockDeleteMessage,
modifyThread: mockModifyThread,
})
const { result } = renderHook(() => useDeleteThread())
@ -74,8 +78,8 @@ describe('useDeleteThread', () => {
await result.current.cleanThread('thread1')
})
expect(mockDeleteThread).toHaveBeenCalled()
expect(mockCreateNewThread).toHaveBeenCalled()
expect(mockDeleteMessage).toHaveBeenCalled()
expect(mockModifyThread).toHaveBeenCalled()
})
it('should handle errors when deleting a thread', async () => {

View File

@ -2,30 +2,25 @@ import { useCallback } from 'react'
import { ExtensionTypeEnum, ConversationalExtension } from '@janhq/core'
import { useAtom, useAtomValue, useSetAtom } from 'jotai'
import { useAtom, useSetAtom } from 'jotai'
import { currentPromptAtom } from '@/containers/Providers/Jotai'
import { toaster } from '@/containers/Toast'
import { useCreateNewThread } from './useCreateNewThread'
import { extensionManager } from '@/extension/ExtensionManager'
import { assistantsAtom } from '@/helpers/atoms/Assistant.atom'
import { deleteChatMessageAtom as deleteChatMessagesAtom } from '@/helpers/atoms/ChatMessage.atom'
import { downloadedModelsAtom } from '@/helpers/atoms/Model.atom'
import {
threadsAtom,
setActiveThreadIdAtom,
deleteThreadStateAtom,
updateThreadAtom,
} from '@/helpers/atoms/Thread.atom'
export default function useDeleteThread() {
const [threads, setThreads] = useAtom(threadsAtom)
const { requestCreateNewThread } = useCreateNewThread()
const assistants = useAtomValue(assistantsAtom)
const models = useAtomValue(downloadedModelsAtom)
const updateThread = useSetAtom(updateThreadAtom)
const setCurrentPrompt = useSetAtom(currentPromptAtom)
const setActiveThreadId = useSetAtom(setActiveThreadIdAtom)
@ -35,43 +30,37 @@ export default function useDeleteThread() {
const cleanThread = useCallback(
async (threadId: string) => {
const thread = threads.find((c) => c.id === threadId)
if (!thread) return
const availableThreads = threads.filter((c) => c.id !== threadId)
setThreads(availableThreads)
// delete the thread state
deleteThreadState(threadId)
const assistantInfo = await extensionManager
const messages = await extensionManager
.get<ConversationalExtension>(ExtensionTypeEnum.Conversational)
?.getThreadAssistant(thread.id)
.catch(console.error)
if (!assistantInfo) return
const model = models.find((c) => c.id === assistantInfo?.model?.id)
requestCreateNewThread(
{
...assistantInfo,
id: assistants[0].id,
name: assistants[0].name,
},
model
? {
...model,
parameters: assistantInfo?.model?.parameters ?? {},
settings: assistantInfo?.model?.settings ?? {},
}
: undefined
)
// Delete this thread
await extensionManager
.get<ConversationalExtension>(ExtensionTypeEnum.Conversational)
?.deleteThread(threadId)
?.listMessages(threadId)
.catch(console.error)
if (messages) {
messages.forEach((message) => {
extensionManager
.get<ConversationalExtension>(ExtensionTypeEnum.Conversational)
?.deleteMessage(threadId, message.id)
.catch(console.error)
})
const thread = threads.find((e) => e.id === threadId)
if (thread) {
const updatedThread = {
...thread,
metadata: {
...thread.metadata,
title: 'New Thread',
lastMessage: '',
},
}
extensionManager
.get<ConversationalExtension>(ExtensionTypeEnum.Conversational)
?.modifyThread(updatedThread)
.catch(console.error)
updateThread(updatedThread)
}
}
deleteMessages(threadId)
},
[assistants, models, requestCreateNewThread, threads]
[deleteMessages, threads, updateThread]
)
const deleteThread = async (threadId: string) => {