Refactor the main bot execution out of index.js
This commit is contained in:
parent
1c49e8b730
commit
cb304f522b
2 changed files with 150 additions and 118 deletions
148
src/execution.ts
Normal file
148
src/execution.ts
Normal file
|
|
@ -0,0 +1,148 @@
|
||||||
|
import DiscordApi from "discord.js";
|
||||||
|
|
||||||
|
import { database, openai } from "./index";
|
||||||
|
import Moderation from "./moderation";
|
||||||
|
import config from "./config";
|
||||||
|
import toOpenAIMessages from "./toOpenAIMessages";
|
||||||
|
|
||||||
|
/** Stores the queue requests on the channels. */
|
||||||
|
const channelsRunning: DiscordApi.Collection<string, DiscordApi.Message[]> = new DiscordApi.Collection();
|
||||||
|
|
||||||
|
/**
|
||||||
|
* gets user remaining limit (or lack of it)
|
||||||
|
* @param user the user to check
|
||||||
|
* @param requestTimestamp the timestamp of the user request
|
||||||
|
* @returns object containing the limit and remaining usage or `false` if there is no limit
|
||||||
|
*/
|
||||||
|
async function getUserLimit(user: string | {id: string}, requestTimestamp: Date) {
|
||||||
|
const userId: string = typeof user === "string" ? user : user.id;
|
||||||
|
|
||||||
|
const userLimits = await database.limits.findUnique({
|
||||||
|
where: {
|
||||||
|
user: BigInt(userId)
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
if (userLimits?.vip) return false;
|
||||||
|
|
||||||
|
const usedLimit = (await database.usage.count({
|
||||||
|
select: { _all: true },
|
||||||
|
where: {
|
||||||
|
user: BigInt(userId),
|
||||||
|
timestamp: {
|
||||||
|
gte: new Date(requestTimestamp.getTime() - 1000 * 60 * 60 * 24 /* 24 hours */)
|
||||||
|
}
|
||||||
|
},
|
||||||
|
}))._all;
|
||||||
|
|
||||||
|
if (!userLimits || !userLimits.limit) return {limit: 25, remaining: 25 - usedLimit};
|
||||||
|
|
||||||
|
return {limit: userLimits.limit, remaining: userLimits.limit - usedLimit};
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Check and queues up the request and runs it if there is nothing in queue.
|
||||||
|
* @param message the message to check and queue
|
||||||
|
*/
|
||||||
|
export async function queueRequest(message: DiscordApi.Message) {
|
||||||
|
|
||||||
|
const userLimit = await getUserLimit(message.author, message.createdAt);
|
||||||
|
|
||||||
|
if (userLimit !== false && userLimit.remaining <= 0) {
|
||||||
|
message.react("🛑");
|
||||||
|
message.author.dmChannel?.send({
|
||||||
|
embeds: [{
|
||||||
|
color: 0xff0000,
|
||||||
|
description: `You've used up your message limit for today, ${userLimit.limit} requrests in last 24 hours`
|
||||||
|
}]
|
||||||
|
});
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
const messagesForChannel = channelsRunning.ensure(
|
||||||
|
message.channelId,
|
||||||
|
() => { return [] as DiscordApi.Message[]; },
|
||||||
|
);
|
||||||
|
const shouldStart = messagesForChannel.length == 0;
|
||||||
|
messagesForChannel.push(message);
|
||||||
|
if (shouldStart)
|
||||||
|
executeFromQueue(message.channelId);
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Executes the queue for the channel
|
||||||
|
* @param channel the channel to run the queue for
|
||||||
|
*/
|
||||||
|
async function executeFromQueue(channel: string) {
|
||||||
|
const channelQueue = channelsRunning.get(channel) as DiscordApi.Message[];
|
||||||
|
const message = channelQueue.at(0) as DiscordApi.Message;
|
||||||
|
|
||||||
|
try {
|
||||||
|
let messages: DiscordApi.Collection<string, DiscordApi.Message> = await message.channel.messages.fetch({ limit: config.limits.messages, cache: false });
|
||||||
|
|
||||||
|
messages = messages.filter(m => message.createdTimestamp - m.createdTimestamp < config.limits.time );
|
||||||
|
|
||||||
|
messages.forEach(m => Moderation.checkMessage(m));
|
||||||
|
|
||||||
|
message.channel.sendTyping();
|
||||||
|
const answer = await openai.createChatCompletion({
|
||||||
|
...config.chatCompletionConfig,
|
||||||
|
messages: toOpenAIMessages(messages),
|
||||||
|
});
|
||||||
|
|
||||||
|
const usage = answer.data.usage;
|
||||||
|
if (usage != undefined) {
|
||||||
|
const channelName: string = message.inGuild() ? `${message.channel.name} (${message.guild.name})` : `@${message.author.tag}`;
|
||||||
|
console.log(`Used ${usage.total_tokens} (${usage.prompt_tokens} + ${usage.completion_tokens}) tokens for ${message.author.tag} (${message.author.id}) in #${channelName}`);
|
||||||
|
|
||||||
|
database.usage.create({
|
||||||
|
data: {
|
||||||
|
timestamp: message.createdAt,
|
||||||
|
user: BigInt(message.author.id),
|
||||||
|
channel: BigInt(message.channelId),
|
||||||
|
guild: message.guildId ? BigInt(message.guildId) : null,
|
||||||
|
usageReguest: usage.prompt_tokens,
|
||||||
|
usageResponse: usage.completion_tokens
|
||||||
|
}
|
||||||
|
}).catch((e => {
|
||||||
|
console.error("Failed to push to a database");
|
||||||
|
console.error(e);
|
||||||
|
}));
|
||||||
|
}
|
||||||
|
|
||||||
|
const answerContent = answer.data.choices[0].message?.content;
|
||||||
|
|
||||||
|
if (answerContent != undefined && answerContent != "") {
|
||||||
|
const response = message.reply({
|
||||||
|
content: answerContent,
|
||||||
|
allowedMentions: {
|
||||||
|
repliedUser: false,
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
Moderation.checkMessage(await response);
|
||||||
|
}
|
||||||
|
else {
|
||||||
|
message.react("😶");
|
||||||
|
}
|
||||||
|
} catch (e) {
|
||||||
|
console.error(`Error ocurred while handling chat completion request (${(e as object).constructor.name}):`);
|
||||||
|
console.error(e);
|
||||||
|
|
||||||
|
message.reply({
|
||||||
|
embeds: [{
|
||||||
|
color: 0xff0000,
|
||||||
|
description: "Something bad happened! :frowning:"
|
||||||
|
}],
|
||||||
|
allowedMentions: {
|
||||||
|
repliedUser: false,
|
||||||
|
}
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
channelQueue.shift();
|
||||||
|
if (channelQueue.length == 0)
|
||||||
|
channelsRunning.delete(channel);
|
||||||
|
else
|
||||||
|
executeFromQueue(channel);
|
||||||
|
}
|
||||||
120
src/index.ts
120
src/index.ts
|
|
@ -3,8 +3,7 @@ import { Configuration as OpenAIApiConfiguration, OpenAIApi } from "openai";
|
||||||
import { PrismaClient } from "@prisma/client";
|
import { PrismaClient } from "@prisma/client";
|
||||||
|
|
||||||
import config from "./config";
|
import config from "./config";
|
||||||
import toOpenAIMessages from "./toOpenAIMessages";
|
import { queueRequest } from "./execution";
|
||||||
import Moderation from "./moderation";
|
|
||||||
|
|
||||||
const discord = new DiscordApi.Client({
|
const discord = new DiscordApi.Client({
|
||||||
intents: [
|
intents: [
|
||||||
|
|
@ -24,126 +23,11 @@ discord.on("ready", async event => {
|
||||||
console.log(`Connected to Discord as ${event.user.tag} (${event.user.id})`);
|
console.log(`Connected to Discord as ${event.user.tag} (${event.user.id})`);
|
||||||
});
|
});
|
||||||
|
|
||||||
const channelsRunning: DiscordApi.Collection<string, DiscordApi.Message[]> = new DiscordApi.Collection();
|
|
||||||
|
|
||||||
discord.on("messageCreate", async message => {
|
discord.on("messageCreate", async message => {
|
||||||
if (message.author.bot) return;
|
if (message.author.bot) return;
|
||||||
if (!message.mentions.has(message.client.user)) return;
|
if (!message.mentions.has(message.client.user)) return;
|
||||||
|
|
||||||
const userLimits = await database.limits.findUnique({
|
queueRequest(message);
|
||||||
where: {
|
|
||||||
user: BigInt(message.author.id)
|
|
||||||
}
|
|
||||||
});
|
|
||||||
|
|
||||||
let limit = 25;
|
|
||||||
|
|
||||||
if (userLimits?.limit) {
|
|
||||||
limit = userLimits.limit;
|
|
||||||
}
|
|
||||||
|
|
||||||
if (!userLimits?.vip) {
|
|
||||||
const usedLimit = await database.usage.count({
|
|
||||||
select: { _all: true },
|
|
||||||
where: {
|
|
||||||
user: BigInt(message.author.id),
|
|
||||||
timestamp: {
|
|
||||||
gte: new Date(message.createdTimestamp - 1000*60*60*24 /*24 hours */)
|
|
||||||
}
|
|
||||||
},
|
|
||||||
});
|
|
||||||
|
|
||||||
if (usedLimit._all >= limit) {
|
|
||||||
message.react("🛑");
|
|
||||||
message.author.dmChannel?.send({
|
|
||||||
embeds: [{
|
|
||||||
color: 0xff0000,
|
|
||||||
description: `You've used up your message limit for today, ${limit} requrests in last 24 hours`
|
|
||||||
}]
|
|
||||||
});
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
const messagesForChannel = channelsRunning.ensure(message.channelId, () => {return [] as DiscordApi.Message[];} );
|
|
||||||
const shouldStart = messagesForChannel.length == 0;
|
|
||||||
messagesForChannel.push(message);
|
|
||||||
if (shouldStart)
|
|
||||||
onMessage(message.channelId);
|
|
||||||
});
|
});
|
||||||
|
|
||||||
async function onMessage(channel: string) {
|
|
||||||
const channelQueue = channelsRunning.get(channel) as DiscordApi.Message[];
|
|
||||||
const message = channelQueue.at(0) as DiscordApi.Message;
|
|
||||||
|
|
||||||
try {
|
|
||||||
let messages: DiscordApi.Collection<string, DiscordApi.Message> = await message.channel.messages.fetch({ limit: config.limits.messages, cache: false });
|
|
||||||
|
|
||||||
messages = messages.filter(m => message.createdTimestamp - m.createdTimestamp < config.limits.time );
|
|
||||||
|
|
||||||
messages.forEach(m => Moderation.checkMessage(m));
|
|
||||||
|
|
||||||
message.channel.sendTyping();
|
|
||||||
const answer = await openai.createChatCompletion({
|
|
||||||
...config.chatCompletionConfig,
|
|
||||||
messages: toOpenAIMessages(messages),
|
|
||||||
});
|
|
||||||
|
|
||||||
const usage = answer.data.usage;
|
|
||||||
if (usage != undefined) {
|
|
||||||
const channelName: string = message.inGuild() ? `${message.channel.name} (${message.guild.name})` : `@${message.author.tag}`;
|
|
||||||
console.log(`Used ${usage.total_tokens} (${usage.prompt_tokens} + ${usage.completion_tokens}) tokens for ${message.author.tag} (${message.author.id}) in #${channelName}`);
|
|
||||||
|
|
||||||
database.usage.create({
|
|
||||||
data: {
|
|
||||||
timestamp: message.createdAt,
|
|
||||||
user: BigInt(message.author.id),
|
|
||||||
channel: BigInt(message.channelId),
|
|
||||||
guild: message.guildId ? BigInt(message.guildId) : null,
|
|
||||||
usageReguest: usage.prompt_tokens,
|
|
||||||
usageResponse: usage.completion_tokens
|
|
||||||
}
|
|
||||||
}).catch((e => {
|
|
||||||
console.error("Failed to push to a database");
|
|
||||||
console.error(e);
|
|
||||||
}));
|
|
||||||
}
|
|
||||||
|
|
||||||
const answerContent = answer.data.choices[0].message?.content;
|
|
||||||
|
|
||||||
if (answerContent != undefined && answerContent != "") {
|
|
||||||
const response = message.reply({
|
|
||||||
content: answerContent,
|
|
||||||
allowedMentions: {
|
|
||||||
repliedUser: false,
|
|
||||||
}
|
|
||||||
});
|
|
||||||
|
|
||||||
Moderation.checkMessage(await response);
|
|
||||||
}
|
|
||||||
else {
|
|
||||||
message.react("😶");
|
|
||||||
}
|
|
||||||
} catch (e) {
|
|
||||||
console.error(`Error ocurred while handling chat completion request (${(e as object).constructor.name}):`);
|
|
||||||
console.error(e);
|
|
||||||
|
|
||||||
message.reply({
|
|
||||||
embeds: [{
|
|
||||||
color: 0xff0000,
|
|
||||||
description: "Something bad happened! :frowning:"
|
|
||||||
}],
|
|
||||||
allowedMentions: {
|
|
||||||
repliedUser: false,
|
|
||||||
}
|
|
||||||
});
|
|
||||||
}
|
|
||||||
|
|
||||||
channelQueue.shift();
|
|
||||||
if (channelQueue.length == 0)
|
|
||||||
channelsRunning.delete(channel);
|
|
||||||
else
|
|
||||||
onMessage(channel);
|
|
||||||
}
|
|
||||||
|
|
||||||
discord.login(config.tokens.Discord);
|
discord.login(config.tokens.Discord);
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue