From 6f9361caee6bbb244837c82813b45037188cbfbd Mon Sep 17 00:00:00 2001 From: darkskygit Date: Tue, 20 May 2025 05:16:45 +0000 Subject: [PATCH] feat(server): trigger workspace embedding (#12328) fix AI-127 ## Summary by CodeRabbit - **New Features** - Added automated event handling for workspace updates and document embedding, streamlining document embedding workflows. - Introduced detection and queuing of documents needing embedding, excluding ignored documents. - **Improvements** - Enhanced performance of embedding-related searches by filtering results at the database level. - Increased concurrency for embedding job processing to improve throughput. - **Bug Fixes** - Improved error handling and fallback for missing document titles during embedding. - Added safeguards to skip invalid embedding jobs based on document identifiers. - **Tests** - Expanded test coverage for document embedding and ignored document filtering. - Updated end-to-end tests to use dynamic content for improved reliability. - Added synchronization waits in document creation utilities to improve test stability. --- .docker/selfhost/schema.json | 4 +- .../copilot-workspace.spec.ts.md | 16 +++++ .../copilot-workspace.spec.ts.snap | Bin 579 -> 606 bytes .../models/copilot-workspace.spec.ts | 63 ++++++++++++++++++ .../server/src/base/job/queue/config.ts | 2 +- .../server/src/models/copilot-context.ts | 8 ++- .../server/src/models/copilot-workspace.ts | 27 ++++++++ .../backend/server/src/models/workspace.ts | 4 ++ .../server/src/plugins/copilot/context/job.ts | 51 +++++++++++--- .../src/plugins/copilot/context/types.ts | 4 ++ .../e2e/chat-with/collections.spec.ts | 26 +++++--- .../e2e/chat-with/tags.spec.ts | 39 ++++++++--- .../e2e/utils/editor-utils.ts | 4 ++ 13 files changed, 213 insertions(+), 35 deletions(-) diff --git a/.docker/selfhost/schema.json b/.docker/selfhost/schema.json index a1b50b112..2ed864dba 100644 --- a/.docker/selfhost/schema.json +++ b/.docker/selfhost/schema.json @@ -48,14 +48,14 @@ }, "queues.copilot": { "type": "object", - "description": "The config for copilot job queue\n@default {\"concurrency\":1}", + "description": "The config for copilot job queue\n@default {\"concurrency\":5}", "properties": { "concurrency": { "type": "number" } }, "default": { - "concurrency": 1 + "concurrency": 5 } }, "queues.doc": { diff --git a/packages/backend/server/src/__tests__/models/__snapshots__/copilot-workspace.spec.ts.md b/packages/backend/server/src/__tests__/models/__snapshots__/copilot-workspace.spec.ts.md index ef8b1f448..315990061 100644 --- a/packages/backend/server/src/__tests__/models/__snapshots__/copilot-workspace.spec.ts.md +++ b/packages/backend/server/src/__tests__/models/__snapshots__/copilot-workspace.spec.ts.md @@ -73,3 +73,19 @@ Generated by [AVA](https://avajs.dev). name: 'file1', }, ] + +> should find docs to embed + + 1 + +> should not find docs to embed + + 0 + +> should find docs to embed + + 1 + +> should not find docs to embed + + 0 diff --git a/packages/backend/server/src/__tests__/models/__snapshots__/copilot-workspace.spec.ts.snap b/packages/backend/server/src/__tests__/models/__snapshots__/copilot-workspace.spec.ts.snap index fae9d358b16b2cbe76aa3ed289e17b9b5afb816f..c1eb1e4f5a18cc64b9a069c106fb5a030b616c68 100644 GIT binary patch literal 606 zcmV-k0-^muRzVQ6WIg)*>2>qtIu3H}5p z+1#?U)E|oo00000000BcRXvZ>KoEVCZ{2~C)6sy0v_e8OC*UZkD4;?J5EK+pq4n-2 zUSjV?-p$2V&>$fpA<2NF6r-m&NH%zHaCdom0nE5`1_ zsq)%KDm`$aoko(y-0+cN!cKH#J)7BN>?T|&rbosmN-}AM)6VdT3$1sbPNn3YKX0Oo z?4cHjRKx=n9W8d`QdYq^;PvMc(WAxzfO`P;sk2AzM@hnuiHcRf-v{l02((M|pacm3 zx&U=v1n-pa)F_vO!2oo54(nr|8dsdO2Uu- zCgDq2WUL96-dc@TRhMbA)IE>>tfDV!I~ePk2)P+4riayym`b(t?!pS5vI>J*l`uZy z+H5Ya^jf95jGC>fu-_HJ&AR~Z19-TNbYo8XX%5^aO0qZ2`F{oQZ43V%HGTm2S$Aum zQ>Ly_v$xWU^C^0E%YC^R6;M!H9OTuq5n_9^GG>a=bYtP6jZ=wd+(_mWPejO694INZ z87&9W$DHgNdik$xyE6dq0h|E%4&WDnCIPo;s{);x0-YeT!BT>DW**j665-U0=T-}~ ze_dgXSH{=vnp;n5=eZGT{Vu&&$LiJbL=_))yh{C@NyN3OxwZ_COSkq=N2+Mro8?8} sMOe8XGN0vvoYA?(fe*D=%sX#$;O1F(y)r2_;r>&wKc^ZwB02{E0LgbE;s5{u literal 579 zcmV-J0=)e}RzVKoEVCZ{2}%r=x*tg@kBMz)?_9K!p$>C@7#p>)lDb z#MzC!n~SfYfC>qTo(kRl4>~G-2PHoOIf)&|3PcCZx$%xYZ)V=xncdStlsGl=AI>eB-7)vXMER>s=#04>P>Bvu{Vqu2XrL2W^${TOx#K$i9k7ruTAfGkSMgC9= zM4ZToljCJXsdW{+0KEECB6`#~1aKF?0d@AN{V+}C2~n}?_xqq75P^1y?w241z!ad) z^XM(BK#g*_x3>qnJcsphK#eQZ1hx06t(|(T3#+A_I$G2_Th^O$n5H%_d#&I{|5EU! zJTo>V%Wke`tEwxEUFp8a{HUK2YkKj&Y`Zf6?*N { const module = await createTestingModule(); t.context.user = module.get(UserModel); t.context.workspace = module.get(WorkspaceModel); + t.context.copilotContext = module.get(CopilotContextModel); t.context.copilotWorkspace = module.get(CopilotWorkspaceConfigModel); t.context.db = module.get(PrismaClient); + t.context.doc = module.get(DocModel); t.context.config = module.get(Config); t.context.module = module; }); @@ -136,6 +144,61 @@ test('should insert and search embedding', async t => { ); } } + + { + const docId = randomUUID(); + await t.context.doc.upsert({ + spaceId: workspace.id, + docId, + blob: Uint8Array.from([1, 2, 3]), + timestamp: Date.now(), + editorId: user.id, + }); + + const toBeEmbedDocIds = await t.context.copilotWorkspace.findDocsToEmbed( + workspace.id + ); + t.snapshot(toBeEmbedDocIds.length, 'should find docs to embed'); + + await t.context.copilotContext.insertWorkspaceEmbedding( + workspace.id, + docId, + [ + { + index: 0, + content: 'content', + embedding: Array.from({ length: 1024 }, () => 1), + }, + ] + ); + + const afterInsertEmbedding = + await t.context.copilotWorkspace.findDocsToEmbed(workspace.id); + t.snapshot(afterInsertEmbedding.length, 'should not find docs to embed'); + } + + { + const docId = randomUUID(); + await t.context.doc.upsert({ + spaceId: workspace.id, + docId, + blob: Uint8Array.from([1, 2, 3]), + timestamp: Date.now(), + editorId: user.id, + }); + + const toBeEmbedDocIds = await t.context.copilotWorkspace.findDocsToEmbed( + workspace.id + ); + t.snapshot(toBeEmbedDocIds.length, 'should find docs to embed'); + + await t.context.copilotWorkspace.updateIgnoredDocs(workspace.id, [docId]); + + const afterAddIgnoreDocs = await t.context.copilotWorkspace.findDocsToEmbed( + workspace.id + ); + t.snapshot(afterAddIgnoreDocs.length, 'should not find docs to embed'); + } }); test('should check embedding table', async t => { diff --git a/packages/backend/server/src/base/job/queue/config.ts b/packages/backend/server/src/base/job/queue/config.ts index 5aeb0f5de..5d294ff01 100644 --- a/packages/backend/server/src/base/job/queue/config.ts +++ b/packages/backend/server/src/base/job/queue/config.ts @@ -48,7 +48,7 @@ defineModuleConfig('job', { 'queues.copilot': { desc: 'The config for copilot job queue', default: { - concurrency: 1, + concurrency: 5, }, schema, }, diff --git a/packages/backend/server/src/models/copilot-context.ts b/packages/backend/server/src/models/copilot-context.ts index 6c868c61d..c2b030c68 100644 --- a/packages/backend/server/src/models/copilot-context.ts +++ b/packages/backend/server/src/models/copilot-context.ts @@ -119,6 +119,11 @@ export class CopilotContextModel extends BaseModel { } async hasWorkspaceEmbedding(workspaceId: string, docIds: string[]) { + const canEmbedding = await this.checkEmbeddingAvailable(); + if (!canEmbedding) { + return new Set(); + } + const existsIds = await this.db.aiWorkspaceEmbedding .findMany({ where: { @@ -238,10 +243,11 @@ export class CopilotContextModel extends BaseModel { WHERE w."workspace_id" = ${workspaceId} AND i."doc_id" IS NULL + AND (w."embedding" <=> ${embedding}::vector) <= ${threshold} ORDER BY "distance" ASC LIMIT ${topK}; `; - return similarityChunks.filter(c => Number(c.distance) <= threshold); + return similarityChunks; } } diff --git a/packages/backend/server/src/models/copilot-workspace.ts b/packages/backend/server/src/models/copilot-workspace.ts index eebeda344..c056ad8b3 100644 --- a/packages/backend/server/src/models/copilot-workspace.ts +++ b/packages/backend/server/src/models/copilot-workspace.ts @@ -34,6 +34,33 @@ export class CopilotWorkspaceConfigModel extends BaseModel { }); } + /** + * find docs to embed, excluding ignored and already embedded docs + * newer docs will be list first + * @param workspaceId id of the workspace + * @returns docIds + */ + @Transactional() + async findDocsToEmbed(workspaceId: string): Promise { + const docIds = await this.db.snapshot + .findMany({ + where: { + workspaceId, + embedding: { + is: null, + }, + }, + select: { id: true }, + }) + .then(r => r.map(doc => doc.id)); + + const skipDocIds = await this.listIgnoredDocIds(workspaceId).then( + r => new Set(r.map(r => r.docId)) + ); + + return docIds.filter(id => !skipDocIds.has(id)); + } + @Transactional() async updateIgnoredDocs( workspaceId: string, diff --git a/packages/backend/server/src/models/workspace.ts b/packages/backend/server/src/models/workspace.ts index 5952da16b..05bb8545d 100644 --- a/packages/backend/server/src/models/workspace.ts +++ b/packages/backend/server/src/models/workspace.ts @@ -7,6 +7,7 @@ import { BaseModel } from './base'; declare global { interface Events { + 'workspace.updated': Workspace; 'workspace.deleted': { id: string; }; @@ -58,6 +59,9 @@ export class WorkspaceModel extends BaseModel { this.logger.debug( `Updated workspace ${workspaceId} with data ${JSON.stringify(data)}` ); + + this.event.emit('workspace.updated', workspace); + return workspace; } diff --git a/packages/backend/server/src/plugins/copilot/context/job.ts b/packages/backend/server/src/plugins/copilot/context/job.ts index c272d2c26..65e69be96 100644 --- a/packages/backend/server/src/plugins/copilot/context/job.ts +++ b/packages/backend/server/src/plugins/copilot/context/job.ts @@ -4,6 +4,7 @@ import { AFFiNELogger, BlobNotFound, Config, + DocNotFound, EventBus, JobQueue, mapAnyError, @@ -87,11 +88,37 @@ export class CopilotContextDocJob { } } - // @OnEvent('doc.indexer.updated') - async addDocEmbeddingQueueFromEvent( - // TODO(@darkskygit): replace this with real event type - doc: { workspaceId: string; docId: string } //Events['doc.indexer.updated'], - ) { + @OnEvent('workspace.updated') + async onWorkspaceConfigUpdate({ + id, + enableDocEmbedding, + }: Events['workspace.updated']) { + if (enableDocEmbedding) { + // trigger workspace embedding + this.event.emit('workspace.embedding', { + workspaceId: id, + }); + } + } + + @OnEvent('workspace.embedding') + async addWorkspaceEmbeddingQueue({ + workspaceId, + }: Events['workspace.embedding']) { + if (!this.supportEmbedding) return; + + const toBeEmbedDocIds = + await this.models.copilotWorkspace.findDocsToEmbed(workspaceId); + for (const docId of toBeEmbedDocIds) { + await this.queue.add('copilot.embedding.docs', { + workspaceId, + docId, + }); + } + } + + @OnEvent('doc.indexer.updated') + async addDocEmbeddingQueueFromEvent(doc: Events['doc.indexer.updated']) { if (!this.supportEmbedding) return; await this.queue.add('copilot.embedding.docs', { @@ -100,11 +127,10 @@ export class CopilotContextDocJob { }); } - // @OnEvent('doc.indexer.deleted') - async deleteDocEmbeddingQueueFromEvent( - // TODO(@darkskygit): replace this with real event type - doc: { workspaceId: string; docId: string } //Events['doc.indexer.deleted'], - ) { + @OnEvent('doc.indexer.deleted') + async deleteDocEmbeddingQueueFromEvent(doc: Events['doc.indexer.deleted']) { + if (!this.supportEmbedding) return; + await this.models.copilotContext.deleteWorkspaceEmbedding( doc.workspaceId, doc.docId @@ -193,13 +219,14 @@ export class CopilotContextDocJob { docId, }: Jobs['copilot.embedding.docs']) { if (!this.supportEmbedding) return; + if (workspaceId === docId || docId.includes('$')) return; try { const content = await this.doc.getFullDocContent(workspaceId, docId); if (content) { // no need to check if embeddings is empty, will throw internally const embeddings = await this.embeddingClient.getFileEmbeddings( - new File([content.summary], `${content.title}.md`) + new File([content.summary], `${content.title || 'Untitled'}.md`) ); for (const chunks of embeddings) { @@ -209,6 +236,8 @@ export class CopilotContextDocJob { chunks ); } + } else if (contextId) { + throw new DocNotFound({ spaceId: workspaceId, docId }); } } catch (error: any) { if (contextId) { diff --git a/packages/backend/server/src/plugins/copilot/context/types.ts b/packages/backend/server/src/plugins/copilot/context/types.ts index 938732502..68524a8cb 100644 --- a/packages/backend/server/src/plugins/copilot/context/types.ts +++ b/packages/backend/server/src/plugins/copilot/context/types.ts @@ -8,6 +8,10 @@ import { parseDoc } from '../../../native'; declare global { interface Events { + 'workspace.embedding': { + workspaceId: string; + }; + 'workspace.doc.embedding': Array<{ workspaceId: string; docId: string; diff --git a/tests/affine-cloud-copilot/e2e/chat-with/collections.spec.ts b/tests/affine-cloud-copilot/e2e/chat-with/collections.spec.ts index 88d2763d1..6035d0a16 100644 --- a/tests/affine-cloud-copilot/e2e/chat-with/collections.spec.ts +++ b/tests/affine-cloud-copilot/e2e/chat-with/collections.spec.ts @@ -22,19 +22,23 @@ test.describe('AIChatWith/Collections', () => { loggedInPage: page, utils, }) => { + const randomStr = Math.random().toString(36).substring(2, 6); // Create two collections await utils.editor.createCollectionAndDoc( page, 'Collection 1', - 'CollectionAAaa is a cute dog' + `Collection${randomStr} is a cute dog` ); await utils.chatPanel.chatWithCollections(page, ['Collection 1']); - await utils.chatPanel.makeChat(page, 'What is CollectionAAaa(Use English)'); + await utils.chatPanel.makeChat( + page, + `What is Collection${randomStr}(Use English)` + ); await utils.chatPanel.waitForHistory(page, [ { role: 'user', - content: 'What is CollectionAAaa(Use English)', + content: `What is Collection${randomStr}(Use English)`, }, { role: 'assistant', @@ -45,7 +49,7 @@ test.describe('AIChatWith/Collections', () => { await expect(async () => { const { content, message } = await utils.chatPanel.getLatestAssistantMessage(page); - expect(content).toMatch(/CollectionAAaa.*dog/); + expect(content).toMatch(new RegExp(`Collection${randomStr}.*dog`)); expect(await message.locator('affine-footnote-node').count()).toBe(1); }).toPass(); }); @@ -54,17 +58,19 @@ test.describe('AIChatWith/Collections', () => { loggedInPage: page, utils, }) => { + const randomStr1 = Math.random().toString(36).substring(2, 6); + const randomStr2 = Math.random().toString(36).substring(2, 6); // Create two collections await utils.editor.createCollectionAndDoc( page, 'Collection 2', - 'CollectionEEee is a cute cat' + `Collection${randomStr1} is a cute cat` ); await utils.editor.createCollectionAndDoc( page, 'Collection 3', - 'CollectionFFff is a cute dog' + `Collection${randomStr2} is a cute dog` ); await utils.chatPanel.chatWithCollections(page, [ @@ -73,12 +79,12 @@ test.describe('AIChatWith/Collections', () => { ]); await utils.chatPanel.makeChat( page, - 'What is CollectionEEee? What is CollectionFFff?(Use English)' + `What is Collection${randomStr1}? What is Collection${randomStr2}?(Use English)` ); await utils.chatPanel.waitForHistory(page, [ { role: 'user', - content: 'What is CollectionEEee? What is CollectionFFff?(Use English)', + content: `What is Collection${randomStr1}? What is Collection${randomStr2}?(Use English)`, }, { role: 'assistant', @@ -89,8 +95,8 @@ test.describe('AIChatWith/Collections', () => { await expect(async () => { const { content, message } = await utils.chatPanel.getLatestAssistantMessage(page); - expect(content).toMatch(/CollectionEEee.*cat/); - expect(content).toMatch(/CollectionFFff.*dog/); + expect(content).toMatch(new RegExp(`Collection${randomStr1}.*cat`)); + expect(content).toMatch(new RegExp(`Collection${randomStr2}.*dog`)); expect(await message.locator('affine-footnote-node').count()).toBe(2); }).toPass(); }); diff --git a/tests/affine-cloud-copilot/e2e/chat-with/tags.spec.ts b/tests/affine-cloud-copilot/e2e/chat-with/tags.spec.ts index b72554ae2..3d256c197 100644 --- a/tests/affine-cloud-copilot/e2e/chat-with/tags.spec.ts +++ b/tests/affine-cloud-copilot/e2e/chat-with/tags.spec.ts @@ -20,13 +20,21 @@ test.describe('AIChatWith/tags', () => { loggedInPage: page, utils, }) => { - await utils.editor.createTagAndDoc(page, 'Tag 1', 'TagAAaa is a cute cat'); + const randomStr = Math.random().toString(36).substring(2, 6); + await utils.editor.createTagAndDoc( + page, + 'Tag 1', + `Tag${randomStr} is a cute cat` + ); await utils.chatPanel.chatWithTags(page, ['Tag 1']); - await utils.chatPanel.makeChat(page, 'What is TagAAaa(Use English)'); + await utils.chatPanel.makeChat( + page, + `What is Tag${randomStr}(Use English)` + ); await utils.chatPanel.waitForHistory(page, [ { role: 'user', - content: 'What is TagAAaa(Use English)', + content: `What is Tag${randomStr}(Use English)`, }, { role: 'assistant', @@ -36,7 +44,7 @@ test.describe('AIChatWith/tags', () => { await expect(async () => { const { content, message } = await utils.chatPanel.getLatestAssistantMessage(page); - expect(content).toMatch(/TagAAaa.*cat/); + expect(content).toMatch(new RegExp(`Tag${randomStr}.*cat`)); await expect(message.locator('affine-footnote-node')).toHaveCount(1); }).toPass(); }); @@ -45,17 +53,28 @@ test.describe('AIChatWith/tags', () => { loggedInPage: page, utils, }) => { - await utils.editor.createTagAndDoc(page, 'Tag 2', 'TagEEee is a cute cat'); - await utils.editor.createTagAndDoc(page, 'Tag 3', 'TagFFff is a cute dog'); + const randomStr1 = Math.random().toString(36).substring(2, 6); + const randomStr2 = Math.random().toString(36).substring(2, 6); + + await utils.editor.createTagAndDoc( + page, + 'Tag 2', + `Tag${randomStr1} is a cute cat` + ); + await utils.editor.createTagAndDoc( + page, + 'Tag 3', + `Tag${randomStr2} is a cute dog` + ); await utils.chatPanel.chatWithTags(page, ['Tag 2', 'Tag 3']); await utils.chatPanel.makeChat( page, - 'What is TagEEee? What is TagFFff?(Use English)' + `What is Tag${randomStr1}? What is Tag${randomStr2}?(Use English)` ); await utils.chatPanel.waitForHistory(page, [ { role: 'user', - content: 'What is TagEEee? What is TagFFff?(Use English)', + content: `What is Tag${randomStr1}? What is Tag${randomStr2}?(Use English)`, }, { role: 'assistant', @@ -65,8 +84,8 @@ test.describe('AIChatWith/tags', () => { await expect(async () => { const { content, message } = await utils.chatPanel.getLatestAssistantMessage(page); - expect(content).toMatch(/TagEEee.*cat/); - expect(content).toMatch(/TagFFff.*dog/); + expect(content).toMatch(new RegExp(`Tag${randomStr1}.*cat`)); + expect(content).toMatch(new RegExp(`Tag${randomStr2}.*dog`)); await expect(message.locator('affine-footnote-node')).toHaveCount(2); }).toPass(); }); diff --git a/tests/affine-cloud-copilot/e2e/utils/editor-utils.ts b/tests/affine-cloud-copilot/e2e/utils/editor-utils.ts index a57b0dc2b..bb661bcc0 100644 --- a/tests/affine-cloud-copilot/e2e/utils/editor-utils.ts +++ b/tests/affine-cloud-copilot/e2e/utils/editor-utils.ts @@ -334,6 +334,8 @@ export class EditorUtils { await page.keyboard.press('Enter'); } } + // sleep 1 sec to wait the doc sync + await page.waitForTimeout(1000); } public static async createTagAndDoc( @@ -362,6 +364,8 @@ export class EditorUtils { await page.keyboard.press('Enter'); } } + // sleep 1 sec to wait the doc sync + await page.waitForTimeout(1000); } public static async selectElementInEdgeless(page: Page, elements: string[]) {