276 lines
9.3 KiB
Python
276 lines
9.3 KiB
Python
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(
|
|
{'value': value},
|
|
agent_id=agent_id,
|
|
kind=kind,
|
|
key=key,
|
|
)
|
|
print(f"Saved agent resource: {agent_id} - {kind} - {key}")
|
|
|
|
async def delete_agent_resource(self, agent_id, kind, key):
|
|
await db.AgentResource.filter(
|
|
agent_id=agent_id,
|
|
kind=kind,
|
|
key=key,
|
|
).delete()
|
|
print(f"Deleted 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 prompt_agent(self, agent_id, instruction, input=None, options={}):
|
|
# Mirror agent.prompt's generator so run ownership and lifecycle stay
|
|
# inside the mixin: callers just iterate and receive the same raw
|
|
# inference.Message objects. get_agent_control is awaited lazily on first
|
|
# iteration, keeping the control object alive for the generator's
|
|
# lifetime.
|
|
|
|
def normalize_tool(t):
|
|
# this function makes a lot of sense in ToolSpec itself...
|
|
from tool import ToolSpec, ToolParameter
|
|
if isinstance(t, dict):
|
|
params = []
|
|
|
|
d = t.get('parameters', {})
|
|
required = d.get('required', [])
|
|
props = d.get('properties', {})
|
|
|
|
for key, p in props.items():
|
|
params.append(ToolParameter(
|
|
name=key,
|
|
type=p.get('type'),
|
|
description=p.get('description'),
|
|
required=key in required,
|
|
))
|
|
return ToolSpec(
|
|
name=t['name'],
|
|
description=t['description'],
|
|
parameters=params,
|
|
)
|
|
else:
|
|
return t
|
|
|
|
if 'tools' in options:
|
|
options = {**options}
|
|
options['tools'] = [
|
|
normalize_tool(t) for t in options['tools']
|
|
]
|
|
|
|
control = await self.get_agent_control(agent_id)
|
|
async for msg in control.prompt(instruction, input, **options):
|
|
yield msg
|
|
|
|
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={}
|
|
)
|