diff --git a/backend/src/app_mixins/chat.py b/backend/src/app_mixins/chat.py index 234dd8d..c3a9f2c 100644 --- a/backend/src/app_mixins/chat.py +++ b/backend/src/app_mixins/chat.py @@ -92,17 +92,6 @@ class ChatMixin: # TODO: Stop repetition. # There's db.Message and chat.Message, also should be combined somehow, or wrap db always. return [message_json(msg) for msg in messages] - - async def get_chat_thread_heads(self, chat_id, limit_message_id=None): - chat = await self.get_chat(chat_id) - messages = await chat.load_thread_heads(limit_message_id) - return [message_json(msg) for msg in messages] - - async def get_thread_messages(self, chat_id, thread_id, limit_message_id=None): - chat = await self.get_chat(chat_id) - thread = chat.get_thread(thread_id) - messages = await thread.load_page(limit_message_id) - return [message_json(msg) for msg in messages] async def create_chat_thread(self, chat_id, parent_id=None, name=None): if tmp.is_tmp_id(chat_id): @@ -116,6 +105,29 @@ class ChatMixin: print(f"Created thread #{thread.id}: {thread.name}") return thread.id + async def get_chat_threads(self, chat_id): + threads = await db.Thread.filter(chat_id=chat_id).all() + return [{ + 'id': it.id, + 'parentId': it.parent_id, + 'name': it.name, + 'description': it.description, + 'firstMessageId': it.first_message_id, + 'lastMessageId': it.last_message_id, + 'messageCount': it.message_count, + } for it in threads] + + async def get_thread_messages(self, chat_id, thread_id, limit_message_id=None): + chat = await self.get_chat(chat_id) + thread = chat.get_thread(thread_id) + messages = await thread.load_page(limit_message_id) + return [message_json(msg) for msg in messages] + + async def get_chat_thread_heads(self, chat_id, limit_message_id=None): + chat = await self.get_chat(chat_id) + messages = await chat.load_thread_heads(limit_message_id) + return [message_json(msg) for msg in messages] + async def fork_chat_thread(self, thread_id, message_id=None, name=None): """Move all messages from message_id onwards into a new child thread. If none is given, this is basically create_chat_thread with an implicit chat id.""" @@ -125,6 +137,15 @@ class ChatMixin: if message_id is not None: await db.Message.filter(thread_id=thread_id, id__gte=message_id).update(thread_id=new_id) + # Notify users. what a heavy operation! + + chat = await db.Chat.get(id=thread.chat_id) + messages = await db.Message.filter(thread_id=new_id).all() + print("Chat: " + repr(chat)) + print("Fork message list: " + repr(messages)) + for msg in messages: + self.notify_message_to_user(chat.user_id, msg, True) + return new_id async def add_agent_to_chat(self, chat_id, agent_id): diff --git a/backend/src/server/websocket.py b/backend/src/server/websocket.py index bb4f36a..3df1a43 100644 --- a/backend/src/server/websocket.py +++ b/backend/src/server/websocket.py @@ -234,6 +234,11 @@ def register(server, app, auth_manager): ) await send({'message': "chat messages retrieved", 'messages': messages_data, 'nonce': nonce}) + elif action == "get-chat-threads": + await send({'message': "threads retrieved", 'nonce': nonce, + 'threads': await app.get_chat_threads(body['chatId']) + }) + elif action == "get-chat-thread-heads": messages_data = await app.get_chat_thread_heads( body['chatId'], diff --git a/frontend/src/chat.js b/frontend/src/chat.js index 0810855..bc199af 100644 --- a/frontend/src/chat.js +++ b/frontend/src/chat.js @@ -4,11 +4,19 @@ import * as backend from './backend'; import * as runs from './runs'; import * as ui from './ui'; +// Known issues: +// - thread loading is fetching all messages everytime +// - fork heads initially shows preview of last message +// - no visual distinction between thread (top level) and forks (children) + // === Thread Helpers === // const THREAD_STALE_MS = 10 * 60 * 1000; function messageTime (message) { + // temporary id override. closely matches but not the same + return message.id; + // Accept the common backend timestamp field names and formats const t = message.timestamp ?? message.time ?? message.createdAt; if (t == null) return null; @@ -22,31 +30,73 @@ function isMessageStale (message) { return time == null || (Date.now() - time) > THREAD_STALE_MS; } -function earliestMessage (...messages) { - let minId = null; +function selectEarliest (messages) { + let selected = null; let minTime = null; - let useTime = true; - for (let msg of messages) { + for (let i = 0; i < messages.length; i++) { + let msg = messages[i]; if (!msg) continue; - if (!minId || msg.id < minId.id) { - minId = msg; - } let msgTime = messageTime(msg); - if (msgTime == null) { - useTime = false; - } else if (!minTime || msgTime < messageTime(minTime)) { - minTime = msg; + if (selected == null || msgTime < minTime) { + selected = i; + minTime = msgTime; } } - return useTime ? minTime : minId; + return selected } export class Thread { - constructor(obj) { - Object.assign(this, {messages: []}, obj || {}); + constructor (obj, session) { + Object.assign(this, { + messages: [], + children: [], + messageCount: 0, + session + }, obj || {}); + + session.threadMap[this.id] = this; + this.attach(this.parentId); + } + + async _container (parent) { + return parent ? + (await this.session.getThread( + parent + )).children + : this.session.threads; + } + + async attach (parent) { + let arr = await this._container(parent); + arr.push(this); + } + + async detach (parent) { + let arr = await this._container(parent); + let index = arr.indexOf(this); + if (index >= 0) arr.splice(index, 1); + else console.warn("Thread not attached", {thread, parent}); + } + + update (obj) { + let initialParent = this.parentId; + + Object.assign(this, obj); + if (this.parentId != initialParent) { + this.detach(initialParent); + this.attach(this.parentId); + } + } + + getFirstMessage () { + if (this.messages.length) { + return this.messages[0]; + } else { + // TODO: check children + } } getLastMessage () { @@ -81,6 +131,11 @@ export class ChatSession { this.chatEntries = ref([]); // real chat data this.threads = []; + // display all the threads in the current thread ancestry + // for now, manually enabled + this.activeThreadIds = []; + // fast access map + this.threadMap = {}; this.editInfo = computed(() => { const parts = [this.creatingThread.value ? 'New Thread' : `Thread ${this.currentThreadId.value}`]; @@ -138,6 +193,7 @@ export class ChatSession { this.chatEntries.value = []; this.threads = []; + this.activeThreadIds = []; // Create an agent name dictionary const agents = await backend.getAgents(); @@ -157,9 +213,13 @@ export class ChatSession { this._unsubscribeRuns = runs.onChatMessage(chatId, event => this.trackRun(event)); // Load threads - const {messages: heads} = await backend.sendAction("get-chat-thread-heads", {chatId}); + let {threads} = await backend.sendAction("get-chat-threads", {chatId}); + for (const it of threads) { + this.getThread(it.id).update(it); + } - // Always good to duplicate when assigning to a ref + // And their head messages + let {messages: heads} = await backend.sendAction("get-chat-thread-heads", {chatId}); for (const it of heads) { this.getThread(it.threadId).messages.push(it); } @@ -167,7 +227,7 @@ export class ChatSession { // Continue the last thread only if it is recent (under 10 minutes old); // otherwise start in "New Thread" state const lastMessage = heads[heads.length - 1]; - if (lastMessage && lastMessage.threadId != null && !isMessageStale(lastMessage)) { + if (lastMessage && !isMessageStale(lastMessage)) { this.activateThread(lastMessage.threadId); } @@ -215,6 +275,7 @@ export class ChatSession { unsent.id = it.id; unsent.threadId = it.threadId; this.addMessage(unsent); + this.computeChatEntries(); return; } else { // Remove from unsent list and add message as received @@ -232,6 +293,7 @@ export class ChatSession { sender: it.sender, content: it.content, threadId: it.threadId, + time: it.time, }); this.computeChatEntries(); @@ -383,6 +445,7 @@ export class ChatSession { // Replace an already-known message in place let existingIndex = list.findIndex(it => it.id == id); if (existingIndex >= 0) { + // TODO: Not copying? does it matter? list[existingIndex] = msg; continue; } @@ -423,29 +486,21 @@ export class ChatSession { getThread (threadId) { // `null` (orphan messages) serializes to the "null" key - let thread = this.threads.find(it => it.id ==threadId); - if (!thread) { - thread = new Thread({ - id: threadId, - loaded: false, - time: null, // TODO - messages: [], - }); - this.threads.push(thread); - } - return thread; + return this.threadMap[threadId] || new Thread({id: threadId}, this); } activateThread (threadId) { if (threadId == null) return; // "New Thread" placeholder isn't clickable // Continuing a thread: it becomes the one open thread - this.currentThreadId.value = threadId; this.creatingThread.value = false; + this.currentThreadId.value = threadId; + this.activeThreadIds = [threadId]; // Load thread messages asynchronously backend.sendAction("get-thread-messages", { chatId: this.chatId, threadId }).then(res => { + // TODO: reassigning all fields? this.addMessageBulk(res.messages.map(it => ({ id: it.id, type: it.kind, @@ -454,9 +509,17 @@ export class ChatSession { sender: it.sender, content: it.content, threadId: it.threadId, + time: it.time, }))); }); + // Activate threads in the ancestry + let thread = this.getThread(threadId); + while (thread.parentId) { + thread = this.getThread(thread.parentId); + this.activeThreadIds.push(thread.id); + } + this.computeChatEntries(); this.rescroll(); } @@ -465,17 +528,27 @@ export class ChatSession { // Pure UI state — nothing is created upstream until a message is sent this.creatingThread.value = true; this.currentThreadId.value = null; + this.activeThreadIds = []; this.computeChatEntries(); this.rescroll(); } + createFork (threadId, messageId) { + backend.sendAction("fork-chat-thread", { + chatId: this.chatId, threadId, messageId + }).then(res => { + // Wait for received message updates??? + ui.alert(`Forked thread ${threadId} into ${res.threadId}`); + }); + } + // === Entry Computation === // // Builds the display list: the active thread's messages in order, every // other thread collapsed into a single preview header - computeChatEntries (messages) { + computeChatEntries () { let entries = []; const senderName = (sender) => { @@ -487,88 +560,196 @@ export class ChatSession { } }; - const getThreadHeader = (thread, last, options={}) => { - let preview = null; - if (last && !options.active) { - let content = String(last.content || '').split('\n')[0].trim(); - if (content.length > 100) content = content.slice(0, 100) + '…'; - preview = {sender: senderName(last.sender), content}; - } - return { - type: 'thread-header', - threadId: thread.id, - active: !!options.active, - messageCount: thread.messages.length, - preview, - }; + const threadPreview = (thread, message) => { + let content = String(message.content || '').trim(); + if (content.length > 60) content = "..." + content.slice(-60); + return {sender: senderName(message.sender), content}; }; - const activeThreadMessages = this.currentThreadId.value == null ? - [] : this.getThread(this.currentThreadId.value).messages; - - let activeIndex = 0; - let threadIndex = 0; - - while (true) { - let active = activeThreadMessages[activeIndex]; - let thread = this.threads[threadIndex]; - - if (!active && !thread) break; - - // Skip active thread - if (thread && thread.id != null && thread.id == this.currentThreadId.value) { - thread = null; - threadIndex++; - - if (!active) continue; + // Sources must have only non-null message elements + class Zipper { + constructor () { + this.sources = []; + this.indices = []; } - - let threadLast = thread ? thread.getLastMessage() : null; - let entry = earliestMessage(active, threadLast); - - // None is earliest? - if (!entry) break; - - if (entry == active) { - // Add thread header before first thread message - if (activeIndex == 0) { - entries.push({ - type: 'thread-header', - threadId: entry.threadId, - active: true, - messageCount: activeThreadMessages.length, - preview: null, - }); + add (source) { + this.sources.push(source); + this.indices.push(0); + } + next () { + let count = this.indices.length; + let values = []; + for (let i = 0; i < count; i++) { + let value = this.sources[i][this.indices[i]]; + // push all values to maintain index parity + values.push(value); } - activeIndex++; - } else if (entry == threadLast) { - // Replace with header - entry = getThreadHeader(thread, threadLast); - threadIndex++; + + // If there are messages, there will be a winner + let nextId = selectEarliest(values); + if (nextId == null) return; + + this.indices[nextId]++ + let entry = values[nextId]; + return entry; } - - // never halt - else break; - - // Update sender name - entry = Object.assign( - {type: entry.kind || "message"}, // default - entry, // value - {sender: senderName(entry.sender)} // override - ); - entries.push(entry); } - if (this.creatingThread.value) { - entries.push({ + + console.log("compute entries", { + threads: this.threads, + activeThreadIds: this.activeThreadIds, + }); + + // Generate thread headers + // only youngest fork is ever shown for non-active threads, + let headers = []; + const addThreadForks = (thread) => { + for (let fork of thread.children) { + // Skip active fork, it's added as a thread header + if (this.activeThreadIds.includes(fork.id)) continue; + + let firstMessage = fork.getFirstMessage(); + let beforeMessage; + + if (firstMessage) { + // Find the last thread message before the first message + // should go backwards for performance + for (let msg of thread.messages) { + if (msg.id < firstMessage) { + beforeMessage = msg; + } + } + } + + // no thread messages? + let attach = + beforeMessage ? {id: beforeMessage.id, position: 'after'} : + firstMessage ? {id: firstMessage.id, position: 'before'} : + null; + + if (!attach) return; + + headers.push({ + // claim just before id + id: 't.' + fork.id, + type: 'thread-header', + threadId: fork.id, + active: false, + messageCount: fork.messageCount, + time: firstMessage?.time, + + attach, + preview: firstMessage && threadPreview(fork, firstMessage), + }); + } + } + const addThreadHeader = (thread) => { + let isActive = this.activeThreadIds.includes(thread.id); + + let referenceMessage; + if (isActive) { + addThreadForks(thread); + referenceMessage = thread.getFirstMessage(); + + // Recursively add active thread headers + for (let fork of thread.children) { + if (this.activeThreadIds.includes(fork.id)) { + addThreadHeader(fork); + } + } + } else { + // Get all descendant messages and find the latest + let stack = [thread]; + let lastMessage = null; + let lastThread = null; + + while (stack.length) { + let next = stack.pop(); + for (let child of next.children) stack.push(child); + for (let msg of next.messages) { + if (!lastMessage || msg.id > lastMessage.id) { + lastMessage = msg; + lastThread = next; + } + } + } + + if (lastMessage) { + thread = lastThread; + referenceMessage = lastMessage; + } + } + + // No children? + if (!referenceMessage) return; + + headers.push({ + // claim just before id + id: 't.' + thread.id, type: 'thread-header', - threadId: null, - active: true, - messageCount: 0, - preview: null, + threadId: thread.id, + active: isActive, + messageCount: isActive ? 0 : thread.messageCount, + time: referenceMessage.time, + + attach: {id: referenceMessage.id, position: 'before'}, + preview: isActive || !referenceMessage ? + null : threadPreview(thread, referenceMessage), }); } + this.threads.forEach(addThreadHeader); + + // Headers must be ordered by their attach id + headers.sort((a, b) => parseInt(a.attach.id) - parseInt(b.attach.id)); + + + // This object will produce ordered messages in O(n) + // assuming sources are ordered + let zipper = new Zipper(); + + // Feed messages from all active threads + for (let threadId of this.activeThreadIds) { + zipper.add(this.getThread(threadId).messages.map( + msg => Object.assign( + // default, value, override + {type: msg.kind || "message"}, msg, + {sender: senderName(msg.sender)} + ) + )); + } + + // Produce entries from zipper, attaching headers as declared + let next; + while (next = zipper.next()) { + let after = []; + + // Collect headers to attach before or after this message + while (headers.length) { + let header = headers[0]; + + // all header entries attach + let attach = + header.attach.id == next.id ? header.attach.position : + header.attach.id < next.id ? "before" : + null; + + if (attach == 'before') entries.push(header); + if (attach == 'after') after.push(header); + + if (attach) headers.splice(0, 1); + else break; + } + + entries.push(next); + for (let header of after) entries.push(header); + } + + // Attach remaining + for (let header of headers) entries.push(header); + + console.log({entries}); this.chatEntries.value = entries; } } \ No newline at end of file diff --git a/frontend/src/views/Chat.vue b/frontend/src/views/Chat.vue index 93c6d4b..c0f4f6c 100644 --- a/frontend/src/views/Chat.vue +++ b/frontend/src/views/Chat.vue @@ -116,7 +116,7 @@ {value: 'pin', title: "Pin message"}, {value: 'edit', title: "Edit in-place"}, {value: 'edit-fork', title: "Edit in new thread"}, - //{value: 'start-thread', title: "Start thread from here"}, + {value: 'fork', title: "Fork thread from here"}, ]; // Long logs can be expanded/collapsed @@ -141,6 +141,10 @@ editMessage(message, true); break; } + case 'fork': { + session.createFork(message.threadId, message.id); + break; + } case 'toggle-log': { expandedLogs.toggle(message.id); break; @@ -192,7 +196,7 @@