From 3c29f622246fd200f9cdc75b61ed1d41cba782ac Mon Sep 17 00:00:00 2001 From: yoyoyohamapi <8338436+yoyoyohamapi@users.noreply.github.com> Date: Tue, 10 Jun 2025 01:54:12 +0000 Subject: [PATCH] refactor(core): hide emebedding status tip if completed (#12720) ## Summary by CodeRabbit - **New Features** - Added real-time embedding status tracking and progress messages to the AI chat composer, with automatic updates every 10 seconds. - **Refactor** - Simplified the embedding status tooltip to display a static message, removing dynamic status updates and hover-based refresh. - **Tests** - Enhanced embedding status tooltip test by creating sample documents and extending visibility timeout to 50 seconds. --- .../core/src/blocksuite/ai/actions/types.ts | 6 ++ .../ai-chat-composer/ai-chat-composer.ts | 70 ++++++++++++++++++- .../ai-chat-input/embedding-status-tooltip.ts | 46 +----------- .../blocksuite/ai/provider/setup-provider.tsx | 24 +++++-- .../e2e/basic/chat.spec.ts | 9 ++- 5 files changed, 101 insertions(+), 54 deletions(-) diff --git a/packages/frontend/core/src/blocksuite/ai/actions/types.ts b/packages/frontend/core/src/blocksuite/ai/actions/types.ts index de8ea8539..cebe62617 100644 --- a/packages/frontend/core/src/blocksuite/ai/actions/types.ts +++ b/packages/frontend/core/src/blocksuite/ai/actions/types.ts @@ -2,6 +2,7 @@ import type { ChatHistoryOrder, ContextMatchedDocChunk, ContextMatchedFileChunk, + ContextWorkspaceEmbeddingStatus, CopilotContextCategory, CopilotContextDoc, CopilotContextFile, @@ -329,6 +330,11 @@ declare global { onPoll: (result: AIDocsAndFilesContext | undefined) => void, abortSignal: AbortSignal ) => Promise; + pollEmbeddingStatus: ( + workspaceId: string, + onPoll: (result: ContextWorkspaceEmbeddingStatus) => void, + abortSignal: AbortSignal + ) => Promise; matchContext: ( content: string, contextId?: string, 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 01a1c5689..61b813b52 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 @@ -2,6 +2,7 @@ import './ai-chat-composer-tip'; import type { ContextEmbedStatus, + ContextWorkspaceEmbeddingStatus, CopilotContextDoc, CopilotContextFile, CopilotDocType, @@ -34,6 +35,8 @@ import type { } from '../ai-chat-input'; import { MAX_IMAGE_COUNT } from '../ai-chat-input/const'; +export const EMBEDDING_STATUS_CHECK_INTERVAL = 10000; + export class AIChatComposer extends SignalWatcher( WithDisposable(ShadowlessElement) ) { @@ -108,6 +111,12 @@ export class AIChatComposer extends SignalWatcher( @state() accessor chips: ChatChip[] = []; + @state() + accessor embeddingProgressText = 'Loading embedding status...'; + + @state() + accessor embeddingCompleted = false; + private _isInitialized = false; private _isLoading = false; @@ -116,6 +125,8 @@ export class AIChatComposer extends SignalWatcher( private _pollAbortController: AbortController | null = null; + private _pollEmbeddingStatusAbortController: AbortController | null = null; + override render() { return html` AI outputs can be misleading or wrong`, - html``, - ]} + this.embeddingCompleted + ? null + : html``, + ].filter(Boolean)} .loop=${false} > @@ -174,10 +189,20 @@ export class AIChatComposer extends SignalWatcher( if (isVisible && !this._isInitialized) { this._initComposer().catch(console.error); } + if (!isVisible) { + this._abortPoll(); + this._abortPollEmbeddingStatus(); + } }) ); } + override disconnectedCallback() { + super.disconnectedCallback(); + this._abortPoll(); + this._abortPollEmbeddingStatus(); + } + protected override willUpdate(_changedProperties: PropertyValues) { if (_changedProperties.has('doc')) { this._resetComposer(); @@ -316,6 +341,40 @@ export class AIChatComposer extends SignalWatcher( ); }; + private readonly _pollEmbeddingStatus = async () => { + if (this._pollEmbeddingStatusAbortController) { + this._pollEmbeddingStatusAbortController.abort(); + } + this._pollEmbeddingStatusAbortController = new AbortController(); + const signal = this._pollEmbeddingStatusAbortController.signal; + + try { + await AIProvider.context?.pollEmbeddingStatus( + this.host.std.workspace.id, + (status: ContextWorkspaceEmbeddingStatus) => { + if (!status) { + this.embeddingProgressText = 'Loading embedding status...'; + this.embeddingCompleted = false; + return; + } + const completed = status.embedded === status.total; + this.embeddingCompleted = completed; + if (completed) { + this.embeddingProgressText = + 'Embedding finished. You are getting the best results!'; + } else { + this.embeddingProgressText = + 'File not embedded yet. Results will improve after embedding.'; + } + }, + signal + ); + } catch { + this.embeddingProgressText = 'Failed to load embedding status...'; + this.embeddingCompleted = false; + } + }; + private readonly _onPoll = ( result?: BlockSuitePresets.AIDocsAndFilesContext ) => { @@ -378,6 +437,11 @@ export class AIChatComposer extends SignalWatcher( this._pollAbortController = null; }; + private readonly _abortPollEmbeddingStatus = () => { + this._pollEmbeddingStatusAbortController?.abort(); + this._pollEmbeddingStatusAbortController = null; + }; + private readonly _initComposer = async () => { if (!this.isVisible.value) return; if (this._isLoading) return; @@ -394,12 +458,14 @@ export class AIChatComposer extends SignalWatcher( if (needPoll) { await this._pollContextDocsAndFiles(); } + await this._pollEmbeddingStatus(); this._isLoading = false; this._isInitialized = true; }; private readonly _resetComposer = () => { this._abortPoll(); + this._abortPollEmbeddingStatus(); this.chips = []; this._contextId = undefined; this._isLoading = false; diff --git a/packages/frontend/core/src/blocksuite/ai/components/ai-chat-input/embedding-status-tooltip.ts b/packages/frontend/core/src/blocksuite/ai/components/ai-chat-input/embedding-status-tooltip.ts index e0b638027..d57ebccda 100644 --- a/packages/frontend/core/src/blocksuite/ai/components/ai-chat-input/embedding-status-tooltip.ts +++ b/packages/frontend/core/src/blocksuite/ai/components/ai-chat-input/embedding-status-tooltip.ts @@ -1,11 +1,7 @@ import { SignalWatcher } from '@blocksuite/affine/global/lit'; import { unsafeCSSVar } from '@blocksuite/affine/shared/theme'; -import type { EditorHost } from '@blocksuite/affine/std'; import { css, html, LitElement } from 'lit'; -import { property, state } from 'lit/decorators.js'; -import { debounce, noop } from 'lodash-es'; - -import { AIProvider } from '../../provider/ai-provider'; +import { property } from 'lit/decorators.js'; export class AIChatEmbeddingStatusTooltip extends SignalWatcher(LitElement) { static override styles = css` @@ -38,47 +34,8 @@ export class AIChatEmbeddingStatusTooltip extends SignalWatcher(LitElement) { `; @property({ attribute: false }) - accessor host!: EditorHost; - - @state() accessor progressText = 'Loading embedding status...'; - override connectedCallback() { - super.connectedCallback(); - this._updateEmbeddingStatus().catch(noop); - } - - private async _updateEmbeddingStatus() { - try { - const status = await AIProvider.embedding?.getEmbeddingStatus( - this.host.std.workspace.id - ); - if (!status) { - this.progressText = 'Loading embedding status...'; - return; - } - const completed = status.embedded === status.total; - if (completed) { - this.progressText = - 'Embedding finished. You are getting the best results!'; - } else { - this.progressText = - 'File not embedded yet. Results will improve after embedding.'; - } - this.requestUpdate(); - } catch { - this.progressText = 'Failed to load embedding status...'; - } - } - - private readonly _handleCheckStatusMouseEnter = debounce( - () => { - this._updateEmbeddingStatus().catch(noop); - }, - 1000, - { leading: true } - ); - override render() { return html`
Check status setTimeout(resolve, interval)); } }, + pollEmbeddingStatus: async ( + workspaceId: string, + onPoll: (result: ContextWorkspaceEmbeddingStatus) => void, + abortSignal: AbortSignal + ) => { + const poll = async () => { + const result = await client.getEmbeddingStatus(workspaceId); + onPoll(result); + }; + + const INTERVAL = 10 * 1000; + + while (!abortSignal.aborted) { + await poll(); + await new Promise(resolve => setTimeout(resolve, INTERVAL)); + } + }, matchContext: async ( content: string, contextId?: string, @@ -792,12 +810,6 @@ Could you make a new website based on these notes and send back just the html fi return client.forkSession(options); }); - AIProvider.provide('embedding', { - getEmbeddingStatus: (workspaceId: string) => { - return client.getEmbeddingStatus(workspaceId); - }, - }); - const disposeRequestLoginHandler = AIProvider.slots.requestLogin.subscribe( () => { globalDialogService.open('sign-in', {}); diff --git a/tests/affine-cloud-copilot/e2e/basic/chat.spec.ts b/tests/affine-cloud-copilot/e2e/basic/chat.spec.ts index d408b95a4..c7260398f 100644 --- a/tests/affine-cloud-copilot/e2e/basic/chat.spec.ts +++ b/tests/affine-cloud-copilot/e2e/basic/chat.spec.ts @@ -19,11 +19,18 @@ test.describe('AIBasic/Chat', () => { test('should display embedding status tooltip', async ({ loggedInPage: page, + utils, }) => { + await utils.editor.createDoc(page, 'Doc 1', 'doc1'); + await utils.editor.createDoc(page, 'Doc 2', 'doc2'); + await utils.editor.createDoc(page, 'Doc 3', 'doc3'); + await utils.editor.createDoc(page, 'Doc 4', 'doc4'); + await utils.editor.createDoc(page, 'Doc 5', 'doc5'); + const check = await page.getByTestId( 'ai-chat-embedding-status-tooltip-check' ); - await expect(check).toBeVisible(); + await expect(check).toBeVisible({ timeout: 50 * 1000 }); await check.hover(); const tooltip = await page.getByTestId('ai-chat-embedding-status-tooltip');