veridian/server/api/topic/[topicId]/message.post.ts

import { modelMessageSchema } from 'ai';
import * as z from 'zod';
import { attachments, messageParts, messages } from '~~/drizzle/schema';
import { db } from '~~/server/lib/db';

export default defineEventHandler(async (event) => {
    await protectRoute(event);

    const topicId = getRouterParam(event, 'topicId')!;
    const userId = event.context.user!.id;

    const topic = await db.query.topics.findFirst({
        where: {
            id: topicId,
            userId,
        },
    });

    if (!topic) {
        throw createError({
            statusCode: 404,
            statusMessage: 'Not Found',
            message: 'Topic not found',
        });
    }

    const result = await readValidatedBody(event, z.object({
        message: z.intersection(
            z.object({
                id: z.string(),
                fileIds: z.array(z.string()).optional(),
            }),
            modelMessageSchema
        ),
    }).safeParse);
    if (!result.success) {
        console.log(result.error);
        throw createError({
            statusCode: 400,
            statusMessage: 'Bad Request',
            message: result.error.issues.map(issue => issue.message).join(', '),
        });
    }

    const { message } = result.data;

    if (['user', 'assistant'].includes(message.role) === false) {
        throw createError({
            statusCode: 400,
            statusMessage: 'Bad Request',
            message: 'Invalid role',
        });
    }

    switch (message.role) {
        case 'user': {
            await db.insert(messages).values({
                // @ts-ignore - drizzle bug
                id: message.id,
                userId,
                topicId,
                parentMessageId: null,
                generationId: null,
                role: message.role,
                content: message.content,
            });

            for (const fileId of message.fileIds || []) {
                await db.insert(attachments).values({
                    userId,
                    topicId,
                    messageId: message.id,
                    fileId,
                });
            }

            const usermessage = await db.query.messages.findFirst({
                where: {
                    id: message.id,
                },
                with: {
                    attachments: {
                        with: {
                            file: true,
                        }
                    }
                },
            });

            if (!usermessage) {
                throw createError({
                    statusCode: 500,
                    statusMessage: 'Failed to insert message',
                    message: 'Failed to insert message',
                });
            }
            topicEvents.emit(topicId, { type: 'MESSAGE_CREATED', payload: usermessage });

            return {
                ok: true
            };
        }
        case 'assistant': {
            const [dbmessage] = await db.insert(messages).values({
                // @ts-ignore - drizzle bug
                id: message.id,
                userId,
                topicId,
                parentMessageId: null,
                generationId: null,
                role: message.role,
                content: message.content,
            }).returning();

            if (!dbmessage) {
                throw createError({
                    statusCode: 500,
                    statusMessage: 'Failed to insert message',
                    message: 'Failed to insert message',
                });
            }

            const [part] = await db.insert(messageParts).values({
                userId,
                topicId,
                messageId: dbmessage.id,
                type: 'text',
                content: message.content,
                providerOptions: null,
                finished: true,
                createdAt: new Date(),
                lastUpdatedAt: new Date(),
            }).returning();

            // TODO: transaction
            if (!part) {
                throw createError({
                    statusCode: 500,
                    statusMessage: 'Failed to insert message part',
                    message: 'Failed to insert message part',
                });
            }

            topicEvents.emit(topicId, { type: 'MESSAGE_CREATED', payload: dbmessage });
            topicEvents.emit(topicId, { type: 'text-start', payload: { messageId: dbmessage.id, part } });
            topicEvents.emit(topicId, { type: 'text-end', payload: { messageId: dbmessage.id, partId: part.id, lastUpdatedAt: new Date(), content: message.content } });

            return {
                ok: true
            };
        }
    }
})