chat: proper fork behavior in backend, wip frontend

backend/src/app_mixins/chat.py

@@ -17,6 +17,25 @@ def message_json(msg):
         "threadId": msg.thread_id,
     }
 
+async def aggregate_thread(thread_id):
+    print(f"aggregate {thread_id}")
+    query = db.Message.filter(thread_id=thread_id)
+
+    # Could do all these in one call?
+    message_count = await query.count()
+    print(f"    count {message_count}")
+    if message_count == 0: return
+
+    first_message = await query.first()
+    last_message = await query.last()
+    print(f"    bounds {first_message} {last_message}")
+
+    return await db.Thread.get(id=thread_id).update(
+        first_message_id=first_message.id,
+        last_message_id=last_message.id,
+        message_count=message_count
+    )
+
 class ChatMixin:
     async def create_chat(self, user_id, name, description, picture, members=[], temporary=False):
         if temporary:
@@ -135,6 +154,12 @@ class ChatMixin:
         new_id = await self.create_chat_thread(thread.chat_id, parent_id=thread_id, name=name or thread.name)
 
         if message_id is not None:
+            # Ensure all forks have the correct information
+            forks = await db.Thread.filter(parent_id=thread_id).all()
+            for fork in forks:
+                await aggregate_thread(fork.id)
+
+            await db.Thread.filter(parent_id=thread_id, first_message_id__gte=message_id).update(parent_id=new_id)
             await db.Message.filter(thread_id=thread_id, id__gte=message_id).update(thread_id=new_id)
 
         # Notify users. what a heavy operation!

backend/src/database.py

@@ -129,7 +129,7 @@ class Thread(models.Model):
     name = fields.CharField(max_length=255, null=True)
     description = fields.TextField(null=True)
 
-    # Automatic summary fields, include parent messages
+    # Automatic summary fields including messages directly in this thread
     first_message = fields.ForeignKeyField("models.Message", related_name=False, null=True, on_delete=fields.SET_NULL)
     last_message = fields.ForeignKeyField("models.Message", related_name=False, null=True, on_delete=fields.SET_NULL)
     # null make life easy

backend/src/inference.py

@@ -4,6 +4,7 @@ import json
 import os
 #import ollama
 import traceback
+import random
 from stream import Stream
 from tool import format_tools_for_engine, build_tool_prompt_prefix, ToolSpec, ToolCall
 
@@ -152,6 +153,8 @@ async def chat_completion(app,
     opts={},
     agent_run=None
 ):
+    # TESTING OVERRIDE
+    return await mock_completion(app, agent_run)
     db_model = (await app.get_model(model_id)).value
 
     model_engine = db_model['engine']
@@ -462,6 +465,32 @@ async def replicate_base_completion(app, model, input, agent_run=None):
         return None
     return content_stream
 
+async def mock_completion(app, agent_run=None):
+    VERBS = ['fix','scan','sync','deploy','parse','render','cache','trace',
+        'merge','purge','index','probe','flush','bind','warm','drain']
+    OBJECTS = ['nodes','packets','tokens','frames','shards','clusters',
+       'buffers','threads','sockets','records','layers','streams',
+       'handles','slots','cursors','segments']
+    DELAY = 0.3
+
+    content_stream = Stream(app)
+    if agent_run:
+        agent_run.emit_event('output', content_stream)
+
+    async def content_task ():
+        await asyncio.sleep(DELAY)
+        content_stream.feed(random.choice(VERBS))
+        await asyncio.sleep(DELAY)
+        content_stream.feed(f" {random.randint(1, 99)} ")
+        await asyncio.sleep(DELAY)
+        content_stream.feed(random.choice(VERBS))
+        await asyncio.sleep(DELAY)
+        content_stream.close()
+        agent_run.emit_event('inference-end', {})
+
+    asyncio.create_task(content_task())
+    if not agent_run: return content_stream
+
 
 def run_replicate_inference_loop(replicate_client, model, input, content_stream, agent_run=None):
     print("START INFERENCE")

frontend/src/chat.js

@@ -255,15 +255,12 @@ export class ChatSession {
 
       let thread = this.getThread(it.threadId);
 
-      for (let entry of thread.messages) {
+      for (let msg of thread.messages) {
-        if (entry.id == it.id) {
+        if (msg.id == it.id) {
-          entry.type = it.kind;
+          msg.type = it.kind;
-          entry.subkind = it.subkind;
+          msg.subkind = it.subkind;
-          entry.format = it.format;
+          msg.format = it.format;
-          entry.content = it.content;
+          msg.content = it.content;
-
-          // If this changes, message has to get moved
-          entry.threadId = it.threadId;
 
           this.computeChatEntries();
           return;
@@ -271,15 +268,20 @@ export class ChatSession {
       }
 
       // Message not found in target thread
-      // Find the old message
+      // Remove from wherever it is, we will add it as new
-      ui.alert("TODO: find updated messages in different threads");
+      let stack = Array.from(this.threads);
-      //let stack = Array.from(this.threads);
+      while (stack.length) {
-      let oldMessage = null;
+        let current = stack.pop();
 
-      // if we didn't find it, add it like a new one
+        let index = current.messages.findIndex(msg => msg.id == it.id);
-      if (oldMessage) {
+        if (index >= 0) {
-        ui.alert("TODO: move a message with updated thread");
+          current.messages.splice(index, 1);
-        return;
+          break;
+        }
+
+        for (let child of current.children) {
+          stack.push(child);
+        }
       }
     }
 
@@ -299,8 +301,8 @@ export class ChatSession {
       }
     }
 
-    // addMessageBulk I think is more correct
+    // addMessageBulk inserts in correct position, unlike addMessage
-    this.addMessage({
+    this.addMessageBulk([{
       id: it.id,
       type: it.kind,
       subkind: it.subkind,
@@ -309,7 +311,7 @@ export class ChatSession {
       content: it.content,
       threadId: it.threadId,
       time: it.time,
-    });
+    }]);
 
     this.computeChatEntries();
     this.rescroll();
@@ -594,30 +596,32 @@ export class ChatSession {
     // Continuing a thread: it becomes the one open thread
     this.creatingThread.value = false;
     this.currentThreadId.value = threadId;
-    this.activeThreadIds = [threadId];
+    this.activeThreadIds = [];
 
-    // Load thread messages asynchronously
+    // Activate and load all threads in the ancestry
-    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,
-        format: it.format,
-        subkind: it.subkind,
-        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) {
+    while (true) {
-      thread = this.getThread(thread.parentId);
       this.activeThreadIds.push(thread.id);
+
+      // Load  messages asynchronously
+      backend.sendAction("get-thread-messages", {
+        chatId: this.chatId, threadId: thread.id
+      }).then(res => {
+        // TODO: reassigning all fields?
+        this.addMessageBulk(res.messages.map(it => ({
+          id: it.id,
+          type: it.kind,
+          format: it.format,
+          subkind: it.subkind,
+          sender: it.sender,
+          content: it.content,
+          threadId: it.threadId,
+          time: it.time,
+        })));
+      });
+
+      if (!thread.parentId) break;
+      thread = this.getThread(thread.parentId);
     }
 
     this.computeChatEntries();
@@ -634,13 +638,13 @@ export class ChatSession {
     this.rescroll();
   }
 
-  createFork (threadId, messageId) {
+  async createFork (threadId, messageId) {
-    backend.sendAction("fork-chat-thread", {
+    let res = await 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}`);
     });
+    let thread = this.getThread(res.threadId);
+    thread.update({parentId: threadId});
+    return thread;
   }
 
 

frontend/src/views/Chat.vue

@@ -142,7 +142,9 @@
         break;
       }
       case 'fork': {
-        session.createFork(message.threadId, message.id);
+        let newThread = await session.createFork(message.threadId, message.id);
+        ui.alert(`Forked thread ${newThread.id} from ${message.threadId}`);
+        session.activateThread(newThread.id);
         break;
       }
       case 'toggle-log': {