feat: queue

src/chat/llama.ts

@@ -6,7 +6,7 @@ const endpoint = evar("LLAMA_CPP_ENDPOINT");
 
 async function* generate(chat: Chat[], system?: string) {
 	for (const c of chat) {
-		if(c.attachment) {
+		if (c.attachment) {
 			throw new Error("Llama3 does not support attachments");
 		}
 	}

src/index.ts

@@ -1,8 +1,16 @@
-import { ActivityType, Client, Events, GatewayIntentBits } from "discord.js";
+import {
+	ActivityType,
+	Attachment,
+	Client,
+	Events,
+	GatewayIntentBits,
+} from "discord.js";
 import { evar } from "./var";
 import { commands } from "./commands";
+import { addChatQueue } from "./queue";
+import { getAttachmentBase64 } from "./chat";
 
-const client = new Client({
+export const client = new Client({
 	intents:
 		GatewayIntentBits.GuildMessages |
 		GatewayIntentBits.Guilds |
@@ -35,13 +43,38 @@ client.once(Events.ClientReady, async () => {
 	await client.application!.commands.set(commands.map((x) => x.builder));
 });
 
+const resolveAttachment = async (attachment: Attachment) => {
+	return {
+		mime: attachment.contentType!,
+		data: await getAttachmentBase64(attachment.url),
+	};
+};
+
 client.on(Events.MessageCreate, async (message) => {
 	if (
 		(!message.content.length && !message.attachments.size) ||
 		message.author?.bot
 	)
 		return;
-	if (!message.inGuild()) return;
+	if (
+		!message.inGuild() ||
+		!message.channel.isTextBased() ||
+		message.channel.isThread() ||
+		message.channel.isVoiceBased()
+	)
+		return;
+	if (!message.channel.topic?.includes("aichat")) return;
+	addChatQueue(message.channel.id, {
+		text: message.content,
+		role: "user",
+		attachment:
+			message.attachments.size > 0
+				? await resolveAttachment(message.attachments.first()!)
+				: undefined,
+		id: message.id,
+	});
+	// loading
+	message.react("🔄");
 });
 
 client.on(Events.InteractionCreate, async (i) => {

src/queue.ts

@@ -0,0 +1,89 @@
+import { BaseGuildTextChannel, Collection } from "discord.js";
+import { Chat, models } from "./chat";
+import { client } from ".";
+
+const channelSpecificHistory = new Collection<string, Chat[]>();
+const chatQueue = new Collection<
+	string,
+	{ queue: (Chat & { id: string })[]; processing: boolean }
+>();
+
+function addHistory(channelId: string, chat: Chat) {
+	channelSpecificHistory.set(channelId, [
+		...(channelSpecificHistory.get(channelId) || []),
+		chat,
+	]);
+}
+
+setInterval(() => {
+	for (const [channelId, { queue, processing }] of chatQueue) {
+		if (processing) continue;
+		if (queue.length === 0) continue;
+		const chat = queue.shift()!;
+		chatQueue.set(channelId, {
+			queue,
+			processing: true,
+		});
+		const model = models.find((x) => x.id === "gemini-1.0-pro")!;
+		model
+			.generate([...(channelSpecificHistory.get(channelId) || []), chat])
+			.then(async (res) => {
+				let msg = undefined;
+				try {
+					const ch = client.channels.cache.get(channelId) as
+						| BaseGuildTextChannel
+						| undefined;
+					msg = await ch?.send({
+						content: "Generating...",
+						reply: {
+							messageReference: chat.id,
+						},
+					});
+					try {
+						const msg = await ch?.messages.fetch(chat.id);
+						msg?.reactions.removeAll();
+					} catch {
+						console.error("Failed to fetch message");
+					}
+					let tokens = 0;
+					let lastTokens = 0;
+					let contents = "";
+					for await (const { tokens: t, content: c } of res) {
+						tokens += t;
+						contents += c;
+						if (tokens - lastTokens > 100) {
+							await msg?.edit(contents);
+							lastTokens = tokens;
+						}
+					}
+					await msg?.edit(contents);
+					addHistory(channelId, chat);
+					addHistory(channelId, {
+						text: contents,
+						role: "assistant",
+					});
+				} catch (e) {
+					console.error(e);
+					if (msg) {
+						try {
+							await msg.edit("Failed to generate message");
+						} catch {}
+					}
+				}
+				chatQueue.set(channelId, {
+					queue,
+					processing: false,
+				});
+			});
+	}
+}, 1000);
+
+export function addChatQueue(channelId: string, chat: Chat & { id: string }) {
+	if (!chatQueue.has(channelId)) {
+		chatQueue.set(channelId, {
+			queue: [],
+			processing: false,
+		});
+	}
+	chatQueue.get(channelId)!.queue.push(chat);
+}