Agent actions and gral function. Need agent classes to separate advanced agent functionality.

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()

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 = {

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

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()

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:

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):

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
+- 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,
+  })

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);
 }