Export to parquet (#151)

Co-authored-by: Eliott C. <[email protected]>
Co-authored-by: Julien Chaumond <[email protected]>
Co-authored-by: Victor Mustar <[email protected]>

.env

@@ -72,3 +72,7 @@ MODELS=`[
 PUBLIC_ORIGIN=#https://hf.co
 PUBLIC_GOOGLE_ANALYTICS_ID=#G-XXXXXXXX / Leave empty to disable
 PUBLIC_DEPRECATED_GOOGLE_ANALYTICS_ID=#UA-XXXXXXXX-X / Leave empty to disable
+
+PARQUET_EXPORT_DATASET=
+PARQUET_EXPORT_HF_TOKEN=
+PARQUET_EXPORT_SECRET=

package-lock.json

@@ -8,6 +8,7 @@
 			"name": "chat-ui",
 			"version": "0.1.0",
 			"dependencies": {
+				"@huggingface/hub": "^0.5.1",
 				"@huggingface/inference": "^2.2.0",
 				"autoprefixer": "^10.4.14",
 				"date-fns": "^2.29.3",
@@ -16,6 +17,7 @@
 				"marked": "^4.3.0",
 				"mongodb": "^5.3.0",
 				"nanoid": "^4.0.2",
+				"parquetjs": "^0.11.2",
 				"postcss": "^8.4.21",
 				"tailwind-scrollbar": "^3.0.0",
 				"tailwindcss": "^3.3.1",
@@ -27,6 +29,7 @@
 				"@sveltejs/kit": "^1.15.10",
 				"@tailwindcss/typography": "^0.5.9",
 				"@types/marked": "^4.0.8",
+				"@types/parquetjs": "^0.10.3",
 				"@typescript-eslint/eslint-plugin": "^5.45.0",
 				"@typescript-eslint/parser": "^5.45.0",
 				"eslint": "^8.28.0",
@@ -473,6 +476,17 @@
 				"node": "^12.22.0 || ^14.17.0 || >=16.0.0"
 			}
 		},
+		"node_modules/@huggingface/hub": {
+			"version": "0.5.1",
+			"resolved": "https://registry.npmjs.org/@huggingface/hub/-/hub-0.5.1.tgz",
+			"integrity": "sha512-ZaE2gY8NY+XwIOL7+gBhPq19PXG4gbGSSJ7zwWLoq6MKP+nsgkQk/c7fBFrxgBwR6lNd0AJMHPRCjwTndqsqWQ==",
+			"dependencies": {
+				"hash-wasm": "^4.9.0"
+			},
+			"engines": {
+				"node": ">=18"
+			}
+		},
 		"node_modules/@huggingface/inference": {
 			"version": "2.2.0",
 			"resolved": "https://registry.npmjs.org/@huggingface/inference/-/inference-2.2.0.tgz",
@@ -893,6 +907,24 @@
 			"resolved": "https://registry.npmjs.org/@types/node/-/node-18.13.0.tgz",
 			"integrity": "sha512-gC3TazRzGoOnoKAhUx+Q0t8S9Tzs74z7m0ipwGpSqQrleP14hKxP4/JUeEQcD3W1/aIpnWl8pHowI7WokuZpXg=="
 		},
+		"node_modules/@types/node-int64": {
+			"version": "0.4.29",
+			"resolved": "https://registry.npmjs.org/@types/node-int64/-/node-int64-0.4.29.tgz",
+			"integrity": "sha512-rHXvenLTj/CcsmNAebaBOhxQ2MqEGl3yXZZcZ21XYR+gzGTTcpOy2N4IxpvTCz48loyQNatHvfn6GhIbbZ1R3Q==",
+			"dev": true,
+			"dependencies": {
+				"@types/node": "*"
+			}
+		},
+		"node_modules/@types/parquetjs": {
+			"version": "0.10.3",
+			"resolved": "https://registry.npmjs.org/@types/parquetjs/-/parquetjs-0.10.3.tgz",
+			"integrity": "sha512-n0xVEor3+3qHfCmFAf0pO4m/Pxc5JEmiVkEWWqJexN+p11/Nr+rqABKcIEj4X6tGKF1cnVIeBqy67mW2Yd+Kbg==",
+			"dev": true,
+			"dependencies": {
+				"@types/node-int64": "*"
+			}
+		},
 		"node_modules/@types/pug": {
 			"version": "2.0.6",
 			"resolved": "https://registry.npmjs.org/@types/pug/-/pug-2.0.6.tgz",
@@ -1248,6 +1280,25 @@
 			"resolved": "https://registry.npmjs.org/balanced-match/-/balanced-match-1.0.2.tgz",
 			"integrity": "sha512-3oSeUO0TMV67hN1AmbXsK4yaqU7tjiHlbxRDZOpH0KW9+CeX4bRAaX0Anxt0tx2MrpRpWwQaPwIlISEJhYU5Pw=="
 		},
+		"node_modules/base64-js": {
+			"version": "1.5.1",
+			"resolved": "https://registry.npmjs.org/base64-js/-/base64-js-1.5.1.tgz",
+			"integrity": "sha512-AKpaYlHn8t4SVbOHCy+b5+KKgvR4vrsD8vbvrbiQJps7fKDTkjkDry6ji0rUJjC0kzbNePLwzxq8iypo41qeWA==",
+			"funding": [
+				{
+					"type": "github",
+					"url": "https://github.com/sponsors/feross"
+				},
+				{
+					"type": "patreon",
+					"url": "https://www.patreon.com/feross"
+				},
+				{
+					"type": "consulting",
+					"url": "https://feross.org/support"
+				}
+			]
+		},
 		"node_modules/binary-extensions": {
 			"version": "2.2.0",
 			"resolved": "https://registry.npmjs.org/binary-extensions/-/binary-extensions-2.2.0.tgz",
@@ -1256,6 +1307,12 @@
 				"node": ">=8"
 			}
 		},
+		"node_modules/bindings": {
+			"version": "1.2.1",
+			"resolved": "https://registry.npmjs.org/bindings/-/bindings-1.2.1.tgz",
+			"integrity": "sha512-u4cBQNepWxYA55FunZSM7wMi55yQaN0otnhhilNoWHq0MfOfJeQx0v0mRRpolGOExPjZcl6FtB0BB8Xkb88F0g==",
+			"optional": true
+		},
 		"node_modules/brace-expansion": {
 			"version": "1.1.11",
 			"resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-1.1.11.tgz",
@@ -1276,6 +1333,14 @@
 				"node": ">=8"
 			}
 		},
+		"node_modules/brotli": {
+			"version": "1.3.3",
+			"resolved": "https://registry.npmjs.org/brotli/-/brotli-1.3.3.tgz",
+			"integrity": "sha512-oTKjJdShmDuGW94SyyaoQvAjf30dZaHnjJ8uAF+u2/vGJkJbJPJAT1gDiOJP5v1Zb6f9KEyW/1HpuaWIXtGHPg==",
+			"dependencies": {
+				"base64-js": "^1.1.2"
+			}
+		},
 		"node_modules/browserslist": {
 			"version": "4.21.5",
 			"resolved": "https://registry.npmjs.org/browserslist/-/browserslist-4.21.5.tgz",
@@ -2204,6 +2269,11 @@
 				"node": ">=8"
 			}
 		},
+		"node_modules/hash-wasm": {
+			"version": "4.9.0",
+			"resolved": "https://registry.npmjs.org/hash-wasm/-/hash-wasm-4.9.0.tgz",
+			"integrity": "sha512-7SW7ejyfnRxuOc7ptQHSf4LDoZaWOivfzqw+5rpcQku0nHfmicPKE51ra9BiRLAmT8+gGLestr1XroUkqdjL6w=="
+		},
 		"node_modules/highlight.js": {
 			"version": "11.7.0",
 			"resolved": "https://registry.npmjs.org/highlight.js/-/highlight.js-11.7.0.tgz",
@@ -2269,6 +2339,11 @@
 			"resolved": "https://registry.npmjs.org/inherits/-/inherits-2.0.4.tgz",
 			"integrity": "sha512-k/vGaX4/Yla3WzyMCvTQOXYeIHvqOKtnqBduzTHpzpQZzAskKMhZ2K+EnBiSM9zGSoIFeMpXKxa4dYeZIQqewQ=="
 		},
+		"node_modules/int53": {
+			"version": "0.2.4",
+			"resolved": "https://registry.npmjs.org/int53/-/int53-0.2.4.tgz",
+			"integrity": "sha512-a5jlKftS7HUOhkUyYD7j2sJ/ZnvWiNlZS1ldR+g1ifQ+/UuZXIE+YTc/lK1qGj/GwAU5F8Z0e1eVq2t1J5Ob2g=="
+		},
 		"node_modules/ip": {
 			"version": "2.0.0",
 			"resolved": "https://registry.npmjs.org/ip/-/ip-2.0.0.tgz",
@@ -2520,6 +2595,16 @@
 				"node": ">=10"
 			}
 		},
+		"node_modules/lzo": {
+			"version": "0.4.11",
+			"resolved": "https://registry.npmjs.org/lzo/-/lzo-0.4.11.tgz",
+			"integrity": "sha512-apQHNoW2Alg72FMqaC/7pn03I7umdgSVFt2KRkCXXils4Z9u3QBh1uOtl2O5WmZIDLd9g6Lu4lIdOLmiSTFVCQ==",
+			"hasInstallScript": true,
+			"optional": true,
+			"dependencies": {
+				"bindings": "~1.2.1"
+			}
+		},
 		"node_modules/magic-string": {
 			"version": "0.30.0",
 			"resolved": "https://registry.npmjs.org/magic-string/-/magic-string-0.30.0.tgz",
@@ -2741,6 +2826,11 @@
 			"integrity": "sha512-Tj+HTDSJJKaZnfiuw+iaF9skdPpTo2GtEly5JHnWV/hfv2Qj/9RKsGISQtLh2ox3l5EAGw487hnBee0sIJ6v2g==",
 			"dev": true
 		},
+		"node_modules/node-int64": {
+			"version": "0.4.0",
+			"resolved": "https://registry.npmjs.org/node-int64/-/node-int64-0.4.0.tgz",
+			"integrity": "sha512-O5lz91xSOeoXP6DulyHfllpq+Eg00MWitZIbtPfoSEvqIHdl5gfcY6hYzDWnj0qD5tz52PI08u9qUvSVeUBeHw=="
+		},
 		"node_modules/node-releases": {
 			"version": "2.0.10",
 			"resolved": "https://registry.npmjs.org/node-releases/-/node-releases-2.0.10.tgz",
@@ -2790,6 +2880,14 @@
 				"node": ">= 6"
 			}
 		},
+		"node_modules/object-stream": {
+			"version": "0.0.1",
+			"resolved": "https://registry.npmjs.org/object-stream/-/object-stream-0.0.1.tgz",
+			"integrity": "sha512-+NPJnRvX9RDMRY9mOWOo/NDppBjbZhXirNNSu2IBnuNboClC9h1ZGHXgHBLDbJMHsxeJDq922aVmG5xs24a/cA==",
+			"engines": {
+				"node": ">=0.10"
+			}
+		},
 		"node_modules/once": {
 			"version": "1.4.0",
 			"resolved": "https://registry.npmjs.org/once/-/once-1.4.0.tgz",
@@ -2872,6 +2970,34 @@
 				"node": ">=6"
 			}
 		},
+		"node_modules/parquetjs": {
+			"version": "0.11.2",
+			"resolved": "https://registry.npmjs.org/parquetjs/-/parquetjs-0.11.2.tgz",
+			"integrity": "sha512-Y6FOc3Oi2AxY4TzJPz7fhICCR8tQNL3p+2xGQoUAMbmlJBR7+JJmMrwuyMjIpDiM7G8Wj/8oqOH4UDUmu4I5ZA==",
+			"dependencies": {
+				"brotli": "^1.3.0",
+				"bson": "^1.0.4",
+				"int53": "^0.2.4",
+				"object-stream": "0.0.1",
+				"snappyjs": "^0.6.0",
+				"thrift": "^0.11.0",
+				"varint": "^5.0.0"
+			},
+			"engines": {
+				"node": ">=7.6"
+			},
+			"optionalDependencies": {
+				"lzo": "^0.4.0"
+			}
+		},
+		"node_modules/parquetjs/node_modules/bson": {
+			"version": "1.1.6",
+			"resolved": "https://registry.npmjs.org/bson/-/bson-1.1.6.tgz",
+			"integrity": "sha512-EvVNVeGo4tHxwi8L6bPj3y3itEvStdwvvlojVxxbyYfoaxJ6keLgrTuKdyfEAszFK+H3olzBuafE0yoh0D1gdg==",
+			"engines": {
+				"node": ">=0.6.19"
+			}
+		},
 		"node_modules/path-exists": {
 			"version": "4.0.0",
 			"resolved": "https://registry.npmjs.org/path-exists/-/path-exists-4.0.0.tgz",
@@ -3197,6 +3323,15 @@
 				"node": ">=6"
 			}
 		},
+		"node_modules/q": {
+			"version": "1.5.1",
+			"resolved": "https://registry.npmjs.org/q/-/q-1.5.1.tgz",
+			"integrity": "sha512-kV/CThkXo6xyFEZUugw/+pIOywXcDbFYgSct5cT3gqlbkBE1SJdwy6UQoZvodiWF/ckQLZyDE/Bu1M6gVu5lVw==",
+			"engines": {
+				"node": ">=0.6.0",
+				"teleport": ">=0.2.0"
+			}
+		},
 		"node_modules/queue-microtask": {
 			"version": "1.2.3",
 			"resolved": "https://registry.npmjs.org/queue-microtask/-/queue-microtask-1.2.3.tgz",
@@ -3473,6 +3608,11 @@
 				"npm": ">= 3.0.0"
 			}
 		},
+		"node_modules/snappyjs": {
+			"version": "0.6.1",
+			"resolved": "https://registry.npmjs.org/snappyjs/-/snappyjs-0.6.1.tgz",
+			"integrity": "sha512-YIK6I2lsH072UE0aOFxxY1dPDCS43I5ktqHpeAsuLNYWkE5pGxRGWfDM4/vSUfNzXjC1Ivzt3qx31PCLmc9yqg=="
+		},
 		"node_modules/socks": {
 			"version": "2.7.1",
 			"resolved": "https://registry.npmjs.org/socks/-/socks-2.7.1.tgz",
@@ -3842,6 +3982,19 @@
 				"node": ">=0.8"
 			}
 		},
+		"node_modules/thrift": {
+			"version": "0.11.0",
+			"resolved": "https://registry.npmjs.org/thrift/-/thrift-0.11.0.tgz",
+			"integrity": "sha512-UpsBhOC45a45TpeHOXE4wwYwL8uD2apbHTbtBvkwtUU4dNwCjC7DpQTjw2Q6eIdfNtw+dKthdwq94uLXTJPfFw==",
+			"dependencies": {
+				"node-int64": "^0.4.0",
+				"q": "^1.5.0",
+				"ws": ">= 2.2.3"
+			},
+			"engines": {
+				"node": ">= 4.1.0"
+			}
+		},
 		"node_modules/tiny-glob": {
 			"version": "0.2.9",
 			"resolved": "https://registry.npmjs.org/tiny-glob/-/tiny-glob-0.2.9.tgz",
@@ -4053,6 +4206,11 @@
 			"resolved": "https://registry.npmjs.org/util-deprecate/-/util-deprecate-1.0.2.tgz",
 			"integrity": "sha512-EPD5q1uXyFxJpCrLnCc1nHnq3gOa6DZBocAIiI2TaSCA7VCJ1UJDMagCzIkXNsUYfD1daK//LTEQ8xiIbrHtcw=="
 		},
+		"node_modules/varint": {
+			"version": "5.0.2",
+			"resolved": "https://registry.npmjs.org/varint/-/varint-5.0.2.tgz",
+			"integrity": "sha512-lKxKYG6H03yCZUpAGOPOsMcGxd1RHCu1iKvEHYDPmTyq2HueGhD73ssNBqqQWfvYs04G9iUFRvmAVLW20Jw6ow=="
+		},
 		"node_modules/vite": {
 			"version": "4.2.1",
 			"resolved": "https://registry.npmjs.org/vite/-/vite-4.2.1.tgz",
@@ -4180,6 +4338,26 @@
 			"resolved": "https://registry.npmjs.org/wrappy/-/wrappy-1.0.2.tgz",
 			"integrity": "sha512-l4Sp/DRseor9wL6EvV2+TuQn63dMkPjZ/sp9XkghTEbV9KlPS1xUsZ3u7/IQO4wxtcFB4bgpQPRcR3QCvezPcQ=="
 		},
+		"node_modules/ws": {
+			"version": "8.13.0",
+			"resolved": "https://registry.npmjs.org/ws/-/ws-8.13.0.tgz",
+			"integrity": "sha512-x9vcZYTrFPC7aSIbj7sRCYo7L/Xb8Iy+pW0ng0wt2vCJv7M9HOMy0UoN3rr+IFC7hb7vXoqS+P9ktyLLLhO+LA==",
+			"engines": {
+				"node": ">=10.0.0"
+			},
+			"peerDependencies": {
+				"bufferutil": "^4.0.1",
+				"utf-8-validate": ">=5.0.2"
+			},
+			"peerDependenciesMeta": {
+				"bufferutil": {
+					"optional": true
+				},
+				"utf-8-validate": {
+					"optional": true
+				}
+			}
+		},
 		"node_modules/yallist": {
 			"version": "4.0.0",
 			"resolved": "https://registry.npmjs.org/yallist/-/yallist-4.0.0.tgz",

package.json

@@ -17,6 +17,7 @@
 		"@sveltejs/kit": "^1.15.10",
 		"@tailwindcss/typography": "^0.5.9",
 		"@types/marked": "^4.0.8",
+		"@types/parquetjs": "^0.10.3",
 		"@typescript-eslint/eslint-plugin": "^5.45.0",
 		"@typescript-eslint/parser": "^5.45.0",
 		"eslint": "^8.28.0",
@@ -35,6 +36,7 @@
 	"type": "module",
 	"dependencies": {
 		"@huggingface/inference": "^2.2.0",
+		"@huggingface/hub": "^0.5.1",
 		"autoprefixer": "^10.4.14",
 		"date-fns": "^2.29.3",
 		"dotenv": "^16.0.3",
@@ -42,6 +44,7 @@
 		"marked": "^4.3.0",
 		"mongodb": "^5.3.0",
 		"nanoid": "^4.0.2",
+		"parquetjs": "^0.11.2",
 		"postcss": "^8.4.21",
 		"tailwind-scrollbar": "^3.0.0",
 		"tailwindcss": "^3.3.1",

src/hooks.server.ts

@@ -14,7 +14,11 @@ export const handle: Handle = async ({ event, resolve }) => {
 
 	event.locals.sessionId = token || crypto.randomUUID();
 
-	if (event.request.method === "POST" && !event.url.pathname.startsWith(`${base}/settings`)) {
+	if (
+		event.request.method === "POST" &&
+		!event.url.pathname.startsWith(`${base}/settings`) &&
+		!event.url.pathname.startsWith(`${base}/admin`)
+	) {
 		const hasAcceptedEthicsModal = await collections.settings.countDocuments({
 			sessionId: event.locals.sessionId,
 			ethicsModalAcceptedAt: { $exists: true },

src/routes/admin/export/+server.ts

@@ -0,0 +1,107 @@
+import {
+	PARQUET_EXPORT_DATASET,
+	PARQUET_EXPORT_HF_TOKEN,
+	PARQUET_EXPORT_SECRET,
+} from "$env/static/private";
+import { collections } from "$lib/server/database.js";
+import type { Message } from "$lib/types/Message.js";
+import { error } from "@sveltejs/kit";
+import { pathToFileURL } from "node:url";
+import { unlink, writeFile } from "node:fs/promises";
+import { uploadFile } from "@huggingface/hub";
+import parquet from "parquetjs";
+import { z } from "zod";
+
+// Triger like this:
+// curl -X POST "http://localhost:5173/chat/admin/export" -H "Authorization: Bearer <PARQUET_EXPORT_SECRET>" -H "Content-Type: application/json" -d '{"model": "OpenAssistant/oasst-sft-6-llama-30b-xor"}'
+
+export async function POST({ request }) {
+	if (!PARQUET_EXPORT_SECRET || !PARQUET_EXPORT_DATASET || !PARQUET_EXPORT_HF_TOKEN) {
+		throw error(500, "Parquet export is not configured.");
+	}
+
+	if (request.headers.get("Authorization") !== `Bearer ${PARQUET_EXPORT_SECRET}`) {
+		throw error(403);
+	}
+
+	const { model } = z
+		.object({
+			model: z.string(),
+		})
+		.parse(await request.json());
+
+	const schema = new parquet.ParquetSchema({
+		title: { type: "UTF8" },
+		created_at: { type: "TIMESTAMP_MILLIS" },
+		updated_at: { type: "TIMESTAMP_MILLIS" },
+		messages: { repeated: true, fields: { from: { type: "UTF8" }, content: { type: "UTF8" } } },
+	});
+
+	const fileName = `/tmp/conversations-${new Date().toJSON().slice(0, 10)}-${Date.now()}.parquet`;
+
+	const writer = await parquet.ParquetWriter.openFile(schema, fileName);
+
+	let count = 0;
+	console.log("Exporting conversations for model", model);
+
+	for await (const conversation of collections.settings.aggregate<{
+		title: string;
+		created_at: Date;
+		updated_at: Date;
+		messages: Message[];
+	}>([
+		{ $match: { shareConversationsWithModelAuthors: true } },
+		{
+			$lookup: {
+				from: "conversations",
+				localField: "sessionId",
+				foreignField: "sessionId",
+				as: "conversations",
+				pipeline: [{ $match: { model } }],
+			},
+		},
+		{ $unwind: "$conversations" },
+		{
+			$project: {
+				title: "$conversations.title",
+				created_at: "$conversations.createdAt",
+				updated_at: "$conversations.updatedAt",
+				messages: "$conversations.messages",
+			},
+		},
+	])) {
+		await writer.appendRow({
+			title: conversation.title,
+			created_at: conversation.created_at,
+			updated_at: conversation.updated_at,
+			messages: conversation.messages.map((message: Message) => ({
+				from: message.from,
+				content: message.content,
+			})),
+		});
+		++count;
+
+		if (count % 1_000 === 0) {
+			console.log("Exported", count, "conversations");
+		}
+	}
+
+	await writer.close();
+
+	console.log("Uploading", fileName, "to Hugging Face Hub");
+
+	await uploadFile({
+		file: pathToFileURL(fileName),
+		credentials: { accessToken: PARQUET_EXPORT_HF_TOKEN },
+		repo: {
+			type: "dataset",
+			name: PARQUET_EXPORT_DATASET,
+		},
+	});
+
+	console.log("Upload done");
+
+	await unlink(fileName);
+
+	return new Response();
+}