feat(server): global embedding gql endpoint (#11809)

fix AI-30
fix AI-31
fix PD-2487

packages/backend/server/src/plugins/copilot/context/job.ts

@@ -13,9 +13,9 @@ import {
 import { DocReader } from '../../../core/doc';
 import { Models } from '../../../models';
 import { CopilotStorage } from '../storage';
+import { readStream } from '../utils';
 import { OpenAIEmbeddingClient } from './embedding';
 import { EmbeddingClient } from './types';
-import { readStream } from './utils';
 
 @Injectable()
 export class CopilotContextDocJob {

packages/backend/server/src/plugins/copilot/context/resolver.ts

@@ -47,10 +47,10 @@ import {
 import { COPILOT_LOCKER, CopilotType } from '../resolver';
 import { ChatSessionService } from '../session';
 import { CopilotStorage } from '../storage';
+import { MAX_EMBEDDABLE_SIZE } from '../types';
+import { readStream } from '../utils';
 import { CopilotContextDocJob } from './job';
 import { CopilotContextService } from './service';
-import { MAX_EMBEDDABLE_SIZE } from './types';
-import { readStream } from './utils';
 
 @InputType()
 class AddContextCategoryInput {

packages/backend/server/src/plugins/copilot/context/types.ts

@@ -1,6 +1,6 @@
 import { File } from 'node:buffer';
 
-import { CopilotContextFileNotSupported, OneMB } from '../../../base';
+import { CopilotContextFileNotSupported } from '../../../base';
 import { Embedding } from '../../../models';
 import { parseDoc } from '../../../native';
 
@@ -46,8 +46,6 @@ declare global {
   }
 }
 
-export const MAX_EMBEDDABLE_SIZE = 50 * OneMB;
-
 export type Chunk = {
   index: number;
   content: string;

packages/backend/server/src/plugins/copilot/context/utils.ts

@@ -1,8 +1,3 @@
-import { Readable } from 'node:stream';
-
-import { readBufferWithLimit } from '../../../base';
-import { MAX_EMBEDDABLE_SIZE } from './types';
-
 export class GqlSignal implements AsyncDisposable {
   readonly abortController = new AbortController();
 
@@ -14,10 +9,3 @@ export class GqlSignal implements AsyncDisposable {
     this.abortController.abort();
   }
 }
-
-export function readStream(
-  readable: Readable,
-  maxSize = MAX_EMBEDDABLE_SIZE
-): Promise<Buffer> {
-  return readBufferWithLimit(readable, maxSize);
-}

packages/backend/server/src/plugins/copilot/index.ts

@@ -7,6 +7,7 @@ import { DocStorageModule } from '../../core/doc';
 import { FeatureModule } from '../../core/features';
 import { PermissionModule } from '../../core/permission';
 import { QuotaModule } from '../../core/quota';
+import { WorkspaceModule } from '../../core/workspaces';
 import {
   CopilotContextDocJob,
   CopilotContextResolver,
@@ -29,6 +30,11 @@ import {
   CopilotTranscriptionService,
 } from './transcript';
 import { CopilotWorkflowExecutors, CopilotWorkflowService } from './workflow';
+import {
+  CopilotWorkspaceEmbeddingConfigResolver,
+  CopilotWorkspaceEmbeddingResolver,
+  CopilotWorkspaceService,
+} from './workspace';
 
 @Module({
   imports: [
@@ -37,6 +43,7 @@ import { CopilotWorkflowExecutors, CopilotWorkflowService } from './workflow';
     QuotaModule,
     PermissionModule,
     ServerConfigModule,
+    WorkspaceModule,
   ],
   providers: [
     // providers
@@ -58,6 +65,10 @@ import { CopilotWorkflowExecutors, CopilotWorkflowService } from './workflow';
     // transcription
     CopilotTranscriptionService,
     CopilotTranscriptionResolver,
+    // workspace embeddings
+    CopilotWorkspaceService,
+    CopilotWorkspaceEmbeddingResolver,
+    CopilotWorkspaceEmbeddingConfigResolver,
     // gql resolvers
     UserCopilotResolver,
     PromptsManagementResolver,

packages/backend/server/src/plugins/copilot/types.ts

@@ -1,6 +1,7 @@
 import { type Tokenizer } from '@affine/server-native';
 import { z } from 'zod';
 
+import { OneMB } from '../../base';
 import { fromModelName } from '../../native';
 import type { ChatPrompt } from './prompt';
 import { PromptMessageSchema, PureMessageSchema } from './providers';
@@ -116,3 +117,5 @@ export type CopilotContextFile = {
   // embedding status
   status: 'in_progress' | 'completed' | 'failed';
 };
+
+export const MAX_EMBEDDABLE_SIZE = 50 * OneMB;

packages/backend/server/src/plugins/copilot/utils.ts

@@ -0,0 +1,11 @@
+import { Readable } from 'node:stream';
+
+import { readBufferWithLimit } from '../../base';
+import { MAX_EMBEDDABLE_SIZE } from './types';
+
+export function readStream(
+  readable: Readable,
+  maxSize = MAX_EMBEDDABLE_SIZE
+): Promise<Buffer> {
+  return readBufferWithLimit(readable, maxSize);
+}

packages/backend/server/src/plugins/copilot/workspace/index.ts

@@ -0,0 +1,5 @@
+export {
+  CopilotWorkspaceEmbeddingConfigResolver,
+  CopilotWorkspaceEmbeddingResolver,
+} from './resolver';
+export { CopilotWorkspaceService } from './service';

packages/backend/server/src/plugins/copilot/workspace/resolver.ts

@@ -0,0 +1,218 @@
+import {
+  Args,
+  Context,
+  Field,
+  Mutation,
+  ObjectType,
+  Parent,
+  ResolveField,
+  Resolver,
+} from '@nestjs/graphql';
+import type { Request } from 'express';
+import { SafeIntResolver } from 'graphql-scalars';
+import GraphQLUpload, {
+  type FileUpload,
+} from 'graphql-upload/GraphQLUpload.mjs';
+
+import {
+  BlobQuotaExceeded,
+  CopilotEmbeddingUnavailable,
+  CopilotFailedToAddWorkspaceFileEmbedding,
+  Mutex,
+  TooManyRequest,
+  UserFriendlyError,
+} from '../../../base';
+import { CurrentUser } from '../../../core/auth';
+import { AccessController } from '../../../core/permission';
+import { WorkspaceType } from '../../../core/workspaces';
+import { CopilotWorkspaceFile, Models } from '../../../models';
+import { COPILOT_LOCKER } from '../resolver';
+import { MAX_EMBEDDABLE_SIZE } from '../types';
+import { CopilotWorkspaceService } from './service';
+
+@ObjectType('CopilotWorkspaceConfig')
+export class CopilotWorkspaceConfigType {
+  @Field(() => String)
+  workspaceId!: string;
+}
+
+@ObjectType('CopilotWorkspaceFile')
+export class CopilotWorkspaceFileType implements CopilotWorkspaceFile {
+  @Field(() => String)
+  workspaceId!: string;
+
+  @Field(() => String)
+  fileId!: string;
+
+  @Field(() => String)
+  fileName!: string;
+
+  @Field(() => String)
+  mimeType!: string;
+
+  @Field(() => SafeIntResolver)
+  size!: number;
+
+  @Field(() => Date)
+  createdAt!: Date;
+}
+
+/**
+ * Workspace embedding config resolver
+ * Public apis rate limit: 10 req/m
+ * Other rate limit: 120 req/m
+ */
+@Resolver(() => WorkspaceType)
+export class CopilotWorkspaceEmbeddingResolver {
+  constructor(private readonly ac: AccessController) {}
+
+  @ResolveField(() => CopilotWorkspaceConfigType, {
+    complexity: 2,
+  })
+  async embedding(
+    @CurrentUser() user: CurrentUser,
+    @Parent() workspace: WorkspaceType
+  ): Promise<CopilotWorkspaceConfigType> {
+    await this.ac
+      .user(user.id)
+      .workspace(workspace.id)
+      .assert('Workspace.Read');
+
+    return { workspaceId: workspace.id };
+  }
+}
+
+@Resolver(() => CopilotWorkspaceConfigType)
+export class CopilotWorkspaceEmbeddingConfigResolver {
+  constructor(
+    private readonly ac: AccessController,
+    private readonly models: Models,
+    private readonly mutex: Mutex,
+    private readonly copilotWorkspace: CopilotWorkspaceService
+  ) {}
+
+  @ResolveField(() => [String], {
+    complexity: 2,
+  })
+  async ignoredDocs(
+    @Parent() config: CopilotWorkspaceConfigType
+  ): Promise<string[]> {
+    return this.models.copilotWorkspace.listIgnoredDocs(config.workspaceId);
+  }
+
+  @Mutation(() => Number, {
+    name: 'updateWorkspaceEmbeddingIgnoredDocs',
+    complexity: 2,
+    description: 'Update ignored docs',
+  })
+  async updateIgnoredDocs(
+    @CurrentUser() user: CurrentUser,
+    @Args('workspaceId', { type: () => String })
+    workspaceId: string,
+    @Args('add', { type: () => [String], nullable: true })
+    add?: string[],
+    @Args('remove', { type: () => [String], nullable: true })
+    remove?: string[]
+  ): Promise<number> {
+    await this.ac
+      .user(user.id)
+      .workspace(workspaceId)
+      .assert('Workspace.Settings.Update');
+    return await this.models.copilotWorkspace.updateIgnoredDocs(
+      workspaceId,
+      add,
+      remove
+    );
+  }
+
+  @ResolveField(() => [CopilotWorkspaceFileType], {
+    complexity: 2,
+  })
+  async files(
+    @Parent() config: CopilotWorkspaceConfigType
+  ): Promise<CopilotWorkspaceFileType[]> {
+    return this.models.copilotWorkspace.listWorkspaceFiles(config.workspaceId);
+  }
+
+  @Mutation(() => CopilotWorkspaceFileType, {
+    name: 'addWorkspaceEmbeddingFiles',
+    complexity: 2,
+    description: 'Update workspace embedding files',
+  })
+  async addFiles(
+    @Context() ctx: { req: Request },
+    @CurrentUser() user: CurrentUser,
+    @Args('workspaceId', { type: () => String })
+    workspaceId: string,
+    @Args({ name: 'blob', type: () => GraphQLUpload })
+    content: FileUpload
+  ): Promise<CopilotWorkspaceFileType> {
+    await this.ac
+      .user(user.id)
+      .workspace(workspaceId)
+      .assert('Workspace.Settings.Update');
+
+    if (!this.copilotWorkspace.canEmbedding) {
+      throw new CopilotEmbeddingUnavailable();
+    }
+
+    const lockFlag = `${COPILOT_LOCKER}:workspace:${workspaceId}`;
+    await using lock = await this.mutex.acquire(lockFlag);
+    if (!lock) {
+      throw new TooManyRequest('Server is busy');
+    }
+
+    const length = Number(ctx.req.headers['content-length']);
+    if (length && length >= MAX_EMBEDDABLE_SIZE) {
+      throw new BlobQuotaExceeded();
+    }
+
+    try {
+      const { blobId, file } = await this.copilotWorkspace.addWorkspaceFile(
+        user.id,
+        workspaceId,
+        content
+      );
+      await this.copilotWorkspace.addWorkspaceFileEmbeddingQueue({
+        userId: user.id,
+        workspaceId,
+        blobId,
+        fileId: file.fileId,
+        fileName: file.fileName,
+      });
+
+      return file;
+    } catch (e: any) {
+      // passthrough user friendly error
+      if (e instanceof UserFriendlyError) {
+        throw e;
+      }
+      throw new CopilotFailedToAddWorkspaceFileEmbedding({
+        message: e.message,
+      });
+    }
+  }
+
+  @Mutation(() => Boolean, {
+    name: 'removeWorkspaceEmbeddingFiles',
+    complexity: 2,
+    description: 'Remove workspace embedding files',
+  })
+  async removeFiles(
+    @CurrentUser() user: CurrentUser,
+    @Args('workspaceId', { type: () => String })
+    workspaceId: string,
+    @Args('fileId', { type: () => String })
+    fileId: string
+  ): Promise<boolean> {
+    await this.ac
+      .user(user.id)
+      .workspace(workspaceId)
+      .assert('Workspace.Settings.Update');
+
+    return await this.models.copilotWorkspace.removeWorkspaceFile(
+      workspaceId,
+      fileId
+    );
+  }
+}

packages/backend/server/src/plugins/copilot/workspace/service.ts

@@ -0,0 +1,87 @@
+import { createHash } from 'node:crypto';
+
+import { Injectable, OnApplicationBootstrap } from '@nestjs/common';
+
+import { FileUpload, JobQueue } from '../../../base';
+import { Models } from '../../../models';
+import { CopilotStorage } from '../storage';
+import { readStream } from '../utils';
+
+declare global {
+  interface Events {
+    'workspace.file.embedding.finished': {
+      jobId: string;
+    };
+    'workspace.file.embedding.failed': {
+      jobId: string;
+    };
+  }
+  interface Jobs {
+    'copilot.workspace.embedding.files': {
+      userId: string;
+      workspaceId: string;
+      blobId: string;
+      fileId: string;
+      fileName: string;
+    };
+  }
+}
+
+@Injectable()
+export class CopilotWorkspaceService implements OnApplicationBootstrap {
+  private supportEmbedding = false;
+
+  constructor(
+    private readonly models: Models,
+    private readonly queue: JobQueue,
+    private readonly storage: CopilotStorage
+  ) {}
+
+  async onApplicationBootstrap() {
+    const supportEmbedding =
+      await this.models.copilotContext.checkEmbeddingAvailable();
+    if (supportEmbedding) {
+      this.supportEmbedding = true;
+    }
+  }
+
+  get canEmbedding() {
+    return this.supportEmbedding;
+  }
+
+  async addWorkspaceFile(
+    userId: string,
+    workspaceId: string,
+    content: FileUpload
+  ) {
+    const fileName = content.filename;
+    const buffer = await readStream(content.createReadStream());
+    const blobId = createHash('sha256').update(buffer).digest('base64url');
+    await this.storage.put(userId, workspaceId, blobId, buffer);
+    const file = await this.models.copilotWorkspace.addFile(workspaceId, {
+      fileName,
+      mimeType: content.mimetype,
+      size: buffer.length,
+    });
+    return { blobId, file };
+  }
+
+  async getWorkspaceFile(workspaceId: string, fileId: string) {
+    return await this.models.copilotWorkspace.getFile(workspaceId, fileId);
+  }
+
+  async addWorkspaceFileEmbeddingQueue(
+    file: Jobs['copilot.workspace.embedding.files']
+  ) {
+    if (!this.supportEmbedding) return;
+
+    const { userId, workspaceId, blobId, fileId, fileName } = file;
+    await this.queue.add('copilot.workspace.embedding.files', {
+      userId,
+      workspaceId,
+      blobId,
+      fileId,
+      fileName,
+    });
+  }
+}