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

fix AI-30 fix AI-31 fix PD-2487
2026-07-13 08:06:24 +08:00 · 2025-04-23 11:25:40 +00:00
parent 5d9a3aac5b
commit 5397fba897
25 changed files with 665 additions and 60 deletions
@@ -0,0 +1,5 @@
+export {
+  CopilotWorkspaceEmbeddingConfigResolver,
+  CopilotWorkspaceEmbeddingResolver,
+} from './resolver';
+export { CopilotWorkspaceService } from './service';
@@ -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
+    );
+  }
+}
@@ -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,
+    });
+  }
+}