Task Management Assistant

Let’s create a Task Management Assistant with AIIDE. This assistant will help you manage your tasks, by creating, updating, and deleting them. Each task also has a custom schedule associated with it.

Prerequisites

Before we start, make sure you have pandas and aiide installed:

pip install pandas aiide

Imports

Let’s start by importing the necessary modules:

from aiide import AIIDE, TOOL
from aiide.tools import TOOL_DEF, DICT, STR
import pandas as pd

AIIDE is used to define and manage the chat instance.
Tool is used to define the tools.
TOOL_DEF, DICT and STR are some of the helpful abstractions for defining the JSON API for tool definition.
Pandas is used to maintian the tasks in a DataFrame, and aiide is used to create the assistant.

Define the Tool

Since the heart of the assistant is task management, let’s define the tool first:

The tool class inherits TOOL and has to have the following methods:

1. tool_def: This method returns the tool’s JSON API. It is called multiple times during generation. You can return different JSON API based on the context. ex: change the enums of the fields etc
2. main: This method is called when the tool is invoked. It takes the inputs defined in the JSON API and returns the output.

Let’s define the Tool

class TaskTool(TOOL):
    """
    Tool to add, edit or delete a task to the list.
    """
    def __init__(self, parent):
        """
        Initialize the tool with the parent AIIDE chat instance.
        """
        self.parent = parent
        # Initialize the environment with the current task list and timestamp
        self.ENV = [
            "Current task list:\n"
            + self.parent.task_df.to_markdown(index=False)
            + "\nThe current timestamp is: "
            + str(pd.Timestamp.now())
        ]

Here we are taking the parent instance(chat instance which will be defined below) as an argument and storing it in the self.parent attribute. It’s better to provide all the tools with the chat instance as the parent, so that they can access the environment and functions of other tools if and when needed.

Environment

Technically, the environment ENV is a list of strings that are associated with each tool.
We can use the environment to provide the LLM with the last couple of snapshots of any data the tool is acting upon.
This way, we won’t overload the llm with the entire history of the data, which might only have slight modifications.
It is also good for cost savings.

Note

Using environment ENV is purely optional. You can still build powerful applications with great DevX without defining any environments.

In this tutorial, we are only going to provide the current task list and the current timestamp as the single(latest) element in the environment.

Tool Definition

tool_def expects OpenAI’s JSON API for the tool definition. We have created a helper function TOOL_DEF to make it easier to define the JSON API.

class TaskTool(TOOL):
    ...
    def tool_def(self):
        """
        Tool definition for the task_management tool.
        """
        return TOOL_DEF(
            name="task_management",
            description="Add, edit or delete a task to the list. task_description and cron_schedule_expression are mandatory for add or edit operation but are optional for delete operation.",
            properties=[
                STR(name="task_id", description="The name of the task. Must be unique"),
                STR(
                    name="operation_type",
                    enums=["add", "edit", "delete"],
                    description="The operation to perform on the task",
                ),
                STR(
                    name="task_description",
                    description="The description of the task. Must be in markdown",
                ),
                STR(
                    name="cron_schedule_expression",
                    description="The cron schedule expression for the task. Format is 'minute hour day month day_of_week'",
                ),
            ],
            required=["task_id", "operation_type"],
        )

Here we are defining the tool with the name task_management and providing a description of what the tool does.
We are defining the properties of the tool:

1. task_id: The name of the task. It must be unique.
2. operation_type: The operation to perform on the task. It can be either add, edit, or delete.
3. task_description: The description of the task. It must be in markdown.
4. cron_schedule_expression: The cron schedule expression for the task. The format is ‘minute hour day month day_of_week’.

We are also defining the required properties of the tool: task_id and operation_type.

Main Function

The main function is the function that is called when the tool is invoked. It takes the inputs defined in the JSON API, performs the operation on the dataframe, updates the ENV and returns the output.

class TaskTool(TOOL):
    ...
        def main(
        self,
        task_id,
        operation_type,
        task_description=None,
        cron_schedule_expression=None,
    ):
        """
        Tool invokation function for the modify_item tool.
        """
        if operation_type == "delete":
            if task_id in self.parent.task_df["item_id"].values:
                self.parent.task_df = self.parent.task_df[
                    self.parent.task_df["item_id"] != task_id
                ]
                self.ENV = [
                    "Current task list:\n"
                    + self.parent.task_df.to_markdown(index=False)
                    + "\nThe current timestamp is: "
                    + str(pd.Timestamp.now())
                ]
                return "Task " + task_id + " has been deleted."
            else:
                return "Task " + task_id + " does not exist in the list."
        else:
            if task_id in self.parent.task_df["item_id"].values:
                self.parent.task_df.loc[
                    self.parent.task_df["item_id"] == task_id, "item_details"
                ] = task_description
                self.parent.task_df.loc[
                    self.parent.task_df["item_id"] == task_id, "cron"
                ] = cron_schedule_expression
            else:
                self.parent.task_df.loc[len(self.parent.task_df.index)] = [
                    task_id,
                    task_description,
                    cron_schedule_expression,
                ]
            self.ENV = [
                "Current task list:\n"
                + self.parent.task_df.to_markdown(index=False)
                + "\nThe current timestamp is: "
                + str(pd.Timestamp.now())
            ]
            return "Task " + task_id + " has been added/edited."

Here we are checking the operation_type and performing the operation on the dataframe accordingly.

Define the Chat Instance

Now that we have defined the tool, let’s define the chat instance:

class Agent(AIIDE):
    """
    AIIDE chat instance for task management.
    """
    def __init__(self, task_df):
        self.task_df = task_df
        # Create an instance of the ItemTool
        self.task_tool = TaskTool(self)
        self.setup(
            messages=[
                {
                    "role": "system",
                    "content": "You are a helpful assistant. You are to assist the user in managing their tasks.",
                }
            ],
            model="gpt-4o",
        )

Here we are defining the chat instance Agent which inherits AIIDE.
We are taking the task dataframe as an argument and storing it in the self.task_df attribute.
We are creating an instance of the TaskTool and storing it in the self.task_tool attribute.
We are setting up the chat instance with a system message and the model type.

Tip

You can use any LLM from the following model providers

Main Loop

Now that we have defined the chat instance, let’s run the main loop:

# Creating a task dataframe with some initial tasks
task_df = pd.DataFrame(
    data={
        "item_id": ["Buy Groceries", "Work on AIIDE"],
        "item_details": ["1. Bread\n2.Veggies\n3.Fruits", "Write unit tests"],
        "cron": ["0 10 * * 6", "0 20 * * 0,2,4,6"],
    }
)
agent = Agent(task_df)

Here we are creating a task dataframe with some initial tasks.
We are creating an instance of the Agent and passing the task dataframe as an argument.

Let’s create a loop to take user input and chat with the assistant:

while True:
    user_message = input("User: ")
    for each in agent.chat(user_message, tools=[agent.task_tool]):
        # Check the type of message and print the response accordingly
        if each["type"] == "text":
            print(each["delta"], end="")
        if each["type"] == "tool_call":
            print(each["name"], each["arguments"])
        if each["type"] == "tool_response":
            print(each["response"])
            print(agent.task_df)
    print("\n")

We call the chat method of the agent instance with the user message and the tools we wish to enable for that interaction.
We are iterating over the response and printing the response accordingly.
AIIDE only supports streaming currently and we iterate over it to get the responses and parse them based on type of response.

1. For response type of text: The LLM is talking to the user, which has delta and content as attributes.
2. For response type of tool_call: The LLM is calling a tool, which has name and arguments as attributes.
3. For response type of tool_response: Response from the tool, which has response as an attribute along with name and arguments.

Complete code for reference

from aiide import AIIDE, TOOL
from aiide.tools import DICT, STR, TOOL_DEF
import pandas as pd

class ItemTool(TOOL):
    """
    Tool to add, edit or delete a task to the list.
    """
    def __init__(self, parent):
        """
        Initialize the tool with the parent AIIDE chat instance.
        """
        self.parent = parent
        # Initialize the environment with the current task list and timestamp
        self.ENV = [
            "Current task list:\n"
            + self.parent.task_df.to_markdown(index=False)
            + "\nThe current timestamp is: "
            + str(pd.Timestamp.now())
        ]

    def tool_def(self):
        """
        Tool definition for the modify_item tool.
        """
        return TOOL_DEF(
            name="modify_item",
            description="Add, edit or delete a task to the list. task_description and cron_schedule_expression are mandatory for add or edit operation but are optional for delete operation.",
            properties=[
                STR(name="task_id", description="The name of the task. Must be unique"),
                STR(
                    name="operation_type",
                    enums=["add", "edit", "delete"],
                    description="The operation to perform on the task",
                ),
                STR(
                    name="task_description",
                    description="The description of the task. Must be in markdown",
                ),
                STR(
                    name="cron_schedule_expression",
                    description="The cron schedule expression for the task. Format is 'minute hour day month day_of_week'",
                ),
            ],
            required=["task_id", "operation_type"],
        )

    def main(
        self,
        task_id,
        operation_type,
        task_description=None,
        cron_schedule_expression=None,
    ):
        """
        Tool invokation function for the modify_item tool.
        """
        if operation_type == "delete":
            if task_id in self.parent.task_df["item_id"].values:
                self.parent.task_df = self.parent.task_df[
                    self.parent.task_df["item_id"] != task_id
                ]
                self.ENV = [
                    "Current task list:\n"
                    + self.parent.task_df.to_markdown(index=False)
                    + "\nThe current timestamp is: "
                    + str(pd.Timestamp.now())
                ]
                return "Task " + task_id + " has been deleted."
            else:
                return "Task " + task_id + " does not exist in the list."
        else:
            if task_id in self.parent.task_df["item_id"].values:
                self.parent.task_df.loc[
                    self.parent.task_df["item_id"] == task_id, "item_details"
                ] = task_description
                self.parent.task_df.loc[
                    self.parent.task_df["item_id"] == task_id, "cron"
                ] = cron_schedule_expression
            else:
                self.parent.task_df.loc[len(self.parent.task_df.index)] = [
                    task_id,
                    task_description,
                    cron_schedule_expression,
                ]
            self.ENV = [
                "Current task list:\n"
                + self.parent.task_df.to_markdown(index=False)
                + "\nThe current timestamp is: "
                + str(pd.Timestamp.now())
            ]
            return "Task " + task_id + " has been added/edited."


class Agent(AIIDE):
    """
    AIIDE chat instance for task management.
    """
    def __init__(self, task_df):
        self.task_df = task_df
        # Create an instance of the ItemTool
        self.item_tool = ItemTool(self)
        self.setup(
            messages=[
                {
                    "role": "system",
                    "content": "You are a helpful assistant. You are to assist the user in managing their tasks.",
                }
            ],
            model="gpt-4o",
        )

# Creating a task dataframe with some initial tasks
task_df = pd.DataFrame(
    data={
        "item_id": ["Buy Groceries", "Work on AIIDE"],
        "item_details": ["1. Bread\n2.Veggies\n3.Fruits", "Write unit tests"],
        "cron": ["0 10 * * 6", "0 20 * * 0,2,4,6"],
    }
)
agent = Agent(task_df)
while True:
    user_message = input("User: ")
    for each in agent.chat(user_message, tools=[agent.item_tool]):
        # Check the type of message and print the response accordingly
        if each["type"] == "text":
            print(each["delta"], end="")
        if each["type"] == "tool_call":
            print(each["name"], each["arguments"])
        if each["type"] == "tool_response":
            print(each["response"])
            print(agent.task_df)
    print("\n")

Conclusion

In this tutorial, we created a Task Management Assistant with AIIDE. We defined a tool to add, edit, or delete tasks to the list and created a chat instance to manage the tasks. We also created a loop to take user input and chat with the assistant.