392 lines
13 KiB
Python
392 lines
13 KiB
Python
import re
|
|
import inspect
|
|
import time
|
|
import asyncio
|
|
import traceback
|
|
|
|
from temporary import generate_tmp_id
|
|
from stream import Stream
|
|
|
|
# Tool related classes must be moved to a different module
|
|
|
|
class AgentField:
|
|
def __init__(self, name, tp, default=None):
|
|
self.name=name
|
|
self.type=tp
|
|
self.default=default
|
|
|
|
# Instances are created and destroyed constantly, agents must not depend solely on in-memory data
|
|
# to be useful and consistent
|
|
class Agent:
|
|
# Class variables, must be defined by subclasses of Agent:
|
|
# - identifier (for kind in creation)
|
|
# - schema: list[AgentField{name, type: str}] (for configuration)
|
|
|
|
def __init__(self, app, data, resources=[]):
|
|
self.app = app
|
|
|
|
# These variables are set by '_set_basic_info'
|
|
self.id = None
|
|
self.name = None
|
|
|
|
# Auto populated based on the schema fields
|
|
agent_class = type(self)
|
|
self.config = {field.name: data.get(field.name, field.default) for field in getattr(agent_class, 'schema_fields', [])}
|
|
self.resources = resources
|
|
|
|
# should be in expert init?
|
|
self.tools = []
|
|
self.remotes = []
|
|
for resource in self.resources:
|
|
if resource.kind == 'tool':
|
|
if resource.value.get('type') == 'native' and 'name' in resource.value:
|
|
tool_cls = next((t for t in self.app.tool_classes if t.tool_name() == resource.value['name']), None)
|
|
if tool_cls:
|
|
tool = tool_cls()
|
|
spec = tool.get_spec().clone()
|
|
spec.name = resource.key
|
|
self.tools.append(spec)
|
|
elif resource.kind == 'remote':
|
|
value = resource.value or {}
|
|
if isinstance(value, str):
|
|
value = {'prefix': value}
|
|
self.remotes.append({
|
|
'name': resource.key,
|
|
'prefix': value.get('prefix'),
|
|
})
|
|
|
|
|
|
def _set_basic_info(self, id, name, user_id):
|
|
self.id = id
|
|
self.name = name
|
|
self.user_id = user_id
|
|
|
|
async def run(self, agent_run):
|
|
"""Core method. Subclasses override this to handle the AgentRun."""
|
|
agent_run.close()
|
|
raise NotImplementedError("Subclasses must implement run()")
|
|
|
|
async def prompt(self, instruction, input=None, **kwargs):
|
|
run = AgentRun.create(
|
|
self.app,
|
|
agent_id=self.id,
|
|
user_id=self.user_id,
|
|
input={
|
|
"instruction": instruction,
|
|
"input": input,
|
|
**kwargs,
|
|
},
|
|
)
|
|
out_stream = Stream(self.app)
|
|
|
|
def on_output_event(stream):
|
|
def on_chunk(chunk):
|
|
out_stream.feed(chunk)
|
|
stream.listen(on_chunk, lambda: None)
|
|
def on_done(data):
|
|
out_stream.close()
|
|
run.on('output', on_output_event)
|
|
run.on('done', on_done)
|
|
|
|
# Change the output shape based on the presence of tools
|
|
if 'tools' in kwargs:
|
|
tool_calls = []
|
|
tool_names = [t['function']['name'] for t in kwargs['tools'] if isinstance(t, dict)]
|
|
def on_tool_call(call):
|
|
print('prompt call', call)
|
|
if call.tool_name in tool_names:
|
|
tool_calls.append(call)
|
|
run.on('tool-call', on_tool_call)
|
|
await self.run(run)
|
|
|
|
# Cannot stream tool calls out, wait for them
|
|
await out_stream.content_wait()
|
|
print('tool names' + repr(tool_names))
|
|
print('tool calls' + repr(tool_calls))
|
|
return {
|
|
'output': out_stream,
|
|
'tool_calls': tool_calls,
|
|
}
|
|
else:
|
|
await self.run(run)
|
|
return out_stream
|
|
|
|
|
|
async def received_message(self, message):
|
|
chat = await self.app.get_chat(message.chat_id)
|
|
thread = chat.get_thread(message.thread_id)
|
|
thread.skip_list += ['log', 'resource']
|
|
await thread.load_page()
|
|
|
|
messages = []
|
|
for msg in thread.messages:
|
|
role = self._role_for_message(msg)
|
|
messages.append({"role": role, "content": msg.content})
|
|
|
|
remotes = []
|
|
def add_remote(name):
|
|
if name not in remotes and (self.user_id, name) in self.app.remotes:
|
|
remotes.append(name)
|
|
|
|
remote_messages = await thread.fetch_kind_page('resource', 'remote')
|
|
for msg in remote_messages:
|
|
add_remote(msg.content)
|
|
|
|
run = AgentRun.create(
|
|
self.app,
|
|
agent_id=self.id,
|
|
user_id=self.user_id,
|
|
input={
|
|
"messages": messages,
|
|
"chat_id": message.chat_id,
|
|
"thread_id": message.thread_id,
|
|
"remotes": remotes,
|
|
},
|
|
)
|
|
|
|
async def on_inference(msg):
|
|
kwargs = {'content': msg.content}
|
|
|
|
if msg.role == 'assistant':
|
|
# kind and subkind are default
|
|
pass
|
|
elif msg.role == 'reasoning':
|
|
kwargs['kind'] = "log"
|
|
kwargs['subkind'] = "reasoning"
|
|
elif msg.role == 'tool':
|
|
kwargs['kind'] = "log"
|
|
kwargs['subkind'] = "tool-call"
|
|
kwargs['content'] = repr(msg.content)
|
|
else:
|
|
# Unknown message
|
|
kwargs['kind'] = "log"
|
|
|
|
new_message = await self.app.send_chat_message(
|
|
sender=self.id,
|
|
chat_id=message.chat_id,
|
|
thread_id=message.thread_id,
|
|
**kwargs
|
|
)
|
|
run.emit_event('chat-message', {
|
|
'messageId': new_message['id'],
|
|
'chatId': message.chat_id,
|
|
'threadId': message.thread_id,
|
|
})
|
|
|
|
await run.on_async('inference', on_inference)
|
|
|
|
await self.run(run)
|
|
return run
|
|
|
|
def _role_for_message(self, message):
|
|
sender = message.sender
|
|
is_digit = isinstance(sender, (int, str)) and str(sender).isdigit()
|
|
if is_digit and str(sender) == str(self.id):
|
|
return 'assistant'
|
|
elif sender == 'user':
|
|
return 'user'
|
|
elif sender.startswith("tmp-"):
|
|
return 'assistant'
|
|
else:
|
|
return 'user'
|
|
|
|
# Returns all resource values matching the kind, and the key when given
|
|
def get_resources(self, kind, key=None):
|
|
return [r for r in self.resources
|
|
if r.kind == kind and (key is None or r.key == key)]
|
|
|
|
def get_tools(self, run=None):
|
|
remotes = list(self.remotes)
|
|
if run is not None and 'remotes' in run.input:
|
|
remotes += list(run.input['remotes'])
|
|
|
|
tools = list(self.tools) + self.app.getRemoteTools(self.user_id, remotes)
|
|
|
|
# Dedupe by tool name (agent-configured vs run-input vs native collisions)
|
|
seen, result = set(), []
|
|
for tool in tools:
|
|
if tool.name in seen:
|
|
continue
|
|
seen.add(tool.name)
|
|
result.append(tool)
|
|
|
|
return result
|
|
|
|
|
|
class DebugAgent(Agent):
|
|
identifier = "debug"
|
|
schema_fields = []
|
|
tools = []
|
|
|
|
async def received_message(self, message):
|
|
chat_view = await self.app.get_chat_view(message.chat_id)
|
|
response = chat_view.debug_string()
|
|
|
|
await self.app.send_chat_message(
|
|
chat_id=message.chat_id,
|
|
sender=self.id,
|
|
content=response,
|
|
thread_id=chat_view.current_thread_id,
|
|
)
|
|
|
|
# Holds parameters, outputs, and details of an agent call
|
|
class AgentRun:
|
|
@classmethod
|
|
def create(cls, app, agent_id, user_id, input=None):
|
|
self = cls()
|
|
self.app = app
|
|
self.id = generate_tmp_id(app)
|
|
self.agent_id = agent_id
|
|
self.user_id = user_id
|
|
self.done = False
|
|
self.start_time = int(time.time())
|
|
self.end_time = None
|
|
|
|
# I/O envelope (no generic data bag)
|
|
self.input = input or {}
|
|
self.events = []
|
|
self.event_listeners = {}
|
|
|
|
# Default Streams
|
|
self.output = None
|
|
self.reasoning = None
|
|
|
|
# Cancellation tracking
|
|
# You should use 'cancelled', task and thread are for
|
|
# external tasks that cannot be cancelled reliably
|
|
self.task = None # For asyncio tasks
|
|
self.thread = None
|
|
self.cancelled = False
|
|
|
|
app.tmp.runs[self.id] = self
|
|
|
|
# Emit run start event
|
|
self.emit_event('created', {
|
|
"agentId": self.agent_id,
|
|
"userId": self.user_id,
|
|
"startTime": self.start_time,
|
|
"input": self.input,
|
|
})
|
|
|
|
return self
|
|
|
|
def on(self, event_name, callback):
|
|
"""Register a callback for a named event.
|
|
|
|
Replays every past event of this type to the callback immediately
|
|
(with only the data payload), so late listeners catch up on history.
|
|
"""
|
|
if event_name not in self.event_listeners:
|
|
self.event_listeners[event_name] = []
|
|
self.event_listeners[event_name].append(callback)
|
|
|
|
# Replay: feed only the data payload of past events
|
|
for event in self.events:
|
|
if event["name"] == event_name:
|
|
callback(event["data"])
|
|
|
|
async def on_async(self, event_name, callback):
|
|
event_loop = asyncio.get_event_loop()
|
|
def new_callback(data):
|
|
event_loop.create_task(callback(data))
|
|
self.on(event_name, new_callback)
|
|
|
|
def emit_event(self, event_name, data=None):
|
|
"""Record an event and notify existing listeners.
|
|
|
|
Listeners receive only the *data* payload, not the full envelope.
|
|
"""
|
|
event = {"name": event_name, "data": data, "time": time.time()}
|
|
self.events.append(event)
|
|
listeners = self.event_listeners.get(event_name, [])
|
|
for cb in listeners:
|
|
try:
|
|
cb(data)
|
|
except Exception as e:
|
|
print(f"ERROR: handling event '{event_name}' in run {self.id}\n" + str(e))
|
|
print(traceback.format_exc())
|
|
|
|
# Broadcast event to websocket listeners
|
|
self.app.emit_event("run-event", {
|
|
"runId": self.id,
|
|
"name": event_name,
|
|
"data": data
|
|
})
|
|
|
|
def add_tools(self, tools):
|
|
if "tools" in self.input:
|
|
self.input["tools"] += tools
|
|
else:
|
|
self.input["tools"] = tools
|
|
|
|
def add_tool_call(self, tool_name, arguments):
|
|
self.emit_event("tool_call", {"name": tool_name, "arguments": arguments})
|
|
|
|
def _serialize_input(self, input):
|
|
"""Convert ToolSpec objects to serializable dictionaries in input data."""
|
|
if not input:
|
|
return input
|
|
|
|
result = {}
|
|
for key, value in input.items():
|
|
if key == "tools" and isinstance(value, list):
|
|
# Convert ToolSpec objects to their OpenAI dictionary format
|
|
result[key] = [tool.to_openai_dict() if hasattr(tool, 'to_openai_dict')
|
|
else tool for tool in value]
|
|
else:
|
|
result[key] = value
|
|
return result
|
|
|
|
def serialize(self):
|
|
"""Serialize the run object to a dictionary for JSON responses."""
|
|
return {
|
|
"id": self.id,
|
|
"agentId": self.agent_id,
|
|
"userId": self.user_id,
|
|
"done": self.done,
|
|
"startTime": self.start_time,
|
|
"endTime": self.end_time,
|
|
"input": self._serialize_input(self.input),
|
|
"events": self.events,
|
|
"output": self.output,
|
|
}
|
|
|
|
def cancel(self):
|
|
"""Cancel the agent run and stop any ongoing inference."""
|
|
if self.cancelled or self.done:
|
|
return
|
|
|
|
self.cancelled = True
|
|
print(f"Cancelling agent run {self.id}")
|
|
|
|
# Cancel asyncio task if exists
|
|
if self.task and not self.task.done():
|
|
self.task.cancel()
|
|
|
|
# Thread cannot be forcefully cancelled, but we can signal
|
|
if self.thread and self.thread.is_alive():
|
|
# The thread should be checking self.cancelled in its loop
|
|
print(f"Signalled thread to stop for run {self.id}")
|
|
|
|
# Close the stream and mark as done
|
|
self.close()
|
|
|
|
# Emit cancellation event
|
|
self.emit_event("cancelled", {})
|
|
|
|
# Broadcast cancellation event
|
|
self.app.emit_event("run-cancelled", {
|
|
"runId": self.id,
|
|
"agentId": self.agent_id,
|
|
"endTime": int(time.time()),
|
|
})
|
|
|
|
def close(self):
|
|
if self.done:
|
|
return
|
|
self.done = True
|
|
self.end_time = int(time.time())
|
|
if self.output: self.output.close()
|
|
if self.reasoning: self.reasoning.close()
|
|
self.emit_event("done", {})
|