334 lines
12 KiB
Python
334 lines
12 KiB
Python
import asyncio
|
|
from datetime import datetime
|
|
|
|
import database as db
|
|
import temporary as tmp
|
|
from chat import Chat
|
|
from stream import Stream
|
|
|
|
def message_json(msg):
|
|
return {
|
|
"id": msg.id,
|
|
"time": msg.time,
|
|
"sender": msg.sender,
|
|
"kind": msg.kind,
|
|
"subkind": msg.subkind,
|
|
"content": msg.content,
|
|
"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:
|
|
chat = tmp.Chat.create(
|
|
self,
|
|
user_id=user_id,
|
|
name=name,
|
|
description=description,
|
|
picture=picture
|
|
)
|
|
else:
|
|
chat = await db.Chat.create(
|
|
user_id=user_id,
|
|
name=name,
|
|
description=description,
|
|
picture=picture
|
|
)
|
|
print(f"Created chat #{chat.id}: {chat.name}")
|
|
|
|
for agent_id in members:
|
|
await self.add_agent_to_chat(chat.id, agent_id)
|
|
|
|
return chat.id
|
|
|
|
# Gets a chat view from a given chat id
|
|
async def get_chat(self, chat_id):
|
|
chat = Chat(self, chat_id)
|
|
await chat.load()
|
|
return chat
|
|
|
|
async def get_chats(self, user_id):
|
|
# Direct filtering on the model class
|
|
chats = await db.Chat.filter(user_id=user_id).all()
|
|
|
|
tmp_chats = self.tmp.get_user_chats(user_id)
|
|
if tmp_chats:
|
|
chats += [self.tmp.chats[chat_id] for chat_id in tmp_chats]
|
|
|
|
return [{
|
|
"id": chat.id,
|
|
"name": chat.name,
|
|
"description": chat.description,
|
|
"picture": chat.picture
|
|
} for chat in chats]
|
|
|
|
async def delete_chat(self, chat_id):
|
|
if tmp.is_tmp_id(chat_id):
|
|
chat = self.tmp.chats[chat_id]
|
|
# TODO: Delete
|
|
|
|
else:
|
|
# Filtering and deleting is a one-liner
|
|
await db.Membership.filter(chat_id=chat_id).delete()
|
|
await db.Chat.filter(id=chat_id).delete()
|
|
# TODO: Delete messages from the chat
|
|
print(f"Deleted chat #{chat_id} and its memberships.")
|
|
|
|
async def get_chat_messages(self, chat_id, earliest_index=None):
|
|
if tmp.is_tmp_id(chat_id):
|
|
chat = self.tmp.chats[chat_id]
|
|
|
|
return [message.to_json() for message in chat.messages]
|
|
|
|
query = db.Message.filter(chat_id=chat_id)
|
|
if earliest_index is not None:
|
|
query = query.filter(id__gt=earliest_index)
|
|
|
|
# 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()
|
|
|
|
# TODO: Stop repetition.
|
|
# There's db.Message and chat.Message, also should be combined somehow, or wrap db always.
|
|
return [message_json(msg) for msg in messages]
|
|
|
|
async def create_chat_thread(self, chat_id, parent_id=None, name=None):
|
|
if tmp.is_tmp_id(chat_id):
|
|
thread = tmp.Thread.create(self, chat_id, name)
|
|
else:
|
|
thread = await db.Thread.create(
|
|
chat_id=chat_id,
|
|
parent_id=parent_id,
|
|
name=name
|
|
)
|
|
print(f"Created thread #{thread.id}: {thread.name}")
|
|
return thread.id
|
|
|
|
async def get_chat_threads(self, chat_id):
|
|
threads = await db.Thread.filter(chat_id=chat_id).all()
|
|
return [{
|
|
'id': it.id,
|
|
'parentId': it.parent_id,
|
|
'name': it.name,
|
|
'description': it.description,
|
|
'firstMessageId': it.first_message_id,
|
|
'lastMessageId': it.last_message_id,
|
|
'messageCount': it.message_count,
|
|
} for it in threads]
|
|
|
|
async def get_thread_messages(self, chat_id, thread_id, limit_message_id=None):
|
|
chat = await self.get_chat(chat_id)
|
|
thread = chat.get_thread(thread_id)
|
|
messages = await thread.load_page(limit_message_id)
|
|
return [message_json(msg) for msg in messages]
|
|
|
|
async def get_chat_thread_heads(self, chat_id, limit_message_id=None):
|
|
chat = await self.get_chat(chat_id)
|
|
messages = await chat.load_thread_heads(limit_message_id)
|
|
return [message_json(msg) for msg in messages]
|
|
|
|
async def fork_chat_thread(self, thread_id, message_id=None, name=None):
|
|
"""Move all messages from message_id onwards into a new child thread.
|
|
If none is given, this is basically create_chat_thread with an implicit chat id."""
|
|
thread = await db.Thread.get(id=thread_id)
|
|
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!
|
|
|
|
chat = await db.Chat.get(id=thread.chat_id)
|
|
messages = await db.Message.filter(thread_id=new_id).all()
|
|
print("Chat: " + repr(chat))
|
|
print("Fork message list: " + repr(messages))
|
|
for msg in messages:
|
|
self.notify_message_to_user(chat.user_id, msg, True)
|
|
|
|
return new_id
|
|
|
|
async def add_agent_to_chat(self, chat_id, agent_id):
|
|
if tmp.is_tmp_id(chat_id):
|
|
chat = self.tmp.chats[chat_id]
|
|
chat.participants.append(agent_id)
|
|
else:
|
|
await db.Membership.create(
|
|
chat_id=chat_id,
|
|
agent_id=agent_id
|
|
)
|
|
print(f"Added agent #{agent_id} to chat #{chat_id}")
|
|
|
|
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_format = format
|
|
format = format and "steam-" + format or "stream"
|
|
initial_content = str(content.id)
|
|
else:
|
|
initial_content = content
|
|
|
|
# Set to true for streams
|
|
updating = False
|
|
if tmp.is_tmp_id(chat_id):
|
|
message = Message(
|
|
tmp.generate_tmp_id(self),
|
|
#chat_id=chat_id,
|
|
time=now,
|
|
#last_edit=now,
|
|
sender=sender,
|
|
kind=kind or 'message',
|
|
subkind=subkind,
|
|
format=format,
|
|
content=initial_content,
|
|
thread_id=thread_id
|
|
)
|
|
message.chat_id = chat_id
|
|
chat = self.tmp.chats[chat_id]
|
|
chat.messages.append(message)
|
|
chat_owner_id = chat.user_id
|
|
else:
|
|
message = await db.Message.create(
|
|
chat_id=chat_id,
|
|
time=now,
|
|
last_edit=now,
|
|
sender=sender,
|
|
kind=kind or 'message',
|
|
subkind=subkind,
|
|
format=format,
|
|
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():
|
|
# Not all messages are visible to agents
|
|
if message.kind == "resource":
|
|
return
|
|
|
|
# Fetch participant agents
|
|
if tmp.is_tmp_id(chat_id):
|
|
agent_ids = self.tmp.chats[chat_id].participants
|
|
else:
|
|
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):
|
|
pass
|
|
async def end_fn ():
|
|
message.format = initial_format
|
|
message.content = content.content
|
|
|
|
if db.is_instance(message):
|
|
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,
|
|
'time': message.time,
|
|
'sender': message.sender,
|
|
'kind': message.kind,
|
|
'content': message.content,
|
|
}
|
|
|
|
async def update_chat_message(self, message_id, content=None, kind=None, subkind=None, format=None):
|
|
if tmp.is_tmp_id(message_id):
|
|
raise Exception("Cannot update temporary messages")
|
|
else:
|
|
message = await db.Message.get_or_none(id=message_id)
|
|
|
|
if not message:
|
|
raise Exception(f"Message not found with id {message_id}")
|
|
|
|
summary = []
|
|
if content is not None:
|
|
message.content = content
|
|
summary.append('content')
|
|
if kind is not None:
|
|
message.kind = kind
|
|
message.subkind = subkind
|
|
summary.append('kind')
|
|
summary.append('subkind')
|
|
if format is not None:
|
|
message.format = format
|
|
summary.append('format')
|
|
|
|
if summary:
|
|
if db.is_instance(message):
|
|
await message.save() # Saves the specific instance
|
|
print(f"Updated message #{message.id} fields: {summary}")
|
|
|
|
def notify_message_to_user(self, user_id, message, update=False):
|
|
if user_id in self.user_connections:
|
|
# TODO: Server's responsibility, and Message.to_json responsibility as well
|
|
user_message = {
|
|
'action': 'chat-message',
|
|
'id': message.id,
|
|
'chatId': message.chat_id,
|
|
'time': message.time,
|
|
'sender': message.sender,
|
|
'kind': message.kind,
|
|
'subkind': message.subkind,
|
|
'format': message.format,
|
|
'content': message.content,
|
|
'threadId': message.thread_id,
|
|
'update': update,
|
|
}
|
|
for channel in self.user_connections[user_id]:
|
|
asyncio.create_task(channel.put(user_message))
|
|
|
|
async def notify_message_to_agent(self, agent_id, message):
|
|
control = await self.get_agent_control(agent_id)
|
|
asyncio.create_task(control.received_message(message)) |