diff --git a/packages/frontend/core/src/blocksuite/ai/actions/types.ts b/packages/frontend/core/src/blocksuite/ai/actions/types.ts index f2c61a3bc3..35a9df92c4 100644 --- a/packages/frontend/core/src/blocksuite/ai/actions/types.ts +++ b/packages/frontend/core/src/blocksuite/ai/actions/types.ts @@ -242,6 +242,11 @@ declare global { ): AIActionTextResponse; } + type AIDocsAndFilesContext = { + docs: CopilotContextDoc[]; + files: CopilotContextFile[]; + }; + interface AIContextService { createContext: ( workspaceId: string, @@ -274,13 +279,14 @@ declare global { workspaceId: string, sessionId: string, contextId: string - ) => Promise< - | { - docs: Array; - files: Array; - } - | undefined - >; + ) => Promise; + pollContextDocsAndFiles: ( + workspaceId: string, + sessionId: string, + contextId: string, + onPoll: (result: AIDocsAndFilesContext | undefined) => void, + abortSignal: AbortSignal + ) => Promise; matchContext: ( contextId: string, content: string, diff --git a/packages/frontend/core/src/blocksuite/ai/chat-panel/chat-context.ts b/packages/frontend/core/src/blocksuite/ai/chat-panel/chat-context.ts index 926368eff9..b314cc6363 100644 --- a/packages/frontend/core/src/blocksuite/ai/chat-panel/chat-context.ts +++ b/packages/frontend/core/src/blocksuite/ai/chat-panel/chat-context.ts @@ -70,13 +70,13 @@ export type ChatBlockMessage = ChatMessage & { avatarUrl?: string; }; -export type ChipState = 'candidate' | 'processing' | 'success' | 'failed'; +export type ChipState = 'candidate' | 'processing' | 'finished' | 'failed'; export interface BaseChip { /** * candidate: the chip is a candidate for the chat * processing: the chip is processing - * success: the chip is successfully processed + * finished: the chip is successfully processed * failed: the chip is failed to process */ state: ChipState; diff --git a/packages/frontend/core/src/blocksuite/ai/chat-panel/chat-panel-chips.ts b/packages/frontend/core/src/blocksuite/ai/chat-panel/chat-panel-chips.ts index e194c75dc1..371bf7b097 100644 --- a/packages/frontend/core/src/blocksuite/ai/chat-panel/chat-panel-chips.ts +++ b/packages/frontend/core/src/blocksuite/ai/chat-panel/chat-panel-chips.ts @@ -68,6 +68,9 @@ export class ChatPanelChips extends WithDisposable(ShadowlessElement) { @property({ attribute: false }) accessor updateContext!: (context: Partial) => void; + @property({ attribute: false }) + accessor pollContextDocsAndFiles!: () => void; + @property({ attribute: false }) accessor docDisplayConfig!: DocDisplayConfig; @@ -201,6 +204,7 @@ export class ChatPanelChips extends WithDisposable(ShadowlessElement) { await this._removeFromContext(chip); } await this._addToContext(chip); + this.pollContextDocsAndFiles(); }; private readonly _updateChip = ( @@ -229,19 +233,24 @@ export class ChatPanelChips extends WithDisposable(ShadowlessElement) { private readonly _removeChip = async (chip: ChatChip) => { if (isDocChip(chip)) { - await this._removeFromContext(chip); - this.updateContext({ - chips: this.chatContextValue.chips.filter(item => { - return !isDocChip(item) || item.docId !== chip.docId; - }), - }); - } else { - await this._removeFromContext(chip); - this.updateContext({ - chips: this.chatContextValue.chips.filter(item => { - return !isFileChip(item) || item.file !== chip.file; - }), - }); + const removed = await this._removeFromContext(chip); + if (removed) { + this.updateContext({ + chips: this.chatContextValue.chips.filter(item => { + return !isDocChip(item) || item.docId !== chip.docId; + }), + }); + } + } + if (isFileChip(chip)) { + const removed = await this._removeFromContext(chip); + if (removed) { + this.updateContext({ + chips: this.chatContextValue.chips.filter(item => { + return !isFileChip(item) || item.file !== chip.file; + }), + }); + } } }; @@ -255,7 +264,8 @@ export class ChatPanelChips extends WithDisposable(ShadowlessElement) { contextId, docId: chip.docId, }); - } else { + } + if (isFileChip(chip)) { try { const blobId = await this.host.doc.blobSync.set(chip.file); const contextFile = await AIProvider.context.addContextFile(chip.file, { @@ -263,7 +273,7 @@ export class ChatPanelChips extends WithDisposable(ShadowlessElement) { blobId, }); this._updateChip(chip, { - state: 'success', + state: contextFile.status, blobId: contextFile.blobId, fileId: contextFile.id, }); @@ -276,22 +286,26 @@ export class ChatPanelChips extends WithDisposable(ShadowlessElement) { } }; - private readonly _removeFromContext = async (chip: ChatChip) => { + private readonly _removeFromContext = async ( + chip: ChatChip + ): Promise => { const contextId = await this.getContextId(); if (!contextId || !AIProvider.context) { - return; + return false; } if (isDocChip(chip)) { - await AIProvider.context.removeContextDoc({ + return await AIProvider.context.removeContextDoc({ contextId, docId: chip.docId, }); - } else if (isFileChip(chip) && chip.fileId) { - await AIProvider.context.removeContextFile({ + } + if (isFileChip(chip) && chip.fileId) { + return await AIProvider.context.removeContextFile({ contextId, fileId: chip.fileId, }); } + return true; }; private readonly _checkTokenLimit = ( @@ -305,7 +319,7 @@ export class ChatPanelChips extends WithDisposable(ShadowlessElement) { if (chip.docId === newChip.docId) { return acc + newTokenCount; } - if (chip.markdown?.value && chip.state === 'success') { + if (chip.markdown?.value && chip.state === 'finished') { const tokenCount = chip.tokenCount ?? estimateTokenCount(chip.markdown.value); return acc + tokenCount; diff --git a/packages/frontend/core/src/blocksuite/ai/chat-panel/chat-panel-input.ts b/packages/frontend/core/src/blocksuite/ai/chat-panel/chat-panel-input.ts index 5427e2a006..5dbe1a1fa1 100644 --- a/packages/frontend/core/src/blocksuite/ai/chat-panel/chat-panel-input.ts +++ b/packages/frontend/core/src/blocksuite/ai/chat-panel/chat-panel-input.ts @@ -226,7 +226,7 @@ export class ChatPanelInput extends SignalWatcher(WithDisposable(LitElement)) { private get _isNetworkDisabled() { return ( !!this.chatContextValue.images.length || - !!this.chatContextValue.chips.filter(chip => chip.state === 'success') + !!this.chatContextValue.chips.filter(chip => chip.state === 'finished') .length ); } @@ -561,7 +561,7 @@ export class ChatPanelInput extends SignalWatcher(WithDisposable(LitElement)) { : []; const contexts = this.chatContextValue.chips.reduce( (acc, chip, index) => { - if (chip.state !== 'success') { + if (chip.state !== 'finished') { return acc; } if (isDocChip(chip) && !!chip.markdown?.value) { diff --git a/packages/frontend/core/src/blocksuite/ai/chat-panel/components/doc-chip.ts b/packages/frontend/core/src/blocksuite/ai/chat-panel/components/doc-chip.ts index a450bf1c44..cf4684646b 100644 --- a/packages/frontend/core/src/blocksuite/ai/chat-panel/components/doc-chip.ts +++ b/packages/frontend/core/src/blocksuite/ai/chat-panel/components/doc-chip.ts @@ -112,7 +112,7 @@ export class ChatPanelDocChip extends SignalWatcher( const markdown = this.chip.markdown ?? new Signal(''); markdown.value = value; this.updateChip(this.chip, { - state: 'success', + state: 'finished', markdown, tokenCount, }); 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 8f8e0a25e4..ad5ef93c0a 100644 --- a/packages/frontend/core/src/blocksuite/ai/chat-panel/index.ts +++ b/packages/frontend/core/src/blocksuite/ai/chat-panel/index.ts @@ -1,6 +1,7 @@ import './chat-panel-input'; import './chat-panel-messages'; +import type { CopilotContextDoc, CopilotContextFile } from '@affine/graphql'; import type { EditorHost } from '@blocksuite/affine/block-std'; import { ShadowlessElement } from '@blocksuite/affine/block-std'; import { SignalWatcher, WithDisposable } from '@blocksuite/affine/global/lit'; @@ -35,7 +36,7 @@ import type { FileChip, } from './chat-context'; import type { ChatPanelMessages } from './chat-panel-messages'; -import { isDocContext } from './components/utils'; +import { isDocChip, isDocContext } from './components/utils'; const DEFAULT_CHAT_CONTEXT_VALUE: ChatContextValue = { quote: '', @@ -174,7 +175,7 @@ export class ChatPanel extends SignalWatcher( this._scrollToEnd(); }; - private readonly _updateChips = async () => { + private readonly _initChips = async () => { // context not initialized, show candidate chip if (!this._chatSessionId || !this._chatContextId) { return; @@ -191,12 +192,16 @@ export class ChatPanel extends SignalWatcher( (a, b) => new Date(a.createdAt).getTime() - new Date(b.createdAt).getTime() ); + let allDone = true; const chips: ChatChip[] = await Promise.all( list.map(async item => { + if (item.status === 'processing') { + allDone = false; + } if (isDocContext(item)) { return { docId: item.id, - state: 'processing', + state: item.status || 'processing', } as DocChip; } const file = await this.host.doc.blobSync.get(item.blobId); @@ -212,7 +217,7 @@ export class ChatPanel extends SignalWatcher( file: new File([file], item.name), blobId: item.blobId, fileId: item.id, - state: item.status === 'finished' ? 'success' : item.status, + state: item.status, tooltip: item.error, } as FileChip; } @@ -223,6 +228,10 @@ export class ChatPanel extends SignalWatcher( ...this.chatContextValue, chips, }; + + if (!allDone) { + await this._pollContextDocsAndFiles(); + } }; private readonly _getSessionId = async () => { @@ -285,6 +294,8 @@ export class ChatPanel extends SignalWatcher( private _width: Signal = signal(undefined); + private _pollAbortController: AbortController | null = null; + private readonly _scrollToEnd = () => { if (!this._wheelTriggered) { this._chatMessages.value?.scrollToEnd(); @@ -345,14 +356,79 @@ export class ChatPanel extends SignalWatcher( this._chatSessionId ); } - await this._updateChips(); + await this._initChips(); } catch (error) { console.error(error); } }; + private readonly _pollContextDocsAndFiles = async () => { + if (!this._chatSessionId || !this._chatContextId || !AIProvider.context) { + return; + } + if (this._pollAbortController) { + // already polling, return + return; + } + this._pollAbortController = new AbortController(); + await AIProvider.context.pollContextDocsAndFiles( + this.doc.workspace.id, + this._chatSessionId, + this._chatContextId, + this._onPoll, + this._pollAbortController.signal + ); + }; + + private readonly _onPoll = ( + result?: BlockSuitePresets.AIDocsAndFilesContext + ) => { + if (!result) { + this._abortPoll(); + return; + } + const { docs = [], files = [] } = result; + const hashMap = new Map(); + let allDone = true; + docs.forEach(doc => { + hashMap.set(doc.id, doc); + if (doc.status === 'processing') { + allDone = false; + } + }); + files.forEach(file => { + hashMap.set(file.id, file); + if (file.status === 'processing') { + allDone = false; + } + }); + const nextChips = this.chatContextValue.chips.map(chip => { + const id = isDocChip(chip) ? chip.docId : chip.fileId; + const item = id && hashMap.get(id); + if (item && item.status) { + return { + ...chip, + state: item.status, + }; + } + return chip; + }); + this.updateContext({ + chips: nextChips, + }); + if (allDone) { + this._abortPoll(); + } + }; + + private readonly _abortPoll = () => { + this._pollAbortController?.abort(); + this._pollAbortController = null; + }; + protected override updated(_changedProperties: PropertyValues) { if (_changedProperties.has('doc')) { + this._abortPoll(); this._chatSessionId = null; this._chatContextId = null; this.chatContextValue = DEFAULT_CHAT_CONTEXT_VALUE; @@ -494,6 +570,7 @@ export class ChatPanel extends SignalWatcher( .chatContextValue=${this.chatContextValue} .getContextId=${this._getContextId} .updateContext=${this.updateContext} + .pollContextDocsAndFiles=${this._pollContextDocsAndFiles} .docDisplayConfig=${this.docDisplayConfig} .docSearchMenuConfig=${this.docSearchMenuConfig} > diff --git a/packages/frontend/core/src/blocksuite/ai/provider/setup-provider.tsx b/packages/frontend/core/src/blocksuite/ai/provider/setup-provider.tsx index 68113f0cf4..88179a6011 100644 --- a/packages/frontend/core/src/blocksuite/ai/provider/setup-provider.tsx +++ b/packages/frontend/core/src/blocksuite/ai/provider/setup-provider.tsx @@ -451,6 +451,38 @@ Could you make a new website based on these notes and send back just the html fi ) => { return client.getContextDocsAndFiles(workspaceId, sessionId, contextId); }, + pollContextDocsAndFiles: async ( + workspaceId: string, + sessionId: string, + contextId: string, + onPoll: ( + result: BlockSuitePresets.AIDocsAndFilesContext | undefined + ) => void, + abortSignal: AbortSignal + ) => { + const poll = async () => { + const result = await client.getContextDocsAndFiles( + workspaceId, + sessionId, + contextId + ); + onPoll(result); + }; + + let attempts = 0; + const MIN_INTERVAL = 1000; + const MAX_INTERVAL = 30 * 1000; + + while (!abortSignal.aborted) { + await poll(); + const interval = Math.min( + MIN_INTERVAL * Math.pow(1.5, attempts), + MAX_INTERVAL + ); + attempts++; + await new Promise(resolve => setTimeout(resolve, interval)); + } + }, matchContext: async ( contextId: string, content: string, diff --git a/tests/affine-cloud-copilot/e2e/copilot.spec.ts b/tests/affine-cloud-copilot/e2e/copilot.spec.ts index 84c60c5862..17fda66701 100644 --- a/tests/affine-cloud-copilot/e2e/copilot.spec.ts +++ b/tests/affine-cloud-copilot/e2e/copilot.spec.ts @@ -1026,7 +1026,7 @@ test.describe('chat with doc', () => { expect(await chipTitle.textContent()).toBe('Untitled'); let chip = await page.getByTestId('chat-panel-chip'); // oxlint-disable-next-line unicorn/prefer-dom-node-dataset - expect(await chip.getAttribute('data-state')).toBe('success'); + expect(await chip.getAttribute('data-state')).toBe('finished'); const editorTitle = await page.locator('doc-title .inline-editor').nth(0); await editorTitle.pressSequentially('AFFiNE AI', { @@ -1048,7 +1048,7 @@ test.describe('chat with doc', () => { expect(await chipTitle.textContent()).toBe('AFFiNE AI'); // oxlint-disable-next-line unicorn/prefer-dom-node-dataset - expect(await chip.getAttribute('data-state')).toBe('success'); + expect(await chip.getAttribute('data-state')).toBe('finished'); await typeChatSequentially(page, 'What is AFiAI?'); await page.keyboard.press('Enter'); @@ -1071,6 +1071,6 @@ test.describe('chat with doc', () => { expect(await chipTitle.textContent()).toBe('AFFiNE AI'); const chip2 = await page.getByTestId('chat-panel-chip'); // oxlint-disable-next-line unicorn/prefer-dom-node-dataset - expect(await chip2.getAttribute('data-state')).toBe('success'); + expect(await chip2.getAttribute('data-state')).toBe('finished'); }); });