98 lines
3.5 KiB
Python
98 lines
3.5 KiB
Python
import re
|
|
import json
|
|
import inspect
|
|
from string import Template
|
|
|
|
import inference
|
|
from inference import LlamaMessage
|
|
from agent import Agent, AgentField
|
|
|
|
CLASS_DESCRIPTION = """
|
|
Expert AI in any field, powered by Llama 3 70b.
|
|
Will get guidelines, documents, and other advanced features in the future.
|
|
|
|
Available models are: [llama/qwen]-[big/small] (replicate), ollama:MODEL_NAME
|
|
""".strip()
|
|
|
|
class ExpertAgent(Agent):
|
|
identifier = "expert"
|
|
description = CLASS_DESCRIPTION
|
|
schema_fields = [
|
|
AgentField('system_prompt', 'string', default="You are an expert AI assistant."),
|
|
AgentField('model', 'string', default="qwen-big"),
|
|
]
|
|
|
|
def __init__(self, app, data):
|
|
Agent.__init__(self, app, data)
|
|
|
|
async def prompt(self, instruction, input=None, **kwargs):
|
|
# Instruction extends system if there's input or messages
|
|
# Otherwise instruction is the user prompt
|
|
# Input is always added as a user message at the end, if present
|
|
|
|
# Use 'messages' as starting point when passed
|
|
messages = []
|
|
if 'messages' in kwargs:
|
|
messages = [LlamaMessage(msg['role'], msg['content']) for msg in kwargs['messages']]
|
|
|
|
# Construct the default messages
|
|
system = self.config['system_prompt']
|
|
prompt = input
|
|
|
|
if input == None and not messages:
|
|
prompt = instruction
|
|
elif instruction:
|
|
system += "\n" + instruction
|
|
|
|
# Add default messages to list
|
|
messages.insert(0, LlamaMessage("system", system))
|
|
if prompt:
|
|
messages.append(LlamaMessage("user", prompt))
|
|
|
|
model = self.config['model']
|
|
|
|
if model.startswith("ollama:"):
|
|
response = await inference.ollama_chat_completion(self.app, messages, model=model[7:])
|
|
elif model == "qwen-big":
|
|
response = await inference.replicate_qwen_completion(self.app, messages, model=inference.Models.qwen3_235b)
|
|
elif model == "qwen-small":
|
|
response = await inference.replicate_qwen_completion(self.app, messages, model=inference.Models.qwen3_8b)
|
|
elif model == "llama-big":
|
|
response = await inference.replicate_llama_completion(self.app, messages, model=inference.Models.llama_70b)
|
|
elif model == "llama-small":
|
|
response = await inference.replicate_llama_completion(self.app, messages, model=inference.Models.llama_8b)
|
|
else:
|
|
# An error here is questionable, model is user configurable
|
|
raise Exception("Unkonwn agent model: " + model)
|
|
|
|
return response
|
|
|
|
async def received_message(self, message):
|
|
chat_view = await self.app.get_chat_view(message.chat_id)
|
|
thread = chat_view.get_thread(message.thread_id)
|
|
|
|
def msg_role(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'
|
|
else:
|
|
raise Exception("Unsupported sender: " + sender)
|
|
|
|
messages = [{'role': msg_role(msg), 'content': msg.content} for msg in thread.messages]
|
|
|
|
response = await self.prompt(
|
|
instruction=None,
|
|
input=None,
|
|
messages=messages
|
|
)
|
|
|
|
await self.app.send_chat_message(
|
|
chat_id=message.chat_id,
|
|
sender=self.id,
|
|
content=response,
|
|
thread_id=message.thread_id,
|
|
) |