diff --git a/backend/app.py b/backend/app.py index 3ecbce7..5206a8c 100644 --- a/backend/app.py +++ b/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: - def __init__(self, messages, prelude=None): - self.messages=messages - self.prelude=prelude - -class App: +class App(ChatMixin): 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( diff --git a/backend/app_mixins/__init__.py b/backend/app_mixins/__init__.py new file mode 100644 index 0000000..7094aec --- /dev/null +++ b/backend/app_mixins/__init__.py @@ -0,0 +1 @@ +from app_mixins.chat import ChatMixin \ No newline at end of file diff --git a/backend/app_mixins/chat.py b/backend/app_mixins/chat.py new file mode 100644 index 0000000..233aafb --- /dev/null +++ b/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))