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 fae9d358b..c1eb1e4f5 100644 Binary files a/packages/backend/server/src/__tests__/models/__snapshots__/copilot-workspace.spec.ts.snap and b/packages/backend/server/src/__tests__/models/__snapshots__/copilot-workspace.spec.ts.snap differ diff --git a/packages/backend/server/src/__tests__/models/copilot-workspace.spec.ts b/packages/backend/server/src/__tests__/models/copilot-workspace.spec.ts index 5d6087836..c99ac6366 100644 --- a/packages/backend/server/src/__tests__/models/copilot-workspace.spec.ts +++ b/packages/backend/server/src/__tests__/models/copilot-workspace.spec.ts @@ -1,8 +1,12 @@ +import { randomUUID } from 'node:crypto'; + import { PrismaClient, User, Workspace } from '@prisma/client'; import ava, { TestFn } from 'ava'; import { Config } from '../../base'; +import { CopilotContextModel } from '../../models/copilot-context'; import { CopilotWorkspaceConfigModel } from '../../models/copilot-workspace'; +import { DocModel } from '../../models/doc'; import { UserModel } from '../../models/user'; import { WorkspaceModel } from '../../models/workspace'; import { createTestingModule, type TestingModule } from '../utils'; @@ -12,8 +16,10 @@ interface Context { config: Config; module: TestingModule; db: PrismaClient; + doc: DocModel; user: UserModel; workspace: WorkspaceModel; + copilotContext: CopilotContextModel; copilotWorkspace: CopilotWorkspaceConfigModel; } @@ -23,8 +29,10 @@ test.before(async t => { 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[]) {