Tortoise based database models

backend/src/app_mixins/chat.py

@@ -1,117 +1,100 @@
 import asyncio
 from datetime import datetime
-from sqlalchemy.sql.expression import select
-
-import database as db
+# 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
 
 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
+    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}")
 
-            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)
+            await self.add_agent_to_chat(chat.id, agent_id)
 
-        return chat_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]
+    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]
 
-    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))
-            # TODO: Delete messages from the chat
-            session.commit()
-            print(f"Deleted chat #{chat_id} and its memberships.")
+    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.")
 
-    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
+    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)
 
-            messages = session.execute(
-                select(db.Message).where(condition).order_by(db.Message.time.desc())
-            ).scalars().all()
+        # Order by time descending, then reverse in Python to get chronological order
+        messages = await query.order_by("-time").all()
+        messages.reverse()
 
-            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]
+        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
+    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
 
-            session.add(thread)
-            session.commit()
-            print(f"Created thread #{thread.id}: {thread.name}")
+    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}")
 
-            return thread.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())
+        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=content,
+            thread_id=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}")
+        # 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
+        
+        self.notify_message_to_user(chat_owner_id, message)
 
-    # 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, kind=None, subkind=None, format=None):
-        with self.sql_session() as session:
-            message = db.Message()
-            message.chat_id = chat_id
-            message.time = int(datetime.now().timestamp())
-            message.last_edit = message.time
-            message.sender = sender
-            message.kind = kind or 'message'
-            message.subkind = subkind
-            message.content = content
-            message.thread_id = thread_id
-
-            session.add(message)
-            session.commit()
-            session.refresh(message)
-
-        # Notify chat owner
-        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()
+        # 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 agent_id != sender:
-                self.notify_message_to_agent(agent_id, message)
+            if str(agent_id) != str(sender):
+                await self.notify_message_to_agent(agent_id, message)
 
         return {
             'id': message.id,
@@ -122,41 +105,35 @@ class ChatMixin:
             'content': message.content,
         }
 
-    # Updates a single chat message.
-    # Most fields are optional. If they're not passed a value, they wont be updated.
-    # When kind is set, subkind is always set too.
-    def update_chat_message(self, message_id, content=None, kind=None, subkind=None, format=None):
-        with self.sql_session() as session:
-            message = session.execute(
-                select(db.Message).where(db.Message.id == message_id)
-            ).scalars().first()
+    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}")
+        if not message:
+            raise Exception(f"Message not found with id {message_id}")
 
-            summary = []
-
-            if content != None:
-                message.content = content
-                summary.append('content')
-            if kind != None:
-                message.kind = kind
-                message.subkind = subkind
-                summary.append('kind')
-                summary.append('subkind')
-            if format != None:
-                message.format = format
-                summary.append('format')
-            #message.last_edit = int(datetime.now().timestamp())
-
-            session.commit()
+        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}")
 
-    def get_chat_view(self, chat_id):
-        return ChatView(self, chat_id)
+    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):
-        # could be moved to user mixin
         if user_id in self.user_connections:
             user_message = {
                 'action': 'chat-message',
@@ -171,7 +148,6 @@ class ChatMixin:
             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))
+    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))