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