feat(server): basic context api (#10056)

fix CLOUD-97
fix CLOUD-98

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

@@ -0,0 +1,260 @@
+import {
+  Args,
+  Field,
+  ID,
+  InputType,
+  Mutation,
+  ObjectType,
+  Parent,
+  registerEnumType,
+  ResolveField,
+  Resolver,
+} from '@nestjs/graphql';
+import { SafeIntResolver } from 'graphql-scalars';
+
+import {
+  CallMetric,
+  CopilotFailedToModifyContext,
+  CopilotSessionNotFound,
+  RequestMutex,
+  Throttle,
+  TooManyRequest,
+} from '../../../base';
+import { CurrentUser } from '../../../core/auth';
+import { COPILOT_LOCKER, CopilotType } from '../resolver';
+import { ChatSessionService } from '../session';
+import { CopilotContextService } from './service';
+import { ContextDoc, type ContextFile, ContextFileStatus } from './types';
+
+@InputType()
+class AddContextDocInput {
+  @Field(() => String)
+  contextId!: string;
+
+  @Field(() => String)
+  docId!: string;
+}
+
+@InputType()
+class RemoveContextFileInput {
+  @Field(() => String)
+  contextId!: string;
+
+  @Field(() => String)
+  fileId!: string;
+}
+
+@ObjectType('CopilotContext')
+export class CopilotContextType {
+  @Field(() => ID)
+  id!: string;
+
+  @Field(() => String)
+  workspaceId!: string;
+}
+
+registerEnumType(ContextFileStatus, { name: 'ContextFileStatus' });
+
+@ObjectType()
+class CopilotContextDoc implements ContextDoc {
+  @Field(() => ID)
+  id!: string;
+
+  @Field(() => SafeIntResolver)
+  createdAt!: number;
+}
+
+@ObjectType()
+class CopilotContextFile implements ContextFile {
+  @Field(() => ID)
+  id!: string;
+
+  @Field(() => String)
+  name!: string;
+
+  @Field(() => SafeIntResolver)
+  chunkSize!: number;
+
+  @Field(() => ContextFileStatus)
+  status!: ContextFileStatus;
+
+  @Field(() => String)
+  blobId!: string;
+
+  @Field(() => SafeIntResolver)
+  createdAt!: number;
+}
+
+@ObjectType()
+class CopilotContextListItem {
+  @Field(() => ID)
+  id!: string;
+
+  @Field(() => SafeIntResolver)
+  createdAt!: number;
+
+  @Field(() => String, { nullable: true })
+  name!: string;
+
+  @Field(() => SafeIntResolver, { nullable: true })
+  chunkSize!: number;
+
+  @Field(() => ContextFileStatus, { nullable: true })
+  status!: ContextFileStatus;
+
+  @Field(() => String, { nullable: true })
+  blobId!: string;
+}
+
+@Throttle()
+@Resolver(() => CopilotType)
+export class CopilotContextRootResolver {
+  constructor(
+    private readonly mutex: RequestMutex,
+    private readonly chatSession: ChatSessionService,
+    private readonly context: CopilotContextService
+  ) {}
+
+  private async checkChatSession(
+    user: CurrentUser,
+    sessionId: string,
+    workspaceId?: string
+  ): Promise<void> {
+    const session = await this.chatSession.get(sessionId);
+    if (
+      !session ||
+      session.config.workspaceId !== workspaceId ||
+      session.config.userId !== user.id
+    ) {
+      throw new CopilotSessionNotFound();
+    }
+  }
+
+  @ResolveField(() => [CopilotContextType], {
+    description: 'Get the context list of a session',
+    complexity: 2,
+  })
+  @CallMetric('ai', 'context_create')
+  async contexts(
+    @Parent() copilot: CopilotType,
+    @CurrentUser() user: CurrentUser,
+    @Args('sessionId') sessionId: string,
+    @Args('contextId', { nullable: true }) contextId?: string
+  ) {
+    const lockFlag = `${COPILOT_LOCKER}:context:${sessionId}`;
+    await using lock = await this.mutex.acquire(lockFlag);
+    if (!lock) {
+      return new TooManyRequest('Server is busy');
+    }
+    await this.checkChatSession(user, sessionId, copilot.workspaceId);
+
+    if (contextId) {
+      const context = await this.context.get(contextId);
+      if (context) return [context];
+    } else {
+      const context = await this.context.getBySessionId(sessionId);
+      if (context) return [context];
+    }
+    return [];
+  }
+
+  @Mutation(() => String, {
+    description: 'Create a context session',
+  })
+  @CallMetric('ai', 'context_create')
+  async createCopilotContext(
+    @CurrentUser() user: CurrentUser,
+    @Args('workspaceId') workspaceId: string,
+    @Args('sessionId') sessionId: string
+  ) {
+    const lockFlag = `${COPILOT_LOCKER}:context:${sessionId}`;
+    await using lock = await this.mutex.acquire(lockFlag);
+    if (!lock) {
+      return new TooManyRequest('Server is busy');
+    }
+    await this.checkChatSession(user, sessionId, workspaceId);
+
+    const context = await this.context.create(sessionId);
+    return context.id;
+  }
+}
+
+@Throttle()
+@Resolver(() => CopilotContextType)
+export class CopilotContextResolver {
+  constructor(
+    private readonly mutex: RequestMutex,
+
+    private readonly context: CopilotContextService
+  ) {}
+
+  @ResolveField(() => [CopilotContextDoc], {
+    description: 'list files in context',
+  })
+  @CallMetric('ai', 'context_file_list')
+  async docs(@Parent() context: CopilotContextType): Promise<ContextDoc[]> {
+    const session = await this.context.get(context.id);
+    return session.listDocs();
+  }
+
+  @Mutation(() => [CopilotContextListItem], {
+    description: 'add a doc to context',
+  })
+  @CallMetric('ai', 'context_doc_add')
+  async addContextDoc(
+    @Args({ name: 'options', type: () => AddContextDocInput })
+    options: AddContextDocInput
+  ) {
+    const lockFlag = `${COPILOT_LOCKER}:context:${options.contextId}`;
+    await using lock = await this.mutex.acquire(lockFlag);
+    if (!lock) {
+      return new TooManyRequest('Server is busy');
+    }
+    const session = await this.context.get(options.contextId);
+
+    try {
+      return await session.addDocRecord(options.docId);
+    } catch (e: any) {
+      throw new CopilotFailedToModifyContext({
+        contextId: options.contextId,
+        message: e.message,
+      });
+    }
+  }
+
+  @Mutation(() => Boolean, {
+    description: 'remove a doc from context',
+  })
+  @CallMetric('ai', 'context_doc_remove')
+  async removeContextDoc(
+    @Args({ name: 'options', type: () => RemoveContextFileInput })
+    options: RemoveContextFileInput
+  ) {
+    const lockFlag = `${COPILOT_LOCKER}:context:${options.contextId}`;
+    await using lock = await this.mutex.acquire(lockFlag);
+    if (!lock) {
+      return new TooManyRequest('Server is busy');
+    }
+    const session = await this.context.get(options.contextId);
+
+    try {
+      return await session.removeDocRecord(options.fileId);
+    } catch (e: any) {
+      throw new CopilotFailedToModifyContext({
+        contextId: options.contextId,
+        message: e.message,
+      });
+    }
+  }
+
+  @ResolveField(() => [CopilotContextFile], {
+    description: 'list files in context',
+  })
+  @CallMetric('ai', 'context_file_list')
+  async files(
+    @Parent() context: CopilotContextType
+  ): Promise<CopilotContextFile[]> {
+    const session = await this.context.get(context.id);
+    return session.listFiles();
+  }
+}