Build a Basic AI Agent from Scratch: Long Task Planning

In the previous part of the Build A Basic AI Agent From Scratch series, we added the essential tools to our agent to allow it to work autonomously for us. We gave it the ability to find files, read and write files, run bash commands and get content from the web. We got a very capable agent with just these tools.

What happens when the agent runs long and complex tasks?

The current agent works very well, but we want our agent to get a lot of work done, and this requires staying on the task for long spans of time. Right now, if we try to give our agent long and complex tasks we will find that it does not think long term, and it stops working after the littlest progress.

This is to be expected because the LLM is trained to behave conversationally. It expects to go back and forth in a question-answer basis. This is fine for a simple chatbot, but our agent needs to be able to get a request and work for a long time on it before returning a result.

Long task planning

The next ability we will give to our agent is the ability to plan for long and complex tasks.

The abilities our agent needs are:

• Understand the goal of the task

• Plan how to tackle the task beforehand

• Break the task into concrete steps

• Keep track of pending, in progress and completed tasks

• If something goes wrong with the current plan, rethink the approach

• Check that everything planned is actually done before stopping

To give our agent these abilities, we will rely on the last part's addition: tools. We will also explain the model how to use long task planning in the model's system prompt.

New tool: Scratchpad

This is a very simple but powerful tool. We are just giving the model a place to write it's thoughts and read them again at a later time.

The main benefit of this tool is that it forces the model to think through the goal and plan the whole approach before starting working on it.

The tool saves the scratchpad content into memory instead of a file or database, which is fine because we don't want to share the scratchpad content between sessions.

Here's the python implementation:

You can find and clone this code in this blog series' Github repo.

New tool: To-do list

A to-do list allows the agent to decompose the work into tasks and keep track of them to know what's left to do (pending), what it's working on currently (in progress) and what is already done (done).

This tool also enforces some good practices: it doesn't allow multiple tasks to be in progress at the same time, it doesn't allow invalid task statuses and it doesn't allow repeated tasks.

Just like the scratchpad, this tool saves the to do list into memory instead of a file or database. This is also fine because we don't want to share the to-do list between agent sessions.

New system prompt

All the strategies for long term task planning that cannot be implemented into tools are explained to the model in the system prompt. Here we will explain to the model how to plan using the process explained in the beginning of the article, and also how to use the new tools to help it in the planning process.

For more details, read the system prompt below.

I also added to the system prompt a little comment explaining to the model that if not stated otherwise, the project it has to work on is in the current directory.

Let's test it!

To test our new and powerful agent, we will have to give it a really hard goal. In my case, I asked it to migrate my static site from using Eleventy to Hugo:

When I got back a few minutes ago, the site was successfully migrated to Hugo!

What You've Built

This agent is now capable enough not just to work on it's own, but also to work on really hard and long tasks. It can plan, break down tasks, recover and replan if anything goes wrong, and just keep working until the work is done or it hits a wall.

What's next?

An agent capable of working for a long time to complete a task is a great thing for productivity. But, if you dispatch this agent and go do something else, it might be editing files and running commands it isn't supposed to without your knowledge. You can't just trust your agent to always behave how you are expecting it to, so you want the agent to ping you before doing something potentially catastrophic. You want to be the human in the loop.