veridian/app/pages/agent/[id]/topic/[topicId].vue

<script setup lang="ts">
import type { Message, MessageEntity } from '~/composables/useChat';
import type { ModelWithProvider, ProviderWithModels } from '~/composables/useModels';

const triplit = useTriplitClient();

const chatPane = ref<HTMLElement | null>(null);
const route = useRoute();
const { sendMessage, regenerateMessage } = useChat(route.params.id as string);
const { getAgent, unsubscribe: unsubscribeAgents } = await useAgents();
const { providers, unsubscribe: unsubscribeModels, allModels } = await useModels();

const agent = computed(() => {
    if (route.params.id === null || typeof route.params.id !== 'string') {
        throw new Error('Invalid agent ID');
    }

    return getAgent(route.params.id)!;
});

const topicQuery = computed(() =>
    triplit
        .query('topics')
        .Where(['id', '=', route.params.topicId])
        .Limit(1)
);

const messagesQuery = computed(() =>
    triplit
        .query('messages')
        .Where(['topicId', '=', route.params.topicId])
        .Order('createdAt', 'ASC')
);

const partsQuery = computed(() =>
    triplit
        .query('message_parts')
        .Where(['topicId', '=', route.params.topicId])
        .Order('createdAt', 'ASC')
        .Include('toolCall')
);

const generationsQuery = computed(() =>
    triplit
        .query('generations')
        .Where(['topicId', '=', route.params.topicId])
);

const [
    { results: rawTopic, unsubscribe: unsubscribeTopic },
    { results: rawMessages, unsubscribe: unsubscribeMessages },
    { results: rawParts, unsubscribe: unsubscribeParts },
    { results: rawGenerations, unsubscribe: unsubscribeGenerations }
] = await Promise.all([
    useQuery('topic', triplit, topicQuery),
    useQuery('messages', triplit, messagesQuery),
    useQuery('parts', triplit, partsQuery),
    useQuery('generations', triplit, generationsQuery),
]);

const topic = computed(() => {
    if (!rawMessages.value || !rawTopic.value || !rawTopic.value[0]) return null;

    // Build the messages tree manually for maximum performance
    const messagesMap = new Map();

    // First pass: Create message objects with parts arrays
    for (const msg of rawMessages.value) {
        messagesMap.set(msg.id, {
            ...msg,
            parts: [],
            children: [],
            generation: rawGenerations.value?.find(g => g.id === msg.generationId) ?? null
        });
    }

    // Second pass: Attach parts to messages
    if (rawParts.value) {
        for (const part of rawParts.value) {
            const msg = messagesMap.get(part.messageId);
            if (msg) {
                // Filter empty parts here if needed, or just push
                if (part.content !== '' || part.toolCall !== null) {
                    msg.parts.push(part);
                }
            }
        }
    }

    // Third pass: Build children relationships
    const rootMessages = [];
    for (const msg of messagesMap.values()) {
        if (msg.parentMessageId && messagesMap.has(msg.parentMessageId)) {
            messagesMap.get(msg.parentMessageId).children.push(msg);
        } else {
            rootMessages.push(msg);
        }
    }

    return {
        ...rawTopic.value[0],
        messages: rootMessages as Message[],
        generations: rawGenerations.value || []
    };
});

if (!topic.value) navigateTo(`/agent/${route.params.id}`);

const activeGeneration = computed(() => {
    if (topic.value === null) return null;
    return topic.value?.generations?.find((generation) => generation.status === 'pending') ?? null;
});

const { scrollToBottom } = useAutoScroll(chatPane);

onMounted(() => {
    scrollToBottom('instant');
});

const handleCancel = async () => {
    await $fetch(`/api/chat/cancel/${activeGeneration.value!.id}`, {
        method: 'POST',
    });
};

const handleSubmit = async (message: string, model: ModelWithProvider | null) => {
    if (!model) {
        console.error('No model selected');
        return;
    }

    const res = await sendMessage(message, topic.value!, focusedMessageTree.value, agent.value!, model.provider, model);
    if (!res.ok) {
        console.error('Failed to send message:', res.error);
        return;
    }

    scrollToBottom('instant');
};

const focusedMessageTree = computed(() => {
    const tree: MessageEntity[] = [];
    for (const message of topic.value?.messages || []) {
        if (message.focusedIndex !== undefined && message.focusedIndex !== null) {
            if (message.focusedIndex === 0) {
                tree.push(message);
                continue;
            }

            tree.push(message.children[message.focusedIndex - 1]!);
        } else {
            tree.push(message);
        }
    }
    return tree;
})

const handleRegenerate = async (message: Message) => {
    if (!agent.value.defaultModelId) {
        console.error('No model selected');
        return;
    }

    let messageId;
    if (
        (message.focusedIndex !== undefined && message.focusedIndex !== null)
        && message.focusedIndex > 0
        && message.children.length > 0
    ) {
        messageId = message.children[message.focusedIndex - 1]!.id;
    } else {
        messageId = message.id;
    }

    const model = allModels.value.find(m => m.id === agent.value.defaultModelId);
    if (!model) {
        console.error('Model not found');
        return;
    }

    const res = await regenerateMessage(messageId, topic.value!, focusedMessageTree.value, agent.value!, model.provider, model);
    if (!res.ok) {
        console.error('Failed to regenerate message:', ChatErrorType[res.error]);
        return;
    }
}

const deeplyDeleteMessage = async (message: MessageEntity) => {
    triplit.delete('messages', message.id);

    if (message.generationId !== null && message.generationId !== undefined) {
        triplit.delete('generations', message.generationId);
    }

    for (const part of message.parts || []) {
        triplit.delete('message_parts', part.id);
    }

    if (topic.value?.messages.filter(m => m.id !== message.id).length === 0) {
        triplit.delete('topics', topic.value!.id);
        return navigateTo(`/agent/${route.params.id}/`);
    }
}

const handleDelete = async (rootMessage: Message) => {
    if (rootMessage.role === 'user') {
        deeplyDeleteMessage(rootMessage);
        return;
    }

    if (rootMessage.deleted === true) {
        const message = rootMessage.children[rootMessage.focusedIndex!];
        if (!message) {
            console.error('Message not found');
            return;
        }

        deeplyDeleteMessage(message);

        if (rootMessage.children.filter(child => child!.id !== message.id).length === 0) {
            deeplyDeleteMessage(rootMessage);
        }
        return;
    }

    // - If the message has children, check if they are all soft deleted
    // - If they are all soft deleted, delete the message
    // - If they are not all soft deleted, mark only this message as deleted
    if (
        (rootMessage.focusedIndex !== undefined && rootMessage.focusedIndex !== null)
        && rootMessage.focusedIndex > 0
        && rootMessage.children.length > 0
    ) {
        // we are a child message
        const message = rootMessage.children[rootMessage.focusedIndex - 1]!;
        if (!message) {
            console.error('Message not found');
            return;
        }

        deeplyDeleteMessage(message);
        return;
    }

    // we have no children
    if (rootMessage.children.length === 0) {
        deeplyDeleteMessage(rootMessage);
        return;
    }

    // we are a root message and we have at least one living child, soft delete
    await triplit.update('messages', rootMessage.id, {
        deleted: true
    });
}

onUnmounted(() => {
    unsubscribeTopic?.();
    unsubscribeAgents?.();
    unsubscribeMessages?.();
    unsubscribeParts?.();
    unsubscribeGenerations?.();
    unsubscribeModels?.();
});
</script>

<template>
    <div class="h-full w-full">
        <!-- chat pane -->
        <div class="flex flex-col w-full px-4 overflow-y-auto h-full"
            style="scrollbar-width: thin; scrollbar-color: #888 transparent;" ref="chatPane">
            <div class="flex-grow w-full flex justify-center">
                <div class="max-w-4xl w-full flex flex-col gap-2 pb-9"
                    v-if="Array.isArray(topic?.messages) && topic.messages.length > 0">
                    <Message v-for="message in topic.messages" @delete="handleDelete(message)" :key="message.id"
                        @regenerate="handleRegenerate(message)" :message="message" />
                </div>
            </div>

            <div class="sticky max-h-full z-10 bottom-0 w-full flex justify-center">
                <div class="pb-4 w-full max-w-4xl bg-[var(--color-neutral)] rounded-t-2xl">
                    <ChatInput class="[view-transition-name:chat-prompt] duration-150 ease-in-out"
                        :loading="activeGeneration !== null" :agent="agent"
                        :providers="providers?.filter(p => p.enabled)" @submit="handleSubmit" @cancel="handleCancel">
                    </ChatInput>
                </div>
            </div>
        </div>
    </div>
</template>