434 lines
16 KiB
TypeScript
434 lines
16 KiB
TypeScript
import type schema from "#triplit/schema";
|
|
import type { Entity } from "@triplit/client";
|
|
import type { ModelMessage } from "ai";
|
|
import { decrypt, base64ToUint8Array } from "~/utils/crypto";
|
|
import { type Result, Ok, Err } 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,
|
|
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) {
|
|
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)
|
|
);
|
|
}
|
|
|
|
try {
|
|
$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 newMessage = await triplit.insert('messages', {
|
|
userId: user.value.id,
|
|
topicId: topic.id,
|
|
createdAt: new Date().toISOString(),
|
|
content: message,
|
|
role: 'user',
|
|
}).catch(error => {
|
|
console.error('Failed to insert message:', error);
|
|
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)
|
|
};
|
|
|
|
/**
|
|
*
|
|
* @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);
|
|
}
|
|
|
|
const autoRename = async (topicId: string, prompt: string) => {
|
|
const { settings } = await useUserSettings();
|
|
|
|
console.log(settings.value);
|
|
|
|
if (!settings.value.systemAssistants.rename.enabled) {
|
|
return false;
|
|
}
|
|
|
|
if (!settings.value.systemAssistants.rename.modelId) {
|
|
return false;
|
|
}
|
|
|
|
await triplit.update('topics', topicId, {
|
|
renaming: true
|
|
});
|
|
|
|
const model = await triplit.fetchOne(triplit.query('models').Where('id', '=', settings.value.systemAssistants.rename.modelId).Include('provider'));
|
|
|
|
if (!model) {
|
|
return false;
|
|
}
|
|
|
|
let providerApiKey: string | undefined = undefined;
|
|
if (model.provider!.config.apiKey !== undefined) {
|
|
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)
|
|
);
|
|
}
|
|
|
|
await $fetch(`/api/topic/auto-rename`, {
|
|
method: 'POST',
|
|
body: JSON.stringify({
|
|
modelId: model.id,
|
|
topicId,
|
|
prompt,
|
|
providerApiKey,
|
|
}),
|
|
});
|
|
|
|
return true;
|
|
}
|
|
|
|
return {
|
|
sendMessage,
|
|
autoRename,
|
|
regenerateMessage,
|
|
createTopic,
|
|
};
|
|
} |