319 lines
12 KiB
Python
319 lines
12 KiB
Python
import re
|
|
import json
|
|
import inspect
|
|
|
|
from tools import Toolset, Action, BasicToolset, SpecialToolset
|
|
from inference import LlamaMessage, replicate_llama_completion
|
|
import gral
|
|
|
|
# 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)
|
|
# tools: list[ToolDescription] (custom actions supported by the class)
|
|
|
|
def __init__(self, app, data):
|
|
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', [])}
|
|
|
|
def _set_basic_info(self, id, name):
|
|
self.id = id
|
|
self.name = name
|
|
|
|
# Receives a single string prompt, and must return a string response
|
|
async def prompt(self, content):
|
|
raise Exception("'prompt' is not implemented for this agent class")
|
|
|
|
# Receives a chat message object, agent can send messages to the chat in response
|
|
async def received_message(self, message):
|
|
raise Exception("'received_message' is not implemented for this agent class")
|
|
|
|
# Executes a tool/action in this agent
|
|
async def execute_action(self, name: str, argument: dict):
|
|
# TODO: For security, check that the method is in the tool list
|
|
method = getattr(self, name)
|
|
result = method(**argument)
|
|
if inspect.isawaitable(result):
|
|
result = await result
|
|
return result
|
|
|
|
class ToolUsingAgent:
|
|
identifier = "tool-using"
|
|
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.",
|
|
])
|
|
),
|
|
]
|
|
|
|
# tools: list[ToolDescription] (actions the agent class support)
|
|
tools = []
|
|
|
|
def __init__(self, app, data):
|
|
Agent.__init__(self, app, data)
|
|
|
|
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 = {}
|
|
|
|
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_prompt_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.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_prompt_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
|
|
|
|
class Llama8bAgent(Agent):
|
|
identifier = "llama-8b"
|
|
schema_fields = [
|
|
AgentField('system_prompt', 'string', default="You are a helpful AI assistant."),
|
|
]
|
|
tools = []
|
|
|
|
def __init__(self, app, data):
|
|
Agent.__init__(self, app, data)
|
|
|
|
async def prompt(self, content):
|
|
messages = [
|
|
LlamaMessage("system", self.config['system_prompt']),
|
|
LlamaMessage("user", content)
|
|
]
|
|
response = await replicate_llama_completion(self.app, messages)
|
|
return response
|
|
|
|
async def received_message(self, message):
|
|
def sender_name(message):
|
|
sender = message['sender']
|
|
is_digit = isinstance(sender, (int, str)) and str(sender).isdigit()
|
|
is_self = is_digit and str(sender) == str(self.id)
|
|
return self.name if is_self else sender
|
|
|
|
chat_context = self.app.get_chat_context(message.chat_id)
|
|
chat_summary = "\n".join(f"[{sender_name(message)}]: {message['content']}" for message in chat_context.messages)
|
|
messages = [
|
|
LlamaMessage("system", f"{self.config['system_prompt']}\nYou are given the context of the conversation, respond with your next message in the conversation."),
|
|
LlamaMessage("user", f"Recent Messages:\n{chat_summary}")
|
|
]
|
|
|
|
response = await replicate_llama_completion(self.app, messages)
|
|
self.app.send_chat_message(
|
|
chat_id=message.chat_id,
|
|
sender=self.id,
|
|
content=response,
|
|
)
|
|
|
|
class BasicAgent(Llama8bAgent):
|
|
identifier = "basic"
|
|
schema_fields = Llama8bAgent.schema_fields
|
|
tools = Llama8bAgent.tools
|
|
|
|
class CoderV1Agent(Llama8bAgent):
|
|
identifier = "coder-v1"
|
|
schema_fields = [
|
|
AgentField('system_prompt', 'string', default="You are an expert Software Engineer profficient in multiple programming languages."),
|
|
]
|
|
# TODO: add gral functions to the tool list
|
|
tools = []
|
|
|
|
async def gral_message(self, message_history: list, file_map: dict):
|
|
return await gral.gral_message(self.app, self.config['system_prompt'], message_history, file_map)
|
|
|
|
async def gral_update_file(self, message_history: list, filepath: str, file_content: str):
|
|
return await gral.gral_update_file(self.app, self.config['system_prompt'], message_history, filepath, file_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
|