Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
merged 2 commits into from
Dec 23, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 4 additions & 2 deletions extensions/conversational-extension/src/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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>
}
/**
Expand Down Expand Up @@ -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'] },
Expand Down
20 changes: 20 additions & 0 deletions web/helpers/atoms/Thread.atom.ts
Original file line number Diff line number Diff line change
Expand Up @@ -87,13 +87,13 @@
* Get the active thread state
*/
export const activeThreadStateAtom = atom<ThreadState | undefined>((get) => {
const threadId = get(activeThreadIdAtom)
if (!threadId) {
console.debug('Active thread id is undefined')
return undefined

Check warning on line 93 in web/helpers/atoms/Thread.atom.ts

View workflow job for this annotation

GitHub Actions / coverage-check

90-93 lines are not covered with tests
}

return get(threadStatesAtom)[threadId]

Check warning on line 96 in web/helpers/atoms/Thread.atom.ts

View workflow job for this annotation

GitHub Actions / coverage-check

96 line is not covered with tests
})

/**
Expand All @@ -107,7 +107,7 @@
return undefined
}

return get(threadModelParamsAtom)[threadId]

Check warning on line 110 in web/helpers/atoms/Thread.atom.ts

View workflow job for this annotation

GitHub Actions / coverage-check

110 line is not covered with tests
}
)
/// End Active Thread Atom
Expand All @@ -125,6 +125,26 @@
*/
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) }

Check warning on line 133 in web/helpers/atoms/Thread.atom.ts

View workflow job for this annotation

GitHub Actions / coverage-check

133 line is not covered with tests

const threadState: ThreadState = {

Check warning on line 135 in web/helpers/atoms/Thread.atom.ts

View workflow job for this annotation

GitHub Actions / coverage-check

135 line is not covered with tests
hasMore: false,
waitingForResponse: false,
lastMessage: undefined,
}
currentState[newThread.id] = threadState
set(threadStatesAtom, currentState)

Check warning on line 141 in web/helpers/atoms/Thread.atom.ts

View workflow job for this annotation

GitHub Actions / coverage-check

140-141 lines are not covered with tests

// add the new thread on top of the thread list to the state
const threads = get(threadsAtom)
set(threadsAtom, [newThread, ...threads])

Check warning on line 145 in web/helpers/atoms/Thread.atom.ts

View workflow job for this annotation

GitHub Actions / coverage-check

144-145 lines are not covered with tests
})

/**
* Remove a thread state from the atom
*/
Expand All @@ -143,13 +163,13 @@
export const updateThreadWaitingForResponseAtom = atom(
null,
(get, set, threadId: string, waitingForResponse: boolean) => {
const currentState = { ...get(threadStatesAtom) }
currentState[threadId] = {

Check warning on line 167 in web/helpers/atoms/Thread.atom.ts

View workflow job for this annotation

GitHub Actions / coverage-check

166-167 lines are not covered with tests
...currentState[threadId],
waitingForResponse,
error: undefined,
}
set(threadStatesAtom, currentState)

Check warning on line 172 in web/helpers/atoms/Thread.atom.ts

View workflow job for this annotation

GitHub Actions / coverage-check

172 line is not covered with tests
}
)

Expand All @@ -159,9 +179,9 @@
export const updateThreadStateLastMessageAtom = atom(
null,
(get, set, threadId: string, lastContent?: ThreadContent[]) => {
const currentState = { ...get(threadStatesAtom) }
const lastMessage = lastContent?.[0]?.text?.value ?? ''
currentState[threadId] = {

Check warning on line 184 in web/helpers/atoms/Thread.atom.ts

View workflow job for this annotation

GitHub Actions / coverage-check

182-184 lines are not covered with tests
...currentState[threadId],
lastMessage,
}
Expand Down
19 changes: 1 addition & 18 deletions web/hooks/useCreateNewThread.ts
Original file line number Diff line number Diff line change
Expand Up @@ -5,12 +5,12 @@
ExtensionTypeEnum,
Thread,
ThreadAssistantInfo,
ThreadState,

Check warning on line 8 in web/hooks/useCreateNewThread.ts

View workflow job for this annotation

GitHub Actions / test-on-ubuntu

'ThreadState' is defined but never used

Check warning on line 8 in web/hooks/useCreateNewThread.ts

View workflow job for this annotation

GitHub Actions / test-on-macos

'ThreadState' is defined but never used

Check warning on line 8 in web/hooks/useCreateNewThread.ts

View workflow job for this annotation

GitHub Actions / test-on-windows-pr

'ThreadState' is defined but never used

Check warning on line 8 in web/hooks/useCreateNewThread.ts

View workflow job for this annotation

GitHub Actions / coverage-check

'ThreadState' is defined but never used
AssistantTool,
Model,
Assistant,
} from '@janhq/core'
import { atom, useAtom, useAtomValue, useSetAtom } from 'jotai'

Check warning on line 13 in web/hooks/useCreateNewThread.ts

View workflow job for this annotation

GitHub Actions / test-on-ubuntu

'atom' is defined but never used

Check warning on line 13 in web/hooks/useCreateNewThread.ts

View workflow job for this annotation

GitHub Actions / test-on-macos

'atom' is defined but never used

Check warning on line 13 in web/hooks/useCreateNewThread.ts

View workflow job for this annotation

GitHub Actions / test-on-windows-pr

'atom' is defined but never used

Check warning on line 13 in web/hooks/useCreateNewThread.ts

View workflow job for this annotation

GitHub Actions / coverage-check

'atom' is defined but never used

import { useDebouncedCallback } from 'use-debounce'

Expand All @@ -33,29 +33,12 @@
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()
Expand Down
20 changes: 12 additions & 8 deletions web/hooks/useDeleteThread.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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())
Expand All @@ -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 () => {
Expand Down
73 changes: 31 additions & 42 deletions web/hooks/useDeleteThread.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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
.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
const messages = 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) => {
Expand Down
Loading