GPTcord/src/execution.ts

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 | DiscordApi.Interaction)[]> = new DiscordApi.Collection();

/**
 * Gets the user that requested the execution
 * @param request The request to get the user from
 * @returns the user or guild member
 */
export function getAuthor(request: DiscordApi.Message | DiscordApi.Interaction) {
  if (request instanceof DiscordApi.Message) return request.member ?? request.author;
  return request.user;
}

/**
 * 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 request the message to check and queue
 */
export async function queueRequest(request: DiscordApi.Message | DiscordApi.Interaction) {
  if (!request.channelId) {
    if (!(request instanceof DiscordApi.Message) && request.isRepliable())
      request.reply("request does not have channelId");
    console.log("There was incoming execution without channelId set, ignoring");
    console.log(request);
    return;
  }

  const userLimit = await getUserLimit(getAuthor(request), request.createdAt);

  if (userLimit !== false && userLimit.remaining <= 0) {
    if (request instanceof DiscordApi.Message) {
      request.react("🛑");
      request.author.dmChannel?.send({
        embeds: [{
          color: 0xff0000,
          description: `You've used up your message limit for today, ${userLimit.limit} requests in last 24 hours`,
        }]
      });
    }
    else if (request.isRepliable()) {
      request.reply({
        content: `You've used up your message limit for today, ${userLimit.limit} requests in last 24 hours`,
        ephemeral: true,
      });
    }
    return;
  }

  const messagesForChannel = channelsRunning.ensure(
    request.channelId,
    () => { return [] as DiscordApi.Message[]; },
  );
  const shouldStart = messagesForChannel.length == 0;
  messagesForChannel.push(request);
  if (shouldStart)
    executeFromQueue(request.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);
}