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

<script setup lang="ts">
import type schema from '#triplit/schema';
import type { Entity } from '@triplit/client';
import type { ModelWithProvider } from '~/composables/useModels';

const rootStart = Date.now();
const triplit = useTriplitClient();

const chatPaneWrapper = ref<HTMLElement | null>(null);
const inputValue = ref('');
const route = useRoute();
const { sendMessage, regenerateMessage } = useChat(route.params.id as string);
const { getAgent, unsubscribe: unsubscribeAgents } = await useAgents();
const { providers, allModels, unsubscribe: unsubscribeModels } = 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?.[0]) return null;

    const messagesMap = new Map();
    const partsByMessage = new Map<string, Entity<typeof schema, 'message_parts'>[]>();

    // Group parts by message ID once
    if (rawParts.value) {
        for (const part of rawParts.value) {
            if (!partsByMessage.has(part.messageId)) {
                partsByMessage.set(part.messageId, []);
            }
            if (part.content !== '' || part.toolCall !== null) {
                partsByMessage.get(part.messageId)!.push(part);
            }
        }
    }

    const generationsMap = new Map(
        rawGenerations.value?.map(g => [g.id, g]) ?? []
    );

    // Single pass to build messages
    for (const msg of rawMessages.value) {
        messagesMap.set(msg.id, {
            ...msg,
            parts: partsByMessage.get(msg.id) ?? [],
            children: [],
            generation: generationsMap.get(msg.generationId!) ?? null
        });
    }

    // Build tree
    const rootMessages: Message[] = [];
    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 || []
    };
});

const submitMessage = 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);
        const chatInput = document.getElementById('chat') as HTMLInputElement | null;
        console.log("chat input", chatInput, message);
        if (chatInput) {
            inputValue.value = message;
            nextTick(() => {
                chatInput.focus();
            });
        }
        return;
    }

    scrollToBottom('instant');
};

const focusedMessageTree = computed(() => {
    const tree: Readonly<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
    });
}

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

const { scrollToBottom } = useAutoScroll(chatPaneWrapper);

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

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

console.log("full page render took", Date.now() - rootStart);

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

<template>
    <!-- chat pane -->
    <div ref="chatPaneWrapper"
        class="chat-scroll-container justify-center w-full h-full [scrollbar-width:thin] [scrollbar-color:#888_transparent] overflow-y-scroll overflow-x-hidden flex justify-center">
        <div class="chatPane max-w-4xl w-full min-h-full flex flex-col px-4">
            <div class="w-full px-px flex flex-col flex-grow gap-2 pb-9">
                <Suspense>
                    <template v-if="Array.isArray(topic?.messages) && topic.messages.length > 0">
                        <Message v-for="message in topic.messages" :key="message.id" :message="message"
                            v-memo="[message.id, message.parts?.length, message.children, message.focusedIndex, message.content]"
                            @delete="handleDelete(message)" @regenerate="handleRegenerate(message)" />
                    </template>
                </Suspense>
            </div>

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