190 lines
6.4 KiB
Python
190 lines
6.4 KiB
Python
import asyncio
|
|
from datetime import datetime
|
|
|
|
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=[]):
|
|
# Tortoise uses .create() for one-line instantiation and saving
|
|
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
|
|
|
|
async def get_chats(self, user_id):
|
|
# Direct filtering on the model class
|
|
chats = await db.Chat.filter(user_id=user_id).all()
|
|
return [{
|
|
"id": chat.id,
|
|
"name": chat.name,
|
|
"description": chat.description,
|
|
"picture": chat.picture
|
|
} for chat in chats]
|
|
|
|
async def delete_chat(self, chat_id):
|
|
# 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):
|
|
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()
|
|
|
|
return [{
|
|
"id": message.id,
|
|
"time": message.time,
|
|
"sender": message.sender,
|
|
"kind": message.kind,
|
|
"content": message.content,
|
|
"threadId": message.thread_id,
|
|
} for message in messages]
|
|
|
|
async def create_chat_thread(self, chat_id, name=None):
|
|
thread = await db.Thread.create(
|
|
chat_id=chat_id,
|
|
name=name
|
|
)
|
|
print(f"Created thread #{thread.id}: {thread.name}")
|
|
return thread.id
|
|
|
|
async def add_agent_to_chat(self, chat_id, agent_id):
|
|
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_content = "..."
|
|
else:
|
|
initial_content = content
|
|
|
|
# Set to true for streams
|
|
updating = False
|
|
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():
|
|
# 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,
|
|
'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):
|
|
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:
|
|
await message.save() # Saves the specific instance
|
|
print(f"Updated message #{message.id} fields: {summary}")
|
|
|
|
async def get_chat_view(self, chat_id):
|
|
chat_view = ChatView(self, chat_id)
|
|
await chat_view.init_load()
|
|
return chat_view
|
|
|
|
def notify_message_to_user(self, user_id, message, update=False):
|
|
if user_id in self.user_connections:
|
|
user_message = {
|
|
'action': 'chat-message',
|
|
'id': message.id,
|
|
'chatId': message.chat_id,
|
|
'time': message.time,
|
|
'sender': message.sender,
|
|
'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))
|
|
|
|
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)) |