chore: enable ai feature in dev (#6618)

This commit is contained in:
darkskygit
2024-04-19 08:12:52 +00:00
parent 6d5dbbd7f4
commit 97669acb40
5 changed files with 35 additions and 22 deletions

View File

@@ -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
}

View File

@@ -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<FeatureType[]> {
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<FeatureType[]> {
const features = await this.feature.getActivatedUserFeatures(userId);
return features.map(f => f.feature.name);
}
}

View File

@@ -231,6 +231,7 @@ export class FeatureService {
feature,
type: FeatureKind.Feature,
},
OR: [{ expiredAt: null }, { expiredAt: { gt: new Date() } }],
},
})
.then(count => count > 0);

View File

@@ -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(

View File

@@ -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;
}