veridian/app/composables/useChat.ts

import type schema from "#triplit/schema";
import type { Entity } from "@triplit/client";
import type { ModelMessage } from "ai";
import { nanoid } from "nanoid";
import { decrypt, base64ToUint8Array } from "~/utils/crypto";
import { type Result, Ok, Err, attempt } from "~~/types/result";
import { assert } from "~~/utils/assert";

export type MessageEntity = Entity<typeof schema, 'messages'> & {
    parts: (Entity<typeof schema, 'message_parts'> & {
        toolCall: Entity<typeof schema, 'tool_calls'> | null
    })[] | undefined
} & { generation: Entity<typeof schema, 'generations'> | null }

export type Message =
    MessageEntity & {
        children: (MessageEntity | undefined)[];
    };

export enum ChatErrorType {
    NoModel = 0,
    NoProvider,
    NoAgent,
    NoUser,
    DatabaseOperationFailed,
    FailedToDecryptProviderApiKey,
    GenerationFailed,
    MarshallFailed,
    NoProviderApiKey,
    NoMessage,
    Unimplemented,
}

export const useChat = (agentId: string) => {
    const triplit = useTriplitClient();

    const createTopic = async () => {
        const { user } = useAuth();
        if (!user.value) {
            console.error('No user');
            return;
        }

        const newTopic = await triplit.insert('topics', {
            name: 'New Topic',
            userId: user.value.id,
            agentId,
            createdAt: new Date().toISOString(),
        });

        assert('flush' in triplit);
        await triplit.flush();

        return newTopic;
    };

    const marshallMessages = (agent: Entity<typeof schema, 'agents'>, messages: Readonly<MessageEntity[]>): Result<ModelMessage[], string> => {
        const marshalledMessages: ModelMessage[] = [];

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

        messages.forEach((message) => {
            switch (message.role) {
                case 'user':
                    marshalledMessages.push({
                        role: 'user',
                        // TODO: when we have images or files, this is where we need to handle them
                        content: message.content,
                    });
                    break;
                case 'assistant':
                    (message.parts || []).forEach((part) => {
                        if (!part) return Err('Part is undefined')

                        switch (part.type) {
                            case 'text':
                            case 'reasoning': {
                                marshalledMessages.push({
                                    role: 'assistant',
                                    content: part.content,
                                });
                                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 = '';

                                switch (part.toolCall.input!.type) {
                                    case 'text':
                                        inputValue = part.toolCall.input!.value;
                                        break;
                                    case 'json':
                                        inputValue = JSON.parse(part.toolCall.input!.value);
                                        break;
                                }

                                marshalledMessages.push({
                                    role: 'assistant',
                                    content: [
                                        {
                                            type: 'tool-call',
                                            toolCallId: part.toolCall.id,
                                            toolName: part.toolCall.toolName,
                                            input: inputValue,
                                        },
                                    ],
                                    providerOptions: part.providerOptions,
                                });

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

                                    if (part.toolCall.error === null || part.toolCall.error === undefined) {
                                        failureType = 'error-text';
                                        failureValue = 'An unknown error occurred';
                                    } else {
                                        switch (part.toolCall.error!.type) {
                                            case 'text':
                                                failureType = 'error-text';
                                                failureValue = part.toolCall.error!.value;
                                                break;
                                            case 'json':
                                                failureType = 'error-json';
                                                failureValue = JSON.stringify(part.toolCall.error!.value, null, 2);
                                                break;
                                        }
                                    }

                                    marshalledMessages.push({
                                        role: 'tool',
                                        content: [
                                            {
                                                type: 'tool-result',
                                                toolCallId: part.toolCall.id,
                                                toolName: part.toolCall.toolName,
                                                output: {
                                                    type: failureType,
                                                    value: failureValue,
                                                },
                                            },
                                        ],
                                        providerOptions: part.providerOptions,
                                    });
                                    break;
                                }

                                if (part.toolCall.status === 'completed') {
                                    marshalledMessages.push({
                                        role: 'tool',
                                        content: [
                                            {
                                                type: 'tool-result',
                                                toolCallId: part.toolCall.id,
                                                toolName: part.toolCall.toolName,
                                                output: {
                                                    type: 'json',
                                                    value: JSON.stringify(part.toolCall.output!.value, null, 2),
                                                },
                                            },
                                        ],
                                        providerOptions: part.providerOptions,
                                    });
                                    break;
                                }
                            } break;
                            default:
                                return Err(`Unknown part type: ${part.type}`)
                        }
                    });
                    break;
                default:
                    return Err(`Unknown message role: ${message.role}`)
            }
        });

        return Ok(marshalledMessages);
    };

    const startGeneration = async (
        messages: ModelMessage[],
        args: Record<string, any>,
        topic: Entity<typeof schema, 'topics'>,
        provider: Entity<typeof schema, 'providers'>,
        model: Entity<typeof schema, 'models'>,
        parentMessageId: string | null = null
    ): Promise<Result<void, ChatErrorType>> => {
        let providerApiKey: string | undefined = undefined;
        if (provider.config.apiKey !== undefined) {
            try {
                const key = await crypto.subtle.importKey(
                    "jwk",
                    JSON.parse(window.localStorage.getItem("encryptionKey")!),
                    "AES-GCM",
                    false,
                    ["encrypt", "decrypt"]
                )

                providerApiKey = await decrypt(
                    key,
                    base64ToUint8Array(provider.config.apiKey)
                );
            } catch (error) {
                console.error('Failed to decrypt provider API key:', error);
                return Err(ChatErrorType.FailedToDecryptProviderApiKey);
            }
        }

        try {
            await $fetch('/api/chat/generate', {
                method: 'POST',
                body: {
                    messages,
                    topicId: topic.id,
                    parentMessageId,
                    model: {
                        providerId: provider.id,
                        modelId: model.id,
                        args,
                    },
                    providerApiKey: providerApiKey,
                },
                headers: {
                    'Content-Type': 'application/json',
                },
            });

            return Ok(undefined);
        } catch (error) {
            console.error('Failed to generate:', error);
            return Err(ChatErrorType.GenerationFailed);
        }
    }

    const sendMessage = async (
        message: string,
        topic: Entity<typeof schema, 'topics'>,
        topicMessages: MessageEntity[],
        agent: Entity<typeof schema, 'agents'>,
        provider: Entity<typeof schema, 'providers'>,
        model: Entity<typeof schema, 'models'>,
    ): Promise<Result<void, ChatErrorType>> => {
        const { user } = useAuth();
        if (!user.value) {
            console.error('No user');
            return Err(ChatErrorType.NoUser);
        }

        const messageId = nanoid();
        const newMessage = await triplit.insert('messages', {
            id: messageId,
            userId: user.value.id,
            topicId: topic.id,
            createdAt: new Date().toISOString(),
            content: message,
            role: 'user',
        }).catch(async error => {
            console.error('Failed to insert message:', error);
            await triplit.delete('messages', messageId);
            return Err(ChatErrorType.DatabaseOperationFailed);
        }) as Message;

        const messages = marshallMessages(
            agent,
            topicMessages.concat(newMessage)
        );

        if (messages.ok === false) {
            console.error('Failed to marshall messages:', messages.error);
            return Err(ChatErrorType.MarshallFailed)
        }

        const args = {
            temperature: 1,
            max_tokens: 100,
            top_p: 1,
            frequency_penalty: 0,
            presence_penalty: 0,
        };

        return startGeneration(messages.data, args, topic, provider, model).then(async res => {
            if (res.ok === false) {
                console.error('Failed to start generation:', res.error);
                await triplit.delete('messages', messageId);
            }

            return res;
        });
    };

    /**
     * 
     * @param messageId The ID of the message we wish to regenerate *for*, this
     * can be either a user's message or an agent's message, and we will find 
     * the message that should be regenerated automatically
     * @param topic 
     * @param topicMessages messages in the current topic, if this does not 
     * include an agent message after the user's message we select, the new
     * message will not be a child of the previous message. However, if this
     * array contains an agent message that is to be regenerated, the new
     * message will have that message's id as its parent. The message
     * reference by messageId *must* be in this array
     * @param agent 
     * @param provider 
     * @param model 
     */
    const regenerateMessage = async (
        messageId: string,
        topic: Entity<typeof schema, 'topics'>,
        topicMessages: MessageEntity[],
        agent: Entity<typeof schema, 'agents'>,
        provider: Entity<typeof schema, 'providers'>,
        model: Entity<typeof schema, 'models'>
    ): Promise<Result<void, ChatErrorType>> => {
        const targetMessage = topicMessages.find(message => message.id === messageId);
        if (!targetMessage) {
            return Err(ChatErrorType.NoMessage);
        }
        let targetMessageIndex = topicMessages.indexOf(targetMessage);

        const args = {
            temperature: 1,
            max_tokens: 100,
            top_p: 1,
            frequency_penalty: 0,
            presence_penalty: 0,
        };

        let parentMessageId = null;
        let focusedMessages;
        if (targetMessage.role === 'user') {
            // we need to find the next agent message
            while (targetMessageIndex < topicMessages.length) {
                const currentMessage = topicMessages[targetMessageIndex];
                if (!currentMessage) break;

                if (currentMessage.role === 'assistant') {
                    parentMessageId = currentMessage.parentMessageId || currentMessage.id;
                    focusedMessages = topicMessages.slice(0, targetMessageIndex).filter(message => message.id !== currentMessage.id);
                    break;
                }
                targetMessageIndex++;
            }
        } else {
            parentMessageId = targetMessage.parentMessageId || topicMessages[targetMessageIndex]!.id;
            focusedMessages = topicMessages.slice(0, targetMessageIndex).filter(message => message.id !== messageId);
        }

        if (focusedMessages === undefined) {
            focusedMessages = topicMessages;
        }

        if (parentMessageId === null) {
            const messages = marshallMessages(
                agent,
                topicMessages
            );
            if (messages.ok === false) {
                console.error('Failed to marshall messages:', messages.error);
                return Err(ChatErrorType.MarshallFailed)
            }

            return startGeneration(messages.data, args, topic, provider, model);
        }

        const focusedMessageIndex = topicMessages.findIndex(m => m.id === parentMessageId);
        if (focusedMessageIndex !== topicMessages.length - 1) {

        }

        const messages = marshallMessages(
            agent,
            focusedMessages
        );
        if (messages.ok === false) {
            console.error('Failed to marshall messages:', messages.error);
            return Err(ChatErrorType.MarshallFailed)
        }

        return startGeneration(messages.data, args, topic, provider, model, parentMessageId);
    }

    enum AutoRenameError {
        AutoRenameDisabled = 0,
        NoModelSelected,
        NoModelFound,
        ModelDisabled,
        DatabaseOperationFailed,
        FailedToDecryptProviderApiKey,
        FailedToGenerate,
    }
    const autoRename = async (topicId: string, prompt: string): Promise<Result<string, AutoRenameError>> => {
        const { settings } = await useUserSettings();

        console.log(settings.value);

        if (!settings.value.systemAssistants.rename.enabled) {
            return Err(AutoRenameError.AutoRenameDisabled);
        }

        if (!settings.value.systemAssistants.rename.modelId) {
            return Err(AutoRenameError.NoModelSelected);
        }

        await triplit.update('topics', topicId, {
            renaming: true
        });

        const modelResult = await attempt(triplit.fetchOne(triplit.query('models').Where('id', '=', settings.value.systemAssistants.rename.modelId).Include('provider')));

        if (modelResult.ok === false) {
            return Err(AutoRenameError.DatabaseOperationFailed);
        }

        const model = modelResult.data;

        if (!model) {
            return Err(AutoRenameError.NoModelFound);
        }

        if (model.enabled === false || model.provider?.enabled === false) {
            return Err(AutoRenameError.ModelDisabled)
        }

        let providerApiKey: string | undefined = undefined;
        if (model.provider!.config.apiKey !== undefined) {
            try {
                const key = await crypto.subtle.importKey(
                    "jwk",
                    JSON.parse(window.localStorage.getItem("encryptionKey")!),
                    "AES-GCM",
                    false,
                    ["encrypt", "decrypt"]
                )

                providerApiKey = await decrypt(
                    key,
                    base64ToUint8Array(model.provider!.config.apiKey)
                );
            } catch (error) {
                console.error('Failed to decrypt provider API key:', error);
                return Err(AutoRenameError.FailedToDecryptProviderApiKey);
            }
        }

        try {
            const res = await $fetch(`/api/topic/auto-rename`, {
                method: 'POST',
                body: JSON.stringify({
                    modelId: model.id,
                    topicId,
                    prompt,
                    providerApiKey,
                }),
            });

            return Ok(res.renameId);
        } catch (error) {
            console.error('Failed to auto-rename:', error);
            return Err(AutoRenameError.FailedToGenerate);
        }
    }

    return {
        sendMessage,
        AutoRenameError,
        autoRename,
        regenerateMessage,
        createTopic,
    };
}