GPTcord/src/execution.ts

import DiscordApi, { GuildTextBasedChannel } from "discord.js";

import { database, openai } from "./index";
import Moderation from "./moderation";
import config from "./config";
import toOpenAIMessages from "./toOpenAIMessages";
import FunctionManager from "./funcitonManager";

type NonNullableInObject<T, V> = { [k in keyof T]: k extends V ? NonNullable<T[k]> : T[k] };
type apiRequest = DiscordApi.Message | DiscordApi.RepliableInteraction;
export type request = apiRequest & NonNullableInObject<apiRequest, "channel" | "channelId">;

/** Stores the queue requests on the channels. */
const channelsRunning: DiscordApi.Collection<string, request[]> = new DiscordApi.Collection();

type ChannelQueue = NonNullable<ReturnType<typeof channelsRunning.get>>;
type RequestMessage = NonNullable<ReturnType<ChannelQueue["at"]>>;

/**
 * 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: apiRequest) {
  if (request instanceof DiscordApi.Message) return 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
 */
export 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};
}

/**
 * gets the timestamp of nth use inside time limit
 * @param user the user or id to check
 * @param requestTimestamp the timestamp of the request (message/interaction createdAt)
 * @param nth which timestamp in time limit to get (orderedd from oldest to newest)
 * @returns `false` if user is vip
 * @returns `null` if there is no request
 * @returns `Date` timestamp of the nth request
 */
export async function getNthUseInLimitTimestamp(user: string | { id: string }, requestTimestamp: Date, nth = 1) {
  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 nthUseInLimit = await database.usage.findFirst({
    where: {
      user: BigInt(userId),
      timestamp: {
        gte: new Date(requestTimestamp.getTime() - 1000 * 60 * 60 * 24 /* 24 hours */)
      }
    },
    orderBy: { timestamp: "asc" },
    skip: nth - 1,
  });

  if (!nthUseInLimit) return null;
  return nthUseInLimit.timestamp;
}

/**
 * Replies to a request
 * @param request the request to reply to
 * @param message the message
 * @param replyOptions additional options if the request is a message
 * @param interactionOptions additional options if the request is an interaction
 * @returns Promise of the done action
 */
function requestReply(
  request: request,
  message: DiscordApi.MessageReplyOptions & DiscordApi.InteractionReplyOptions,
  // TODO: add support for these below
  replyOptions: DiscordApi.MessageReplyOptions = {},
  interactionOptions: DiscordApi.InteractionReplyOptions = {},
) {
  if (request instanceof DiscordApi.Message) {
    return request.reply(Object.assign(message, replyOptions));
  }
  else {
    if (!request.deferred)
      return request.reply(Object.assign(message, interactionOptions));
    else if (request.replied)
      return request.followUp(Object.assign(message, interactionOptions));
    else
      return request.editReply(Object.assign(message, interactionOptions));
  }
}

/**
 * Checks if the request can be replied in a channel
 * @param request the request to check
 * @returns if the request can be replied to
 */
function canReplyToRequest(request: apiRequest) {
  return (request.guild?.members.me?.permissionsIn(request.channel as GuildTextBasedChannel).has("SendMessages") ?? true);
}

/**
 * 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: apiRequest) {
  if (!request.channelId) {
    if (request instanceof DiscordApi.Message)
      request.reply("request does not have channelId");
    else if (request.isRepliable())
      request.reply("request does not have channelId");
    console.log("There was incoming execution without channelId set, ignoring");
    console.log(request);
    return;
  }

  if (!canReplyToRequest(request)) return;

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

  if (userLimit !== false && userLimit.remaining <= 0) {
    if (request instanceof DiscordApi.Message) {
      request.react("🛑").catch(/*it's okay*/);
      if (!request.author.dmChannel) await request.author.createDM();
      request.author.dmChannel?.send({
        embeds: [{
          color: 0xff0000,
          description:
            "You've used up your message limit for today,\n" +
            `${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 []; },
  );
  const shouldStart = messagesForChannel.length == 0;
  messagesForChannel.push(request as request);
  if (shouldStart)
    executeFromQueue(request.channelId);
}

/**
 * Logs used tokens to the terminal and to the database
 * @param answer the response that OpenAI returned
 * @param message the message that initiated the execution
 * @param functionRan counter of how many function have been ran
 */
function logUsedTokens(
    answer: Awaited<ReturnType<typeof openai.createChatCompletion>>,
    message: RequestMessage,
    functionRan: number,
  ) {
    const usage = answer.data.usage;
    const functionName = answer.data.choices[0].message?.function_call?.name;
    if (usage != undefined) {
      const channelName: string = !message.channel.isDMBased() ? `${message.channel.name} (${message.guild?.name})` : `@${getAuthor(message).tag}`;
      console.log(`Used ${usage.total_tokens} (${usage.prompt_tokens} + ${usage.completion_tokens}) tokens for ${getAuthor(message).tag} (${getAuthor(message).id}) in #${channelName}${functionName ? " [Function: " + functionName + "]" : ""}`);

      database.usage.create({
        data: {
          timestamp: message.createdAt,
          user: BigInt(getAuthor(message).id),
          channel: BigInt(message.channelId),
          guild: message.guildId ? BigInt(message.guildId) : null,
          usageReguest: usage.prompt_tokens,
          usageResponse: usage.completion_tokens,
          functionName: functionName ?? null,
          functionRan: functionName ? functionRan : 0,
        }
      }).catch((e => {
        console.error("Failed to push to a database");
        console.error(e);
      }));
    }
}

/**
 * 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 ChannelQueue;
  const message = channelQueue.at(0) as RequestMessage;
  let functionRanCounter = 0;

  // ignore if we can't even send anything to reply
  if (!canReplyToRequest(message)) return;

  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));

    if (message instanceof DiscordApi.Message) {
      message.channel.sendTyping();
    }
    else if (message.isRepliable()) {
      message.deferReply();
    }

    const OpenAImessages = toOpenAIMessages(messages);
    let answer = await openai.createChatCompletion({
      ...config.chatCompletionConfig,
      messages: OpenAImessages,
      // FIXME: don't use new instance of FunctionManager
      functions: new FunctionManager().getFunctions(),
    });

    logUsedTokens(answer, message, ++functionRanCounter);

    let generatedMessage = answer.data.choices[0].message;
    if (!generatedMessage) throw new Error("empty message received");

    // handle function calls
    while (generatedMessage.function_call) {
      OpenAImessages.push(generatedMessage);
      OpenAImessages.push({
        role: "function",
        name: generatedMessage.function_call.name,
        // FIXME: don't use new instance of FunctionManager
        content: new FunctionManager().handleFunction(generatedMessage.function_call),
      });
      answer = await openai.createChatCompletion({
        ...config.chatCompletionConfig,
        messages: OpenAImessages,
        // FIXME: don't use new instance of FunctionManager
        functions: new FunctionManager().getFunctions(),
      });

      logUsedTokens(answer, message, ++functionRanCounter);

      generatedMessage = answer.data.choices[0].message;
      if (!generatedMessage) throw new Error("empty message received");
    }

    const answerContent = answer.data.choices[0].message?.content;

    if (answerContent == undefined || answerContent == "") {
      if (message instanceof DiscordApi.Message) message.react("😶").catch(/*it's okay*/);
    }
    else {
      const answerMessagesContent :string[] = [""];
      for (const i in answerContent.split(/\n\n/)) {
        if (answerMessagesContent[answerMessagesContent.length-1].length + i.length < 2000) {
          answerMessagesContent[answerMessagesContent.length-1] += "\n\n" + i;
        }
        else {
          answerMessagesContent.push(i);
        }
      }

      for (const i in answerMessagesContent) {
        const response = requestReply(message, {content: i}, {allowedMentions: { repliedUser: false }});

        response.then(rval => Moderation.checkMessage(rval));
        await response;
      }
    }
  } catch (e) {
    console.error(`Error ocurred while handling chat completion request (${(e as object).constructor.name}):`);
    console.error(e);

    let errorText = "\n";

    if (e instanceof Error) {
      errorText += e.message;
    }
    else errorText = "";

    requestReply(
      message,
      {
        embeds: [{
          color: 0xff0000,
          description: "Something bad happened! :frowning:" + errorText
        }],
      },
      {allowedMentions: { repliedUser: false } },
      { ephemeral: true },
    );
  }

  channelQueue.shift();
  if (channelQueue.length == 0)
    channelsRunning.delete(channel);
  else
    executeFromQueue(channel);
}