65 lines
1.8 KiB
TypeScript
65 lines
1.8 KiB
TypeScript
import { create } from 'zustand'
|
|
import { ThreadMessage } from '@janhq/core'
|
|
import { createJSONStorage, persist } from 'zustand/middleware'
|
|
import { localStoregeKey } from '@/constants/localStorage'
|
|
import {
|
|
createMessage,
|
|
deleteMessage as deleteMessageExt,
|
|
} from '@/services/messages'
|
|
|
|
type MessageState = {
|
|
messages: Record<string, ThreadMessage[]>
|
|
getMessages: (threadId: string) => ThreadMessage[]
|
|
setMessages: (threadId: string, messages: ThreadMessage[]) => void
|
|
addMessage: (message: ThreadMessage) => void
|
|
deleteMessage: (threadId: string, messageId: string) => void
|
|
}
|
|
|
|
export const useMessages = create<MessageState>()(
|
|
persist(
|
|
(set, get) => ({
|
|
messages: {},
|
|
getMessages: (threadId) => {
|
|
return get().messages[threadId] || []
|
|
},
|
|
setMessages: (threadId, messages) => {
|
|
set((state) => ({
|
|
messages: {
|
|
...state.messages,
|
|
[threadId]: messages,
|
|
},
|
|
}))
|
|
},
|
|
addMessage: (message) => {
|
|
createMessage(message).then((createdMessage) => {
|
|
set((state) => ({
|
|
messages: {
|
|
...state.messages,
|
|
[message.thread_id]: [
|
|
...(state.messages[message.thread_id] || []),
|
|
createdMessage,
|
|
],
|
|
},
|
|
}))
|
|
})
|
|
},
|
|
deleteMessage: (threadId, messageId) => {
|
|
deleteMessageExt(threadId, messageId)
|
|
set((state) => ({
|
|
messages: {
|
|
...state.messages,
|
|
[threadId]:
|
|
state.messages[threadId]?.filter(
|
|
(message) => message.id !== messageId
|
|
) || [],
|
|
},
|
|
}))
|
|
},
|
|
}),
|
|
{
|
|
name: localStoregeKey.messages,
|
|
storage: createJSONStorage(() => localStorage),
|
|
}
|
|
)
|
|
)
|