veridian/utils/message.ts

import type { AssistantContent, FilePart, ImagePart, JSONValue, ModelMessage, ToolContent } from "ai";
import { ToolCallType } from "~~/drizzle/schema";
import { Err, Ok, type Result } from "~~/types/result";
import type { Message, MessageEntity } from "~/composables/useChat";
import type { Agent } from "~/composables/useAgents";
import { resolveFileUrl } from "~~/utils/url";

export const buildMessageTree = (flatMessages: MessageEntity[]) => {
    const messagesMap = new Map(flatMessages.map(m => [m.id, { ...m, children: [] as MessageEntity[] }]));
    const roots = [] as Message[];

    for (const msg of messagesMap.values()) {
        if (msg.parentMessageId && messagesMap.has(msg.parentMessageId)) {
            messagesMap.get(msg.parentMessageId)!.children.push(msg);
        } else {
            roots.push(msg);
        }
    }
    return roots;
};

export const buildFocusedMessageTree = (messages: Readonly<Message[]>): MessageEntity[] => {
    let focusedMessageTree: MessageEntity[] = [];

    for (const message of messages) {
        if (message.activeChildId && message.children.length > 0) {
            const activeChild = message.children.find(c => c?.id === message.activeChildId);
            if (activeChild) {
                focusedMessageTree.push(activeChild);
                continue;
            }
        }

        focusedMessageTree.push(message);
    }

    return focusedMessageTree;
}

export interface MarshallOptions {
    signFileUrl?: (fileKey: string) => string;
}

export const marshallMessages = (agent: Agent, messages: Readonly<MessageEntity[]>, opts?: MarshallOptions): Result<ModelMessage[], string> => {
    const marshalledMessages: ModelMessage[] = [];

    if (agent && agent.systemPrompt) {
        marshalledMessages.push({
            role: 'system',
            content: agent.systemPrompt,
        });
    }

    for (const message of messages) {
        switch (message.role) {
            case 'user': {
                const attachments = message.attachments.map(attachment => {
                    let url = resolveFileUrl(attachment.file.url);

                    if (opts?.signFileUrl && attachment.file.url.startsWith('/api/files/')) {
                        const fileKey = attachment.file.url.slice('/api/files/'.length);
                        const token = opts.signFileUrl(fileKey);
                        url = `${url}?${token}`;
                    }

                    if (attachment.file.mimeType.startsWith('image/')) {
                        return {
                            type: 'image',
                            image: new URL(url),
                        };
                    }

                    return {
                        type: 'file',
                        data: new URL(url),
                        filename: attachment.file.name,
                        mediaType: attachment.file.mimeType,
                    };
                }) as (FilePart | ImagePart)[];

                let messageDate = new Date(message.createdAt);
                const prompt = `[${messageDate.toDateString()} ${messageDate.toLocaleTimeString()}]: ${message.content!}`;

                marshalledMessages.push({
                    role: 'user',
                    content: attachments ? [
                        {
                            type: 'text',
                            text: `[${messageDate.toDateString()} ${messageDate.toLocaleTimeString()}]: ${message.content!}`
                        },
                        ...attachments,
                    ] : prompt,
                });
                break;
            }
            case 'assistant':
                let assistantPart: AssistantContent = [];
                let toolParts: ToolContent = [];

                for (const part of (message.parts || [])) {
                    if (!part) return Err('Part is undefined');

                    switch (part.type) {
                        case 'text':
                        case 'reasoning': {
                            if (toolParts.length > 0) {
                                marshalledMessages.push({
                                    role: 'assistant',
                                    content: assistantPart,
                                });
                                marshalledMessages.push({
                                    role: 'tool',
                                    content: toolParts,
                                });
                                toolParts = [];
                                assistantPart = [];
                            }

                            if (part.providerOptions || part.content) assistantPart.push({
                                type: part.type,
                                text: part.content || '',
                                providerOptions: part.providerOptions ? part.providerOptions as Record<string, any> : undefined,
                            })
                            break;
                        }
                        case 'tool-call': {
                            if (part.toolCall === null) return Err('Tool call is null');

                            if (part.toolCall.status === 'pending') {
                                return Err('Marshalling tool call that is still pending. This is likely a UI bug if this happens.');
                            }

                            let inputValue: string | object = '';

                            switch (part.toolCall.input!.type) {
                                case ToolCallType.Text:
                                    inputValue = part.toolCall.input!.value;
                                    break;
                                case ToolCallType.Json:
                                    if (typeof part.toolCall.input!.value === 'string') {
                                        inputValue = JSON.parse(part.toolCall.input!.value);
                                    } else {
                                        inputValue = part.toolCall.input!.value;
                                    }
                                    break;
                            }

                            assistantPart.push({
                                type: 'tool-call',
                                toolCallId: part.toolCall.id!,
                                toolName: part.toolCall.toolName!,
                                input: inputValue,
                                providerOptions: part.providerOptions ? part.providerOptions as Record<string, any> : undefined,
                            })

                            if (part.toolCall.status === 'failed') {
                                let failureType: 'error-text' | 'error-json';
                                let failureValue: string | JSONValue;

                                if (part.toolCall.error === null || part.toolCall.error === undefined) {
                                    failureType = 'error-text';
                                    failureValue = 'An unknown error occurred';
                                } else {
                                    switch (part.toolCall.error!.type) {
                                        case ToolCallType.Text:
                                            failureType = 'error-text';
                                            failureValue = part.toolCall.error!.value;
                                            break;
                                        case ToolCallType.Json:
                                            failureType = 'error-json';
                                            if (typeof part.toolCall.error!.value === 'string') {
                                                failureValue = JSON.parse(part.toolCall.error!.value);
                                            } else {
                                                failureValue = part.toolCall.error!.value;
                                            }
                                            break;
                                    }

                                    failureType = 'error-json';
                                    // failureValue = JSON.stringify(part.toolCall.error!.value);
                                    if (typeof part.toolCall.error!.value === 'string') {
                                        failureValue = JSON.parse(part.toolCall.error!.value);
                                    } else {
                                        failureValue = part.toolCall.error!.value;
                                    }
                                }

                                toolParts.push({
                                    type: 'tool-result',
                                    toolCallId: part.toolCall.id!,
                                    toolName: part.toolCall.toolName!,
                                    // @ts-expect-error - This is a type error, because typescript cant provie that the value must be a string when the type is error-text
                                    output: {
                                        type: failureType,
                                        value: failureValue,
                                    },
                                    providerOptions: part.providerOptions ? part.providerOptions as Record<string, any> : undefined,
                                });
                                break;
                            }

                            if (part.toolCall.status === 'completed') {
                                let outputType: 'text' | 'json';
                                let outputValue: string;

                                switch (part.toolCall.output!.type) {
                                    case ToolCallType.Text:
                                        outputType = 'text';
                                        outputValue = part.toolCall.output!.value;
                                        break;
                                    case ToolCallType.Json:
                                        outputType = 'json';
                                        if (typeof part.toolCall.output!.value === 'string') {
                                            outputValue = JSON.parse(part.toolCall.output!.value);
                                        } else {
                                            outputValue = part.toolCall.output!.value;
                                        }
                                        break;
                                }

                                toolParts.push({
                                    type: 'tool-result',
                                    toolCallId: part.toolCall.id!,
                                    toolName: part.toolCall.toolName!,
                                    output: {
                                        type: outputType,
                                        value: outputValue,
                                    },
                                    providerOptions: part.providerOptions ? part.providerOptions as Record<string, any> : undefined,
                                });
                                break;
                            }
                        } break;
                        default:
                            return Err(`Unknown part type: ${part.type}`);
                    }
                }

                if (assistantPart.length > 0) marshalledMessages.push({
                    role: 'assistant',
                    content: assistantPart,
                });
                if (toolParts.length > 0) marshalledMessages.push({
                    role: 'tool',
                    content: toolParts,
                });
                break;
            default:
                return Err(`Unknown message role: ${message.role}`);
        }
    }

    return Ok(marshalledMessages);
};