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")