365 lines
11 KiB
TypeScript
365 lines
11 KiB
TypeScript
import { BroadcastChannel } from 'broadcast-channel';
|
|
import EventEmitter from 'events';
|
|
import MiniSearch, { SearchResult } from 'minisearch'
|
|
import { v4 as uuidv4 } from 'uuid';
|
|
import { Chat, deserializeChat, getOpenAIMessageFromMessage, Message, Parameters, serializeChat, UserSubmittedMessage } from './types';
|
|
import { MessageTree } from './message-tree';
|
|
import { createStreamingChatCompletion } from './openai';
|
|
import { createTitle } from './titles';
|
|
import { ellipsize, sleep } from './utils';
|
|
import * as idb from './idb';
|
|
import { selectMessagesToSendSafely } from './tokenizer';
|
|
|
|
export const channel = new BroadcastChannel('chats');
|
|
|
|
export class ChatManager extends EventEmitter {
|
|
public chats = new Map<string, Chat>();
|
|
public search = new Search(this.chats);
|
|
private loaded = false;
|
|
private changed = false;
|
|
private activeReplies = new Map<string, Message>();
|
|
|
|
constructor() {
|
|
super();
|
|
this.load();
|
|
|
|
this.on('update', () => {
|
|
this.changed = true;
|
|
});
|
|
|
|
channel.onmessage = (message: {
|
|
type: 'chat-update',
|
|
data: string,
|
|
}) => {
|
|
const chat = deserializeChat(message.data);
|
|
const id = chat.id;
|
|
this.chats.set(id, chat);
|
|
this.emit(id);
|
|
};
|
|
|
|
(async () => {
|
|
while (true) {
|
|
await sleep(100);
|
|
|
|
if (this.loaded && this.changed) {
|
|
this.changed = false;
|
|
await this.save();
|
|
}
|
|
}
|
|
})();
|
|
}
|
|
|
|
public async createChat(): Promise<string> {
|
|
const id = uuidv4();
|
|
|
|
const chat: Chat = {
|
|
id,
|
|
messages: new MessageTree(),
|
|
created: Date.now(),
|
|
updated: Date.now(),
|
|
};
|
|
|
|
this.chats.set(id, chat);
|
|
this.search.update(chat);
|
|
channel.postMessage({ type: 'chat-update', data: serializeChat(chat) });
|
|
|
|
return id;
|
|
}
|
|
|
|
public async sendMessage(message: UserSubmittedMessage) {
|
|
const chat = this.chats.get(message.chatID);
|
|
|
|
if (!chat) {
|
|
throw new Error('Chat not found');
|
|
}
|
|
|
|
const newMessage: Message = {
|
|
id: uuidv4(),
|
|
parentID: message.parentID,
|
|
chatID: chat.id,
|
|
timestamp: Date.now(),
|
|
role: 'user',
|
|
content: message.content,
|
|
done: true,
|
|
};
|
|
|
|
chat.messages.addMessage(newMessage);
|
|
chat.updated = Date.now();
|
|
|
|
this.emit(chat.id);
|
|
this.emit('messages', [newMessage]);
|
|
channel.postMessage({ type: 'chat-update', data: serializeChat(chat) });
|
|
|
|
const messages: Message[] = message.parentID
|
|
? chat.messages.getMessageChainTo(message.parentID)
|
|
: [];
|
|
messages.push(newMessage);
|
|
|
|
await this.getReply(messages, message.requestedParameters);
|
|
}
|
|
|
|
public async regenerate(message: Message, requestedParameters: Parameters) {
|
|
const chat = this.chats.get(message.chatID);
|
|
|
|
if (!chat) {
|
|
throw new Error('Chat not found');
|
|
}
|
|
|
|
const messages: Message[] = message.parentID
|
|
? chat.messages.getMessageChainTo(message.parentID)
|
|
: [];
|
|
|
|
await this.getReply(messages, requestedParameters);
|
|
}
|
|
|
|
private async getReply(messages: Message[], requestedParameters: Parameters) {
|
|
const latestMessage = messages[messages.length - 1];
|
|
const chat = this.chats.get(latestMessage.chatID);
|
|
|
|
if (!chat) {
|
|
throw new Error('Chat not found');
|
|
}
|
|
|
|
const reply: Message = {
|
|
id: uuidv4(),
|
|
parentID: latestMessage.id,
|
|
chatID: latestMessage.chatID,
|
|
timestamp: Date.now(),
|
|
role: 'assistant',
|
|
content: '',
|
|
done: false,
|
|
};
|
|
this.activeReplies.set(reply.id, reply);
|
|
|
|
chat.messages.addMessage(reply);
|
|
chat.updated = Date.now();
|
|
|
|
this.emit(chat.id);
|
|
channel.postMessage({ type: 'chat-update', data: serializeChat(chat) });
|
|
|
|
const messagesToSend = selectMessagesToSendSafely(messages.map(getOpenAIMessageFromMessage));
|
|
|
|
const { emitter, cancel } = await createStreamingChatCompletion(messagesToSend, requestedParameters);
|
|
|
|
let lastChunkReceivedAt = Date.now();
|
|
|
|
const onError = () => {
|
|
if (reply.done) {
|
|
return;
|
|
}
|
|
clearInterval(timer);
|
|
cancel();
|
|
reply.content += "\n\nI'm sorry, I'm having trouble connecting to OpenAI. Please make sure you've entered your OpenAI API key correctly and try again.";
|
|
reply.content = reply.content.trim();
|
|
reply.done = true;
|
|
this.activeReplies.delete(reply.id);
|
|
chat.messages.updateMessage(reply);
|
|
chat.updated = Date.now();
|
|
this.emit(chat.id);
|
|
this.emit('messages', [reply]);
|
|
channel.postMessage({ type: 'chat-update', data: serializeChat(chat) });
|
|
};
|
|
|
|
let timer = setInterval(() => {
|
|
const sinceLastChunk = Date.now() - lastChunkReceivedAt;
|
|
if (sinceLastChunk > 10000 && !reply.done) {
|
|
onError();
|
|
}
|
|
}, 2000);
|
|
|
|
emitter.on('error', () => {
|
|
if (!reply.content && !reply.done) {
|
|
lastChunkReceivedAt = Date.now();
|
|
onError();
|
|
}
|
|
});
|
|
|
|
emitter.on('data', (data: string) => {
|
|
if (reply.done) {
|
|
cancel();
|
|
return;
|
|
}
|
|
lastChunkReceivedAt = Date.now();
|
|
reply.content = data;
|
|
chat.messages.updateMessage(reply);
|
|
this.emit(chat.id);
|
|
channel.postMessage({ type: 'chat-update', data: serializeChat(chat) });
|
|
});
|
|
|
|
emitter.on('done', async () => {
|
|
if (reply.done) {
|
|
return;
|
|
}
|
|
clearInterval(timer);
|
|
lastChunkReceivedAt = Date.now();
|
|
reply.done = true;
|
|
this.activeReplies.delete(reply.id);
|
|
chat.messages.updateMessage(reply);
|
|
chat.updated = Date.now();
|
|
this.emit(chat.id);
|
|
this.emit('messages', [reply]);
|
|
this.emit('update');
|
|
channel.postMessage({ type: 'chat-update', data: serializeChat(chat) });
|
|
setTimeout(() => this.search.update(chat), 500);
|
|
|
|
if (!chat.title) {
|
|
chat.title = await createTitle(chat, requestedParameters.apiKey);
|
|
if (chat.title) {
|
|
this.emit(chat.id);
|
|
this.emit('title', chat.id, chat.title);
|
|
this.emit('update');
|
|
channel.postMessage({ type: 'chat-update', data: serializeChat(chat) });
|
|
setTimeout(() => this.search.update(chat), 500);
|
|
}
|
|
}
|
|
});
|
|
}
|
|
|
|
private async save() {
|
|
const serialized = Array.from(this.chats.values())
|
|
.map((c) => {
|
|
const serialized = { ...c } as any;
|
|
serialized.messages = c.messages.serialize();
|
|
return serialized;
|
|
});
|
|
await idb.set('chats', serialized);
|
|
}
|
|
|
|
public cancelReply(id: string) {
|
|
const reply = this.activeReplies.get(id);
|
|
if (reply) {
|
|
reply.done = true;
|
|
this.activeReplies.delete(reply.id);
|
|
|
|
const chat = this.chats.get(reply.chatID);
|
|
const message = chat?.messages.get(id);
|
|
if (message) {
|
|
message.done = true;
|
|
this.emit(reply.chatID);
|
|
this.emit('messages', [reply]);
|
|
this.emit('update');
|
|
channel.postMessage({ type: 'chat-update', data: serializeChat(chat!) });
|
|
}
|
|
} else {
|
|
console.log('failed to find reply');
|
|
}
|
|
}
|
|
|
|
private async load() {
|
|
const serialized = await idb.get('chats');
|
|
if (serialized) {
|
|
for (const chat of serialized) {
|
|
const messages = new MessageTree();
|
|
for (const m of chat.messages) {
|
|
messages.addMessage(m);
|
|
}
|
|
chat.messages = messages;
|
|
this.loadChat(chat);
|
|
}
|
|
this.emit('update');
|
|
}
|
|
this.loaded = true;
|
|
}
|
|
|
|
public loadChat(chat: Chat) {
|
|
if (!chat?.id) {
|
|
return;
|
|
}
|
|
|
|
const existing = this.chats.get(chat.id);
|
|
if (existing && existing.title && !chat.title) {
|
|
chat.title = existing.title;
|
|
}
|
|
|
|
chat.created = chat.messages.first?.timestamp || 0;
|
|
chat.updated = chat.messages.mostRecentLeaf().timestamp;
|
|
|
|
this.chats.set(chat.id, chat);
|
|
this.search.update(chat);
|
|
this.emit(chat.id);
|
|
}
|
|
|
|
public get(id: string): Chat | undefined {
|
|
return this.chats.get(id);
|
|
}
|
|
}
|
|
|
|
export class Search {
|
|
private index = new MiniSearch({
|
|
fields: ['value'],
|
|
storeFields: ['id', 'value'],
|
|
});
|
|
|
|
constructor(private chats: Map<string, Chat>) {
|
|
}
|
|
|
|
public update(chat: Chat) {
|
|
const messages = chat.messages.serialize()
|
|
.map((m: Message) => m.content)
|
|
.join('\n\n');
|
|
const doc = {
|
|
id: chat.id,
|
|
value: chat.title + '\n\n' + messages,
|
|
};
|
|
if (!this.index.has(chat.id)) {
|
|
this.index.add(doc);
|
|
} else {
|
|
this.index.replace(doc);
|
|
}
|
|
}
|
|
|
|
public query(query: string) {
|
|
if (!query?.trim().length) {
|
|
const searchResults = Array.from(this.chats.values())
|
|
.sort((a, b) => b.updated - a.updated)
|
|
.slice(0, 10);
|
|
const results = this.processSearchResults(searchResults);
|
|
return results;
|
|
}
|
|
|
|
let searchResults = this.index.search(query, { fuzzy: 0.2 });
|
|
let output = this.processSearchResults(searchResults);
|
|
|
|
if (!output.length) {
|
|
searchResults = this.index.search(query, { prefix: true });
|
|
output = this.processSearchResults(searchResults);
|
|
}
|
|
|
|
return output;
|
|
}
|
|
|
|
private processSearchResults(searchResults: SearchResult[] | Chat[]) {
|
|
const output: any[] = [];
|
|
for (const item of searchResults) {
|
|
const chatID = item.id;
|
|
let chat = this.chats.get(chatID);
|
|
if (!chat) {
|
|
continue;
|
|
}
|
|
|
|
chat = { ...chat };
|
|
|
|
let description = chat.messages?.first?.content || '';
|
|
description = ellipsize(description, 400);
|
|
|
|
if (!chat.title) {
|
|
chat.title = ellipsize(description, 100);
|
|
}
|
|
|
|
if (!chat.title || !description) {
|
|
continue;
|
|
}
|
|
|
|
output.push({
|
|
chatID,
|
|
title: chat.title,
|
|
description,
|
|
});
|
|
}
|
|
return output;
|
|
}
|
|
}
|
|
|
|
const chatManager = new ChatManager();
|
|
export default chatManager;
|