Develop with Native SDK

🚀 Develop and customize new agents

This guide will walk you through creating and running your own agents for AIOS with the native SDK.

Set up the Base Agent Class

First, in the entry.py in your agent folder, create your agent class by inheriting from BaseAgent:

from cerebrum.agents.base import BaseAgent
from cerebrum.llm.communication import LLMQuery
import json

class DemoAgent(BaseAgent):
    def __init__(self, agent_name, task_input, config_):
        super().__init__(agent_name, task_input, config_)

        self.plan_max_fail_times = 3
        self.tool_call_max_fail_times = 3

        self.start_time = None
        self.end_time = None
        self.request_waiting_times: list = []
        self.request_turnaround_times: list = []
        self.task_input = task_input
        self.messages = []
        self.workflow_mode = "manual"  # (manual, automatic)
        self.rounds = 0

Import Query Functions

AIOS provides several Query classes for different types of interactions and use the Response class in here to receive results from the AIOS kernel.

Query Class	Arguments	Output
LLMQuery	messages: List, tools: List, action_type: str, message_return_type: str	response: Response
MemoryQuery	TBD	response: Response
StorageQuery	TBD	response: Response
ToolQuery	tool_calls: List	response: Response

Here's how to import a specific query

from cerebrum.llm.communication import LLMQuery  # Using LLMQuery as an example

Construct system instructions

Here's how to set up your agent's system instructions and you need to put this function inside your agent class

def build_system_instruction(self):
    prefix = "".join(["".join(self.config["description"])])

    plan_instruction = "".join(
        [
            f"You are given the available tools from the tool list: {json.dumps(self.tool_info)} to help you solve problems. ",
            "Generate a plan with comprehensive yet minimal steps to fulfill the task. ",
            "The plan must follow the json format as below: ",
            "[",
            '{"action_type": "action_type_value", "action": "action_value","tool_use": [tool_name1, tool_name2,...]}',
            '{"action_type": "action_type_value", "action": "action_value", "tool_use": [tool_name1, tool_name2,...]}',
            "...",
            "]",
            "In each step of the planned plan, identify tools to use and recognize no tool is necessary. ",
            "Followings are some plan examples. ",
            "[" "[",
            '{"action_type": "tool_use", "action": "gather information from arxiv. ", "tool_use": ["arxiv"]},',
            '{"action_type": "chat", "action": "write a summarization based on the gathered information. ", "tool_use": []}',
            "];",
            "[",
            '{"action_type": "tool_use", "action": "gather information from arxiv. ", "tool_use": ["arxiv"]},',
            '{"action_type": "chat", "action": "understand the current methods and propose ideas that can improve ", "tool_use": []}',
            "]",
            "]",
        ]
    )

    if self.workflow_mode == "manual":
        self.messages.append({"role": "system", "content": prefix})

    else:
        assert self.workflow_mode == "automatic"
        self.messages.append({"role": "system", "content": prefix})
        self.messages.append({"role": "user", "content": plan_instruction})

Create Workflows

You can create a workflow for the agent to execute its task and you need to put this function inside your agent class.

Manual workflow example:

def manual_workflow(self):
    workflow = [
        {
            "action_type": "tool_use",
            "action": "Search for relevant papers",
            "tool_use": ["demo_author/arxiv"],
        },
        {
            "action_type": "chat",
            "action": "Provide responses based on the user's query",
            "tool_use": [],
        },
    ]
    return workflow

Implement the Run Method

Finally, implement the run method to execute your agent's workflow and you need to put this function inside your agent class.

def run(self):
    self.build_system_instruction()

    task_input = self.task_input

    self.messages.append({"role": "user", "content": task_input})

    workflow = None

    if self.workflow_mode == "automatic":
        workflow = self.automatic_workflow()
        self.messages = self.messages[:1]  # clear long context

    else:
        assert self.workflow_mode == "manual"
        workflow = self.manual_workflow()

    self.messages.append(
        {
            "role": "user",
            "content": f"[Thinking]: The workflow generated for the problem is {json.dumps(workflow)}. Follow the workflow to solve the problem step by step. ",
        }
    )

    try:
        if workflow:
            final_result = ""

            for i, step in enumerate(workflow):
                action_type = step["action_type"]
                action = step["action"]
                tool_use = step["tool_use"]

                prompt = f"At step {i + 1}, you need to: {action}. "
                self.messages.append({"role": "user", "content": prompt})

                if tool_use:
                    selected_tools = self.pre_select_tools(tool_use)

                else:
                    selected_tools = None

                response = self.send_request(
                    agent_name=self.agent_name,
                    query=LLMQuery(
                        messages=self.messages,
                        tools=selected_tools,
                        action_type=action_type,
                    ),
                )["response"]
                
                self.messages.append({"role": "assistant", "content": response.response_message})

                self.rounds += 1


            final_result = self.messages[-1]["content"]
            
            return {
                "agent_name": self.agent_name,
                "result": final_result,
                "rounds": self.rounds,
            }

        else:
            return {
                "agent_name": self.agent_name,
                "result": "Failed to generate a valid workflow in the given times.",
                "rounds": self.rounds,
            }
            
    except Exception as e:
        return {}

Run the Agent

To test your agent, use the run-agent command:

run-agent \
--llm_name <llm_name> \
--llm_backend <llm_backend> \
--agent_name_or_path <agent_name_or_path> \
--local_agent <yes/no> \
--task <task_input> \
--aios_kernel_url <aios_kernel_url>

Parameters

• <llm_name>: The name of the language model you want to use

• <llm_backend>: The backend service for the language model

• <agent_name_or_path>:

  • For hub agents: The name of the agent in the format of author/agent

  • For local agents: Absolute path to the agent folder

• <local_agent>: Use local agent? Default is no (yes/no)

• <task_input>: The task you want your agent to complete

• <aios_kernel_url>: The URL that is connected to the AIOS kernel

Examples

Run a hub agent:

run-agent \
--llm_name gpt-4 \
--llm_backend openai \
--agent_name_or_path demo_author/demo_agent \
--local_agent no \
--task "Your task here" \
--aios_kernel_url http://localhost:8000

Run a local agent:

run-agent \
--llm_name gpt-4 \
--llm_backend openai \
--agent_name_or_path /absolute/path/to/local/agent \
--local_agent yes \
--task "Your task here" \
--aios_kernel_url http://localhost:8000

Using the source code:

You can also run an agent using the source code of the run-agent command, which iscerebrum/example/run_agent:

python cerebrum/example/run_agent.py \
--llm_name <llm_name> \
--llm_backend <llm_backend> \
--agent_name_or_path <agent_name_or_path> \
--local_agent <yes/no> \
--task <task_input> \
--aios_kernel_url <aios_kernel_url>

🔧Develop and Customize New Tools

Tool Structure

Similar as developing new agents, developing tools also need to follow a simple directory structure:

demo_author/
└── demo_tool/
    │── entry.py      # Contains your tool's main logic
    └── config.json   # Tool configuration and metadata

Setting up config.json

Your tool needs a configuration file that describes its properties. Here's an example of how to set it up:

{
    "name": "demo_tool",
    "description": [
        "The arxiv tool that can be used to search for papers on arxiv"
    ],
    "meta": {
        "author": "demo_author",
        "version": "1.0.6",
        "license": "CC0"
    },
    "build": {
        "entry": "tool.py",
        "module": "DemoTool"
    }
}

Create Tool Class

In entry.py, you'll need to implement a tool class which is identified in the config.json with two essential methods:

1. get_tool_call_format: Defines how LLMs should interact with your tool

2. run: Contains your tool's main functionality

Here's an example:

class Arxiv:
    def get_tool_call_format(self):
        tool_call_format = {
            "type": "function",
            "function": {
                "name": "demo_author/arxiv",
                "description": "Query articles or topics in arxiv",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "query": {
                            "type": "string",
                            "description": "Input query that describes what to search in arxiv"
                        }
                    },
                    "required": [
                        "query"
                    ]
                }
            }
        }
        return tool_call_format

    def run(self, params: dict):
        """
        Main tool logic goes here.
        Args:
            params: Dictionary containing tool parameters
        Returns:
            Your tool's output
        """
        # Your code here
        result = do_something(params['param_name'])
        return result

Integration Tips

When integrating your tool for the agents you develop:

• Use absolute paths to reference your tool in agent configurations

• Example: /path/to/your/tools/example/your_tool instead of just author/tool_name