streamed messages

This commit is contained in:
Shivering Ocelot committed 2026-01-17 12:11:13 -05:00
1 parent f2a6fc0cd5
commit 7114d91403
7 files changed
+107 -34

No files matched your search

+5 -2
View File
@@ -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:])
+48 -11
View File
@@ -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))
+32 -14
View File
@@ -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):
+2 -2
View File
@@ -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'])
+6 -3
View File
@@ -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({
+13 -1
View File
@@ -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");
}
}
}
+1 -1
View File
@@ -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");