Temporary agents and chats

backend/src/app_mixins/chat.py

@@ -2,18 +2,27 @@ import asyncio
 from datetime import datetime
 
 import database as db 
+import temporary as tmp
 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
-        )
+    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:
@@ -24,6 +33,11 @@ class ChatMixin:
     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,
@@ -32,13 +46,23 @@ class ChatMixin:
         } 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.")
+        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)
@@ -58,18 +82,25 @@ class ChatMixin:
         } 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
-        )
+        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,
+                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
-        )
+        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):
@@ -82,27 +113,49 @@ class ChatMixin:
 
         # 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
-        )
+        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
+            # 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)
+            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)
@@ -120,7 +173,8 @@ class ChatMixin:
                 notify_users()
             async def end_fn ():
                 message.content = content.content
-                await message.save()
+                if db.is_instance(message):
+                    await message.save()
                 await notify_agents()
                 notify_users()
 
@@ -142,7 +196,10 @@ class ChatMixin:
         }
 
     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 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}")
@@ -161,7 +218,8 @@ class ChatMixin:
             summary.append('format')
         
         if summary:
-            await message.save() # Saves the specific instance
+            if db.is_instance(message):
+                await message.save() # Saves the specific instance
             print(f"Updated message #{message.id} fields: {summary}")
 
     async def get_chat_view(self, chat_id):