From 97669acb402bbce2e96be82536919d8be6fac60d Mon Sep 17 00:00:00 2001 From: darkskygit Date: Fri, 19 Apr 2024 08:12:52 +0000 Subject: [PATCH] chore: enable ai feature in dev (#6618) --- .../backend/server/src/core/auth/service.ts | 1 + .../server/src/core/features/management.ts | 32 +++++++++++++++---- .../server/src/core/features/service.ts | 1 + .../backend/server/src/core/quota/service.ts | 15 +++------ .../server/src/plugins/copilot/session.ts | 8 ++--- 5 files changed, 35 insertions(+), 22 deletions(-) diff --git a/packages/backend/server/src/core/auth/service.ts b/packages/backend/server/src/core/auth/service.ts index 1c4eacb205..639b2a5f9b 100644 --- a/packages/backend/server/src/core/auth/service.ts +++ b/packages/backend/server/src/core/auth/service.ts @@ -88,6 +88,7 @@ export class AuthService implements OnApplicationBootstrap { }); } await this.quota.switchUserQuota(devUser.id, QuotaType.ProPlanV1); + await this.feature.addCopilot(devUser.id); } catch (e) { // ignore } diff --git a/packages/backend/server/src/core/features/management.ts b/packages/backend/server/src/core/features/management.ts index 756b125e23..b0c23bd2f3 100644 --- a/packages/backend/server/src/core/features/management.ts +++ b/packages/backend/server/src/core/features/management.ts @@ -108,6 +108,32 @@ export class FeatureManagementService { } } + // ======== CopilotFeature ======== + async addCopilot(userId: string, reason = 'Copilot plan user') { + return this.feature.addUserFeature( + userId, + FeatureType.UnlimitedCopilot, + reason + ); + } + + async removeCopilot(userId: string) { + return this.feature.removeUserFeature(userId, FeatureType.UnlimitedCopilot); + } + + async isCopilotUser(userId: string) { + return await this.feature.hasUserFeature( + userId, + FeatureType.UnlimitedCopilot + ); + } + + // ======== User Feature ======== + async getActivatedUserFeatures(userId: string): Promise { + const features = await this.feature.getActivatedUserFeatures(userId); + return features.map(f => f.feature.name); + } + // ======== Workspace Feature ======== async addWorkspaceFeatures( workspaceId: string, @@ -147,10 +173,4 @@ export class FeatureManagementService { async listFeatureWorkspaces(feature: FeatureType) { return this.feature.listFeatureWorkspaces(feature); } - - // ======== User Feature ======== - async getActivatedUserFeatures(userId: string): Promise { - const features = await this.feature.getActivatedUserFeatures(userId); - return features.map(f => f.feature.name); - } } diff --git a/packages/backend/server/src/core/features/service.ts b/packages/backend/server/src/core/features/service.ts index d59dc3d548..4cb4a62da6 100644 --- a/packages/backend/server/src/core/features/service.ts +++ b/packages/backend/server/src/core/features/service.ts @@ -231,6 +231,7 @@ export class FeatureService { feature, type: FeatureKind.Feature, }, + OR: [{ expiredAt: null }, { expiredAt: { gt: new Date() } }], }, }) .then(count => count > 0); diff --git a/packages/backend/server/src/core/quota/service.ts b/packages/backend/server/src/core/quota/service.ts index 03b8022800..b19a2729d1 100644 --- a/packages/backend/server/src/core/quota/service.ts +++ b/packages/backend/server/src/core/quota/service.ts @@ -4,7 +4,7 @@ import { PrismaClient } from '@prisma/client'; import type { EventPayload } from '../../fundamentals'; import { OnEvent, PrismaTransaction } from '../../fundamentals'; import { SubscriptionPlan } from '../../plugins/payment/types'; -import { FeatureKind, FeatureService, FeatureType } from '../features'; +import { FeatureKind, FeatureManagementService } from '../features'; import { QuotaConfig } from './quota'; import { QuotaType } from './types'; @@ -12,7 +12,7 @@ import { QuotaType } from './types'; export class QuotaService { constructor( private readonly prisma: PrismaClient, - private readonly feature: FeatureService + private readonly feature: FeatureManagementService ) {} // get activated user quota @@ -167,11 +167,7 @@ export class QuotaService { }: EventPayload<'user.subscription.activated'>) { switch (plan) { case SubscriptionPlan.AI: - await this.feature.addUserFeature( - userId, - FeatureType.UnlimitedCopilot, - 'subscription activated' - ); + await this.feature.addCopilot(userId, 'subscription activated'); break; case SubscriptionPlan.Pro: await this.switchUserQuota( @@ -192,10 +188,7 @@ export class QuotaService { }: EventPayload<'user.subscription.canceled'>) { switch (plan) { case SubscriptionPlan.AI: - await this.feature.removeUserFeature( - userId, - FeatureType.UnlimitedCopilot - ); + await this.feature.removeCopilot(userId); break; case SubscriptionPlan.Pro: await this.switchUserQuota( diff --git a/packages/backend/server/src/plugins/copilot/session.ts b/packages/backend/server/src/plugins/copilot/session.ts index 4d100c7e34..367f3b984d 100644 --- a/packages/backend/server/src/plugins/copilot/session.ts +++ b/packages/backend/server/src/plugins/copilot/session.ts @@ -3,7 +3,7 @@ import { randomUUID } from 'node:crypto'; import { Injectable, Logger } from '@nestjs/common'; import { AiPromptRole, PrismaClient } from '@prisma/client'; -import { FeatureManagementService, FeatureType } from '../../core/features'; +import { FeatureManagementService } from '../../core/features'; import { QuotaService } from '../../core/quota'; import { PaymentRequiredException } from '../../fundamentals'; import { ChatMessageCache } from './message'; @@ -379,12 +379,10 @@ export class ChatSessionService { } async getQuota(userId: string) { - const hasCopilotFeature = await this.feature - .getActivatedUserFeatures(userId) - .then(f => f.includes(FeatureType.UnlimitedCopilot)); + const isCopilotUser = await this.feature.isCopilotUser(userId); let limit: number | undefined; - if (!hasCopilotFeature) { + if (!isCopilotUser) { const quota = await this.quota.getUserQuota(userId); limit = quota.feature.copilotActionLimit; }