From e6a576551ad3c93c9068375784c5813009d7fbfe Mon Sep 17 00:00:00 2001 From: darkskygit Date: Wed, 10 Apr 2024 11:15:31 +0000 Subject: [PATCH] feat: add copilot impl (#6230) fix CLOUD-22 fix CLOUD-24 --- .github/actions/deploy/deploy.mjs | 6 +- .../graphql/templates/copilot-secret.yaml | 9 + .../charts/graphql/templates/deployment.yaml | 7 + .../helm/affine/charts/graphql/values.yaml | 5 + .github/workflows/deploy.yml | 2 +- .../backend/server/src/core/quota/quota.ts | 7 + .../backend/server/src/core/quota/schema.ts | 28 ++- .../backend/server/src/core/quota/storage.ts | 3 + .../backend/server/src/core/quota/types.ts | 7 + .../server/src/plugins/copilot/index.ts | 17 +- .../src/plugins/copilot/providers/index.ts | 10 +- .../src/plugins/copilot/providers/openai.ts | 195 ++++++++++++++++++ .../server/src/plugins/copilot/session.ts | 64 ++++++ .../server/src/plugins/copilot/types.ts | 24 ++- packages/backend/server/src/schema.gql | 2 + packages/backend/server/tests/quota.spec.ts | 50 +++-- .../src/graphql/create-copilot-session.gql | 3 + .../get-copilot-anonymous-histories.gql | 17 ++ .../get-copilot-anonymous-sessions.gql | 6 + .../src/graphql/get-copilot-histories.gql | 19 ++ .../src/graphql/get-copilot-sessions.gql | 8 + .../frontend/graphql/src/graphql/index.ts | 85 ++++++++ packages/frontend/graphql/src/schema.ts | 129 ++++++++++++ 23 files changed, 669 insertions(+), 34 deletions(-) create mode 100644 .github/helm/affine/charts/graphql/templates/copilot-secret.yaml create mode 100644 packages/backend/server/src/plugins/copilot/providers/openai.ts create mode 100644 packages/frontend/graphql/src/graphql/create-copilot-session.gql 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 create mode 100644 packages/frontend/graphql/src/graphql/get-copilot-histories.gql create mode 100644 packages/frontend/graphql/src/graphql/get-copilot-sessions.gql diff --git a/.github/actions/deploy/deploy.mjs b/.github/actions/deploy/deploy.mjs index 8940592502..33b5707c8c 100644 --- a/.github/actions/deploy/deploy.mjs +++ b/.github/actions/deploy/deploy.mjs @@ -13,8 +13,8 @@ const { R2_ACCOUNT_ID, R2_ACCESS_KEY_ID, R2_SECRET_ACCESS_KEY, - ENABLE_CAPTCHA, CAPTCHA_TURNSTILE_SECRET, + COPILOT_OPENAI_API_KEY, MAILER_SENDER, MAILER_USER, MAILER_PASSWORD, @@ -97,8 +97,10 @@ const createHelmCommand = ({ isDryRun }) => { `--set graphql.replicaCount=${graphqlReplicaCount}`, `--set-string graphql.image.tag="${imageTag}"`, `--set graphql.app.host=${host}`, - `--set graphql.app.captcha.enabled=${ENABLE_CAPTCHA}`, + `--set graphql.app.captcha.enabled=true`, `--set-string graphql.app.captcha.turnstile.secret="${CAPTCHA_TURNSTILE_SECRET}"`, + `--set graphql.app.copilot.enabled=true`, + `--set-string graphql.app.copilot.openai.key="${COPILOT_OPENAI_API_KEY}"`, `--set graphql.app.objectStorage.r2.enabled=true`, `--set-string graphql.app.objectStorage.r2.accountId="${R2_ACCOUNT_ID}"`, `--set-string graphql.app.objectStorage.r2.accessKeyId="${R2_ACCESS_KEY_ID}"`, diff --git a/.github/helm/affine/charts/graphql/templates/copilot-secret.yaml b/.github/helm/affine/charts/graphql/templates/copilot-secret.yaml new file mode 100644 index 0000000000..277b1ff965 --- /dev/null +++ b/.github/helm/affine/charts/graphql/templates/copilot-secret.yaml @@ -0,0 +1,9 @@ +{{- if .Values.app.copilot.enabled -}} +apiVersion: v1 +kind: Secret +metadata: + name: "{{ .Values.app.copilot.secretName }}" +type: Opaque +data: + openaiSecret: {{ .Values.app.copilot.openai.key | b64enc }} +{{- end }} diff --git a/.github/helm/affine/charts/graphql/templates/deployment.yaml b/.github/helm/affine/charts/graphql/templates/deployment.yaml index 5553d2f1a6..faa9b02fb6 100644 --- a/.github/helm/affine/charts/graphql/templates/deployment.yaml +++ b/.github/helm/affine/charts/graphql/templates/deployment.yaml @@ -148,6 +148,13 @@ spec: name: "{{ .Values.app.captcha.secretName }}" key: turnstileSecret {{ end }} + {{ if .Values.app.copilot.enabled }} + - name: COPILOT_OPENAI_API_KEY + valueFrom: + secretKeyRef: + name: "{{ .Values.app.copilot.secretName }}" + key: openaiSecret + {{ end }} {{ if .Values.app.oauth.google.enabled }} - name: OAUTH_GOOGLE_ENABLED value: "true" diff --git a/.github/helm/affine/charts/graphql/values.yaml b/.github/helm/affine/charts/graphql/values.yaml index f4ca76a970..625cb1b7e9 100644 --- a/.github/helm/affine/charts/graphql/values.yaml +++ b/.github/helm/affine/charts/graphql/values.yaml @@ -24,6 +24,11 @@ app: secretName: captcha turnstile: secret: '' + copilot: + enable: false + secretName: copilot + openai: + key: '' objectStorage: r2: enabled: false diff --git a/.github/workflows/deploy.yml b/.github/workflows/deploy.yml index 8faac5215f..bc3253da6a 100644 --- a/.github/workflows/deploy.yml +++ b/.github/workflows/deploy.yml @@ -133,8 +133,8 @@ jobs: R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }} R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }} R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} - ENABLE_CAPTCHA: true CAPTCHA_TURNSTILE_SECRET: ${{ secrets.CAPTCHA_TURNSTILE_SECRET }} + COPILOT_OPENAI_API_KEY: ${{ secrets.COPILOT_OPENAI_API_KEY }} MAILER_SENDER: ${{ secrets.OAUTH_EMAIL_SENDER }} MAILER_USER: ${{ secrets.OAUTH_EMAIL_LOGIN }} MAILER_PASSWORD: ${{ secrets.OAUTH_EMAIL_PASSWORD }} diff --git a/packages/backend/server/src/core/quota/quota.ts b/packages/backend/server/src/core/quota/quota.ts index 3f481de06d..61422a7c42 100644 --- a/packages/backend/server/src/core/quota/quota.ts +++ b/packages/backend/server/src/core/quota/quota.ts @@ -79,6 +79,10 @@ export class QuotaConfig { return this.config.configs.memberLimit; } + get copilotActionLimit() { + return this.config.configs.copilotActionLimit || undefined; + } + get humanReadable() { return { name: this.config.configs.name, @@ -86,6 +90,9 @@ export class QuotaConfig { storageQuota: formatSize(this.storageQuota), historyPeriod: formatDate(this.historyPeriod), memberLimit: this.memberLimit.toString(), + copilotActionLimit: this.copilotActionLimit + ? `${this.copilotActionLimit} times` + : 'Unlimited', }; } } diff --git a/packages/backend/server/src/core/quota/schema.ts b/packages/backend/server/src/core/quota/schema.ts index 5c607f8a21..5776b98481 100644 --- a/packages/backend/server/src/core/quota/schema.ts +++ b/packages/backend/server/src/core/quota/schema.ts @@ -93,11 +93,35 @@ export const Quotas: Quota[] = [ memberLimit: 3, }, }, + { + feature: QuotaType.FreePlanV1, + type: FeatureKind.Quota, + version: 4, + configs: { + // quota name + name: 'Free', + // single blob limit 10MB + blobLimit: 10 * OneMB, + // server limit will larger then client to handle a edge case: + // when a user downgrades from pro to free, he can still continue + // to upload previously added files that exceed the free limit + // NOTE: this is a product decision, may change in future + businessBlobLimit: 100 * OneMB, + // total blob limit 10GB + storageQuota: 10 * OneGB, + // history period of validity 7 days + historyPeriod: 7 * OneDay, + // member limit 3 + memberLimit: 3, + // copilot action limit 10 + copilotActionLimit: 10, + }, + }, ]; export const Quota_FreePlanV1_1 = { - feature: Quotas[4].feature, - version: Quotas[4].version, + feature: Quotas[5].feature, + version: Quotas[5].version, }; export const Quota_ProPlanV1 = { diff --git a/packages/backend/server/src/core/quota/storage.ts b/packages/backend/server/src/core/quota/storage.ts index f3ddd2e60d..7ec0d510e6 100644 --- a/packages/backend/server/src/core/quota/storage.ts +++ b/packages/backend/server/src/core/quota/storage.ts @@ -33,6 +33,7 @@ export class QuotaManagementService { storageQuota: quota.feature.storageQuota, historyPeriod: quota.feature.historyPeriod, memberLimit: quota.feature.memberLimit, + copilotActionLimit: quota.feature.copilotActionLimit, }; } @@ -72,6 +73,7 @@ export class QuotaManagementService { historyPeriod, memberLimit, storageQuota, + copilotActionLimit, humanReadable, }, } = await this.quota.getUserQuota(owner.id); @@ -85,6 +87,7 @@ export class QuotaManagementService { historyPeriod, memberLimit, storageQuota, + copilotActionLimit, humanReadable, usedSize, }; diff --git a/packages/backend/server/src/core/quota/types.ts b/packages/backend/server/src/core/quota/types.ts index 8bc5854066..800b87f751 100644 --- a/packages/backend/server/src/core/quota/types.ts +++ b/packages/backend/server/src/core/quota/types.ts @@ -34,6 +34,7 @@ const quotaPlan = z.object({ historyPeriod: z.number().positive().int(), memberLimit: z.number().positive().int(), businessBlobLimit: z.number().positive().int().nullish(), + copilotActionLimit: z.number().positive().int().nullish(), }), }); @@ -65,6 +66,9 @@ export class HumanReadableQuotaType { @Field(() => String) memberLimit!: string; + + @Field(() => String, { nullable: true }) + copilotActionLimit?: string; } @ObjectType() @@ -84,6 +88,9 @@ export class QuotaQueryType { @Field(() => SafeIntResolver) storageQuota!: number; + @Field(() => SafeIntResolver, { nullable: true }) + copilotActionLimit?: number; + @Field(() => HumanReadableQuotaType) humanReadable!: HumanReadableQuotaType; diff --git a/packages/backend/server/src/plugins/copilot/index.ts b/packages/backend/server/src/plugins/copilot/index.ts index 0cb50c5ae4..732109abff 100644 --- a/packages/backend/server/src/plugins/copilot/index.ts +++ b/packages/backend/server/src/plugins/copilot/index.ts @@ -1,12 +1,25 @@ import { ServerFeature } from '../../core/config'; +import { PermissionService } from '../../core/workspaces/permission'; import { Plugin } from '../registry'; import { PromptService } from './prompt'; -import { assertProvidersConfigs, CopilotProviderService } from './providers'; +import { + assertProvidersConfigs, + CopilotProviderService, + OpenAIProvider, + registerCopilotProvider, +} from './providers'; import { ChatSessionService } from './session'; +registerCopilotProvider(OpenAIProvider); + @Plugin({ name: 'copilot', - providers: [ChatSessionService, PromptService, CopilotProviderService], + providers: [ + PermissionService, + ChatSessionService, + PromptService, + CopilotProviderService, + ], contributesTo: ServerFeature.Copilot, if: config => { if (config.flavor.graphql) { diff --git a/packages/backend/server/src/plugins/copilot/providers/index.ts b/packages/backend/server/src/plugins/copilot/providers/index.ts index 2b66669d88..52164d2a3d 100644 --- a/packages/backend/server/src/plugins/copilot/providers/index.ts +++ b/packages/backend/server/src/plugins/copilot/providers/index.ts @@ -5,9 +5,9 @@ import { Injectable, Logger } from '@nestjs/common'; import { Config } from '../../../fundamentals'; import { CapabilityToCopilotProvider, + CopilotCapability, CopilotConfig, CopilotProvider, - CopilotProviderCapability, CopilotProviderType, } from '../types'; @@ -19,7 +19,7 @@ interface CopilotProviderDefinition { // type of the provider readonly type: CopilotProviderType; // capabilities of the provider, like text to text, text to image, etc. - readonly capabilities: CopilotProviderCapability[]; + readonly capabilities: CopilotCapability[]; // asserts that the config is valid for this provider assetsConfig(config: C): boolean; } @@ -32,7 +32,7 @@ const COPILOT_PROVIDER = new Map< // map of capabilities to providers const PROVIDER_CAPABILITY_MAP = new Map< - CopilotProviderCapability, + CopilotCapability, CopilotProviderType[] >(); @@ -116,7 +116,7 @@ export class CopilotProviderService { return this.cachedProviders.get(provider)!; } - getProviderByCapability( + getProviderByCapability( capability: C, prefer?: CopilotProviderType ): CapabilityToCopilotProvider[C] | null { @@ -133,3 +133,5 @@ export class CopilotProviderService { return null; } } + +export { OpenAIProvider } from './openai'; diff --git a/packages/backend/server/src/plugins/copilot/providers/openai.ts b/packages/backend/server/src/plugins/copilot/providers/openai.ts new file mode 100644 index 0000000000..0a863430f3 --- /dev/null +++ b/packages/backend/server/src/plugins/copilot/providers/openai.ts @@ -0,0 +1,195 @@ +import assert from 'node:assert'; + +import { ClientOptions, OpenAI } from 'openai'; + +import { + ChatMessage, + ChatMessageRole, + CopilotCapability, + CopilotProviderType, + CopilotTextToEmbeddingProvider, + CopilotTextToTextProvider, +} from '../types'; + +export class OpenAIProvider + implements CopilotTextToTextProvider, CopilotTextToEmbeddingProvider +{ + static readonly type = CopilotProviderType.OpenAI; + static readonly capabilities = [ + CopilotCapability.TextToText, + CopilotCapability.TextToEmbedding, + CopilotCapability.TextToImage, + ]; + + readonly availableModels = [ + // text to text + 'gpt-4-vision-preview', + 'gpt-4-turbo-preview', + 'gpt-3.5-turbo', + // embeddings + 'text-embedding-3-large', + 'text-embedding-3-small', + 'text-embedding-ada-002', + // moderation + 'text-moderation-latest', + 'text-moderation-stable', + ]; + + private readonly instance: OpenAI; + + constructor(config: ClientOptions) { + assert(OpenAIProvider.assetsConfig(config)); + this.instance = new OpenAI(config); + } + + static assetsConfig(config: ClientOptions) { + return !!config.apiKey; + } + + getCapabilities(): CopilotCapability[] { + return OpenAIProvider.capabilities; + } + + private chatToGPTMessage(messages: ChatMessage[]) { + // filter redundant fields + return messages.map(message => ({ + role: message.role, + content: message.content, + })); + } + + private checkParams({ + messages, + embeddings, + model, + }: { + messages?: ChatMessage[]; + embeddings?: string[]; + model: string; + }) { + if (!this.availableModels.includes(model)) { + throw new Error(`Invalid model: ${model}`); + } + if (Array.isArray(messages) && messages.length > 0) { + if ( + messages.some( + m => + // check non-object + typeof m !== 'object' || + !m || + // check content + typeof m.content !== 'string' || + !m.content || + !m.content.trim() + ) + ) { + throw new Error('Empty message content'); + } + if ( + messages.some( + m => + typeof m.role !== 'string' || + !m.role || + !ChatMessageRole.includes(m.role) + ) + ) { + throw new Error('Invalid message role'); + } + } else if ( + Array.isArray(embeddings) && + embeddings.some(e => typeof e !== 'string' || !e || !e.trim()) + ) { + throw new Error('Invalid embedding'); + } + } + + // ====== text to text ====== + + async generateText( + messages: ChatMessage[], + model: string = 'gpt-3.5-turbo', + options: { + temperature?: number; + maxTokens?: number; + signal?: AbortSignal; + user?: string; + } = {} + ): Promise { + this.checkParams({ messages, model }); + const result = await this.instance.chat.completions.create( + { + messages: this.chatToGPTMessage(messages), + model: model, + temperature: options.temperature || 0, + max_tokens: options.maxTokens || 4096, + user: options.user, + }, + { signal: options.signal } + ); + const { content } = result.choices[0].message; + if (!content) { + throw new Error('Failed to generate text'); + } + return content; + } + + async *generateTextStream( + messages: ChatMessage[], + model: string, + options: { + temperature?: number; + maxTokens?: number; + signal?: AbortSignal; + user?: string; + } = {} + ): AsyncIterable { + this.checkParams({ messages, model }); + const result = await this.instance.chat.completions.create( + { + stream: true, + messages: this.chatToGPTMessage(messages), + model: model, + temperature: options.temperature || 0, + max_tokens: options.maxTokens || 4096, + user: options.user, + }, + { + signal: options.signal, + } + ); + + for await (const message of result) { + const content = message.choices[0].delta.content; + if (content) { + yield content; + if (options.signal?.aborted) { + result.controller.abort(); + break; + } + } + } + } + + // ====== text to embedding ====== + + async generateEmbedding( + messages: string | string[], + model: string, + options: { + dimensions: number; + signal?: AbortSignal; + user?: string; + } = { dimensions: 256 } + ): Promise { + messages = Array.isArray(messages) ? messages : [messages]; + this.checkParams({ embeddings: messages, model }); + + const result = await this.instance.embeddings.create({ + model: model, + input: messages, + dimensions: options.dimensions, + user: options.user, + }); + return result.data.map(e => e.embedding); + } +} diff --git a/packages/backend/server/src/plugins/copilot/session.ts b/packages/backend/server/src/plugins/copilot/session.ts index e762cbe552..6ac0691f17 100644 --- a/packages/backend/server/src/plugins/copilot/session.ts +++ b/packages/backend/server/src/plugins/copilot/session.ts @@ -5,8 +5,11 @@ import { PrismaClient } from '@prisma/client'; import { ChatPrompt, PromptService } from './prompt'; import { + AvailableModel, + ChatHistory, ChatMessage, ChatMessageSchema, + getTokenEncoder, PromptMessage, PromptParams, } from './types'; @@ -27,6 +30,13 @@ export interface ChatSessionState messages: ChatMessage[]; } +export type ListHistoriesOptions = { + action: boolean | undefined; + limit: number | undefined; + skip: number | undefined; + sessionId: string | undefined; +}; + export class ChatSession implements AsyncDisposable { constructor( private readonly state: ChatSessionState, @@ -39,6 +49,13 @@ export class ChatSession implements AsyncDisposable { } push(message: ChatMessage) { + if ( + this.state.prompt.action && + this.state.messages.length > 0 && + message.role === 'user' + ) { + throw new Error('Action has been taken, no more messages allowed'); + } this.state.messages.push(message); } @@ -167,6 +184,53 @@ export class ChatSessionService { }); } + async listHistories( + workspaceId: string, + docId: string, + options: ListHistoriesOptions + ): Promise { + return await this.db.aiSession + .findMany({ + where: { + workspaceId: workspaceId, + docId: workspaceId === docId ? undefined : docId, + prompt: { action: { not: null } }, + id: options.sessionId ? { equals: options.sessionId } : undefined, + }, + select: { + id: true, + prompt: true, + messages: { + select: { + role: true, + content: true, + }, + orderBy: { + createdAt: 'asc', + }, + }, + }, + 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 }; + } + return undefined; + }) + .filter((v): v is NonNullable => !!v) + ); + } + async create(options: ChatSessionOptions): Promise { const sessionId = randomUUID(); const prompt = await this.prompt.get(options.promptName); diff --git a/packages/backend/server/src/plugins/copilot/types.ts b/packages/backend/server/src/plugins/copilot/types.ts index f48ef60bb0..34c6c996f3 100644 --- a/packages/backend/server/src/plugins/copilot/types.ts +++ b/packages/backend/server/src/plugins/copilot/types.ts @@ -90,7 +90,7 @@ export enum CopilotProviderType { OpenAI = 'openai', } -export enum CopilotProviderCapability { +export enum CopilotCapability { TextToText = 'text-to-text', TextToEmbedding = 'text-to-embedding', TextToImage = 'text-to-image', @@ -98,7 +98,7 @@ export enum CopilotProviderCapability { } export interface CopilotProvider { - getCapabilities(): CopilotProviderCapability[]; + getCapabilities(): CopilotCapability[]; } export interface CopilotTextToTextProvider extends CopilotProvider { @@ -124,15 +124,25 @@ export interface CopilotTextToTextProvider extends CopilotProvider { ): AsyncIterable; } -export interface CopilotTextToEmbeddingProvider extends CopilotProvider {} +export interface CopilotTextToEmbeddingProvider extends CopilotProvider { + generateEmbedding( + messages: string[] | string, + model: string, + options: { + dimensions: number; + signal?: AbortSignal; + user?: string; + } + ): Promise; +} export interface CopilotTextToImageProvider extends CopilotProvider {} export interface CopilotImageToImageProvider extends CopilotProvider {} export type CapabilityToCopilotProvider = { - [CopilotProviderCapability.TextToText]: CopilotTextToTextProvider; - [CopilotProviderCapability.TextToEmbedding]: CopilotTextToEmbeddingProvider; - [CopilotProviderCapability.TextToImage]: CopilotTextToImageProvider; - [CopilotProviderCapability.ImageToImage]: CopilotImageToImageProvider; + [CopilotCapability.TextToText]: CopilotTextToTextProvider; + [CopilotCapability.TextToEmbedding]: CopilotTextToEmbeddingProvider; + [CopilotCapability.TextToImage]: CopilotTextToImageProvider; + [CopilotCapability.ImageToImage]: CopilotImageToImageProvider; }; diff --git a/packages/backend/server/src/schema.gql b/packages/backend/server/src/schema.gql index 9cb954954c..b1537af422 100644 --- a/packages/backend/server/src/schema.gql +++ b/packages/backend/server/src/schema.gql @@ -38,6 +38,7 @@ enum FeatureType { type HumanReadableQuotaType { blobLimit: String! + copilotActionLimit: String historyPeriod: String! memberLimit: String! name: String! @@ -224,6 +225,7 @@ type Query { type QuotaQueryType { blobLimit: SafeInt! + copilotActionLimit: SafeInt historyPeriod: SafeInt! humanReadable: HumanReadableQuotaType! memberLimit: SafeInt! diff --git a/packages/backend/server/tests/quota.spec.ts b/packages/backend/server/tests/quota.spec.ts index 7c55329a77..58fa7b33db 100644 --- a/packages/backend/server/tests/quota.spec.ts +++ b/packages/backend/server/tests/quota.spec.ts @@ -18,7 +18,7 @@ import { createTestingModule } from './utils'; const test = ava as TestFn<{ auth: AuthService; quota: QuotaService; - storageQuota: QuotaManagementService; + quotaManager: QuotaManagementService; module: TestingModule; }>; @@ -28,12 +28,12 @@ test.beforeEach(async t => { }); const quota = module.get(QuotaService); - const storageQuota = module.get(QuotaManagementService); + const quotaManager = module.get(QuotaManagementService); const auth = module.get(AuthService); t.context.module = module; t.context.quota = quota; - t.context.storageQuota = storageQuota; + t.context.quotaManager = quotaManager; t.context.auth = auth; }); @@ -49,7 +49,7 @@ test('should be able to set quota', async t => { const q1 = await quota.getUserQuota(u1.id); t.truthy(q1, 'should have quota'); t.is(q1?.feature.name, QuotaType.FreePlanV1, 'should be free plan'); - t.is(q1?.feature.version, 3, 'should be version 3'); + t.is(q1?.feature.version, 4, 'should be version 4'); await quota.switchUserQuota(u1.id, QuotaType.ProPlanV1); @@ -61,35 +61,35 @@ test('should be able to set quota', async t => { }); test('should be able to check storage quota', async t => { - const { auth, quota, storageQuota } = t.context; + const { auth, quota, quotaManager } = t.context; const u1 = await auth.signUp('DarkSky', 'darksky@example.org', '123456'); - const q1 = await storageQuota.getUserQuota(u1.id); - t.is(q1?.blobLimit, Quotas[4].configs.blobLimit, 'should be free plan'); - t.is(q1?.storageQuota, Quotas[4].configs.storageQuota, 'should be free plan'); + const q1 = await quotaManager.getUserQuota(u1.id); + t.is(q1?.blobLimit, Quotas[5].configs.blobLimit, 'should be free plan'); + t.is(q1?.storageQuota, Quotas[5].configs.storageQuota, 'should be free plan'); await quota.switchUserQuota(u1.id, QuotaType.ProPlanV1); - const q2 = await storageQuota.getUserQuota(u1.id); + const q2 = await quotaManager.getUserQuota(u1.id); t.is(q2?.blobLimit, Quotas[1].configs.blobLimit, 'should be pro plan'); t.is(q2?.storageQuota, Quotas[1].configs.storageQuota, 'should be pro plan'); }); test('should be able revert quota', async t => { - const { auth, quota, storageQuota } = t.context; + const { auth, quota, quotaManager } = t.context; const u1 = await auth.signUp('DarkSky', 'darksky@example.org', '123456'); - const q1 = await storageQuota.getUserQuota(u1.id); - t.is(q1?.blobLimit, Quotas[4].configs.blobLimit, 'should be free plan'); - t.is(q1?.storageQuota, Quotas[4].configs.storageQuota, 'should be free plan'); + const q1 = await quotaManager.getUserQuota(u1.id); + t.is(q1?.blobLimit, Quotas[5].configs.blobLimit, 'should be free plan'); + t.is(q1?.storageQuota, Quotas[5].configs.storageQuota, 'should be free plan'); await quota.switchUserQuota(u1.id, QuotaType.ProPlanV1); - const q2 = await storageQuota.getUserQuota(u1.id); + const q2 = await quotaManager.getUserQuota(u1.id); t.is(q2?.blobLimit, Quotas[1].configs.blobLimit, 'should be pro plan'); t.is(q2?.storageQuota, Quotas[1].configs.storageQuota, 'should be pro plan'); await quota.switchUserQuota(u1.id, QuotaType.FreePlanV1); - const q3 = await storageQuota.getUserQuota(u1.id); - t.is(q3?.blobLimit, Quotas[4].configs.blobLimit, 'should be free plan'); + const q3 = await quotaManager.getUserQuota(u1.id); + t.is(q3?.blobLimit, Quotas[5].configs.blobLimit, 'should be free plan'); const quotas = await quota.getUserQuotas(u1.id); t.is(quotas.length, 3, 'should have 3 quotas'); @@ -100,3 +100,21 @@ test('should be able revert quota', async t => { t.is(quotas[1].activated, false, 'should be activated'); t.is(quotas[2].activated, true, 'should be activated'); }); + +test('should be able to check quota', async t => { + const { auth, quotaManager } = t.context; + const u1 = await auth.signUp('DarkSky', 'darksky@example.org', '123456'); + + const q1 = await quotaManager.getUserQuota(u1.id); + const freePlan = Quotas[5].configs; + t.assert(q1, 'should have quota'); + t.is(q1.blobLimit, freePlan.blobLimit, 'should be free plan'); + t.is(q1.storageQuota, freePlan.storageQuota, 'should be free plan'); + t.is(q1.historyPeriod, freePlan.historyPeriod, 'should be free plan'); + t.is(q1.memberLimit, freePlan.memberLimit, 'should be free plan'); + t.is( + q1.copilotActionLimit!, + freePlan.copilotActionLimit!, + 'should be free plan' + ); +}); diff --git a/packages/frontend/graphql/src/graphql/create-copilot-session.gql b/packages/frontend/graphql/src/graphql/create-copilot-session.gql new file mode 100644 index 0000000000..01056a7f2a --- /dev/null +++ b/packages/frontend/graphql/src/graphql/create-copilot-session.gql @@ -0,0 +1,3 @@ +mutation createCopilotSession($options: CreateChatSessionInput!) { + createCopilotSession(options: $options) +} 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 0000000000..f04af20b34 --- /dev/null +++ b/packages/frontend/graphql/src/graphql/get-copilot-anonymous-histories.gql @@ -0,0 +1,17 @@ +query getCopilotAnonymousHistories( + $workspaceId: String! + $docId: String + $options: QueryChatHistoriesInput +) { + copilotAnonymous(workspaceId: $workspaceId) { + histories(docId: $docId, options: $options) { + sessionId + tokens + messages { + role + content + attachments + } + } + } +} 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 0000000000..57c4f77a5a --- /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/get-copilot-histories.gql b/packages/frontend/graphql/src/graphql/get-copilot-histories.gql new file mode 100644 index 0000000000..75541dfd36 --- /dev/null +++ b/packages/frontend/graphql/src/graphql/get-copilot-histories.gql @@ -0,0 +1,19 @@ +query getCopilotHistories( + $workspaceId: String! + $docId: String + $options: QueryChatHistoriesInput +) { + currentUser { + copilot(workspaceId: $workspaceId) { + histories(docId: $docId, options: $options) { + sessionId + tokens + messages { + role + content + attachments + } + } + } + } +} diff --git a/packages/frontend/graphql/src/graphql/get-copilot-sessions.gql b/packages/frontend/graphql/src/graphql/get-copilot-sessions.gql new file mode 100644 index 0000000000..1c065f8d1d --- /dev/null +++ b/packages/frontend/graphql/src/graphql/get-copilot-sessions.gql @@ -0,0 +1,8 @@ +query getCopilotSessions($workspaceId: String!) { + currentUser { + copilot(workspaceId: $workspaceId) { + chats + actions + } + } +} diff --git a/packages/frontend/graphql/src/graphql/index.ts b/packages/frontend/graphql/src/graphql/index.ts index 4169577125..0672676298 100644 --- a/packages/frontend/graphql/src/graphql/index.ts +++ b/packages/frontend/graphql/src/graphql/index.ts @@ -144,6 +144,17 @@ mutation createCheckoutSession($input: CreateCheckoutSessionInput!) { }`, }; +export const createCopilotSessionMutation = { + id: 'createCopilotSessionMutation' as const, + operationName: 'createCopilotSession', + definitionName: 'createCopilotSession', + containsFile: false, + query: ` +mutation createCopilotSession($options: CreateChatSessionInput!) { + createCopilotSession(options: $options) +}`, +}; + export const createCustomerPortalMutation = { id: 'createCustomerPortalMutation' as const, operationName: 'createCustomerPortal', @@ -240,6 +251,80 @@ 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 + } + } + } +}`, +}; + +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', + definitionName: 'currentUser', + containsFile: false, + query: ` +query getCopilotHistories($workspaceId: String!, $docId: String, $options: QueryChatHistoriesInput) { + currentUser { + copilot(workspaceId: $workspaceId) { + histories(docId: $docId, options: $options) { + sessionId + tokens + messages { + role + content + attachments + } + } + } + } +}`, +}; + +export const getCopilotSessionsQuery = { + id: 'getCopilotSessionsQuery' as const, + operationName: 'getCopilotSessions', + definitionName: 'currentUser', + containsFile: false, + query: ` +query getCopilotSessions($workspaceId: String!) { + currentUser { + copilot(workspaceId: $workspaceId) { + chats + actions + } + } +}`, +}; + export const getCurrentUserQuery = { id: 'getCurrentUserQuery' as const, operationName: 'getCurrentUser', diff --git a/packages/frontend/graphql/src/schema.ts b/packages/frontend/graphql/src/schema.ts index c8fb983623..37acadc048 100644 --- a/packages/frontend/graphql/src/schema.ts +++ b/packages/frontend/graphql/src/schema.ts @@ -34,6 +34,14 @@ export interface Scalars { Upload: { input: File; output: File }; } +export interface CreateChatSessionInput { + action: Scalars['Boolean']['input']; + docId: Scalars['String']['input']; + model: Scalars['String']['input']; + promptName: Scalars['String']['input']; + workspaceId: Scalars['String']['input']; +} + export interface CreateCheckoutSessionInput { coupon: InputMaybe; idempotencyKey: Scalars['String']['input']; @@ -76,6 +84,13 @@ export enum PublicPageMode { Page = 'Page', } +export interface QueryChatHistoriesInput { + action: InputMaybe; + limit: InputMaybe; + sessionId: InputMaybe; + skip: InputMaybe; +} + export enum ServerDeploymentType { Affine = 'Affine', Selfhosted = 'Selfhosted', @@ -217,6 +232,15 @@ export type CreateCheckoutSessionMutation = { createCheckoutSession: string; }; +export type CreateCopilotSessionMutationVariables = Exact<{ + options: CreateChatSessionInput; +}>; + +export type CreateCopilotSessionMutation = { + __typename?: 'Mutation'; + createCopilotSession: string; +}; + export type CreateCustomerPortalMutationVariables = Exact<{ [key: string]: never; }>; @@ -309,6 +333,86 @@ 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; + }>; + }>; + }; +}; + +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; + options: InputMaybe; +}>; + +export type GetCopilotHistoriesQuery = { + __typename?: 'Query'; + currentUser: { + __typename?: 'UserType'; + copilot: { + __typename?: 'Copilot'; + histories: Array<{ + __typename?: 'CopilotHistories'; + sessionId: string; + tokens: number; + messages: Array<{ + __typename?: 'ChatMessage'; + role: string; + content: string; + attachments: Array | null; + }>; + }>; + }; + } | null; +}; + +export type GetCopilotSessionsQueryVariables = Exact<{ + workspaceId: Scalars['String']['input']; +}>; + +export type GetCopilotSessionsQuery = { + __typename?: 'Query'; + currentUser: { + __typename?: 'UserType'; + copilot: { + __typename?: 'Copilot'; + chats: Array; + actions: Array; + }; + } | null; +}; + export type GetCurrentUserQueryVariables = Exact<{ [key: string]: never }>; export type GetCurrentUserQuery = { @@ -953,6 +1057,26 @@ export type Queries = variables: EarlyAccessUsersQueryVariables; response: EarlyAccessUsersQuery; } + | { + name: 'getCopilotAnonymousHistoriesQuery'; + variables: GetCopilotAnonymousHistoriesQueryVariables; + response: GetCopilotAnonymousHistoriesQuery; + } + | { + name: 'getCopilotAnonymousSessionsQuery'; + variables: GetCopilotAnonymousSessionsQueryVariables; + response: GetCopilotAnonymousSessionsQuery; + } + | { + name: 'getCopilotHistoriesQuery'; + variables: GetCopilotHistoriesQueryVariables; + response: GetCopilotHistoriesQuery; + } + | { + name: 'getCopilotSessionsQuery'; + variables: GetCopilotSessionsQueryVariables; + response: GetCopilotSessionsQuery; + } | { name: 'getCurrentUserQuery'; variables: GetCurrentUserQueryVariables; @@ -1110,6 +1234,11 @@ export type Mutations = variables: CreateCheckoutSessionMutationVariables; response: CreateCheckoutSessionMutation; } + | { + name: 'createCopilotSessionMutation'; + variables: CreateCopilotSessionMutationVariables; + response: CreateCopilotSessionMutation; + } | { name: 'createCustomerPortalMutation'; variables: CreateCustomerPortalMutationVariables;