refactor(core): workspace mutation effect (#12488)

### TL;DR

* refactor: workspace embedding mutation effect
* tests: error display for workspace embedding
This commit is contained in:
yoyoyohamapi
2025-05-26 07:17:37 +00:00
parent da22391910
commit c06c72e108
6 changed files with 231 additions and 172 deletions
@@ -11,15 +11,8 @@ import {
onStart,
smartRetry,
} from '@toeverything/infra';
import { EMPTY, interval, of, Subject } from 'rxjs';
import {
concatMap,
exhaustMap,
mergeMap,
switchMap,
takeUntil,
tap,
} from 'rxjs/operators';
import { EMPTY, interval, Subject } from 'rxjs';
import { exhaustMap, mergeMap, switchMap, takeUntil } from 'rxjs/operators';
import { COUNT_PER_PAGE } from '../constants';
import type { EmbeddingStore } from '../stores/embedding';
@@ -110,31 +103,13 @@ export class Embedding extends Entity {
})
);
setEnabled = effect(
exhaustMap((enabled: boolean) => {
return fromPromise(signal =>
this.store.updateEnabled(
this.workspaceService.workspace.id,
enabled,
signal
)
).pipe(
smartRetry(),
concatMap(() => {
this.getEnabled();
return EMPTY;
}),
catchErrorInto(this.error$, error => {
logger.error(
'Failed to update workspace doc embedding enabled',
error
);
}),
onStart(() => this.isEnabledLoading$.setValue(true)),
onComplete(() => this.isEnabledLoading$.setValue(false))
);
})
);
setEnabled = (enabled: boolean) => {
return this.store
.updateEnabled(this.workspaceService.workspace.id, enabled)
.then(() => {
this.getEnabled();
});
};
getIgnoredDocs = effect(
exhaustMap(() => {
@@ -158,30 +133,19 @@ export class Embedding extends Entity {
})
);
updateIgnoredDocs = effect(
exhaustMap(({ add, remove }: { add: string[]; remove: string[] }) => {
return fromPromise(signal =>
this.store.updateIgnoredDocs(
this.workspaceService.workspace.id,
add,
remove,
signal
)
).pipe(
smartRetry(),
concatMap(() => {
this.getIgnoredDocs();
return EMPTY;
}),
catchErrorInto(this.error$, error => {
logger.error(
'Failed to update workspace doc embedding ignored docs',
error
);
})
);
})
);
updateIgnoredDocs = ({
add,
remove,
}: {
add: string[];
remove: string[];
}) => {
return this.store
.updateIgnoredDocs(this.workspaceService.workspace.id, add, remove)
.then(() => {
this.getIgnoredDocs();
});
};
getAttachments = effect(
exhaustMap((pagination: PaginationInput) => {
@@ -219,95 +183,60 @@ export class Embedding extends Entity {
})
);
addAttachments = effect(
// Support parallel upload
mergeMap((files: File[]) => {
const generateLocalId = () =>
Math.random().toString(36).slice(2) + Date.now();
const localAttachments: LocalAttachmentFile[] = files.map(file => ({
localId: generateLocalId(),
fileName: file.name,
mimeType: file.type,
size: file.size,
createdAt: file.lastModified,
status: 'uploading',
}));
addAttachments = (files: File[]) => {
const generateLocalId = () =>
Math.random().toString(36).slice(2) + Date.now();
const localAttachments: LocalAttachmentFile[] = files.map(file => ({
localId: generateLocalId(),
fileName: file.name,
mimeType: file.type,
size: file.size,
createdAt: file.lastModified,
status: 'uploading',
}));
return of({ files, localAttachments }).pipe(
// Refresh uploading attachments immediately
tap(({ localAttachments }) => {
this.uploadingAttachments$.next([
...localAttachments,
...this.uploadingAttachments$.value,
]);
}),
// Uploading embedding files
switchMap(({ files }) => {
return fromPromise(signal =>
this.store.addEmbeddingFiles(
this.workspaceService.workspace.id,
files,
signal
)
);
}),
// Refresh uploading attachments
tap(() => {
this.uploadingAttachments$.next(
this.uploadingAttachments$.value.filter(
att => !localAttachments.some(l => l.localId === att.localId)
)
);
this.getAttachments({ first: COUNT_PER_PAGE, after: null });
}),
catchErrorInto(this.error$, error => {
this.uploadingAttachments$.next(
this.uploadingAttachments$.value.map(att =>
localAttachments.some(l => l.localId === att.localId)
? { ...att, status: 'error', errorMessage: String(error) }
: att
)
);
logger.error(
'Failed to add workspace doc embedding attachments',
error
);
})
);
})
);
this.uploadingAttachments$.next([
...localAttachments,
...this.uploadingAttachments$.value,
]);
removeAttachment = effect(
exhaustMap((id: string) => {
const localIndex = this.uploadingAttachments$.value.findIndex(
att => att.localId === id
);
if (localIndex !== -1) {
this.store
.addEmbeddingFiles(this.workspaceService.workspace.id, files)
.then(() => {
this.uploadingAttachments$.next(
this.uploadingAttachments$.value.filter(att => att.localId !== id)
this.uploadingAttachments$.value.filter(
att => !localAttachments.some(l => l.localId === att.localId)
)
);
return EMPTY;
}
return fromPromise(signal =>
this.store.removeEmbeddingFile(
this.workspaceService.workspace.id,
id,
signal
)
).pipe(
concatMap(() => {
this.getAttachments({ first: COUNT_PER_PAGE, after: null });
return EMPTY;
}),
catchErrorInto(this.error$, error => {
logger.error(
'Failed to remove workspace doc embedding attachment',
error
);
})
this.getAttachments({ first: COUNT_PER_PAGE, after: null });
})
.catch(error => {
this.uploadingAttachments$.next(
this.uploadingAttachments$.value.map(att =>
localAttachments.some(l => l.localId === att.localId)
? { ...att, status: 'error', errorMessage: String(error) }
: att
)
);
});
};
removeAttachment = (id: string) => {
const localIndex = this.uploadingAttachments$.value.findIndex(
att => att.localId === id
);
if (localIndex !== -1) {
this.uploadingAttachments$.next(
this.uploadingAttachments$.value.filter(att => att.localId !== id)
);
})
);
return Promise.resolve();
}
return this.store
.removeEmbeddingFile(this.workspaceService.workspace.id, id)
.then(() => {
this.getAttachments({ first: COUNT_PER_PAGE, after: null });
});
};
startEmbeddingProgressPolling() {
this.stopEmbeddingProgressPolling();
@@ -355,10 +284,6 @@ export class Embedding extends Entity {
this.getEnabled.unsubscribe();
this.getAttachments.unsubscribe();
this.getIgnoredDocs.unsubscribe();
this.updateIgnoredDocs.unsubscribe();
this.addAttachments.unsubscribe();
this.removeAttachment.unsubscribe();
this.setEnabled.unsubscribe();
this.stopEmbeddingProgress$.next();
this.getEmbeddingProgress.unsubscribe();
}
@@ -1,4 +1,4 @@
import { Button, Switch } from '@affine/component';
import { Button, notify, Switch } from '@affine/component';
import {
SettingHeader,
SettingRow,
@@ -6,6 +6,7 @@ import {
} from '@affine/component/setting-components';
import { Upload } from '@affine/core/components/pure/file-upload';
import { WorkspaceDialogService } from '@affine/core/modules/dialogs';
import { UserFriendlyError } from '@affine/error';
import { useI18n } from '@affine/i18n';
import track from '@affine/track';
import { useLiveData, useService } from '@toeverything/infra';
@@ -48,9 +49,19 @@ export const EmbeddingSettings: React.FC<EmbeddingSettingsProps> = () => {
control: 'Workspace embedding',
option: checked ? 'on' : 'off',
});
embeddingService.embedding.setEnabled(checked);
embeddingService.embedding.setEnabled(checked).catch(error => {
const err = UserFriendlyError.fromAny(error);
notify.error({
title:
t[
'com.affine.settings.workspace.indexer-embedding.embedding.switch.error'
](),
message: t[`error.${err.name}`](err.data),
});
});
},
[embeddingService.embedding]
[embeddingService.embedding, t]
);
const handleAttachmentUpload = useCallback(
@@ -67,9 +78,18 @@ export const EmbeddingSettings: React.FC<EmbeddingSettingsProps> = () => {
const handleAttachmentsDelete = useCallback(
(fileId: string) => {
embeddingService.embedding.removeAttachment(fileId);
embeddingService.embedding.removeAttachment(fileId).catch(error => {
const err = UserFriendlyError.fromAny(error);
notify.error({
title:
t[
'com.affine.settings.workspace.indexer-embedding.embedding.remove-attachment.error'
](),
message: t[`error.${err.name}`](err.data),
});
});
},
[embeddingService.embedding]
[embeddingService.embedding, t]
);
const handleAttachmentsPageChange = useCallback(
@@ -102,7 +122,18 @@ export const EmbeddingSettings: React.FC<EmbeddingSettingsProps> = () => {
});
const add = selectedIds.filter(id => !initialIds?.includes(id));
const remove = initialIds?.filter(id => !selectedIds.includes(id));
embeddingService.embedding.updateIgnoredDocs({ add, remove });
embeddingService.embedding
.updateIgnoredDocs({ add, remove })
.catch(error => {
const err = UserFriendlyError.fromAny(error);
notify.error({
title:
t[
'com.affine.settings.workspace.indexer-embedding.embedding.update-ignored-docs.error'
](),
message: t[`error.${err.name}`](err.data),
});
});
}
);
}, [
@@ -110,6 +141,7 @@ export const EmbeddingSettings: React.FC<EmbeddingSettingsProps> = () => {
isIgnoredDocsLoading,
workspaceDialogService,
embeddingService.embedding,
t,
]);
useEffect(() => {
+12
View File
@@ -6247,6 +6247,18 @@ export function useAFFiNEI18N(): {
* `AI can call files embedded in the workspace.`
*/
["com.affine.settings.workspace.indexer-embedding.embedding.switch.description"](): string;
/**
* `Failed to update workspace doc embedding enabled`
*/
["com.affine.settings.workspace.indexer-embedding.embedding.switch.error"](): string;
/**
* `Failed to remove attachment from embedding`
*/
["com.affine.settings.workspace.indexer-embedding.embedding.remove-attachment.error"](): string;
/**
* `Failed to update ignored docs`
*/
["com.affine.settings.workspace.indexer-embedding.embedding.update-ignored-docs.error"](): string;
/**
* `Embedding progress`
*/
@@ -1560,6 +1560,9 @@
"com.affine.settings.workspace.indexer-embedding.embedding.select-doc": "Select doc",
"com.affine.settings.workspace.indexer-embedding.embedding.switch.title": "Workspace Embedding",
"com.affine.settings.workspace.indexer-embedding.embedding.switch.description": "AI can call files embedded in the workspace.",
"com.affine.settings.workspace.indexer-embedding.embedding.switch.error": "Failed to update workspace doc embedding enabled",
"com.affine.settings.workspace.indexer-embedding.embedding.remove-attachment.error": "Failed to remove attachment from embedding",
"com.affine.settings.workspace.indexer-embedding.embedding.update-ignored-docs.error": "Failed to update ignored docs",
"com.affine.settings.workspace.indexer-embedding.embedding.progress.title": "Embedding progress",
"com.affine.settings.workspace.indexer-embedding.embedding.progress.syncing": "Syncing",
"com.affine.settings.workspace.indexer-embedding.embedding.progress.synced": "Synced",
@@ -43,6 +43,23 @@ test.describe('AISettings/Embedding', () => {
await utils.settings.waitForWorkspaceEmbeddingSwitchToBe(page, true);
});
test('should show error message if enable workspace embedding failed', async ({
loggedInPage: page,
utils,
}) => {
await utils.settings.enableWorkspaceEmbedding(page);
await utils.settings.disableWorkspaceEmbedding(page);
await utils.settings.waitForWorkspaceEmbeddingSwitchToBe(page, false);
await page.context().setOffline(true);
await utils.settings.enableWorkspaceEmbedding(page, false);
await expect(
page.getByText(/Failed to update workspace doc embedding enabled/i)
).toBeVisible();
await page.context().setOffline(false);
});
test('should show embedding progress', async ({
loggedInPage: page,
utils,
@@ -297,6 +314,36 @@ test.describe('AISettings/Embedding', () => {
await utils.settings.removeAttachment(page, 'document1.txt');
});
test('should show error message if remove attachment failed', async ({
loggedInPage: page,
utils,
}) => {
await utils.settings.enableWorkspaceEmbedding(page);
const textContent = 'WorkspaceEBEEE is a cute cat';
const attachments = [
{
name: 'document1.txt',
mimeType: 'text/plain',
buffer: Buffer.from(textContent),
},
];
await utils.settings.uploadWorkspaceEmbedding(page, attachments);
const attachmentList = await page.getByTestId(
'workspace-embedding-setting-attachment-list'
);
await expect(
attachmentList.getByTestId('workspace-embedding-setting-attachment-item')
).toHaveCount(1);
await page.context().setOffline(true);
await utils.settings.clickRemoveAttachment(page, 'document1.txt');
await expect(
page.getByText(/Failed to remove attachment from embedding/i)
).toBeVisible();
await page.context().setOffline(false);
});
test('should support remove error attachment directly', async ({
loggedInPage: page,
utils,
@@ -393,4 +440,23 @@ test.describe('AISettings/Embedding', () => {
expect(content).toMatch(/I dont know/i);
}).toPass({ timeout: 20000 });
});
test('should show error message if update ignored docs failed', async ({
loggedInPage: page,
utils,
}) => {
await utils.settings.enableWorkspaceEmbedding(page);
await utils.settings.closeSettingsPanel(page);
await utils.editor.createDoc(page, 'Test Doc', 'HelloWorld');
// Ignore docs
await utils.settings.openSettingsPanel(page);
await page.context().setOffline(true);
await utils.settings.ignoreDocForEmbedding(page, 'Test Doc', false);
await expect(
page.getByText(/Failed to update ignored docs/i)
).toBeVisible();
await page.context().setOffline(false);
});
});
@@ -46,20 +46,30 @@ export class SettingsPanelUtils {
await input.click();
}
public static async enableWorkspaceEmbedding(page: Page) {
public static async enableWorkspaceEmbedding(
page: Page,
waitForEnabled = true
) {
const enabled = await this.isWorkspaceEmbeddingEnabled(page);
if (!enabled) {
await this.toggleWorkspaceEmbedding(page);
}
await this.waitForWorkspaceEmbeddingSwitchToBe(page, true);
if (waitForEnabled) {
await this.waitForWorkspaceEmbeddingSwitchToBe(page, true);
}
}
public static async disableWorkspaceEmbedding(page: Page) {
public static async disableWorkspaceEmbedding(
page: Page,
waitForDisabled = true
) {
const enabled = await this.isWorkspaceEmbeddingEnabled(page);
if (enabled) {
await this.toggleWorkspaceEmbedding(page);
}
await this.waitForWorkspaceEmbeddingSwitchToBe(page, false);
if (waitForDisabled) {
await this.waitForWorkspaceEmbeddingSwitchToBe(page, false);
}
}
public static async uploadWorkspaceEmbedding(
@@ -102,7 +112,7 @@ export class SettingsPanelUtils {
}
}
public static async removeAttachment(
public static async clickRemoveAttachment(
page: Page,
attachment: string,
shouldConfirm = true
@@ -116,6 +126,14 @@ export class SettingsPanelUtils {
if (shouldConfirm) {
await page.getByTestId('confirm-modal-confirm').click();
}
}
public static async removeAttachment(
page: Page,
attachment: string,
shouldConfirm = true
) {
await this.clickRemoveAttachment(page, attachment, shouldConfirm);
await page
.getByTestId('workspace-embedding-setting-attachment-item')
.filter({ hasText: attachment })
@@ -124,7 +142,11 @@ export class SettingsPanelUtils {
});
}
public static async ignoreDocForEmbedding(page: Page, doc: string) {
public static async ignoreDocForEmbedding(
page: Page,
doc: string,
shouldWaitForRefresh = true
) {
// Open Dos Searcher
const ignoreDocsButton = await page.getByTestId(
'workspace-embedding-setting-ignore-docs-button'
@@ -137,25 +159,24 @@ export class SettingsPanelUtils {
await searchInput.focus();
await page.keyboard.insertText(doc);
const pageListItem = searcher.getByTestId('page-list-item');
const pageListItem = searcher.getByTestId('doc-list-item');
await expect(pageListItem).toHaveCount(1);
const pageListItemTitle = pageListItem.getByTestId(
'page-list-item-title-text'
);
const pageListItemTitle = pageListItem.getByTestId('doc-list-item-title');
await expect(pageListItemTitle).toHaveText(doc);
await pageListItem.getByTestId('affine-checkbox').check();
await pageListItem.click();
await searcher.getByTestId('doc-selector-confirm-button').click();
const ignoredDocs = await page.getByTestId(
'workspace-embedding-setting-ignore-docs-list'
);
await expect(
ignoredDocs
.getByTestId('workspace-embedding-setting-ignore-docs-list-item')
.filter({ hasText: doc })
).toBeVisible();
if (shouldWaitForRefresh) {
const ignoredDocs = await page.getByTestId(
'workspace-embedding-setting-ignore-docs-list'
);
await expect(
ignoredDocs
.getByTestId('workspace-embedding-setting-ignore-docs-list-item')
.filter({ hasText: doc })
).toBeVisible();
}
}
public static async clearAllIgnoredDocs(page: Page) {