feat: ditch triplit, move to postgresql + drizzle orm

This commit is contained in:
Zoe
2026-04-08 17:03:07 -05:00
parent c341c96798
commit e2e3ac6e86
121 changed files with 6680 additions and 4373 deletions
+341 -419
View File
@@ -1,33 +1,27 @@
import type schema from "#triplit/schema";
import type { Entity } from "@triplit/client";
import type { FilePart, ImagePart, 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";
import { nanoid } from 'nanoid';
import * as schema from '~~/drizzle/schema';
import { Err, Ok, type Result } from '~~/types/result';
import { buildMessageTree } from '~~/utils/message';
export type BaseMessage = {
content: string;
fileIds: string[];
}
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 }
& { attachments: Entity<typeof schema, 'attachments'>[] }
export type ToolCall = typeof schema.toolCalls.$inferSelect
export type Message =
MessageEntity & {
children: (MessageEntity | undefined)[];
}
export type MessagePart = typeof schema.messageParts.$inferSelect & { toolCall: ToolCall | null }
export type MessageEntity = typeof schema.messages.$inferSelect & { parts: MessagePart[] | undefined } & { generation: typeof schema.generations.$inferSelect | null } & { attachments: (typeof schema.attachments.$inferSelect & { file: typeof schema.files.$inferSelect })[] }
export type Message = MessageEntity & { children: (MessageEntity | undefined)[] }
export enum ChatErrorType {
NoModel = 0,
NoProvider,
NoAgent,
NoUser,
NoTopic,
DatabaseOperationFailed,
FailedToDecryptProviderApiKey,
GenerationFailed,
@@ -37,315 +31,347 @@ export enum ChatErrorType {
Unimplemented,
}
export const useChat = (agentId: string) => {
const triplit = useTriplitClient();
type Event = { type: string; payload: any; timestamp: number };
const createTopic = async () => {
const { user } = useAuth();
if (!user.value) {
console.error('No user');
const TARGET_UPDATES_PER_SECOND = 24;
export const useChat = async (topicId: MaybeRef<string>, connect = true) => {
let data = useState<Topic & { messages: Message[] } | undefined>('useChat:data', () => undefined);
let sse: EventSource | undefined;
const textDeltaBuffer: Map<string, Map<string, string>> = new Map();
let flushTimeout: ReturnType<typeof setTimeout> | undefined;
let flushCbs: Set<() => void> = new Set();
const flushTextDeltas = () => {
if (!data.value) return;
for (const [messageId, parts] of textDeltaBuffer) {
const msg = data.value.messages.find(m => m.id === messageId);
if (!msg) continue;
for (const [partId, content] of parts) {
console.log("content", content);
const part = msg.parts?.find(p => p.id === partId);
if (part) {
part.content += content;
}
}
}
textDeltaBuffer.clear();
flushTimeout = undefined;
for (const cb of flushCbs) {
cb();
}
flushCbs.clear();
};
const nextFlush = (cb: () => void) => {
if (flushTimeout) {
flushCbs.add(cb);
return;
}
const newTopic = await triplit.insert('topics', {
name: 'New Topic',
userId: user.value.id,
agentId,
createdAt: new Date().toISOString(),
});
cb();
}
assert('flush' in triplit);
await triplit.flush();
return newTopic;
const scheduleFlush = () => {
if (flushTimeout) return;
flushTimeout = setTimeout(flushTextDeltas, 1000 / TARGET_UPDATES_PER_SECOND);
};
const marshallMessages = (agent: Entity<typeof schema, 'agents'>, messages: Readonly<MessageEntity[]>): Result<ModelMessage[], string> => {
const marshalledMessages: ModelMessage[] = [];
const processEvent = (event: Event) => {
if (!data.value) return;
if (agent && agent.systemPrompt) {
marshalledMessages.push({
role: 'system',
content: agent.systemPrompt,
});
const { type, payload } = event;
switch (type) {
case 'MESSAGE_CREATED': {
// Push the new message if it doesn't exist (prevents duplicates from HTTP vs SSE)
if (!data.value.messages.find(m => m.id === payload.id)) {
data.value.messages.push({ ...payload, parts: [] });
} else {
// replace the existing data with the new data
const index = data.value.messages.findIndex(m => m.id === payload.id);
if (index !== -1) {
data.value.messages[index] = { ...payload, parts: [] };
}
}
break;
}
case 'MESSAGE_UPDATED': {
const msgIndex = data.value.messages.findIndex(m => m.id === payload.id);
if (msgIndex !== -1) {
data.value.messages[msgIndex] = { ...data.value.messages[msgIndex], ...payload };
}
break;
}
case 'MESSAGE_DELETED': {
data.value.messages = data.value.messages.filter(m => m.id !== payload.id);
break;
}
case 'text-start': {
const msg = data.value.messages.find(m => m.id === payload.messageId);
if (msg) {
if (!msg.parts) msg.parts = [];
msg.parts.push(payload.part);
}
break;
}
case 'text-delta': {
const msg = data.value.messages.find(m => m.id === payload.messageId);
if (!msg) break;
if (!textDeltaBuffer.has(payload.messageId)) {
textDeltaBuffer.set(payload.messageId, new Map());
}
const messageBuffer = textDeltaBuffer.get(payload.messageId)!;
const existing = messageBuffer.get(payload.partId) || '';
messageBuffer.set(payload.partId, existing + payload.content);
console.log("messageBuffer", messageBuffer, existing + payload.content);
scheduleFlush();
break;
}
case 'text-end': {
const msg = data.value.messages.find(m => m.id === payload.messageId);
const part = msg?.parts?.find(p => p.id === payload.partId);
if (part) {
nextFlush(() => {
part.content = payload.content;
part.finished = true;
});
}
break;
}
case 'tool-call-start': {
const msg = data.value.messages.find(m => m.id === payload.messageId);
if (msg) {
if (!msg.parts) msg.parts = [];
msg.parts.push(payload.part);
}
break;
}
case 'tool-call-delta': {
const msg = data.value.messages.find(m => m.id === payload.messageId);
const part = msg?.parts?.find(p => p.toolCallId === payload.toolCallId);
if (part?.toolCall) {
if (payload.input) part.toolCall.input = payload.input;
if (payload.output) part.toolCall.output = payload.output;
if (payload.error) part.toolCall.error = payload.error;
if (payload.status) part.toolCall.status = payload.status;
}
break;
}
case 'generation-complete': {
const msg = data.value.messages.find(m => m.id === payload.messageId);
if (msg && msg.generation) {
msg.generation.status = 'completed';
msg.generation.tokens = payload.tokens;
}
break;
}
case 'generation-failed': {
const msg = data.value.messages.find(m => m.id === payload.messageId);
if (msg && msg.generation) {
msg.generation.status = 'failed';
msg.generation.error = payload.error;
}
break;
}
case 'topic_updated': {
data.value = { ...data.value, ...payload };
break;
}
}
}
const connectSSE = async () => {
const { promise, resolve } = Promise.withResolvers<void>();
const id = unref(topicId);
const lastUpdate = data.value?.messages.at(-1)?.updatedAt || '0';
const msgCount = data.value?.messages.length || 0;
if (sse && sse.readyState !== EventSource.CLOSED) {
if (!(new URL(sse.url).pathname.startsWith(`/api/topic/${id}`))) {
sse.close();
} else {
return;
}
}
messages.forEach((message) => {
switch (message.role) {
case 'user':
const attachments = message.attachments.map(attachment => {
if (attachment.mimeType.startsWith('image/')) {
return {
type: 'image',
image: attachment.url,
};
}
sse = new EventSource(`/api/topic/${id}?lastUpdate=${lastUpdate}&count=${msgCount}`);
return {
type: 'file',
data: attachment.url,
filename: attachment.name,
mediaType: attachment.mimeType,
};
}) as (FilePart | ImagePart)[];
sse.addEventListener("error", (e) => {
console.error('sse error', e);
setTimeout(connectSSE, 1500);
})
marshalledMessages.push({
role: 'user',
// TODO: when we have images or files, this is where we need to handle them
content: [
{
type: 'text',
text: message.content
},
...attachments,
],
});
break;
case 'assistant':
(message.parts || []).forEach((part) => {
if (!part) return Err('Part is undefined')
sse.addEventListener("message", (e) => {
const raw = JSON.parse(e.data);
const events: Event[] = Array.isArray(raw) ? raw : [raw];
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')
for (const { type, payload, timestamp } of events) {
if (!type || !payload) continue;
if (part.toolCall.status === 'pending') {
return Err('Marshalling tool call that is still pending. This is likely a UI bug if this happens.')
}
if (type === 'initial_state') {
data.value = payload;
resolve()
continue;
}
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}`)
processEvent({ type, payload, timestamp });
}
});
return Ok(marshalledMessages);
};
await promise;
}
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);
if (connect) {
onBeforeUnmount(() => {
if (flushTimeout) {
clearTimeout(flushTimeout);
flushTextDeltas();
}
if (sse) {
console.log('[SSE] Closing connection');
sse.close();
sse = undefined;
}
});
if (import.meta.server) {
const id = unref(topicId);
const { data: ssrData } = await useFetch<Topic & { messages: Message[] }>(`/api/topic/${id}`);
data.value = ssrData.value;
} else {
// if (data.value === undefined) {
// const id = unref(topicId);
// const { data: ssrData } = await useFetch<Topic & { messages: Message[] }>(`/api/topic/${id}`);
// data.value = ssrData.value;
// }
await connectSSE();
}
}
const topic = computed(() => ({
...data.value,
messages: buildMessageTree(data.value?.messages || [])
}));
const getProviderAPIKey = async (provider: Provider): Promise<Result<string | undefined, ChatErrorType>> => {
if (provider.config.apiKey === undefined) {
return Ok(undefined);
}
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',
},
});
const key = await crypto.subtle.importKey(
"jwk",
JSON.parse(window.localStorage.getItem("encryptionKey")!),
"AES-GCM",
false,
["encrypt", "decrypt"]
)
return Ok(undefined);
return Ok(await decrypt(
key,
base64ToUint8Array(provider.config.apiKey)
));
} catch (error) {
console.error('Failed to generate:', error);
return Err(ChatErrorType.GenerationFailed);
console.error('Failed to decrypt provider API key:', error);
return Err(ChatErrorType.FailedToDecryptProviderApiKey);
}
}
const sendMessage = async (
message: BaseMessage,
topic: Entity<typeof schema, 'topics'>,
topicMessages: MessageEntity[],
agent: Entity<typeof schema, 'agents'>,
provider: Entity<typeof schema, 'providers'>,
model: Entity<typeof schema, 'models'>,
baseMessage: BaseMessage,
onRequest?: () => void,
): Promise<Result<void, ChatErrorType>> => {
console.log("sendMessage", baseMessage);
const { user } = useAuth();
if (!user.value) {
console.error('No user');
return Err(ChatErrorType.NoUser);
}
if (!user.value) return Err(ChatErrorType.NoUser);
const messageId = nanoid();
const attachmentsPromise = message.fileIds.map(async fileId => {
const file = await triplit.fetchOne(triplit.query('files').Where('id', '=', fileId));
assert(file !== null);
return await triplit.insert('attachments', {
userId: user.value!.id,
topicId: topic.id,
messageId: messageId,
fileId: fileId,
name: file.name,
mimeType: file.mimeType,
url: file.url,
createdAt: file.createdAt,
})!;
});
const newMessage = await triplit.insert('messages', {
id: messageId,
userId: user.value.id,
topicId: topic.id,
createdAt: new Date().toISOString(),
content: message.content,
role: 'user',
}).catch(async error => {
console.error('Failed to insert message:', error);
await triplit.delete('messages', messageId);
return Err(ChatErrorType.DatabaseOperationFailed);
}) as Message;
const attachments = await Promise.all(attachmentsPromise);
newMessage.attachments = attachments;
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);
try {
const message = {
id: nanoid(),
role: 'user',
content: baseMessage.content,
fileIds: baseMessage.fileIds,
}
return res;
});
console.log("sendMessage", message, baseMessage.content, baseMessage.fileIds);
await $fetch(`/api/topic/${unref(topicId)}/message`, {
method: 'POST',
body: {
message
},
onRequest() {
if (data.value) {
data.value!.messages.push({
topicId: unref(topicId),
userId: user.value!.id,
// TODO
attachments: [],
parts: undefined,
generation: null,
parentMessageId: null,
generationId: null,
focusedIndex: null,
deleted: null,
createdAt: new Date(),
updatedAt: new Date(),
children: [],
...message
} as Message);
}
onRequest?.();
},
onResponseError() {
data.value!.messages = data.value!.messages.filter(m => m.id !== message.id);
},
});
return Ok(undefined);
} catch (error) {
console.error('Failed to send message:', error);
return Err(ChatErrorType.GenerationFailed);
}
};
const startGeneration = async (model: ModelWithProvider) => {
console.log(model.provider);
const providerApiKeyRes = await getProviderAPIKey(model.provider);
if (providerApiKeyRes.ok === false) {
return providerApiKeyRes;
}
const providerApiKey = providerApiKeyRes.data;
await $fetch(`/api/topic/${unref(topicId)}/chat`, {
method: 'POST',
body: {
modelId: model.id,
providerApiKey,
},
});
}
/**
*
* @param messageId The ID of the message we wish to regenerate *for*, this
@@ -364,11 +390,8 @@ export const useChat = (agentId: string) => {
*/
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'>
model: ModelWithProvider
): Promise<Result<void, ChatErrorType>> => {
const targetMessage = topicMessages.find(message => message.id === messageId);
if (!targetMessage) {
@@ -376,16 +399,14 @@ export const useChat = (agentId: string) => {
}
let targetMessageIndex = topicMessages.indexOf(targetMessage);
const args = {
temperature: 1,
max_tokens: 100,
top_p: 1,
frequency_penalty: 0,
presence_penalty: 0,
};
const providerApiKeyRes = await getProviderAPIKey(model.provider);
if (providerApiKeyRes.ok === false) {
return providerApiKeyRes;
}
const providerApiKey = providerApiKeyRes.data;
let parentMessageId = null;
let focusedMessages;
let parentMessageId = undefined;
let focusedMessages: MessageEntity[] | undefined;
if (targetMessage.role === 'user') {
// we need to find the next agent message
while (targetMessageIndex < topicMessages.length) {
@@ -408,130 +429,31 @@ export const useChat = (agentId: string) => {
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 } = 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);
}
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);
}
}
await triplit.update('topics', topicId, {
renaming: true
await $fetch(`/api/topic/${unref(topicId)}/chat`, {
method: 'POST',
body: {
modelId: model.id,
providerApiKey,
parentMessageId,
},
});
try {
// @ts-ignore - excessive stack depth
const res = await $fetch(`/api/auto-rename/${topicId}`, {
method: 'POST',
body: JSON.stringify({
modelId: model.id,
prompt,
providerApiKey,
}),
}) as { ok: true, renameId: string } | { ok: false, code: string };
if (!res.ok) {
console.error('Failed to auto-rename:', res.code);
return Err(AutoRenameError.FailedToGenerate);
}
return Ok(undefined);
}
return Ok(res.renameId);
} catch (error) {
triplit.update('topics', topicId, {
renaming: false,
});
const patchMessageLocally = (id: string, updates: Partial<Message>) => {
if (!data.value) return;
console.error('Failed to auto-rename:', error);
return Err(AutoRenameError.FailedToGenerate);
}
data.value.messages = data.value.messages.map(m =>
m.id === id ? { ...m, ...updates } : m
);
}
return {
topic,
sendMessage,
AutoRenameError,
autoRename,
startGeneration,
regenerateMessage,
createTopic,
patchMessageLocally
};
}
}