88 lines
2 KiB
TypeScript
88 lines
2 KiB
TypeScript
import { atom, getDefaultStore } from "jotai"
|
|
import { atomWithStorage } from "jotai/utils"
|
|
|
|
import {
|
|
getInitialActiveSessionId,
|
|
writeStoredSessionId,
|
|
} from "@/features/chat/state"
|
|
|
|
export interface ChatAttachment {
|
|
type: "image" | "audio" | "video" | "file"
|
|
url: string
|
|
filename?: string
|
|
contentType?: string
|
|
}
|
|
|
|
export type AssistantMessageKind = "normal" | "thought"
|
|
|
|
export interface ChatMessage {
|
|
id: string
|
|
role: "user" | "assistant"
|
|
content: string
|
|
timestamp: number | string
|
|
kind?: AssistantMessageKind
|
|
attachments?: ChatAttachment[]
|
|
}
|
|
|
|
export interface ContextUsage {
|
|
used_tokens: number
|
|
total_tokens: number
|
|
compress_at_tokens: number
|
|
used_percent: number
|
|
}
|
|
|
|
export type ConnectionState =
|
|
| "disconnected"
|
|
| "connecting"
|
|
| "connected"
|
|
| "error"
|
|
|
|
export interface ChatStoreState {
|
|
messages: ChatMessage[]
|
|
connectionState: ConnectionState
|
|
isTyping: boolean
|
|
activeSessionId: string
|
|
hasHydratedActiveSession: boolean
|
|
contextUsage?: ContextUsage
|
|
}
|
|
|
|
type ChatStorePatch = Partial<ChatStoreState>
|
|
|
|
const SHOW_THOUGHTS_STORAGE_KEY = "picoclaw:chat-show-thoughts"
|
|
|
|
const DEFAULT_CHAT_STATE: ChatStoreState = {
|
|
messages: [],
|
|
connectionState: "disconnected",
|
|
isTyping: false,
|
|
activeSessionId: getInitialActiveSessionId(),
|
|
hasHydratedActiveSession: false,
|
|
}
|
|
|
|
export const chatAtom = atom<ChatStoreState>(DEFAULT_CHAT_STATE)
|
|
export const showThoughtsAtom = atomWithStorage<boolean>(
|
|
SHOW_THOUGHTS_STORAGE_KEY,
|
|
true,
|
|
)
|
|
|
|
const store = getDefaultStore()
|
|
|
|
export function getChatState() {
|
|
return store.get(chatAtom)
|
|
}
|
|
|
|
export function updateChatStore(
|
|
patch:
|
|
| ChatStorePatch
|
|
| ((prev: ChatStoreState) => ChatStorePatch | ChatStoreState),
|
|
) {
|
|
store.set(chatAtom, (prev) => {
|
|
const nextPatch = typeof patch === "function" ? patch(prev) : patch
|
|
const next = { ...prev, ...nextPatch }
|
|
|
|
if (next.activeSessionId !== prev.activeSessionId) {
|
|
writeStoredSessionId(next.activeSessionId)
|
|
}
|
|
|
|
return next
|
|
})
|
|
}
|