diff --git a/config.example.js b/config.example.js index 69a01bd2..b64a5c8a 100644 --- a/config.example.js +++ b/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 }; diff --git a/src/app/RirikoAI-NLP.js b/src/app/RirikoAI-NLP.js index d3b19052..a370d866 100644 --- a/src/app/RirikoAI-NLP.js +++ b/src/app/RirikoAI-NLP.js @@ -1,17 +1,14 @@ /** * @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, @@ -19,6 +16,9 @@ const { 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 diff --git a/src/commands/slash/StableDiffusion/imagine.js b/src/commands/slash/StableDiffusion/imagine.js index b9d2e44c..1af8e162 100644 --- a/src/commands/slash/StableDiffusion/imagine.js +++ b/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"); diff --git a/src/helpers/commandUsage.js b/src/helpers/commandUsage.js new file mode 100644 index 00000000..9b3e7515 --- /dev/null +++ b/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, +};