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