From 7114d91403feff9740f52f90263929b2c6379bbd Mon Sep 17 00:00:00 2001 From: Shivering Ocelot Date: Sat, 17 Jan 2026 12:11:13 -0500 Subject: [PATCH] streamed messages --- backend/src/agents/expert.py | 7 ++-- backend/src/app_mixins/chat.py | 59 +++++++++++++++++++++++++++------- backend/src/inference.py | 46 ++++++++++++++++++-------- backend/src/server.py | 4 +-- backend/src/stream.py | 9 ++++-- frontend/src/views/Chat.vue | 14 +++++++- frontend/src/views/Models.vue | 2 +- 7 files changed, 107 insertions(+), 34 deletions(-) diff --git a/backend/src/agents/expert.py b/backend/src/agents/expert.py index d63830b..48fa08d 100644 --- a/backend/src/agents/expert.py +++ b/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:]) diff --git a/backend/src/app_mixins/chat.py b/backend/src/app_mixins/chat.py index 63ba66c..2444fa5 100644 --- a/backend/src/app_mixins/chat.py +++ b/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,22 +90,48 @@ class ChatMixin: kind=kind or 'message', subkind=subkind, format=format, - content=content, + content=initial_content, thread_id=thread_id ) # Fetch chat owner using values_list for efficiency chat = await db.Chat.get(id=chat_id).prefetch_related("user") chat_owner_id = chat.user_id + + # 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) - self.notify_message_to_user(chat_owner_id, message) + def notify_users(): + nonlocal updating + self.notify_message_to_user(chat_owner_id, message, updating) + updating = True - # 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) + 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)) diff --git a/backend/src/inference.py b/backend/src/inference.py index 2cac32b..b6ecef0 100644 --- a/backend/src/inference.py +++ b/backend/src/inference.py @@ -26,20 +26,37 @@ class ModelCache: last_used = None last_model = None -# llama.cpp +# llama.cpp async def local_chat_completion(app, messages, model, opts={}): - 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] - ) + content_stream = Stream(app) - # TODO: KeyError: 'text' - return response["choices"][0]["text"] + 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 + ) + + # 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, + ) + + 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): diff --git a/backend/src/server.py b/backend/src/server.py index f746518..5538f5f 100644 --- a/backend/src/server.py +++ b/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']) diff --git a/backend/src/stream.py b/backend/src/stream.py index 0546a8f..333ac4e 100644 --- a/backend/src/stream.py +++ b/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({ diff --git a/frontend/src/views/Chat.vue b/frontend/src/views/Chat.vue index c1422ce..15d217f 100644 --- a/frontend/src/views/Chat.vue +++ b/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"); } } } diff --git a/frontend/src/views/Models.vue b/frontend/src/views/Models.vue index 0048be8..0449d67 100644 --- a/frontend/src/views/Models.vue +++ b/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");