From 4f78a8aa4872cae8e4d275cbca7c9d02aab70525 Mon Sep 17 00:00:00 2001 From: Innei Date: Thu, 31 Jul 2025 00:16:40 +0800 Subject: [PATCH] refactor: ai module and implement rich text (#4282) --- apps/desktop/layer/renderer/package.json | 4 + .../ai/chat/__internal__/AIChatContext.ts | 23 +- .../src/modules/ai/chat/__internal__/hooks.ts | 104 ++ .../chat/__internal__/slices/block.slice.ts | 109 ++ .../ai/chat/__internal__/slices/chat.slice.ts | 544 ++++++++++ .../ai/chat/__internal__/slices/index.ts | 5 + .../src/modules/ai/chat/__internal__/store.ts | 216 +--- .../modules/ai/chat/__internal__/transport.ts | 14 + .../src/modules/ai/chat/__internal__/types.ts | 7 +- .../src/modules/ai/chat/__internal__/utils.ts | 72 -- .../src/modules/ai/chat/atoms/session.ts | 42 - .../ai/chat/components/AIChatContextBar.tsx | 16 +- .../modules/ai/chat/components/AIChatRoot.tsx | 216 +--- .../ai/chat/components/CollapsibleError.tsx | 100 -- .../components/ToolInvocationComponent.tsx | 2 +- .../{ => layouts}/AIChatSendButton.tsx | 0 .../chat/components/layouts}/ChatHeader.tsx | 29 +- .../components/{ => layouts}/ChatInput.tsx | 74 +- .../components/layouts}/ChatInterface.tsx | 81 +- .../components/layouts}/ChatMoreDropdown.tsx | 49 +- .../components/layouts/CollapsibleError.tsx | 229 +++++ .../components/layouts/LexicalRichEditor.tsx | 227 +++++ .../{ => layouts}/WelcomeScreen.tsx | 10 +- .../chat/components/message/AIChatMessage.tsx | 27 +- .../components/message/AIDataBlockPart.tsx | 199 ++++ .../components/message/AIMessageParts.tsx | 21 +- .../components/message/EditableMessage.tsx | 6 +- .../modules/ai/chat/hooks/useChatHistory.ts | 14 +- .../modules/ai/chat/hooks/useLoadMessages.ts | 26 +- .../modules/ai/chat/hooks/useSaveMessages.ts | 15 +- .../src/modules/ai/chat/services/index.ts | 357 ++++++- .../src/modules/ai/chat/types/ChatSession.ts | 2 +- .../modules/ai/chat/utils/lexical-markdown.ts | 86 ++ .../modules/app-layout/ai/AIChatLayout.tsx | 13 +- .../entry-content/EntryLayoutContent.tsx | 11 +- .../entry-column/EntryColumnLayout.tsx | 8 +- .../entry-column/components/DateItem.tsx | 10 +- .../entry-column/layouts/EntryItemWrapper.tsx | 2 +- .../entry-content/EntryContent.ai.tsx | 2 +- .../entry-content/EntryContent.legacy.tsx | 2 - .../entry-content/ai/AIChatBottom.tsx | 31 - .../entry-content/ai/AIChatContainer.tsx | 198 ---- .../entry-content/ai/AIChatInput.tsx | 128 --- .../entry-content/ai/AIPanelHeader.tsx | 61 -- .../components/entry-content/ai/context.ts | 13 - .../components/entry-content/ai/index.tsx | 326 ------ .../components/layouts/ArticleLayout.tsx | 4 +- eslint.config.mjs | 1 + packages/internal/database/src/db.ts | 5 + .../src/drizzle/0033_shiny_sebastian_shaw.sql | 25 + .../src/drizzle/meta/0033_snapshot.json | 950 ++++++++++++++++++ .../database/src/drizzle/meta/_journal.json | 7 + .../database/src/drizzle/migrations.js | 2 + .../internal/database/src/schemas/index.ts | 117 ++- packages/internal/tracker/src/op/index.ts | 1 - pnpm-lock.yaml | 320 +++++- 56 files changed, 3523 insertions(+), 1640 deletions(-) create mode 100644 apps/desktop/layer/renderer/src/modules/ai/chat/__internal__/hooks.ts create mode 100644 apps/desktop/layer/renderer/src/modules/ai/chat/__internal__/slices/block.slice.ts create mode 100644 apps/desktop/layer/renderer/src/modules/ai/chat/__internal__/slices/chat.slice.ts create mode 100644 apps/desktop/layer/renderer/src/modules/ai/chat/__internal__/slices/index.ts create mode 100644 apps/desktop/layer/renderer/src/modules/ai/chat/__internal__/transport.ts delete mode 100644 apps/desktop/layer/renderer/src/modules/ai/chat/__internal__/utils.ts delete mode 100644 apps/desktop/layer/renderer/src/modules/ai/chat/components/CollapsibleError.tsx rename apps/desktop/layer/renderer/src/modules/ai/chat/components/{ => layouts}/AIChatSendButton.tsx (100%) rename apps/desktop/layer/renderer/src/modules/{app-layout/ai/components => ai/chat/components/layouts}/ChatHeader.tsx (72%) rename apps/desktop/layer/renderer/src/modules/ai/chat/components/{ => layouts}/ChatInput.tsx (55%) rename apps/desktop/layer/renderer/src/modules/{app-layout/ai/components => ai/chat/components/layouts}/ChatInterface.tsx (76%) rename apps/desktop/layer/renderer/src/modules/{app-layout/ai/components => ai/chat/components/layouts}/ChatMoreDropdown.tsx (80%) create mode 100644 apps/desktop/layer/renderer/src/modules/ai/chat/components/layouts/CollapsibleError.tsx create mode 100644 apps/desktop/layer/renderer/src/modules/ai/chat/components/layouts/LexicalRichEditor.tsx rename apps/desktop/layer/renderer/src/modules/ai/chat/components/{ => layouts}/WelcomeScreen.tsx (93%) create mode 100644 apps/desktop/layer/renderer/src/modules/ai/chat/components/message/AIDataBlockPart.tsx create mode 100644 apps/desktop/layer/renderer/src/modules/ai/chat/utils/lexical-markdown.ts delete mode 100644 apps/desktop/layer/renderer/src/modules/entry-content/components/entry-content/ai/AIChatBottom.tsx delete mode 100644 apps/desktop/layer/renderer/src/modules/entry-content/components/entry-content/ai/AIChatContainer.tsx delete mode 100644 apps/desktop/layer/renderer/src/modules/entry-content/components/entry-content/ai/AIChatInput.tsx delete mode 100644 apps/desktop/layer/renderer/src/modules/entry-content/components/entry-content/ai/AIPanelHeader.tsx delete mode 100644 apps/desktop/layer/renderer/src/modules/entry-content/components/entry-content/ai/context.ts delete mode 100644 apps/desktop/layer/renderer/src/modules/entry-content/components/entry-content/ai/index.tsx create mode 100644 packages/internal/database/src/drizzle/0033_shiny_sebastian_shaw.sql create mode 100644 packages/internal/database/src/drizzle/meta/0033_snapshot.json diff --git a/apps/desktop/layer/renderer/package.json b/apps/desktop/layer/renderer/package.json index 85841d36f..2b84fb6ab 100644 --- a/apps/desktop/layer/renderer/package.json +++ b/apps/desktop/layer/renderer/package.json @@ -5,6 +5,7 @@ "main": "./dist/main/index.js", "scripts": { "build:web": "cd ../.. && pnpm build:web", + "db:generate": "pnpm --filter @follow/database run generate", "dev": "cd ../.. && pnpm dev:web", "dev:ssl": "cd ../.. && SSL=true pnpm dev:web", "generate-pwa-assets": "pwa-assets-generator public/icon.svg", @@ -28,6 +29,8 @@ "@hcaptcha/react-hcaptcha": "1.12.0", "@headlessui/react": "2.2.4", "@hookform/resolvers": "5.1.1", + "@lexical/markdown": "0.33.1", + "@lexical/react": "0.33.1", "@lottiefiles/dotlottie-react": "0.14.2", "@openpanel/web": "1.0.1", "@radix-ui/react-avatar": "1.1.10", @@ -73,6 +76,7 @@ "immer": "10.1.1", "jotai": "2.12.5", "lethargy": "1.0.9", + "lexical": "0.33.1", "masonic": "4.1.0", "mdast-util-gfm-table": "2.0.0", "mdast-util-to-markdown": "2.1.2", diff --git a/apps/desktop/layer/renderer/src/modules/ai/chat/__internal__/AIChatContext.ts b/apps/desktop/layer/renderer/src/modules/ai/chat/__internal__/AIChatContext.ts index 43ad7d9eb..81bae0ffb 100644 --- a/apps/desktop/layer/renderer/src/modules/ai/chat/__internal__/AIChatContext.ts +++ b/apps/desktop/layer/renderer/src/modules/ai/chat/__internal__/AIChatContext.ts @@ -1,15 +1,8 @@ -import type { UIMessage, UseChatHelpers } from "@ai-sdk/react" -import type { UIDataTypes } from "ai" import { createContext, use } from "react" import type { StoreApi } from "zustand" import type { UseBoundStoreWithEqualityFn } from "zustand/traditional" -import type { AiChatContextStore } from "./store" -import type { BizUIMetadata, BizUITools } from "./types" - -export const AIChatContext = createContext< - UseChatHelpers> ->(null!) +import type { AiChatStore } from "./store" export type AIPanelRefs = { panelRef: React.RefObject @@ -18,15 +11,14 @@ export type AIPanelRefs = { export const AIPanelRefsContext = createContext(null!) -export const AIChatContextStoreContext = createContext< - UseBoundStoreWithEqualityFn> ->(null!) +export const AIChatStoreContext = createContext>>( + null!, +) -// Hook to access AI chat context information export const useAIChatStore = () => { - const store = use(AIChatContextStoreContext) + const store = use(AIChatStoreContext) if (!store && import.meta.env.DEV) { - throw new Error("useAIChatStore must be used within a AIChatContextStoreContext") + throw new Error("useAIChatStore must be used within a AIChatStoreContext") } return store } @@ -34,9 +26,8 @@ export const useAIChatStore = () => { // Session methods context for managing chat session actions export interface AIChatSessionMethods { handleTitleGenerated: (title: string) => Promise - handleFirstMessage: () => Promise + handleNewChat: () => void - handleSwitchRoom: (roomId: string) => Promise } export const AIChatSessionMethodsContext = createContext(null!) diff --git a/apps/desktop/layer/renderer/src/modules/ai/chat/__internal__/hooks.ts b/apps/desktop/layer/renderer/src/modules/ai/chat/__internal__/hooks.ts new file mode 100644 index 000000000..b370e100c --- /dev/null +++ b/apps/desktop/layer/renderer/src/modules/ai/chat/__internal__/hooks.ts @@ -0,0 +1,104 @@ +import { useAIChatStore } from "./AIChatContext" + +/** + * Hook to get the current room ID (chat ID) from the AI chat store + */ +export const useCurrentChatId = () => { + const store = useAIChatStore() + return store((state) => state.chatId) +} + +/** + * Hook to get the current chat title from the AI chat store + */ +export const useCurrentTitle = () => { + const store = useAIChatStore() + return store((state) => state.currentTitle) +} + +/** + * Hook to get the setter for current room ID + */ +export const useSetCurrentChatId = () => { + const store = useAIChatStore() + return (chatId: string | null) => { + const actions = store.getState().chatActions + if (chatId && chatId !== actions.getCurrentChatId()) { + // If we need to switch to a different room, create a new chat + actions.newChat() + } + } +} + +/** + * Hook to get the setter for current title + */ +export const useSetCurrentTitle = () => { + const store = useAIChatStore() + return store.getState().chatActions.setCurrentTitle +} + +/** + * Hook to get the chat actions + */ +export const useChatActions = () => { + const store = useAIChatStore() + return store((state) => state.chatActions) +} + +/** + * Hook to get the block actions + */ +export const useBlockActions = () => { + const store = useAIChatStore() + return store((state) => state.blockActions) +} + +/** + * Hook to get the chat instance + */ +export const useChatInstance = () => { + const store = useAIChatStore() + return store((state) => state.chatInstance) +} + +/** + * Hook to get the current messages + */ +export const useMessages = () => { + const store = useAIChatStore() + return store((state) => state.messages) +} + +/** + * Hook to check if the chat has messages + */ +export const useHasMessages = () => { + const store = useAIChatStore() + return store((state) => state.messages.length > 0) +} + +export const useChatBlockActions = () => useAIChatStore()((state) => state.blockActions) +/** + * Hook to get the chat status + */ +export const useChatStatus = () => { + const store = useAIChatStore() + return store((state) => state.status) +} + +/** + * Hook to get the chat error + */ +export const useChatError = () => { + const store = useAIChatStore() + return store((state) => state.error) +} + +/** + * Hook to get the streaming status + */ +export const useIsStreaming = () => { + const store = useAIChatStore() + return store((state) => state.isStreaming) +} diff --git a/apps/desktop/layer/renderer/src/modules/ai/chat/__internal__/slices/block.slice.ts b/apps/desktop/layer/renderer/src/modules/ai/chat/__internal__/slices/block.slice.ts new file mode 100644 index 000000000..a5615b9ae --- /dev/null +++ b/apps/desktop/layer/renderer/src/modules/ai/chat/__internal__/slices/block.slice.ts @@ -0,0 +1,109 @@ +import { produce } from "immer" +import { nanoid } from "nanoid" +import type { StateCreator } from "zustand" + +import type { AIChatContextBlock } from "../types" + +export interface BlockSlice { + blocks: AIChatContextBlock[] + blockActions: BlockSliceAction +} + +export const createBlockSlice: ( + initialBlocks?: AIChatContextBlock[], +) => StateCreator = + (initialBlocks?: AIChatContextBlock[]) => + (...params) => { + const defaultBlocks: AIChatContextBlock[] = initialBlocks || [] + + return { + blocks: defaultBlocks, + blockActions: new BlockSliceAction(params), + } + } + +class BlockSliceAction { + constructor(private params: Parameters>) {} + + get set() { + return this.params[0] + } + + get get() { + return this.params[1] + } + addBlock(block: Omit) { + const currentBlocks = this.get().blocks + + // Only allow one mainEntry + if (block.type === "mainEntry" && currentBlocks.some((b) => b.type === "mainEntry")) { + return + } + + // Only allow one selectedText + if (block.type === "selectedText" && currentBlocks.some((b) => b.type === "selectedText")) { + return + } + + // Prevent duplicate referEntry or referFeed + if ( + block.type === "referEntry" && + block.value && + currentBlocks.some((b) => b.type === "referEntry" && b.value === block.value) + ) { + return + } + + if ( + block.type === "referFeed" && + block.value && + currentBlocks.some((b) => b.type === "referFeed" && b.value === block.value) + ) { + return + } + + this.set( + produce((state: BlockSlice) => { + state.blocks.push({ ...block, id: nanoid(8) }) + }), + ) + } + + removeBlock(id: string) { + this.set( + produce((state: BlockSlice) => { + state.blocks = state.blocks.filter((block) => block.id !== id) + }), + ) + } + + updateBlock(id: string, updates: Partial) { + this.set( + produce((state: BlockSlice) => { + state.blocks = state.blocks.map((block) => + block.id === id ? { ...block, ...updates } : block, + ) + }), + ) + } + + clearBlocks() { + this.set( + produce((state: BlockSlice) => { + state.blocks = [] + }), + ) + } + + resetContext() { + this.set( + produce((state: BlockSlice) => { + state.blocks = [] + }), + ) + } + + getBlocks() { + return this.get().blocks + } +} diff --git a/apps/desktop/layer/renderer/src/modules/ai/chat/__internal__/slices/chat.slice.ts b/apps/desktop/layer/renderer/src/modules/ai/chat/__internal__/slices/chat.slice.ts new file mode 100644 index 000000000..cb1551c22 --- /dev/null +++ b/apps/desktop/layer/renderer/src/modules/ai/chat/__internal__/slices/chat.slice.ts @@ -0,0 +1,544 @@ +import type { ChatInit, ChatState, ChatStatus } from "ai" +import { AbstractChat } from "ai" +import { throttle } from "es-toolkit/compat" +import { nanoid } from "nanoid" +import type { StateCreator } from "zustand" + +import { AIPersistService } from "../../services" +import { generateChatTitle } from "../../utils/titleGeneration" +import { createChatTransport } from "../transport" +import type { BizUIMessage } from "../types" + +// Event types and payloads +interface ChatStateEvents { + messages: { messages: UI_MESSAGE[] } + status: { status: ChatStatus } + error: { error: Error | undefined } +} + +type ChatStateEventType = keyof ChatStateEvents + +// Event emitter for AI chat state changes with typed payloads +class ChatStateEventEmitter { + #listeners = new Map void>>() + + on( + event: T, + listener: (payload: ChatStateEvents[T]) => void, + ): () => void { + if (!this.#listeners.has(event)) { + this.#listeners.set(event, new Set()) + } + this.#listeners.get(event)!.add(listener) + + // Return unsubscribe function + return () => { + this.#listeners.get(event)?.delete(listener) + } + } + + emit(event: T, payload: ChatStateEvents[T]): void { + this.#listeners.get(event)?.forEach((listener) => { + try { + listener(payload) + } catch (error) { + console.error(`Error in chat state listener for event ${event}:`, error) + } + }) + } + + clear(): void { + this.#listeners.clear() + } +} + +// Zustand Chat State that implements AI SDK ChatState interface +class ZustandChatState implements ChatState { + #messages: UI_MESSAGE[] + #status: ChatStatus = "ready" + #error: Error | undefined = undefined + #eventEmitter = new ChatStateEventEmitter() + + constructor( + initialMessages: UI_MESSAGE[] = [], + private updateZustandState: (updater: (state: ChatSlice) => ChatSlice) => void, + private chatId: string, + ) { + this.#messages = initialMessages + this.#setupEventHandlers() + } + + #setupEventHandlers(): void { + // Setup event handlers for automatic Zustand synchronization + this.#eventEmitter.on("messages", ({ messages }) => { + this.updateZustandState((state) => ({ + ...state, + messages: [...messages], + })) + }) + + this.#eventEmitter.on("status", ({ status }) => { + this.updateZustandState((state) => ({ + ...state, + status, + isStreaming: status === "streaming", + })) + }) + + this.#eventEmitter.on("error", ({ error }) => { + this.updateZustandState((state) => ({ + ...state, + error, + })) + }) + } + + get status(): ChatStatus { + return this.#status + } + + set status(newStatus: ChatStatus) { + if (this.#status === newStatus) return + + this.#status = newStatus + this.#eventEmitter.emit("status", { status: newStatus }) + } + + get error(): Error | undefined { + return this.#error + } + + set error(newError: Error | undefined) { + if (this.#error === newError) return + + this.#error = newError + this.#eventEmitter.emit("error", { error: newError }) + } + + get messages(): UI_MESSAGE[] { + return this.#messages + } + + set messages(newMessages: UI_MESSAGE[]) { + this.#messages = [...newMessages] + this.#eventEmitter.emit("messages", { messages: this.#messages }) + + // Auto-persist messages when they change + this.#persistMessages() + } + + pushMessage = (message: UI_MESSAGE) => { + this.messages = this.#messages.concat(message) + } + + popMessage = () => { + if (this.#messages.length === 0) return + + this.messages = this.#messages.slice(0, -1) + } + + replaceMessage = (index: number, message: UI_MESSAGE) => { + if (index < 0 || index >= this.#messages.length) return + + this.messages = [ + ...this.#messages.slice(0, index), + // Deep clone the message to ensure React detects changes + this.snapshot(message), + ...this.#messages.slice(index + 1), + ] + } + + snapshot = (value: T): T => structuredClone(value) + + // Callback registration methods with proper AI SDK compatibility + registerMessagesCallback = (onChange: () => void, throttleWaitMs?: number): (() => void) => { + const callback = throttleWaitMs ? throttle(onChange, throttleWaitMs) : onChange + + // Convert payload-based event to AI SDK expected callback format + return this.#eventEmitter.on("messages", () => callback()) + } + + registerStatusCallback = (onChange: () => void): (() => void) => { + // Convert payload-based event to AI SDK expected callback format + return this.#eventEmitter.on("status", () => onChange()) + } + + registerErrorCallback = (onChange: () => void): (() => void) => { + // Convert payload-based event to AI SDK expected callback format + return this.#eventEmitter.on("error", () => onChange()) + } + + // Internal event subscription with payload access + onMessagesChange = (listener: (messages: UI_MESSAGE[]) => void): (() => void) => { + return this.#eventEmitter.on("messages", ({ messages }) => listener(messages)) + } + + onStatusChange = (listener: (status: ChatStatus) => void): (() => void) => { + return this.#eventEmitter.on("status", ({ status }) => listener(status)) + } + + onErrorChange = (listener: (error: Error | undefined) => void): (() => void) => { + return this.#eventEmitter.on("error", ({ error }) => listener(error)) + } + + // Persistence methods + #persistMessages = throttle( + async () => { + // Skip if no messages + if (this.#messages.length === 0) return + + try { + await AIPersistService.ensureSession(this.chatId, "New Chat") + // Save messages using incremental updates + await AIPersistService.upsertMessages(this.chatId, this.#messages) + } catch (error) { + console.error("Failed to persist messages:", error) + } + }, + 100, + { leading: true, trailing: true }, + ) + + // Cleanup method + destroy(): void { + this.#eventEmitter.clear() + } +} + +// Custom Chat class that uses Zustand-integrated state +export class ZustandChat extends AbstractChat { + override state: ZustandChatState + #unsubscribeFns: (() => void)[] = [] + + constructor( + { messages, ...init }: ChatInit, + updateZustandState: (updater: (state: ChatSlice) => ChatSlice) => void, + ) { + const state = new ZustandChatState(messages, updateZustandState, init.id || "") + super({ ...init, state }) + this.state = state + } + + // Public getter for state access + get chatState() { + return this.state + } + + // Cleanup method + destroy(): void { + // Unsubscribe from AI SDK callbacks + this.#unsubscribeFns.forEach((unsubscribe) => unsubscribe()) + this.#unsubscribeFns = [] + + this.state.destroy() + } + + protected override setStatus({ status, error }: { status: ChatStatus; error?: Error }): void { + super.setStatus({ status, error }) + this.state.status = status + this.state.error = error + } +} + +// Zustand slice interface +export interface ChatSlice { + // Chat state (mirrored from ChatState) + chatId: string + messages: BizUIMessage[] + status: ChatStatus + error: Error | undefined + isStreaming: boolean + + // UI state + currentTitle: string | undefined + + // AI SDK Chat instance + chatInstance: ZustandChat + + // Actions + chatActions: ChatSliceActions +} + +export const createChatSlice: StateCreator = (...params) => { + const [set, get] = params + const chatId = nanoid() + + // Create chat instance with Zustand integration + const chatInstance = new ZustandChat( + { + id: chatId, + messages: [], + transport: createChatTransport(), + onFinish: async (options) => { + const { message } = options + + // Only trigger title generation for assistant messages (AI responses) + if (message.role !== "assistant") return + + // Get current messages to check if this is the first AI response + const allMessages = chatInstance.chatState.messages + + // Check if we have exactly 2 messages (1 user + 1 assistant = first exchange) + // Or if we have 2+ messages and this is the first assistant message + const assistantMessages = allMessages.filter((m) => m.role === "assistant") + const isFirstAIResponse = assistantMessages.length === 1 + + if (isFirstAIResponse && allMessages.length >= 2) { + try { + // Generate title using the first user message and first AI response + const firstExchange = allMessages.slice(0, 2) + + const title = await generateChatTitle(firstExchange) + + if (title && chatId) { + try { + await AIPersistService.updateSessionTitle(chatId, title) + if (get().chatId === chatId) { + chatActions.setCurrentTitle(title) + } + } catch (error) { + console.error("Failed to update session title:", error) + } + } + } catch (error) { + console.error("Failed to generate chat title:", error) + } + } + }, + }, + set, // Pass set function for state updates + ) + + // Create actions + const chatActions = new ChatSliceActions(params, chatInstance) + + return { + // Chat state + chatId, + messages: [], + status: "ready", + error: undefined, + isStreaming: false, + currentTitle: undefined, + chatInstance, + chatActions, + } +} + +class ChatSliceActions { + constructor( + private params: Parameters>, + private chatInstance: ZustandChat, + ) {} + + get set() { + return this.params[0] + } + + get get() { + return this.params[1] + } + + // Direct message management methods (delegating to chat instance state) + setMessages = ( + messagesParam: BizUIMessage[] | ((messages: BizUIMessage[]) => BizUIMessage[]), + ) => { + if (typeof messagesParam === "function") { + this.chatInstance.chatState.messages = messagesParam(this.chatInstance.chatState.messages) + } else { + this.chatInstance.chatState.messages = messagesParam + } + } + + pushMessage = (message: BizUIMessage) => { + this.chatInstance.chatState.pushMessage(message) + } + + popMessage = () => { + this.chatInstance.chatState.popMessage() + } + + replaceMessage = (index: number, message: BizUIMessage) => { + this.chatInstance.chatState.replaceMessage(index, message) + } + + updateMessage = (id: string, updates: Partial) => { + const messageIndex = this.chatInstance.chatState.messages.findIndex( + (msg: BizUIMessage) => msg.id === id, + ) + if (messageIndex !== -1) { + const message = this.chatInstance.chatState.messages[messageIndex] + if (message) { + const updatedMessage = { ...message, ...updates } + this.replaceMessage(messageIndex, updatedMessage) + } + } + } + + // Getter + getMessages = (): BizUIMessage[] => { + return this.chatInstance.chatState.messages + } + + // Status management (delegating to chat instance state) + setStatus = (status: ChatStatus) => { + this.chatInstance.chatState.status = status + } + + setError = (error: Error | undefined) => { + this.chatInstance.chatState.error = error + } + + setStreaming = (streaming: boolean) => { + this.chatInstance.chatState.status = streaming ? "streaming" : "ready" + } + + // Title management + setCurrentTitle = (title: string | undefined) => { + this.set((state) => ({ ...state, currentTitle: title })) + } + + getCurrentTitle = (): string | undefined => { + return this.get().currentTitle + } + + getCurrentChatId = (): string | null => { + return this.get().chatId + } + + // Core chat actions using AI SDK AbstractChat methods + sendMessage = async (message: string | BizUIMessage) => { + try { + // Convert string to message object if needed + const messageObj = + typeof message === "string" + ? ({ parts: [{ type: "text", text: message }] } as Parameters< + typeof this.chatInstance.sendMessage + >[0]) + : (message as Parameters[0]) + + // Use the AI SDK's sendMessage method + const response = await this.chatInstance.sendMessage(messageObj) + return response + } catch (error) { + this.setError(error as Error) + throw error + } + } + + regenerate = async ({ messageId }: { messageId: string }) => { + try { + // Use the AI SDK's regenerate method + const response = await this.chatInstance.regenerate({ messageId }) + return response + } catch (error) { + this.setError(error as Error) + throw error + } + } + + stop = () => { + // Use AI SDK's stop method + this.chatInstance.stop() + } + + resumeStream = async () => { + try { + // Use AI SDK's resumeStream method + await this.chatInstance.resumeStream() + } catch (error) { + this.setError(error as Error) + throw error + } + } + + resetChat = () => { + // Reset through the chat instance state + this.chatInstance.chatState.messages = [] + this.chatInstance.chatState.error = undefined + this.chatInstance.chatState.status = "ready" + // Reset title + this.setCurrentTitle(undefined) + } + + newChat = () => { + const newChatId = nanoid() + + // Cleanup old chat instance + this.chatInstance.destroy() + + // Create new chat instance + const newChatInstance = new ZustandChat( + { + id: newChatId, + messages: [], + transport: createChatTransport(), + }, + this.set, + ) + + // Update store state + this.set((state) => ({ + ...state, + chatId: newChatId, + messages: [], + status: "ready" as ChatStatus, + error: undefined, + isStreaming: false, + currentTitle: undefined, + chatInstance: newChatInstance, + })) + + // Update the reference + this.chatInstance = newChatInstance + } + + switchToChat = async (chatId: string) => { + try { + // Set loading state (using ready as there's no loading status in ChatStatus) + this.setStatus("ready") + this.setError(undefined) + + // Load messages from persistence service + const messages = await AIPersistService.loadUIMessages(chatId) + + // Load chat session details to get title (direct SQL query) + const chatSession = await AIPersistService.getChatSession(chatId) + + // Cleanup old chat instance + this.chatInstance.destroy() + + // Create new chat instance with loaded messages + const newChatInstance = new ZustandChat( + { + id: chatId, + messages, + transport: createChatTransport(), + }, + this.set, + ) + + // Update store state + this.set((state) => ({ + ...state, + chatId, + messages: [...messages], + status: "ready" as ChatStatus, + error: undefined, + isStreaming: false, + currentTitle: chatSession?.title || undefined, + chatInstance: newChatInstance, + })) + + // Update the reference + this.chatInstance = newChatInstance + } catch (error) { + console.error("Failed to switch to chat:", error) + this.setError(error as Error) + this.setStatus("ready") + throw error + } + } +} + +export { type ChatStatus } from "ai" diff --git a/apps/desktop/layer/renderer/src/modules/ai/chat/__internal__/slices/index.ts b/apps/desktop/layer/renderer/src/modules/ai/chat/__internal__/slices/index.ts new file mode 100644 index 000000000..4d5de9862 --- /dev/null +++ b/apps/desktop/layer/renderer/src/modules/ai/chat/__internal__/slices/index.ts @@ -0,0 +1,5 @@ +export { + type BlockSlice as ContextSlice, + createBlockSlice as createContextSlice, +} from "./block.slice" +export { type ChatSlice, type ChatStatus, createChatSlice } from "./chat.slice" diff --git a/apps/desktop/layer/renderer/src/modules/ai/chat/__internal__/store.ts b/apps/desktop/layer/renderer/src/modules/ai/chat/__internal__/store.ts index d623edfbb..6e21cd957 100644 --- a/apps/desktop/layer/renderer/src/modules/ai/chat/__internal__/store.ts +++ b/apps/desktop/layer/renderer/src/modules/ai/chat/__internal__/store.ts @@ -1,203 +1,29 @@ -import { produce } from "immer" -import { nanoid } from "nanoid" import { createWithEqualityFn } from "zustand/traditional" -import type { AIChatContextBlock, AIChatContextInfo } from "./types" -import { blocksToContextInfo, contextInfoToBlocks } from "./utils" +import type { BlockSlice } from "./slices/block.slice" +import { createBlockSlice } from "./slices/block.slice" +import type { ChatSlice } from "./slices/chat.slice" +import { createChatSlice } from "./slices/chat.slice" +import type { AIChatStoreInitial } from "./types" -export interface AiChatContextStore { - state: AIChatContextInfo - blocks: AIChatContextBlock[] - // Block management methods - addBlock: (block: Omit) => void - removeBlock: (id: string) => void - updateBlock: (id: string, updates: Partial) => void - clearBlocks: () => void - // Context info management methods - setMainEntryId: (entryId?: string) => void - addReferEntryId: (entryId: string) => void - removeReferEntryId: (entryId: string) => void - addReferFeedId: (feedId: string) => void - removeReferFeedId: (feedId: string) => void - setSelectedText: (selectedText?: string) => void - // Legacy compatibility - setEntryId: (entryId?: string) => void - setFeedId: (feedId?: string) => void - reset: () => void - // Sync methods - syncBlocksToContext: () => void - syncContextToBlocks: () => void -} - -export const createAIChatContextStore = (initialState?: Partial) => { - const defaultState: AIChatContextInfo = { - mainEntryId: undefined, - referEntryIds: [], - referFeedIds: [], - selectedText: undefined, +export type AiChatStore = BlockSlice & + ChatSlice & { + reset: () => void } - const defaultBlocks: AIChatContextBlock[] = [] +export const createAIChatStore = (initialState?: Partial) => { + return createWithEqualityFn((...a) => { + const blockSlice = createBlockSlice(initialState?.blocks)(...a) + const chatSlice = createChatSlice(...a) - const state: AIChatContextInfo = { - mainEntryId: initialState?.mainEntryId, - referEntryIds: initialState?.referEntryIds || [], - referFeedIds: initialState?.referFeedIds || [], - selectedText: initialState?.selectedText, - } + return { + ...blockSlice, + ...chatSlice, - return createWithEqualityFn((set, get) => ({ - state, - blocks: defaultBlocks, - - // Block management methods - addBlock: (block: Omit) => { - // Check for uniqueness constraints - const currentBlocks = get().blocks - - // Only allow one mainEntry - if (block.type === "mainEntry" && currentBlocks.some((b) => b.type === "mainEntry")) { - return - } - - // Only allow one selectedText - if (block.type === "selectedText" && currentBlocks.some((b) => b.type === "selectedText")) { - return - } - - // Prevent duplicate referEntry or referFeed - if ( - block.type === "referEntry" && - block.value && - currentBlocks.some((b) => b.type === "referEntry" && b.value === block.value) - ) { - return - } - - if ( - block.type === "referFeed" && - block.value && - currentBlocks.some((b) => b.type === "referFeed" && b.value === block.value) - ) { - return - } - - set( - produce((s) => { - s.blocks.push({ ...block, id: nanoid(8) }) - }), - ) - - // Sync to context info - get().syncBlocksToContext() - }, - - removeBlock: (id: string) => { - set( - produce((s: AiChatContextStore) => { - s.blocks = s.blocks.filter((block) => block.id !== id) - }), - ) - - // Sync to context info - get().syncBlocksToContext() - }, - - updateBlock: (id: string, updates: Partial) => { - set( - produce((s: AiChatContextStore) => { - s.blocks = s.blocks.map((block) => (block.id === id ? { ...block, ...updates } : block)) - }), - ) - - // Sync to context info - get().syncBlocksToContext() - }, - - clearBlocks: () => { - set( - produce((s: AiChatContextStore) => { - s.blocks = defaultBlocks - }), - ) - get().syncBlocksToContext() - }, - - // Context info management methods - setMainEntryId: (entryId?: string) => { - set((s) => ({ state: { ...s.state, mainEntryId: entryId } })) - get().syncContextToBlocks() - }, - - addReferEntryId: (entryId: string) => { - set((s) => ({ - state: { - ...s.state, - referEntryIds: [...(s.state.referEntryIds || []), entryId], - }, - })) - get().syncContextToBlocks() - }, - - removeReferEntryId: (entryId: string) => { - set((s) => ({ - state: { - ...s.state, - referEntryIds: (s.state.referEntryIds || []).filter((id) => id !== entryId), - }, - })) - get().syncContextToBlocks() - }, - - addReferFeedId: (feedId: string) => { - set((s) => ({ - state: { - ...s.state, - referFeedIds: [...(s.state.referFeedIds || []), feedId], - }, - })) - get().syncContextToBlocks() - }, - - removeReferFeedId: (feedId: string) => { - set((s) => ({ - state: { - ...s.state, - referFeedIds: (s.state.referFeedIds || []).filter((id) => id !== feedId), - }, - })) - get().syncContextToBlocks() - }, - - setSelectedText: (selectedText?: string) => { - set((s) => ({ state: { ...s.state, selectedText } })) - get().syncContextToBlocks() - }, - - // Legacy compatibility - setEntryId: (entryId?: string) => { - get().setMainEntryId(entryId) - }, - - setFeedId: (feedId?: string) => { - if (feedId) { - get().addReferFeedId(feedId) - } - }, - - reset: () => { - set(() => ({ state: defaultState, blocks: defaultBlocks })) - }, - - // Sync methods - syncBlocksToContext: () => { - const contextInfo = blocksToContextInfo(get().blocks) - set((s) => ({ state: { ...s.state, ...contextInfo } })) - }, - - syncContextToBlocks: () => { - const blocks = contextInfoToBlocks(get().state) - set(() => ({ blocks })) - }, - })) + reset: () => { + blockSlice.blockActions.resetContext() + chatSlice.chatActions.resetChat() + }, + } + }) } diff --git a/apps/desktop/layer/renderer/src/modules/ai/chat/__internal__/transport.ts b/apps/desktop/layer/renderer/src/modules/ai/chat/__internal__/transport.ts new file mode 100644 index 000000000..664281202 --- /dev/null +++ b/apps/desktop/layer/renderer/src/modules/ai/chat/__internal__/transport.ts @@ -0,0 +1,14 @@ +import { env } from "@follow/shared/env.desktop" +import { DefaultChatTransport } from "ai" + +/** + * Create a chat transport for AI SDK + * This is used by the AbstractChat instance to communicate with AI providers + */ +export function createChatTransport() { + return new DefaultChatTransport({ + // Custom fetch configuration + api: `${env.VITE_API_URL}/ai/chat`, + credentials: "include", + }) +} diff --git a/apps/desktop/layer/renderer/src/modules/ai/chat/__internal__/types.ts b/apps/desktop/layer/renderer/src/modules/ai/chat/__internal__/types.ts index 116611252..bbc1f4829 100644 --- a/apps/desktop/layer/renderer/src/modules/ai/chat/__internal__/types.ts +++ b/apps/desktop/layer/renderer/src/modules/ai/chat/__internal__/types.ts @@ -6,11 +6,8 @@ export interface AIChatContextBlock { value: string } -export interface AIChatContextInfo { - mainEntryId?: string - referEntryIds?: string[] - referFeedIds?: string[] - selectedText?: string +export interface AIChatStoreInitial { + blocks: AIChatContextBlock[] } export interface AIChatContextBlocks { diff --git a/apps/desktop/layer/renderer/src/modules/ai/chat/__internal__/utils.ts b/apps/desktop/layer/renderer/src/modules/ai/chat/__internal__/utils.ts deleted file mode 100644 index 3b4bb7e28..000000000 --- a/apps/desktop/layer/renderer/src/modules/ai/chat/__internal__/utils.ts +++ /dev/null @@ -1,72 +0,0 @@ -import type { AIChatContextBlock, AIChatContextInfo } from "./types" - -export const contextInfoToBlocks = (contextInfo: AIChatContextInfo): AIChatContextBlock[] => { - const blocks: AIChatContextBlock[] = [] - - // Main entry (single) - if (contextInfo.mainEntryId) { - blocks.push({ - id: `main-entry-${contextInfo.mainEntryId}`, - type: "mainEntry", - value: contextInfo.mainEntryId, - }) - } - - // Reference entries (multiple) - contextInfo.referEntryIds?.forEach((entryId) => { - blocks.push({ - id: `refer-entry-${entryId}`, - type: "referEntry", - value: entryId, - }) - }) - - // Reference feeds (multiple) - contextInfo.referFeedIds?.forEach((feedId) => { - blocks.push({ - id: `refer-feed-${feedId}`, - type: "referFeed", - value: feedId, - }) - }) - - // Selected text (single, auto-detected) - if (contextInfo.selectedText) { - blocks.push({ - id: "selected-text", - type: "selectedText", - value: contextInfo.selectedText, - }) - } - - return blocks -} - -export const blocksToContextInfo = (blocks: AIChatContextBlock[]): AIChatContextInfo => { - const contextInfo: AIChatContextInfo = {} - - blocks.forEach((block) => { - switch (block.type) { - case "mainEntry": { - contextInfo.mainEntryId = block.value - break - } - case "referEntry": { - if (!contextInfo.referEntryIds) contextInfo.referEntryIds = [] - contextInfo.referEntryIds.push(block.value) - break - } - case "referFeed": { - if (!contextInfo.referFeedIds) contextInfo.referFeedIds = [] - contextInfo.referFeedIds.push(block.value) - break - } - case "selectedText": { - contextInfo.selectedText = block.value - break - } - } - }) - - return contextInfo -} diff --git a/apps/desktop/layer/renderer/src/modules/ai/chat/atoms/session.ts b/apps/desktop/layer/renderer/src/modules/ai/chat/atoms/session.ts index b7f3e9502..4a0205cf3 100644 --- a/apps/desktop/layer/renderer/src/modules/ai/chat/atoms/session.ts +++ b/apps/desktop/layer/renderer/src/modules/ai/chat/atoms/session.ts @@ -1,49 +1,7 @@ import { atom } from "jotai" -import { useCallback } from "react" import { createAtomHooks } from "~/lib/jotai" -// Session state atoms with hooks -export const [, , useCurrentRoomId, useSetCurrentRoomId, , setCurrentRoomId] = createAtomHooks( - atom(null), -) - -export const [, , useCurrentTitle, useSetCurrentTitle, , setCurrentTitle] = - createAtomHooks(atom()) - -export const [, , useSessionPersisted, useSetSessionPersisted, , setSessionPersisted] = - createAtomHooks(atom(false)) - // Edit state management for messages export const [, , useEditingMessageId, useSetEditingMessageId, , setEditingMessageId] = createAtomHooks(atom(null)) - -// Combined hook for all session state -export const useSessionState = () => { - return { - currentRoomId: useCurrentRoomId(), - currentTitle: useCurrentTitle(), - sessionPersisted: useSessionPersisted(), - editingMessageId: useEditingMessageId(), - } -} - -// Hook for session state setters -export const useSessionSetters = () => { - const setCurrentRoomIdAtom = useSetCurrentRoomId() - const setCurrentTitleAtom = useSetCurrentTitle() - const setSessionPersistedAtom = useSetSessionPersisted() - - return useCallback( - (updates: { - currentRoomId?: string | null - currentTitle?: string | undefined - sessionPersisted?: boolean - }) => { - if (updates.currentRoomId !== undefined) setCurrentRoomIdAtom(updates.currentRoomId) - if (updates.currentTitle !== undefined) setCurrentTitleAtom(updates.currentTitle) - if (updates.sessionPersisted !== undefined) setSessionPersistedAtom(updates.sessionPersisted) - }, - [setCurrentRoomIdAtom, setCurrentTitleAtom, setSessionPersistedAtom], - ) -} diff --git a/apps/desktop/layer/renderer/src/modules/ai/chat/components/AIChatContextBar.tsx b/apps/desktop/layer/renderer/src/modules/ai/chat/components/AIChatContextBar.tsx index c5fd372cf..673e4d677 100644 --- a/apps/desktop/layer/renderer/src/modules/ai/chat/components/AIChatContextBar.tsx +++ b/apps/desktop/layer/renderer/src/modules/ai/chat/components/AIChatContextBar.tsx @@ -25,10 +25,12 @@ import { useRouteParamsSelector } from "~/hooks/biz/useRouteParams" import { useAIChatStore } from "~/modules/ai/chat/__internal__/AIChatContext" import type { AIChatContextBlock } from "~/modules/ai/chat/__internal__/types" +import { useChatBlockActions } from "../__internal__/hooks" + export const AIChatContextBar: Component<{ onSendShortcut?: (prompt: string) => void }> = memo( ({ className, onSendShortcut }) => { const blocks = useAIChatStore()((s) => s.blocks) - const { addBlock } = useAIChatStore()() + const blockActions = useChatBlockActions() const { shortcuts } = useAISettingValue() // Filter enabled shortcuts @@ -47,7 +49,7 @@ export const AIChatContextBar: Component<{ onSendShortcut?: (prompt: string) => - addBlock({ + blockActions.addBlock({ type: "referEntry", value: entryId, }) @@ -65,7 +67,7 @@ export const AIChatContextBar: Component<{ onSendShortcut?: (prompt: string) => - addBlock({ + blockActions.addBlock({ type: "referEntry", value: entryId, }) @@ -81,7 +83,7 @@ export const AIChatContextBar: Component<{ onSendShortcut?: (prompt: string) => - addBlock({ + blockActions.addBlock({ type: "referFeed", value: feedId, }) @@ -227,7 +229,7 @@ const PickerList = ({ const CurrentFeedEntriesPickerList: FC<{ onSelect: (entryId: string) => void }> = ({ onSelect, }) => { - const mainEntryId = useAIChatStore()((s) => s.state.mainEntryId) + const mainEntryId = useAIChatStore()((s) => s.blocks.find((b) => b.type === "mainEntry")?.value) const feedId = useEntry(mainEntryId, (e) => e?.feedId) const entryIds = useEntryIdsByFeedId(feedId!) @@ -329,7 +331,7 @@ const FeedPickerItem: FC<{ } const ContextBlock: FC<{ block: AIChatContextBlock }> = ({ block }) => { - const { removeBlock } = useAIChatStore()() + const blockActions = useChatBlockActions() const getBlockIcon = () => { switch (block.type) { @@ -409,7 +411,7 @@ const ContextBlock: FC<{ block: AIChatContextBlock }> = ({ block }) => { {canRemove && (