Add Subscriptions for the Chatbot Feature

package.json

@@ -24,12 +24,14 @@
     "migration:revert": "npm run typeorm migration:revert -- -d src/config/typeorm.config.ts"
   },
   "dependencies": {
+    "@golevelup/nestjs-stripe": "^0.9.2",
     "@google/generative-ai": "^0.21.0",
     "@langchain/core": "^0.3.40",
     "@langchain/google-genai": "^0.1.8",
     "@langchain/groq": "^0.1.3",
     "@nestjs-modules/mailer": "^2.0.2",
     "@nestjs/common": "^10.4.4",
+    "@nestjs/config": "^4.0.0",
     "@nestjs/core": "^10.4.4",
     "@nestjs/jwt": "^10.2.0",
     "@nestjs/mapped-types": "^2.0.5",
@@ -49,6 +51,7 @@
     "passport-local": "^1.0.0",
     "reflect-metadata": "^0.2.2",
     "rxjs": "^7.8.1",
+    "stripe": "^17.6.0",
     "typeorm": "^0.3.20",
     "uuid": "^9.0.1"
   },

src/app.module.ts

@@ -27,7 +27,9 @@ import { DefinitionReportsModule } from './definition-reports/definition-reports
 import { DefinitionReport } from './typeorm/entities/definition-report';
 import { RagModule } from './rag/rag.module';
 import { ModelService } from './rag/services/model/model.service';
-import { VectorStoreService } from './rag/services/vector-store/vector-store.service';
+import { SubscriptionsModule } from "./subscriptions/subscriptions.module";
+import { ConfigModule } from "@nestjs/config";
+
 
 dotenv.config({ path: './safe/.env' });
 
@@ -81,6 +83,9 @@ dotenv.config({ path: './safe/.env' });
         },
       }),
     }),
+    ConfigModule.forRoot({
+      isGlobal: true,
+    }),
     PassportModule.register({ session: true }),
     UsersModule,
     CountriesModule,
@@ -91,6 +96,7 @@ dotenv.config({ path: './safe/.env' });
     AuthModule,
     DefinitionReportsModule,
     RagModule,
+    SubscriptionsModule,
   ],
   controllers: [AppController, AuthController, DefinitionReportsController],
   providers: [AppService, ModelService],

src/authentication/controllers/authentication/auth.controller.ts

@@ -10,7 +10,7 @@ import {
 import { Request as ExpressRequest, Response, NextFunction } from 'express';
 import { Session } from 'express-session';
 import { Throttle } from '@nestjs/throttler';
-import { AuthenticatedGuard, LocalAuthGuard } from '../../../utils/local.guard';
+import { AuthenticatedGuard, LocalAuthGuard } from '../../../utils/guards/local.guard';
 
 @Controller('auth')
 export class AuthController {

src/definition-likes-dislikes/controllers/definition-likes-dislikes/definition-likes-dislikes.controller.ts

@@ -9,7 +9,7 @@ import {
   Logger, Body
 } from "@nestjs/common";
 import { DefinitionLikesDislikesService } from '../../services/definition-likes-dislikes/definition-likes-dislikes.service';
-import { AuthenticatedGuard, LocalAuthGuard } from '../../../utils/local.guard';
+import { AuthenticatedGuard, LocalAuthGuard } from '../../../utils/guards/local.guard';
 import { RequestType } from 'express-serve-static-core';
 import { Throttle } from '@nestjs/throttler';
 import { UserRequest } from '../../../utils/types';

src/definition-reports/controllers/definition-reports/definition-reports.controller.ts

@@ -13,7 +13,7 @@ import {
 import { DefinitionReportsService } from '../../services/definition-reports/definition-reports.service';
 import { CreateDefinitionReportDto } from '../../dtos/create-definition-report.dto';
 import { UpdateDefinitionReportDto } from '../../dtos/update-definition-report.dto';
-import { AuthenticatedGuard } from '../../../utils/local.guard';
+import { AuthenticatedGuard } from '../../../utils/guards/local.guard';
 import { RequestType } from 'express-serve-static-core';
 import { DefinitionReport } from '../../../typeorm/entities/definition-report';
 

src/definitions/controllers/definitions/definitions.controller.ts

@@ -16,7 +16,7 @@ import { CreateDefinitionDto } from '../../dtos/create-definition.dto';
 import { UpdateDefinitionDto } from '../../dtos/update-definition.dto';
 import { DeleteResult, UpdateResult } from 'typeorm';
 import { Throttle } from '@nestjs/throttler';
-import { AuthenticatedGuard } from '../../../utils/local.guard';
+import { AuthenticatedGuard } from '../../../utils/guards/local.guard';
 import { RequestType } from 'express-serve-static-core';
 import { Country } from '../../../typeorm/entities/country';
 

src/rag/controllers/rag/rag.controller.ts

@@ -1,15 +1,31 @@
-import { Body, Controller, HttpException, HttpStatus, Post } from "@nestjs/common";
+import { Body, Controller, HttpException, HttpStatus, Post, Req, UseGuards } from "@nestjs/common";
 import { RagRequest, RagResponse } from "../../dtos/rag.dto";
 import { RagService } from "../../services/rag/rag.service";
+import { AuthenticatedGuard } from "../../../utils/guards/local.guard";
+import { SubscriptionGuard } from "../../../utils/guards/subscription.guard";
+import { RequestType } from "express-serve-static-core";
 
 @Controller('rag')
 export class RagController {
   constructor(private readonly ragService: RagService) {}
 
   @Post()
-  async handleQuery(@Body() request: RagRequest): Promise<RagResponse> {
+  async handleQuery(@Body() request: RagRequest, @Req() req: RequestType): Promise<RagResponse> {
     try {
-      return await this.ragService.processQuery(request);
+      return await this.ragService.processQuery(request, req.user.id);
+    } catch (error) {
+      throw new HttpException(
+        'Failed to process RAG query: ' + error.message,
+        HttpStatus.INTERNAL_SERVER_ERROR
+      );
+    }
+  }
+
+  @Post('premium')
+  @UseGuards(AuthenticatedGuard, SubscriptionGuard)
+  async handlePremiumQuery(@Body() request: RagRequest, @Req() req: RequestType): Promise<RagResponse> {
+    try {
+      return await this.ragService.processQuery(request, req.user.id);
     } catch (error) {
       throw new HttpException(
         'Failed to process RAG query: ' + error.message,

src/rag/dtos/rag.dto.ts

@@ -1,4 +1,4 @@
-export type AiModel = 'groq' | 'gemini' | 'gpt4' | 'mistral';
+export type AiModel = 'groq' | 'gemini' | 'groq-3'| 'gemini-pro' | 'gpt4' | 'mistral';
 
 export class RagRequest {
   query: string;

src/rag/rag.module.ts

@@ -6,10 +6,12 @@ import { Word } from "../typeorm/entities/word";
 import { Definition } from "../typeorm/entities/definition";
 import { ModelService } from "./services/model/model.service";
 import { VectorStoreService } from "./services/vector-store/vector-store.service";
+import { SubscriptionsModule } from "../subscriptions/subscriptions.module";
 
 @Module({
   imports: [
     TypeOrmModule.forFeature([Word, Definition]),
+    SubscriptionsModule
   ],
   controllers: [RagController],
   providers: [RagService, ModelService, VectorStoreService]

src/rag/services/model/model.service.ts

@@ -43,6 +43,8 @@ export class ModelService {
           return this.generateGroqResponse(fullPrompt);
         case 'gemini':
           return this.generateGeminiResponse(fullPrompt);
+        case 'groq-3':
+        case 'gemini-pro':
         case 'gpt4':
         case 'mistral':
           throw new Error(`Model ${model} not implemented yet`);

src/rag/services/rag/rag.service.ts

@@ -1,4 +1,4 @@
-import { Injectable } from '@nestjs/common';
+import { ForbiddenException, Injectable } from "@nestjs/common";
 import { InjectRepository } from "@nestjs/typeorm";
 import { Word } from "../../../typeorm/entities/word";
 import { Like, Repository } from "typeorm";
@@ -7,6 +7,8 @@ import { In } from 'typeorm';
 import { RagRequest, RagResponse } from "../../dtos/rag.dto";
 import { VectorStoreService } from "../vector-store/vector-store.service";
 import { ModelService } from "../model/model.service";
+import { SubscriptionsService } from "../../../subscriptions/services/subscriptions/subscriptions.service";
+import { Subscription } from "../../../typeorm/entities/subscription";
 
 interface CitedWord {
   id: number;
@@ -23,12 +25,24 @@ export class RagService {
     @InjectRepository(Word)
     private wordRepository: Repository<Word>,
     @InjectRepository(Definition)
-    private definitionRepository: Repository<Definition>,
     private modelService: ModelService,
     private vectorStore: VectorStoreService,
+    private subscriptionsService: SubscriptionsService
   ) {}
 
-  async processQuery(request: RagRequest): Promise<RagResponse> {
+  private readonly modelTiers = {
+    free: ['groq', 'gemini'],
+    premium: ['groq-3', 'gemini-pro', 'gpt4', 'mistral']
+  };
+
+  async processQuery(request: RagRequest, userId: number): Promise<RagResponse> {
+    const subscription: Subscription = await this.subscriptionsService.getActiveSubscription(userId);
+    const tier: 'premium' | 'free' = subscription?.status === 'active' ? 'premium' : 'free';
+
+    if (!this.modelTiers[tier].includes(request.model)) {
+      throw new ForbiddenException(`Model ${request.model} requires premium subscription`);
+    }
+
     // 1. Extract words from the query
     const words: string[] = this.extractWords(request.query);
 

src/subscriptions/controllers/subscriptions/subscriptions.controller.spec.ts

@@ -0,0 +1,18 @@
+import { Test, TestingModule } from '@nestjs/testing';
+import { SubscriptionsController } from "./subscriptions.controller";
+
+describe('SubscriptionsController', () => {
+  let controller: SubscriptionsController;
+
+  beforeEach(async () => {
+    const module: TestingModule = await Test.createTestingModule({
+      controllers: [SubscriptionsController],
+    }).compile();
+
+    controller = module.get<SubscriptionsController>(SubscriptionsController);
+  });
+
+  it('should be defined', () => {
+    expect(controller).toBeDefined();
+  });
+});