move chat app functions to mixin

backend/app.py

@@ -16,13 +16,9 @@ from server import start_server
 from chat import ChatView, Message, MESSAGE_PAGE_SIZE
 
 from agents.conversation import ConversationAgent
+from app_mixins import ChatMixin
 
-class ChatContext:
+class App(ChatMixin):
-    def __init__(self, messages, prelude=None):
-        self.messages=messages
-        self.prelude=prelude
-
-class App:
     def __init__(self, data_directory, replicate_api_token=None):
         self.data_directory = data_directory
         self.user_connections = {}
@@ -106,25 +102,6 @@ class App:
 
             return agent.id
 
-    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
-
     # Updates the agent's basic fields, and applies field updates in its data
     def update_agent(self, agent_id, name=None, description=None, picture=None, agentname=None, data=None):
         with self.sql_session() as session:
@@ -159,16 +136,6 @@ class App:
             session.commit()
             print(f"Updated agent #{agent.id} fields: {summary}")
 
-    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 get_agents(self, user_id):
         with self.sql_session() as session:
             agents = session.execute(select(db.Agent).where(db.Agent.user_id == user_id)).scalars().all()
@@ -193,47 +160,6 @@ class App:
             "description": getattr(agent_class, 'description', None),
         }
 
-    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.Chat()
-            thread.chat_id = chat_id
-            thread.name = name
-
-            session.add(thread)
-            session.commit()
-            print(f"Created thread #{thread.id}: {thread.name}")
-
-            thread_id = thread.id
-        return thread_id
-
     def verify_user_credentials(self, username, password):
         with self.sql_session() as session:
             user = session.execute(
@@ -265,57 +191,6 @@ class App:
                 print(f"Signup error: '{email}' or '{username}' are already used")
         return user_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_context(self, chat_id, earliest_index=None):
-        return ChatContext(self.get_chat_messages(chat_id, earliest_index))
-
-    def get_chat_view(self, chat_id):
-        return ChatView(self, chat_id)
-
     def get_agent_control(self, agent_id):
         with self.sql_session() as session:
             agent = session.execute(
@@ -365,25 +240,6 @@ class App:
         response = await control.execute_action(action, args)
         return response
 
-    def notify_message_to_user(self, user_id, message):
-        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):
-        control = self.get_agent_control(agent_id)
-        asyncio.create_task(control.received_message(message))
-
     def generate_default_user_data(self, user_id):
         chat_id = self.create_chat(
             user_id=user_id,
@@ -403,18 +259,6 @@ class App:
             data={}
         )
     
-    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}")
-
-            return thread.id
-    
     def get_agent_resource(self, agent_id, kind, key):
         with self.sql_session() as session:
             stored_value = session.execute(select(db.AgentResource.value).where(

backend/app_mixins/__init__.py

@@ -0,0 +1 @@
+from app_mixins.chat import ChatMixin

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))