diff --git a/packages/frontend/core/src/blocksuite/ai/chat-panel/index.ts b/packages/frontend/core/src/blocksuite/ai/chat-panel/index.ts index 9e65935e0..c9376de68 100644 --- a/packages/frontend/core/src/blocksuite/ai/chat-panel/index.ts +++ b/packages/frontend/core/src/blocksuite/ai/chat-panel/index.ts @@ -1,3 +1,4 @@ +import type { AIDraftService } from '@affine/core/modules/ai-button'; import type { WorkspaceDialogService } from '@affine/core/modules/dialogs'; import type { FeatureFlagService } from '@affine/core/modules/feature-flag'; import type { AppThemeService } from '@affine/core/modules/theme'; @@ -115,6 +116,9 @@ export class ChatPanel extends SignalWatcher( @property({ attribute: false }) accessor notificationService!: NotificationService; + @property({ attribute: false }) + accessor aiDraftService!: AIDraftService; + @state() accessor session: CopilotChatHistoryFragment | null | undefined; @@ -408,6 +412,7 @@ export class ChatPanel extends SignalWatcher( .affineWorkspaceDialogService=${this.affineWorkspaceDialogService} .affineThemeService=${this.affineThemeService} .notificationService=${this.notificationService} + .aiDraftService=${this.aiDraftService} .onEmbeddingProgressChange=${this.onEmbeddingProgressChange} .onContextChange=${this.onContextChange} .width=${this.sidebarWidth} diff --git a/packages/frontend/core/src/blocksuite/ai/components/ai-chat-composer/ai-chat-composer.ts b/packages/frontend/core/src/blocksuite/ai/components/ai-chat-composer/ai-chat-composer.ts index 49affd381..d2d521957 100644 --- a/packages/frontend/core/src/blocksuite/ai/components/ai-chat-composer/ai-chat-composer.ts +++ b/packages/frontend/core/src/blocksuite/ai/components/ai-chat-composer/ai-chat-composer.ts @@ -1,5 +1,6 @@ import './ai-chat-composer-tip'; +import type { AIDraftService } from '@affine/core/modules/ai-button'; import type { WorkspaceDialogService } from '@affine/core/modules/dialogs'; import type { ContextEmbedStatus, @@ -116,6 +117,9 @@ export class AIChatComposer extends SignalWatcher( @property({ attribute: false }) accessor notificationService!: NotificationService; + @property({ attribute: false }) + accessor aiDraftService!: AIDraftService; + @state() accessor chips: ChatChip[] = []; @@ -161,6 +165,7 @@ export class AIChatComposer extends SignalWatcher( .reasoningConfig=${this.reasoningConfig} .docDisplayConfig=${this.docDisplayConfig} .searchMenuConfig=${this.searchMenuConfig} + .aiDraftService=${this.aiDraftService} .portalContainer=${this.portalContainer} .onChatSuccess=${this.onChatSuccess} .trackOptions=${this.trackOptions} diff --git a/packages/frontend/core/src/blocksuite/ai/components/ai-chat-content/ai-chat-content.ts b/packages/frontend/core/src/blocksuite/ai/components/ai-chat-content/ai-chat-content.ts index 7a4dab495..79c110e97 100644 --- a/packages/frontend/core/src/blocksuite/ai/components/ai-chat-content/ai-chat-content.ts +++ b/packages/frontend/core/src/blocksuite/ai/components/ai-chat-content/ai-chat-content.ts @@ -1,3 +1,5 @@ +import type { AIDraftService } from '@affine/core/modules/ai-button'; +import type { AIDraftState } from '@affine/core/modules/ai-button/services/ai-draft'; import type { WorkspaceDialogService } from '@affine/core/modules/dialogs'; import type { FeatureFlagService } from '@affine/core/modules/feature-flag'; import type { AppThemeService } from '@affine/core/modules/theme'; @@ -15,6 +17,7 @@ import { property, state } from 'lit/decorators.js'; import { classMap } from 'lit/directives/class-map.js'; import { createRef, type Ref, ref } from 'lit/directives/ref.js'; import { styleMap } from 'lit/directives/style-map.js'; +import { pick } from 'lodash-es'; import { HISTORY_IMAGE_ACTIONS } from '../../chat-panel/const'; import { type AIChatParams, AIProvider } from '../../provider/ai-provider'; @@ -149,6 +152,9 @@ export class AIChatContent extends SignalWatcher( @property({ attribute: false }) accessor notificationService!: NotificationService; + @property({ attribute: false }) + accessor aiDraftService!: AIDraftService; + @property({ attribute: false }) accessor onEmbeddingProgressChange: | ((count: Record) => void) @@ -263,6 +269,19 @@ export class AIChatContent extends SignalWatcher( private readonly updateContext = (context: Partial) => { this.chatContextValue = { ...this.chatContextValue, ...context }; this.onContextChange?.(context); + this.updateDraft(context).catch(console.error); + }; + + private readonly updateDraft = async (context: Partial) => { + const draft: Partial = pick(context, [ + 'quote', + 'images', + 'markdown', + ]); + if (!Object.keys(draft).length) { + return; + } + await this.aiDraftService.setDraft(draft); }; private readonly initChatContent = async () => { @@ -322,8 +341,19 @@ export class AIChatContent extends SignalWatcher( override connectedCallback() { super.connectedCallback(); + this.initChatContent().catch(console.error); + this.aiDraftService + .getDraft() + .then(draft => { + this.chatContextValue = { + ...this.chatContextValue, + ...draft, + }; + }) + .catch(console.error); + this._disposables.add( AIProvider.slots.actions.subscribe(({ event }) => { const { status } = this.chatContextValue; @@ -403,6 +433,7 @@ export class AIChatContent extends SignalWatcher( .searchMenuConfig=${this.searchMenuConfig} .affineWorkspaceDialogService=${this.affineWorkspaceDialogService} .notificationService=${this.notificationService} + .aiDraftService=${this.aiDraftService} .trackOptions=${{ where: 'chat-panel', control: 'chat-send', diff --git a/packages/frontend/core/src/blocksuite/ai/components/ai-chat-input/ai-chat-input.ts b/packages/frontend/core/src/blocksuite/ai/components/ai-chat-input/ai-chat-input.ts index e7d87621d..80e3a56fb 100644 --- a/packages/frontend/core/src/blocksuite/ai/components/ai-chat-input/ai-chat-input.ts +++ b/packages/frontend/core/src/blocksuite/ai/components/ai-chat-input/ai-chat-input.ts @@ -1,10 +1,11 @@ +import type { AIDraftService } from '@affine/core/modules/ai-button'; import type { CopilotChatHistoryFragment } from '@affine/graphql'; import { SignalWatcher, WithDisposable } from '@blocksuite/affine/global/lit'; import { unsafeCSSVar, unsafeCSSVarV2 } from '@blocksuite/affine/shared/theme'; import type { EditorHost } from '@blocksuite/affine/std'; import { ShadowlessElement } from '@blocksuite/affine/std'; import { ArrowUpBigIcon, CloseIcon } from '@blocksuite/icons/lit'; -import { css, html, nothing } from 'lit'; +import { css, html, nothing, type PropertyValues } from 'lit'; import { property, query, state } from 'lit/decorators.js'; import { repeat } from 'lit/directives/repeat.js'; import { styleMap } from 'lit/directives/style-map.js'; @@ -351,6 +352,9 @@ export class AIChatInput extends SignalWatcher( @property({ attribute: false }) accessor searchMenuConfig!: SearchMenuConfig; + @property({ attribute: false }) + accessor aiDraftService!: AIDraftService; + @property({ attribute: false }) accessor isRootSession: boolean = true; @@ -379,6 +383,7 @@ export class AIChatInput extends SignalWatcher( override connectedCallback() { super.connectedCallback(); + this._disposables.add( AIProvider.slots.requestSendWithChat.subscribe( (params: AISendParams | null) => { @@ -399,6 +404,17 @@ export class AIChatInput extends SignalWatcher( ); } + protected override firstUpdated(changedProperties: PropertyValues): void { + super.firstUpdated(changedProperties); + this.aiDraftService + .getDraft() + .then(draft => { + this.textarea.value = draft.input; + this.isInputEmpty = !this.textarea.value.trim(); + }) + .catch(console.error); + } + protected override render() { const { images, status } = this.chatContextValue; const hasImages = images.length > 0; @@ -506,9 +522,11 @@ export class AIChatInput extends SignalWatcher( } }; - private readonly _handleInput = () => { + private readonly _handleInput = async () => { const { textarea } = this; - this.isInputEmpty = !textarea.value.trim(); + const value = textarea.value.trim(); + this.isInputEmpty = !value; + textarea.style.height = 'auto'; textarea.style.height = textarea.scrollHeight + 'px'; let imagesHeight = this.imagePreviewGrid?.scrollHeight ?? 0; @@ -517,6 +535,10 @@ export class AIChatInput extends SignalWatcher( textarea.style.height = '148px'; textarea.style.overflowY = 'scroll'; } + + await this.aiDraftService.setDraft({ + input: value, + }); }; private readonly _handleKeyDown = async (evt: KeyboardEvent) => { @@ -572,6 +594,9 @@ export class AIChatInput extends SignalWatcher( this.textarea.style.height = 'unset'; await this.send(value); + await this.aiDraftService.setDraft({ + input: '', + }); }; private readonly _handleModelChange = (modelId: string) => { diff --git a/packages/frontend/core/src/desktop/pages/workspace/chat/index.tsx b/packages/frontend/core/src/desktop/pages/workspace/chat/index.tsx index 933a9f610..174dafc81 100644 --- a/packages/frontend/core/src/desktop/pages/workspace/chat/index.tsx +++ b/packages/frontend/core/src/desktop/pages/workspace/chat/index.tsx @@ -11,6 +11,7 @@ import { getViewManager } from '@affine/core/blocksuite/manager/view'; import { NotificationServiceImpl } from '@affine/core/blocksuite/view-extensions/editor-view/notification-service'; import { useAIChatConfig } from '@affine/core/components/hooks/affine/use-ai-chat-config'; import { useAISpecs } from '@affine/core/components/hooks/affine/use-ai-specs'; +import { AIDraftService } from '@affine/core/modules/ai-button'; import { EventSourceService, FetchService, @@ -221,6 +222,7 @@ export const Component = () => { confirmModal.closeConfirmModal, confirmModal.openConfirmModal ); + content.aiDraftService = framework.get(AIDraftService); content.createSession = createSession; content.onOpenDoc = onOpenDoc; diff --git a/packages/frontend/core/src/desktop/pages/workspace/detail-page/tabs/chat.tsx b/packages/frontend/core/src/desktop/pages/workspace/detail-page/tabs/chat.tsx index a5030e4d3..98fadd7ce 100644 --- a/packages/frontend/core/src/desktop/pages/workspace/detail-page/tabs/chat.tsx +++ b/packages/frontend/core/src/desktop/pages/workspace/detail-page/tabs/chat.tsx @@ -4,6 +4,7 @@ import type { AffineEditorContainer } from '@affine/core/blocksuite/block-suite- import { NotificationServiceImpl } from '@affine/core/blocksuite/view-extensions/editor-view/notification-service'; import { useAIChatConfig } from '@affine/core/components/hooks/affine/use-ai-chat-config'; import { useAISpecs } from '@affine/core/components/hooks/affine/use-ai-specs'; +import { AIDraftService } from '@affine/core/modules/ai-button'; import { WorkspaceDialogService } from '@affine/core/modules/dialogs'; import { FeatureFlagService } from '@affine/core/modules/feature-flag'; import { AppThemeService } from '@affine/core/modules/theme'; @@ -95,6 +96,7 @@ export const EditorChatPanel = forwardRef(function EditorChatPanel( confirmModal.closeConfirmModal, confirmModal.openConfirmModal ); + chatPanelRef.current.aiDraftService = framework.get(AIDraftService); containerRef.current?.append(chatPanelRef.current); } else { diff --git a/packages/frontend/core/src/modules/ai-button/index.ts b/packages/frontend/core/src/modules/ai-button/index.ts index 66403c8fa..03c76dcf6 100644 --- a/packages/frontend/core/src/modules/ai-button/index.ts +++ b/packages/frontend/core/src/modules/ai-button/index.ts @@ -1,12 +1,15 @@ export { AIButtonProvider } from './provider/ai-button'; export { AIButtonService } from './services/ai-button'; +export { AIDraftService } from './services/ai-draft'; import type { Framework } from '@toeverything/infra'; import { FeatureFlagService } from '../feature-flag'; -import { GlobalStateService } from '../storage'; +import { CacheStorage, GlobalStateService } from '../storage'; +import { WorkspaceScope } from '../workspace'; import { AIButtonProvider } from './provider/ai-button'; import { AIButtonService } from './services/ai-button'; +import { AIDraftService } from './services/ai-draft'; import { AINetworkSearchService } from './services/network-search'; import { AIPlaygroundService } from './services/playground'; import { AIReasoningService } from './services/reasoning'; @@ -31,3 +34,9 @@ export function configureAIReasoningModule(framework: Framework) { export function configureAIPlaygroundModule(framework: Framework) { framework.service(AIPlaygroundService, [FeatureFlagService]); } + +export function configureAIDraftModule(framework: Framework) { + framework + .scope(WorkspaceScope) + .service(AIDraftService, [GlobalStateService, CacheStorage]); +} diff --git a/packages/frontend/core/src/modules/ai-button/services/ai-draft.ts b/packages/frontend/core/src/modules/ai-button/services/ai-draft.ts new file mode 100644 index 000000000..12a4a07ee --- /dev/null +++ b/packages/frontend/core/src/modules/ai-button/services/ai-draft.ts @@ -0,0 +1,159 @@ +import { Service } from '@toeverything/infra'; + +import type { CacheStorage, GlobalStateService } from '../../storage'; + +const AI_DRAFTS_KEY = 'AIDrafts'; +const AI_DRAFT_FILES_PREFIX = 'AIDraftFile:'; + +export interface CacheFile { + name: string; + size: number; + type: string; + cacheKey: string; +} + +export interface AIDraftState { + input: string; + quote: string; + markdown: string; + images: File[]; +} + +export interface AIDraftGlobal { + input: string; + quote: string; + markdown: string; + images: CacheFile[]; +} + +const DEFAULT_VALUE = { + input: '', + quote: '', + markdown: '', + images: [], +}; + +export class AIDraftService extends Service { + private state: AIDraftState | null = null; + + constructor( + private readonly globalStateService: GlobalStateService, + private readonly cacheStorage: CacheStorage + ) { + super(); + } + + setDraft = async (data: Partial) => { + const state = await this.getState(); + const newState = { + ...state, + ...data, + }; + this.state = newState; + await this.saveDraft(newState); + }; + + getDraft = async () => { + const state = await this.getState(); + return state; + }; + + private readonly saveDraft = async (state: AIDraftState) => { + const draft = + this.globalStateService.globalState.get(AI_DRAFTS_KEY) || + DEFAULT_VALUE; + + const addedImages = state.images.filter(image => { + return !draft.images.some(cacheImage => { + return cacheImage.cacheKey === this.getCacheKey(image); + }); + }); + const removedImages = draft.images.filter(cacheImage => { + return !state.images.some(image => { + return cacheImage.cacheKey === this.getCacheKey(image); + }); + }); + + const cacheKeys = removedImages.map(image => image.cacheKey); + await this.removeFilesFromCache(cacheKeys); + await this.addFilesToCache(addedImages); + + this.globalStateService.globalState.set(AI_DRAFTS_KEY, { + input: state.input, + quote: state.quote, + markdown: state.markdown, + images: state.images.map(image => { + return { + name: image.name, + size: image.size, + type: image.type, + cacheKey: this.getCacheKey(image), + }; + }), + }); + }; + + private readonly initState = async () => { + if (this.state) { + return; + } + const draft = + this.globalStateService.globalState.get(AI_DRAFTS_KEY); + if (draft) { + const images = await this.restoreFilesFromData(draft.images); + this.state = { + input: draft.input, + quote: draft.quote, + markdown: draft.markdown, + images, + }; + } else { + this.state = DEFAULT_VALUE; + } + }; + + private readonly getState = async () => { + await this.initState(); + return this.state as AIDraftState; + }; + + private readonly getCacheKey = (file: File) => { + return AI_DRAFT_FILES_PREFIX + file.name + file.size; + }; + + private readonly addFilesToCache = async (files: File[]) => { + for (const file of files) { + const arrayBuffer = await file.arrayBuffer(); + const cacheKey = this.getCacheKey(file); + await this.cacheStorage.set(cacheKey, arrayBuffer); + } + }; + + private readonly removeFilesFromCache = async (cacheKeys: string[]) => { + for (const cacheKey of cacheKeys) { + await this.cacheStorage.del(cacheKey); + } + }; + + private readonly restoreFilesFromData = async ( + cacheFiles: CacheFile[] + ): Promise => { + const files: File[] = []; + for (const cacheFile of cacheFiles) { + try { + const arrayBuffer = await this.cacheStorage.get( + cacheFile.cacheKey + ); + if (arrayBuffer) { + const file = new File([arrayBuffer], cacheFile.name, { + type: cacheFile.type, + }); + files.push(file); + } + } catch (error) { + console.warn(`Failed to restore file ${cacheFile.name}:`, error); + } + } + return files; + }; +} diff --git a/packages/frontend/core/src/modules/index.ts b/packages/frontend/core/src/modules/index.ts index 461a81218..99a7d476e 100644 --- a/packages/frontend/core/src/modules/index.ts +++ b/packages/frontend/core/src/modules/index.ts @@ -3,6 +3,7 @@ import { type Framework } from '@toeverything/infra'; import { configureAIButtonModule, + configureAIDraftModule, configureAINetworkSearchModule, configureAIPlaygroundModule, configureAIReasoningModule, @@ -110,6 +111,7 @@ export function configureCommonModules(framework: Framework) { configureAIReasoningModule(framework); configureAIPlaygroundModule(framework); configureAIButtonModule(framework); + configureAIDraftModule(framework); configureTemplateDocModule(framework); configureBlobManagementModule(framework); configureMediaModule(framework);