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

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

export const listByChat = query({
	args: { chatId: v.id('chats') },
	returns: v.array(
		v.object({
			_id: v.id('messages'),
			_creationTime: v.number(),
			chatId: v.id('chats'),
			role: v.union(v.literal('user'), v.literal('assistant')),
			content: v.string(),
			imageBase64: v.optional(v.string()),
			imageMediaType: v.optional(v.string()),
			followUpOptions: v.optional(v.array(v.string())),
			source: v.union(v.literal('telegram'), v.literal('web')),
			createdAt: v.number(),
			isStreaming: v.optional(v.boolean())
		})
	),
	handler: async (ctx, args) => {
		return await ctx.db
			.query('messages')
			.withIndex('by_chat_id_and_created_at', (q) => q.eq('chatId', args.chatId))
			.order('asc')
			.collect();
	}
});

export const create = mutation({
	args: {
		chatId: v.id('chats'),
		role: v.union(v.literal('user'), v.literal('assistant')),
		content: v.string(),
		source: v.union(v.literal('telegram'), v.literal('web')),
		imageBase64: v.optional(v.string()),
		imageMediaType: v.optional(v.string()),
		followUpOptions: v.optional(v.array(v.string())),
		isStreaming: v.optional(v.boolean())
	},
	returns: v.id('messages'),
	handler: async (ctx, args) => {
		const messageId = await ctx.db.insert('messages', {
			chatId: args.chatId,
			role: args.role,
			content: args.content,
			source: args.source,
			imageBase64: args.imageBase64,
			imageMediaType: args.imageMediaType,
			followUpOptions: args.followUpOptions,
			createdAt: Date.now(),
			isStreaming: args.isStreaming
		});

		if (args.source === 'web' && args.role === 'user') {
			const chat = await ctx.db.get(args.chatId);
			if (chat) {
				await ctx.db.insert('pendingGenerations', {
					userId: chat.userId,
					chatId: args.chatId,
					userMessage: args.content,
					createdAt: Date.now()
				});
			}
		}

		return messageId;
	}
});

export const update = mutation({
	args: {
		messageId: v.id('messages'),
		content: v.optional(v.string()),
		followUpOptions: v.optional(v.array(v.string())),
		isStreaming: v.optional(v.boolean())
	},
	returns: v.null(),
	handler: async (ctx, args) => {
		const updates: {
			content?: string;
			followUpOptions?: string[];
			isStreaming?: boolean;
		} = {};

		if (args.content !== undefined) {
			updates.content = args.content;
		}
		if (args.followUpOptions !== undefined) {
			updates.followUpOptions = args.followUpOptions;
		}
		if (args.isStreaming !== undefined) {
			updates.isStreaming = args.isStreaming;
		}

		await ctx.db.patch(args.messageId, updates);
		return null;
	}
});

export const getHistoryForAI = query({
	args: { chatId: v.id('chats'), limit: v.optional(v.number()) },
	returns: v.array(
		v.object({
			role: v.union(v.literal('user'), v.literal('assistant')),
			content: v.string()
		})
	),
	handler: async (ctx, args) => {
		const messages = await ctx.db
			.query('messages')
			.withIndex('by_chat_id_and_created_at', (q) => q.eq('chatId', args.chatId))
			.order('asc')
			.collect();

		const limit = args.limit ?? 50;
		const limited = messages.slice(-limit);

		return limited.map((m) => ({
			role: m.role,
			content: m.content
		}));
	}
});

export const getLastAssistantMessage = query({
	args: { chatId: v.id('chats') },
	returns: v.union(
		v.object({
			_id: v.id('messages'),
			_creationTime: v.number(),
			chatId: v.id('chats'),
			role: v.union(v.literal('user'), v.literal('assistant')),
			content: v.string(),
			imageBase64: v.optional(v.string()),
			imageMediaType: v.optional(v.string()),
			followUpOptions: v.optional(v.array(v.string())),
			source: v.union(v.literal('telegram'), v.literal('web')),
			createdAt: v.number(),
			isStreaming: v.optional(v.boolean())
		}),
		v.null()
	),
	handler: async (ctx, args) => {
		const messages = await ctx.db
			.query('messages')
			.withIndex('by_chat_id_and_created_at', (q) => q.eq('chatId', args.chatId))
			.order('desc')
			.collect();

		return messages.find((m) => m.role === 'assistant') ?? null;
	}
});

export const getChatImages = query({
	args: { chatId: v.id('chats') },
	returns: v.array(
		v.object({
			base64: v.string(),
			mediaType: v.string()
		})
	),
	handler: async (ctx, args) => {
		const messages = await ctx.db
			.query('messages')
			.withIndex('by_chat_id', (q) => q.eq('chatId', args.chatId))
			.collect();

		return messages
			.filter((m) => m.imageBase64 && m.imageMediaType)
			.map((m) => ({
				base64: m.imageBase64!,
				mediaType: m.imageMediaType!
			}));
	}
});