implement chat message sending

.gitignore

@@ -1,2 +1,3 @@
 .venv
+__pycache__
 horde.db

backend/app.py

@@ -1,7 +1,9 @@
 from database import create_sql_engine, User, Agent, Chat, Message, Membership
+from server import start_server
 from sqlalchemy.orm import Session
 from sqlalchemy.sql.expression import select
 from quart import Quart, Response, request, send_from_directory, jsonify, websocket
+from datetime import datetime
 import asyncio
 import json
 import sys
@@ -10,10 +12,24 @@ class App:
     def __init__(self, data_directory):
         self.data_directory = data_directory
         self.sql_engine = create_sql_engine(self.data_directory)
+        self.user_connections = {}
 
     def sql_session(self):
         return Session(self.sql_engine)
 
+    def add_user_connection(self, user_id, channel):
+        if user_id not in self.user_connections:
+            self.user_connections[user_id] = []
+        if channel not in self.user_connections[user_id]:
+            self.user_connections[user_id].append(channel)
+
+    def remove_user_connection(self, user_id, channel):
+        if user_id in self.user_connections:
+            if channel in self.user_connections[user_id]:
+                self.user_connections[user_id].remove(channel)
+                if not self.user_connections[user_id]:  # Check if the list is empty
+                    del self.user_connections[user_id]
+
     def create_agent(self, user_id, name, description, picture, kind, data):
         with self.sql_session() as session:
             agent = Agent()
@@ -106,7 +122,7 @@ class App:
             else: 
                 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 = Membership()
@@ -117,6 +133,55 @@ class App:
             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'):
+        with self.sql_session() as session:
+            message = Message()
+            message.chat_id = chat_id
+            message.time = int(datetime.now().timestamp())
+            message.sender = sender
+            message.kind = 'text'  # default kind
+            message.content = content
+
+            session.add(message)
+            session.commit()
+
+            # Notify chat owner
+            if sender != 'user':
+                chat_owner = session.execute(select(Chat.user_id).where(Chat.id == chat_id)).scalar()
+                self.notify_message_to_user(chat_owner, message)
+
+            # Notify participant agents
+            agent_ids = session.execute(select(Membership.agent_id).where(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 notify_message_to_user(self, user_id, message):
+        if user_id in self.user_connections:
+            user_message = {
+                'action': 'chat-message',
+                'id': message.id,
+                'time': message.time,
+                'sender': message.sender,
+                'kind': message.kind,
+                'content': message.content
+            }
+            for channel in self.user_connections[user_id]:
+                asyncio.create_task(channel.put(user_message))
+
+    def notify_message_to_agent(self, agent_id, message):
+        pass # TODO
+
     def generate_default_user_data(self, user_id):
         agent_id = self.create_agent(
             user_id=user_id,
@@ -138,115 +203,6 @@ class App:
         # Add the agent to the chat
         self.add_agent_to_chat(chat_id, agent_id)
 
-def start_server(app):
-    server = Quart(__name__, static_folder='../frontend/build/web')
-
-    @server.route('/')
-    async def index():
-        return await send_from_directory(server.static_folder, 'index.html')
-
-    @server.route('/<path:path>')
-    async def serve_static(path):
-        return await send_from_directory(server.static_folder, path)
-
-    @server.route('/user/signup', methods=['POST'])
-    async def signup():
-        data = await request.get_json()
-
-        user_id = None
-
-        username = data['username']
-        email = data['email']
-        password = data['password']
-
-        user_id = app.create_new_user(username, email, password)
-        if user_id is None:
-            return {"error": "existing credentials"}
-        
-        app.generate_default_user_data(user_id)
-
-        # token = auth_manager.dump_token(user_id)
-        return {"userId": user_id}
-
-    @server.websocket('/ws')
-    async def ws():
-        print("incoming ws")
-
-        # Authenticate before dedicating any resources
-        auth_req = json.loads(await websocket.receive())
-        if "username" not in auth_req or "password" not in auth_req:
-            print("Invalid auth request:", auth_req)
-            await websocket.send(json.dumps({'error': 'invalid auth'}))
-            return
-
-        user_id = app.verify_user_credentials(auth_req['username'], auth_req['password'])
-        if user_id is None:
-            await websocket.send(json.dumps({'error': 'invalid credentials'}))
-            return
-
-        send_queue = asyncio.Queue()
-
-        async def sender_task():
-            nonlocal send_queue
-            while True:
-                msg = await send_queue.get()
-                await websocket.send(json.dumps(msg))
-
-        async def send(msg):
-            await send_queue.put(msg)
-
-        asyncio.create_task(sender_task())
-
-        while True:
-            req = json.loads(await websocket.receive())
-            action = req['action']
-            body = req.get('body', {})
-            nonce = req.get('nonce', None)
-
-            if action == "ping":
-                print("ping")
-                await send({'message': "pong"})
-
-            elif action == "create-agent":
-                agent_id = app.create_agent(
-                    user_id=user_id,
-                    name=body['name'],
-                    description=body['description'],
-                    picture=body.get('picture'),
-                    kind=body['kind'],
-                    data=body['data']
-                )
-                await send({'message': "agent created", 'agentId': agent_id, 'nonce': nonce})
-
-            elif action == "get-agents":
-                agents_data = app.get_agents(user_id)
-                await send({'message': "agents retrieved", 'agents': agents_data, 'nonce': nonce})
-
-            elif action == "create-chat":
-                chat_id = app.create_chat(
-                    user_id=user_id,
-                    name=body['name'],
-                    description=body.get('description'),
-                    picture=body.get('picture')
-                )
-                await send({'message': "chat created", 'chatId': chat_id, 'nonce': nonce})
-
-            elif action == "get-chats":
-                chats_data = app.get_chats(user_id)
-                await send({'message': "chats retrieved", 'chats': chats_data, 'nonce': nonce})
-
-            elif action == "get-chat-messages":
-                messages_data = app.get_chat_messages(
-                    body['chatId'],
-                    earliest_index=body.get('earliestIndex', 0)
-                )
-                await send({'message': "chat messages retrieved", 'messages': messages_data, 'nonce': nonce})
-
-            else:
-                print("Unknown message: " + action)
-
-    server.run(debug=True)
-
 if __name__ == "__main__":
     if len(sys.argv) > 1:
         data_directory = sys.argv[1]

backend/server.py

@@ -0,0 +1,124 @@
+from quart import Quart, Response, request, send_from_directory, websocket
+import json
+import asyncio
+
+def start_server(app):
+    server = Quart(__name__, static_folder='../frontend/build/web')
+
+    @server.route('/')
+    async def index():
+        return await send_from_directory(server.static_folder, 'index.html')
+
+    @server.route('/<path:path>')
+    async def serve_static(path):
+        return await send_from_directory(server.static_folder, path)
+
+    @server.route('/user/signup', methods=['POST'])
+    async def signup():
+        data = await request.get_json()
+
+        user_id = None
+
+        username = data['username']
+        email = data['email']
+        password = data['password']
+
+        user_id = app.create_new_user(username, email, password)
+        if user_id is None:
+            return {"error": "existing credentials"}
+        
+        app.generate_default_user_data(user_id)
+
+        # token = auth_manager.dump_token(user_id)
+        return {"userId": user_id}
+
+    @server.websocket('/ws')
+    async def ws():
+        print("incoming ws")
+
+        # Authenticate before dedicating any resources
+        auth_req = json.loads(await websocket.receive())
+        if "username" not in auth_req or "password" not in auth_req:
+            print("Invalid auth request:", auth_req)
+            await websocket.send(json.dumps({'error': 'invalid auth'}))
+            return
+
+        user_id = app.verify_user_credentials(auth_req['username'], auth_req['password'])
+        if user_id is None:
+            await websocket.send(json.dumps({'error': 'invalid credentials'}))
+            return
+
+        send_queue = asyncio.Queue()
+        app.add_user_connection(user_id, send_queue)
+
+        async def sender_task():
+            nonlocal send_queue
+            while True:
+                msg = await send_queue.get()
+                await websocket.send(json.dumps(msg))
+
+        async def send(msg):
+            await send_queue.put(msg)
+
+        asyncio.create_task(sender_task())
+
+        try:
+            while True:
+                req = json.loads(await websocket.receive())
+                action = req['action']
+                body = req.get('body', {})
+                nonce = req.get('nonce', None)
+
+                if action == "ping":
+                    print("ping")
+                    await send({'message': "pong"})
+
+                elif action == "create-agent":
+                    agent_id = app.create_agent(
+                        user_id=user_id,
+                        name=body['name'],
+                        description=body['description'],
+                        picture=body.get('picture'),
+                        kind=body['kind'],
+                        data=body['data']
+                    )
+                    await send({'message': "agent created", 'agentId': agent_id, 'nonce': nonce})
+
+                elif action == "get-agents":
+                    agents_data = app.get_agents(user_id)
+                    await send({'message': "agents retrieved", 'agents': agents_data, 'nonce': nonce})
+
+                elif action == "create-chat":
+                    chat_id = app.create_chat(
+                        user_id=user_id,
+                        name=body['name'],
+                        description=body.get('description'),
+                        picture=body.get('picture')
+                    )
+                    await send({'message': "chat created", 'chatId': chat_id, 'nonce': nonce})
+
+                elif action == "get-chats":
+                    chats_data = app.get_chats(user_id)
+                    await send({'message': "chats retrieved", 'chats': chats_data, 'nonce': nonce})
+
+                elif action == "get-chat-messages":
+                    messages_data = app.get_chat_messages(
+                        body['chatId'],
+                        earliest_index=body.get('earliestIndex', 0)
+                    )
+                    await send({'message': "chat messages retrieved", 'messages': messages_data, 'nonce': nonce})
+
+                elif action == "send-chat-message":
+                    message = app.send_chat_message(
+                        chat_id=body['chatId'],
+                        content=body['content'],
+                    )
+                    await send({'message': "message sent", 'messageId': message['id'], 'nonce': nonce})
+
+                else:
+                    print("Unknown message: " + action)
+
+        finally:
+            app.remove_user_connection(user_id, send_queue)
+
+    server.run(debug=True)

testing/suite.js

@@ -5,6 +5,9 @@ const ws = new WebSocket('ws://localhost:5000/ws');
 var nonceCounter = 1;
 var nonceCallbacks = {};
 
+var messageQueue = [];
+var messageResolverQueue = [];
+
 // Setup
 
 ws.on('open', function open() {
@@ -26,7 +29,6 @@ ws.on('message', function incoming(rawMessage) {
         } else {
             messageQueue.push(parsed);
         }
-        console.log(parsed);
     }
 });
 
@@ -84,9 +86,11 @@ async function testChats () {
     let response = await send('get-chats')
     let chat = response.chats[0];
 
-    let msgRes = await send('chat-send-message', {
-        chatId: chat.id,
-        content: "Hello World!",
+    let msgRes = await send('send-chat-message', {
+        body: {
+            chatId: chat.id,
+            content: "Hello World!",
+        }
     });
     console.log('message sent', msgRes);