679 lines
26 KiB
TypeScript
679 lines
26 KiB
TypeScript
import { nanoid } from 'nanoid';
|
|
import { toRaw } from 'vue';
|
|
import * as schema from '~~/drizzle/schema';
|
|
import { Err, Ok, type Result } from '~~/types/result';
|
|
import { buildMessageTree } from '~~/utils/message';
|
|
import { appendTextDelta, serializeTopicSnapshot } from '~~/utils/topic-sync';
|
|
|
|
export type BaseMessage = {
|
|
content: string;
|
|
fileIds: string[];
|
|
files?: {
|
|
id: string;
|
|
name: string;
|
|
mimeType: string;
|
|
url: string;
|
|
size?: number;
|
|
}[];
|
|
}
|
|
|
|
export type ToolCall = typeof schema.toolCalls.$inferSelect
|
|
|
|
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,
|
|
MarshallFailed,
|
|
NoProviderApiKey,
|
|
NoMessage,
|
|
Unimplemented,
|
|
}
|
|
|
|
type Event = { type: string; payload: any; timestamp: number; replay?: boolean };
|
|
|
|
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;
|
|
let disposed = false;
|
|
let connectionAttempt = 0;
|
|
let reconnectTimeout: ReturnType<typeof setTimeout> | undefined;
|
|
const { promise: connectionReady, resolve: markConnected } = Promise.withResolvers<void>();
|
|
|
|
const textDeltaBuffer: Map<string, Map<string, {
|
|
content: string;
|
|
lastUpdatedAt: 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) {
|
|
const part = msg.parts?.find(p => p.id === partId);
|
|
if (part) {
|
|
part.content = content.content;
|
|
part.lastUpdatedAt = new Date(content.lastUpdatedAt);
|
|
}
|
|
}
|
|
}
|
|
|
|
textDeltaBuffer.clear();
|
|
flushTimeout = undefined;
|
|
|
|
for (const cb of flushCbs) {
|
|
cb();
|
|
}
|
|
flushCbs.clear();
|
|
};
|
|
|
|
const nextFlush = (cb: () => void) => {
|
|
if (flushTimeout) {
|
|
flushCbs.add(cb);
|
|
return;
|
|
}
|
|
|
|
cb();
|
|
}
|
|
|
|
const scheduleFlush = () => {
|
|
if (flushTimeout) return;
|
|
flushTimeout = setTimeout(flushTextDeltas, 1000 / TARGET_UPDATES_PER_SECOND);
|
|
};
|
|
|
|
const resetTextDeltas = () => {
|
|
if (flushTimeout !== undefined) clearTimeout(flushTimeout);
|
|
flushTimeout = undefined;
|
|
textDeltaBuffer.clear();
|
|
flushCbs.clear();
|
|
};
|
|
|
|
const scheduleReconnect = (forceSnapshot = false) => {
|
|
sse?.close();
|
|
sse = undefined;
|
|
if (disposed || reconnectTimeout !== undefined) return;
|
|
reconnectTimeout = setTimeout(() => {
|
|
reconnectTimeout = undefined;
|
|
void connectSSE(forceSnapshot);
|
|
}, 1500);
|
|
};
|
|
|
|
const processEvent = (event: Event) => {
|
|
if (!data.value) return;
|
|
|
|
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: payload.parts ?? [], attachments: payload.attachments ?? [] });
|
|
} else {
|
|
// replace the existing data with the new data
|
|
const index = data.value.messages.findIndex(m => m.id === payload.id);
|
|
if (index !== -1 && !event.replay) {
|
|
const existing = data.value.messages[index]!;
|
|
data.value.messages[index] = {
|
|
...payload,
|
|
parts: payload.parts ?? [],
|
|
attachments: payload.attachments ?? existing.attachments ?? [],
|
|
};
|
|
}
|
|
}
|
|
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 = [];
|
|
if (!msg.parts.some(part => part.id === payload.part.id)) {
|
|
msg.parts.push(payload.part);
|
|
}
|
|
}
|
|
break;
|
|
}
|
|
|
|
case 'text-delta': {
|
|
const msg = data.value.messages.find(m => m.id === payload.messageId);
|
|
if (!msg) break;
|
|
const part = msg.parts?.find(part => part.id === payload.partId);
|
|
if (!part || part.finished) break;
|
|
|
|
if (!textDeltaBuffer.has(payload.messageId)) {
|
|
textDeltaBuffer.set(payload.messageId, new Map());
|
|
}
|
|
const messageBuffer = textDeltaBuffer.get(payload.messageId)!;
|
|
const existing = messageBuffer.get(payload.partId)?.content ?? part.content ?? '';
|
|
if (typeof payload.offset === 'number' && payload.offset > existing.length) {
|
|
scheduleReconnect(true);
|
|
break;
|
|
}
|
|
const content = appendTextDelta(existing, payload.content, payload.offset);
|
|
if (content === existing) break;
|
|
messageBuffer.set(payload.partId, { content, lastUpdatedAt: payload.lastUpdatedAt });
|
|
|
|
scheduleFlush();
|
|
break;
|
|
}
|
|
|
|
case 'text-delete': {
|
|
const msg = data.value.messages.find(m => m.id === payload.messageId);
|
|
if (msg?.parts) msg.parts = msg.parts.filter(part => part.id !== payload.partId);
|
|
textDeltaBuffer.get(payload.messageId)?.delete(payload.partId);
|
|
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 && part.type !== 'tool-call') {
|
|
nextFlush(() => {
|
|
part.content = payload.content;
|
|
if (payload.lastUpdatedAt) {
|
|
part.lastUpdatedAt = payload.lastUpdatedAt;
|
|
}
|
|
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 = [];
|
|
}
|
|
const part = { ...payload.part };
|
|
if (!part.createdAt) {
|
|
part.createdAt = new Date().toISOString();
|
|
}
|
|
if (!msg.parts.some(existing => existing.id === part.id)) {
|
|
msg.parts.push(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;
|
|
if (payload.status === 'completed' || payload.status === 'failed') {
|
|
part.lastUpdatedAt = payload.lastUpdatedAt || new Date().toISOString();
|
|
part.finished = true;
|
|
}
|
|
}
|
|
}
|
|
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 (forceSnapshot = false) => {
|
|
if (disposed) return;
|
|
const attempt = ++connectionAttempt;
|
|
const id = unref(topicId);
|
|
sse?.close();
|
|
if (flushTimeout !== undefined) {
|
|
clearTimeout(flushTimeout);
|
|
flushTextDeltas();
|
|
}
|
|
let offeredSnapshot: string | undefined;
|
|
const query = new URLSearchParams();
|
|
if (!forceSnapshot && data.value?.id === id && globalThis.crypto?.subtle) {
|
|
try {
|
|
offeredSnapshot = serializeTopicSnapshot(toRaw(data.value));
|
|
const digest = await crypto.subtle.digest('SHA-256', new TextEncoder().encode(offeredSnapshot));
|
|
query.set('snapshot', Array.from(new Uint8Array(digest), byte => byte.toString(16).padStart(2, '0')).join(''));
|
|
} catch {
|
|
// Fingerprinting is an optimization; a full snapshot remains safe.
|
|
offeredSnapshot = undefined;
|
|
}
|
|
}
|
|
if (disposed || attempt !== connectionAttempt) return;
|
|
|
|
const source = new EventSource(`/api/topic/${id}?${query}`);
|
|
sse = source;
|
|
source.addEventListener('error', () => {
|
|
if (sse === source) scheduleReconnect();
|
|
});
|
|
|
|
source.addEventListener('message', (e) => {
|
|
if (disposed || sse !== source) return;
|
|
const raw = JSON.parse(e.data);
|
|
const events: Event[] = Array.isArray(raw) ? raw : [raw];
|
|
|
|
for (const event of events) {
|
|
const { type, payload } = event;
|
|
if (!type || !payload) continue;
|
|
|
|
if (type === 'synced') {
|
|
// Optimistic edits can occur while the handshake is in flight.
|
|
// Only retain the exact state that we offered to the server.
|
|
if (offeredSnapshot !== serializeTopicSnapshot(toRaw(data.value))) {
|
|
scheduleReconnect(true);
|
|
return;
|
|
}
|
|
offeredSnapshot = undefined;
|
|
markConnected();
|
|
continue;
|
|
}
|
|
|
|
if (type === 'initial_state') {
|
|
resetTextDeltas();
|
|
data.value = payload;
|
|
offeredSnapshot = undefined;
|
|
markConnected();
|
|
continue;
|
|
}
|
|
|
|
processEvent(event);
|
|
}
|
|
});
|
|
|
|
await connectionReady;
|
|
};
|
|
|
|
|
|
if (connect) {
|
|
onBeforeUnmount(() => {
|
|
disposed = true;
|
|
connectionAttempt++;
|
|
if (reconnectTimeout !== undefined) clearTimeout(reconnectTimeout);
|
|
resetTextDeltas();
|
|
markConnected();
|
|
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 {
|
|
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 {
|
|
const key = await crypto.subtle.importKey(
|
|
"jwk",
|
|
JSON.parse(window.localStorage.getItem("encryptionKey")!),
|
|
"AES-GCM",
|
|
false,
|
|
["encrypt", "decrypt"]
|
|
)
|
|
|
|
return Ok(await decrypt(
|
|
key,
|
|
base64ToUint8Array(provider.config.apiKey)
|
|
));
|
|
} catch (error) {
|
|
console.error('Failed to decrypt provider API key:', error);
|
|
return Err(ChatErrorType.FailedToDecryptProviderApiKey);
|
|
}
|
|
}
|
|
|
|
const sendMessage = async (
|
|
baseMessage: BaseMessage,
|
|
onRequest?: () => void,
|
|
role: 'user' | 'assistant' = 'user',
|
|
): Promise<Result<void, ChatErrorType>> => {
|
|
const { user } = useAuth();
|
|
if (!user.value) return Err(ChatErrorType.NoUser);
|
|
|
|
try {
|
|
const message = {
|
|
id: nanoid(),
|
|
role,
|
|
}
|
|
|
|
await $fetch(`/api/topic/${unref(topicId)}/message`, {
|
|
method: 'POST',
|
|
body: {
|
|
message: {
|
|
...message,
|
|
content: baseMessage.content,
|
|
fileIds: baseMessage.fileIds,
|
|
}
|
|
},
|
|
onRequest() {
|
|
if (data.value) {
|
|
switch (role) {
|
|
case 'user': {
|
|
const optimisticAttachments = baseMessage.files?.map(file => ({
|
|
id: nanoid(),
|
|
userId: user.value!.id,
|
|
topicId: unref(topicId),
|
|
messageId: message.id,
|
|
fileId: file.id,
|
|
createdAt: new Date(),
|
|
file: {
|
|
id: file.id,
|
|
userId: user.value!.id,
|
|
name: file.name,
|
|
mimeType: file.mimeType,
|
|
size: file.size ?? 0,
|
|
url: file.url,
|
|
createdAt: new Date(),
|
|
},
|
|
})) ?? [];
|
|
|
|
data.value!.messages.push({
|
|
topicId: unref(topicId),
|
|
userId: user.value!.id,
|
|
attachments: optimisticAttachments,
|
|
parts: undefined,
|
|
generation: null,
|
|
parentMessageId: null,
|
|
generationId: null,
|
|
activeChildId: null,
|
|
children: [],
|
|
deleted: false,
|
|
content: baseMessage.content,
|
|
createdAt: new Date(),
|
|
updatedAt: new Date(),
|
|
...message
|
|
});
|
|
break;
|
|
}
|
|
case 'assistant': {
|
|
data.value!.messages.push({
|
|
topicId: unref(topicId),
|
|
userId: user.value!.id,
|
|
// TODO
|
|
attachments: [],
|
|
generation: null,
|
|
parentMessageId: null,
|
|
generationId: null,
|
|
activeChildId: null,
|
|
children: [],
|
|
deleted: false,
|
|
content: null,
|
|
createdAt: new Date(),
|
|
updatedAt: new Date(),
|
|
parts: [
|
|
{
|
|
id: 'bogus-id',
|
|
userId: user.value!.id,
|
|
topicId: unref(topicId),
|
|
messageId: message.id,
|
|
type: 'text',
|
|
content: baseMessage.content,
|
|
finished: true,
|
|
createdAt: new Date(),
|
|
lastUpdatedAt: new Date(),
|
|
providerOptions: null,
|
|
toolCallId: null,
|
|
toolCall: null,
|
|
}
|
|
],
|
|
...message
|
|
});
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
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, parentMessageId?: string, agentIdOverride?: string) => {
|
|
const providerApiKeyRes = await getProviderAPIKey(model.provider);
|
|
if (providerApiKeyRes.ok === false) {
|
|
return providerApiKeyRes;
|
|
}
|
|
const providerApiKey = providerApiKeyRes.data;
|
|
|
|
const { agents } = await useAgents();
|
|
|
|
const agent = agents.value.find(agent => agent.id === (agentIdOverride || topic.value?.agentId));
|
|
if (!agent) {
|
|
return Err(ChatErrorType.NoAgent);
|
|
}
|
|
|
|
await $fetch(`/api/topic/${unref(topicId)}/chat`, {
|
|
method: 'POST',
|
|
body: {
|
|
modelId: model.id,
|
|
providerApiKey,
|
|
parentMessageId,
|
|
},
|
|
});
|
|
}
|
|
|
|
/**
|
|
*
|
|
* @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,
|
|
topicMessages: MessageEntity[],
|
|
model: ModelWithProvider
|
|
): Promise<Result<void, ChatErrorType>> => {
|
|
const targetMessage = topicMessages.find(message => message.id === messageId);
|
|
if (!targetMessage) {
|
|
return Err(ChatErrorType.NoMessage);
|
|
}
|
|
let targetMessageIndex = topicMessages.indexOf(targetMessage);
|
|
|
|
let parentMessageId = undefined;
|
|
let focusedMessages: MessageEntity[] | undefined;
|
|
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;
|
|
}
|
|
|
|
startGeneration(model, parentMessageId);
|
|
|
|
return Ok(undefined);
|
|
}
|
|
|
|
const patchMessageLocally = (id: string, updates: Partial<Message>) => {
|
|
if (!data.value) return;
|
|
|
|
data.value.messages = data.value.messages.map(m =>
|
|
m.id === id ? { ...m, ...updates } : m
|
|
);
|
|
}
|
|
|
|
const deleteMessagesFocusedDescendant = async (messageId: string): Promise<Result<void, Error>> => {
|
|
if (!data.value) {
|
|
return Err(new Error("No data"));
|
|
}
|
|
|
|
const msg = data.value.messages.find((m) => {
|
|
return m.id === messageId;
|
|
});
|
|
if (!msg) {
|
|
return Err(new Error("Message not found"));
|
|
}
|
|
|
|
const previousMessages = data.value.messages;
|
|
|
|
const res = await $fetch(`/api/messages/${messageId}`, {
|
|
method: "DELETE",
|
|
onRequest() {
|
|
if (!data.value) {
|
|
return;
|
|
}
|
|
|
|
if (msg.parentMessageId) {
|
|
const parent = data.value.messages.find((m) => {
|
|
return m.id === msg.parentMessageId;
|
|
});
|
|
data.value.messages = data.value.messages.filter((m) => {
|
|
return m.id !== msg.id;
|
|
});
|
|
if (parent) {
|
|
const remainingChildren = data.value.messages.filter((m) => {
|
|
return m.parentMessageId === parent.id;
|
|
});
|
|
if (parent.deleted && remainingChildren.length === 0) {
|
|
data.value.messages = data.value.messages.filter((m) => {
|
|
return m.id !== parent.id;
|
|
});
|
|
} else if (parent.activeChildId === msg.id) {
|
|
parent.activeChildId = remainingChildren.length > 0
|
|
? (parent.deleted ? remainingChildren[0]!.id : remainingChildren[remainingChildren.length - 1]!.id)
|
|
: null;
|
|
}
|
|
}
|
|
return;
|
|
}
|
|
|
|
const children = data.value.messages.filter((m) => {
|
|
return m.parentMessageId === msg.id;
|
|
});
|
|
if (children.length === 0) {
|
|
data.value.messages = data.value.messages.filter((m) => {
|
|
return m.id !== msg.id;
|
|
});
|
|
return;
|
|
}
|
|
|
|
msg.deleted = true;
|
|
msg.activeChildId = children[0]?.id ?? null;
|
|
},
|
|
onResponseError() {
|
|
if (data.value) {
|
|
data.value.messages = previousMessages;
|
|
}
|
|
},
|
|
});
|
|
|
|
if (!res.ok) {
|
|
return Err(new Error("Failed to delete message"));
|
|
}
|
|
|
|
return Ok(undefined);
|
|
};
|
|
|
|
return {
|
|
topic,
|
|
sendMessage,
|
|
startGeneration,
|
|
regenerateMessage,
|
|
patchMessageLocally,
|
|
deleteMessagesFocusedDescendant
|
|
};
|
|
}
|