feat(core): embedding progress (#12367)

### TL;DR

feat: show embedding progress in settings panel

![截屏2025-05-19 20.25.19.png](https://graphite-user-uploaded-assets-prod.s3.amazonaws.com/MyktQ6Qwc7H6TiRCFoYN/59d8f9ef-0876-4ed5-9c09-db12686adb47.png)

### What changed

* show embedding progress in settings panel
* polling embedding status based on RxJS

<!-- This is an auto-generated comment: release notes by coderabbit.ai -->
## Summary by CodeRabbit

- **New Features**
  - Added real-time embedding progress tracking and display in embedding settings, including a visual progress bar and status messages.
  - Introduced localized text for embedding progress statuses.
  - Added an optional test ID attribute to the progress bar component for improved testing.
- **Style**
  - Added new styles for embedding progress UI elements.
- **Tests**
  - Added an end-to-end test to verify embedding progress is displayed correctly in the settings UI.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
This commit is contained in:
yoyoyohamapi
2025-05-21 05:07:13 +00:00
parent 8f352580a7
commit d70f09b498
9 changed files with 225 additions and 5 deletions

View File

@@ -11,8 +11,14 @@ import {
onStart,
smartRetry,
} from '@toeverything/infra';
import { EMPTY } from 'rxjs';
import { concatMap, exhaustMap, mergeMap } from 'rxjs/operators';
import { EMPTY, interval, Subject } from 'rxjs';
import {
concatMap,
exhaustMap,
mergeMap,
switchMap,
takeUntil,
} from 'rxjs/operators';
import { COUNT_PER_PAGE } from '../constants';
import type { EmbeddingStore } from '../stores/embedding';
@@ -35,6 +41,11 @@ interface Attachments {
type IgnoredDocs = IgnoredDoc[];
interface EmbeddingProgress {
embedded: number;
total: number;
}
export class Embedding extends Entity {
enabled$ = new LiveData<boolean>(false);
error$ = new LiveData<any>(null);
@@ -50,6 +61,11 @@ export class Embedding extends Entity {
isEnabledLoading$ = new LiveData(false);
isAttachmentsLoading$ = new LiveData(false);
isIgnoredDocsLoading$ = new LiveData(false);
embeddingProgress$ = new LiveData<EmbeddingProgress | null>(null);
isEmbeddingProgressLoading$ = new LiveData(false);
private readonly EMBEDDING_PROGRESS_POLL_INTERVAL = 3000;
private readonly stopEmbeddingProgress$ = new Subject<void>();
constructor(
private readonly workspaceService: WorkspaceService,
@@ -59,6 +75,7 @@ export class Embedding extends Entity {
this.getEnabled();
this.getAttachments({ first: COUNT_PER_PAGE, after: null });
this.getIgnoredDocs();
this.getEmbeddingProgress();
}
getEnabled = effect(
@@ -228,6 +245,48 @@ export class Embedding extends Entity {
})
);
startEmbeddingProgressPolling() {
this.stopEmbeddingProgressPolling();
this.getEmbeddingProgress();
}
stopEmbeddingProgressPolling() {
this.stopEmbeddingProgress$.next();
}
getEmbeddingProgress = effect(
exhaustMap(() => {
return interval(this.EMBEDDING_PROGRESS_POLL_INTERVAL).pipe(
takeUntil(this.stopEmbeddingProgress$),
switchMap(() =>
fromPromise(signal =>
this.store.getEmbeddingProgress(
this.workspaceService.workspace.id,
signal
)
).pipe(
smartRetry(),
mergeMap(value => {
this.embeddingProgress$.next(value);
if (value && value.embedded === value.total) {
this.stopEmbeddingProgressPolling();
}
return EMPTY;
}),
catchErrorInto(this.error$, error => {
logger.error(
'Failed to fetch workspace embedding progress',
error
);
}),
onStart(() => this.isEmbeddingProgressLoading$.setValue(true)),
onComplete(() => this.isEmbeddingProgressLoading$.setValue(false))
)
)
);
})
);
override dispose(): void {
this.getEnabled.unsubscribe();
this.getAttachments.unsubscribe();
@@ -236,5 +295,7 @@ export class Embedding extends Entity {
this.addAttachments.unsubscribe();
this.removeAttachment.unsubscribe();
this.setEnabled.unsubscribe();
this.stopEmbeddingProgress$.next();
this.getEmbeddingProgress.unsubscribe();
}
}

View File

@@ -5,6 +5,7 @@ import {
getAllWorkspaceEmbeddingIgnoredDocsQuery,
getWorkspaceConfigQuery,
getWorkspaceEmbeddingFilesQuery,
getWorkspaceEmbeddingStatusQuery,
type PaginationInput,
removeWorkspaceEmbeddingFilesMutation,
removeWorkspaceEmbeddingIgnoredDocsMutation,
@@ -175,4 +176,19 @@ export class EmbeddingStore extends Store {
});
return data.workspace.embedding.files;
}
async getEmbeddingProgress(workspaceId: string, signal?: AbortSignal) {
if (!this.workspaceServerService.server) {
throw new Error('No Server');
}
const data = await this.workspaceServerService.server.gql({
query: getWorkspaceEmbeddingStatusQuery,
variables: {
workspaceId,
},
context: { signal },
});
return data.queryWorkspaceEmbeddingStatus;
}
}

View File

@@ -0,0 +1,60 @@
import { Progress } from '@affine/component';
import { useI18n } from '@affine/i18n';
import { embeddingProgress, embeddingProgressTitle } from './styles-css';
interface EmbeddingProgressProps {
status: {
embedded: number;
total: number;
} | null;
}
const EmbeddingProgress: React.FC<EmbeddingProgressProps> = ({ status }) => {
const t = useI18n();
const loading = status === null;
const percent = loading
? 0
: status.total === 0
? 1
: status.embedded / status.total;
const progress = Math.round(percent * 100);
const synced = percent === 1;
return (
<div className={embeddingProgress} data-testid="embedding-progress-wrapper">
<div
className={embeddingProgressTitle}
data-testid="embedding-progress-title"
data-progress={loading ? 'loading' : synced ? 'synced' : 'syncing'}
>
<div>
{loading
? t[
'com.affine.settings.workspace.indexer-embedding.embedding.progress.loading-sync-status'
]()
: synced
? t[
'com.affine.settings.workspace.indexer-embedding.embedding.progress.synced'
]()
: t[
'com.affine.settings.workspace.indexer-embedding.embedding.progress.syncing'
]()}
</div>
{loading ? null : (
<div data-testid="embedding-progress-count">{`${status.embedded}/${status.total}`}</div>
)}
</div>
<Progress
testId="embedding-progress"
value={progress}
readonly
style={{ visibility: loading ? 'hidden' : 'visible' }}
/>
</div>
);
};
export default EmbeddingProgress;

View File

@@ -9,10 +9,11 @@ import { WorkspaceDialogService } from '@affine/core/modules/dialogs';
import { useI18n } from '@affine/i18n';
import { useLiveData, useService } from '@toeverything/infra';
import type React from 'react';
import { useCallback, useMemo } from 'react';
import { useCallback, useEffect, useMemo } from 'react';
import { EmbeddingService } from '../services/embedding';
import { Attachments } from './attachments';
import EmbeddingProgress from './embedding-progress';
import { IgnoredDocs } from './ignored-docs';
interface EmbeddingSettingsProps {}
@@ -23,6 +24,10 @@ export const EmbeddingSettings: React.FC<EmbeddingSettingsProps> = () => {
const embeddingEnabled = useLiveData(embeddingService.embedding.enabled$);
const attachments = useLiveData(embeddingService.embedding.attachments$);
const ignoredDocs = useLiveData(embeddingService.embedding.ignoredDocs$);
const embeddingProgress = useLiveData(
embeddingService.embedding.embeddingProgress$
);
const isIgnoredDocsLoading = useLiveData(
embeddingService.embedding.isIgnoredDocsLoading$
);
@@ -34,7 +39,6 @@ export const EmbeddingSettings: React.FC<EmbeddingSettingsProps> = () => {
[attachments]
);
const ignoredDocNodes = ignoredDocs;
const workspaceDialogService = useService(WorkspaceDialogService);
const handleEmbeddingToggle = useCallback(
@@ -94,6 +98,13 @@ export const EmbeddingSettings: React.FC<EmbeddingSettingsProps> = () => {
embeddingService.embedding,
]);
useEffect(() => {
embeddingService.embedding.startEmbeddingProgressPolling();
return () => {
embeddingService.embedding.stopEmbeddingProgressPolling();
};
}, [embeddingService.embedding]);
return (
<>
<SettingHeader
@@ -120,6 +131,15 @@ export const EmbeddingSettings: React.FC<EmbeddingSettingsProps> = () => {
/>
</SettingRow>
<SettingRow
name={t[
'com.affine.settings.workspace.indexer-embedding.embedding.progress.title'
]()}
style={{ marginBottom: '0px' }}
/>
<EmbeddingProgress status={embeddingProgress} />
<SettingRow
name={t[
'com.affine.settings.workspace.indexer-embedding.embedding.additional-attachments.title'

View File

@@ -98,3 +98,22 @@ export const docItemInfo = css({
gap: '12px',
alignItems: 'center',
});
export const embeddingProgress = css({
display: 'flex',
flexDirection: 'column',
alignItems: 'center',
width: '100%',
paddingBottom: '16px',
fontSize: '14px',
fontWeight: 400,
color: cssVar('textSecondaryColor'),
});
export const embeddingProgressTitle = css({
textAlign: 'left',
width: '100%',
display: 'flex',
alignItems: 'center',
justifyContent: 'space-between',
});