src/lib/server/api/routes/groups/misc.ts

import { Elysia } from "elysia";
import { authPlugin } from "../../authPlugin";
import { requiresUser } from "$lib/server/auth";
import { collections } from "$lib/server/database";
import { authCondition } from "$lib/server/auth";
import { config } from "$lib/server/config";
import { Client } from "@gradio/client";
import yazl from "yazl";
import { downloadFile } from "$lib/server/files/downloadFile";
import mimeTypes from "mime-types";
import { logger } from "$lib/server/logger";
import { getMetacognitiveConfig } from "$lib/server/metacognitiveConfig";

export interface FeatureFlags {
	enableAssistants: boolean;
	loginEnabled: boolean;
	loginRequired: boolean;
	guestMode: boolean;
	isAdmin: boolean;
}

export type ApiReturnType = Awaited<ReturnType<typeof Client.prototype.view_api>>;

export const misc = new Elysia()
	.use(authPlugin)
	.get("/public-config", async () => config.getPublicConfig())
	.get("/metacognitive-config", async () => {
		const metacogConfig = getMetacognitiveConfig();
		return {
			frequencies: metacogConfig.frequencies,
			comprehensionPrompts: metacogConfig.comprehensionPrompts,
			perspectivePrompts: metacogConfig.perspectivePrompts,
			enabled: metacogConfig.enabled,
		};
	})
	.get("/feature-flags", async ({ locals }) => {
		let loginRequired = false;
		const messagesBeforeLogin = config.MESSAGES_BEFORE_LOGIN
			? parseInt(config.MESSAGES_BEFORE_LOGIN)
			: 0;
		const nConversations = await collections.conversations.countDocuments(authCondition(locals));

		if (requiresUser && !locals.user) {
			if (messagesBeforeLogin === 0) {
				loginRequired = true;
			} else if (nConversations >= messagesBeforeLogin) {
				loginRequired = true;
			} else {
				// get the number of messages where `from === "assistant"` across all conversations.
				const totalMessages =
					(
						await collections.conversations
							.aggregate([
								{ $match: { ...authCondition(locals), "messages.from": "assistant" } },
								{ $project: { messages: 1 } },
								{ $limit: messagesBeforeLogin + 1 },
								{ $unwind: "$messages" },
								{ $match: { "messages.from": "assistant" } },
								{ $count: "messages" },
							])
							.toArray()
					)[0]?.messages ?? 0;

				loginRequired = totalMessages >= messagesBeforeLogin;
			}
		}

		return {
			enableAssistants: config.ENABLE_ASSISTANTS === "true",
			loginEnabled: requiresUser, // misnomer, this is actually whether the feature is available, not required
			loginRequired,
			guestMode: requiresUser && messagesBeforeLogin > 0,
			isAdmin: locals.isAdmin,
		} satisfies FeatureFlags;
	})
	.get("/export", async ({ locals }) => {
		if (!locals.user) {
			throw new Error("Not logged in");
		}

		if (!locals.isAdmin) {
			throw new Error("Not admin");
		}

		if (config.ENABLE_DATA_EXPORT !== "true") {
			throw new Error("Data export is not enabled");
		}

		const nExports = await collections.messageEvents.countDocuments({
			userId: locals.user._id,
			type: "export",
			expiresAt: { $gt: new Date() },
		});

		if (nExports >= 1) {
			throw new Error(
				"You have already exported your data recently. Please wait 1 hour before exporting again."
			);
		}

		const stats: {
			nConversations: number;
			nMessages: number;
			nFiles: number;
			nAssistants: number;
			nAvatars: number;
		} = {
			nConversations: 0,
			nMessages: 0,
			nFiles: 0,
			nAssistants: 0,
			nAvatars: 0,
		};

		const zipfile = new yazl.ZipFile();

		const promises = [
			collections.conversations
				.find({ ...authCondition(locals) })
				.toArray()
				.then(async (conversations) => {
					const formattedConversations = await Promise.all(
						conversations.map(async (conversation) => {
							stats.nConversations++;
							const hashes: string[] = [];
							conversation.messages.forEach(async (message) => {
								stats.nMessages++;
								if (message.files) {
									message.files.forEach((file) => {
										hashes.push(file.value);
									});
								}
							});
							const files = await Promise.all(
								hashes.map(async (hash) => {
									try {
										const fileData = await downloadFile(hash, conversation._id);
										return fileData;
									} catch {
										return null;
									}
								})
							);

							const filenames: string[] = [];
							files.forEach((file) => {
								if (!file) return;

								const extension = mimeTypes.extension(file.mime) || null;
								const convId = conversation._id.toString();
								const fileId = file.name.split("-")[1].slice(0, 8);
								const fileName = `file-${convId}-${fileId}` + (extension ? `.${extension}` : "");
								filenames.push(fileName);
								zipfile.addBuffer(Buffer.from(file.value, "base64"), fileName);
								stats.nFiles++;
							});

							return {
								...conversation,
								messages: conversation.messages.map((message) => {
									return {
										...message,
										files: filenames,
										updates: undefined,
									};
								}),
							};
						})
					);

					zipfile.addBuffer(
						Buffer.from(JSON.stringify(formattedConversations, null, 2)),
						"conversations.json"
					);
				}),
			collections.assistants
				.find({ createdById: locals.user._id })
				.toArray()
				.then(async (assistants) => {
					const formattedAssistants = await Promise.all(
						assistants.map(async (assistant) => {
							if (assistant.avatar) {
								const fileId = collections.bucket.find({ filename: assistant._id.toString() });

								const content = await fileId.next().then(async (file) => {
									if (!file?._id) return;

									const fileStream = collections.bucket.openDownloadStream(file?._id);

									const fileBuffer = await new Promise<Buffer>((resolve, reject) => {
										const chunks: Uint8Array[] = [];
										fileStream.on("data", (chunk) => chunks.push(chunk));
										fileStream.on("error", reject);
										fileStream.on("end", () => resolve(Buffer.concat(chunks)));
									});

									return fileBuffer;
								});

								if (!content) return;

								zipfile.addBuffer(content, `avatar-${assistant._id.toString()}.jpg`);
								stats.nAvatars++;
							}

							stats.nAssistants++;

							return {
								_id: assistant._id.toString(),
								name: assistant.name,
								createdById: assistant.createdById.toString(),
								createdByName: assistant.createdByName,
								avatar: `avatar-${assistant._id.toString()}.jpg`,
								modelId: assistant.modelId,
								preprompt: assistant.preprompt,
								description: assistant.description,
								dynamicPrompt: assistant.dynamicPrompt,
								exampleInputs: assistant.exampleInputs,
								generateSettings: assistant.generateSettings,
								createdAt: assistant.createdAt.toISOString(),
								updatedAt: assistant.updatedAt.toISOString(),
							};
						})
					);

					zipfile.addBuffer(
						Buffer.from(JSON.stringify(formattedAssistants, null, 2)),
						"assistants.json"
					);
				}),
		];

		Promise.all(promises).then(async () => {
			logger.info(
				{
					userId: locals.user?._id,
					...stats,
				},
				"Exported user data"
			);
			zipfile.end();
			if (locals.user?._id) {
				await collections.messageEvents.insertOne({
					userId: locals.user?._id,
					type: "export",
					createdAt: new Date(),
					expiresAt: new Date(Date.now() + 1000 * 60 * 60), // 1 hour
				});
			}
		});

		// @ts-expect-error - zipfile.outputStream is not typed correctly
		return new Response(zipfile.outputStream, {
			headers: {
				"Content-Type": "application/zip",
				"Content-Disposition": 'attachment; filename="export.zip"',
			},
		});
	});