From 2199bad6fd6fff16c9d9d49e5329a4125ef9609b Mon Sep 17 00:00:00 2001 From: DarkSky Date: Tue, 2 Apr 2024 17:58:51 +0800 Subject: [PATCH] feat: public endpoint --- .github/actions/deploy/deploy.mjs | 1 + .../charts/graphql/templates/deployment.yaml | 2 + .../helm/affine/charts/graphql/values.yaml | 1 + .../backend/server/src/config/affine.env.ts | 4 + .../server/src/fundamentals/config/def.ts | 1 + .../server/src/fundamentals/config/default.ts | 1 + .../server/src/fundamentals/graphql/index.ts | 10 ++ .../server/src/plugins/copilot/controller.ts | 8 +- .../server/src/plugins/copilot/resolver.ts | 105 +++++++++++------- .../server/src/plugins/copilot/session.ts | 59 +++++++++- packages/backend/server/src/schema.gql | 1 + .../get-copilot-anonymous-histories.gql | 18 +++ .../get-copilot-anonymous-sessions.gql | 6 + .../frontend/graphql/src/graphql/index.ts | 36 ++++++ packages/frontend/graphql/src/schema.ts | 48 ++++++++ 15 files changed, 253 insertions(+), 48 deletions(-) create mode 100644 packages/frontend/graphql/src/graphql/get-copilot-anonymous-histories.gql create mode 100644 packages/frontend/graphql/src/graphql/get-copilot-anonymous-sessions.gql diff --git a/.github/actions/deploy/deploy.mjs b/.github/actions/deploy/deploy.mjs index e1583cf5409ae..03f22f5ab66a5 100644 --- a/.github/actions/deploy/deploy.mjs +++ b/.github/actions/deploy/deploy.mjs @@ -118,6 +118,7 @@ const createHelmCommand = ({ isDryRun }) => { `--set graphql.app.experimental.enableJwstCodec=${namespace === 'dev'}`, `--set graphql.app.features.earlyAccessPreview=false`, `--set graphql.app.features.syncClientVersionCheck=true`, + `--set graphql.app.features.copilotAuthorization=true`, `--set sync.replicaCount=${syncReplicaCount}`, `--set-string sync.image.tag="${imageTag}"`, ...serviceAnnotations, diff --git a/.github/helm/affine/charts/graphql/templates/deployment.yaml b/.github/helm/affine/charts/graphql/templates/deployment.yaml index 56f575206df4a..01b4bc46aba66 100644 --- a/.github/helm/affine/charts/graphql/templates/deployment.yaml +++ b/.github/helm/affine/charts/graphql/templates/deployment.yaml @@ -83,6 +83,8 @@ spec: value: "{{ .Values.app.features.earlyAccessPreview }}" - name: FEATURES_SYNC_CLIENT_VERSION_CHECK value: "{{ .Values.app.features.syncClientVersionCheck }}" + - name: FEATURES_COPILOT_SKIP_AUTHORIZATION + value: "{{ .Values.app.features.copilotAuthorization }}" - name: MAILER_HOST valueFrom: secretKeyRef: diff --git a/.github/helm/affine/charts/graphql/values.yaml b/.github/helm/affine/charts/graphql/values.yaml index 625cb1b7e9083..a1ebc682ee0fb 100644 --- a/.github/helm/affine/charts/graphql/values.yaml +++ b/.github/helm/affine/charts/graphql/values.yaml @@ -62,6 +62,7 @@ app: features: earlyAccessPreview: false syncClientVersionCheck: false + copilotAuthorization: false serviceAccount: create: true diff --git a/packages/backend/server/src/config/affine.env.ts b/packages/backend/server/src/config/affine.env.ts index c393f6b89dee7..845a6d23f59aa 100644 --- a/packages/backend/server/src/config/affine.env.ts +++ b/packages/backend/server/src/config/affine.env.ts @@ -38,5 +38,9 @@ AFFiNE.ENV_MAP = { 'featureFlags.syncClientVersionCheck', 'boolean', ], + FEATURES_COPILOT_SKIP_AUTHORIZATION: [ + 'featureFlags.copilotAuthorization', + 'boolean', + ], TELEMETRY_ENABLE: ['telemetry.enabled', 'boolean'], }; diff --git a/packages/backend/server/src/fundamentals/config/def.ts b/packages/backend/server/src/fundamentals/config/def.ts index c4c110be3b556..9f8e7b0e7acd7 100644 --- a/packages/backend/server/src/fundamentals/config/def.ts +++ b/packages/backend/server/src/fundamentals/config/def.ts @@ -185,6 +185,7 @@ export interface AFFiNEConfig { featureFlags: { earlyAccessPreview: boolean; syncClientVersionCheck: boolean; + copilotAuthorization: boolean; }; /** diff --git a/packages/backend/server/src/fundamentals/config/default.ts b/packages/backend/server/src/fundamentals/config/default.ts index d4ff5b1f8fdc6..da94f05979ab7 100644 --- a/packages/backend/server/src/fundamentals/config/default.ts +++ b/packages/backend/server/src/fundamentals/config/default.ts @@ -119,6 +119,7 @@ export const getDefaultAFFiNEConfig: () => AFFiNEConfig = () => { featureFlags: { earlyAccessPreview: false, syncClientVersionCheck: false, + copilotAuthorization: false, }, https: false, host: 'localhost', diff --git a/packages/backend/server/src/fundamentals/graphql/index.ts b/packages/backend/server/src/fundamentals/graphql/index.ts index 04b5c6a1d2f11..aa46be9533ce0 100644 --- a/packages/backend/server/src/fundamentals/graphql/index.ts +++ b/packages/backend/server/src/fundamentals/graphql/index.ts @@ -23,6 +23,15 @@ export type GraphqlContext = { GraphQLModule.forRootAsync({ driver: ApolloDriver, useFactory: (config: Config) => { + const copilotAuthorization = config.featureFlags.copilotAuthorization; + const cors = { + cors: { + origin: [ + 'https://try-blocksuite.vercel.app/', + 'http://localhost:5173/', + ], + }, + }; return { ...config.graphql, path: `${config.path}/graphql`, @@ -78,6 +87,7 @@ export type GraphqlContext = { return formattedError; }, + ...(copilotAuthorization ? cors : {}), }; }, inject: [Config], diff --git a/packages/backend/server/src/plugins/copilot/controller.ts b/packages/backend/server/src/plugins/copilot/controller.ts index bc527ebd7998e..d099bfafafb91 100644 --- a/packages/backend/server/src/plugins/copilot/controller.ts +++ b/packages/backend/server/src/plugins/copilot/controller.ts @@ -82,7 +82,7 @@ export class CopilotController { @Public() @Get('/chat/:sessionId') async chat( - @CurrentUser() user: CurrentUser, + @CurrentUser() user: CurrentUser | undefined, @Req() req: Request, @Param('sessionId') sessionId: string, @Query('message') message: string | undefined, @@ -110,7 +110,7 @@ export class CopilotController { session.model, { signal: req.signal, - user: user.id, + user: user?.id, } ); @@ -132,7 +132,7 @@ export class CopilotController { @Public() @Sse('/chat/:sessionId/stream') async chatStream( - @CurrentUser() user: CurrentUser, + @CurrentUser() user: CurrentUser | undefined, @Req() req: Request, @Param('sessionId') sessionId: string, @Query('message') message: string | undefined, @@ -157,7 +157,7 @@ export class CopilotController { return from( provider.generateTextStream(session.finish(params), session.model, { signal: req.signal, - user: user.id, + user: user?.id, }) ).pipe( connect(shared$ => diff --git a/packages/backend/server/src/plugins/copilot/resolver.ts b/packages/backend/server/src/plugins/copilot/resolver.ts index f3e57ab9072d9..ee81cd1585414 100644 --- a/packages/backend/server/src/plugins/copilot/resolver.ts +++ b/packages/backend/server/src/plugins/copilot/resolver.ts @@ -1,4 +1,4 @@ -import { Logger } from '@nestjs/common'; +import { ForbiddenException, Logger } from '@nestjs/common'; import { Args, Field, @@ -7,13 +7,14 @@ import { Mutation, ObjectType, Parent, + Query, registerEnumType, ResolveField, Resolver, } from '@nestjs/graphql'; import { SafeIntResolver } from 'graphql-scalars'; -import { CurrentUser } from '../../core/auth'; +import { CurrentUser, Public } from '../../core/auth'; import { QuotaService } from '../../core/quota'; import { UserType } from '../../core/user'; import { PermissionService } from '../../core/workspaces/permission'; @@ -158,8 +159,11 @@ export class CopilotResolver { }) async getQuota( @Parent() _copilot: CopilotType, - @CurrentUser() user: CurrentUser + @CurrentUser() user: CurrentUser | undefined ) { + // TODO(@darkskygit): remove this after the feature is stable + if (!user) return { used: 0 }; + const quota = await this.quota.getUserQuota(user.id); const limit = quota.feature.copilotActionLimit; @@ -182,11 +186,11 @@ export class CopilotResolver { }) async chats( @Parent() copilot: CopilotType, - @CurrentUser() user: CurrentUser + @CurrentUser() user?: CurrentUser ) { if (!copilot.workspaceId) return []; - await this.permissions.checkCloudWorkspace(copilot.workspaceId, user.id); - return await this.chatSession.listSessions(user.id, copilot.workspaceId); + await this.permissions.checkCloudWorkspace(copilot.workspaceId, user?.id); + return await this.chatSession.listSessions(user?.id); } @ResolveField(() => [String], { @@ -195,11 +199,11 @@ export class CopilotResolver { }) async actions( @Parent() copilot: CopilotType, - @CurrentUser() user: CurrentUser + @CurrentUser() user?: CurrentUser ) { if (!copilot.workspaceId) return []; - await this.permissions.checkCloudWorkspace(copilot.workspaceId, user.id); - return await this.chatSession.listSessions(user.id, copilot.workspaceId, { + await this.permissions.checkCloudWorkspace(copilot.workspaceId, user?.id); + return await this.chatSession.listSessions(user?.id, copilot.workspaceId, { action: true, }); } @@ -207,7 +211,7 @@ export class CopilotResolver { @ResolveField(() => [CopilotHistoriesType], {}) async histories( @Parent() copilot: CopilotType, - @CurrentUser() user: CurrentUser, + @CurrentUser() user?: CurrentUser, @Args('docId', { nullable: true }) docId?: string, @Args({ name: 'options', @@ -217,70 +221,88 @@ export class CopilotResolver { 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); + // todo(@darkskygit): remove this after the feature is stable + const publishable = AFFiNE.featureFlags.copilotAuthorization; + if (user) { + if (!workspaceId) { + return []; + } + if (docId) { + await this.permissions.checkCloudPagePermission( + workspaceId, + docId, + user.id + ); + } else { + await this.permissions.checkCloudWorkspace(workspaceId, user.id); + } + } else if (!publishable) { + return new ForbiddenException('Login required'); } return await this.chatSession.listHistories( - user.id, + user?.id, workspaceId, docId, options ); } + @Public() @Mutation(() => String, { description: 'Create a chat session', }) async createCopilotSession( - @CurrentUser() user: CurrentUser, + @CurrentUser() user: CurrentUser | undefined, @Args({ name: 'options', type: () => CreateChatSessionInput }) options: CreateChatSessionInput ) { - await this.permissions.checkCloudPagePermission( - options.workspaceId, - options.docId, - user.id - ); - const lockFlag = `${COPILOT_LOCKER}:session:${user.id}:${options.workspaceId}`; + // todo(@darkskygit): remove this after the feature is stable + const publishable = AFFiNE.featureFlags.copilotAuthorization; + if (!user && !publishable) { + return new ForbiddenException('Login required'); + } + + const lockFlag = `${COPILOT_LOCKER}: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( - { workspaceId: undefined }, - 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.` + if (options.action && user) { + const { limit, used } = await this.getQuota( + { workspaceId: undefined }, + 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, + // todo: force user to be logged in + userId: user?.id ?? '', }); return session; } + @Public() @Mutation(() => String, { description: 'Create a chat message', }) async createCopilotMessage( - @CurrentUser() user: CurrentUser, + @CurrentUser() user: CurrentUser | undefined, @Args({ name: 'options', type: () => CreateChatMessageInput }) options: CreateChatMessageInput ) { + // todo(@darkskygit): remove this after the feature is stable + const publishable = AFFiNE.featureFlags.copilotAuthorization; + if (!user && !publishable) { + return new ForbiddenException('Login required'); + } + const lockFlag = `${COPILOT_LOCKER}:message:${user?.id}:${options.sessionId}`; await using lock = await this.mutex.lock(lockFlag); if (!lock) { @@ -314,4 +336,11 @@ export class UserCopilotResolver { } return { workspaceId }; } + + @Public() + @Query(() => CopilotType) + async copilotAnonymous(@Args('workspaceId') workspaceId: string) { + if (!AFFiNE.featureFlags.copilotAuthorization) return; + return { workspaceId }; + } } diff --git a/packages/backend/server/src/plugins/copilot/session.ts b/packages/backend/server/src/plugins/copilot/session.ts index 6fca6d688edac..551cc70e3e46a 100644 --- a/packages/backend/server/src/plugins/copilot/session.ts +++ b/packages/backend/server/src/plugins/copilot/session.ts @@ -112,6 +112,8 @@ export class ChatSession implements AsyncDisposable { @Injectable() export class ChatSessionService { private readonly logger = new Logger(ChatSessionService.name); + // NOTE: only used for anonymous session in development + private readonly unsavedSessions = new Map(); constructor( private readonly db: PrismaClient, @@ -120,6 +122,13 @@ export class ChatSessionService { ) {} private async setSession(state: ChatSessionState): Promise { + if (!state.userId && AFFiNE.featureFlags.copilotAuthorization) { + // todo(@darkskygit): allow anonymous session in development + // remove this after the feature is stable + this.unsavedSessions.set(state.sessionId, state); + return state.sessionId; + } + return await this.db.$transaction(async tx => { let sessionId = state.sessionId; @@ -141,7 +150,6 @@ export class ChatSessionService { await tx.aiSession.upsert({ where: { id: sessionId, - userId: state.userId, }, update: { messages: { @@ -203,8 +211,15 @@ export class ChatSessionService { }, }) .then(async session => { - if (!session) return; - + if (!session) { + const publishable = AFFiNE.featureFlags.copilotAuthorization; + if (publishable) { + // todo(@darkskygit): allow anonymous session in development + // remove this after the feature is stable + return this.unsavedSessions.get(sessionId); + } + return; + } const messages = ChatMessageSchema.array().safeParse(session.messages); return { @@ -229,16 +244,25 @@ export class ChatSessionService { } async countUserActions(userId: string): Promise { + // NOTE: only used for anonymous session in development + if (!userId && AFFiNE.featureFlags.copilotAuthorization) { + return this.unsavedSessions.size; + } return await this.db.aiSession.count({ where: { userId, prompt: { action: { not: null } } }, }); } async listSessions( - userId: string, - workspaceId: string, + userId: string | undefined, + workspaceId?: string, options?: { docId?: string; action?: boolean } ): Promise { + // NOTE: only used for anonymous session in development + if (!userId && AFFiNE.featureFlags.copilotAuthorization) { + return Array.from(this.unsavedSessions.keys()); + } + return await this.db.aiSession .findMany({ where: { @@ -255,11 +279,34 @@ export class ChatSessionService { } async listHistories( - userId: string, + userId: string | undefined, workspaceId?: string, docId?: string, options?: ListHistoriesOptions ): Promise { + // NOTE: only used for anonymous session in development + if (!userId && AFFiNE.featureFlags.copilotAuthorization) { + return [...this.unsavedSessions.values()] + .map(state => { + const ret = ChatMessageSchema.array().safeParse(state.messages); + if (ret.success) { + const tokens = this.calculateTokenSize( + state.messages, + state.prompt.model as AvailableModel + ); + return { + sessionId: state.sessionId, + action: state.prompt.action, + tokens, + messages: ret.data, + }; + } + console.error('Unexpected error in listHistories', ret.error); + return undefined; + }) + .filter((v): v is NonNullable => !!v); + } + return await this.db.aiSession .findMany({ where: { diff --git a/packages/backend/server/src/schema.gql b/packages/backend/server/src/schema.gql index ba1e50db182f4..30e5cbff331a3 100644 --- a/packages/backend/server/src/schema.gql +++ b/packages/backend/server/src/schema.gql @@ -249,6 +249,7 @@ enum PublicPageMode { type Query { checkBlobSize(size: SafeInt!, workspaceId: String!): WorkspaceBlobSizes! @deprecated(reason: "no more needed") collectAllBlobSizes: WorkspaceBlobSizes! @deprecated(reason: "use `user.storageUsage` instead") + copilotAnonymous(workspaceId: String!): Copilot! """Get current user""" currentUser: UserType diff --git a/packages/frontend/graphql/src/graphql/get-copilot-anonymous-histories.gql b/packages/frontend/graphql/src/graphql/get-copilot-anonymous-histories.gql new file mode 100644 index 0000000000000..f7d91d2a2ef00 --- /dev/null +++ b/packages/frontend/graphql/src/graphql/get-copilot-anonymous-histories.gql @@ -0,0 +1,18 @@ +query getCopilotAnonymousHistories( + $workspaceId: String! + $docId: String + $options: QueryChatHistoriesInput +) { + copilotAnonymous(workspaceId: $workspaceId) { + histories(docId: $docId, options: $options) { + sessionId + tokens + messages { + role + content + attachments + createdAt + } + } + } +} diff --git a/packages/frontend/graphql/src/graphql/get-copilot-anonymous-sessions.gql b/packages/frontend/graphql/src/graphql/get-copilot-anonymous-sessions.gql new file mode 100644 index 0000000000000..57c4f77a5a470 --- /dev/null +++ b/packages/frontend/graphql/src/graphql/get-copilot-anonymous-sessions.gql @@ -0,0 +1,6 @@ +query getCopilotAnonymousSessions($workspaceId: String!) { + copilotAnonymous(workspaceId: $workspaceId) { + chats + actions + } +} diff --git a/packages/frontend/graphql/src/graphql/index.ts b/packages/frontend/graphql/src/graphql/index.ts index cdb7e924774e1..91c925ee0587a 100644 --- a/packages/frontend/graphql/src/graphql/index.ts +++ b/packages/frontend/graphql/src/graphql/index.ts @@ -251,6 +251,42 @@ mutation removeEarlyAccess($email: String!) { }`, }; +export const getCopilotAnonymousHistoriesQuery = { + id: 'getCopilotAnonymousHistoriesQuery' as const, + operationName: 'getCopilotAnonymousHistories', + definitionName: 'copilotAnonymous', + containsFile: false, + query: ` +query getCopilotAnonymousHistories($workspaceId: String!, $docId: String, $options: QueryChatHistoriesInput) { + copilotAnonymous(workspaceId: $workspaceId) { + histories(docId: $docId, options: $options) { + sessionId + tokens + messages { + role + content + attachments + createdAt + } + } + } +}`, +}; + +export const getCopilotAnonymousSessionsQuery = { + id: 'getCopilotAnonymousSessionsQuery' as const, + operationName: 'getCopilotAnonymousSessions', + definitionName: 'copilotAnonymous', + containsFile: false, + query: ` +query getCopilotAnonymousSessions($workspaceId: String!) { + copilotAnonymous(workspaceId: $workspaceId) { + chats + actions + } +}`, +}; + export const getCopilotHistoriesQuery = { id: 'getCopilotHistoriesQuery' as const, operationName: 'getCopilotHistories', diff --git a/packages/frontend/graphql/src/schema.ts b/packages/frontend/graphql/src/schema.ts index 035d5dd09c4c3..d641ef2725b19 100644 --- a/packages/frontend/graphql/src/schema.ts +++ b/packages/frontend/graphql/src/schema.ts @@ -341,6 +341,44 @@ export type PasswordLimitsFragment = { maxLength: number; }; +export type GetCopilotAnonymousHistoriesQueryVariables = Exact<{ + workspaceId: Scalars['String']['input']; + docId: InputMaybe; + options: InputMaybe; +}>; + +export type GetCopilotAnonymousHistoriesQuery = { + __typename?: 'Query'; + copilotAnonymous: { + __typename?: 'Copilot'; + histories: Array<{ + __typename?: 'CopilotHistories'; + sessionId: string; + tokens: number; + messages: Array<{ + __typename?: 'ChatMessage'; + role: string; + content: string; + attachments: Array | null; + createdAt: string | null; + }>; + }>; + }; +}; + +export type GetCopilotAnonymousSessionsQueryVariables = Exact<{ + workspaceId: Scalars['String']['input']; +}>; + +export type GetCopilotAnonymousSessionsQuery = { + __typename?: 'Query'; + copilotAnonymous: { + __typename?: 'Copilot'; + chats: Array; + actions: Array; + }; +}; + export type GetCopilotHistoriesQueryVariables = Exact<{ workspaceId: Scalars['String']['input']; docId: InputMaybe; @@ -1045,6 +1083,16 @@ export type Queries = variables: EarlyAccessUsersQueryVariables; response: EarlyAccessUsersQuery; } + | { + name: 'getCopilotAnonymousHistoriesQuery'; + variables: GetCopilotAnonymousHistoriesQueryVariables; + response: GetCopilotAnonymousHistoriesQuery; + } + | { + name: 'getCopilotAnonymousSessionsQuery'; + variables: GetCopilotAnonymousSessionsQueryVariables; + response: GetCopilotAnonymousSessionsQuery; + } | { name: 'getCopilotHistoriesQuery'; variables: GetCopilotHistoriesQueryVariables;