backend: implement basic functionality in app class. testing: organize and improve javascript test swuite

backend/app.py

@@ -1,11 +1,10 @@
-import sys
-import json
-import asyncio
-
-from database import create_sql_engine, User, Agent
+from database import create_sql_engine, User, Agent, Chat, Message, Membership
 from sqlalchemy.orm import Session
 from sqlalchemy.sql.expression import select
 from quart import Quart, Response, request, send_from_directory, jsonify, websocket
+import asyncio
+import json
+import sys
 
 class App:
     def __init__(self, data_directory):
@@ -15,7 +14,129 @@ class App:
     def sql_session(self):
         return Session(self.sql_engine)
 
-app = App("./")
+    def create_agent(self, user_id, name, description, picture, kind, data):
+        with self.sql_session() as session:
+            agent = Agent()
+            agent.user_id = user_id
+            agent.name = name
+            agent.description = description
+            agent.picture = picture
+            agent.kind = kind
+            agent.data = data
+
+            session.add(agent)
+            session.commit()
+            print(f"Created agent #{agent.id}: {agent.name}")
+            return agent.id
+
+    def create_chat(self, user_id, name, description, picture):
+        with self.sql_session() as session:
+            chat = 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}")
+            return chat.id
+
+    def get_chats(self, user_id):
+        with self.sql_session() as session:
+            chats = session.execute(select(Chat).where(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(Agent).where(Agent.user_id == user_id)).scalars().all()
+            return [{
+                "id": agent.id,
+                "name": agent.name,
+                "description": agent.description,
+                "picture": agent.picture
+            } for agent in agents]
+
+    def get_chat_messages(self, chat_id, earliest_index=0):
+        with self.sql_session() as session:
+            messages = session.execute(
+                select(Message).where(Message.chat_id == chat_id).order_by(Message.id.desc())
+            ).scalars().fetchmany(MESSAGE_PAGE_SIZE)
+
+            return [{
+                "id": message.id,
+                "time": message.time,
+                "sender": message.sender,
+                "kind": message.kind,
+                "content": message.content
+            } for message in messages]
+
+    def verify_user_credentials(self, username, password):
+        with self.sql_session() as session:
+            user = session.execute(
+                select(User).where(User.username == username)
+            ).scalars().first()
+
+            if user and user.password == password:
+                return user.id
+            return None
+
+    def create_new_user(self, username, email, password):
+        with self.sql_session() as session:
+            result = session.execute(
+                select(User.id).where(User.email == email or User.username == username)
+            ).first()
+
+            if result is None:
+                user = User()
+                user.email = email
+                user.username = username
+                user.password = password
+
+                session.add(user)
+                session.commit()
+                
+                user_id = user.id
+                print(f"Created user #{user_id}: {username} ({email})")
+            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()
+            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}")
+
+    def generate_default_user_data(self, user_id):
+        agent_id = self.create_agent(
+            user_id=user_id,
+            name="Intro",
+            description="Introduces you to the application.",
+            picture=None,
+            kind="intro",
+            data={}
+        )
+
+        # Create a chat with that agent in it
+        chat_id = self.create_chat(
+            user_id=user_id,
+            name="Initial Chat",
+            description=None,
+            picture=None
+        )
+
+        # 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')
@@ -38,31 +159,14 @@ def start_server(app):
         email = data['email']
         password = data['password']
 
-        with app.sql_session() as session:
-            result = session.execute(
-                select(User.id).where(User.email == email or User.username == username)
-            ).first()
-
-            if result is None:
-                user = User()
-                user.email = email
-                user.username = username
-                user.password = password
-
-                session.add(user)
-                session.commit()
-                
-                user_id = user.id
-                print(f"Created user #{user_id}: {username} ({email})")
-            else:
-                user_id = None
-                print(f"Signup error: '{email}' or '{username}' are already used")
-
+        user_id = app.create_new_user(username, email, password)
         if user_id is None:
             return {"error": "existing credentials"}
-        else:
-            # token = auth_manager.dump_token(user_id)
-            return {"userId": user_id}
+        
+        app.generate_default_user_data(user_id)
+
+        # token = auth_manager.dump_token(user_id)
+        return {"userId": user_id}
 
     @server.websocket('/ws')
     async def ws():
@@ -70,14 +174,15 @@ def start_server(app):
 
         # Authenticate before dedicating any resources
         auth_req = json.loads(await websocket.receive())
-        if "username" not in auth_req:
+        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
-        
-        # TODO: Actually authenticate credentials against stored users
-        # TODO: Get the user id from the credentials, use them in create-agent
-        print(f"User {auth_req['username']} connected")
+
+        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()
 
@@ -94,30 +199,49 @@ def start_server(app):
 
         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({'content': "pong"})
+                await send({'message': "pong"})
+
             elif action == "create-agent":
-                with app.sql_session() as session:
-                    agent = 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})
 
-                    agent.user_id = None
-                    agent.name = body['name']
-                    agent.description = body['description']
-                    agent.picture = body.get('picture')
-                    agent.kind = body['kind']
-                    agent.data = body['data']
+            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})
 
-                    session.add(agent)
-                    session.commit()
-                    
-                    print(f"Created user #{agent.id}: {agent.name}")
-                    await send({'content': "agent created", 'agentId': agent.id, 'nonce': nonce})
             else:
                 print("Unknown message: " + action)