streamed messages

backend/src/agents/expert.py

@@ -53,11 +53,14 @@ class ExpertAgent(Agent):
 
         if type(model) is int:
             db_model = await self.app.get_model(model)
+            model_engine = db_model.value["engine"]
             
-            if db_model.value["engine"] == "llama-cpp":
+            if model_engine == "llama-cpp":
                 response = await inference.local_chat_completion(self.app, messages, model=db_model.value['model'])
+            elif model_engine == "replicate":
+                response = await inference.replicate_llama_completion(self.app, messages, model=db_model.value['model'])
             else:
-                raise Exception("Unknown model: " + repr(model))
+                raise Exception("Unsuported model: " + repr(db_model.value))
 
         elif model.startswith("ollama:"):
             response = await inference.ollama_chat_completion(self.app, messages, model=model[7:])

backend/src/app_mixins/chat.py

@@ -1,8 +1,9 @@
 import asyncio
 from datetime import datetime
-# Assuming db is your database.py containing the Tortoise models from the previous step
+
 import database as db 
 from chat import ChatView, Message, MESSAGE_PAGE_SIZE
+from stream import Stream
 
 class ChatMixin:
     async def create_chat(self, user_id, name, description, picture, members=[]):
@@ -42,8 +43,9 @@ class ChatMixin:
         if earliest_index is not None:
             query = query.filter(id__gt=earliest_index)
 
-        # Order by time descending, then reverse in Python to get chronological order
-        messages = await query.order_by("-time").all()
+        # Order quey by time descending (id orders better)
+        # then reverse in Python to get chronological order
+        messages = await query.order_by("-id").all()
         messages.reverse()
 
         return [{
@@ -72,6 +74,14 @@ class ChatMixin:
 
     async def send_chat_message(self, chat_id, content, sender='user', thread_id=None, kind=None, subkind=None, format=None):
         now = int(datetime.now().timestamp())
+
+        if isinstance(content, Stream):
+            initial_content = "..."
+        else:
+            initial_content = content
+
+        # Set to true for streams
+        updating = False
         message = await db.Message.create(
             chat_id=chat_id,
             time=now,
@@ -80,7 +90,7 @@ class ChatMixin:
             kind=kind or 'message',
             subkind=subkind,
             format=format,
-            content=content,
+            content=initial_content,
             thread_id=thread_id
         )
 
@@ -88,14 +98,40 @@ class ChatMixin:
         chat = await db.Chat.get(id=chat_id).prefetch_related("user")
         chat_owner_id = chat.user_id
 
-        self.notify_message_to_user(chat_owner_id, message)
+        # Notification order depends on whether the content is a stream
 
+        async def notify_agents():
             # Fetch participant agents
             agent_ids = await db.Membership.filter(chat_id=chat_id).values_list("agent_id", flat=True)
             for agent_id in agent_ids:
                 if str(agent_id) != str(sender):
                     await self.notify_message_to_agent(agent_id, message)
         
+        def notify_users():
+            nonlocal updating
+            self.notify_message_to_user(chat_owner_id, message, updating)
+            updating = True
+
+        if isinstance(content, Stream):
+            notify_users()
+
+            def chunk_fn (chunk):
+                message.content = content.content + "..."
+                notify_users()
+            async def end_fn ():
+                message.content = content.content
+                await message.save()
+                await notify_agents()
+                notify_users()
+
+            content.listen(chunk_fn=chunk_fn, end_fn=end_fn)
+        else:
+            notify_users()
+            await notify_agents()
+        
+        # TODO: Chat preview based on the message
+
+        # TODO: return for who lmao
         return {
             'id': message.id,
             'chatId': message.chat_id,
@@ -133,7 +169,7 @@ class ChatMixin:
         await chat_view.init_load()
         return chat_view
 
-    def notify_message_to_user(self, user_id, message):
+    def notify_message_to_user(self, user_id, message, update=False):
         if user_id in self.user_connections:
             user_message = {
                 'action': 'chat-message',
@@ -144,6 +180,7 @@ class ChatMixin:
                 'kind': message.kind,
                 'content': message.content,
                 'threadId': message.thread_id,
+                'update': update,
             }
             for channel in self.user_connections[user_id]:
                 asyncio.create_task(channel.put(user_message))

backend/src/inference.py

@@ -28,18 +28,35 @@ class ModelCache:
 
 # llama.cpp 
 async def local_chat_completion(app, messages, model, opts={}):
+    content_stream = Stream(app)
+
+    def worker():
         llm = Llama(    
             model_path=f"{app.data_directory}/models/{model}",
             # n_gpu_layers=-1, # Uncomment to use GPU acceleration
             # seed=1337, # Uncomment to set a specific seed
             # n_ctx=2048, # Uncomment to increase the context window
         )
-    response = llm.create_chat_completion(
-        messages = [{'role': msg.role, 'content': msg.content} for msg in messages]
+
+        # With stream=True, returns an iterator that infers tokens on demand (openai compatible)
+        llm_output = llm.create_chat_completion(
+            messages = [{'role': msg.role, 'content': msg.content} for msg in messages],
+            stream = True,
         )
 
-    # TODO: KeyError: 'text'
-    return response["choices"][0]["text"]
+        for chunk in llm_output:
+            delta = chunk['choices'][0]['delta']
+            if 'content' in delta:
+                chunk_text = delta['content']
+                content_stream.feed(chunk_text)
+            #chunk_text = chunk['choices'][0]['text']
+
+        content_stream.close()
+
+    thread = threading.Thread(target=worker)
+    thread.start()
+
+    return content_stream
 
 async def ollama_list(app, messages, model, opts={}):
     # TODO: test the output, and use this
@@ -112,8 +129,9 @@ async def replicate_base_completion(app, model, input):
         }
     )
     thread.start()
-    response = await content_stream.content_wait()
-    return response.lstrip()
+    #response = await content_stream.content_wait()
+    #return response.lstrip()
+    return content_stream
 
 
 def run_replicate_inference_loop(replicate_client, model, input, content_stream):

backend/src/server.py

@@ -170,7 +170,7 @@ def start_server(app):
 
                 elif action == "get-models":
                     await send({
-                        'models': [{"id": model.id, "name": model.name} for model in await app.get_models()],
+                        'models': [{"id": model.id, "name": model.name, "value": model.value} for model in await app.get_models()],
                         'nonce': nonce
                     })
 
@@ -262,7 +262,7 @@ def start_server(app):
                         body.get('input'),
                         body.get('options', {})
                     )
-                    await send({'message': "propmt sent", 'response': response, 'nonce': nonce})
+                    await send({'message': "prompt sent", 'response': response, 'nonce': nonce})
 
                 elif action == "call-agent-action":
                     result = await app.call_agent_action(body['agentId'], body['agentAction'], body['arguments'])

backend/src/stream.py

@@ -28,6 +28,7 @@ def make_runnable(fn_obj):
     There can also be arguments coming along with the signal, those are passed to the function.
     """
     signal_queue = asyncio.Queue()
+    loop = asyncio.get_event_loop()
 
     async def signal_listener():
         while True:
@@ -40,7 +41,9 @@ def make_runnable(fn_obj):
     asyncio.create_task(signal_listener())
 
     # Return a method to put items on the signal queue
-    return lambda *args: signal_queue.put_nowait(args)
+    def run_it(*args):
+        loop.call_soon_threadsafe(signal_queue.put_nowait, args)
+    return run_it
 
 class Stream:
     def __init__(self, app):
@@ -57,8 +60,8 @@ class Stream:
 
     def listen (self, chunk_fn, end_fn):
         nest_asyncio.apply()
-        #chunk_fn = make_runnable(chunk_fn)
-        #end_fn = make_runnable(end_fn)
+        chunk_fn = make_runnable(chunk_fn)
+        end_fn = make_runnable(end_fn)
 
         if self.open:
             self.listeners.append({

frontend/src/views/Chat.vue

@@ -76,6 +76,18 @@
     });
 
     unsubscribeChat = backend.onChatMessages(chatId, it => {
+      if (it.update) {
+        // Find message in list. If not found, add it
+        for (let entry of messages.value) {
+          if (entry.id == it.id) {
+            entry.type = it.kind;
+            entry.content = it.content;
+            entry.threadId = it.threadId;
+            return;
+          }
+        }
+      }
+
       if (it.sender == 'user') {
         let unsent = unsentMessages.value.shift();
         if (unsent) {
@@ -85,7 +97,7 @@
             return;
           } else {
             // Remove from unsent list and add message as received
-            ui.alert("Unsent messages received differently");
+            ui.alert("TODO: update unsent message");
           }
         }
       }

frontend/src/views/Models.vue

@@ -12,7 +12,7 @@
   const models = ref([]);
   const createExpanded = ref(false);
 
-  const editModel = ref({kind: 'text', name: null, inputValue: '{}'});
+  const editModel = ref({kind: 'text', name: null, inputValue: '{"engine":"llama-cpp","model":"modelfile.gguf"}'});
 
   onMounted(async () => {
     let res = await backend.sendAction("get-models");