move backend source to src

backend/src/app.py

@@ -0,0 +1,300 @@
+from datetime import datetime
+import asyncio
+import json
+import sys
+import os
+
+from sqlalchemy.orm import Session
+from sqlalchemy.sql.expression import select
+from quart import Quart, Response, request, send_from_directory, jsonify, websocket
+from qdrant_client import QdrantClient
+import replicate
+
+import database as db
+import agent as agent_mod
+from server import start_server
+from chat import ChatView, Message, MESSAGE_PAGE_SIZE
+
+from agents.conversation import ConversationAgent
+from app_mixins import ChatMixin
+
+class App(ChatMixin):
+    def __init__(self, data_directory, replicate_api_token=None):
+        self.data_directory = data_directory
+        self.user_connections = {}
+        self.remote_connections = {}
+        self.streams = {}
+
+        # Initialize Clients
+        self.sql_engine = db.create_sql_engine(self.data_directory)
+        self.qdrant_client = QdrantClient(path=f"{self.data_directory}/qdrant_data")
+
+        if replicate_api_token:
+            self.replicate_client = replicate.Client(api_token=replicate_api_token)
+        else:
+            raise Exception("Horde only works with a replicate connection. Please provide an API Token.")
+
+        # `Agent` is the kind 'basic', will be class `BasicAgent` in the future
+        self.agent_classes = [
+            agent_mod.BasicAgent,
+            ConversationAgent,
+            agent_mod.LlamaAgent,
+            agent_mod.CoderV1Agent,
+            agent_mod.ToolUsingAgent,
+            agent_mod.DebugAgent,
+        ]
+
+    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 add_remote_connection(self, user_id, remotename, channel):
+        if user_id not in self.remote_connections:
+            self.remote_connections[user_id] = {}
+        user_connections = self.remote_connections[user_id]
+
+        if channel_list not in user_connections[remotename]:
+            user_connections[remotename] = []
+        user_connections[remotename].append(channel)
+
+    def create_agent(self, user_id, name, agentname, description, picture, kind, data, main_chat_id=None):
+        with self.sql_session() as session:
+            if main_chat_id is None:
+                chat = db.Chat()
+                chat.user_id = user_id
+                chat.name = name
+
+                session.add(chat)
+                session.commit()
+                main_chat_id = chat.id
+
+            agent = db.Agent()
+            agent.user_id = user_id
+            agent.name = name
+            agent.agentname = agentname
+            agent.main_chat_id = main_chat_id
+            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}")
+
+            membership = db.Membership()
+            membership.chat_id = main_chat_id
+            membership.agent_id = agent.id
+            session.add(membership)
+            session.commit()
+
+            return agent.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:
+            agent = session.execute(
+                select(db.Agent).where(db.Agent.id == agent_id)
+            ).scalars().first()
+
+            if not agent:
+                raise Exception(f"Agent not found with id {agent_id}")
+            
+            summary = []
+
+            if name != None:
+                agent.name = name
+                summary.append('name')
+            if description != None:
+                agent.description = description
+                summary.append('description')
+            if picture != None:
+                agent.picture = picture
+                summary.append('picture')
+            if agentname != None:
+                agent.agentname = agentname
+                summary.append('agentname')
+
+            if data:
+                # Assign a copy to trigger a value change
+                agent.data = agent.data.copy() if agent.data else {}
+                agent.data.update(data)
+                summary.extend([f"data.{name}" for name in data.keys()])
+
+            session.commit()
+            print(f"Updated agent #{agent.id} fields: {summary}")
+
+    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()
+            return [{
+                "id": agent.id,
+                "name": agent.name,
+                "agentname": agent.agentname,
+                "description": agent.description,
+                "picture": agent.picture,
+                "kind": agent.kind,
+            } for agent in agents]
+
+    def get_agent_kinds(self):
+        return [kind.identifier for kind in self.agent_classes]
+
+    def get_agent_kind_info(self, kind):
+        agent_class = next((it for it in self.agent_classes if it.identifier == kind), None)
+        if agent_class is None:
+            raise Exception(f"invalid agent kind: {kind}")
+        return {
+            "name": agent_class.identifier,
+            "description": getattr(agent_class, 'description', None),
+        }
+
+    def verify_user_credentials(self, username, password):
+        with self.sql_session() as session:
+            user = session.execute(
+                select(db.User).where(db.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(db.User.id).where(db.User.email == email or db.User.username == username)
+            ).first()
+
+            if result is None:
+                user = db.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 get_agent_control(self, agent_id):
+        with self.sql_session() as session:
+            agent = session.execute(
+                select(db.Agent).where(db.Agent.id == agent_id)
+            ).scalars().first()
+
+            if not agent:
+                raise Exception(f"Agent not found with id {agent_id}")
+            
+            agent_class = next((kind for kind in self.agent_classes if kind.identifier == agent.kind), None)
+            if agent_class is None:
+                raise Exception(f"invalid agent kind: {agent.kind}")
+        
+            control = agent_class(app, agent.data)
+            control._set_basic_info(agent_id, agent.name)
+            return control
+
+    def get_agent_config_fields(self, agent_id):
+        with self.sql_session() as session:
+            agent = session.execute(
+                select(db.Agent).where(db.Agent.id == agent_id)
+            ).scalars().first()
+            if not agent:
+                raise Exception(f"Agent not found with id {agent_id}")
+
+            agent_class = next((it for it in self.agent_classes if it.identifier == agent.kind), None)
+            if not agent_class:
+                raise Exception(f"Invalid agent kind: {agent.kind}")
+
+            data = agent.data or {}
+            fields = [{
+                'name': field.name,
+                'type': field.type,
+                'value': data.get(field.name, field.default)
+                } for field in agent_class.schema_fields]
+            
+            return fields
+        return {}
+
+    async def send_agent_prompt(self, agent_id, message):
+        control = self.get_agent_control(agent_id)
+        response = await control.prompt(message)
+        return response
+
+    async def call_agent_action(self, agent_id, action, args):
+        control = self.get_agent_control(agent_id)
+        response = await control.execute_action(action, args)
+        return response
+
+    def generate_default_user_data(self, user_id):
+        chat_id = self.create_chat(
+            user_id=user_id,
+            name="Initial Chat",
+            description="First user chat with Intro.",
+            picture=None,
+            members=[],
+        )
+        agent_id = self.create_agent(
+            user_id=user_id,
+            name="Intro",
+            agentname="intro",
+            main_chat_id=chat_id,
+            description="Introduces you to the application.",
+            picture=None,
+            kind="conversation", # Should be 'intro'
+            data={}
+        )
+    
+    def get_agent_resource(self, agent_id, kind, key):
+        with self.sql_session() as session:
+            stored_value = session.execute(select(db.AgentResource.value).where(
+                db.AgentResource.agent_id == agent_id and db.AgentResource.kind == kind and db.AgentResource.key == key
+            )).scalar()
+            return stored_value
+
+    def save_agent_resource(self, agent_id, kind, key, value):
+        with self.sql_session() as session:
+            existing_resource = session.execute(select(db.AgentResource).where(
+                db.AgentResource.agent_id == agent_id and db.AgentResource.kind == kind and db.AgentResource.key == key
+            )).scalars().first()
+
+            if existing_resource:
+                existing_resource.value = value
+            else:
+                new_resource = db.AgentResource()
+                new_resource.agent_id = agent_id
+                new_resource.kind = kind
+                new_resource.key = key
+                new_resource.value = value
+
+                session.add(new_resource)
+            session.commit()
+            print(f"Saved agent resource: {agent_id} - {kind} - {key} = {value}")
+
+
+if __name__ == "__main__":
+    if len(sys.argv) > 1:
+        data_directory = sys.argv[1]
+    else:
+        print("No directory specified. Using current directory as default.")
+        data_directory = "."
+
+    app = App(
+        data_directory=data_directory,
+        replicate_api_token=os.getenv('REPLICATE_API_TOKEN'),
+    )
+    start_server(app)