Implement basic agent and toolset functionality

backend/app.py

@@ -9,9 +9,9 @@ from sqlalchemy.sql.expression import select
 from quart import Quart, Response, request, send_from_directory, jsonify, websocket
 import replicate
 
-from database import create_sql_engine, User, Agent, Chat, Message, Membership
+import database as db
 from server import start_server
-from agent import AgentControl
+from agent import Agent
 
 MESSAGE_PAGE_SIZE = 32
 
@@ -23,7 +23,7 @@ class ChatContext:
 class App:
     def __init__(self, data_directory, replicate_api_token=None):
         self.data_directory = data_directory
-        self.sql_engine = create_sql_engine(self.data_directory)
+        self.sql_engine = db.create_sql_engine(self.data_directory)
         self.user_connections = {}
         self.streams = {}
 
@@ -50,7 +50,7 @@ class App:
 
     def create_agent(self, user_id, name, description, picture, kind, data):
         with self.sql_session() as session:
-            agent = Agent()
+            agent = db.Agent()
             agent.user_id = user_id
             agent.name = name
             agent.description = description
@@ -65,7 +65,7 @@ class App:
 
     def create_chat(self, user_id, name, description, picture):
         with self.sql_session() as session:
-            chat = Chat()
+            chat = db.Chat()
             chat.user_id = user_id
             chat.name = name
             chat.description = description
@@ -78,7 +78,7 @@ class App:
 
     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()
+            chats = session.execute(select(db.Chat).where(db.Chat.user_id == user_id)).scalars().all()
             return [{
                 "id": chat.id,
                 "name": chat.name,
@@ -88,7 +88,7 @@ class App:
 
     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()
+            agents = session.execute(select(db.Agent).where(db.Agent.user_id == user_id)).scalars().all()
             return [{
                 "id": agent.id,
                 "name": agent.name,
@@ -99,7 +99,7 @@ class App:
     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.time.asc())
+                select(db.Message).where(db.Message.chat_id == chat_id).order_by(db.Message.time.asc())
             ).scalars().fetchmany(MESSAGE_PAGE_SIZE)
 
             return [{
@@ -113,7 +113,7 @@ class App:
     def verify_user_credentials(self, username, password):
         with self.sql_session() as session:
             user = session.execute(
-                select(User).where(User.username == username)
+                select(db.User).where(db.User.username == username)
             ).scalars().first()
 
             if user and user.password == password:
@@ -123,11 +123,11 @@ class App:
     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)
+                select(db.User.id).where(db.User.email == email or db.User.username == username)
             ).first()
 
             if result is None:
-                user = User()
+                user = db.User()
                 user.email = email
                 user.username = username
                 user.password = password
@@ -143,7 +143,7 @@ class App:
 
     def add_agent_to_chat(self, chat_id, agent_id):
         with self.sql_session() as session:
-            membership = Membership()
+            membership = db.Membership()
             membership.chat_id = chat_id
             membership.agent_id = agent_id
             
@@ -154,7 +154,7 @@ class App:
     # 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 = db.Message()
             message.chat_id = chat_id
             message.time = int(datetime.now().timestamp())
             message.sender = sender
@@ -166,11 +166,11 @@ class App:
 
             # Notify chat owner
             if sender != 'user':
-                chat_owner = session.execute(select(Chat.user_id).where(Chat.id == chat_id)).scalar()
+                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(Membership.agent_id).where(Membership.chat_id == chat_id)).scalars().all()
+            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)
@@ -190,13 +190,13 @@ class App:
     def get_agent_control(self, agent_id):
         with self.sql_session() as session:
             agent = session.execute(
-                select(Agent).where(Agent.id == agent_id)
+                select(db.Agent).where(db.Agent.id == agent_id)
             ).scalars().first()
 
             if not agent:
                 raise Exception(f"Agent not found with id {agent_id}")
         
-            control = AgentControl.create_from_data(app, agent_id, agent.kind, agent.data)
+            control = Agent.create_from_data(app, agent_id, agent.kind, agent.data)
             return control
 
     async def send_agent_prompt(self, agent_id, message):