97 lines
3.5 KiB
Python
97 lines
3.5 KiB
Python
import inference
|
|
from agent import Agent, AgentField
|
|
|
|
import asyncio
|
|
|
|
CLASS_DESCRIPTION = """
|
|
Expert AI with customizable system prompt.
|
|
Will get guidelines, documents, and other advanced features in the future.
|
|
""".strip()
|
|
|
|
class ExpertAgent(Agent):
|
|
identifier = "expert"
|
|
description = CLASS_DESCRIPTION
|
|
schema_fields = [
|
|
AgentField('system_prompt', 'string', default="You are an expert AI assistant."),
|
|
AgentField('model', 'model_id', default=None),
|
|
AgentField('temperature', 'number', default=None),
|
|
AgentField('repetition_penalty', 'number', default=None),
|
|
]
|
|
|
|
def __init__(self, *args):
|
|
Agent.__init__(self, *args)
|
|
|
|
async def run(self, run):
|
|
opts = {}
|
|
|
|
# Is this roon using *Agent's* tools, as opposed to run tools
|
|
using_agent_tools = False
|
|
if 'tools' not in run.input:
|
|
# Agent tools will fully replace run tools if none are present,
|
|
# otherwise they won't be used
|
|
if self.tools:
|
|
run.input['tools'] = self.tools
|
|
using_agent_tools = True
|
|
|
|
if self.config.get("temperature") is not None:
|
|
opts["temperature"] = self.config["temperature"]
|
|
if self.config.get("repetition_penalty") is not None:
|
|
opts["repetition_penalty"] = self.config["repetition_penalty"]
|
|
|
|
# Build message list based on input
|
|
messages = inference.build_messages(
|
|
system_prompt=self.config.get("system_prompt", ""),
|
|
instruction=run.input.get("instruction"),
|
|
input=run.input.get("input"),
|
|
messages=run.input.get("messages"),
|
|
)
|
|
|
|
# Collect inference output. These are resused across calls
|
|
tool_calls = []
|
|
output_stream = None
|
|
inference_end = asyncio.Event()
|
|
def on_output(stream):
|
|
nonlocal output_stream
|
|
output_stream = stream
|
|
run.on('output', on_output)
|
|
run.on('inference-end', lambda event: inference_end.set())
|
|
run.on('tool-call', lambda tc: print(f"[expert] tool call: {tc}") or tool_calls.append(tc))
|
|
|
|
while True:
|
|
# Reset data
|
|
inference_end.clear()
|
|
output_stream = None
|
|
|
|
# Call inference, passing the run to get events out
|
|
await inference.chat_completion(
|
|
self.app,
|
|
model_id=self.config.get("model"),
|
|
messages=messages,
|
|
tools=run.input.get("tools", []),
|
|
opts=opts,
|
|
agent_run=run,
|
|
)
|
|
|
|
await inference_end.wait()
|
|
|
|
messages.append(inference.LlamaMessage('assistant', output_stream and await output_stream.content_wait() or ""))
|
|
print(f"[expert] collected tool calls: {len(tool_calls)}, using? {using_agent_tools}")
|
|
|
|
if using_agent_tools and tool_calls:
|
|
print("Will execute tool calls")
|
|
for tc in tool_calls:
|
|
tool = next(t for t in self.tools if t.name == tc.tool_name)
|
|
if not tool:
|
|
print(f"Bad agent tool name: {tc.name}")
|
|
continue
|
|
tc.result = await tool.execute_async(**tc.arguments)
|
|
messages.append(inference.LlamaMessage('tool', tc))
|
|
tool_calls = []
|
|
# Buffer time to let openrouter rest?
|
|
await asyncio.sleep(2)
|
|
else:
|
|
print("No tool calls to execute, finish")
|
|
break
|
|
|
|
run.close()
|