Implement daily limit for chatbot and stablediffusion

config.example.js

@@ -43,15 +43,22 @@ module.exports = {
       stableDiffusion2_1:
         "stability-ai/stable-diffusion:db21e45d3f7023abc2a46ee38a23973f6dce16bb082a930b0c49861f96d1e5bf",
     },
+
+    // These are all the StableDiffusion options. If you have problems generating images, try changing these (follow the settings in
+    // the model's page (https://replicate.com/cjwbw/anything-v3.0, https://replicate.com/cjwbw/eimis_anime_diffusion, etc.)
     NegativePrompt:
       "lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, " +
       "low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry, artist name",
     NumInferenceSteps: 10,
     Width: 512,
-    Height: 640,
+    Height: 512,
     GuidanceScale: 7,
     Scheduler: "DPMSolverMultistep",
     DisableSafetyCheck: false,
+
+    // Enable daily limit for each member. Put numbers of uses per day in the DailyLimit config below (e.g: 15, 30)
+    // Leave this to false if you don't want to enable daily limit.
+    DailyLimit: false,
   },
 
   // ------------------------------------------ Ririko AI / Chatbot ----------------------------------------------------
@@ -92,6 +99,10 @@ module.exports = {
       "Human: Play https://youtube.com/watch?v=Lh63pBzylFg",
       "Friend: Sure! Now Playing 🎵 https://youtube.com/watch?v=Lh63pBzylFg 🎵",
     ],
+
+    // Enable daily limit for each member. Put numbers of uses per day in the DailyLimit config below (e.g: 15, 30)
+    // Leave this to false if you don't want to enable daily limit.
+    DailyLimit: false,
   },
 
   // ------------------------------------------------ Database ---------------------------------------------------------
@@ -263,5 +274,5 @@ module.exports = {
     LogDir: "logs",
   },
 
-  VERSION: "5", // DO NOT TOUCH
+  VERSION: "6", // DO NOT TOUCH
 };

src/app/RirikoAI-NLP.js

@@ -1,24 +1,24 @@
 /**
  * @author earnestangel https://github.com/RirikoAI/RirikoBot
  */
-const { Configuration, OpenAIApi } = require("openai");
 const colors = require("colors");
-const config = require("config");
 
 const { OpenAIProvider } = require("./Providers/AI/OpenAIProvider");
 const { NLPCloudProvider } = require("app/Providers/AI/NLPCloudProvider");
 
 const getconfig = require("helpers/getconfig");
-const { AIProvider } = require("helpers/getconfig");
-const { AIPersonality } = require("helpers/getconfig");
-const { AIPrompts } = require("../helpers/getconfig");
+const { AIProvider, AIPersonality, AIPrompts } = require("helpers/getconfig");
+
 const {
   findChatHistory,
   addChatHistory,
   updateChatHistory,
   deleteChatHistory,
 } = require("./Schemas/ChatHistory");
 
+const { getAndIncrementUsageCount } = require("helpers/commandUsage");
+const { AI } = require("config");
+
 /**
  * Now, this is going to be an awesome AI that can remember past conversations by saving it into the
  * "brain"
@@ -114,6 +114,17 @@ class RirikoAINLP {
    */
   async handleMessage(message) {
     if (message.content.substring(0, 1) === this.prefix) {
+      if (AI.DailyLimit !== false)
+        try {
+          const usageCount = await getAndIncrementUsageCount(
+            message.member.user.id,
+            AI.DailyLimit,
+            "ai"
+          );
+        } catch (e) {
+          return await message.reply(e.message);
+        }
+
       await message.channel.sendTyping();
 
       const prompt = message.content.substring(1); //remove the prefix from the message

src/commands/slash/StableDiffusion/imagine.js

@@ -15,6 +15,7 @@ import language from "languages/en";
 import { StableDiffusion } from "config";
 import axios, { get } from "axios";
 import { AttachmentBuilder } from "discord.js";
+import { getAndIncrementUsageCount } from "helpers/commandUsage";
 
 module.exports = {
   name: "imagine",
@@ -34,7 +35,19 @@ module.exports = {
 };
 
 async function imagineCommand(client, interaction, args, prefix) {
+  if (StableDiffusion.DailyLimit !== false)
+    try {
+      const usageCount = await getAndIncrementUsageCount(
+        interaction.member.user.id,
+        StableDiffusion.DailyLimit,
+        "imagine"
+      );
+    } catch (e) {
+      return await interaction.reply(e.message);
+    }
+
   await interaction.deferReply();
+
   const replicate = createReplicateClient();
 
   const userPrompt = interaction.options.getString("prompt");

src/helpers/commandUsage.js

@@ -0,0 +1,36 @@
+const mongoose = require("mongoose");
+
+const commandUsageSchema = new mongoose.Schema({
+  memberId: { type: String, required: true },
+  commandName: { type: String, required: true },
+  usageCount: { type: Number, default: 0 },
+  createdAt: { type: Date, default: Date.now, expires: 24 * 60 * 60 }, // Set expiration time to 24 hours
+});
+
+const CommandUsageModel = mongoose.model("CommandUsage", commandUsageSchema);
+
+async function getAndIncrementUsageCount(memberId, limit, commandName) {
+  let commandUsage = await CommandUsageModel.findOne({
+    memberId: memberId,
+    commandName: commandName,
+  });
+
+  if (commandUsage?.usageCount >= limit) {
+    throw new Error(
+      `You have reached the usage limit for this command. [${limit} uses per day]`
+    );
+  }
+
+  commandUsage = await CommandUsageModel.findOneAndUpdate(
+    { memberId: memberId, commandName: commandName },
+    { $inc: { usageCount: 1 }, lastUsageTimestamp: new Date() },
+    { upsert: true, new: true }
+  );
+
+  return commandUsage.usageCount;
+}
+
+module.exports = {
+  CommandUsageModel,
+  getAndIncrementUsageCount,
+};