stealth-ai-relay/frontend/src/lib/convex/users.ts

import { v } from 'convex/values';
import { mutation, query } from './_generated/server';

const DEFAULT_MODEL = 'gemini-3-pro-preview';

export const getById = query({
	args: { userId: v.id('users') },
	returns: v.union(
		v.object({
			_id: v.id('users'),
			_creationTime: v.number(),
			telegramId: v.int64(),
			telegramChatId: v.optional(v.int64()),
			geminiApiKey: v.optional(v.string()),
			systemPrompt: v.optional(v.string()),
			followUpPrompt: v.optional(v.string()),
			model: v.string(),
			followUpModel: v.optional(v.string()),
			activeChatId: v.optional(v.id('chats')),
			ragCollectionMode: v.optional(
				v.object({
					ragDatabaseId: v.id('ragDatabases'),
					activeSince: v.number()
				})
			),
			injectCollectionMode: v.optional(
				v.object({
					injectDatabaseId: v.id('injectDatabases'),
					activeSince: v.number()
				})
			)
		}),
		v.null()
	),
	handler: async (ctx, args) => {
		return await ctx.db.get(args.userId);
	}
});

export const getByTelegramId = query({
	args: { telegramId: v.int64() },
	returns: v.union(
		v.object({
			_id: v.id('users'),
			_creationTime: v.number(),
			telegramId: v.int64(),
			telegramChatId: v.optional(v.int64()),
			geminiApiKey: v.optional(v.string()),
			systemPrompt: v.optional(v.string()),
			followUpPrompt: v.optional(v.string()),
			model: v.string(),
			followUpModel: v.optional(v.string()),
			activeChatId: v.optional(v.id('chats')),
			ragCollectionMode: v.optional(
				v.object({
					ragDatabaseId: v.id('ragDatabases'),
					activeSince: v.number()
				})
			),
			injectCollectionMode: v.optional(
				v.object({
					injectDatabaseId: v.id('injectDatabases'),
					activeSince: v.number()
				})
			)
		}),
		v.null()
	),
	handler: async (ctx, args) => {
		return await ctx.db
			.query('users')
			.withIndex('by_telegram_id', (q) => q.eq('telegramId', args.telegramId))
			.unique();
	}
});

export const getOrCreate = mutation({
	args: { telegramId: v.int64(), telegramChatId: v.optional(v.int64()) },
	returns: v.id('users'),
	handler: async (ctx, args) => {
		const existing = await ctx.db
			.query('users')
			.withIndex('by_telegram_id', (q) => q.eq('telegramId', args.telegramId))
			.unique();

		if (existing) {
			if (args.telegramChatId && existing.telegramChatId !== args.telegramChatId) {
				await ctx.db.patch(existing._id, { telegramChatId: args.telegramChatId });
			}
			return existing._id;
		}

		return await ctx.db.insert('users', {
			telegramId: args.telegramId,
			telegramChatId: args.telegramChatId,
			model: DEFAULT_MODEL
		});
	}
});

export const setApiKey = mutation({
	args: { userId: v.id('users'), apiKey: v.string() },
	returns: v.null(),
	handler: async (ctx, args) => {
		await ctx.db.patch(args.userId, { geminiApiKey: args.apiKey });
		return null;
	}
});

export const setSystemPrompt = mutation({
	args: { userId: v.id('users'), prompt: v.string() },
	returns: v.null(),
	handler: async (ctx, args) => {
		await ctx.db.patch(args.userId, { systemPrompt: args.prompt });
		return null;
	}
});

export const setFollowUpPrompt = mutation({
	args: { userId: v.id('users'), prompt: v.string() },
	returns: v.null(),
	handler: async (ctx, args) => {
		await ctx.db.patch(args.userId, { followUpPrompt: args.prompt });
		return null;
	}
});

export const setModel = mutation({
	args: { userId: v.id('users'), model: v.string() },
	returns: v.null(),
	handler: async (ctx, args) => {
		await ctx.db.patch(args.userId, { model: args.model });
		return null;
	}
});

export const setFollowUpModel = mutation({
	args: { userId: v.id('users'), model: v.string() },
	returns: v.null(),
	handler: async (ctx, args) => {
		await ctx.db.patch(args.userId, { followUpModel: args.model });
		return null;
	}
});

export const setActiveChat = mutation({
	args: { userId: v.id('users'), chatId: v.id('chats') },
	returns: v.null(),
	handler: async (ctx, args) => {
		await ctx.db.patch(args.userId, { activeChatId: args.chatId });
		return null;
	}
});

export const startRagCollectionMode = mutation({
	args: { userId: v.id('users'), ragDatabaseId: v.id('ragDatabases') },
	returns: v.null(),
	handler: async (ctx, args) => {
		await ctx.db.patch(args.userId, {
			ragCollectionMode: {
				ragDatabaseId: args.ragDatabaseId,
				activeSince: Date.now()
			}
		});
		return null;
	}
});

export const stopRagCollectionMode = mutation({
	args: { userId: v.id('users') },
	returns: v.null(),
	handler: async (ctx, args) => {
		await ctx.db.patch(args.userId, { ragCollectionMode: undefined });
		return null;
	}
});

export const getRagCollectionMode = query({
	args: { userId: v.id('users') },
	returns: v.union(
		v.object({
			ragDatabaseId: v.id('ragDatabases'),
			activeSince: v.number()
		}),
		v.null()
	),
	handler: async (ctx, args) => {
		const user = await ctx.db.get(args.userId);
		return user?.ragCollectionMode ?? null;
	}
});

export const startInjectCollectionMode = mutation({
	args: { userId: v.id('users'), injectDatabaseId: v.id('injectDatabases') },
	returns: v.null(),
	handler: async (ctx, args) => {
		await ctx.db.patch(args.userId, {
			injectCollectionMode: {
				injectDatabaseId: args.injectDatabaseId,
				activeSince: Date.now()
			}
		});
		return null;
	}
});

export const stopInjectCollectionMode = mutation({
	args: { userId: v.id('users') },
	returns: v.null(),
	handler: async (ctx, args) => {
		await ctx.db.patch(args.userId, { injectCollectionMode: undefined });
		return null;
	}
});

export const getInjectCollectionMode = query({
	args: { userId: v.id('users') },
	returns: v.union(
		v.object({
			injectDatabaseId: v.id('injectDatabases'),
			activeSince: v.number()
		}),
		v.null()
	),
	handler: async (ctx, args) => {
		const user = await ctx.db.get(args.userId);
		return user?.injectCollectionMode ?? null;
	}
});