feat: copilot controller (#6272)

fix CLOUD-27

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

@@ -0,0 +1,151 @@
+import {
+  BadRequestException,
+  Controller,
+  Get,
+  InternalServerErrorException,
+  Param,
+  Query,
+  Req,
+  Sse,
+} from '@nestjs/common';
+import {
+  concatMap,
+  connect,
+  EMPTY,
+  from,
+  map,
+  merge,
+  Observable,
+  switchMap,
+  toArray,
+} from 'rxjs';
+
+import { Public } from '../../core/auth';
+import { CurrentUser } from '../../core/auth/current-user';
+import { CopilotProviderService } from './providers';
+import { ChatSessionService } from './session';
+import { CopilotCapability } from './types';
+
+export interface ChatEvent {
+  data: string;
+  id?: string;
+}
+
+@Controller('/api/copilot')
+export class CopilotController {
+  constructor(
+    private readonly chatSession: ChatSessionService,
+    private readonly provider: CopilotProviderService
+  ) {}
+
+  @Public()
+  @Get('/chat/:sessionId')
+  async chat(
+    @CurrentUser() user: CurrentUser,
+    @Req() req: Request,
+    @Param('sessionId') sessionId: string,
+    @Query('message') content: string,
+    @Query() params: Record<string, string | string[]>
+  ): Promise<string> {
+    const provider = this.provider.getProviderByCapability(
+      CopilotCapability.TextToText
+    );
+    if (!provider) {
+      throw new InternalServerErrorException('No provider available');
+    }
+    const session = await this.chatSession.get(sessionId);
+    if (!session) {
+      throw new BadRequestException('Session not found');
+    }
+    if (!content || !content.trim()) {
+      throw new BadRequestException('Message is empty');
+    }
+    session.push({
+      role: 'user',
+      content: decodeURIComponent(content),
+      createdAt: new Date(),
+    });
+
+    try {
+      delete params.message;
+      const content = await provider.generateText(
+        session.finish(params),
+        session.model,
+        {
+          signal: req.signal,
+          user: user.id,
+        }
+      );
+
+      session.push({
+        role: 'assistant',
+        content,
+        createdAt: new Date(),
+      });
+      await session.save();
+
+      return content;
+    } catch (e: any) {
+      throw new InternalServerErrorException(
+        e.message || "Couldn't generate text"
+      );
+    }
+  }
+
+  @Public()
+  @Sse('/chat/:sessionId/stream')
+  async chatStream(
+    @CurrentUser() user: CurrentUser,
+    @Req() req: Request,
+    @Param('sessionId') sessionId: string,
+    @Query('message') content: string,
+    @Query() params: Record<string, string>
+  ): Promise<Observable<ChatEvent>> {
+    const provider = this.provider.getProviderByCapability(
+      CopilotCapability.TextToText
+    );
+    if (!provider) {
+      throw new InternalServerErrorException('No provider available');
+    }
+    const session = await this.chatSession.get(sessionId);
+    if (!session) {
+      throw new BadRequestException('Session not found');
+    }
+    if (!content || !content.trim()) {
+      throw new BadRequestException('Message is empty');
+    }
+    session.push({
+      role: 'user',
+      content: decodeURIComponent(content),
+      createdAt: new Date(),
+    });
+
+    delete params.message;
+    return from(
+      provider.generateTextStream(session.finish(params), session.model, {
+        signal: req.signal,
+        user: user.id,
+      })
+    ).pipe(
+      connect(shared$ =>
+        merge(
+          // actual chat event stream
+          shared$.pipe(map(data => ({ id: sessionId, data }))),
+          // save the generated text to the session
+          shared$.pipe(
+            toArray(),
+            concatMap(values => {
+              session.push({
+                role: 'assistant',
+                content: values.join(''),
+                createdAt: new Date(),
+              });
+              return from(session.save());
+            }),
+            switchMap(() => EMPTY)
+          )
+        )
+      )
+    );
+  }
+}

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

@@ -1,6 +1,8 @@
 import { ServerFeature } from '../../core/config';
+import { QuotaService } from '../../core/quota';
 import { PermissionService } from '../../core/workspaces/permission';
 import { Plugin } from '../registry';
+import { CopilotController } from './controller';
 import { PromptService } from './prompt';
 import {
   assertProvidersConfigs,
@@ -8,6 +10,7 @@ import {
   OpenAIProvider,
   registerCopilotProvider,
 } from './providers';
+import { CopilotResolver, UserCopilotResolver } from './resolver';
 import { ChatSessionService } from './session';
 
 registerCopilotProvider(OpenAIProvider);
@@ -16,10 +19,14 @@ registerCopilotProvider(OpenAIProvider);
   name: 'copilot',
   providers: [
     PermissionService,
+    QuotaService,
     ChatSessionService,
+    CopilotResolver,
+    UserCopilotResolver,
     PromptService,
     CopilotProviderService,
   ],
+  controllers: [CopilotController],
   contributesTo: ServerFeature.Copilot,
   if: config => {
     if (config.flavor.graphql) {

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

@@ -38,16 +38,16 @@ export class ChatPrompt {
   ) {
     return new ChatPrompt(
       options.name,
-      options.action,
-      options.model,
+      options.action || undefined,
+      options.model || undefined,
       options.messages
     );
   }
 
   constructor(
     public readonly name: string,
-    public readonly action: string | null,
-    public readonly model: string | null,
+    public readonly action: string | undefined,
+    public readonly model: string | undefined,
     private readonly messages: PromptMessage[]
   ) {
     this.encoder = getTokenEncoder(model);

packages/backend/server/src/plugins/copilot/providers/openai.ts

@@ -3,14 +3,16 @@ import assert from 'node:assert';
 import { ClientOptions, OpenAI } from 'openai';
 
 import {
-  ChatMessage,
   ChatMessageRole,
   CopilotCapability,
   CopilotProviderType,
   CopilotTextToEmbeddingProvider,
   CopilotTextToTextProvider,
+  PromptMessage,
 } from '../types';
 
+const DEFAULT_DIMENSIONS = 256;
+
 export class OpenAIProvider
   implements CopilotTextToTextProvider, CopilotTextToEmbeddingProvider
 {
@@ -50,7 +52,7 @@ export class OpenAIProvider
     return OpenAIProvider.capabilities;
   }
 
-  private chatToGPTMessage(messages: ChatMessage[]) {
+  private chatToGPTMessage(messages: PromptMessage[]) {
     // filter redundant fields
     return messages.map(message => ({
       role: message.role,
@@ -63,7 +65,7 @@ export class OpenAIProvider
     embeddings,
     model,
   }: {
-    messages?: ChatMessage[];
+    messages?: PromptMessage[];
     embeddings?: string[];
     model: string;
   }) {
@@ -106,7 +108,7 @@ export class OpenAIProvider
   // ====== text to text ======
 
   async generateText(
-    messages: ChatMessage[],
+    messages: PromptMessage[],
     model: string = 'gpt-3.5-turbo',
     options: {
       temperature?: number;
@@ -134,8 +136,8 @@ export class OpenAIProvider
   }
 
   async *generateTextStream(
-    messages: ChatMessage[],
-    model: string,
+    messages: PromptMessage[],
+    model: string = 'gpt-3.5-turbo',
     options: {
       temperature?: number;
       maxTokens?: number;
@@ -179,7 +181,7 @@ export class OpenAIProvider
       dimensions: number;
       signal?: AbortSignal;
       user?: string;
-    } = { dimensions: 256 }
+    } = { dimensions: DEFAULT_DIMENSIONS }
   ): Promise<number[][]> {
     messages = Array.isArray(messages) ? messages : [messages];
     this.checkParams({ embeddings: messages, model });
@@ -187,7 +189,7 @@ export class OpenAIProvider
     const result = await this.instance.embeddings.create({
       model: model,
       input: messages,
-      dimensions: options.dimensions,
+      dimensions: options.dimensions || DEFAULT_DIMENSIONS,
       user: options.user,
     });
     return result.data.map(e => e.embedding);

packages/backend/server/src/plugins/copilot/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 { CurrentUser, Public } from '../../core/auth';
+import { QuotaService } from '../../core/quota';
+import { UserType } from '../../core/user';
+import { PermissionService } from '../../core/workspaces/permission';
+import {
+  MutexService,
+  PaymentRequiredException,
+  TooManyRequestsException,
+} from '../../fundamentals';
+import { ChatSessionService, ListHistoriesOptions } from './session';
+import { AvailableModels, type ChatHistory, type ChatMessage } from './types';
+
+registerEnumType(AvailableModels, { name: 'CopilotModel' });
+
+// ================== Input Types ==================
+
+@InputType()
+class CreateChatSessionInput {
+  @Field(() => String)
+  workspaceId!: string;
+
+  @Field(() => String)
+  docId!: string;
+
+  @Field(() => String, {
+    description: 'An mark identifying which view to use to display the session',
+    nullable: true,
+  })
+  action!: string | undefined;
+
+  @Field(() => String, {
+    description: 'The prompt name to use for the session',
+  })
+  promptName!: string;
+}
+
+@InputType()
+class QueryChatHistoriesInput implements Partial<ListHistoriesOptions> {
+  @Field(() => Boolean, { nullable: true })
+  action: boolean | undefined;
+
+  @Field(() => Number, { nullable: true })
+  limit: number | undefined;
+
+  @Field(() => Number, { nullable: true })
+  skip: number | undefined;
+
+  @Field(() => String, { nullable: true })
+  sessionId: string | undefined;
+}
+
+// ================== Return Types ==================
+
+@ObjectType('ChatMessage')
+class ChatMessageType implements Partial<ChatMessage> {
+  @Field(() => String)
+  role!: 'system' | 'assistant' | 'user';
+
+  @Field(() => String)
+  content!: string;
+
+  @Field(() => [String], { nullable: true })
+  attachments!: string[];
+
+  @Field(() => Date, { nullable: true })
+  createdAt!: Date | undefined;
+}
+
+@ObjectType('CopilotHistories')
+class CopilotHistoriesType implements Partial<ChatHistory> {
+  @Field(() => String)
+  sessionId!: string;
+
+  @Field(() => String, {
+    description: 'An mark identifying which view to use to display the session',
+  })
+  action!: string;
+
+  @Field(() => Number, {
+    description: 'The number of tokens used in the session',
+  })
+  tokens!: number;
+
+  @Field(() => [ChatMessageType])
+  messages!: ChatMessageType[];
+}
+
+@ObjectType('CopilotQuota')
+class CopilotQuotaType {
+  @Field(() => SafeIntResolver)
+  limit!: number;
+
+  @Field(() => SafeIntResolver)
+  used!: number;
+}
+
+// ================== Resolver ==================
+
+@ObjectType('Copilot')
+export class CopilotType {
+  @Field(() => ID, { nullable: true })
+  workspaceId!: string | undefined;
+}
+
+@Resolver(() => CopilotType)
+export class CopilotResolver {
+  constructor(
+    private readonly permissions: PermissionService,
+    private readonly quota: QuotaService,
+    private readonly mutex: MutexService,
+    private readonly chatSession: ChatSessionService
+  ) {}
+
+  @ResolveField(() => CopilotQuotaType, {
+    name: 'quota',
+    description: 'Get the quota of the user in the workspace',
+    complexity: 2,
+  })
+  async getQuota(@CurrentUser() user: CurrentUser) {
+    const quota = await this.quota.getUserQuota(user.id);
+    const limit = quota.feature.copilotActionLimit;
+
+    const actions = await this.chatSession.countUserActions(user.id);
+    const chats = await this.chatSession
+      .listHistories(user.id)
+      .then(histories =>
+        histories.reduce(
+          (acc, h) => acc + h.messages.filter(m => m.role === 'user').length,
+          0
+        )
+      );
+
+    return { limit, used: actions + chats };
+  }
+
+  @ResolveField(() => [String], {
+    description: 'Get the session list of chats in the workspace',
+    complexity: 2,
+  })
+  async chats(
+    @Parent() copilot: CopilotType,
+    @CurrentUser() user: CurrentUser
+  ) {
+    if (!copilot.workspaceId) return [];
+    await this.permissions.checkCloudWorkspace(copilot.workspaceId, user.id);
+    return await this.chatSession.listSessions(user.id, copilot.workspaceId);
+  }
+
+  @ResolveField(() => [String], {
+    description: 'Get the session list of actions in the workspace',
+    complexity: 2,
+  })
+  async actions(
+    @Parent() copilot: CopilotType,
+    @CurrentUser() user: CurrentUser
+  ) {
+    if (!copilot.workspaceId) return [];
+    await this.permissions.checkCloudWorkspace(copilot.workspaceId, user.id);
+    return await this.chatSession.listSessions(user.id, copilot.workspaceId, {
+      action: true,
+    });
+  }
+
+  @ResolveField(() => [CopilotHistoriesType], {})
+  async histories(
+    @Parent() copilot: CopilotType,
+    @CurrentUser() user: CurrentUser,
+    @Args('docId', { nullable: true }) docId?: string,
+    @Args({
+      name: 'options',
+      type: () => QueryChatHistoriesInput,
+      nullable: true,
+    })
+    options?: QueryChatHistoriesInput
+  ) {
+    const workspaceId = copilot.workspaceId;
+    if (!workspaceId) {
+      return [];
+    } else if (docId) {
+      await this.permissions.checkCloudPagePermission(
+        workspaceId,
+        docId,
+        user.id
+      );
+    } else {
+      await this.permissions.checkCloudWorkspace(workspaceId, user.id);
+    }
+
+    return await this.chatSession.listHistories(
+      user.id,
+      workspaceId,
+      docId,
+      options
+    );
+  }
+
+  @Public()
+  @Mutation(() => String, {
+    description: 'Create a chat session',
+  })
+  async createCopilotSession(
+    @CurrentUser() user: CurrentUser,
+    @Args({ name: 'options', type: () => CreateChatSessionInput })
+    options: CreateChatSessionInput
+  ) {
+    await this.permissions.checkCloudPagePermission(
+      options.workspaceId,
+      options.docId,
+      user.id
+    );
+    const lockFlag = `session:${user.id}:${options.workspaceId}`;
+    await using lock = await this.mutex.lock(lockFlag);
+    if (!lock) {
+      return new TooManyRequestsException('Server is busy');
+    }
+
+    const { limit, used } = await this.getQuota(user);
+    if (limit && Number.isFinite(limit) && used >= limit) {
+      return new PaymentRequiredException(
+        `You have reached the limit of actions in this workspace, please upgrade your plan.`
+      );
+    }
+
+    const session = await this.chatSession.create({
+      ...options,
+      userId: user.id,
+    });
+    return session;
+  }
+}
+
+@Resolver(() => UserType)
+export class UserCopilotResolver {
+  constructor(private readonly permissions: PermissionService) {}
+
+  @ResolveField(() => CopilotType)
+  async copilot(
+    @CurrentUser() user: CurrentUser,
+    @Args('workspaceId', { nullable: true }) workspaceId?: string
+  ) {
+    if (workspaceId) {
+      await this.permissions.checkCloudWorkspace(workspaceId, user.id);
+    }
+    return { workspaceId };
+  }
+}

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

@@ -11,6 +11,7 @@ import {
   ChatMessageSchema,
   getTokenEncoder,
   PromptMessage,
+  PromptMessageSchema,
   PromptParams,
 } from './types';
 
@@ -105,37 +106,62 @@ export class ChatSession implements AsyncDisposable {
 @Injectable()
 export class ChatSessionService {
   private readonly logger = new Logger(ChatSessionService.name);
+
   constructor(
     private readonly db: PrismaClient,
     private readonly prompt: PromptService
   ) {}
 
-  private async setSession(state: ChatSessionState): Promise<void> {
-    await this.db.aiSession.upsert({
-      where: {
-        id: state.sessionId,
-      },
-      update: {
-        messages: {
-          create: state.messages.map((m, idx) => ({ idx, ...m })),
-        },
-      },
-      create: {
-        id: state.sessionId,
-        messages: { create: state.messages },
-        // connect
-        user: { connect: { id: state.userId } },
-        workspace: { connect: { id: state.workspaceId } },
-        doc: {
-          connect: {
-            id_workspaceId: {
-              id: state.docId,
+  private async setSession(state: ChatSessionState): Promise<string> {
+    return await this.db.$transaction(async tx => {
+      let sessionId = state.sessionId;
+
+      // find existing session if session is chat session
+      if (!state.prompt.action) {
+        const { id } =
+          (await tx.aiSession.findFirst({
+            where: {
+              userId: state.userId,
               workspaceId: state.workspaceId,
+              docId: state.docId,
+              prompt: { action: { equals: null } },
             },
+            select: { id: true },
+          })) || {};
+        if (id) sessionId = id;
+      }
+
+      await tx.aiSession.upsert({
+        where: {
+          id: sessionId,
+          userId: state.userId,
+        },
+        update: {
+          messages: {
+            // delete old messages
+            deleteMany: {},
+            create: state.messages.map(m => ({
+              ...m,
+              params: m.params || undefined,
+            })),
           },
         },
-        prompt: { connect: { name: state.prompt.name } },
-      },
+        create: {
+          id: sessionId,
+          workspaceId: state.workspaceId,
+          docId: state.docId,
+          messages: {
+            create: state.messages.map(m => ({
+              ...m,
+              params: m.params || undefined,
+            })),
+          },
+          // connect
+          user: { connect: { id: state.userId } },
+          prompt: { connect: { name: state.prompt.name } },
+        },
+      });
+      return sessionId;
     });
   }
 
@@ -171,6 +197,7 @@ export class ChatSessionService {
       })
       .then(async session => {
         if (!session) return;
+
         const messages = ChatMessageSchema.array().safeParse(session.messages);
 
         return {
@@ -184,18 +211,58 @@ export class ChatSessionService {
       });
   }
 
-  async listHistories(
+  private calculateTokenSize(
+    messages: PromptMessage[],
+    model: AvailableModel
+  ): number {
+    const encoder = getTokenEncoder(model);
+    return messages
+      .map(m => encoder?.encode_ordinary(m.content).length || 0)
+      .reduce((total, length) => total + length, 0);
+  }
+
+  async countUserActions(userId: string): Promise<number> {
+    return await this.db.aiSession.count({
+      where: { userId, prompt: { action: { not: null } } },
+    });
+  }
+
+  async listSessions(
+    userId: string,
     workspaceId: string,
-    docId: string,
-    options: ListHistoriesOptions
+    options?: { docId?: string; action?: boolean }
+  ): Promise<string[]> {
+    return await this.db.aiSession
+      .findMany({
+        where: {
+          userId,
+          workspaceId,
+          docId: workspaceId === options?.docId ? undefined : options?.docId,
+          prompt: {
+            action: options?.action ? { not: null } : null,
+          },
+        },
+        select: { id: true },
+      })
+      .then(sessions => sessions.map(({ id }) => id));
+  }
+
+  async listHistories(
+    userId: string,
+    workspaceId?: string,
+    docId?: string,
+    options?: ListHistoriesOptions
   ): Promise<ChatHistory[]> {
     return await this.db.aiSession
       .findMany({
         where: {
+          userId,
           workspaceId: workspaceId,
           docId: workspaceId === docId ? undefined : docId,
-          prompt: { action: { not: null } },
-          id: options.sessionId ? { equals: options.sessionId } : undefined,
+          prompt: {
+            action: options?.action ? { not: null } : null,
+          },
+          id: options?.sessionId ? { equals: options.sessionId } : undefined,
         },
         select: {
           id: true,
@@ -210,20 +277,33 @@ export class ChatSessionService {
             },
           },
         },
-        take: options.limit,
-        skip: options.skip,
+        take: options?.limit,
+        skip: options?.skip,
         orderBy: { createdAt: 'desc' },
       })
       .then(sessions =>
         sessions
           .map(({ id, prompt, messages }) => {
-            const ret = ChatMessageSchema.array().safeParse(messages);
-            if (ret.success) {
-              const encoder = getTokenEncoder(prompt.model as AvailableModel);
-              const tokens = ret.data
-                .map(m => encoder?.encode_ordinary(m.content).length || 0)
-                .reduce((total, length) => total + length, 0);
-              return { sessionId: id, tokens, messages: ret.data };
+            try {
+              const ret = PromptMessageSchema.array().safeParse(messages);
+              if (ret.success) {
+                const tokens = this.calculateTokenSize(
+                  ret.data,
+                  prompt.model as AvailableModel
+                );
+                return {
+                  sessionId: id,
+                  action: prompt.action || undefined,
+                  tokens,
+                  messages: ret.data,
+                };
+              } else {
+                this.logger.error(
+                  `Unexpected message schema: ${JSON.stringify(ret.error)}`
+                );
+              }
+            } catch (e) {
+              this.logger.error('Unexpected error in listHistories', e);
             }
             return undefined;
           })
@@ -238,8 +318,12 @@ export class ChatSessionService {
       this.logger.error(`Prompt not found: ${options.promptName}`);
       throw new Error('Prompt not found');
     }
-    await this.setSession({ ...options, sessionId, prompt, messages: [] });
-    return sessionId;
+    return await this.setSession({
+      ...options,
+      sessionId,
+      prompt,
+      messages: [],
+    });
   }
 
   /**

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

@@ -76,8 +76,9 @@ export type ChatMessage = z.infer<typeof ChatMessageSchema>;
 export const ChatHistorySchema = z
   .object({
     sessionId: z.string(),
+    action: z.string().optional(),
     tokens: z.number(),
-    messages: z.array(ChatMessageSchema),
+    messages: z.array(PromptMessageSchema.or(ChatMessageSchema)),
   })
   .strict();
 
@@ -104,8 +105,8 @@ export interface CopilotProvider {
 export interface CopilotTextToTextProvider extends CopilotProvider {
   generateText(
     messages: PromptMessage[],
-    model: string,
-    options: {
+    model?: string,
+    options?: {
       temperature?: number;
       maxTokens?: number;
       signal?: AbortSignal;
@@ -114,8 +115,8 @@ export interface CopilotTextToTextProvider extends CopilotProvider {
   ): Promise<string>;
   generateTextStream(
     messages: PromptMessage[],
-    model: string,
-    options: {
+    model?: string,
+    options?: {
       temperature?: number;
       maxTokens?: number;
       signal?: AbortSignal;