mirror of
https://github.com/toeverything/AFFiNE.git
synced 2026-02-12 12:28:42 +00:00
chore: enable ai feature in dev (#6618)
This commit is contained in:
@@ -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
|
||||
}
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -231,6 +231,7 @@ export class FeatureService {
|
||||
feature,
|
||||
type: FeatureKind.Feature,
|
||||
},
|
||||
OR: [{ expiredAt: null }, { expiredAt: { gt: new Date() } }],
|
||||
},
|
||||
})
|
||||
.then(count => count > 0);
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user