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

backend/app.py

@@ -9,9 +9,9 @@ from sqlalchemy.sql.expression import select
 from quart import Quart, Response, request, send_from_directory, jsonify, websocket
 import replicate
 
-from database import create_sql_engine, User, Agent, Chat, Message, Membership
+import database as db
 from server import start_server
-from agent import AgentControl
+from agent import Agent
 
 MESSAGE_PAGE_SIZE = 32
 
@@ -23,7 +23,7 @@ class ChatContext:
 class App:
     def __init__(self, data_directory, replicate_api_token=None):
         self.data_directory = data_directory
-        self.sql_engine = create_sql_engine(self.data_directory)
+        self.sql_engine = db.create_sql_engine(self.data_directory)
         self.user_connections = {}
         self.streams = {}
 
@@ -50,7 +50,7 @@ class App:
 
     def create_agent(self, user_id, name, description, picture, kind, data):
         with self.sql_session() as session:
-            agent = Agent()
+            agent = db.Agent()
             agent.user_id = user_id
             agent.name = name
             agent.description = description
@@ -65,7 +65,7 @@ class App:
 
     def create_chat(self, user_id, name, description, picture):
         with self.sql_session() as session:
-            chat = Chat()
+            chat = db.Chat()
             chat.user_id = user_id
             chat.name = name
             chat.description = description
@@ -78,7 +78,7 @@ class App:
 
     def get_chats(self, user_id):
         with self.sql_session() as session:
-            chats = session.execute(select(Chat).where(Chat.user_id == user_id)).scalars().all()
+            chats = session.execute(select(db.Chat).where(db.Chat.user_id == user_id)).scalars().all()
             return [{
                 "id": chat.id,
                 "name": chat.name,
@@ -88,7 +88,7 @@ class App:
 
     def get_agents(self, user_id):
         with self.sql_session() as session:
-            agents = session.execute(select(Agent).where(Agent.user_id == user_id)).scalars().all()
+            agents = session.execute(select(db.Agent).where(db.Agent.user_id == user_id)).scalars().all()
             return [{
                 "id": agent.id,
                 "name": agent.name,
@@ -99,7 +99,7 @@ class App:
     def get_chat_messages(self, chat_id, earliest_index=0):
         with self.sql_session() as session:
             messages = session.execute(
-                select(Message).where(Message.chat_id == chat_id).order_by(Message.time.asc())
+                select(db.Message).where(db.Message.chat_id == chat_id).order_by(db.Message.time.asc())
             ).scalars().fetchmany(MESSAGE_PAGE_SIZE)
 
             return [{
@@ -113,7 +113,7 @@ class App:
     def verify_user_credentials(self, username, password):
         with self.sql_session() as session:
             user = session.execute(
-                select(User).where(User.username == username)
+                select(db.User).where(db.User.username == username)
             ).scalars().first()
 
             if user and user.password == password:
@@ -123,11 +123,11 @@ class App:
     def create_new_user(self, username, email, password):
         with self.sql_session() as session:
             result = session.execute(
-                select(User.id).where(User.email == email or User.username == username)
+                select(db.User.id).where(db.User.email == email or db.User.username == username)
             ).first()
 
             if result is None:
-                user = User()
+                user = db.User()
                 user.email = email
                 user.username = username
                 user.password = password
@@ -143,7 +143,7 @@ class App:
 
     def add_agent_to_chat(self, chat_id, agent_id):
         with self.sql_session() as session:
-            membership = Membership()
+            membership = db.Membership()
             membership.chat_id = chat_id
             membership.agent_id = agent_id
             
@@ -154,7 +154,7 @@ class App:
     # If an integer is passed in sender, it is interpreted as an agent id
     def send_chat_message(self, chat_id, content, sender='user'):
         with self.sql_session() as session:
-            message = Message()
+            message = db.Message()
             message.chat_id = chat_id
             message.time = int(datetime.now().timestamp())
             message.sender = sender
@@ -166,11 +166,11 @@ class App:
 
             # Notify chat owner
             if sender != 'user':
-                chat_owner = session.execute(select(Chat.user_id).where(Chat.id == chat_id)).scalar()
+                chat_owner = session.execute(select(db.Chat.user_id).where(db.Chat.id == chat_id)).scalar()
                 self.notify_message_to_user(chat_owner, message)
 
             # Notify participant agents
-            agent_ids = session.execute(select(Membership.agent_id).where(Membership.chat_id == chat_id)).scalars().all()
+            agent_ids = session.execute(select(db.Membership.agent_id).where(db.Membership.chat_id == chat_id)).scalars().all()
             for agent_id in agent_ids:
                 if agent_id != sender:
                     self.notify_message_to_agent(agent_id, message)
@@ -190,13 +190,13 @@ class App:
     def get_agent_control(self, agent_id):
         with self.sql_session() as session:
             agent = session.execute(
-                select(Agent).where(Agent.id == agent_id)
+                select(db.Agent).where(db.Agent.id == agent_id)
             ).scalars().first()
 
             if not agent:
                 raise Exception(f"Agent not found with id {agent_id}")
         
-            control = AgentControl.create_from_data(app, agent_id, agent.kind, agent.data)
+            control = Agent.create_from_data(app, agent_id, agent.kind, agent.data)
             return control
 
     async def send_agent_prompt(self, agent_id, message):

backend/tools.py

@@ -0,0 +1,100 @@
+
+import inspect
+
+# Do not call methods from instances of this class, they must not get overridden
+class Toolset:
+    # Fields: description, identifier, _tools: list[ToolDescription]
+    def __init__(self):
+        # TODO: receive app, assign _tools
+        pass
+
+    def get_namespace_definitions(toolset):
+        header = toolset.description
+        if toolset.identifier:
+            header = f"Namespace {toolset.identifier}: {toolset.description}"
+        lines = [header]
+ 
+        for tool in toolset._tools:
+            full_name = tool.identifier
+            if toolset.identifier:
+                full_name = toolset.identifier + "." + full_name
+            lines.append(f"- {tool.description}\n!!{full_name} {tool.schema}")
+        return "\n".join(lines)
+    
+    async def execute_action(toolset, name: str, argument: dict):
+        method = getattr(toolset, name)
+        result = method(**argument)
+        if inspect.isawaitable(result):
+            result = await result
+
+class ToolDescription:
+    def __init__(self, identifier, description, schema):
+        self.identifier=identifier
+        self.description=description
+        self.schema=schema
+
+# An action is an object representing a command invokation
+class Action:
+    def __init__(self, toolset_name: str, tool_name: str, args: dict, is_query: bool = False):
+        self.toolset_name = toolset_name
+        self.tool_name = tool_name
+        self.args = args
+        self.is_query = is_query
+    
+    def is_special(self, name: str):
+        return self.toolset_name == "" and self.tool_name == name
+
+    def __str__(self):
+        return f"{self.is_query and 'query' or ''}{self.toolset_name}.{self.tool_name} {self.args}"
+    
+    def __repr__(self):
+        return f"Action({self.__str__()})"
+
+class SpecialToolset(Toolset):
+    def __init__(self):
+        self.identifier = ""
+        self.description = "Special commands that don't receive arguments"
+        self._tools = [
+            ToolDescription("resume", "Resume the conversation with the command outputs", ''),
+            ToolDescription("respond", "When used, all text after this command will be used as the actual response", ''),
+        ]
+
+    def resume(self):
+        pass
+
+    def respond(self):
+        pass
+
+class BasicToolset(Toolset):
+    def __init__(self):
+        self.identifier = "util"
+        self.description = "Basic assistant utility tools"
+        self._tools = [
+            ToolDescription("google_search", "Executes a google search", '{"query":string}'),
+            ToolDescription("fetch", "Fetch the contents of a url", '{"url":string}'),
+            ToolDescription("eval", "Evaluates a python expression", '{"expression":string}'),
+        ]
+
+    def google_search(self, query):
+        return "We're sorry. Google is not accessible at this time"
+
+    def fetch(self, url):
+        return "No internet connection."
+    
+    def eval(self, expression):
+        try:
+            return eval(expression)
+        except Exception as e:
+            return str(e)
+
+class HubControlsToolset(Toolset):
+    def __init__(self):
+        self.identifier = "hub"
+        self.description = "Controls for the Hub system"
+        self._tools = [
+            ToolDescription("get_chats", "Get all the existing chats", '{}'),
+            ToolDescription("get_agents", "Get all the existing bots", '{}'),
+            ToolDescription("create_chat", "Creates a new chat", '{"name":string,"agent_ids":int[]}'),
+            ToolDescription("get_agent_data", "Gets the configuration data of an agent", '{"agent_id":int}'),
+            ToolDescription("set_agent_data", "Get all the existing bots", '{"agent_id":int,"data":object}'),
+        ]