move chat app functions to mixin

backend/app_mixins/chat.py

@@ -0,0 +1,145 @@
+import asyncio
+from datetime import datetime
+from sqlalchemy.sql.expression import select
+
+import database as db
+from chat import ChatView, Message, MESSAGE_PAGE_SIZE
+
+class ChatMixin:
+    def create_chat(self, user_id, name, description, picture, members=[]):
+        with self.sql_session() as session:
+            chat = db.Chat()
+            chat.user_id = user_id
+            chat.name = name
+            chat.description = description
+            chat.picture = picture
+
+            session.add(chat)
+            session.commit()
+            print(f"Created chat #{chat.id}: {chat.name}")
+
+            chat_id = chat.id
+        
+        for agent_id in members:
+            self.add_agent_to_chat(chat_id, agent_id)
+
+        return chat_id
+
+    def get_chats(self, user_id):
+        with self.sql_session() as session:
+            chats = session.execute(select(db.Chat).where(db.Chat.user_id == user_id)).scalars().all()
+            return [{
+                "id": chat.id,
+                "name": chat.name,
+                "description": chat.description,
+                "picture": chat.picture
+            } for chat in chats]
+
+    def delete_chat(self, chat_id):
+        with self.sql_session() as session:
+            session.execute(db.Membership.__table__.delete().where(db.Membership.chat_id == chat_id))
+            session.execute(db.Chat.__table__.delete().where(db.Chat.id == chat_id))
+            session.commit()
+            print(f"Deleted chat #{chat_id} and its memberships.")
+
+    def get_chat_messages(self, chat_id, earliest_index=None):
+        with self.sql_session() as session:
+            condition = (db.Message.chat_id == chat_id)
+            if earliest_index is not None:
+                condition = condition and db.Message.id < earliest_index
+
+            messages = session.execute(
+                select(db.Message).where(condition).order_by(db.Message.time.desc())
+            ).scalars().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]
+    
+    def create_chat_thread(self, chat_id, name=None):
+        with self.sql_session() as session:
+            thread = db.Thread()
+            thread.chat_id = chat_id
+            thread.name = name
+
+            session.add(thread)
+            session.commit()
+            print(f"Created thread #{thread.id}: {thread.name}")
+
+            return thread.id
+
+    def add_agent_to_chat(self, chat_id, agent_id):
+        with self.sql_session() as session:
+            membership = db.Membership()
+            membership.chat_id = chat_id
+            membership.agent_id = agent_id
+            
+            session.add(membership)
+            session.commit()
+            print(f"Added agent #{agent_id} to chat #{chat_id}")
+
+    # If an integer is passed in sender, it is interpreted as an agent id
+    def send_chat_message(self, chat_id, content, sender='user', thread_id=None):
+        with self.sql_session() as session:
+            message = db.Message()
+            message.chat_id = chat_id
+            message.time = int(datetime.now().timestamp())
+            message.sender = sender
+            message.kind = 'text'  # default kind
+            message.content = content
+            message.thread_id = thread_id
+
+            session.add(message)
+            session.commit()
+            session.refresh(message)
+
+        # Notify chat owner
+        if sender != 'user':
+            chat_owner = session.execute(select(db.Chat.user_id).where(db.Chat.id == chat_id)).scalar()
+            self.notify_message_to_user(chat_owner, message)
+
+        # Notify participant agents
+        agent_ids = session.execute(select(db.Membership.agent_id).where(db.Membership.chat_id == chat_id)).scalars().all()
+        for agent_id in agent_ids:
+            if agent_id != sender:
+                self.notify_message_to_agent(agent_id, message)
+
+        return {
+            'id': message.id,
+            'chatId': message.chat_id,
+            'time': message.time,
+            'sender': message.sender,
+            'kind': message.kind,
+            'content': message.content,
+        }
+
+    def get_chat_view(self, chat_id):
+        return ChatView(self, chat_id)
+
+    def notify_message_to_user(self, user_id, message):
+        # could be moved to user mixin
+        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,
+            }
+            for channel in self.user_connections[user_id]:
+                asyncio.create_task(channel.put(user_message))
+
+    def notify_message_to_agent(self, agent_id, message):
+        # could be moved to agent mixin
+        control = self.get_agent_control(agent_id)
+        asyncio.create_task(control.received_message(message))