feat(server): trigger workspace embedding (#12328)

fix AI-127

<!-- This is an auto-generated comment: release notes by coderabbit.ai -->
## 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.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

.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": {

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

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 => {

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,
   },

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;
   }
 }

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<string[]> {
+    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,

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;
   }
 

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) {

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;

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();
   });

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();
   });

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[]) {