[Assistants] trending feature (#938)

* [Assistants] trending feature

* cleaner sort query

* create & use `AssistantStats` collection

* fix grouping

* rm usage of migration

* ♻️ Refacto lock usage / AssistantStat

* 🩹 Fix last24HoursCount.count => last24HoursCount

* 🐛 Fix collection name

* 🐛 Fix DB query

* 🐛 Fix rate limit on assitant creation

* 🐛 Fix aggregation query

* ✨ Only run refreshAssistants if assistants are enabled

* hide UI until we have enough stats collection

---------

Co-authored-by: coyotte508 <coyotte508@gmail.com>

src/hooks.server.ts

@@ -1,6 +1,7 @@
 import {
 	ADMIN_API_SECRET,
 	COOKIE_NAME,
+	ENABLE_ASSISTANTS,
 	EXPOSE_API,
 	MESSAGES_BEFORE_LOGIN,
 	PARQUET_EXPORT_SECRET,
@@ -19,9 +20,13 @@ import { sha256 } from "$lib/utils/sha256";
 import { addWeeks } from "date-fns";
 import { checkAndRunMigrations } from "$lib/migrations/migrations";
 import { building } from "$app/environment";
+import { refreshAssistantsCounts } from "$lib/assistantStats/refresh-assistants-counts";
 
 if (!building) {
 	await checkAndRunMigrations();
+	if (ENABLE_ASSISTANTS) {
+		refreshAssistantsCounts();
+	}
 }
 
 export const handle: Handle = async ({ event, resolve }) => {

src/lib/assistantStats/refresh-assistants-counts.ts

@@ -0,0 +1,84 @@
+import { client, collections } from "$lib/server/database";
+import { acquireLock, refreshLock } from "$lib/migrations/lock";
+import type { ObjectId } from "mongodb";
+import { subDays } from "date-fns";
+
+const LOCK_KEY = "assistants.count";
+
+let hasLock = false;
+let lockId: ObjectId | null = null;
+
+async function refreshAssistantsCountsHelper() {
+	if (!hasLock) {
+		return;
+	}
+
+	try {
+		await client.withSession((session) =>
+			session.withTransaction(async () => {
+				await collections.assistants
+					.aggregate([
+						{ $project: { _id: 1 } },
+						{ $set: { last24HoursCount: 0 } },
+						{
+							$unionWith: {
+								coll: "assistants.stats",
+								pipeline: [
+									{ $match: { "date.at": { $gte: subDays(new Date(), 1) }, "date.span": "hour" } },
+									{
+										$group: {
+											_id: "$assistantId",
+											last24HoursCount: { $sum: "$count" },
+										},
+									},
+								],
+							},
+						},
+						{
+							$group: {
+								_id: "$_id",
+								last24HoursCount: { $sum: "$last24HoursCount" },
+							},
+						},
+						{
+							$merge: {
+								into: "assistants",
+								on: "_id",
+								whenMatched: "merge",
+								whenNotMatched: "discard",
+							},
+						},
+					])
+					.next();
+			})
+		);
+	} catch (e) {
+		console.log("Refresh assistants counter failed!");
+		console.error(e);
+	}
+}
+
+async function maintainLock() {
+	if (hasLock && lockId) {
+		hasLock = await refreshLock(LOCK_KEY, lockId);
+
+		if (!hasLock) {
+			lockId = null;
+		}
+	} else if (!hasLock) {
+		lockId = (await acquireLock(LOCK_KEY)) || null;
+		hasLock = !!lockId;
+	}
+
+	setTimeout(maintainLock, 10_000);
+}
+
+export function refreshAssistantsCounts() {
+	const ONE_HOUR_MS = 3_600_000;
+
+	maintainLock().then(() => {
+		refreshAssistantsCountsHelper();
+
+		setInterval(refreshAssistantsCountsHelper, ONE_HOUR_MS);
+	});
+}

src/lib/migrations/lock.ts

@@ -1,36 +1,45 @@
 import { collections } from "$lib/server/database";
+import { ObjectId } from "mongodb";
 
-export async function acquireLock(key = "migrations") {
+/**
+ * Returns the lock id if the lock was acquired, false otherwise
+ */
+export async function acquireLock(key: string): Promise<ObjectId | false> {
 	try {
+		const id = new ObjectId();
+
 		const insert = await collections.semaphores.insertOne({
+			_id: id,
 			key,
 			createdAt: new Date(),
 			updatedAt: new Date(),
 		});
 
-		return !!insert.acknowledged; // true if the document was inserted
+		return insert.acknowledged ? id : false; // true if the document was inserted
 	} catch (e) {
 		// unique index violation, so there must already be a lock
 		return false;
 	}
 }
 
-export async function releaseLock(key = "migrations") {
+export async function releaseLock(key: string, lockId: ObjectId) {
 	await collections.semaphores.deleteOne({
+		_id: lockId,
 		key,
 	});
 }
 
-export async function isDBLocked(key = "migrations"): Promise<boolean> {
+export async function isDBLocked(key: string): Promise<boolean> {
 	const res = await collections.semaphores.countDocuments({
 		key,
 	});
 	return res > 0;
 }
 
-export async function refreshLock(key = "migrations") {
-	await collections.semaphores.updateOne(
+export async function refreshLock(key: string, lockId: ObjectId): Promise<boolean> {
+	const result = await collections.semaphores.updateOne(
 		{
+			_id: lockId,
 			key,
 		},
 		{
@@ -39,4 +48,6 @@ export async function refreshLock(key = "migrations") {
 			},
 		}
 	);
+
+	return result.matchedCount > 0;
 }

src/lib/migrations/migrations.spec.ts

@@ -1,8 +1,10 @@
-import { afterEach, describe, expect, it } from "vitest";
+import { afterEach, assert, describe, expect, it } from "vitest";
 import { migrations } from "./routines";
 import { acquireLock, isDBLocked, refreshLock, releaseLock } from "./lock";
 import { collections } from "$lib/server/database";
 
+const LOCK_KEY = "migrations";
+
 describe("migrations", () => {
 	it("should not have duplicates guid", async () => {
 		const guids = migrations.map((m) => m._id.toString());
@@ -11,7 +13,7 @@ describe("migrations", () => {
 	});
 
 	it("should acquire only one lock on DB", async () => {
-		const results = await Promise.all(new Array(1000).fill(0).map(() => acquireLock()));
+		const results = await Promise.all(new Array(1000).fill(0).map(() => acquireLock(LOCK_KEY)));
 		const locks = results.filter((r) => r);
 
 		const semaphores = await collections.semaphores.find({}).toArray();
@@ -23,21 +25,24 @@ describe("migrations", () => {
 	});
 
 	it("should read the lock correctly", async () => {
-		expect(await acquireLock()).toBe(true);
-		expect(await isDBLocked()).toBe(true);
-		expect(await acquireLock()).toBe(false);
-		await releaseLock();
-		expect(await isDBLocked()).toBe(false);
+		const lockId = await acquireLock(LOCK_KEY);
+		assert(lockId);
+		expect(await isDBLocked(LOCK_KEY)).toBe(true);
+		expect(!!(await acquireLock(LOCK_KEY))).toBe(false);
+		await releaseLock(LOCK_KEY, lockId);
+		expect(await isDBLocked(LOCK_KEY)).toBe(false);
 	});
 
 	it("should refresh the lock", async () => {
-		await acquireLock();
+		const lockId = await acquireLock(LOCK_KEY);
+
+		assert(lockId);
 
 		// get the updatedAt time
 
 		const updatedAtInitially = (await collections.semaphores.findOne({}))?.updatedAt;
 
-		await refreshLock();
+		await refreshLock(LOCK_KEY, lockId);
 
 		const updatedAtAfterRefresh = (await collections.semaphores.findOne({}))?.updatedAt;
 

src/lib/migrations/migrations.ts

@@ -3,6 +3,8 @@ import { migrations } from "./routines";
 import { acquireLock, releaseLock, isDBLocked, refreshLock } from "./lock";
 import { isHuggingChat } from "$lib/utils/isHuggingChat";
 
+const LOCK_KEY = "migrations";
+
 export async function checkAndRunMigrations() {
 	// make sure all GUIDs are unique
 	if (new Set(migrations.map((m) => m._id.toString())).size !== migrations.length) {
@@ -25,16 +27,17 @@ export async function checkAndRunMigrations() {
 	// connect to the database
 	const connectedClient = await client.connect();
 
-	const hasLock = await acquireLock();
+	const lockId = await acquireLock(LOCK_KEY);
 
-	if (!hasLock) {
+	if (!lockId) {
 		// another instance already has the lock, so we exit early
 		console.log(
 			"[MIGRATIONS] Another instance already has the lock. Waiting for DB to be unlocked."
 		);
 
+		// Todo: is this necessary? Can we just return?
 		// block until the lock is released
-		while (await isDBLocked()) {
+		while (await isDBLocked(LOCK_KEY)) {
 			await new Promise((resolve) => setTimeout(resolve, 1000));
 		}
 		return;
@@ -43,7 +46,7 @@ export async function checkAndRunMigrations() {
 	// once here, we have the lock
 	// make sure to refresh it regularly while it's running
 	const refreshInterval = setInterval(async () => {
-		await refreshLock();
+		await refreshLock(LOCK_KEY, lockId);
 	}, 1000 * 10);
 
 	// iterate over all migrations
@@ -112,5 +115,5 @@ export async function checkAndRunMigrations() {
 	console.log("[MIGRATIONS] All migrations applied. Releasing lock");
 
 	clearInterval(refreshInterval);
-	await releaseLock();
+	await releaseLock(LOCK_KEY, lockId);
 }

src/lib/server/database.ts

@@ -12,6 +12,7 @@ import type { Report } from "$lib/types/Report";
 import type { ConversationStats } from "$lib/types/ConversationStats";
 import type { MigrationResult } from "$lib/types/MigrationResult";
 import type { Semaphore } from "$lib/types/Semaphore";
+import type { AssistantStats } from "$lib/types/AssistantStats";
 
 if (!MONGODB_URL) {
 	throw new Error(
@@ -32,6 +33,7 @@ export function getCollections(mongoClient: MongoClient) {
 	const conversations = db.collection<Conversation>("conversations");
 	const conversationStats = db.collection<ConversationStats>(CONVERSATION_STATS_COLLECTION);
 	const assistants = db.collection<Assistant>("assistants");
+	const assistantStats = db.collection<AssistantStats>("assistants.stats");
 	const reports = db.collection<Report>("reports");
 	const sharedConversations = db.collection<SharedConversation>("sharedConversations");
 	const abortedGenerations = db.collection<AbortedGeneration>("abortedGenerations");
@@ -47,6 +49,7 @@ export function getCollections(mongoClient: MongoClient) {
 		conversations,
 		conversationStats,
 		assistants,
+		assistantStats,
 		reports,
 		sharedConversations,
 		abortedGenerations,
@@ -67,6 +70,7 @@ const {
 	conversations,
 	conversationStats,
 	assistants,
+	assistantStats,
 	reports,
 	sharedConversations,
 	abortedGenerations,
@@ -143,6 +147,11 @@ client.on("open", () => {
 	assistants.createIndex({ featured: 1, userCount: -1 }).catch(console.error);
 	assistants.createIndex({ modelId: 1, userCount: -1 }).catch(console.error);
 	assistants.createIndex({ searchTokens: 1 }).catch(console.error);
+	assistants.createIndex({ last24HoursCount: 1 }).catch(console.error);
+	assistantStats
+		// Order of keys is important for the queries
+		.createIndex({ "date.span": 1, "date.at": 1, assistantId: 1 }, { unique: true })
+		.catch(console.error);
 	reports.createIndex({ assistantId: 1 }).catch(console.error);
 	reports.createIndex({ createdBy: 1, assistantId: 1 }).catch(console.error);
 

src/lib/types/Assistant.ts

@@ -27,4 +27,11 @@ export interface Assistant extends Timestamps {
 	};
 	dynamicPrompt?: boolean;
 	searchTokens: string[];
+	last24HoursCount: number;
+}
+
+// eslint-disable-next-line no-shadow
+export enum SortKey {
+	POPULAR = "popular",
+	TRENDING = "trending",
 }

src/lib/types/AssistantStats.ts

@@ -0,0 +1,11 @@
+import type { Timestamps } from "./Timestamps";
+import type { Assistant } from "./Assistant";
+
+export interface AssistantStats extends Timestamps {
+	assistantId: Assistant["_id"];
+	date: {
+		at: Date;
+		span: "hour";
+	};
+	count: number;
+}

src/routes/assistants/+page.server.ts

@@ -1,7 +1,7 @@
 import { base } from "$app/paths";
 import { ENABLE_ASSISTANTS } from "$env/static/private";
 import { collections } from "$lib/server/database.js";
-import type { Assistant } from "$lib/types/Assistant";
+import { SortKey, type Assistant } from "$lib/types/Assistant";
 import type { User } from "$lib/types/User";
 import { generateQueryTokens } from "$lib/utils/searchTokens.js";
 import { error, redirect } from "@sveltejs/kit";
@@ -18,6 +18,7 @@ export const load = async ({ url, locals }) => {
 	const pageIndex = parseInt(url.searchParams.get("p") ?? "0");
 	const username = url.searchParams.get("user");
 	const query = url.searchParams.get("q")?.trim() ?? null;
+	const sort = url.searchParams.get("sort")?.trim() ?? SortKey.POPULAR;
 	const createdByCurrentUser = locals.user?.username && locals.user.username === username;
 
 	let user: Pick<User, "_id"> | null = null;
@@ -41,7 +42,10 @@ export const load = async ({ url, locals }) => {
 	const assistants = await collections.assistants
 		.find(filter)
 		.skip(NUM_PER_PAGE * pageIndex)
-		.sort({ userCount: -1 })
+		.sort({
+			...(sort === SortKey.TRENDING && { last24HoursCount: -1 }),
+			userCount: -1,
+		})
 		.limit(NUM_PER_PAGE)
 		.toArray();
 
@@ -53,5 +57,6 @@ export const load = async ({ url, locals }) => {
 		numTotalItems,
 		numItemsPerPage: NUM_PER_PAGE,
 		query,
+		sort,
 	};
 };

src/routes/assistants/+page.svelte

@@ -22,6 +22,7 @@
 	import { useSettingsStore } from "$lib/stores/settings";
 	import IconInternet from "$lib/components/icons/IconInternet.svelte";
 	import { isDesktop } from "$lib/utils/isDesktop";
+	import { SortKey } from "$lib/types/Assistant";
 
 	export let data: PageData;
 
@@ -32,6 +33,7 @@
 	let filterInputEl: HTMLInputElement;
 	let filterValue = data.query;
 	let isFilterInPorgress = false;
+	let sortValue = data.sort as SortKey;
 
 	const onModelChange = (e: Event) => {
 		const newUrl = getHref($page.url, {
@@ -71,6 +73,14 @@
 		}
 	}, SEARCH_DEBOUNCE_DELAY);
 
+	const sortAssistants = () => {
+		const newUrl = getHref($page.url, {
+			newKeys: { sort: sortValue },
+			existingKeys: { behaviour: "delete", keys: ["p"] },
+		});
+		goto(newUrl);
+	};
+
 	const settings = useSettingsStore();
 </script>
 
@@ -130,7 +140,7 @@
 			</a>
 		</div>
 
-		<div class="mt-7 flex items-center gap-x-2 text-sm">
+		<div class="mt-7 flex flex-wrap items-center gap-x-2 gap-y-3 text-sm">
 			{#if assistantsCreator && !createdByMe}
 				<div
 					class="flex items-center gap-1.5 rounded-full border border-gray-300 bg-gray-50 px-3 py-1 dark:border-gray-600 dark:bg-gray-700 dark:text-white"
@@ -198,6 +208,14 @@
 					type="search"
 				/>
 			</div>
+			<select
+				bind:value={sortValue}
+				on:change={sortAssistants}
+				class="hidden rounded-lg border border-gray-300 bg-gray-50 px-2 py-1 text-sm text-gray-900 focus:border-blue-700 focus:ring-blue-700 dark:border-gray-600 dark:bg-gray-700 dark:text-white dark:placeholder-gray-400"
+			>
+				<option value={SortKey.POPULAR}>{SortKey.POPULAR}</option>
+				<option value={SortKey.TRENDING}>{SortKey.TRENDING}</option>
+			</select>
 		</div>
 
 		<div class="mt-8 grid grid-cols-2 gap-3 sm:gap-5 md:grid-cols-3 lg:grid-cols-4">

src/routes/conversation/[id]/+server.ts

@@ -1,4 +1,5 @@
 import { MESSAGES_BEFORE_LOGIN, ENABLE_ASSISTANTS_RAG } from "$env/static/private";
+import { startOfHour } from "date-fns";
 import { authCondition, requiresUser } from "$lib/server/auth";
 import { collections } from "$lib/server/database";
 import { models } from "$lib/server/models";
@@ -510,6 +511,14 @@ export async function POST({ request, locals, params, getClientAddress }) {
 		},
 	});
 
+	if (conv.assistantId) {
+		await collections.assistantStats.updateOne(
+			{ assistantId: conv.assistantId, "date.at": startOfHour(new Date()), "date.span": "hour" },
+			{ $inc: { count: 1 } },
+			{ upsert: true }
+		);
+	}
+
 	// Todo: maybe we should wait for the message to be saved before ending the response - in case of errors
 	return new Response(stream, {
 		headers: {

src/routes/settings/(nav)/assistants/new/+page.server.ts

@@ -82,7 +82,9 @@ export const actions: Actions = {
 			return fail(400, { error: true, errors });
 		}
 
-		const assistantsCount = await collections.assistants.countDocuments(authCondition(locals));
+		const createdById = locals.user?._id ?? locals.sessionId;
+
+		const assistantsCount = await collections.assistants.countDocuments({ createdById });
 
 		if (usageLimits?.assistants && assistantsCount > usageLimits.assistants) {
 			const errors = [
@@ -94,8 +96,6 @@ export const actions: Actions = {
 			return fail(400, { error: true, errors });
 		}
 
-		const createdById = locals.user?._id ?? locals.sessionId;
-
 		const newAssistantId = new ObjectId();
 
 		const exampleInputs: string[] = [
@@ -139,6 +139,7 @@ export const actions: Actions = {
 			},
 			dynamicPrompt: parse.data.dynamicPrompt,
 			searchTokens: generateSearchTokens(parse.data.name),
+			last24HoursCount: 0,
 			generateSettings: {
 				temperature: parse.data.temperature,
 				top_p: parse.data.top_p,