258 lines
8.9 KiB
Python
258 lines
8.9 KiB
Python
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
|
|
import replicate
|
|
|
|
from database import create_sql_engine, User, Agent, Chat, Message, Membership
|
|
from server import start_server
|
|
from agent import AgentControl
|
|
|
|
MESSAGE_PAGE_SIZE = 32
|
|
|
|
class ChatContext:
|
|
def __init__(self, messages, prelude=None):
|
|
self.messages=messages
|
|
self.prelude=prelude
|
|
|
|
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.user_connections = {}
|
|
self.streams = {}
|
|
|
|
if replicate_api_token:
|
|
self.replicate_client = replicate.Client(api_token=replicate_api_token)
|
|
else:
|
|
raise Exception("Horde only works with a replicate client. Please provide an API Token.")
|
|
|
|
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()
|
|
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.time.asc())
|
|
).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}")
|
|
|
|
# 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 get_chat_context(self, chat_id):
|
|
return ChatContext(self.get_chat_messages(chat_id))
|
|
|
|
def get_agent_control(self, agent_id):
|
|
with self.sql_session() as session:
|
|
agent = session.execute(
|
|
select(Agent).where(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)
|
|
return control
|
|
|
|
async def send_agent_prompt(self, agent_id, message):
|
|
control = self.get_agent_control(agent_id)
|
|
response = await control.prompt(message)
|
|
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,
|
|
'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):
|
|
control = self.get_agent_control(agent_id)
|
|
asyncio.create_task(control.received_message(message))
|
|
|
|
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)
|
|
|
|
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 = "."
|
|
|
|
print("api token", os.getenv('REPLICATE_API_TOKEN'))
|
|
|
|
app = App(
|
|
data_directory=data_directory,
|
|
replicate_api_token=os.getenv('REPLICATE_API_TOKEN'),
|
|
)
|
|
start_server(app) |