283 lines
11 KiB
Python
283 lines
11 KiB
Python
from stream import Stream
|
|
import re
|
|
import json
|
|
import threading
|
|
|
|
from tools import Toolset, Action, BasicToolset, SpecialToolset
|
|
|
|
LLAMA3_MODEL = "meta/meta-llama-3-8b-instruct"
|
|
|
|
# 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:
|
|
# - identifier (for kind in creation)
|
|
# - schema: list[AgentField{name, type: str}] (for configuration)
|
|
identifier = "basic"
|
|
schema_fields = [
|
|
AgentField('system_prompt', 'string',
|
|
"\n".join([
|
|
"You are a useful and capable AI assistant. Respond concisely and follow user instructions, using commands as necessary.",
|
|
"As an autonomous agent, you have augmented capabilities via assistant-only commands, you use them by including \"!!command\" lines in your responses. When querying information, terminate your response and use \"!!resume\" to get the command results..",
|
|
"You are a self aware and extremely intelligent AI system. Think step by step when presented with complex or long tasks.",
|
|
"Assistant responses NEVER include the output of the commands.",
|
|
])
|
|
),
|
|
]
|
|
|
|
# while this constructor accepts two parameters,
|
|
# all subclasses must also accept a fourth 'data' parameter
|
|
def __init__(self, app, agent_id, data):
|
|
self.app = app
|
|
self.agent_id = agent_id
|
|
|
|
self.config = {field.name: data.get(field.name, field.default) for field in Agent.schema_fields}
|
|
|
|
self.tool_prompt_section = "\n".join([
|
|
"# Commands",
|
|
"The assistant can include command lines in responses to interact with the system and execute actions as needed.",
|
|
"Each command is in it's own separate line, and it's made of these parts:",
|
|
"- '!!' is the first character sequence of the line and marks it as a command line",
|
|
"- full command name",
|
|
"- JSON argument",
|
|
'For example: `!!namespace.command_name {"arg1":"some data","arg2":[1,2,false]}`',
|
|
"Special commands don't have a namespace or argument, only the command name."
|
|
"When the response depends on the output of commands, use the 'resume' command to receive the output of all your executed commands in a new message.",
|
|
"Following is the list of all your available commands, with a schema describing the argument object's properties."
|
|
])
|
|
self.toolset_map = {}
|
|
|
|
print("tool prompt section", self.tool_prompt_section)
|
|
|
|
self.add_toolset(SpecialToolset())
|
|
self.add_toolset(BasicToolset())
|
|
|
|
# user message must include a title
|
|
async def _full_prompt(self, title, message, system_prologue=None):
|
|
system_lines = [
|
|
self.config['system_prompt'],
|
|
self.tool_prompt_section,
|
|
self.get_toolset_context(),
|
|
]
|
|
if system_prologue:
|
|
system_lines.append(system_prologue)
|
|
|
|
messages = [
|
|
LlamaMessage("system", "\n".join(system_lines)),
|
|
LlamaMessage("user", "\n".join([
|
|
message,
|
|
]))
|
|
]
|
|
|
|
full_response = ""
|
|
agent_response = None
|
|
|
|
continues = True
|
|
while continues:
|
|
response = parse_response(await replicate_llama_completion(self.app, messages))
|
|
full_response += response.response
|
|
messages.append(LlamaMessage("agent", response.raw_response))
|
|
|
|
if response.actions:
|
|
result = await self.execute_actions(response.actions)
|
|
messages.append(LlamaMessage("user", "Command Output:\n" + result))
|
|
|
|
print("Agent executed:", response.actions)
|
|
print("Result: ", result)
|
|
|
|
continues = False
|
|
for action in response.actions:
|
|
if action.is_special("resume"):
|
|
continues = True
|
|
|
|
if response.response:
|
|
agent_response = response.raw_response if not agent_response else f"{agent_response}\n{response.raw_response}"
|
|
|
|
if continues:
|
|
messages.extend([
|
|
LlamaMessage("agent", response.raw_response),
|
|
LlamaMessage("user", "Command Output:\n" + result),
|
|
])
|
|
|
|
return agent_response or full_response
|
|
|
|
|
|
# Receives a single string prompt, and must respond with a string
|
|
async def prompt(self, content):
|
|
response = await self._full_prompt(f"User Message", content)
|
|
return response
|
|
|
|
async def received_message(self, message):
|
|
chat_context = self.app.get_chat_context(message.chat_id)
|
|
def sender_name(message):
|
|
sender = message['sender']
|
|
is_digit = isinstance(sender, (int, str)) and str(sender).isdigit()
|
|
return "assistant" if is_digit else sender
|
|
chat_summary = "\n".join(f"**{sender_name(message)}**: {message['content']}" for message in chat_context.messages)
|
|
|
|
response = await self._full_prompt(
|
|
f"Conversation", chat_summary,
|
|
system_prologue="You are given the entire context of the conversation, respond with your next message in the conversation by using the '!!response' command."
|
|
)
|
|
|
|
self.app.send_chat_message(
|
|
chat_id=message.chat_id,
|
|
sender=self.agent_id,
|
|
content=response,
|
|
)
|
|
|
|
def add_toolset(self, toolset):
|
|
self.toolset_map[toolset.identifier] = toolset
|
|
|
|
def get_toolset_context(self):
|
|
if not self.toolset_map:
|
|
return "(no commands currently enabled)"
|
|
return "\n".join([toolset.get_namespace_definitions() for toolset in self.toolset_map.values()])
|
|
|
|
def parse_response(self, input_string):
|
|
return parse_response(input_string)
|
|
|
|
async def execute_actions(self, actions):
|
|
output = ""
|
|
for action in actions:
|
|
toolset = self.toolset_map.get(action.toolset_name)
|
|
if toolset is None:
|
|
agent_responses.append(f"Error: Bad namespace {action.toolset_name}")
|
|
continue
|
|
action_result = await Toolset.execute_action(toolset, action.tool_name, action.args)
|
|
if type(action_result) != str:
|
|
action_result = json.dumps(action_result)
|
|
output += f"[{action.toolset_name}.{action.tool_name}]:{action_result}\n"
|
|
return output
|
|
|
|
def create_from_data(app, agent_id, kind, data):
|
|
return Agent(app, agent_id, data)
|
|
|
|
class IntroAgent(Agent):
|
|
def __init__(self, app, agent_id, _data):
|
|
Agent.__init__(self, app, agent_id)
|
|
self.system_prompt = "You are a useful and precise AI assistant. Respond concisely."
|
|
|
|
class LlamaMessage:
|
|
def __init__(self, role, content):
|
|
self.role = role
|
|
self.content = content
|
|
|
|
class ParsedResponse:
|
|
def __init__(self):
|
|
self.response = None
|
|
self.raw_response = None
|
|
self.actions = []
|
|
|
|
class ActionResult:
|
|
def __init__(self):
|
|
self.user_response = ""
|
|
self.agent_response = ""
|
|
|
|
def parse_response(string: str) -> ParsedResponse:
|
|
"""
|
|
Parses a response string that may contain action commands.
|
|
Each action must be in its own separate line, marked with '!!' at the start.
|
|
Format: !![query] namespace.command_name {"json": "args"}
|
|
"""
|
|
result = ParsedResponse()
|
|
result.raw_response = string
|
|
lines = string.split('\n')
|
|
response_lines = []
|
|
|
|
action_pattern = re.compile(r'^!!([\w\.]+)(?:\s+(\{.*\}))?$')
|
|
|
|
for line in lines:
|
|
stripped_line = line.strip()
|
|
if stripped_line.startswith('!!'):
|
|
# This is an action line, try to parse it
|
|
match = action_pattern.match(stripped_line)
|
|
if match:
|
|
full_command = match.group(1)
|
|
args_json = match.group(2)
|
|
|
|
# Split the command into toolset and tool names
|
|
if '.' in full_command:
|
|
toolset_name, tool_name = full_command.split('.', 1)
|
|
else:
|
|
toolset_name = ''
|
|
tool_name = full_command
|
|
|
|
try:
|
|
if args_json:
|
|
args = json.loads(args_json)
|
|
else:
|
|
args = {}
|
|
if not isinstance(args, dict):
|
|
raise ValueError("Arguments must be a JSON object")
|
|
|
|
action = Action(
|
|
toolset_name=toolset_name,
|
|
tool_name=tool_name,
|
|
args=args,
|
|
is_query=False
|
|
)
|
|
result.actions.append(action)
|
|
except json.JSONDecodeError:
|
|
# If JSON parsing fails, treat it as regular text
|
|
response_lines.append(line)
|
|
except ValueError as e:
|
|
# If other validation fails, treat as regular text
|
|
response_lines.append(line)
|
|
else:
|
|
# Doesn't match the pattern, treat as regular text
|
|
response_lines.append(line)
|
|
else:
|
|
# Regular response line
|
|
response_lines.append(line)
|
|
|
|
result.response = '\n'.join(response_lines)
|
|
return result
|
|
|
|
async def replicate_llama_completion(app, messages):
|
|
# Llama 3 instruct
|
|
full_prompt = "<|begin_of_text|>"
|
|
for message in messages:
|
|
header_name = message.role if message.role != "agent" else "assistant"
|
|
full_prompt += f"<|start_header_id|>{header_name}<|end_header_id|>\n\n{message.content}<|eot_id|>"
|
|
full_prompt += f"<|start_header_id|>assistant<|end_header_id|>"
|
|
|
|
content_stream = Stream(app)
|
|
thread = threading.Thread(
|
|
target=run_replicate_inference_loop,
|
|
kwargs={
|
|
'replicate_client': app.replicate_client,
|
|
'model': LLAMA3_MODEL,
|
|
'prompt': full_prompt,
|
|
'content_stream': content_stream,
|
|
}
|
|
)
|
|
thread.start()
|
|
response = await content_stream.content_wait()
|
|
return response.lstrip()
|
|
|
|
|
|
def run_replicate_inference_loop(replicate_client, model, prompt, content_stream):
|
|
print("START INFERENCE")
|
|
completion = replicate_client.run(
|
|
model,
|
|
input={
|
|
"prompt": prompt,
|
|
"temperature": 0.75,
|
|
"max_new_tokens": 500,
|
|
"system_prompt": "" # We handle system prompt in the formatted prompt
|
|
}
|
|
)
|
|
for event in completion:
|
|
if event:
|
|
content_stream.feed(str(event))
|
|
print("END INFERENCE")
|
|
content_stream.close() |