Implement basic agent and toolset functionality

backend/agent.py

@@ -1,22 +1,104 @@
 from stream import Stream
+import re
+import json
 import threading
 
+from tools import Toolset, Action, BasicToolset, SpecialToolset
+
 LLAMA3_MODEL = "meta/meta-llama-3-8b-instruct"
 
-# Instances of this class lives very briefly, created when needed and immediately disposed.
-# Subclasses must be careful in not setting up long term resources
-class AgentControl:
+# Tool related classes must be moved to a different module
+
+
+# Instances are created and destroyed constantly, agents must not depend solely on in-memory data
+# to be useful and consistent
+class Agent:
+    # while this constructor accepts two parameters,
+    # all subclasses must also accept a fourth 'data' parameter
     def __init__(self, app, agent_id):
         self.app = app
         self.agent_id = agent_id
-        self.system_prompt = "You are a useful and precise AI assistant. Respond concisely."
+
+        # Basic implementation, will be moved to a subclass
+        self.system_prompt = "\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.",
+        ])
+
+        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 = {}
+
+        print("tool prompt section", self.tool_prompt_section)
+
+        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.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_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 replicate_llama_completion(self.app, [
-            LlamaMessage("system", self.system_prompt),
-            LlamaMessage("user", content)
-        ])
+        response = await self._full_prompt(f"User Message", content) 
         return response
 
     async def received_message(self, message):
@@ -27,35 +109,131 @@ class AgentControl:
             return "assistant" if is_digit else sender
         chat_summary = "\n".join(f"**{sender_name(message)}**: {message['content']}" for message in chat_context.messages)
 
-        prologue = "\nYou are given the entire context of the conversation, respond with your next message in the conversation."
-        print(chat_summary)
-
-        response = await replicate_llama_completion(self.app, [
-            LlamaMessage("system", self.system_prompt + prologue),
-            LlamaMessage("user", chat_summary)
-        ])
+        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.agent_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_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
     
     def create_from_data(app, agent_id, kind, data):
-        return AgentControl(app, agent_id)
+        return Agent(app, agent_id)
+
+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
+        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
+
 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)
@@ -68,10 +246,8 @@ async def replicate_llama_completion(app, messages):
             'content_stream': content_stream,
         }
     )
-
     thread.start()
     response = await content_stream.content_wait()
-    
     return response.lstrip()