.

src/chat/chat.module.ts

@@ -4,9 +4,10 @@ import { ChatController } from './chat.controller'
 import { TypeOrmModule } from '@nestjs/typeorm'
 import { ChatHistory } from './entities/chat.entity'
 import { TokenUsage } from './entities/token-usage.entity'
+import { MembershipModule } from '../membership/membership.module'
 
 @Module({
-    imports: [TypeOrmModule.forFeature([ChatHistory, TokenUsage])],
+    imports: [TypeOrmModule.forFeature([ChatHistory, TokenUsage]),MembershipModule],
     providers: [ChatService],
     controllers: [ChatController]
 })

src/chat/chat.service.ts

@@ -1,4 +1,4 @@
-import { Injectable, Logger } from '@nestjs/common'
+import { ForbiddenException, Injectable, Logger } from '@nestjs/common'
 import { Observable } from 'rxjs'
 import { ChatGPTAPI, ChatMessage } from '../chatapi'
 import type { RequestProps } from './types'
@@ -8,6 +8,8 @@ import { ChatHistory } from './entities/chat.entity'
 import { InjectRepository } from '@nestjs/typeorm'
 import { TokenUsage } from './entities/token-usage.entity'
 import { format } from 'date-fns'
+import { MembershipService } from '../membership/membership.service'
+import { MemberType } from '../membership/entities/membership.entity'
 
 @Injectable()
 export class ChatService {
@@ -15,7 +17,8 @@ export class ChatService {
         @InjectRepository(ChatHistory)
         private readonly chatHistoryRepository: Repository<ChatHistory>,
         @InjectRepository(TokenUsage)
-        private readonly tokenUsageRepository: Repository<TokenUsage>
+        private readonly tokenUsageRepository: Repository<TokenUsage>,
+        private readonly membershipService: MembershipService
     ) {}
 
     public chat(req, res): Observable<any> {
@@ -51,7 +54,16 @@ export class ChatService {
 
     public async chat1(req, res) {
         res.setHeader('Content-type', 'application/octet-stream')
-
+        const membership = await this.membershipService.getMembership(req.user.id)
+        if (!membership) {
+            throw new ForbiddenException('请先成为会员')
+        }
+        if (membership.memberType == MemberType.Trial && membership.tokenLeft <= 0) {
+            throw new ForbiddenException('您的试用额度已用完，请升级会员')
+        }
+        if (membership.isExpired) {
+            throw new ForbiddenException('您的会员已过期，请即时续费')
+        }
         try {
             const promptTime = new Date()
             const { prompt, options = {}, systemMessage, temperature, top_p } = req.body as RequestProps
@@ -119,5 +131,7 @@ export class ChatService {
         } else {
             await this.tokenUsageRepository.save(new TokenUsage({ userId, usage, date }))
         }
+
+        this.membershipService.saveUsage(userId, usage)
     }
 }

src/membership/membership.service.ts

@@ -1,7 +1,7 @@
 import { Injectable, InternalServerErrorException, NotFoundException } from '@nestjs/common'
 import { InjectRepository } from '@nestjs/typeorm'
 import { MemberType, Membership } from './entities/membership.entity'
-import { Repository } from 'typeorm'
+import { Repository, LessThan } from 'typeorm'
 import { MemberPlan } from './entities/memberPlan.entity'
 import { addDays } from 'date-fns'
 import { MemberPlanDto } from './dto/memberPlan.dto'
@@ -128,4 +128,22 @@ export class MembershipService {
         await this.memberOrderRepository.save(order)
         await this.renewMembership(order.userId, order.planId)
     }
+
+    async saveUsage(userId: number, usage: number) {
+        const membership = await this.memberShipRepository.findOneBy({
+            userId: userId
+        })
+        if (!membership) {
+            throw new NotFoundException(`Membership #${userId} not found`)
+        }
+        membership.tokenLeft = membership.tokenLeft - usage
+        await this.memberShipRepository.update(
+            {
+                userId
+            },
+            {
+                tokenLeft: () => `tokenLeft - ${usage}`
+            }
+        )
+    }
 }

src/weixin/weixin.controller.ts

@@ -1,7 +1,7 @@
 import { Controller, Get, Query } from '@nestjs/common'
 import { WeixinService } from './weixin.service'
 import { Public } from 'src/auth/public.decorator'
-import { ApiTags } from '@nestjs/swagger'
+import { ApiQuery, ApiTags } from '@nestjs/swagger'
 
 @ApiTags('weixin')
 @Controller('/weixin')
@@ -15,6 +15,7 @@ export class WeixinController {
 
     @Public()
     @Get('/redirectUrl')
+    @ApiQuery({ name: 'url' })
     public getRedirectUrl(@Query() { url, state }) {
         return this.weixinService.getRedirectUrl(url, state)
     }

src/weixin/weixin.service.spec.ts

@@ -23,7 +23,7 @@ describe('WeixinService', () => {
     describe('getRedirectUrl', () => {
         it('getRedirectUrl', async () => {
             expect(weixinService).toBeDefined()
-            console.log(weixinService.getRedirectUrl())
+            console.log(weixinService.getRedirectUrl(''))
         })
     })
 })