v0.1.0
This commit is contained in:
Cogent Apps
GitHub
261
src/chat-manager.ts
Normal file
261
src/chat-manager.ts
Normal file
@@ -0,0 +1,261 @@
|
||||
import { BroadcastChannel } from 'broadcast-channel';
|
||||
import EventEmitter from 'events';
|
||||
import MiniSearch, { SearchResult } from 'minisearch'
|
||||
import { v4 as uuidv4 } from 'uuid';
|
||||
import { Chat, getOpenAIMessageFromMessage, Message, 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';
|
||||
|
||||
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;
|
||||
|
||||
constructor() {
|
||||
super();
|
||||
this.load();
|
||||
|
||||
this.on('update', () => {
|
||||
this.changed = true;
|
||||
});
|
||||
|
||||
channel.onmessage = (message: {
|
||||
type: 'chat-update',
|
||||
data: Chat,
|
||||
}) => {
|
||||
const id = message.data.id;
|
||||
this.chats.set(id, message.data);
|
||||
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: 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,
|
||||
};
|
||||
|
||||
const reply: Message = {
|
||||
id: uuidv4(),
|
||||
parentID: newMessage.id,
|
||||
chatID: chat.id,
|
||||
timestamp: Date.now(),
|
||||
role: 'assistant',
|
||||
content: '',
|
||||
done: false,
|
||||
};
|
||||
|
||||
chat.messages.addMessage(newMessage);
|
||||
chat.messages.addMessage(reply);
|
||||
chat.updated = Date.now();
|
||||
|
||||
this.emit(chat.id);
|
||||
this.emit('messages', [newMessage]);
|
||||
channel.postMessage({ type: 'chat-update', data: chat });
|
||||
|
||||
const messages: Message[] = message.parentID
|
||||
? chat.messages.getMessageChainTo(message.parentID)
|
||||
: [];
|
||||
messages.push(newMessage);
|
||||
|
||||
const response = await createStreamingChatCompletion(messages.map(getOpenAIMessageFromMessage),
|
||||
message.requestedParameters);
|
||||
|
||||
response.on('error', () => {
|
||||
if (!reply.content) {
|
||||
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;
|
||||
chat.messages.updateMessage(reply);
|
||||
chat.updated = Date.now();
|
||||
this.emit(chat.id);
|
||||
this.emit('messages', [reply]);
|
||||
channel.postMessage({ type: 'chat-update', data: chat });
|
||||
}
|
||||
})
|
||||
|
||||
response.on('data', (data: string) => {
|
||||
reply.content = data;
|
||||
chat.messages.updateMessage(reply);
|
||||
this.emit(chat.id);
|
||||
channel.postMessage({ type: 'chat-update', data: chat });
|
||||
});
|
||||
|
||||
response.on('done', async () => {
|
||||
reply.done = true;
|
||||
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: chat });
|
||||
setTimeout(() => this.search.update(chat), 500);
|
||||
|
||||
if (!chat.title) {
|
||||
chat.title = await createTitle(chat, message.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: 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);
|
||||
}
|
||||
|
||||
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;
|
||||
}
|
||||
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;
|
||||
const chat = this.chats.get(chatID);
|
||||
if (!chat) {
|
||||
continue;
|
||||
}
|
||||
|
||||
let description = chat.messages?.first?.content || '';
|
||||
description = ellipsize(description, 400);
|
||||
if (!chat.title || !description) {
|
||||
continue;
|
||||
}
|
||||
|
||||
output.push({
|
||||
chatID,
|
||||
title: chat.title,
|
||||
description,
|
||||
});
|
||||
}
|
||||
return output;
|
||||
}
|
||||
}
|
||||
|
||||
export default new ChatManager();
|
||||
Reference in New Issue
Block a user