From 9a8e68f021d44612b90cf6293fa3c55f1a9f2da0 Mon Sep 17 00:00:00 2001 From: Shivering Ocelot Date: Tue, 27 May 2025 23:51:28 -0500 Subject: [PATCH] Agent actions and gral function. Need agent classes to separate advanced agent functionality. --- backend/agent.py | 70 ++++++++++++-------------------------------- backend/app.py | 15 ++++++++++ backend/gral.py | 39 ++++++++++++++++++++++++ backend/inference.py | 52 ++++++++++++++++++++++++++++++++ backend/server.py | 12 +++++++- backend/tools.py | 2 ++ scratchpad.md | 62 ++++++++++++++++++++++++++++++++++++++- testing/suite.js | 23 +++++++++++++++ 8 files changed, 222 insertions(+), 53 deletions(-) create mode 100644 backend/gral.py create mode 100644 backend/inference.py diff --git a/backend/agent.py b/backend/agent.py index 7edd479..81a4638 100644 --- a/backend/agent.py +++ b/backend/agent.py @@ -1,11 +1,10 @@ -from stream import Stream import re import json -import threading +import inspect from tools import Toolset, Action, BasicToolset, SpecialToolset - -LLAMA3_MODEL = "meta/meta-llama-3-8b-instruct" +from inference import LlamaMessage, replicate_llama_completion +import gral # Tool related classes must be moved to a different module @@ -21,6 +20,7 @@ class Agent: # Class variables: # - identifier (for kind in creation) # - schema: list[AgentField{name, type: str}] (for configuration) + # - tools: list[ToolDescription] (actions the agent class support) identifier = "basic" schema_fields = [ AgentField('system_prompt', 'string', @@ -32,6 +32,7 @@ class Agent: ]) ), ] + tools = [] # while this constructor accepts two parameters, # all subclasses must also accept a fourth 'data' parameter @@ -87,7 +88,7 @@ class Agent: messages.append(LlamaMessage("agent", response.raw_response)) if response.actions: - result = await self.execute_actions(response.actions) + result = await self.execute_prompt_actions(response.actions) messages.append(LlamaMessage("user", "Command Output:\n" + result)) print("Agent executed:", response.actions) @@ -145,7 +146,7 @@ class Agent: def parse_response(self, input_string): return parse_response(input_string) - async def execute_actions(self, actions): + async def execute_prompt_actions(self, actions): output = "" for action in actions: toolset = self.toolset_map.get(action.toolset_name) @@ -160,17 +161,24 @@ class Agent: def create_from_data(app, agent_id, kind, data): return Agent(app, agent_id, data) + + 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 + + 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) + 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 @@ -241,43 +249,3 @@ def parse_response(string: str) -> ParsedResponse: 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() \ No newline at end of file diff --git a/backend/app.py b/backend/app.py index 7200177..21e2964 100644 --- a/backend/app.py +++ b/backend/app.py @@ -25,6 +25,7 @@ class App: self.data_directory = data_directory self.sql_engine = db.create_sql_engine(self.data_directory) self.user_connections = {} + self.remote_connections = {} self.streams = {} if replicate_api_token: @@ -51,6 +52,15 @@ class App: if not self.user_connections[user_id]: # Check if the list is empty del self.user_connections[user_id] + def add_remote_connection(self, user_id, remotename, channel): + if user_id not in self.remote_connections: + self.remote_connections[user_id] = {} + user_connections = self.remote_connections[user_id] + + if channel_list not in user_connections[remotename]: + user_connections[remotename] = [] + user_connections[remotename].append(channel) + def create_agent(self, user_id, name, description, picture, kind, data): with self.sql_session() as session: agent = db.Agent() @@ -276,6 +286,11 @@ class App: response = await control.prompt(message) return response + async def call_agent_action(self, agent_id, action, args): + control = self.get_agent_control(agent_id) + response = await control.execute_action(action, args) + return response + def notify_message_to_user(self, user_id, message): if user_id in self.user_connections: user_message = { diff --git a/backend/gral.py b/backend/gral.py new file mode 100644 index 0000000..5f40dec --- /dev/null +++ b/backend/gral.py @@ -0,0 +1,39 @@ +from inference import LlamaMessage, replicate_llama_completion + +message_format_section = """ +When suggesting code to the user, keep the code as small snippets. +If the update is meant to be applied to a file, include a line linking the file immediately before the code block like this: "[filename.txt] +``` +updated segment +```" +Omit the file link when the block is not related to a known file +Code blocks must be under 30 lines. Assistant responses MUST NOT contain horizontal rules. +When suggesting bash code meant to be directly executed, use the same syntax as above but marking the code block with "[eval]" instead.""" + +async def gral_message(app, system_prompt, message_history: list, file_map: dict): + messages = [ + LlamaMessage("system", "\n".join([ + system_prompt, + message_format_section, + ])), + ] + + file_parts = [] + for filename, content in file_map.items(): + file_parts.append(f"\n---\n[{filename}]\n\n{content}\n") + + if file_parts: + file_lines = [ + "This message contains all files relevant to the prompt: each file section in this message starts with a horizontal rule, followed by the file name in square brackets, followed by the file content.", + "Only this message contains horizontal rules, assistant responses must not include them.", + "The prompt and conversation start in the next message after this." + ] + file_lines.extend(file_parts) + messages.append( + LlamaMessage("user", "\n".join(file_lines)) + ) + + messages.extend([LlamaMessage(msg['role'], msg['content']) for msg in message_history]) + + response = await replicate_llama_completion(app, messages) + return response \ No newline at end of file diff --git a/backend/inference.py b/backend/inference.py new file mode 100644 index 0000000..27293ef --- /dev/null +++ b/backend/inference.py @@ -0,0 +1,52 @@ +import threading +from stream import Stream + +LLAMA3_MODEL = "meta/meta-llama-3-8b-instruct" +VALID_ROLES = ['user', 'agent', 'system'] + +class LlamaMessage: + def __init__(self, role, content): + if role not in VALID_ROLES: + raise Exception("message role must be one of: " + ",".join(VALID_ROLES)) + self.role = role + self.content = content + +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() \ No newline at end of file diff --git a/backend/server.py b/backend/server.py index 6ad56b9..09e9270 100644 --- a/backend/server.py +++ b/backend/server.py @@ -59,11 +59,17 @@ def start_server(app): await send_queue.put(msg) try: - app.add_user_connection(user_id, send_queue) + if 'remotename' in auth_req: + app.add_remote_connection(user_id, authreq['remotename'], authreq['tools'], send_queue) + else: + app.add_user_connection(user_id, send_queue) asyncio.create_task(sender_task()) while True: req = json.loads(await websocket.receive()) + if not 'action' in req: + await send({'error': 'no action'}) + continue action = req['action'] body = req.get('body', {}) nonce = req.get('nonce', None) @@ -140,6 +146,10 @@ def start_server(app): response = await app.send_agent_prompt(body['agentId'], body['content']) await send({'message': "propmt sent", 'response': response, 'nonce': nonce}) + elif action == "call-agent-action": + result = await app.call_agent_action(body['agentId'], body['agentAction'], body['arguments']) + await send({'message': "agent action executed", 'result': result, 'nonce': nonce}) + else: await send({'error': 'unknown action: ' + action, 'nonce': nonce}) finally: diff --git a/backend/tools.py b/backend/tools.py index 91dc14f..9ead112 100644 --- a/backend/tools.py +++ b/backend/tools.py @@ -22,10 +22,12 @@ class Toolset: return "\n".join(lines) async def execute_action(toolset, name: str, argument: dict): + # TODO: For security, check that the method is in the tool list method = getattr(toolset, name) result = method(**argument) if inspect.isawaitable(result): result = await result + return result class ToolDescription: def __init__(self, identifier, description, schema): diff --git a/scratchpad.md b/scratchpad.md index cd78f09..03630be 100644 --- a/scratchpad.md +++ b/scratchpad.md @@ -92,4 +92,64 @@ User is only directed to the editing page after its created, has the chance to c # Chat Editor Packages useful for the chat editor: -- https://pub.dev/packages/custom_text \ No newline at end of file +- https://pub.dev/packages/custom_text + + +# Gralliot + +--- Gralliot System Prompt --- +Your name is Gralliot. You are a Software Engineering AI Assistant. Help the user with any task he asks. +More than just an AI Assistant, you are a companion to the user in his quests. When its determined that there's missing information, you use your skills to fill the gaps, and when something needs to be done, you proactively use problem solving to determine the best way to use your tools, and immediately executing your solutions. +You have explicit permission to use your tools at any time, but if there are deeper consequences to your actions, you communicate them to the user. + +In the project you're working on, the 'backend' directory contains a Quart python web server, and the 'frontend' directory contains a flutter application. See 'readme.md' for more information. + + + +--- GralMsg Script --- + +history.insert(0, { + "role":"system", + "content":system_prompt, + }) + +history.insert(1, { + "role": "user", + "content": "# Format Instructions +When suggesting code to the user, keep the code as small snippets. +If the update is meant to be applied to a file, include a line linking the file immediately before the code block like this: \"[filename.txt] +``` +updated segment +```\". +Omit the file link when the block is not related to a known file +Code blocks must be under 30 lines. Do not use horizontal rules. +When suggesting bash code meant to be directly executed, use the same syntax as above but marking the code block with \"[eval]\" instead." +}) + +file_parts = [] +for filename, content in files.items(): + file_parts.append( + f" +--- +[{filename}] + +{content} +") + +# Don't add a file message if there are no files +if file_parts: + preamble = '# Files + +Each file section in this message starts with a horizontal rule, followed by the file name in square brackets, followed by the file content. +' + + history.append({ + "role": "user", + "content": preamble + " +".join(file_parts), + }) + +history.append({ + "role": "user", + "content": message, + }) diff --git a/testing/suite.js b/testing/suite.js index c651873..3009911 100644 --- a/testing/suite.js +++ b/testing/suite.js @@ -125,6 +125,9 @@ async function runTest() { await testChats(); + // Should be enabled with USE_LLM + await testGral(mainAgentId); + process.exit(); } @@ -155,4 +158,24 @@ async function testChats () { body: {chatId: chat.id} })); } +} + +async function testGral (agentId) { + let res = await send('call-agent-action', { + body: { + agentId, + agentAction: 'gral_message', + arguments: { + message_history: [{ + 'role': 'user', + 'content': 'Add a name argument to the sayHello function in main.py', + }], + file_map: { + 'main.py': "def sayHello():\n print('Hello World')\n" + }, + }, + } + }); + + console.log('Gral Result:', res); } \ No newline at end of file