chat,frontend: better chat entry algorithm (wip)

backend/src/app_mixins/chat.py

@@ -93,17 +93,6 @@ class ChatMixin:
         # 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):
             thread = tmp.Thread.create(self, chat_id, name)
@@ -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):

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'],

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) {
+function selectEarliest (messages) {
-  let minId = null;
+  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) {
+    if (selected == null || msgTime < minTime) {
-      useTime = false;
+      selected = i;
-    } else if (!minTime || msgTime < messageTime(minTime)) {
+      minTime = msgTime;
-      minTime = msg;
     }
   }
 
-  return useTime ? minTime : minId;
+  return selected
 }
 
 export class Thread {
-  constructor(obj) {
+  constructor (obj, session) {
-    Object.assign(this, {messages: []}, obj || {});
+    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);
+    return this.threadMap[threadId] || new Thread({id: threadId}, this);
-    if (!thread) {
-      thread = new Thread({
-        id: threadId,
-        loaded: false,
-        time: null, // TODO
-        messages: [],
-      });
-      this.threads.push(thread);
-    }
-    return thread;
   }
 
   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={}) => {
+    const threadPreview = (thread, message) => {
-      let preview = null;
+      let content =  String(message.content || '').trim();
-      if (last && !options.active) {
+      if (content.length > 60) content = "..." + content.slice(-60);
-        let content = String(last.content || '').split('\n')[0].trim();
+      return {sender: senderName(message.sender), content};
-        if (content.length > 100) content = content.slice(0, 100) + '…';
+    };
-        preview = {sender: senderName(last.sender), content};
+
+    // Sources must have only non-null message elements
+    class Zipper {
+      constructor () {
+        this.sources = [];
+        this.indices = [];
       }
-      return {
+      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);
+        }
+
+        // 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;
+      }
+    }
+
+
+    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: thread.id,
-        active: !!options.active,
+        active: isActive,
-        messageCount: thread.messages.length,
+        messageCount: isActive ? 0 : thread.messageCount,
-        preview,
+        time: referenceMessage.time,
-      };
-    };
 
-    const activeThreadMessages = this.currentThreadId.value == null ?
+        attach: {id: referenceMessage.id, position: 'before'},
-      [] : this.getThread(this.currentThreadId.value).messages;
+        preview: isActive || !referenceMessage ?
-
+          null : threadPreview(thread, referenceMessage),
-    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;
-      }
-
-      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,
       });
     }
-        activeIndex++;
+
-      } else if (entry == threadLast) {
+    this.threads.forEach(addThreadHeader);
-        // Replace with header
+
-        entry = getThreadHeader(thread, threadLast);
+    // Headers must be ordered by their attach id
-        threadIndex++;
+    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)}
+        )
+      ));
     }
 
-      // never halt
+    // 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;
-
-      // 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(next);
-      entries.push({
+      for (let header of after) entries.push(header);
-        type: 'thread-header',
-        threadId: null,
-        active: true,
-        messageCount: 0,
-        preview: null,
-      });
     }
 
+    // Attach remaining
+    for (let header of headers) entries.push(header);
+
+    console.log({entries});
     this.chatEntries.value = entries;
   }
 }

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 @@
     <div class="messages-container" ref="chatContainerElem">
       <Message 
         v-for="msg in chatEntries" 
-        :key="msg.id ?? ('thread-' + msg.threadId)" 
+        :key="msg.id" 
         :message="msg" 
         @open-menu="chatEntryContextMenu"
         @open-thread="session.activateThread($event)"