feat: agent app mixin
This commit is contained in:
Shivering Ocelot committed
1 parent
7e68348368
commit
75e49fcbcd
3 files changed
+233
-222
No files matched your search
+2
-222
@@ -21,9 +21,9 @@ from server import start_server
|
||||
from chat import ChatView, Message, MESSAGE_PAGE_SIZE
|
||||
from remote import RemoteConnection, RemoteToolSpec, parse_function_specs
|
||||
|
||||
from app_mixins import ChatMixin
|
||||
from app_mixins import ChatMixin, AgentMixin
|
||||
|
||||
class App(ChatMixin):
|
||||
class App(ChatMixin, AgentMixin):
|
||||
def __init__(self, data_directory, replicate_token=None, openrouter_token=None):
|
||||
self.data_directory = data_directory
|
||||
self.user_connections = {}
|
||||
@@ -155,127 +155,6 @@ class App(ChatMixin):
|
||||
else:
|
||||
RemoteConnection.reject(request_id, error);
|
||||
|
||||
# --- Agent Methods (Asynchronous) ---
|
||||
|
||||
async def get_agent(self, agent_id):
|
||||
if tmp.is_tmp_id(agent_id):
|
||||
agent = self.tmp.agents.get(agent_id)
|
||||
else:
|
||||
agent = await db.Agent.get_or_none(id=agent_id)
|
||||
|
||||
if not agent:
|
||||
raise Exception(f"Agent not found with id {agent_id}")
|
||||
return agent
|
||||
|
||||
async def agent_id_from_name(self, agentname: str):
|
||||
agent_id = await db.Agent.get_or_none(agentname=agentname).values_list("id", flat=True)
|
||||
if not agent_id:
|
||||
raise Exception(f"Agent not found with agentname {agentname}")
|
||||
return agent_id
|
||||
|
||||
async def create_agent(self, user_id, name, agentname, description, picture, kind, data, main_chat_id=None, temporary=False):
|
||||
print("create_agent; user_id, name:", user_id, name)
|
||||
|
||||
if temporary:
|
||||
agent = tmp.Agent.create(
|
||||
self,
|
||||
user_id=user_id,
|
||||
name=name,
|
||||
agentname=agentname,
|
||||
main_chat_id=main_chat_id,
|
||||
description=description,
|
||||
picture=picture,
|
||||
kind=kind,
|
||||
data=data or {}
|
||||
)
|
||||
|
||||
if main_chat_id is not None:
|
||||
if tmp.is_tmp_id(main_chat_id):
|
||||
self.tmp.chats[main_chat_id].participants.append(agent.id)
|
||||
else:
|
||||
print(f"ERROR: cannot use {main_chat_id} as main chat for temporary agent {agentname or agent.id}")
|
||||
agent.main_chat_id = None
|
||||
|
||||
return agent.id
|
||||
|
||||
if main_chat_id is None:
|
||||
chat = await db.Chat.create(
|
||||
user_id=user_id,
|
||||
name=name
|
||||
)
|
||||
main_chat_id = chat.id
|
||||
|
||||
agent = await db.Agent.create(
|
||||
user_id=user_id,
|
||||
name=name,
|
||||
agentname=agentname,
|
||||
main_chat_id=main_chat_id,
|
||||
description=description,
|
||||
picture=picture,
|
||||
kind=kind,
|
||||
data=data or {}
|
||||
)
|
||||
print(f"Created agent #{agent.id}: {agent.name}")
|
||||
|
||||
print(f"... adding agent {repr(agent.id)} to chat {repr(main_chat_id)}")
|
||||
await db.Membership.create(
|
||||
chat_id=main_chat_id,
|
||||
agent_id=agent.id
|
||||
)
|
||||
return agent.id
|
||||
|
||||
async def update_agent(self, agent_id, name=None, description=None, picture=None, agentname=None, data=None):
|
||||
agent = await self.get_agent(agent_id)
|
||||
|
||||
summary = []
|
||||
if name is not None:
|
||||
agent.name = name
|
||||
summary.append('name')
|
||||
if description is not None:
|
||||
agent.description = description
|
||||
summary.append('description')
|
||||
if picture is not None:
|
||||
agent.picture = picture
|
||||
summary.append('picture')
|
||||
if agentname is not None:
|
||||
agent.agentname = agentname
|
||||
summary.append('agentname')
|
||||
|
||||
if data:
|
||||
# Tortoise tracks changes in JSONField automatically
|
||||
agent.data.update(data)
|
||||
summary.extend([f"data.{k}" for k in data.keys()])
|
||||
|
||||
if db.is_instance(agent):
|
||||
await agent.save()
|
||||
print(f"Updated agent #{agent.id} fields: {summary}")
|
||||
|
||||
async def get_agents(self, user_id):
|
||||
agents = await db.Agent.filter(user_id=user_id).all()
|
||||
if user_id in self.tmp.user_agents:
|
||||
agents += [self.tmp.agents[agent_id] for agent_id in self.tmp.user_agents[user_id]]
|
||||
|
||||
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),
|
||||
}
|
||||
|
||||
# --- User Methods (Asynchronous) ---
|
||||
|
||||
async def verify_user_id(self, user_id):
|
||||
@@ -303,82 +182,8 @@ class App(ChatMixin):
|
||||
print(f"Signup error: '{email}' or '{username}' are already used")
|
||||
return None
|
||||
|
||||
# --- Agent Control & Actions ---
|
||||
|
||||
async def get_agent_control(self, agent_id):
|
||||
agent = await self.get_agent(agent_id)
|
||||
resources = await self.get_agent_resources(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(self, agent.data, resources)
|
||||
control._set_basic_info(agent_id, agent.name, agent.user_id)
|
||||
return control
|
||||
|
||||
async def get_agent_config_fields(self, agent_id):
|
||||
agent = await self.get_agent(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 {}
|
||||
return [{
|
||||
'name': field.name,
|
||||
'type': field.type,
|
||||
'value': data.get(field.name, field.default)
|
||||
} for field in agent_class.schema_fields]
|
||||
|
||||
async def send_agent_prompt(self, agent_id, instruction, input=None, options={}):
|
||||
control = await self.get_agent_control(agent_id)
|
||||
response = await control.prompt(instruction, input, **options)
|
||||
return response
|
||||
|
||||
async def call_agent_action(self, agent_id, action, args):
|
||||
control = await self.get_agent_control(agent_id)
|
||||
response = await control.execute_action(action, args)
|
||||
return response
|
||||
|
||||
async def generate_default_user_data(self, user_id):
|
||||
chat_id = await self.create_chat(
|
||||
user_id=user_id,
|
||||
name="Initial Chat",
|
||||
description="First user chat with Intro.",
|
||||
picture=None,
|
||||
members=[],
|
||||
)
|
||||
await 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",
|
||||
data={}
|
||||
)
|
||||
|
||||
# --- Resource Management ---
|
||||
|
||||
async def get_agent_resources(self, agent_id, kind=None, key=None):
|
||||
filter = {"agent_id": agent_id}
|
||||
if kind: filter['kind'] = kind
|
||||
if key: filter['key'] = key
|
||||
|
||||
values = await db.AgentResource.filter(**filter)
|
||||
return values
|
||||
|
||||
async def save_agent_resource(self, agent_id, kind, key, value):
|
||||
await db.AgentResource.update_or_create(
|
||||
agent_id=agent_id,
|
||||
kind=kind,
|
||||
key=key,
|
||||
value=value,
|
||||
)
|
||||
print(f"Saved agent resource: {agent_id} - {kind} - {key}")
|
||||
|
||||
async def get_models(self, kind=None):
|
||||
return await db.Model.all()
|
||||
async def get_model(self, model_id):
|
||||
@@ -408,31 +213,6 @@ class App(ChatMixin):
|
||||
print(f"Saved model: {model.id} - {name}")
|
||||
return model
|
||||
|
||||
async def create_agent_run(self, agent_id):
|
||||
agent = await self.get_agent(agent_id)
|
||||
run = agent_mod.AgentRun.create(
|
||||
self,
|
||||
agent_id=agent_id,
|
||||
user_id=agent.user_id,
|
||||
)
|
||||
return run.id
|
||||
|
||||
async def cancel_agent_run(self, run_id):
|
||||
"""Cancel an ongoing agent run by ID."""
|
||||
run = self.tmp.runs.get(run_id)
|
||||
if run:
|
||||
run.cancel()
|
||||
return True
|
||||
return False
|
||||
|
||||
async def get_agent_runs(self, agent_id=None, user_id=None):
|
||||
runs = list(self.tmp.runs.values())
|
||||
if agent_id:
|
||||
runs = [r for r in runs if r.agent_id == agent_id]
|
||||
elif user_id:
|
||||
runs = [r for r in runs if r.user_id == user_id]
|
||||
return [r.serialize() for r in runs]
|
||||
|
||||
# --- Main Entry Point ---
|
||||
|
||||
async def main():
|
||||
|
||||
@@ -1 +1,2 @@
|
||||
from app_mixins.chat import ChatMixin
|
||||
from app_mixins.agent import AgentMixin
|
||||
@@ -0,0 +1,230 @@
|
||||
import asyncio
|
||||
from datetime import datetime
|
||||
|
||||
import database as db
|
||||
import temporary as tmp
|
||||
import agent as agent_mod
|
||||
import agents as agents_mod
|
||||
|
||||
class AgentMixin:
|
||||
# --- Base ---
|
||||
|
||||
async def get_agent(self, agent_id):
|
||||
if tmp.is_tmp_id(agent_id):
|
||||
agent = self.tmp.agents.get(agent_id)
|
||||
else:
|
||||
agent = await db.Agent.get_or_none(id=agent_id)
|
||||
|
||||
if not agent:
|
||||
raise Exception(f"Agent not found with id {agent_id}")
|
||||
return agent
|
||||
|
||||
async def agent_id_from_name(self, agentname: str):
|
||||
agent_id = await db.Agent.get_or_none(agentname=agentname).values_list("id", flat=True)
|
||||
if not agent_id:
|
||||
raise Exception(f"Agent not found with agentname {agentname}")
|
||||
return agent_id
|
||||
|
||||
async def create_agent(self, user_id, name, agentname, description, picture, kind, data, main_chat_id=None, temporary=False):
|
||||
print("create_agent; user_id, name:", user_id, name)
|
||||
|
||||
if temporary:
|
||||
agent = tmp.Agent.create(
|
||||
self,
|
||||
user_id=user_id,
|
||||
name=name,
|
||||
agentname=agentname,
|
||||
main_chat_id=main_chat_id,
|
||||
description=description,
|
||||
picture=picture,
|
||||
kind=kind,
|
||||
data=data or {}
|
||||
)
|
||||
|
||||
if main_chat_id is not None:
|
||||
if tmp.is_tmp_id(main_chat_id):
|
||||
self.tmp.chats[main_chat_id].participants.append(agent.id)
|
||||
else:
|
||||
print(f"ERROR: cannot use {main_chat_id} as main chat for temporary agent {agentname or agent.id}")
|
||||
agent.main_chat_id = None
|
||||
|
||||
return agent.id
|
||||
|
||||
if main_chat_id is None:
|
||||
chat = await db.Chat.create(
|
||||
user_id=user_id,
|
||||
name=name
|
||||
)
|
||||
main_chat_id = chat.id
|
||||
|
||||
agent = await db.Agent.create(
|
||||
user_id=user_id,
|
||||
name=name,
|
||||
agentname=agentname,
|
||||
main_chat_id=main_chat_id,
|
||||
description=description,
|
||||
picture=picture,
|
||||
kind=kind,
|
||||
data=data or {}
|
||||
)
|
||||
print(f"Created agent #{agent.id}: {agent.name}")
|
||||
|
||||
print(f"... adding agent {repr(agent.id)} to chat {repr(main_chat_id)}")
|
||||
await db.Membership.create(
|
||||
chat_id=main_chat_id,
|
||||
agent_id=agent.id
|
||||
)
|
||||
return agent.id
|
||||
|
||||
async def update_agent(self, agent_id, name=None, description=None, picture=None, agentname=None, data=None):
|
||||
agent = await self.get_agent(agent_id)
|
||||
|
||||
summary = []
|
||||
if name is not None:
|
||||
agent.name = name
|
||||
summary.append('name')
|
||||
if description is not None:
|
||||
agent.description = description
|
||||
summary.append('description')
|
||||
if picture is not None:
|
||||
agent.picture = picture
|
||||
summary.append('picture')
|
||||
if agentname is not None:
|
||||
agent.agentname = agentname
|
||||
summary.append('agentname')
|
||||
|
||||
if data:
|
||||
# Tortoise tracks changes in JSONField automatically
|
||||
agent.data.update(data)
|
||||
summary.extend([f"data.{k}" for k in data.keys()])
|
||||
|
||||
if db.is_instance(agent):
|
||||
await agent.save()
|
||||
print(f"Updated agent #{agent.id} fields: {summary}")
|
||||
|
||||
async def get_agents(self, user_id):
|
||||
agents = await db.Agent.filter(user_id=user_id).all()
|
||||
if user_id in self.tmp.user_agents:
|
||||
agents += [self.tmp.agents[agent_id] for agent_id in self.tmp.user_agents[user_id]]
|
||||
|
||||
return [{
|
||||
"id": agent.id,
|
||||
"name": agent.name,
|
||||
"agentname": agent.agentname,
|
||||
"description": agent.description,
|
||||
"picture": agent.picture,
|
||||
"kind": agent.kind,
|
||||
} for agent in agents]
|
||||
|
||||
async def get_agent_resources(self, agent_id, kind=None, key=None):
|
||||
filter = {"agent_id": agent_id}
|
||||
if kind: filter['kind'] = kind
|
||||
if key: filter['key'] = key
|
||||
|
||||
values = await db.AgentResource.filter(**filter)
|
||||
return values
|
||||
|
||||
async def save_agent_resource(self, agent_id, kind, key, value):
|
||||
await db.AgentResource.update_or_create(
|
||||
agent_id=agent_id,
|
||||
kind=kind,
|
||||
key=key,
|
||||
value=value,
|
||||
)
|
||||
print(f"Saved agent resource: {agent_id} - {kind} - {key}")
|
||||
|
||||
# --- Control & Actions ---
|
||||
|
||||
async def get_agent_control(self, agent_id):
|
||||
agent = await self.get_agent(agent_id)
|
||||
resources = await self.get_agent_resources(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(self, agent.data, resources)
|
||||
control._set_basic_info(agent_id, agent.name, agent.user_id)
|
||||
return control
|
||||
|
||||
async def create_agent_run(self, agent_id):
|
||||
agent = await self.get_agent(agent_id)
|
||||
run = agent_mod.AgentRun.create(
|
||||
self,
|
||||
agent_id=agent_id,
|
||||
user_id=agent.user_id,
|
||||
)
|
||||
return run.id
|
||||
|
||||
async def cancel_agent_run(self, run_id):
|
||||
"""Cancel an ongoing agent run by ID."""
|
||||
run = self.tmp.runs.get(run_id)
|
||||
if run:
|
||||
run.cancel()
|
||||
return True
|
||||
return False
|
||||
|
||||
async def get_agent_runs(self, agent_id=None, user_id=None):
|
||||
runs = list(self.tmp.runs.values())
|
||||
if agent_id:
|
||||
runs = [r for r in runs if r.agent_id == agent_id]
|
||||
elif user_id:
|
||||
runs = [r for r in runs if r.user_id == user_id]
|
||||
return [r.serialize() for r in runs]
|
||||
|
||||
async def send_agent_prompt(self, agent_id, instruction, input=None, options={}):
|
||||
control = await self.get_agent_control(agent_id)
|
||||
response = await control.prompt(instruction, input, **options)
|
||||
return response
|
||||
|
||||
async def call_agent_action(self, agent_id, action, args):
|
||||
control = await self.get_agent_control(agent_id)
|
||||
response = await control.execute_action(action, args)
|
||||
return response
|
||||
|
||||
# --- Information ---
|
||||
|
||||
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),
|
||||
}
|
||||
|
||||
async def get_agent_config_fields(self, agent_id):
|
||||
agent = await self.get_agent(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 {}
|
||||
return [{
|
||||
'name': field.name,
|
||||
'type': field.type,
|
||||
'value': data.get(field.name, field.default)
|
||||
} for field in agent_class.schema_fields]
|
||||
|
||||
async def generate_default_user_data(self, user_id):
|
||||
chat_id = await self.create_chat(
|
||||
user_id=user_id,
|
||||
name="Initial Chat",
|
||||
description="First user chat with Intro.",
|
||||
picture=None,
|
||||
members=[],
|
||||
)
|
||||
await 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",
|
||||
data={}
|
||||
)
|
||||
Reference in new issue
Block a user