49 lines
1.6 KiB
Python
49 lines
1.6 KiB
Python
import inference
|
|
from agent import Agent, AgentField
|
|
|
|
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, app, data):
|
|
Agent.__init__(self, app, data)
|
|
|
|
async def run(self, run):
|
|
opts = {}
|
|
|
|
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"]
|
|
|
|
# 1. 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"),
|
|
)
|
|
|
|
# 2. Call inference, passing the run to feed run.stream and tool calls
|
|
await inference.chat_completion(
|
|
self.app,
|
|
model_id=self.config.get("model"),
|
|
messages=messages,
|
|
tools=run.input.get("tools", []),
|
|
opts=opts,
|
|
agent_run=run,
|
|
)
|
|
|
|
run.stream.listen(lambda chunk: None, lambda: run.close())
|